okami.mom/albedo

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

Aly Raffaufrefactor(tests): share repeated E2E fixtures5bdc0ee

main
branchesmainall refs →
10.7 KiB266 linesraw
1"""Concurrent child kernel boot stays responsive and forwards all answers."""
2
3import json
4import os
5import sqlite3
6import sys
7import time
8import unittest
9import urllib.error
10
11from harness import (
12    Albedo,
13    Provider,
14    exclusive,
15    latest_user,
16    operation_id,
17    python,
18    release_fifo,
19    text,
20    wait_until,
21)
22
23HEAVY = os.environ.get("ALBEDO_E2E_HEAVY") == "1"
24ROOTS = 4 if HEAVY else 2
25CHILDREN = 12 if HEAVY else 4
26
27
28class AgentsSwarmTests(unittest.TestCase):
29    def test_children_boot_without_blocking_actors_or_losing_answers(self):
30        def script(request):
31            last = request["messages"][-1]
32            if last.get("role") == "user" and 'kind="task"' in latest_user(request):
33                return python("print('cell ran')")
34            return text("done")
35
36        provider = Provider(script)
37        try:
38            with Albedo(provider) as app:
39                roots = [app.session() for _ in range(ROOTS)]
40                children = []
41                started = time.monotonic()
42                for root in roots:
43                    for number in range(CHILDREN):
44                        child = operation_id()
45                        with app.api(
46                            f"/sessions/{child}",
47                            {
48                                "kind": "child",
49                                "parent_id": root,
50                                "address": f"sp-{number}",
51                                "name": f"sp-{number}",
52                                "initial_input_id": operation_id(),
53                                "task": f"count to {number}",
54                            },
55                            method="PUT",
56                            headers={"If-None-Match": "*"},
57                        ) as response:
58                            self.assertEqual(json.load(response)["id"], child)
59                        children.append(child)
60                self.assertLess(
61                    time.monotonic() - started, 30, "spawn must not await kernel boot"
62                )
63                slowest = 0.0
64                deadline = time.monotonic() + 300
65                while time.monotonic() < deadline:
66                    answers = self.answers(app, roots)
67                    tick = time.monotonic()
68                    for root in roots:
69                        with app.api(
70                            f"/sessions?scope=all&family_id={root}"
71                        ) as response:
72                            self.assertEqual(
73                                len(json.load(response)["items"]), CHILDREN + 1
74                            )
75                    for child in children[::6]:
76                        with app.api(f"/sessions/{child}?tail=0") as response:
77                            json.load(response)
78                    slowest = max(slowest, time.monotonic() - tick)
79                    if sum(answers) == ROOTS * CHILDREN:
80                        break
81                    time.sleep(0.5)
82                self.assertEqual(self.answers(app, roots), [CHILDREN] * ROOTS)
83                self.assertLess(
84                    slowest, 5, "status and tree calls must stay responsive"
85                )
86                # Fast kernels can finish before the first status poll. Their
87                # completed Python outputs prove every child actually booted.
88                completed = [
89                    record
90                    for record in provider.requests
91                    if record["request"]["messages"][-1].get("role") == "tool"
92                    and "cell ran"
93                    in record["request"]["messages"][-1].get("content", "")
94                ]
95                self.assertEqual(len(completed), ROOTS * CHILDREN)
96                with app.api("/server") as response:
97                    self.assertEqual(json.load(response)["state"], "ready")
98        finally:
99            provider.close()
100
101    @exclusive
102    def test_closed_child_does_not_start_when_its_kernel_finishes_booting(self):
103        def script(request):
104            if request["messages"][-1].get("role") != "user":
105                return text("done")
106            if latest_user(request) == "spawn held child":
107                return python(
108                    "models = await agents.models()\n"
109                    "child = await agents.self.spawn('held task', name='held', model=models[0])\n"
110                    "print(child.id)"
111                )
112            if latest_user(request) == "close held child":
113                return python(
114                    "await child.close()\nprint((await mail.submit(child, 'late task')).status)"
115                )
116            return text("done")
117
118        def prepare(app):
119            # The parent boots normally; creating this FIFO holds later launches.
120            wrapper = app.home / "bin" / "python3"
121            wrapper.parent.mkdir()
122            gate = app.workspace / "boot-gate"
123            wrapper.write_text(
124                f"#!{sys.executable}\n"
125                "import os, sys\n"
126                "from pathlib import Path\n"
127                f"gate = Path({str(gate)!r})\n"
128                "if ('launch' in sys.argv or 'start' in sys.argv) and gate.exists():\n"
129                f"    Path({str(app.workspace / 'boot-ready')!r}).touch()\n"
130                "    with gate.open('rb') as reader: reader.read(1)\n"
131                f"os.execv({sys.executable!r}, [{sys.executable!r}, *sys.argv[1:]])\n"
132            )
133            wrapper.chmod(0o700)
134            app.env["PATH"] = str(wrapper.parent) + os.pathsep + app.env["PATH"]
135
136        provider = Provider(script)
137        self.addCleanup(provider.close)
138        with Albedo(provider, prepare=prepare) as app:
139            parent = app.session()
140            app.prompt(parent, "prime parent kernel").close()
141            app.idle(parent)
142            gate = app.workspace / "boot-gate"
143            os.mkfifo(gate)
144            app.prompt(parent, "spawn held child").close()
145            app.idle(parent)
146            with app.api(f"/sessions?scope=all&family_id={parent}") as response:
147                child = next(
148                    item["id"]
149                    for item in json.load(response)["items"]
150                    if item["parent_id"] == parent
151                )
152            with sqlite3.connect(app.home / "albedo.sqlite") as database:
153                self.assertEqual(
154                    database.execute(
155                        "SELECT count(*) FROM mail WHERE recipient=? AND delivered_at IS NULL",
156                        (child,),
157                    ).fetchone()[0],
158                    1,
159                )
160            self.assertEqual(app.history(child)["items"], [])
161            wait_until(
162                (app.workspace / "boot-ready").exists, 10, "child did not start booting"
163            )
164            app.prompt(parent, "close held child").close()
165            app.idle(parent)
166            wait_until(
167                lambda: release_fifo(gate), 10, "child did not reach its boot gate"
168            )
169            gate.unlink()
170            # Wait for the child's preparation to finish, not an arbitrary sleep.
171            app.idle(child)
172            self.assertEqual(app.history(child)["items"], [])
173            with sqlite3.connect(app.home / "albedo.sqlite") as database:
174                self.assertIsNotNone(
175                    database.execute(
176                        "SELECT closed_at FROM session_family WHERE session=?",
177                        (child,),
178                    ).fetchone()[0]
179                )
180                self.assertEqual(
181                    database.execute(
182                        "SELECT count(*) FROM mail WHERE recipient=? AND delivered_at IS NULL",
183                        (child,),
184                    ).fetchone()[0],
185                    2,
186                )
187            self.assertFalse(
188                any(
189                    "held task" in latest_user(record["request"])
190                    for record in provider.requests
191                )
192            )
193            # A manual chat can still visit a closed child without consuming its letters.
194            app.prompt(child, "manual visit").close()
195            app.idle(child)
196            self.assertTrue(
197                any(
198                    latest_user(record["request"]) == "manual visit"
199                    for record in provider.requests
200                )
201            )
202            # Closing one child must not prevent a later sibling from starting.
203            sibling = operation_id()
204            with app.api(
205                f"/sessions/{sibling}",
206                {
207                    "kind": "child",
208                    "parent_id": parent,
209                    "address": "sibling",
210                    "name": "sibling",
211                    "initial_input_id": operation_id(),
212                    "task": "normal task",
213                },
214                method="PUT",
215                headers={"If-None-Match": "*"},
216            ) as response:
217                self.assertEqual(json.load(response)["id"], sibling)
218            app.idle(sibling)
219            self.assertTrue(
220                any(
221                    "normal task" in latest_user(record["request"])
222                    for record in provider.requests
223                )
224            )
225
226    def test_explicit_session_provider_binds_and_unknown_profile_is_rejected(self):
227        provider = Provider(lambda _request: text("ok"))
228        try:
229            with Albedo(provider) as app:
230                session = app.session()
231                with app.api("/sessions") as response:
232                    info = next(
233                        item
234                        for item in json.load(response)["items"]
235                        if item["id"] == session
236                    )
237                self.assertEqual(info["provider_profile"], app.profile)
238                with self.assertRaises(urllib.error.HTTPError) as rejected:
239                    app.api(
240                        f"/sessions/{operation_id()}",
241                        {
242                            "kind": "new",
243                            "workspace": str(app.workspace),
244                            "provider_profile": "not-configured",
245                        },
246                        method="PUT",
247                        headers={"If-None-Match": "*"},
248                    ).close()
249                self.assertEqual(rejected.exception.code, 400)
250        finally:
251            provider.close()
252
253    @staticmethod
254    def answers(app, roots):
255        counts = []
256        for root in roots:
257            with app.api(f"/sessions/{root}/history?limit=200") as response:
258                items = json.load(response)["items"]
259            counts.append(
260                sum(
261                    item["kind"] == "user"
262                    and (item.get("mail") or {}).get("kind") == "unreviewed"
263                    for item in items
264                )
265            )
266        return counts