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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382"""One WebSocket must isolate subscriptions and preserve the SSE replay cursor."""
import base64import hashlibimport http.clientimport jsonimport secretsimport socketimport structimport timeimport unittestimport urllib.parse
from harness import Albedo, Provider, text
class WebSocket: """A masked stdlib client for the daemon's JSON event socket."""
def __init__(self, app): 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() self.socket.sendall( ( f"GET /events 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): self.reader.close() self.socket.close()
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(self, **message): self.frame(1, json.dumps(message).encode())
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(self): message = bytearray() while True: 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] if length > 1048576: raise AssertionError("event frame exceeded 1 MiB") payload = self.read_exact(length) if opcode == 8: raise AssertionError(f"WebSocket closed: {payload!r}") if opcode == 9: self.frame(10, payload) continue if opcode == 10: continue if opcode not in (0, 1): raise AssertionError(f"unexpected opcode {opcode}") message.extend(payload) if len(message) > 1048576: raise AssertionError("event message exceeded 1 MiB") if flags & 128: return json.loads(message)
def user_texts(batch): return [ part["text"] for event in batch["events"] if event["type"] == "message" and event["data"]["entry"]["kind"] == "user" for part in event["data"]["entry"]["content"] if part["kind"] == "text" ]
class EventsTest(unittest.TestCase): def setUp(self): self.provider = Provider(lambda _: text("socket answer")) self.addCleanup(self.provider.close)
def connect(self, app): connection = WebSocket(app) self.addCleanup(connection.close) return connection
def until(self, connection, predicate, *, reject=()): deadline = time.monotonic() + 20 while time.monotonic() < deadline: reply = connection.receive() self.assertNotIn("failure", reply, reply) self.assertNotIn(reply.get("key"), reject, reply) if predicate(reply): return reply self.fail("event socket did not deliver the expected reply")
def test_multiplexes_sessions_and_list_and_ends_only_the_deleted_subscription(self): with Albedo(self.provider) as app: sessions = {"left": app.session(), "right": app.session()} connection = self.connect(app) for key, session in sessions.items(): connection.send(key=key, type="session", session_id=session, tail=0) connection.send(key="roots", type="sessions", scope="roots") initial = {} while len(initial) < 3: reply = connection.receive() self.assertNotIn("failure", reply, reply) initial[reply["key"]] = reply["batch"] for key, session in sessions.items(): self.assertEqual(initial[key]["snapshot"]["id"], session) self.assertEqual( [event["type"] for event in initial["roots"]["events"]], ["ready", "reset"], )
for key, session in sessions.items(): app.prompt(session, f"question for {key}").close() observed = {key: [] for key in sessions} completed = set() listed = set()
def both_completed(reply): key = reply["key"] if key in sessions: batch = reply["batch"] observed[key].extend(user_texts(batch)) if any(e["type"] == "turn_completed" for e in batch["events"]): completed.add(key) elif key == "roots": for event in reply["batch"]["events"]: if event["type"] == "activity": listed.add(event["data"]["session_id"]) elif event["type"] == "invalidate": listed.update(event["data"]["session_ids"]) return len(completed) == 2 and listed.issuperset(sessions.values())
self.until(connection, both_completed) self.assertEqual( observed, {key: [f"question for {key}"] for key in sessions} ) for session in sessions.values(): app.idle(session)
connection.send(key="left", type="stop") self.until(connection, lambda reply: reply == {"key": "left", "end": True}) app.prompt(sessions["left"], "must not reach stopped subscription").close() app.idle(sessions["left"]) app.prompt(sessions["right"], "right remains subscribed").close() self.until( connection, lambda reply: ( reply.get("key") == "right" and "right remains subscribed" in user_texts(reply["batch"]) ), reject=("left",), ) app.idle(sessions["right"]) resource = f"/sessions/{sessions['right']}?view=configuration" with app.api(resource) as response: validator = response.getheader("ETag") response.read() app.api(resource, method="DELETE", headers={"If-Match": validator}).close() self.until( connection, lambda reply: reply == {"key": "right", "end": True}, reject=("left",), ) connection.send(type="ping") self.until( connection, lambda reply: reply == {"type": "pong"}, reject=("left", "right"), )
def test_resume_uses_the_same_cursor_without_reset_or_duplicate_messages(self): with Albedo(self.provider) as app: session = app.session() app.prompt(session, "already consumed").close() app.idle(session) connection = self.connect(app) connection.send(key="before", type="session", session_id=session, tail=0) before = self.until(connection, lambda reply: reply.get("key") == "before")[ "batch" ] connection.close()
app.prompt(session, "while disconnected").close() app.idle(session) resumed = self.connect(app) resumed.send( key="after", type="session", session_id=session, tail=0, after_generation=before["generation"], after_seq=before["cursor"], ) consumed = before["cursor"] observed = []
def completed(reply): nonlocal consumed self.assertEqual(reply["key"], "after") batch = reply["batch"] self.assertEqual(batch["generation"], before["generation"]) self.assertNotIn("reset", [event["type"] for event in batch["events"]]) sequences = [event["sequence"] for event in batch["events"]] self.assertEqual( sequences, list(range(consumed + 1, consumed + 1 + len(sequences))) ) self.assertEqual( batch["cursor"], sequences[-1] if sequences else consumed ) consumed = batch["cursor"] observed.extend(user_texts(batch)) return any( event["type"] == "turn_completed" for event in batch["events"] )
self.until(resumed, completed) self.assertEqual(observed, ["while disconnected"])
def test_subscription_limits_and_cursor_failures_leave_other_keys_live(self): with Albedo(self.provider) as app: session = app.session() connection = self.connect(app) connection.send(key="keep", type="session", session_id=session, tail=0) self.until(connection, lambda reply: reply.get("key") == "keep") connection.send( key="bad-cursor", type="session", session_id=session, tail=0, after_generation="invalid", after_seq=0, ) failure = connection.receive() self.assertEqual(failure["key"], "bad-cursor") self.assertIn("failure", failure) self.assertNotIn("batch", failure) for payload, key in [ (b"{", ""), (b'{"key":"unknown","type":"unknown"}', "unknown"), (b'{"key":"duplicate","type":"ping","type":"ping"}', "duplicate"), (b'{"key":"extra","type":"stop","extra":true}', "extra"), ]: with self.subTest(command=payload): connection.frame(1, payload) failure = connection.receive() self.assertEqual(failure["key"], key) self.assertEqual(failure["failure"]["code"], "invalid_request") self.assertNotIn("batch", failure)
# Every key follows the same session: removing one must not unwatch another. others = {f"copy-{index}" for index in range(63)} for key in sorted(others): connection.send(key=key, type="session", session_id=session, tail=0) ready = set()
def all_ready(reply): self.assertIn(reply["key"], others) self.assertEqual(reply["batch"]["snapshot"]["id"], session) ready.add(reply["key"]) return ready == others
self.until(connection, all_ready) connection.send( key="over-limit", type="session", session_id=session, tail=0 ) failure = connection.receive() self.assertEqual(failure["key"], "over-limit") self.assertEqual(failure["failure"]["code"], "follow_limit") self.assertNotIn("batch", failure)
for key in sorted(others): connection.send(key=key, type="stop") stopped = set()
def all_stopped(reply): self.assertEqual(reply, {"key": reply["key"], "end": True}) self.assertIn(reply["key"], others) stopped.add(reply["key"]) return stopped == others
self.until(connection, all_stopped) connection.send( key="replacement", type="session", session_id=session, tail=0 ) self.until( connection, lambda reply: reply.get("key") == "replacement", reject=others, ) app.prompt(session, "survives per-key failures and unsubscribe").close() observed = set()
def both_live(reply): self.assertIn(reply["key"], {"keep", "replacement"}) if "survives per-key failures and unsubscribe" in user_texts( reply["batch"] ): observed.add(reply["key"]) return observed == {"keep", "replacement"}
self.until(connection, both_live, reject=others) app.idle(session)
def test_upgrade_requires_authentication_allowed_origin_and_websocket_headers(self): with Albedo(self.provider) as app: address = urllib.parse.urlsplit(app.base) token = {"Authorization": "Bearer " + app.connection["token"]} upgrade = { "Connection": "Upgrade", "Upgrade": "websocket", "Sec-WebSocket-Key": base64.b64encode(secrets.token_bytes(16)).decode(), "Sec-WebSocket-Version": "13", } for name, headers, status in [ ("missing token", upgrade, 401), ( "forbidden origin", {**upgrade, **token, "Origin": "https://evil.example"}, 403, ), ("missing upgrade", token, 400), ]: with self.subTest(boundary=name): connection = http.client.HTTPConnection( address.hostname, address.port, timeout=20 ) try: connection.request("GET", "/events", headers=headers) response = connection.getresponse() self.assertEqual(response.status, status) response.read() finally: connection.close()