okami.mom/albedo

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

Aly Raffaufrefactor(test): share inspection and stream supportcba5ff8

main
branchesmainall refs →
3.6 KiB79 linesraw
1"""Cancellation releases execution workers and long tool turns complete."""
2
3from concurrent.futures import ThreadPoolExecutor
4import json
5import time
6
7from harness import exclusive
8from stream_support import StreamProbe
9
10from integration_support import IntegrationScenario
11
12
13class ExecutionRecoveryTests(IntegrationScenario):
14    # exclusive: enables the existing native inspection seam
15    @exclusive
16    def test_interrupt_before_first_token_stops_worker_without_killing_it(self):
17        for protocol in ("responses", "chat_completions"):
18            with self.subTest(protocol=protocol):
19                app = self.app_for(
20                    protocol, prepare=lambda app: app.env.update(ALBEDO_INSPECT="1")
21                )
22                session = app.session()
23                start = len(self.provider.requests)
24                with app.prompt(session, "no first token") as response:
25                    input_id = json.load(response)["id"]
26                self.wait_for_request(start)
27                captured = self.read(app, f"/sessions/{session}?tail=0")
28                probe = StreamProbe(app)
29                ready = app.workspace / f"monitor-{protocol}"
30                arguments = (
31                    "["
32                    + ",".join(
33                        "<<" + json.dumps(value) + ">>"
34                        for value in [session, captured["status"]["run_id"], str(ready)]
35                    )
36                    + "]"
37                )
38                with ThreadPoolExecutor(max_workers=1) as pool:
39                    exited = pool.submit(probe.call, "worker_exit_json", arguments)
40                    deadline = time.monotonic() + 30
41                    while not ready.exists() and time.monotonic() < deadline:
42                        time.sleep(0.01)
43                    self.assertTrue(ready.exists(), "worker monitor was not installed")
44                    app.api(
45                        f"/sessions/{session}/interrupt",
46                        {
47                            "run_id": captured["status"]["run_id"],
48                            "through_input_order": captured["input_order"],
49                        },
50                    ).close()
51                    self.assertEqual(exited.result(timeout=40), {"reason": "normal"})
52                app.idle(session)
53                receipt = self.read(app, f"/sessions/{session}/inputs/{input_id}")
54                self.assertEqual(receipt["turn"]["state"], "interrupted")
55                # The interrupted session takes a new turn and commits it.
56                self.send(app, session, "continue from branch")
57                self.assertTrue(
58                    any(
59                        e["kind"] == "assistant" and self.entry_text(e) == "finished"
60                        for e in self.history(app, session)
61                    )
62                )
63
64    def test_more_than_hundred_tool_turns_complete(self):
65        for protocol in ("responses", "chat_completions"):
66            with self.subTest(protocol=protocol):
67                app = self.app_for(protocol)
68                session = app.session()
69                start = len(self.provider.requests)
70                self.send(app, session, "continue for over 100 turns")
71                self.assertEqual(len(self.records(start)), 106)
72                events = self.history(app, session)
73                self.assertEqual(sum(e["kind"] == "tool_result" for e in events), 105)
74                self.assertTrue(
75                    any(
76                        e["kind"] == "assistant" and self.entry_text(e) == "finished"
77                        for e in events
78                    )
79                )