okami.mom/albedo

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

nekomimi.petfeat: multiple images, collapsed pastes, and an attachment strip in the prompt6e82811

main
branchesmainall refs →
12.5 KiB327 linesraw
1"""Admission and consumption failures leave recoverable, atomic identified decisions."""
2
3import base64
4import json
5import random
6import struct
7import zlib
8import sqlite3
9import time
10import urllib.error
11
12from harness import exclusive, operation_id
13from image_limits_test import png
14from operations_test import InputScenario
15
16SEVEN_DAYS_MS = 7 * 24 * 60 * 60 * 1000
17
18
19class OperationLifecycleTests(InputScenario):
20    def database(self):
21        return sqlite3.connect(self.app.home / "albedo.sqlite", timeout=10)
22
23    def create(self, session_id, intent):
24        with self.app.api(
25            f"/sessions/{session_id}",
26            intent,
27            method="PUT",
28            headers={"If-None-Match": "*"},
29        ) as response:
30            return json.load(response)
31
32    @exclusive
33    def test_creation_failure_rolls_back_session_and_receipt(self):
34        session_id = operation_id()
35        intent = {
36            "kind": "new",
37            "workspace": str(self.app.workspace),
38            "provider_profile": self.app.profile,
39        }
40        with self.database() as database:
41            before = database.execute("SELECT count(*) FROM sessions").fetchone()[0]
42            database.execute(
43                "CREATE TRIGGER fail_creation BEFORE INSERT ON sessions BEGIN SELECT RAISE(ABORT,'fixture'); END"
44            )
45        with self.assertRaises(urllib.error.HTTPError):
46            self.create(session_id, intent)
47        with self.database() as database:
48            self.assertEqual(
49                database.execute("SELECT count(*) FROM sessions").fetchone()[0], before
50            )
51            self.assertEqual(
52                database.execute(
53                    "SELECT count(*) FROM operations WHERE id=?", (session_id,)
54                ).fetchone()[0],
55                0,
56            )
57            database.execute("DROP TRIGGER fail_creation")
58        created = self.create(session_id, intent)
59        self.assertEqual(created["id"], session_id)
60        with self.app.api(f"/sessions/{session_id}") as response:
61            self.assertEqual(json.load(response)["creation"], created["creation"])
62
63    @exclusive
64    def test_admission_failure_can_retry_original_id(self):
65        session = self.app.session()
66        input_id = operation_id()
67        intent = {"kind": "message", "text": "durable input"}
68        with self.database() as database:
69            database.execute(
70                "CREATE TRIGGER fail_admission BEFORE INSERT ON pending_inputs BEGIN SELECT RAISE(ABORT,'fixture'); END"
71            )
72        with self.assertRaises(urllib.error.HTTPError):
73            self.put_input(session, input_id, intent)
74        with self.database() as database:
75            self.assertEqual(
76                database.execute(
77                    "SELECT count(*) FROM operations WHERE id=?", (input_id,)
78                ).fetchone()[0],
79                0,
80            )
81            self.assertEqual(
82                database.execute(
83                    "SELECT count(*) FROM transcript WHERE session=?", (session,)
84                ).fetchone()[0],
85                0,
86            )
87            database.execute("DROP TRIGGER fail_admission")
88        accepted = self.put_input(session, input_id, intent)
89        self.committed(session, input_id)
90        self.app.idle(session)
91        self.assertEqual(
92            self.put_input(session, input_id, intent)["acceptance_order"],
93            accepted["acceptance_order"],
94        )
95        self.assertEqual(
96            [entry["input_id"] for entry in self.users(session)], [input_id]
97        )
98
99    @exclusive
100    def test_consumption_failure_keeps_pending_input_until_storage_recovers(self):
101        session = self.app.session()
102        input_id = operation_id()
103        intent = {
104            "kind": "message",
105            "text": "recover after storage repair",
106            "images": [
107                {
108                    "mime_type": "image/png",
109                    "data": base64.b64encode(png(2, 2)).decode(),
110                }
111            ],
112        }
113        with self.database() as database:
114            database.execute(
115                "CREATE TRIGGER fail_consumption BEFORE INSERT ON submission_events BEGIN SELECT RAISE(ABORT,'fixture'); END"
116            )
117        accepted = self.put_input(session, input_id, intent)
118        self.assertEqual(accepted["admission"], "accepted")
119        self.assertEqual(self.receipt(session, input_id)["delivery"], "pending")
120        with self.database() as database:
121            self.assertEqual(
122                database.execute(
123                    "SELECT count(*) FROM pending_inputs WHERE operation_id=?",
124                    (input_id,),
125                ).fetchone()[0],
126                1,
127            )
128            self.assertEqual(
129                database.execute(
130                    "SELECT count(*) FROM transcript WHERE session=?", (session,)
131                ).fetchone()[0],
132                0,
133            )
134            self.assertEqual(
135                database.execute("SELECT count(*) FROM images").fetchone()[0], 0
136            )
137            self.assertEqual(
138                database.execute("SELECT count(*) FROM submission_events").fetchone()[
139                    0
140                ],
141                0,
142            )
143        self.app.restart(crash=True)
144        self.assertEqual(
145            self.put_input(session, input_id, intent)["acceptance_order"],
146            accepted["acceptance_order"],
147        )
148        self.assertEqual(self.receipt(session, input_id)["delivery"], "pending")
149        with self.database() as database:
150            database.execute("DROP TRIGGER fail_consumption")
151        self.committed(session, input_id)
152        self.app.idle(session)
153        self.app.restart(crash=True)
154        self.assertEqual(self.receipt(session, input_id)["delivery"], "committed")
155        self.conflict(session, input_id, {**intent, "text": "changed after restart"})
156        self.assertEqual(
157            [entry["input_id"] for entry in self.users(session)], [input_id]
158        )
159
160    @exclusive
161    def test_repeated_image_commits_grow_only_compact_metadata(self):
162        session = self.app.session()
163        width, height = 256, 256
164
165        def chunk(kind, data):
166            body = kind + data
167            return (
168                struct.pack(">I", len(data))
169                + body
170                + struct.pack(">I", zlib.crc32(body))
171            )
172
173        pixels = random.Random(1).randbytes(width * height * 3)
174        rows = b"".join(
175            b"\x00" + pixels[start : start + width * 3]
176            for start in range(0, len(pixels), width * 3)
177        )
178        image = (
179            b"\x89PNG\r\n\x1a\n"
180            + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
181            + chunk(b"IDAT", zlib.compress(rows))
182            + chunk(b"IEND", b"")
183        )
184        attachment = {
185            "mime_type": "image/png",
186            "data": base64.b64encode(image).decode(),
187        }
188
189        def durable_bytes():
190            with self.database() as database:
191                return sum(
192                    database.execute(
193                        f"SELECT coalesce(sum(length({column})),0) FROM {table}"
194                    ).fetchone()[0]
195                    for table, column in (
196                        ("images", "data"),
197                        ("transcript", "payload"),
198                        ("submission_events", "payload"),
199                        ("continuation_markers", "payload"),
200                        ("pending_inputs", "payload"),
201                        ("operations", "response"),
202                    )
203                )
204
205        inputs, sizes = [], []
206        for index in range(3):
207            input_id = operation_id()
208            inputs.append(input_id)
209            self.put_input(
210                session,
211                input_id,
212                {
213                    "kind": "message",
214                    "text": "same image",
215                    "client_id": f"image-echo-{index}",
216                    "images": [attachment],
217                },
218            )
219            self.committed(session, input_id)
220            self.app.idle(session)
221            sizes.append(durable_bytes())
222        self.assertLess(sizes[2] - sizes[0], len(image))
223        self.app.restart(crash=True)
224        users = self.users(session)
225        self.assertEqual([entry["input_id"] for entry in users], inputs)
226        for index, input_id in enumerate(inputs):
227            self.assertEqual(
228                self.receipt(session, input_id)["client_id"], f"image-echo-{index}"
229            )
230        metadata = [
231            part["image"]
232            for entry in users
233            for part in entry["content"]
234            if part["kind"] == "image"
235        ]
236        self.assertEqual(
237            [
238                (item["width"], item["height"], item["original_bytes"])
239                for item in metadata
240            ],
241            [(width, height, len(image))] * 3,
242        )
243
244    def test_expired_and_future_ids_cannot_create_work(self):
245        now = int(time.time() * 1000)
246        for timestamp, status, code in (
247            (now - SEVEN_DAYS_MS - 10000, 410, "identity_expired"),
248            (now + 6 * 60 * 1000, 400, "identity_future"),
249        ):
250            session_id = operation_id(timestamp)
251            with self.assertRaises(urllib.error.HTTPError) as failure:
252                self.create(
253                    session_id, {"kind": "new", "workspace": str(self.app.workspace)}
254                )
255            self.assertEqual(failure.exception.code, status)
256            self.assertEqual(json.load(failure.exception)["code"], code)
257            with self.database() as database:
258                self.assertEqual(
259                    database.execute(
260                        "SELECT count(*) FROM operations WHERE id=?", (session_id,)
261                    ).fetchone()[0],
262                    0,
263                )
264
265    @exclusive
266    def test_rejected_creation_replays_after_workspace_is_repaired(self):
267        session_id = operation_id()
268        workspace = self.app.workspace / "missing"
269        intent = {
270            "kind": "new",
271            "workspace": str(workspace),
272            "provider_profile": self.app.profile,
273        }
274        with self.assertRaises(urllib.error.HTTPError) as original:
275            self.create(session_id, intent)
276        problem = json.load(original.exception)
277        with self.assertRaises(urllib.error.HTTPError) as lookup:
278            self.app.api(f"/sessions/{session_id}")
279        self.assertEqual(
280            json.load(lookup.exception)["decision"]["admission"], "rejected"
281        )
282        workspace.mkdir()
283        self.app.restart(crash=True)
284        with self.assertRaises(urllib.error.HTTPError) as repeated:
285            self.create(session_id, intent)
286        self.assertEqual(repeated.exception.code, original.exception.code)
287        self.assertEqual(json.load(repeated.exception), problem)
288        with self.database() as database:
289            self.assertEqual(
290                database.execute("SELECT count(*) FROM sessions").fetchone()[0], 0
291            )
292
293    @exclusive
294    def test_failed_cancellation_does_not_discard_pending_input(self):
295        session = self.app.session()
296        input_id = operation_id()
297        with self.database() as database:
298            database.execute(
299                "CREATE TRIGGER hold_consumption BEFORE INSERT ON transcript WHEN NEW.row_class='user' BEGIN SELECT RAISE(ABORT,'fixture'); END"
300            )
301        self.put_input(
302            session,
303            input_id,
304            {"kind": "message", "text": "keep until cancellation commits"},
305        )
306        with self.database() as database:
307            database.execute(
308                "CREATE TRIGGER fail_cancellation BEFORE UPDATE ON operations WHEN NEW.delivery='cancelled' BEGIN SELECT RAISE(ABORT,'fixture'); END"
309            )
310        resource = f"/sessions/{session}/inputs/{input_id}/cancel"
311        with self.assertRaises(urllib.error.HTTPError):
312            self.app.api(resource, {})
313        self.assertEqual(self.receipt(session, input_id)["delivery"], "pending")
314        with self.database() as database:
315            database.execute("DROP TRIGGER fail_cancellation")
316        with self.app.api(resource, {}) as response:
317            self.assertEqual(json.load(response)["result"], "cancelled")
318        self.app.restart(crash=True)
319        self.assertEqual(self.receipt(session, input_id)["delivery"], "cancelled")
320        with self.database() as database:
321            self.assertEqual(
322                database.execute(
323                    "SELECT count(*) FROM pending_inputs WHERE operation_id=?",
324                    (input_id,),
325                ).fetchone()[0],
326                0,
327            )