okami.mom/albedo

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

Aly Raffauffix(session): use monotonic time for retry deadlines31c9e74

main
branchesmainall refs →
23.4 KiB601 linesraw
1"""Identified inputs retain their admission and deliver once through failure and restart."""
2
3from concurrent.futures import ThreadPoolExecutor
4import base64
5import json
6import os
7from pathlib import Path
8import shlex
9import shutil
10import sqlite3
11import threading
12import time
13import unittest
14import urllib.error
15
16from harness import Albedo, Provider, exclusive, operation_id, python, text
17from image_limits_test import png
18
19
20class InputScenario(unittest.TestCase):
21    def setUp(self):
22        self.provider = Provider(lambda _: text("done"))
23        self.addCleanup(self.provider.close)
24        self.app = Albedo(self.provider, protocol="responses")
25        self.app.__enter__()
26        self.addCleanup(self.app.__exit__, None, None, None)
27
28    def put_input(self, session, input_id, intent):
29        with self.app.api(
30            f"/sessions/{session}/inputs/{input_id}", intent, method="PUT"
31        ) as response:
32            return json.load(response)
33
34    def receipt(self, session, input_id):
35        with self.app.api(f"/sessions/{session}/inputs/{input_id}") as response:
36            return json.load(response)
37
38    def committed(self, session, input_id):
39        deadline = time.monotonic() + 30
40        while time.monotonic() < deadline:
41            receipt = self.receipt(session, input_id)
42            if receipt["delivery"] == "committed":
43                return receipt
44            time.sleep(0.03)
45        self.fail(f"input remained pending: {receipt}")
46
47    def duplicates(self, session, input_id, intent):
48        with ThreadPoolExecutor(max_workers=6) as executor:
49            results = list(
50                executor.map(
51                    lambda _: self.put_input(session, input_id, intent), range(6)
52                )
53            )
54        self.assertTrue(
55            all(
56                result["id"] == input_id and result["admission"] == "accepted"
57                for result in results
58            )
59        )
60        self.assertEqual(len({result["acceptance_order"] for result in results}), 1)
61        return results[0]
62
63    def conflict(self, session, input_id, intent):
64        with self.assertRaises(urllib.error.HTTPError) as failure:
65            self.put_input(session, input_id, intent)
66        self.assertEqual(failure.exception.code, 409)
67        self.assertEqual(json.load(failure.exception)["code"], "input_conflict")
68
69    def users(self, session):
70        return [
71            entry
72            for entry in self.app.history(session)["items"]
73            if entry["kind"] == "user"
74        ]
75
76    def skill_intent(self, session, arguments):
77        with self.app.api(f"/sessions/{session}/catalog?kind=skills") as response:
78            discovery = json.load(response)["discovery"]
79        candidate = next(
80            row for row in discovery["candidates"] if row["title"] == "retry-skill"
81        )
82        return {
83            "kind": "skill",
84            "candidate_id": candidate["id"],
85            "catalog_revision": discovery["revision"],
86            "arguments": arguments,
87        }
88
89    def write_skill(self):
90        skill = self.app.workspace / ".agents/skills/retry-skill/SKILL.md"
91        skill.parent.mkdir(parents=True)
92        skill.write_text(
93            "---\nname: retry-skill\ndescription: test activation\n---\nOriginal saved instructions.\n"
94        )
95        return skill
96
97    def assert_user_order(self, session, inputs):
98        users = self.users(session)
99        self.assertEqual(
100            [entry["input_id"] for entry in users], [input_id for input_id, _ in inputs]
101        )
102        self.assertEqual(
103            [
104                part["text"]
105                for entry in users
106                for part in entry["content"]
107                if part["kind"] == "text"
108            ],
109            [intent["text"] for _, intent in inputs],
110        )
111
112
113class OperationsTests(InputScenario):
114    def test_targeted_cancel_preserves_shared_turn_and_cannot_interrupt_a_later_turn(
115        self,
116    ):
117        first_entered, second_entered, later_entered = [
118            threading.Event() for _ in range(3)
119        ]
120        first_release, second_release, later_release = [
121            threading.Event() for _ in range(3)
122        ]
123        for gate in (first_release, second_release, later_release):
124            self.addCleanup(gate.set)
125        calls = 0
126
127        def reply(_request):
128            nonlocal calls
129            calls += 1
130            entered, release = [
131                (first_entered, first_release),
132                (second_entered, second_release),
133                (later_entered, later_release),
134            ][calls - 1]
135            entered.set()
136            if not release.wait(20):
137                raise AssertionError(
138                    "targeted cancellation scenario did not release provider"
139                )
140            return python("1") if calls == 1 else text("completed")
141
142        self.provider.script = reply
143        session = self.app.session()
144        identities = [operation_id() for _ in range(3)]
145        self.put_input(
146            session, identities[0], {"kind": "message", "text": "start shared turn"}
147        )
148        self.assertTrue(first_entered.wait(10))
149        for index, identity in enumerate(identities[1:]):
150            self.put_input(
151                session,
152                identity,
153                {"kind": "message", "text": f"join shared turn {index}"},
154            )
155        first_release.set()
156        self.assertTrue(second_entered.wait(15))
157        receipts = [self.committed(session, identity) for identity in identities]
158        turn_ids = {receipt["turn"]["id"] for receipt in receipts}
159        self.assertEqual(len(turn_ids), 1)
160        with self.app.api(
161            f"/sessions/{session}/inputs/{identities[1]}/cancel", {}
162        ) as response:
163            cancellation = json.load(response)
164        self.assertEqual(cancellation["result"], "shared_running")
165        with self.app.api(f"/sessions/{session}?tail=0") as response:
166            self.assertFalse(json.load(response)["status"]["interrupt_requested"])
167        second_release.set()
168        self.app.idle(session)
169        self.assertTrue(
170            all(
171                self.receipt(session, identity)["turn"]["state"] == "completed"
172                for identity in identities
173            )
174        )
175        later = operation_id()
176        self.put_input(
177            session, later, {"kind": "message", "text": "a later independent turn"}
178        )
179        self.assertTrue(later_entered.wait(10))
180        with self.app.api(
181            f"/sessions/{session}/inputs/{identities[0]}/cancel", {}
182        ) as response:
183            self.assertEqual(json.load(response)["result"], "not_pending")
184        with self.app.api(f"/sessions/{session}?tail=0") as response:
185            snapshot = json.load(response)
186        self.assertFalse(snapshot["status"]["interrupt_requested"])
187        self.assertNotIn(snapshot["status"]["run_id"], turn_ids)
188        later_release.set()
189        self.app.idle(session)
190        self.assertEqual(self.receipt(session, later)["turn"]["state"], "completed")
191
192    def test_concurrent_creation_has_one_resource_and_immutable_provenance(self):
193        session_id = operation_id()
194        resource = f"/sessions/{session_id}"
195        intent = {
196            "kind": "new",
197            "workspace": str(self.app.workspace),
198            "provider_profile": self.app.profile,
199        }
200
201        def create(_):
202            try:
203                with self.app.api(
204                    resource, intent, method="PUT", headers={"If-None-Match": "*"}
205                ) as response:
206                    return response.status, json.load(response)
207            except urllib.error.HTTPError as failure:
208                self.assertEqual(failure.code, 412)
209                return failure.code, None
210
211        with ThreadPoolExecutor(max_workers=6) as executor:
212            results = list(executor.map(create, range(6)))
213        self.assertEqual([status for status, _ in results].count(201), 1)
214        with self.app.api(resource) as response:
215            created = json.load(response)
216        self.assertEqual(
217            created["creation"]["submitted"]["workspace"], intent["workspace"]
218        )
219        self.assertEqual(
220            created["creation"]["resolved"],
221            {
222                key: created[key]
223                for key in ("workspace", "provider_profile", "model", "effort")
224            },
225        )
226        configuration = created["configuration_resource"]
227        with self.app.api(
228            configuration["url"],
229            {"name": "edited after creation"},
230            method="PATCH",
231            headers={"If-Match": configuration["etag"]},
232        ) as response:
233            response.read()
234        with self.app.api(resource) as response:
235            edited = json.load(response)
236        self.assertEqual(edited["creation"], created["creation"])
237        self.assertEqual(edited["name"], "edited after creation")
238        with self.app.api("/sessions?scope=all") as response:
239            matches = [
240                row for row in json.load(response)["items"] if row["id"] == session_id
241            ]
242        self.assertEqual(len(matches), 1)
243
244    def test_submission_identity_and_durable_echo(self):
245        session = self.app.session()
246        first = operation_id()
247        intent = {"kind": "message", "text": "same message", "client_id": "first"}
248        self.duplicates(session, first, intent)
249        duplicate = self.put_input(session, first, {**intent, "client_id": "retry"})
250        self.assertEqual(duplicate["client_id"], "first")
251        self.committed(session, first)
252        self.app.idle(session)
253        self.conflict(session, first, {**intent, "text": "changed message"})
254        self.conflict(session, first, {"kind": "continue"})
255        second = operation_id()
256        self.put_input(session, second, {**intent, "client_id": "second"})
257        self.committed(session, second)
258        self.app.idle(session)
259        self.assert_user_order(session, [(first, intent), (second, intent)])
260
261    def test_image_duplicates_and_changed_image_conflict(self):
262        session = self.app.session()
263        input_id = operation_id()
264        intent = {
265            "kind": "message",
266            "text": "inspect image",
267            "images": [
268                {
269                    "mime_type": "image/png",
270                    "data": base64.b64encode(png(2, 2)).decode(),
271                }
272            ],
273        }
274        self.duplicates(session, input_id, intent)
275        self.committed(session, input_id)
276        self.app.idle(session)
277        self.conflict(
278            session,
279            input_id,
280            {
281                **intent,
282                "images": [
283                    {
284                        **intent["images"][0],
285                        "data": base64.b64encode(png(3, 2)).decode(),
286                    }
287                ],
288            },
289        )
290        users = self.users(session)
291        self.assertEqual([entry["input_id"] for entry in users], [input_id])
292        image = next(
293            part["image"] for part in users[0]["content"] if part["kind"] == "image"
294        )
295        self.assertEqual(image["width"], 2)
296
297    def test_continue_is_one_admission_without_an_extra_user_row(self):
298        session = self.app.session()
299        self.app.prompt(session, "start").close()
300        self.app.idle(session)
301        input_id = operation_id()
302        self.duplicates(session, input_id, {"kind": "continue"})
303        self.committed(session, input_id)
304        self.app.idle(session)
305        self.assertEqual(len(self.users(session)), 1)
306        self.assertEqual(len(self.provider.requests), 2)
307        self.assertIsNotNone(self.receipt(session, input_id)["transcript_position"])
308
309    def test_skill_admission_uses_resolved_content_once(self):
310        skill = self.write_skill()
311        session = self.app.session()
312        input_id = operation_id()
313        intent = self.skill_intent(session, "do this")
314        accepted = self.duplicates(session, input_id, intent)
315        self.committed(session, input_id)
316        self.app.idle(session)
317        skill.unlink()
318        duplicate = self.put_input(session, input_id, intent)
319        self.assertEqual(duplicate["acceptance_order"], accepted["acceptance_order"])
320        self.conflict(session, input_id, {**intent, "arguments": "other"})
321        self.assertIn(
322            "Original saved instructions.", json.dumps(self.provider.requests)
323        )
324        self.assertEqual(
325            [entry["input_id"] for entry in self.users(session)], [input_id]
326        )
327
328
329@exclusive
330class WaitingOperationsTests(InputScenario):
331    def setUp(self):
332        self.gate = threading.Event()
333        self.started = threading.Event()
334        self.first = True
335
336        def reply(_):
337            if self.first:
338                self.first = False
339                self.started.set()
340                self.gate.wait(40)
341            return text("done")
342
343        self.provider = Provider(reply)
344        self.addCleanup(self.provider.close)
345        self.addCleanup(self.gate.set)
346        self.app = Albedo(self.provider, protocol="responses")
347        self.app.__enter__()
348        self.addCleanup(self.app.__exit__, None, None, None)
349
350    def test_full_queue_duplicate_bypasses_limit_and_interrupt_cancels(self):
351        session = self.app.session()
352        self.app.prompt(session, "hold active generation").close()
353        self.assertTrue(self.started.wait(10))
354        inputs = [
355            (operation_id(), {"kind": "message", "text": f"waiting {index}"})
356            for index in range(32)
357        ]
358        results = [
359            self.put_input(session, input_id, intent) for input_id, intent in inputs
360        ]
361        self.assertTrue(all(result["delivery"] == "pending" for result in results))
362        self.assertEqual(
363            self.put_input(session, *inputs[0])["acceptance_order"],
364            results[0]["acceptance_order"],
365        )
366        with self.assertRaises(urllib.error.HTTPError) as full:
367            self.put_input(
368                session, operation_id(), {"kind": "message", "text": "queue overflow"}
369            )
370        self.assertEqual(full.exception.code, 429)
371        with self.app.api(f"/sessions/{session}?tail=0") as response:
372            snapshot = json.load(response)
373        self.app.api(
374            f"/sessions/{session}/interrupt",
375            {
376                "run_id": snapshot["status"]["run_id"],
377                "through_input_order": snapshot["input_order"],
378            },
379        ).close()
380        self.gate.set()
381        self.app.idle(session)
382        self.app.restart(crash=True)
383        for (input_id, intent), accepted in zip(inputs, results):
384            self.assertEqual(self.receipt(session, input_id)["delivery"], "cancelled")
385            self.assertEqual(
386                self.put_input(session, input_id, intent)["acceptance_order"],
387                accepted["acceptance_order"],
388            )
389
390    def test_crash_recovers_order_images_and_resolved_skill_content(self):
391        skill = self.write_skill()
392        session = self.app.session()
393        with self.app.prompt(session, "hold active generation") as response:
394            active_input = json.load(response)["id"]
395        self.assertTrue(self.started.wait(10))
396        inputs = [
397            (operation_id(), {"kind": "message", "text": "first waiting"}),
398            (
399                operation_id(),
400                {
401                    "kind": "message",
402                    "text": "second image",
403                    "images": [
404                        {
405                            "mime_type": "image/png",
406                            "data": base64.b64encode(png(2, 2)).decode(),
407                        }
408                    ],
409                },
410            ),
411            (operation_id(), self.skill_intent(session, "saved arguments")),
412        ]
413        accepted = [
414            self.put_input(session, input_id, intent) for input_id, intent in inputs
415        ]
416        self.assertTrue(
417            all(
418                self.receipt(session, input_id)["delivery"] == "pending"
419                for input_id, _ in inputs
420            )
421        )
422        skill.unlink()
423        self.app.restart(crash=True)
424        self.assertEqual(
425            self.receipt(session, active_input)["turn"]["state"], "abandoned"
426        )
427        self.gate.set()
428        for input_id, _ in inputs:
429            self.committed(session, input_id)
430        self.app.idle(session)
431        self.app.restart(crash=True)
432        for (input_id, intent), result in zip(inputs, accepted):
433            self.assertEqual(
434                self.put_input(session, input_id, intent)["acceptance_order"],
435                result["acceptance_order"],
436            )
437        users = [
438            entry
439            for entry in self.users(session)
440            if entry["input_id"] in {input_id for input_id, _ in inputs}
441        ]
442        self.assertEqual(
443            [entry["input_id"] for entry in users], [input_id for input_id, _ in inputs]
444        )
445        image = next(
446            part["image"] for part in users[1]["content"] if part["kind"] == "image"
447        )
448        self.assertEqual(image["width"], 2)
449        self.assertIn(
450            "Original saved instructions.", json.dumps(self.provider.requests)
451        )
452
453
454@exclusive
455class BlockedOperationsTests(InputScenario):
456    def test_full_blocked_queue_survives_restart_and_replays_rejection(self):
457        session = self.app.session()
458        inputs = [
459            (operation_id(), {"kind": "message", "text": f"blocked {index}"})
460            for index in range(32)
461        ]
462        with sqlite3.connect(self.app.home / "albedo.sqlite") as database:
463            database.execute(
464                "CREATE TRIGGER hold_inputs BEFORE INSERT ON transcript WHEN NEW.row_class='user' BEGIN SELECT RAISE(ABORT,'fixture'); END"
465            )
466        accepted = [
467            self.put_input(session, input_id, intent) for input_id, intent in inputs
468        ]
469        deadline = time.monotonic() + 10
470        while time.monotonic() < deadline:
471            if self.receipt(session, inputs[0][0])["blocking_reason"]:
472                break
473            time.sleep(0.03)
474        else:
475            self.fail("waiting inputs were never blocked")
476        with self.app.api(f"/sessions/{session}?tail=0") as response:
477            snapshot = json.load(response)
478        self.assertEqual(
479            snapshot["status"]["blocking_reason"]["detail"],
480            self.receipt(session, inputs[0][0])["blocking_reason"]["detail"],
481        )
482        rejected_id = operation_id()
483        rejected = {"kind": "message", "text": "blocked overflow"}
484        with self.assertRaises(urllib.error.HTTPError) as failure:
485            self.put_input(session, rejected_id, rejected)
486        self.assertEqual(failure.exception.code, 429)
487        rejection = json.load(failure.exception)
488        self.assertEqual(self.receipt(session, rejected_id)["admission"], "rejected")
489        self.app.restart(crash=True)
490        self.assertEqual(
491            self.put_input(session, *inputs[-1])["acceptance_order"],
492            accepted[-1]["acceptance_order"],
493        )
494        with self.assertRaises(urllib.error.HTTPError) as restored_full:
495            self.put_input(session, operation_id(), {"kind": "continue"})
496        self.assertEqual(restored_full.exception.code, 429)
497        with sqlite3.connect(self.app.home / "albedo.sqlite") as database:
498            database.execute("DROP TRIGGER hold_inputs")
499        for input_id, _ in inputs:
500            self.committed(session, input_id)
501        self.app.idle(session)
502        self.assert_user_order(session, inputs)
503        with self.assertRaises(urllib.error.HTTPError) as replay:
504            self.put_input(session, rejected_id, rejected)
505        self.assertEqual(replay.exception.code, 429)
506        self.assertEqual(json.load(replay.exception), rejection)
507
508
509# exclusive: changes daemon PATH to gate kernel startup; also tests restart
510@exclusive
511class PreparingOperationsTests(InputScenario):
512    opened: Path
513    release: Path
514
515    def setUp(self):
516        provider = Provider(lambda request: text("done"))
517        self.addCleanup(provider.close)
518
519        def prepare(app):
520            self.opened = app.root / "kernel-opening"
521            self.release = app.root / "kernel-release"
522            commands = app.root / "kernel-commands"
523            commands.mkdir()
524            executable = shutil.which("python3")
525            self.assertIsNotNone(executable)
526            wrapper = commands / "python3"
527            wrapper.write_text(
528                "#!/bin/sh\n"
529                # a detached kernel opens through its bridge's start or launch
530                + 'case "$*" in *albedo_kernel.py*|*"albedo_bridge.py start"*|*"albedo_bridge.py launch"*)\n'
531                + "touch "
532                + shlex.quote(str(self.opened))
533                + "\n"
534                + "while [ ! -e "
535                + shlex.quote(str(self.release))
536                + " ]; do sleep 0.05; done\n"
537                + ";; esac\nexec "
538                + shlex.quote(str(executable))
539                + ' "$@"\n'
540            )
541            wrapper.chmod(0o755)
542            app.env["PATH"] = str(commands) + os.pathsep + app.env["PATH"]
543
544        self.app = Albedo(provider, protocol="responses", prepare=prepare)
545        self.app.__enter__()
546        self.addCleanup(self.app.__exit__, None, None, None)
547
548    def test_kernel_preparation_admission_survives_restart(self):
549        session = self.app.session()
550        input_id = operation_id()
551        intent = {"kind": "message", "text": "stored while kernel opens"}
552        accepted = self.put_input(session, input_id, intent)
553        deadline = time.monotonic() + 10
554        while not self.opened.exists() and time.monotonic() < deadline:
555            time.sleep(0.03)
556        self.assertTrue(self.opened.exists())
557        self.assertEqual(self.receipt(session, input_id)["delivery"], "pending")
558        self.app.restart(crash=True, prepare=lambda _: self.release.touch())
559        self.assertEqual(
560            self.put_input(session, input_id, intent)["acceptance_order"],
561            accepted["acceptance_order"],
562        )
563        self.committed(session, input_id)
564        self.app.idle(session)
565        self.assertEqual(
566            [entry["input_id"] for entry in self.users(session)], [input_id]
567        )
568
569    def test_full_kernel_opening_queue_delivers_once_and_replays_rejection(self):
570        session = self.app.session()
571        inputs = [
572            (operation_id(), {"kind": "message", "text": f"opening {index}"})
573            for index in range(32)
574        ]
575        results = [
576            self.put_input(session, input_id, intent) for input_id, intent in inputs
577        ]
578        deadline = time.monotonic() + 10
579        while not self.opened.exists() and time.monotonic() < deadline:
580            time.sleep(0.03)
581        self.assertTrue(self.opened.exists())
582        self.assertEqual(
583            self.put_input(session, *inputs[0])["acceptance_order"],
584            results[0]["acceptance_order"],
585        )
586        rejected_id = operation_id()
587        intent = {"kind": "continue"}
588        with self.assertRaises(urllib.error.HTTPError) as failure:
589            self.put_input(session, rejected_id, intent)
590        self.assertEqual(failure.exception.code, 429)
591        rejection = json.load(failure.exception)
592        self.assertEqual(self.receipt(session, rejected_id)["admission"], "rejected")
593        self.release.touch()
594        for input_id, _ in inputs:
595            self.committed(session, input_id)
596        self.app.idle(session)
597        self.assert_user_order(session, inputs)
598        with self.assertRaises(urllib.error.HTTPError) as replay:
599            self.put_input(session, rejected_id, intent)
600        self.assertEqual(replay.exception.code, 429)
601        self.assertEqual(json.load(replay.exception), rejection)