okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffaufrefactor(test): share inspection and stream supportcba5ff8

main
branchesmainall refs →
8.4 KiB191 linesraw
1"""Unread TCP streams cannot retain unbounded activity or block model completion."""
2
3import json
4import threading
5import unittest
6
7from harness import Albedo, Provider, exclusive, text
8from stream_support import StreamProbe, unread_stream
9
10
11# exclusive: enables daemon inspection at startup and measures global subscribers
12@exclusive
13class StreamPressureTest(unittest.TestCase):
14    def setUp(self):
15        self.provider = Provider(lambda _: text("model completed under pressure"))
16        self.addCleanup(self.provider.close)
17        self.app = Albedo(
18            self.provider,
19            prepare=lambda app: app.daemon.env.update(ALBEDO_INSPECT="1"),
20        )
21        self.app.__enter__()
22        self.addCleanup(self.app.__exit__, None, None, None)
23        self.probe = StreamProbe(self.app)
24
25    def test_stalled_agents_feed_is_bounded_while_responsive_feed_and_model_complete(
26        self,
27    ):
28        sessions = [self.app.session() for _ in range(8)]
29        stalled = unread_stream(self.app, "/sessions?scope=all")
30        self.addCleanup(stalled.close)
31        observed = []
32        ready = threading.Event()
33        completed = threading.Event()
34
35        def consume():
36            with self.app.api(
37                "/sessions?scope=all", headers={"Accept": "text/event-stream"}
38            ) as response:
39                for line in response:
40                    if line.startswith(b"data: "):
41                        events = json.loads(line[6:])["events"]
42                        observed.extend(events)
43                        ready.set()
44                        if any(
45                            event["type"] == "activity"
46                            and any(
47                                row["text"] == "model completed under pressure"
48                                for row in event["data"]["activity"]["lines"]
49                            )
50                            for event in events
51                        ):
52                            completed.set()
53                            return
54
55        threading.Thread(target=consume, daemon=True).start()
56        self.assertTrue(ready.wait(5), "responsive feed never subscribed")
57        ids = "[" + ",".join(f"<<{json.dumps(session)}>>" for session in sessions) + "]"
58        measured = self.probe.call("pressure_json", f"[{ids}]")
59        self.assertTrue(measured["pressure"], measured)
60        self.assertEqual(measured["blocked"], 1, measured)
61        self.app.prompt(
62            sessions[0], "complete while another subscriber is stalled"
63        ).close()
64        self.app.idle(sessions[0])
65        self.assertTrue(completed.wait(5), "responsive feed lost model completion")
66        self.assertLessEqual(measured["count"], 256, measured)
67        self.assertLessEqual(measured["bytes"], 1048576, measured)
68        self.assertLessEqual(measured["mailbox_bytes"], 1024, measured)
69        self.assertLessEqual(measured["mailbox"], 3, measured)
70        pressure_events = [
71            event
72            for event in observed
73            if event["type"] == "activity"
74            and event["data"]["activity"]["lines"]
75            and event["data"]["activity"]["lines"][0]["text"].startswith("pressure:")
76        ]
77        self.assertEqual(len(pressure_events), measured["published"])
78        for session in sessions:
79            self.assertEqual(
80                [
81                    int(event["data"]["activity"]["lines"][0]["text"].split(":", 3)[2])
82                    for event in pressure_events
83                    if event["data"]["session_id"] == session
84                ],
85                list(range(1, measured["published"] // len(sessions) + 1)),
86            )
87        self.assertFalse(any(event["type"] == "overflow" for event in observed))
88        with self.app.api(f"/sessions/{sessions[0]}?tail=0") as response:
89            snapshot = json.load(response)
90        self.assertEqual(snapshot["status"]["phase"], "idle")
91        stalled.close()
92        self.assertEqual(
93            self.probe.call("removed_json"), {"subscribers": 0, "queues": 0}
94        )
95
96    def test_idle_agents_disconnect_removes_subscription_without_publication(self):
97        connection = unread_stream(self.app, "/sessions?scope=all")
98        self.assertEqual(self.probe.call("subscriptions_json")["subscribers"], 1)
99        connection.close()
100        self.assertEqual(
101            self.probe.call("removed_json"), {"subscribers": 0, "queues": 0}
102        )
103
104    def test_agent_count_byte_and_oversize_limits_emit_only_overflow_then_close(self):
105        session = self.app.session()
106        for name, count, size in (
107            ("count", 257, 16),
108            ("bytes", 128, 16384),
109            ("oversize", 1, 1048577),
110        ):
111            with (
112                self.subTest(limit=name),
113                self.app.api(
114                    "/sessions?scope=all", headers={"Accept": "text/event-stream"}
115                ) as response,
116            ):
117                initial = next(line for line in response if line.startswith(b"data: "))
118                self.assertEqual(json.loads(initial[6:])["events"][0]["type"], "ready")
119                measured = self.probe.call(
120                    "burst_json", f"[[<<{json.dumps(session)}>>],{count},{size}]"
121                )
122                batches = [
123                    json.loads(line[6:])
124                    for line in response
125                    if line.startswith(b"data: ")
126                ]
127                self.assertEqual(
128                    batches, [{"events": [{"type": "overflow", "data": {}}]}]
129                )
130                self.assertEqual(measured["overflow"], 1, measured)
131                self.assertLessEqual(measured["count"], 256, measured)
132                self.assertLessEqual(measured["bytes"], 1048576, measured)
133                if name == "count":
134                    self.assertEqual(measured["count"], 256, measured)
135                elif name == "bytes":
136                    self.assertLess(measured["count"], 256, measured)
137                else:
138                    self.assertEqual(measured["bytes"], 0, measured)
139            self.assertEqual(
140                self.probe.call("removed_json"), {"subscribers": 0, "queues": 0}
141            )
142        self.app.prompt(session, "model continues after overflow").close()
143        self.app.idle(session)
144        with self.app.api(f"/sessions/{session}?tail=0") as response:
145            snapshot = json.load(response)
146        self.assertEqual(snapshot["status"]["phase"], "idle")
147
148    def test_stalled_session_has_one_wake_and_reconnect_resets_durable_history(self):
149        session = self.app.session()
150        before = self.app.stream_page(session)
151        connection = unread_stream(self.app, f"/sessions/{session}")
152        self.addCleanup(connection.close)
153        self.provider.chunk_size = 4096
154        self.provider.script = lambda _: text("session pressure " * 250000, delay=0.002)
155        self.app.prompt(session, "durable prompt under session pressure").close()
156        measured = self.probe.call(
157            "session_pressure_json", f"[<<{json.dumps(session)}>>]"
158        )
159        self.assertEqual(measured["blocked"], 1, measured)
160        self.assertLessEqual(measured["wakes"], 1, measured)
161        self.app.idle(session)
162        recovered = self.app.stream_page(session, before, tail=2)
163        self.assertEqual(recovered["generation"], before["generation"])
164        self.assertEqual(recovered["events"][0]["type"], "reset")
165        snapshot = recovered["snapshot"]
166        self.assertIsNotNone(snapshot["history"]["older"])
167        earliest = min(entry["position"] for entry in snapshot["history"]["items"])
168        with self.app.api(
169            f"/sessions/{session}/history?before={earliest}&limit=2"
170        ) as response:
171            earlier = json.load(response)
172        self.assertIn(
173            "durable prompt under session pressure",
174            [
175                part["text"]
176                for entry in earlier["items"]
177                if entry["kind"] == "user"
178                for part in entry["content"]
179                if part["kind"] == "text"
180            ],
181        )
182        self.assertLessEqual(len(json.dumps(recovered).encode()), 1048576)
183
184    def test_idle_session_disconnect_removes_watcher_without_publication(self):
185        session = self.app.session()
186        connection = unread_stream(self.app, f"/sessions/{session}")
187        connection.close()
188        measured = self.probe.call(
189            "session_removed_json", f"[<<{json.dumps(session)}>>]"
190        )
191        self.assertEqual(measured["watchers"], 0, measured)