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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227"""Interactive pseudo-terminal lifecycle and WebSocket streaming."""
import base64import hashlibimport jsonimport secretsimport socketimport structimport timeimport unittestimport urllib.parseimport urllib.request
from harness import Albedo
class TerminalWebSocket: def __init__(self, app, terminal_id, **params): address = urllib.parse.urlsplit(app.base) self.socket = socket.create_connection((address.hostname, address.port), 20) 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 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) 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()
if __name__ == "__main__": unittest.main()