okami.mom/albedo

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

dawnimages: refuse a PNG or JPEG cut short where it is attached, so it never reaches history79b55b3

main
branchesmainall refs →
16.5 KiB382 linesraw
1"""Session setup, ingress failures, workspace repair, and image inspection."""
2
3import base64
4import json
5from pathlib import Path
6import sqlite3
7import time
8import urllib.error
9import urllib.parse
10
11from harness import exclusive, operation_id
12
13from image_limits_test import png
14from integration_support import IntegrationScenario
15
16
17class IntegrationTest(IntegrationScenario):
18    # exclusive: changes scheduler timing and installs a global database trigger
19    @exclusive
20    def test_scheduler_reports_a_failed_occurrence_write_and_retries(self):
21        def prepare(app):
22            app.env["ALBEDO_SCHEDULE_TICK_MS"] = "50"
23
24        app = self.app_for("chat_completions", prepare=prepare)
25        session = app.session()
26        with sqlite3.connect(app.home / "albedo.sqlite") as db:
27            db.execute(
28                "CREATE TRIGGER reject_advance BEFORE DELETE ON schedules "
29                "BEGIN SELECT RAISE(FAIL, 'fixture advance failed'); END"
30            )
31            db.execute(
32                "INSERT INTO schedules(session,kind,prompt,next_at) VALUES(?,'once',?,0)",
33                (session, "scheduled fixture"),
34            )
35        deadline = time.monotonic() + 15
36        while "fixture advance failed" not in (app.home / "daemon.log").read_text():
37            self.assertLess(time.monotonic(), deadline)
38            time.sleep(0.05)
39        self.assertIn(
40            "schedule occurrence could not be advanced",
41            (app.home / "daemon.log").read_text(),
42        )
43        with sqlite3.connect(app.home / "albedo.sqlite") as db:
44            self.assertEqual(
45                db.execute("SELECT COUNT(*) FROM schedules").fetchone()[0], 1
46            )
47            db.execute("DROP TRIGGER reject_advance")
48        while True:
49            with sqlite3.connect(app.home / "albedo.sqlite") as db:
50                if db.execute("SELECT COUNT(*) FROM schedules").fetchone()[0] == 0:
51                    break
52            self.assertLess(time.monotonic(), deadline)
53            time.sleep(0.05)
54        app.idle(session)
55        self.assertTrue(self.provider.requests)
56
57    # exclusive: renames the global session_family table
58    @exclusive
59    def test_deletion_refuses_a_failed_child_lookup(self):
60        app = self.app_for("chat_completions")
61        session = app.session()
62        with app.api(f"/sessions/{session}?view=configuration") as response:
63            configuration = json.load(response)
64            etag = response.headers["ETag"]
65        with sqlite3.connect(app.home / "albedo.sqlite") as db:
66            db.execute("ALTER TABLE session_family RENAME TO unavailable_family")
67        for scope in ("leaf", "subtree"):
68            with self.subTest(scope=scope):
69                with self.assertRaises(urllib.error.HTTPError) as caught:
70                    app.api(
71                        f"/sessions/{session}?view=configuration&scope={scope}"
72                        f"&family_revision={configuration['family_revision']}",
73                        method="DELETE",
74                        headers={"If-Match": etag},
75                    ).close()
76                self.assertGreaterEqual(caught.exception.code, 400)
77                with sqlite3.connect(app.home / "albedo.sqlite") as db:
78                    self.assertIsNotNone(
79                        db.execute(
80                            "SELECT id FROM sessions WHERE id=?", (session,)
81                        ).fetchone()
82                    )
83        with sqlite3.connect(app.home / "albedo.sqlite") as db:
84            db.execute("ALTER TABLE unavailable_family RENAME TO session_family")
85        with app.api(
86            f"/sessions/{session}?view=configuration&scope=subtree"
87            f"&family_revision={configuration['family_revision']}",
88            method="DELETE",
89            headers={"If-Match": etag},
90        ) as response:
91            self.assertEqual(json.load(response)["deleted_ids"], [session])
92
93    # exclusive: tests unconfigured startup and legacy migration
94    @exclusive
95    def test_unconfigured_startup_and_legacy_provider_migration(self):
96        def prepare(app):
97            workspace = app.root / "legacy-workspace"
98            workspace.mkdir(exist_ok=True)
99            with sqlite3.connect(app.home / "albedo.sqlite") as db:
100                db.execute(
101                    "CREATE TABLE sessions(id TEXT PRIMARY KEY,cwd TEXT NOT NULL,model TEXT NOT NULL,protocol TEXT NOT NULL,stage TEXT NOT NULL DEFAULT 'idle')"
102                )
103                for protocol in ("responses", "chat_completions"):
104                    db.execute(
105                        "INSERT INTO sessions(id,cwd,model,protocol) VALUES(?,?,?,?)",
106                        (
107                            "legacy-" + protocol,
108                            str(workspace),
109                            "legacy-model",
110                            protocol,
111                        ),
112                    )
113
114        app = self.app_for("responses", prepare=prepare, providers={})
115        health = self.read(app, "/server")
116        self.assertEqual(health["state"], "ready")
117        unconfigured = self.read(app, "/sessions?scope=all")["items"]
118        self.assertEqual(
119            {info["id"] for info in unconfigured},
120            {"legacy-responses", "legacy-chat_completions"},
121        )
122        for info in unconfigured:
123            self.assertIsNone(info["provider_profile"])
124            snapshot = self.read(app, f"/sessions/{info['id']}?tail=0")
125            self.assertIsNone(snapshot["provider_profile"])
126            resource = snapshot["configuration_resource"]
127            self.assertIsNone(resource["value"]["provider_profile"])
128            self.assertIsNone(self.read(app, resource["url"])["provider_profile"])
129        with self.assertRaises(urllib.error.HTTPError) as caught:
130            app.api(
131                f"/sessions/{operation_id()}",
132                {"kind": "new", "workspace": str(app.workspace)},
133                method="PUT",
134                headers={"If-None-Match": "*"},
135            ).close()
136        self.assertEqual(caught.exception.code, 409)
137        self.assertIn("/login", json.load(caught.exception)["detail"])
138        for protocol in ("responses", "chat_completions"):
139            with self.subTest(protocol=protocol):
140                legacy = "legacy-" + protocol
141                self.settings(app, protocol)
142                config = json.loads((app.home / "config.json").read_text())
143                config["providers"]["default"] = {
144                    "baseUrl": self.provider.url + "/legacy/v1",
145                    "apiKey": "legacy-key",
146                    "model": "legacy-model",
147                    "protocol": protocol,
148                }
149                (app.home / "config.json").write_text(json.dumps(config))
150                start = len(self.provider.requests)
151                self.send(app, legacy, "inspect the legacy workspace")
152                requests = self.records(start)
153                self.assertTrue(requests)
154                self.assertTrue(all("/legacy/v1/" in r["path"] for r in requests))
155                self.assertTrue(
156                    all(r["authorization"] == "Bearer legacy-key" for r in requests)
157                )
158                session = self.default_session(app)
159                info = next(
160                    s
161                    for s in self.read(app, "/sessions?scope=all")["items"]
162                    if s["id"] == session
163                )
164                self.assertEqual(
165                    (info["provider_profile"], info["automatic_name"]),
166                    ("alpha", "new session"),
167                )
168                pid = app.connection["pid"]
169                app.cli("sessions")
170                self.assertEqual(
171                    json.loads((app.home / "daemon.json").read_text())["pid"], pid
172                )
173                self.restart(app)
174                info = next(
175                    s
176                    for s in self.read(app, "/sessions?scope=all")["items"]
177                    if s["id"] == legacy
178                )
179                self.assertEqual(
180                    (info["provider_profile"], info["automatic_name"]),
181                    ("default", "inspect the legacy workspace"),
182                )
183
184    def test_image_metadata_is_verified_and_only_metadata_is_exposed(self):
185        for protocol in ("responses", "chat_completions"):
186            with self.subTest(protocol=protocol):
187                app = self.app_for(protocol)
188                session = app.session()
189                png = "iVBORw0KGgoAAAANSUhEUgAAAAIAAAADCAIAAAA2iEnWAAAAAElFTkSuQmCC"
190                with self.assertRaises(urllib.error.HTTPError) as caught:
191                    app.api(
192                        f"/sessions/{session}/inputs/{operation_id()}",
193                        {
194                            "kind": "message",
195                            "text": "reject this",
196                            "images": [
197                                {"mime_type": "image/png", "data": "not base64"}
198                            ],
199                        },
200                        method="PUT",
201                    ).close()
202                self.assertEqual(caught.exception.code, 400)
203                start = len(self.provider.requests)
204                app.api(
205                    f"/sessions/{session}/inputs/{operation_id()}",
206                    {
207                        "kind": "message",
208                        "text": "describe this image",
209                        "images": [{"mime_type": "image/png", "data": png}],
210                    },
211                    method="PUT",
212                ).close()
213                app.idle(session)
214                self.send(app, session, "confirm image history")
215                requests = self.records(start)
216                self.assertEqual(len(requests), 3)
217                self.assertTrue(
218                    all(
219                        "data:image/png;base64," + png in json.dumps(r["request"])
220                        for r in requests
221                    )
222                )
223                image = next(
224                    part["image"]
225                    for e in self.history(app, session)
226                    for part in e["content"]
227                    if part["kind"] == "image"
228                )
229                self.assertEqual(
230                    (
231                        image["mime_type"],
232                        image["width"],
233                        image["height"],
234                        image["original_bytes"],
235                    ),
236                    ("image/png", 2, 3, 45),
237                )
238                self.assertNotIn(png, json.dumps(image))
239
240    def test_message_images_follow_their_text_in_attachment_order(self):
241        app = self.app_for("responses")
242        session = app.session()
243        first, second = (
244            base64.b64encode(png(width, height)).decode()
245            for width, height in ((2, 3), (4, 5))
246        )
247        upload = [
248            {"mime_type": "image/png", "data": first},
249            {"mime_type": "image/png", "data": second},
250        ]
251        with self.assertRaises(urllib.error.HTTPError) as caught:
252            app.api(
253                f"/sessions/{session}/inputs/{operation_id()}",
254                {"kind": "message", "text": "too many", "images": upload * 6},
255                method="PUT",
256            ).close()
257        self.assertEqual(caught.exception.code, 400)
258        start = len(self.provider.requests)
259        text = "[Image #1] is before, then I wrote this, and [Image #2] is after"
260        app.api(
261            f"/sessions/{session}/inputs/{operation_id()}",
262            {"kind": "message", "text": text, "images": upload},
263            method="PUT",
264        ).close()
265        app.idle(session)
266        request = json.dumps(self.records(start)[0]["request"])
267        positions = [
268            request.index(needle)
269            for needle in (text, "base64," + first, "base64," + second)
270        ]
271        self.assertEqual(positions, sorted(positions))
272        user = next(e for e in self.history(app, session) if e["kind"] == "user")
273        self.assertEqual(self.entry_text(user), text)
274        self.assertEqual(
275            [
276                (part["image"]["width"], part["image"]["height"])
277                for part in user["content"]
278                if part["kind"] == "image"
279            ],
280            [(2, 3), (4, 5)],
281        )
282
283    def test_short_pastes_are_inlined_and_long_ones_saved_for_the_model(self):
284        app = self.app_for("responses")
285        session = app.session()
286        short = "\n".join(f"short line {n}" for n in range(20))
287        long = "\n".join(f"long line {n}" for n in range(150))
288        text = "read [Paste #1, +20 lines] and then [Paste #2, +150 lines] please"
289        start = len(self.provider.requests)
290        input_id = operation_id()
291        app.api(
292            f"/sessions/{session}/inputs/{input_id}",
293            {"kind": "message", "text": text, "pastes": [short, long]},
294            method="PUT",
295        ).close()
296        app.idle(session)
297        saved = app.home / "pastes" / session / f"{input_id}-2.md"
298        self.assertEqual(saved.read_text(), long)
299        request = json.dumps(self.records(start)[0]["request"])
300        sent = f"read {short} and then [Paste #2, 150 lines, saved to {saved}] please"
301        self.assertIn(json.dumps(sent)[1:-1], request)
302        user = next(e for e in self.history(app, session) if e["kind"] == "user")
303        self.assertEqual(self.entry_text(user), text)
304
305    def test_busy_messages_join_next_model_request_in_order(self):
306        app = self.app_for("responses")
307        session = app.session()
308        start = len(self.provider.requests)
309        app.prompt(session, "first task").close()
310        self.wait_for_request(start)
311        app.prompt(session, "queued direction").close()
312        app.prompt(session, "another direction").close()
313        app.idle(session)
314        requests = self.records(start)
315        self.assertEqual(len(requests), 2)
316        first, second = (json.dumps(r["request"]) for r in requests)
317        self.assertNotIn("queued direction", first)
318        self.assertIn("queued direction", second)
319        self.assertIn("another direction", second)
320        self.assertLess(
321            second.index("queued direction"), second.index("another direction")
322        )
323        self.assertTrue(
324            any(
325                e["kind"] == "user" and self.entry_text(e) == "queued direction"
326                for e in self.history(app, session)
327            )
328        )
329
330    # exclusive: restarts the daemon
331    @exclusive
332    def test_workspace_repair_validates_path_and_preserves_history(self):
333        app = self.app_for("responses")
334        old = app.root / "original-workspace"
335        old.mkdir()
336        moved = app.root / "moved-workspace"
337        session = app.session(old)
338        self.send(app, session, "workspace probe before move")
339        self.assertEqual(Path((old / "cwd-probe").read_text()).resolve(), old.resolve())
340        old.rename(moved)
341        pid = app.connection["pid"]
342        before = self.history(app, session)
343        identity = operation_id()
344        app.api(
345            f"/sessions/{session}/inputs/{identity}",
346            {"kind": "message", "text": "workspace probe after move"},
347            method="PUT",
348        ).close()
349        receipt = self.read(app, f"/sessions/{session}/inputs/{identity}")
350        self.assertEqual(receipt["delivery"], "pending")
351        for invalid in ("relative", str(old), str(moved / "cwd-probe")):
352            with self.assertRaises(urllib.error.HTTPError) as caught:
353                self.change(app, session, {"workspace": invalid})
354            self.assertIn(caught.exception.code, (400, 409))
355        self.change(app, session, {"workspace": str(moved)})
356        self.assertEqual(
357            self.read(app, f"/sessions/{session}")["workspace"], str(moved)
358        )
359        self.assertEqual(json.loads((app.home / "daemon.json").read_text())["pid"], pid)
360        app.idle(session, timeout=35)
361        receipt = self.read(app, f"/sessions/{session}/inputs/{identity}")
362        self.assertEqual(receipt["delivery"], "committed")
363        self.assertEqual(
364            Path((moved / "cwd-probe").read_text()).resolve(), moved.resolve()
365        )
366        self.assertFalse(old.exists())
367        after = self.history(app, session)
368        self.assertEqual(
369            [e for e in before if e["kind"] == "user"],
370            [e for e in after if e["kind"] == "user"][:-1],
371        )
372        with sqlite3.connect(app.home / "albedo.sqlite") as db:
373            self.assertEqual(
374                db.execute(
375                    "SELECT cwd FROM sessions WHERE id=?", (session,)
376                ).fetchone()[0],
377                str(moved),
378            )
379        self.restart(app)
380        self.assertEqual(
381            self.read(app, f"/sessions/{session}")["workspace"], str(moved)
382        )