okami.mom/albedo

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

Aly Raffaufrefactor(tests): share repeated E2E fixtures5bdc0ee

main
branchesmainall refs →
10.1 KiB270 linesraw
1"""Unread background job results wake idle sessions once, even across idle reaping."""
2
3import json
4import os
5import time
6import unittest
7import urllib.error
8
9from harness import (
10    Albedo,
11    Provider,
12    Reply,
13    exclusive,
14    latest_user,
15    python,
16    release_fifo,
17    text,
18)
19
20# A daemon under this limit sweeps every half second, so the detached job
21# below outlives several sweeps.
22IDLE_SECONDS = 2
23
24
25def wait_for(predicate, timeout=40):
26    deadline = time.monotonic() + timeout
27    while time.monotonic() < deadline:
28        result = predicate()
29        if result:
30            return result
31        time.sleep(0.1)
32    raise AssertionError("background job did not wake session")
33
34
35class JobWakeCase(unittest.TestCase):
36    idle_seconds = None
37
38    def setUp(self):
39        self.job_gate = None
40
41        def script(request):
42            last = request["messages"][-1]
43            user = latest_user(request)
44            if last.get("role") == "user" and "start a slow job" in user:
45                seconds = IDLE_SECONDS * 2.5 if "detached" in user else 1.2
46                program = (
47                    f"open({str(self.job_gate)!r}, 'rb').read(1); print('wake-done')"
48                    if self.job_gate is not None
49                    else f"import time; time.sleep({seconds}); print('wake-done')"
50                )
51                code = (
52                    f"import sys\njob = run(sys.executable, '-c', {program!r})\njob.id"
53                )
54                return python(code)
55            if last.get("role") == "user" and "await a slow job past its cell" in user:
56                program = (
57                    f"open({str(self.job_gate)!r}, 'rb').read(1); print('wake-done')"
58                )
59                code = f"import sys\njob = run(sys.executable, '-c', {program!r})\nawait job"
60                return Reply(
61                    "python", tool_arguments={"code": code, "timeout_ms": 1500}
62                )
63            return text("finished")
64
65        def prepare(app):
66            if self.idle_seconds:
67                app.env["ALBEDO_IDLE_SECONDS"] = str(self.idle_seconds)
68
69        self.provider = Provider(script)
70        self.addCleanup(self.provider.close)
71        self.app = Albedo(self.provider, prepare=prepare)
72        self.app.__enter__()
73        self.addCleanup(self.app.__exit__, None, None, None)
74
75    def users(self):
76        return [latest_user(record["request"]) for record in self.provider.requests]
77
78
79class JobWakeTests(JobWakeCase):
80    def test_unread_job_wakes_once_with_live_source_and_durable_turn(self):
81        session = self.app.session()
82        gate = self.app.workspace / "job-release"
83        self.job_gate = gate
84        os.mkfifo(gate)
85        self.app.prompt(session, "start a slow job").close()
86        self.app.idle(session)
87        self.assertEqual(len(self.users()), 2)
88        before = self.app.stream_page(session)
89
90        wait_for(lambda: release_fifo(gate))
91        wait_for(lambda: len(self.users()) >= 3)
92        self.app.idle(session)
93        self.assertEqual(len(self.users()), 3)
94        wake = self.users()[2]
95        self.assertIn("background job finished", wake)
96        self.assertIn("jobs[", wake)
97        self.assertIn("output.read", wake)
98        live = self.app.stream_page(session, before)["events"]
99        notices = [
100            event["data"]["entry"]
101            for event in live
102            if event["type"] == "message"
103            and event["data"]["entry"]["kind"] == "note"
104            and any(
105                part["kind"] == "text" and "job finished" in part["text"]
106                for part in event["data"]["entry"]["content"]
107            )
108        ]
109        self.assertEqual(
110            len(notices),
111            1,
112            {"live": live, "history": self.app.history(session)},
113        )
114        # The transcript shows the one-line display; the model's notice with the
115        # handle and advice stays in the request above.
116        shown = "".join(
117            part["text"] for part in notices[0]["content"] if part["kind"] == "text"
118        )
119        self.assertTrue(shown.startswith("job finished (exit_code=0"), shown)
120        self.assertNotIn("\n", shown)
121        self.assertNotIn("jobs[", shown)
122        durable = [
123            entry
124            for entry in self.app.history(session)["items"]
125            if entry["kind"] == "note"
126            and any(
127                part["kind"] == "text" and "job finished" in part["text"]
128                for part in entry["content"]
129            )
130        ]
131        self.assertEqual(len(durable), 1)
132        self.assertEqual(durable[0]["id"], notices[0]["id"])
133        self.assertEqual(durable[0]["turn_id"], notices[0]["turn_id"])
134        self.assertIsNotNone(durable[0]["turn_id"])
135        self.assertEqual(durable[0]["input_id"], notices[0]["input_id"])
136        for entry in [notices[0], durable[0]]:
137            self.assertEqual(
138                [
139                    part["value"]
140                    for part in entry["content"]
141                    if part["kind"] == "json" and part["field"] == "origin"
142                ],
143                ["job"],
144            )
145        time.sleep(2.5)
146        self.assertEqual(len(self.users()), 3)
147
148    def start_blocked_job(self, prompt="start a slow job"):
149        session = self.app.session()
150        self.job_gate = self.app.workspace / "job-release"
151        os.mkfifo(self.job_gate)
152        self.app.prompt(session, prompt).close()
153        self.app.idle(session)
154        self.assertEqual(len(self.users()), 2)
155        return session
156
157    def jobs_resource(self, session):
158        return json.loads(
159            self.app.api(f"/extensions/run/sessions/{session}/jobs").read()
160        )
161
162    def test_running_job_observation_and_page_clear_after_completion(self):
163        session = self.start_blocked_job()
164        kernel = json.loads(self.app.api(f"/sessions/{session}?tail=0").read())[
165            "kernel"
166        ]
167        self.assertEqual(kernel["live_job_count"], 1)
168        running = kernel["running_jobs"]
169        self.assertEqual(len(running), 1)
170        self.assertIsInstance(running[0]["pid"], int)
171        self.assertIn(str(self.job_gate), running[0]["command"])
172        # The row says when the owner recorded the job, in wall-clock
173        # milliseconds, so a client can tell fresh work from settled work.
174        started = running[0]["started_at"]
175        self.assertIsInstance(started, int)
176        self.assertGreater(started, 0)
177        self.assertLess(abs(started - time.time() * 1000), 60_000)
178
179        resource = self.jobs_resource(session)
180        self.assertEqual(resource["items"], running)
181        self.assertEqual(resource["live_job_count"], 1)
182        page = resource["page"]
183        self.assertEqual(page["title"], "jobs")
184        self.assertEqual(page["summary"], "1 background job running")
185        self.assertEqual(len(page["rows"]), 1)
186        self.assertIn(str(self.job_gate), page["rows"][0]["text"])
187        # A row resource is a validated representation or nothing; clients
188        # reject one without an etag.
189        row_resource = page["rows"][0]["resource"]
190        self.assertTrue(
191            row_resource is None or isinstance(row_resource.get("etag"), str),
192            row_resource,
193        )
194        stop = page["actions"][0]["operation"]
195        self.assertEqual(stop["method"], "POST")
196        self.assertEqual(
197            stop["path_template"],
198            "/extensions/run/sessions/{session_id}/jobs/{job_id}/stop",
199        )
200        self.assertEqual(
201            stop["path"],
202            {
203                "session_id": {"source": "literal", "value": session},
204                "job_id": {"source": "row", "pointer": "/id"},
205            },
206        )
207
208        gate = self.job_gate
209        assert gate is not None
210
211        wait_for(lambda: release_fifo(gate))
212        wait_for(lambda: len(self.users()) >= 3)
213        self.app.idle(session)
214        after = json.loads(self.app.api(f"/sessions/{session}?tail=0").read())["kernel"]
215        self.assertEqual(after["live_job_count"], 0)
216        self.assertEqual(after["running_jobs"], [])
217        self.assertEqual(self.jobs_resource(session)["items"], [])
218
219    def test_an_await_cut_short_by_the_cell_deadline_still_wakes(self):
220        session = self.start_blocked_job("await a slow job past its cell deadline")
221        statuses = [
222            json.loads(part["value"])["status"]
223            for entry in self.app.history(session)["items"]
224            if entry["kind"] == "tool_result"
225            for part in entry["content"]
226            if part["kind"] == "json" and part["field"] == "result"
227        ]
228        self.assertEqual(statuses, ["interrupted"])
229        gate = self.job_gate
230        assert gate is not None
231        wait_for(lambda: release_fifo(gate))
232        wait_for(lambda: len(self.users()) >= 3)
233        self.assertIn("background job finished", self.users()[2])
234
235    def test_stopping_one_job_and_rejecting_unknown_job(self):
236        session = self.start_blocked_job()
237        resource = self.jobs_resource(session)
238        job_id = resource["items"][0]["id"]
239        stopped = json.loads(
240            self.app.api(
241                f"/extensions/run/sessions/{session}/jobs/{job_id}/stop",
242                method="POST",
243            ).read()
244        )
245        self.assertEqual(stopped, {"stopped": job_id})
246        wait_for(lambda: self.jobs_resource(session)["live_job_count"] == 0)
247        self.assertEqual(self.jobs_resource(session)["items"], [])
248        with self.assertRaises(urllib.error.HTTPError) as missing:
249            self.app.api(
250                f"/extensions/run/sessions/{session}/jobs/missing/stop",
251                method="POST",
252            )
253        self.assertEqual(missing.exception.code, 404)
254
255
256# exclusive: changes daemon idle sweep timing
257@exclusive
258class IdleSweepTests(JobWakeCase):
259    idle_seconds = IDLE_SECONDS
260
261    def test_live_job_survives_idle_kernel_sweep(self):
262        session = self.app.session()
263        self.app.prompt(session, "start a slow job, detached").close()
264        self.app.idle(session)
265        self.assertEqual(len(self.users()), 2)
266        # No daemon contact while detached: a live job must pin its kernel.
267        wait_for(lambda: len(self.users()) >= 3)
268        self.app.idle(session)
269        self.assertEqual(len(self.users()), 3)
270        self.assertIn("background job finished", self.users()[2])