Something went wrong. Try again.
small gleam coding and persistent agent daemon with a detachable cli and a canvas ui forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373"""Interactive pseudo-terminal lifecycle and WebSocket streaming."""
import base64import hashlibimport jsonimport secretsimport socketimport structimport timeimport unittestimport urllib.parseimport urllib.request
from harness import Albedo, exclusivefrom inspect_support import Inspection
class TerminalWebSocket: def __init__(self, app, terminal_id, timeout=20, **params): address = urllib.parse.urlsplit(app.base) self.socket = socket.create_connection( (address.hostname, address.port), timeout ) self.reader = self.socket.makefile("rb") key = base64.b64encode(secrets.token_bytes(16)).decode() query = urllib.parse.urlencode(params) path = f"/terminals/{terminal_id}" + (f"?{query}" if query else "") self.socket.sendall( ( f"GET {path} HTTP/1.1\r\nHost: {address.netloc}\r\n" "Connection: Upgrade\r\nUpgrade: websocket\r\n" f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n" f"Authorization: Bearer {app.connection['token']}\r\n\r\n" ).encode() ) try: status = self.reader.readline() if b" 101 " not in status: raise AssertionError(f"WebSocket upgrade failed: {status!r}") headers = {} while (line := self.reader.readline()) not in (b"\r\n", b""): name, value = line.decode().split(":", 1) headers[name.lower()] = value.strip() expected = base64.b64encode( hashlib.sha1( (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode() ).digest() ).decode() if headers.get("sec-websocket-accept") != expected: raise AssertionError(f"bad WebSocket accept: {headers}") except BaseException: self.close() raise
def close(self): try: self.reader.close() self.socket.close() except OSError: pass
def frame(self, opcode, payload): mask = secrets.token_bytes(4) size = len(payload) header = bytes([0x80 | opcode]) if size < 126: header += bytes([0x80 | size]) elif size < 65536: header += bytes([0x80 | 126]) + struct.pack("!H", size) else: header += bytes([0x80 | 127]) + struct.pack("!Q", size) self.socket.sendall( header + mask + bytes(value ^ mask[i % 4] for i, value in enumerate(payload)) )
def send_text(self, **message): self.frame(1, json.dumps(message).encode("utf-8"))
def send_bytes(self, data): self.frame(2, data)
def read_exact(self, size): value = self.reader.read(size) if len(value) != size: raise AssertionError("WebSocket closed before a complete frame") return value
def receive_frame(self, timeout=10): self.socket.settimeout(timeout) flags, length = self.read_exact(2) opcode = flags & 15 if length & 128: raise AssertionError("server sent a masked frame") if length == 126: length = struct.unpack("!H", self.read_exact(2))[0] elif length == 127: length = struct.unpack("!Q", self.read_exact(8))[0] payload = self.read_exact(length) if opcode == 9: self.frame(10, payload) return self.receive_frame(timeout) if opcode == 10: return self.receive_frame(timeout) return opcode, payload
class TerminalsTest(unittest.TestCase): def setUp(self): self.app = Albedo(providers={}).__enter__() self.addCleanup(self.app.__exit__, None, None, None)
def test_terminal_session_io_and_lifecycle(self): tid = "term-test-1" ws = TerminalWebSocket(self.app, tid, cols=80, rows=24, shell="/bin/sh") self.addCleanup(ws.close)
# 1. Expect initial ready frame (text) opcode, payload = ws.receive_frame() self.assertEqual(opcode, 1) ready = json.loads(payload.decode("utf-8")) self.assertEqual(ready["type"], "ready", f"terminal returned error: {ready}") self.assertEqual(ready["id"], tid) self.assertEqual(ready["cols"], 80) self.assertEqual(ready["rows"], 24) self.assertEqual(ready["state"], "running")
# 2. Send command to PTY ws.send_bytes(b"echo ALBEDO_TERM_OK\r\n")
# 3. Read output frames until our string appears output = bytearray() deadline = time.monotonic() + 10.0 while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: output.extend(chunk) if b"ALBEDO_TERM_OK" in output: break self.assertIn(b"ALBEDO_TERM_OK", output)
# 4. Test resize command ws.send_text(type="resize", cols=100, rows=30) ws.send_bytes(b"stty size\r\n") deadline = time.monotonic() + 10.0 size_output = bytearray() while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: size_output.extend(chunk) if b"30 100" in size_output: break self.assertIn(b"30 100", size_output)
# 5. Test signal (SIGINT stops running cat) ws.send_bytes(b"cat\r\n") time.sleep(0.05) ws.send_text(type="signal", signal="SIGINT") ws.send_bytes(b"echo AFTER_SIGINT\r\n") deadline = time.monotonic() + 10.0 sig_output = bytearray() while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: sig_output.extend(chunk) if b"AFTER_SIGINT" in sig_output: break self.assertIn(b"AFTER_SIGINT", sig_output) ws.close()
# 6. Reconnect to the same terminal: scrollback history should be preserved ws2 = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws2.close) opcode, payload = ws2.receive_frame() self.assertEqual(opcode, 1) ready2 = json.loads(payload.decode("utf-8")) self.assertEqual(ready2["id"], tid)
# Following frame is binary history opcode, history = ws2.receive_frame() self.assertEqual(opcode, 2) self.assertIn(b"ALBEDO_TERM_OK", history)
# 7. Test in-stream restart observer = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(observer.close) self.assertEqual(observer.receive_frame()[0], 1) self.assertEqual(observer.receive_frame()[0], 2)
ws2.send_text(type="restart", cols=80, rows=24) while True: opcode, payload = ws2.receive_frame(timeout=5) if opcode == 1: break restarted = json.loads(payload.decode("utf-8")) self.assertEqual(restarted["type"], "ready") self.assertEqual(restarted["id"], tid)
while True: opcode, payload = observer.receive_frame(timeout=5) if opcode == 1: observer_event = json.loads(payload.decode("utf-8")) break self.assertEqual(observer_event["type"], "closed") ws2.close()
# 8. Explicit delete endpoint req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200) res = json.load(resp) self.assertEqual(res["closed"], tid)
# 9. Connecting after deletion gives fresh terminal with no old history ws3 = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws3.close) opcode, payload = ws3.receive_frame() self.assertEqual(opcode, 1) ready3 = json.loads(payload.decode("utf-8")) self.assertEqual(ready3["type"], "ready") ws3.send_bytes(b"echo FRESH_AFTER_DELETE\r\n") fresh_out = bytearray() deadline = time.monotonic() + 10.0 while time.monotonic() < deadline: opcode, chunk = ws3.receive_frame(timeout=5) if opcode == 2: fresh_out.extend(chunk) if b"FRESH_AFTER_DELETE" in fresh_out: break self.assertIn(b"FRESH_AFTER_DELETE", fresh_out) self.assertNotIn(b"AFTER_SIGINT", fresh_out) ws3.close()
def test_reconnect_sends_ready_before_live_output(self): tid = "term-reconnect-order" ws = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws.close) opcode, _ = ws.receive_frame() self.assertEqual(opcode, 1)
for attempt in range(16): ws.send_bytes( b"i=0; while [ $i -lt 10000 ]; do echo RECONNECT_OUTPUT; " b"i=$((i+1)); done\r\n" ) output = bytearray() while output.count(b"RECONNECT_OUTPUT") < 2: opcode, payload = ws.receive_frame(timeout=1) if opcode == 2: output.extend(payload) ws.close() ws = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws.close) opcode, payload = ws.receive_frame() self.assertEqual( opcode, 1, f"output arrived before ready on attach {attempt}: {payload[:80]!r}", ) ready = json.loads(payload.decode("utf-8")) self.assertEqual(ready["type"], "ready")
req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200)
# exclusive: ALBEDO_INSPECT must be set before the daemon boots so the test can# read terminal mailbox lengths over erlang distribution (erlang:process_info)@exclusiveclass TerminalPressureTest(unittest.TestCase): def setUp(self): self.app = Albedo( providers={}, prepare=lambda app: app.daemon.env.update(ALBEDO_INSPECT="1"), ).__enter__() self.addCleanup(self.app.__exit__, None, None, None)
def mailbox_lengths(self): """Queue lengths of every terminal worker and the sockets watching it.""" return json.loads( Inspection(self.app).evaluate( timeout=10, expression="Length = fun(P) -> case Call(erlang, process_info, " "[P, message_queue_len]) of {message_queue_len, N} -> N; " "_ -> 0 end end, " "Monitored = fun(P) -> case Call(erlang, process_info, " "[P, monitors]) of {monitors, Ms} -> " "[Q || {process, Q} <- Ms, is_pid(Q)]; _ -> [] end end, " "Manager = Call(erlang, whereis, [albedo_terminals_manager]), " "Workers = Monitored(Manager), " 'io:format("~w", [[[Length(W) | [Length(P) || P <- Monitored(W)]]' " || W <- Workers]])", ) )
def test_a_flooding_terminal_stays_bounded_and_accepts_new_sockets(self): tid = "term-pressure" unread = TerminalWebSocket(self.app, tid, timeout=5, shell="/bin/sh") self.addCleanup(unread.close) self.assertEqual(unread.receive_frame(timeout=5)[0], 1) unread.send_bytes(b"yes PRESSURE_OUTPUT\r\n") self.assertEqual(unread.receive_frame(timeout=5)[0], 2)
# From here the unread socket stops reading. Its send buffer fills # within a moment, after which an unbounded worker or socket mailbox # only grows; the first sample already sees it. for _ in range(3): time.sleep(1) lengths = self.mailbox_lengths() self.assertTrue(lengths, "no terminal worker found") self.assertLessEqual( max(max(row) for row in lengths), 100, f"terminal mailboxes [worker, watchers...] grew: {lengths}", )
# Each new socket must get through Mist's 500 ms init window while # the producer keeps writing. deadline = time.monotonic() + 15 for attempt in range(16): self.assertLess(time.monotonic(), deadline, f"attach {attempt} too slow") ws = TerminalWebSocket(self.app, tid, timeout=3, shell="/bin/sh") self.addCleanup(ws.close) opcode, payload = ws.receive_frame(timeout=3) self.assertEqual(opcode, 1, f"attach {attempt}: {payload[:80]!r}") self.assertEqual(json.loads(payload)["type"], "ready") ws.close()
# Resize and signal reach the PTY ahead of the flood: Ctrl-C stops # the producer and the shell reports the new size. control = TerminalWebSocket(self.app, tid, timeout=3, shell="/bin/sh") self.addCleanup(control.close) self.assertEqual(control.receive_frame(timeout=3)[0], 1) control.send_text(type="resize", cols=100, rows=30) control.send_text(type="signal", signal="SIGINT") control.send_bytes(b"stty size\r\n") self.assertIn(b"30 100", self.read_until(control, b"30 100"))
# Once the unread socket reads again it is told what it missed. self.assertIn(b"[albedo: skipped", self.read_until(unread, b"[albedo: skipped"))
req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200)
def read_until(self, ws, marker, seconds=10): """Read output frames until marker appears; return the last frames read.""" window = b"" deadline = time.monotonic() + seconds while marker not in window and time.monotonic() < deadline: opcode, payload = ws.receive_frame(timeout=3) if opcode == 2: window = window[-len(marker) :] + payload return window
if __name__ == "__main__": unittest.main()