"""Unread TCP streams cannot retain unbounded activity or block model completion.""" import json import socket import subprocess import threading import unittest from harness import Albedo, Provider, ROOT, exclusive, text class StreamProbe: def __init__(self, app): self.app = app source = ROOT / "test/e2e/albedo_stream_pressure_probe.erl" subprocess.run( ["erlc", "-Werror", "-o", str(app.root), str(source)], check=True, timeout=30, ) def call(self, function, arguments="[]"): app = self.app cookie = (app.home / "inspect.cookie").read_text().strip() node = f"albedo_{app.daemon._pid}@{socket.gethostname().split('.')[0]}" beam = app.root / "albedo_stream_pressure_probe.beam" expression = ( f"Node = '{node}', " f"{{ok, Binary}} = file:read_file({json.dumps(str(beam))}), " "{module, albedo_stream_pressure_probe} = " "rpc:call(Node, code, load_binary, " '[albedo_stream_pressure_probe, "probe.erl", Binary]), ' f"io:put_chars(rpc:call(Node, albedo_stream_pressure_probe, {function}, {arguments})), halt()." ) result = subprocess.run( [ "erl", "+S", "2:2", "-sname", f"pressure_probe_{app.daemon._pid}", "-setcookie", cookie, "-noshell", "-eval", expression, ], capture_output=True, text=True, check=False, timeout=40, ) if result.returncode != 0: raise AssertionError(result.stdout + result.stderr) return json.loads(result.stdout) def unread_stream(app, path): connection = socket.socket() connection.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1024) connection.settimeout(15) connection.connect(("127.0.0.1", app.connection["port"])) connection.sendall( ( f"GET {path} HTTP/1.1\r\nHost: localhost\r\n" f"Authorization: Bearer {app.connection['token']}\r\n\r\n" ).encode() ) # Read only the readiness batch, then leave all later socket data unread. initial = b"" while b"data: " not in initial or b"\n\n" not in initial.split(b"data: ", 1)[1]: part = connection.recv(1) if not part: raise AssertionError("stream closed before readiness") initial += part return connection @exclusive class StreamPressureTest(unittest.TestCase): def setUp(self): self.provider = Provider(lambda _: text("model completed under pressure")) self.addCleanup(self.provider.close) self.app = Albedo( self.provider, prepare=lambda app: app.daemon.env.update(ALBEDO_INSPECT="1"), ) self.app.__enter__() self.addCleanup(self.app.__exit__, None, None, None) self.probe = StreamProbe(self.app) def test_stalled_agents_feed_is_bounded_while_responsive_feed_and_model_complete( self, ): sessions = [self.app.session() for _ in range(8)] stalled = unread_stream(self.app, "/agents/stream") self.addCleanup(stalled.close) observed = [] ready = threading.Event() completed = threading.Event() def consume(): with self.app.api("/agents/stream") as response: for line in response: if line.startswith(b"data: "): events = json.loads(line[6:])["events"] observed.extend(events) ready.set() if any( event.get("text") == "model completed under pressure" for event in events ): completed.set() return threading.Thread(target=consume, daemon=True).start() self.assertTrue(ready.wait(5), "responsive feed never subscribed") ids = "[" + ",".join(f"<<{json.dumps(session)}>>" for session in sessions) + "]" measured = self.probe.call("pressure_json", f"[{ids}]") self.assertTrue(measured["pressure"], measured) self.assertEqual(measured["blocked"], 1, measured) self.app.prompt( sessions[0], "complete while another subscriber is stalled" ).close() self.app.idle(sessions[0]) self.assertTrue(completed.wait(5), "responsive feed lost model completion") self.assertLessEqual(measured["count"], 256, measured) self.assertLessEqual(measured["bytes"], 1048576, measured) self.assertLessEqual(measured["mailbox_bytes"], 1024, measured) self.assertLessEqual(measured["mailbox"], 3, measured) pressure_events = [ event for event in observed if event.get("text", "").startswith("pressure:") ] self.assertEqual(len(pressure_events), measured["published"]) for session in sessions: self.assertEqual( [ int(event["text"].split(":", 3)[2]) for event in pressure_events if event["session"] == session ], list(range(1, measured["published"] // len(sessions) + 1)), ) self.assertFalse(any(event["type"] == "overflow" for event in observed)) with self.app.api(f"/agents?session={sessions[0]}") as response: snapshot = json.load(response) self.assertFalse(snapshot["nodes"][0]["running"]) stalled.close() self.assertEqual( self.probe.call("removed_json"), {"subscribers": 0, "queues": 0} ) def test_idle_agents_disconnect_removes_subscription_without_publication(self): connection = unread_stream(self.app, "/agents/stream") self.assertEqual(self.probe.call("subscriptions_json")["subscribers"], 1) connection.close() self.assertEqual( self.probe.call("removed_json"), {"subscribers": 0, "queues": 0} ) def test_agent_count_byte_and_oversize_limits_emit_only_overflow_then_close(self): session = self.app.session() for name, count, size in ( ("count", 257, 16), ("bytes", 128, 16384), ("oversize", 1, 1048577), ): with self.subTest(limit=name), self.app.api("/agents/stream") as response: initial = next(line for line in response if line.startswith(b"data: ")) self.assertEqual(json.loads(initial[6:])["events"], []) measured = self.probe.call( "burst_json", f"[[<<{json.dumps(session)}>>],{count},{size}]" ) batches = [ json.loads(line[6:]) for line in response if line.startswith(b"data: ") ] self.assertEqual(batches, [{"events": [{"type": "overflow"}]}]) self.assertEqual(measured["overflow"], 1, measured) self.assertLessEqual(measured["count"], 256, measured) self.assertLessEqual(measured["bytes"], 1048576, measured) if name == "count": self.assertEqual(measured["count"], 256, measured) elif name == "bytes": self.assertLess(measured["count"], 256, measured) else: self.assertEqual(measured["bytes"], 0, measured) self.assertEqual( self.probe.call("removed_json"), {"subscribers": 0, "queues": 0} ) self.app.prompt(session, "model continues after overflow").close() self.app.idle(session) with self.app.api(f"/agents?session={session}") as response: snapshot = json.load(response) self.assertFalse(snapshot["nodes"][0]["running"]) def test_stalled_session_has_one_wake_and_reconnect_resets_durable_history(self): session = self.app.session() before = self.app.stream_page(session) connection = unread_stream(self.app, f"/sessions/{session}/stream") self.addCleanup(connection.close) self.provider.chunk_size = 4096 self.provider.script = lambda _: text("session pressure " * 250000, delay=0.002) self.app.prompt(session, "durable prompt under session pressure").close() measured = self.probe.call( "session_pressure_json", f"[<<{json.dumps(session)}>>]" ) self.assertEqual(measured["blocked"], 1, measured) self.assertLessEqual(measured["wakes"], 1, measured) self.app.idle(session) recovered = self.app.stream_page(session, before, tail=2) self.assertEqual(recovered["generation"], before["generation"]) self.assertEqual(recovered["events"][0]["type"], "reset") self.assertIn( "durable prompt under session pressure", [ event.get("text") for event in recovered["events"] if event["type"] == "user" ], ) def test_idle_session_disconnect_removes_watcher_without_publication(self): session = self.app.session() connection = unread_stream(self.app, f"/sessions/{session}/stream") connection.close() measured = self.probe.call( "session_removed_json", f"[<<{json.dumps(session)}>>]" ) self.assertEqual(measured["watchers"], 0, measured)