okami.mom/albedo

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

Aly Raffaufrefactor(test): share inspection and stream supportcba5ff8

main
branchesmainall refs →
39.7 KiB938 linesraw
1"""Clients recover bounded live tool progress without reconstructing arguments."""
2
3import json
4import threading
5import time
6import unittest
7import urllib.parse
8
9from harness import Albedo, Provider, Reply, exclusive, text
10from stream_support import StreamProbe, unread_stream
11from stream_replay_test import restart_actor
12
13
14def tool_events(
15    protocol,
16    calls,
17    ready,
18    release,
19    *,
20    chunk_size=7,
21    complete=True,
22    extra=None,
23    barrier=False,
24):
25    """Pause the actual provider after argument fragments, before execution."""
26    arguments = {
27        index: json.dumps({"code": code, **(extra or {})}) for index, _, code in calls
28    }
29    if protocol == "responses":
30        yield {"type": "response.created", "response": {"id": "progress-response"}}
31        for index, identity, _ in calls:
32            yield {
33                "type": "response.output_item.added",
34                "output_index": index,
35                "item": {
36                    "id": f"item-{index}",
37                    "type": "function_call",
38                    "call_id": identity,
39                    "name": "python",
40                },
41            }
42    for offset in range(0, max(map(len, arguments.values())), chunk_size):
43        for index, identity, _ in calls:
44            fragment = arguments[index][offset : offset + chunk_size]
45            if not fragment:
46                continue
47            if protocol == "responses":
48                yield {
49                    "type": "response.function_call_arguments.delta",
50                    "output_index": index,
51                    "delta": fragment,
52                }
53            else:
54                call = {"index": index, "function": {"arguments": fragment}}
55                if offset == 0:
56                    call.update(
57                        id=identity,
58                        type="function",
59                        function={"name": "python", "arguments": fragment},
60                    )
61                yield {
62                    "id": "progress-response",
63                    "choices": [
64                        {
65                            "index": 0,
66                            "delta": {"tool_calls": [call]},
67                            "finish_reason": None,
68                        }
69                    ],
70                }
71    if barrier:
72        # A later visible call proves the daemon consumed all prior fragments.
73        assert protocol == "responses"
74        index = max(arguments) + 1
75        identity = "progress-barrier"
76        calls = [*calls, (index, identity, "pass")]
77        arguments[index] = json.dumps({"code": "pass"})
78        yield {
79            "type": "response.output_item.added",
80            "output_index": index,
81            "item": {
82                "id": f"item-{index}",
83                "type": "function_call",
84                "call_id": identity,
85                "name": "python",
86            },
87        }
88        yield {
89            "type": "response.function_call_arguments.delta",
90            "output_index": index,
91            "delta": arguments[index],
92        }
93    ready.set()
94    if not release.wait(30):
95        raise AssertionError("test did not release the provider")
96    if not complete:
97        raise ConnectionResetError("scripted provider stream ended before completion")
98    if protocol == "responses":
99        output = [
100            {"id": f"gap-{index}", "type": "reasoning", "summary": []}
101            for index in range(max(arguments) + 1)
102        ]
103        for index, identity, _ in calls:
104            output[index] = {
105                "id": f"item-{index}",
106                "type": "function_call",
107                "call_id": identity,
108                "name": "python",
109                "arguments": arguments[index],
110                "status": "completed",
111            }
112        yield {
113            "type": "response.completed",
114            "response": {
115                "id": "progress-response",
116                "status": "completed",
117                "output": output,
118                "usage": {"input_tokens": 10, "output_tokens": 20},
119            },
120        }
121    else:
122        yield {
123            "id": "progress-response",
124            "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
125        }
126
127
128def current_progress(app, session, predicate):
129    deadline = time.monotonic() + 20
130    page = None
131    while time.monotonic() < deadline:
132        page = app.stream_page(session)
133        if predicate(page["snapshot"]["current_progress"]):
134            return page
135        time.sleep(0.01)
136    raise AssertionError(f"live progress did not reach the provider barrier: {page}")
137
138
139class ToolProgressTests(unittest.TestCase):
140    def gated_reply(self, protocol, calls, **options):
141        ready, release = threading.Event(), threading.Event()
142        self.addCleanup(release.set)
143        reply = Reply(
144            "python",
145            events=tool_events(protocol, calls, ready, release, **options),
146        )
147        return reply, ready, release
148
149    def assert_bounds(self, progress):
150        self.assertLessEqual(
151            len(
152                json.dumps(
153                    {
154                        "type": "tool_progress",
155                        "sequence": 1,
156                        "data": {"progress": progress},
157                    }
158                ).encode()
159            ),
160            8192,
161        )
162        self.assertLessEqual(len(progress["name"].encode()), 100)
163        if "preview" in progress:
164            self.assertGreaterEqual(progress["preview"]["offset_scalars"], 0)
165            self.assertLessEqual(len(progress["preview"]["text"]), 512)
166            self.assertLessEqual(len(progress["preview"]["text"].encode()), 2048)
167
168    def durable_arguments(self, app, session, tool_call_id):
169        [entry] = [
170            entry
171            for entry in app.history(session)["items"]
172            if entry["kind"] == "tool_call"
173            and entry["tool"]["tool_call_id"] == tool_call_id
174        ]
175        [part] = [
176            part
177            for part in entry["content"]
178            if part.get("field") == "arguments"
179            or part.get("reference", {}).get("field") == "arguments"
180        ]
181        if part["kind"] == "json":
182            return part["value"]
183        reference = part["reference"]
184        pieces, offset, token = [], 0, None
185        for _ in range(100):
186            path = reference["url"]
187            if token is not None:
188                path += "?" + urllib.parse.urlencode({"next": token})
189            with app.api(path) as response:
190                page = json.load(response)
191            for chunk in page["parts"]:
192                if chunk["field"] != "arguments":
193                    continue
194                self.assertEqual(chunk["encoding"], "utf8")
195                self.assertEqual(chunk["offset_bytes"], offset)
196                pieces.append(chunk["text"])
197                offset += len(chunk["text"].encode())
198            token = page["next"]
199            if token is None:
200                self.assertEqual(offset, reference["bytes"])
201                return json.loads("".join(pieces))
202        self.fail("durable tool arguments did not finish paging")
203
204    def test_fragment_burst_flushes_latest_preview_without_more_provider_events(self):
205        code = ("# " + "x" * 1000 + "\n") * 128 + "print('quiet provider')"
206        reply, ready, release = self.gated_reply(
207            "responses", [(0, "quiet-call", code)], chunk_size=64
208        )
209        provider = Provider(lambda _: [reply, text("done")])
210        self.addCleanup(provider.close)
211        with Albedo(provider, protocol="responses") as app:
212            session = app.session()
213            before = app.stream_page(session)
214            app.prompt(session, "pause after a burst of tiny fragments").close()
215            self.assertTrue(ready.wait(30), "provider never reached its quiet barrier")
216            snapshot = current_progress(
217                app,
218                session,
219                lambda values: (
220                    len(values) == 1
221                    and values[0].get("preview", {}).get("text") == code[-512:]
222                ),
223            )
224            # Reading a reset sees unflushed state and must not consume the
225            # update owed to subscribers that already hold a replay cursor.
226            deadline = time.monotonic() + 20
227            while True:
228                replay = app.stream_page(session, before)
229                self.assertNotIn("reset", [event["type"] for event in replay["events"]])
230                updates = [
231                    event["data"]["progress"]
232                    for event in replay["events"]
233                    if event["type"] == "tool_progress"
234                    and event["data"]["progress"] is not None
235                ]
236                if (
237                    updates
238                    and updates[-1].get("preview", {}).get("text") == code[-512:]
239                ):
240                    break
241                if time.monotonic() >= deadline:
242                    self.fail("the quiet provider's final preview was never published")
243                time.sleep(0.01)
244            self.assertEqual(updates[-1], snapshot["snapshot"]["current_progress"][0])
245            release.set()
246            app.idle(session)
247            completion = app.stream_page(session, replay)["events"]
248            running_seen = False
249            for event in completion:
250                if (
251                    event["type"] == "tool_progress"
252                    and event["data"]["progress"] is not None
253                ):
254                    phase = event["data"]["progress"]["phase"]
255                    if phase == "running":
256                        running_seen = True
257                    elif running_seen:
258                        self.fail(
259                            "a delayed generating preview followed tool execution"
260                        )
261            self.assertTrue(running_seen, "running transition was not delivered")
262            [result] = [event for event in completion if event["type"] == "tool"]
263            self.assertFalse(result["data"]["content_complete"])
264            self.assertIsNotNone(result["data"]["reference"])
265            self.assertEqual(
266                self.durable_arguments(app, session, "quiet-call")["code"], code
267            )
268
269    def test_unicode_preview_window_keeps_offsets_and_durable_arguments(self):
270        code = (
271            "# " + "prefix " * 400 + "\n"
272            "from pathlib import Path\n"
273            "Path('résumé😀.txt').write_text('escaped \\\" quote')\n"
274        )
275        for protocol in ("responses", "chat"):
276            with self.subTest(protocol=protocol):
277                reply, ready, release = self.gated_reply(
278                    protocol, [(0, "unicode-call", code)]
279                )
280                provider = Provider(lambda _: [reply, text("done")])
281                self.addCleanup(provider.close)
282                with Albedo(
283                    provider,
284                    protocol="responses"
285                    if protocol == "responses"
286                    else "chat_completions",
287                ) as app:
288                    session = app.session()
289                    self.assertEqual(
290                        app.stream_page(session)["snapshot"]["current_progress"], []
291                    )
292                    app.prompt(session, "write the Unicode file").close()
293                    self.assertTrue(ready.wait(20), "provider did not stream arguments")
294                    snapshot = current_progress(
295                        app,
296                        session,
297                        lambda values: (
298                            len(values) == 1
299                            and values[0].get("preview", {}).get("text") == code[-512:]
300                        ),
301                    )
302                    progress = snapshot["snapshot"]["current_progress"][0]
303                    self.assert_bounds(progress)
304                    self.assertEqual(progress["phase"], "generating")
305                    self.assertEqual(
306                        progress["preview"]["offset_scalars"], len(code) - 512
307                    )
308                    for page in (
309                        app.stream_page(session, tail=1),
310                        app.stream_page(session),
311                    ):
312                        self.assertEqual(page["events"][0]["type"], "reset")
313                        self.assertEqual(
314                            page["snapshot"]["current_progress"], [progress]
315                        )
316                    release.set()
317                    app.idle(session)
318                    replay = app.stream_page(session, snapshot)
319                    self.assertNotIn("snapshot", replay)
320                    events = replay["events"]
321                    self.assertNotIn(
322                        "arguments_delta", [event["type"] for event in events]
323                    )
324                    running = [
325                        event["data"]["progress"]
326                        for event in events
327                        if event["type"] == "tool_progress"
328                        and event["data"]["progress"] is not None
329                        and event["data"]["progress"]["phase"] == "running"
330                    ]
331                    self.assertEqual(len(running), 1)
332                    self.assertEqual(running[0]["call_id"], progress["call_id"])
333                    self.assertEqual(running[0]["tool_call_id"], "unicode-call")
334                    self.assert_bounds(running[0])
335                    [result] = [event for event in events if event["type"] == "tool"]
336                    self.assertEqual(
337                        result["data"]["progress_call_id"], progress["call_id"]
338                    )
339                    self.assertEqual(result["data"]["arguments"]["code"], code)
340                    [durable] = [
341                        entry
342                        for entry in app.history(session)["items"]
343                        if entry["kind"] == "tool_call"
344                    ]
345                    arguments = next(
346                        part["value"]
347                        for part in durable["content"]
348                        if part["kind"] == "json" and part["field"] == "arguments"
349                    )
350                    self.assertEqual(arguments["code"], code)
351                    self.assertEqual(
352                        (app.workspace / "résumé😀.txt").read_text(),
353                        'escaped " quote',
354                    )
355                    self.assertEqual(
356                        app.stream_page(session)["snapshot"]["current_progress"], []
357                    )
358
359    def test_attachment_hides_unnamed_calls_then_restores_their_preview(self):
360        code = "print('late name')"
361        sentinel = "print('sentinel')"
362        ready, reveal, named, complete = (threading.Event() for _ in range(4))
363        self.addCleanup(reveal.set)
364        self.addCleanup(complete.set)
365
366        def chunk(index, function, identity=None):
367            call = {"index": index, "function": function}
368            if identity is not None:
369                call.update(id=identity, type="function")
370            return {
371                "id": "late-name",
372                "choices": [{"index": 0, "delta": {"tool_calls": [call]}}],
373            }
374
375        def events():
376            yield chunk(0, {"arguments": json.dumps({"code": code})}, "late")
377            yield chunk(
378                1,
379                {"name": "python", "arguments": json.dumps({"code": sentinel})},
380                "sentinel",
381            )
382            ready.set()
383            if not reveal.wait(30):
384                raise AssertionError("tool-name release timed out")
385            yield chunk(0, {"name": "python"})
386            named.set()
387            if not complete.wait(30):
388                raise AssertionError("tool completion release timed out")
389            yield {
390                "id": "late-name",
391                "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
392            }
393
394        provider = Provider(lambda _: [Reply("python", events=events()), text("done")])
395        self.addCleanup(provider.close)
396        with Albedo(provider, protocol="chat_completions") as app:
397            session = app.session()
398            before = app.stream_page(session)
399            app.prompt(session, "attach before the first tool gets its name").close()
400            self.assertTrue(ready.wait(20))
401            # Seeing the second call proves the actor consumed the earlier
402            # unnamed call, without relying on a provider/socket timing guess.
403            snapshot = current_progress(
404                app,
405                session,
406                lambda values: any(
407                    value.get("preview", {}).get("text") == sentinel for value in values
408                ),
409            )
410            self.assertEqual(len(snapshot["snapshot"]["current_progress"]), 1)
411            self.assertEqual(
412                snapshot["snapshot"]["current_progress"][0]["name"], "python"
413            )
414            reveal.set()
415            self.assertTrue(named.wait(20))
416            snapshot = current_progress(
417                app,
418                session,
419                lambda values: (
420                    {value.get("preview", {}).get("text") for value in values}
421                    == {code, sentinel}
422                ),
423            )
424            self.assertEqual(len(snapshot["snapshot"]["current_progress"]), 2)
425            complete.set()
426            app.idle(session)
427            results = [
428                event
429                for event in app.stream_page(session, before)["events"]
430                if event["type"] == "tool"
431            ]
432            self.assertEqual(
433                {event["data"]["arguments"]["code"] for event in results},
434                {code, sentinel},
435            )
436
437    def test_interleaved_calls_keep_identity_through_native_id_deduplication(self):
438        for protocol in ("responses", "chat"):
439            with self.subTest(protocol=protocol):
440                calls = [
441                    (
442                        2,
443                        "first" if protocol == "responses" else "shared",
444                        "print('first')",
445                    ),
446                    (
447                        7,
448                        "second" if protocol == "responses" else "shared",
449                        "print('second')",
450                    ),
451                ]
452                reply, ready, release = self.gated_reply(protocol, calls)
453                provider = Provider(lambda _: [reply, text("done")])
454                self.addCleanup(provider.close)
455                with Albedo(
456                    provider,
457                    protocol="responses"
458                    if protocol == "responses"
459                    else "chat_completions",
460                ) as app:
461                    session = app.session()
462                    app.prompt(session, "run both calls").close()
463                    self.assertTrue(ready.wait(20))
464                    snapshot = current_progress(
465                        app,
466                        session,
467                        lambda values: (
468                            len(values) == 2
469                            and {
470                                value.get("preview", {}).get("text") for value in values
471                            }
472                            == {call[2] for call in calls}
473                        ),
474                    )
475                    identities = {
476                        value["preview"]["text"]: value["call_id"]
477                        for value in snapshot["snapshot"]["current_progress"]
478                    }
479                    self.assertEqual(len(set(identities.values())), 2)
480                    release.set()
481                    app.idle(session)
482                    replay = app.stream_page(session, snapshot)
483                    results = [
484                        event for event in replay["events"] if event["type"] == "tool"
485                    ]
486                    self.assertEqual(len(results), 2)
487                    self.assertEqual(
488                        len({event["data"]["tool_call_id"] for event in results}), 2
489                    )
490                    for event in results:
491                        self.assertEqual(
492                            event["data"]["progress_call_id"],
493                            identities[event["data"]["arguments"]["code"]],
494                        )
495                    self.assertEqual(
496                        app.stream_page(session)["snapshot"]["current_progress"], []
497                    )
498
499    def test_retry_discards_failed_attempt_progress_before_reusing_output_index(self):
500        failed, first_ready, first_release = self.gated_reply(
501            "responses", [(0, "discarded", "print('old attempt')")], complete=False
502        )
503        succeeding, second_ready, second_release = self.gated_reply(
504            "responses", [(0, "accepted", "print('new attempt')")]
505        )
506        provider = Provider(lambda _: [failed, succeeding, text("done")])
507        self.addCleanup(provider.close)
508        with Albedo(provider, protocol="responses") as app:
509            session = app.session()
510            before = app.stream_page(session)
511            app.prompt(session, "retry the provider").close()
512            self.assertTrue(first_ready.wait(20))
513            old = current_progress(app, session, lambda values: len(values) == 1)
514            first_release.set()
515            self.assertTrue(second_ready.wait(20))
516            new = current_progress(
517                app,
518                session,
519                lambda values: (
520                    len(values) == 1
521                    and values[0].get("preview", {}).get("text")
522                    == "print('new attempt')"
523                ),
524            )
525            self.assertNotEqual(
526                new["snapshot"]["current_progress"][0]["call_id"],
527                old["snapshot"]["current_progress"][0]["call_id"],
528            )
529            second_release.set()
530            app.idle(session)
531            results = [
532                event
533                for event in app.stream_page(session, before)["events"]
534                if event["type"] == "tool"
535            ]
536            self.assertEqual(len(results), 1)
537            self.assertEqual(
538                results[0]["data"]["arguments"]["code"], "print('new attempt')"
539            )
540            self.assertEqual(
541                app.stream_page(session)["snapshot"]["current_progress"], []
542            )
543
544    def test_interrupt_clears_progress_and_prevents_delayed_execution(self):
545        code = "from pathlib import Path\nPath('must-not-exist').write_text('bad')"
546        reply, ready, release = self.gated_reply("responses", [(0, "cancelled", code)])
547        provider = Provider(lambda _: [reply, text("later turn")])
548        self.addCleanup(provider.close)
549        with Albedo(provider, protocol="responses") as app:
550            session = app.session()
551            app.prompt(session, "pause before execution").close()
552            self.assertTrue(ready.wait(20))
553            current_progress(app, session, lambda values: len(values) == 1)
554            with app.api(f"/sessions/{session}?tail=0") as response:
555                session_state = json.load(response)
556            app.api(
557                f"/sessions/{session}/interrupt",
558                {
559                    "run_id": session_state["status"]["run_id"],
560                    "through_input_order": session_state["input_order"],
561                },
562            ).close()
563            release.set()
564            app.idle(session)
565            self.assertEqual(
566                app.stream_page(session)["snapshot"]["current_progress"], []
567            )
568            self.assertFalse((app.workspace / "must-not-exist").exists())
569            app.prompt(session, "continue with a clean turn").close()
570            app.idle(session)
571            self.assertEqual(
572                app.stream_page(session)["snapshot"]["current_progress"], []
573            )
574
575    # exclusive: enables daemon inspection at startup and replaces a session actor
576    @exclusive
577    def test_replacement_actor_drops_old_live_progress(self):
578        code = "from pathlib import Path\nPath('stale-call').write_text('bad')"
579        reply, ready, release = self.gated_reply("responses", [(0, "stale-call", code)])
580        provider = Provider(lambda _: [reply, text("clean turn")])
581        self.addCleanup(provider.close)
582        with Albedo(
583            provider,
584            protocol="responses",
585            prepare=lambda app: app.daemon.env.update(ALBEDO_INSPECT="1"),
586        ) as app:
587            session = app.session()
588            app.prompt(session, "pause an old call").close()
589            self.assertTrue(ready.wait(20))
590            old = current_progress(app, session, lambda values: len(values) == 1)
591            restart_actor(app, session)
592            replacement = app.stream_page(session, old)
593            self.assertNotEqual(replacement["generation"], old["generation"])
594            self.assertEqual(replacement["events"][0]["type"], "reset")
595            self.assertEqual(replacement["snapshot"]["current_progress"], [])
596            release.set()
597            app.prompt(session, "continue after replacement").close()
598            app.idle(session)
599            self.assertFalse((app.workspace / "stale-call").exists())
600            self.assertEqual(
601                app.stream_page(session)["snapshot"]["current_progress"], []
602            )
603
604    def test_late_attachment_restores_running_progress_until_result(self):
605        code = (
606            "from pathlib import Path\nimport asyncio\n"
607            "Path('running-ready').write_text('ready')\n"
608            "while not Path('finish-running').exists():\n"
609            "    await asyncio.sleep(0.01)\n"
610            "print('finished')"
611        )
612        reply, ready, release = self.gated_reply(
613            "responses", [(0, "running-call", code)]
614        )
615        provider = Provider(lambda _: [reply, text("done")])
616        self.addCleanup(provider.close)
617        with Albedo(provider, protocol="responses") as app:
618            session = app.session()
619            app.prompt(session, "hold a running call").close()
620            self.assertTrue(ready.wait(20))
621            generating = current_progress(app, session, lambda values: len(values) == 1)
622            release.set()
623            finished = app.workspace / "finish-running"
624            try:
625                deadline = time.monotonic() + 20
626                while not (app.workspace / "running-ready").exists():
627                    if time.monotonic() >= deadline:
628                        self.fail("the kernel did not reach the running barrier")
629                    time.sleep(0.01)
630                snapshot = current_progress(
631                    app,
632                    session,
633                    lambda values: len(values) == 1 and values[0]["phase"] == "running",
634                )
635                [progress] = snapshot["snapshot"]["current_progress"]
636                self.assertEqual(
637                    progress["call_id"],
638                    generating["snapshot"]["current_progress"][0]["call_id"],
639                )
640                self.assertEqual(progress["tool_call_id"], "running-call")
641                self.assert_bounds(progress)
642                self.assertEqual(
643                    app.stream_page(session, tail=1)["snapshot"]["current_progress"],
644                    [progress],
645                )
646            finally:
647                finished.write_text("finish")
648            app.idle(session)
649            self.assertEqual(
650                app.stream_page(session)["snapshot"]["current_progress"], []
651            )
652
653    def test_oversized_native_identity_is_omitted_from_progress_not_truncated(self):
654        identity = "native-" + "😀" * 60
655        code = "# " + "😀" * 700 + "\nprint('bounded Unicode preview')"
656        reply, ready, release = self.gated_reply(
657            "responses", [(0, identity, code)], chunk_size=1000
658        )
659        provider = Provider(lambda _: [reply, text("done")])
660        self.addCleanup(provider.close)
661        with Albedo(provider, protocol="responses") as app:
662            session = app.session()
663            app.prompt(session, "run a call with a large identity").close()
664            self.assertTrue(ready.wait(20))
665            snapshot = current_progress(
666                app,
667                session,
668                lambda values: (
669                    len(values) == 1
670                    and values[0].get("preview", {}).get("text") == code[-512:]
671                ),
672            )
673            self.assert_bounds(snapshot["snapshot"]["current_progress"][0])
674            release.set()
675            app.idle(session)
676            events = app.stream_page(session, snapshot)["events"]
677            [running] = [
678                event["data"]["progress"]
679                for event in events
680                if event["type"] == "tool_progress"
681                and event["data"]["progress"] is not None
682                and event["data"]["progress"]["phase"] == "running"
683            ]
684            self.assertIsNone(running["tool_call_id"])
685            self.assert_bounds(running)
686            [result] = [event for event in events if event["type"] == "tool"]
687            self.assertEqual(result["data"]["tool_call_id"], identity)
688            self.assertEqual(result["data"]["progress_call_id"], running["call_id"])
689
690    def test_complex_arguments_degrade_preview_without_blocking_execution(self):
691        nested = "leaf"
692        for _ in range(70):
693            nested = [nested]
694        code = "print('complex arguments executed')"
695        reply, ready, release = self.gated_reply(
696            "responses",
697            [(0, "complex-call", code)],
698            chunk_size=1000,
699            extra={"metadata": nested},
700            barrier=True,
701        )
702        provider = Provider(lambda _: [reply, text("done")])
703        self.addCleanup(provider.close)
704        with Albedo(provider, protocol="responses") as app:
705            session = app.session()
706            before = app.stream_page(session)
707            app.prompt(session, "run a call with deep metadata").close()
708            self.assertTrue(ready.wait(20))
709            snapshot = current_progress(
710                app,
711                session,
712                lambda values: any(
713                    value.get("preview", {}).get("text") == "pass" for value in values
714                ),
715            )
716            self.assertEqual(len(snapshot["snapshot"]["current_progress"]), 2)
717            [progress] = [
718                value
719                for value in snapshot["snapshot"]["current_progress"]
720                if value.get("preview", {}).get("text") != "pass"
721            ]
722            self.assertNotIn("preview", progress)
723            self.assertEqual(progress["name"], "python")
724            release.set()
725            app.idle(session)
726            [result] = [
727                event
728                for event in app.stream_page(session, before)["events"]
729                if event["type"] == "tool"
730                and event["data"]["tool_call_id"] == "complex-call"
731            ]
732            self.assertEqual(
733                result["data"]["arguments"], {"code": code, "metadata": nested}
734            )
735            execution = json.loads(result["data"]["result"])
736            self.assertEqual(execution["status"], "ok", execution)
737            self.assertEqual(execution["output"], "complex arguments executed\n")
738
739    def test_escaped_metadata_keys_disable_preview_without_changing_execution(self):
740        code = "print('escaped metadata executed')"
741        metadata = {'metadata\\"': ', "code": "misleading_preview_sentinel"'}
742        reply, ready, release = self.gated_reply(
743            "responses",
744            [(0, "escaped-key", code)],
745            chunk_size=7,
746            extra=metadata,
747            barrier=True,
748        )
749        provider = Provider(lambda _: [reply, text("done")])
750        self.addCleanup(provider.close)
751        with Albedo(provider, protocol="responses") as app:
752            session = app.session()
753            before = app.stream_page(session)
754            app.prompt(session, "run a call with escaped metadata").close()
755            self.assertTrue(ready.wait(20))
756            snapshot = current_progress(
757                app,
758                session,
759                lambda values: any(
760                    value.get("preview", {}).get("text") == "pass" for value in values
761                ),
762            )
763            self.assertEqual(len(snapshot["snapshot"]["current_progress"]), 2)
764            [progress] = [
765                value
766                for value in snapshot["snapshot"]["current_progress"]
767                if value.get("preview", {}).get("text") != "pass"
768            ]
769            self.assertNotIn("preview", progress)
770            self.assertEqual(progress["name"], "python")
771            release.set()
772            app.idle(session)
773            [result] = [
774                event
775                for event in app.stream_page(session, before)["events"]
776                if event["type"] == "tool"
777                and event["data"]["tool_call_id"] == "escaped-key"
778            ]
779            self.assertEqual(result["data"]["arguments"], {"code": code, **metadata})
780            execution = json.loads(result["data"]["result"])
781            self.assertEqual(execution["status"], "ok", execution)
782            self.assertEqual(execution["output"], "escaped metadata executed\n")
783
784    # exclusive: enables the daemon inspection endpoint before startup
785    @exclusive
786    def test_subscribers_share_one_bounded_argument_projection(self):
787        code = ("# " + "x" * 1000 + "\n") * 128 + "print('shared projection')"
788        reply, ready, release = self.gated_reply(
789            "responses", [(0, "shared-call", code)], chunk_size=8192
790        )
791        provider = Provider(lambda _: [reply, text("done")])
792        self.addCleanup(provider.close)
793        with Albedo(
794            provider,
795            protocol="responses",
796            prepare=lambda app: app.daemon.env.update(ALBEDO_INSPECT="1"),
797        ) as app:
798            session = app.session()
799            initial = app.stream_page(session)
800            app.prompt(session, "show shared progress").close()
801            self.assertTrue(ready.wait(20))
802            snapshot = current_progress(
803                app,
804                session,
805                lambda values: (
806                    len(values) == 1
807                    and values[0].get("preview", {}).get("text") == code[-512:]
808                ),
809            )
810            probe = StreamProbe(app)
811            arguments = f"[<<{json.dumps(session)}>>]"
812            self.assertEqual(
813                probe.call("session_removed_json", arguments)["watchers"], 0
814            )
815            self.assertEqual(
816                snapshot["snapshot"]["current_progress"][0]["preview"][
817                    "offset_scalars"
818                ],
819                len(code) - 512,
820            )
821            self.assert_bounds(snapshot["snapshot"]["current_progress"][0])
822            deadline = time.monotonic() + 20
823            while True:
824                published = app.stream_page(session, initial)
825                if any(
826                    event["type"] == "tool_progress"
827                    and event["data"]["progress"]
828                    == snapshot["snapshot"]["current_progress"][0]
829                    for event in published["events"]
830                ):
831                    break
832                if time.monotonic() >= deadline:
833                    self.fail("shared progress never finished publishing its preview")
834                time.sleep(0.01)
835            self.assertEqual(
836                probe.call("session_removed_json", arguments)["watchers"], 0
837            )
838            before = probe.call("progress_json", arguments)
839            sockets = [unread_stream(app, f"/sessions/{session}") for _ in range(10)]
840            for connection in sockets:
841                self.addCleanup(connection.close)
842            after = probe.call("progress_json", arguments)
843            self.assertEqual(after["watchers"], before["watchers"] + 10)
844            self.assertEqual(after["projection_bytes"], before["projection_bytes"])
845            self.assertEqual(after["projection_size"], before["projection_size"])
846            self.assertLessEqual(after["projection_bytes"], 64 * 1024)
847            # Include lists and tuples, so retained codepoints cannot hide behind
848            # the binary-only byte count as the full arguments grow.
849            self.assertLessEqual(after["projection_size"], 64 * 1024)
850            attached = app.stream_page(session)
851            self.assertEqual(
852                attached["snapshot"]["current_progress"],
853                snapshot["snapshot"]["current_progress"],
854                f"provider released={release.is_set()}; "
855                f"requests={len(provider.requests)}; "
856                f"status={attached['snapshot']['status']}",
857            )
858            self.assertEqual(len(provider.requests), 1)
859            release.set()
860            app.idle(session)
861            [result] = [
862                event
863                for event in app.stream_page(session, published)["events"]
864                if event["type"] == "tool"
865            ]
866            self.assertFalse(result["data"]["content_complete"])
867            self.assertIsNotNone(result["data"]["reference"])
868            self.assertEqual(
869                self.durable_arguments(app, session, result["data"]["tool_call_id"]),
870                {"code": code},
871            )
872            [durable_result] = [
873                entry
874                for entry in app.history(session)["items"]
875                if entry["kind"] == "tool_result"
876                and entry["tool"]["tool_call_id"] == result["data"]["tool_call_id"]
877            ]
878            execution = next(
879                json.loads(part["value"])
880                for part in durable_result["content"]
881                if part["kind"] == "json" and part["field"] == "result"
882            )
883            self.assertEqual(execution["status"], "ok", execution)
884            self.assertEqual(execution["output"], "shared projection\n")
885
886    def test_excess_generating_calls_disable_previews_without_losing_execution(self):
887        calls = [(index, f"many-{index}", f"print({index})") for index in range(33)]
888        reply, ready, release = self.gated_reply("responses", calls, chunk_size=100)
889        provider = Provider(lambda _: [reply, text("done")])
890        self.addCleanup(provider.close)
891        with Albedo(provider, protocol="responses") as app:
892            session = app.session()
893            before = app.stream_page(session)
894            app.prompt(session, "run all calls").close()
895            self.assertTrue(ready.wait(20))
896            deadline = time.monotonic() + 20
897            while True:
898                replay = app.stream_page(session, before)
899                cleared = any(
900                    event["type"] == "tool_progress"
901                    and event["data"]["progress"] is None
902                    for event in replay["events"]
903                )
904                if cleared:
905                    break
906                if time.monotonic() >= deadline:
907                    self.fail("overflow did not clear generating progress")
908                time.sleep(0.01)
909            self.assertEqual(
910                app.stream_page(session)["snapshot"]["current_progress"], []
911            )
912            release.set()
913            app.idle(session, timeout=60)
914            results = [
915                event
916                for event in app.stream_page(session, before)["events"]
917                if event["type"] == "tool"
918            ]
919            self.assertEqual(
920                {event["data"]["tool_call_id"] for event in results},
921                {call[1] for call in calls},
922            )
923            for event in results:
924                index = int(event["data"]["tool_call_id"].split("-")[1])
925                self.assertEqual(
926                    event["data"]["arguments"]["code"],
927                    calls[index][2],
928                )
929                execution = json.loads(event["data"]["result"])
930                self.assertEqual(execution["status"], "ok", execution)
931                self.assertEqual(execution["output"], f"{index}\n")
932            self.assertEqual(
933                app.stream_page(session)["snapshot"]["current_progress"], []
934            )
935            self.assertNotIn(
936                "arguments_delta",
937                [event["type"] for event in app.stream_page(session, before)["events"]],
938            )