okami.mom/albedo

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

Aly Raffauffix: restore host staging progress and correct configuration597cb84

main
branchesmainall refs →
8.8 KiB233 linesraw
1"""Deleting a session removes every row any table keeps for it.
2
3The delete route clears a list of tables the daemon names itself, so a table an
4extension adds is forgotten until someone remembers to add it. The check reads
5the schema, so it covers a table nobody has listed yet.
6"""
7
8import json
9import sqlite3
10import threading
11import unittest
12import urllib.parse
13import urllib.error
14
15from harness import Albedo, Provider, exclusive, operation_id, text
16
17# Paperclip provenance and retained input outcomes outlive the session.
18KEPT = {"paperclips", "input_turns"}
19
20
21def reply(request):
22    return text("older conversation summary")
23
24
25class SessionDeleteTests(unittest.TestCase):
26    def setUp(self):
27        self.provider = Provider(reply)
28        self.addCleanup(self.provider.close)
29        self.app = Albedo(self.provider, protocol="responses")
30        self.app.__enter__()
31        self.addCleanup(self.app.__exit__, None, None, None)
32
33    def database(self):
34        return sqlite3.connect(
35            f"file:{self.app.home / 'albedo.sqlite'}?mode=ro", uri=True
36        )
37
38    def keyed_tables(self, db):
39        tables = [
40            row[0]
41            for row in db.execute("SELECT name FROM sqlite_master WHERE type='table'")
42        ]
43        return [
44            table
45            for table in tables
46            if table not in KEPT
47            and any(
48                column[1] == "session"
49                for column in db.execute(f"PRAGMA table_info({table})")
50            )
51        ]
52
53    def rows(self, db, table, session):
54        return db.execute(
55            f"SELECT count(*) FROM {table} WHERE session=?", (session,)
56        ).fetchone()[0]
57
58    def compacted_session(self, strategy):
59        session = self.app.session()
60        # Long enough that a forced compaction has something to evict.
61        for prompt in ("first", "second", "third", "fourth"):
62            self.app.prompt(session, prompt + " " + "x" * 6000).close()
63            self.app.idle(session)
64        self.compact(session, strategy)
65        return session
66
67    def compact(self, session, strategy):
68        with self.app.api(
69            f"/sessions/{session}/compaction",
70            {"strategy": strategy},
71        ) as response:
72            self.assertEqual(json.load(response)["state"], "compacted")
73        self.app.idle(session)
74
75    def delete_session(self, session):
76        resource = f"/sessions/{session}?view=configuration"
77        with self.app.api(resource) as response:
78            validator = response.getheader("ETag")
79            response.read()
80        with self.app.api(
81            resource, method="DELETE", headers={"If-Match": validator}
82        ) as response:
83            self.assertEqual(json.load(response)["state"], "complete")
84
85    def assert_delete_clears(self, strategy, table, then=None):
86        session = self.compacted_session(strategy)
87        with self.database() as db:
88            self.assertGreater(self.rows(db, table, session), 0, table)
89            self.assertIn(table, self.keyed_tables(db))
90        if then:
91            self.compact(session, then)
92            with self.database() as db:
93                self.assertGreater(self.rows(db, table, session), 0, table)
94        self.delete_session(session)
95        with self.database() as db:
96            left = {
97                name: self.rows(db, name, session) for name in self.keyed_tables(db)
98            }
99            self.assertEqual({k: v for k, v in left.items() if v}, {})
100            orphans = db.execute(
101                "SELECT count(*) FROM lcm_compaction_edge WHERE child NOT IN "
102                "(SELECT id FROM lcm_compaction_node) OR parent NOT IN "
103                "(SELECT id FROM lcm_compaction_node)"
104            ).fetchone()[0]
105            self.assertEqual(orphans, 0)
106
107    def test_leaf_deletion_refuses_busy_and_stale_observations(self):
108        entered, release = threading.Event(), threading.Event()
109        self.addCleanup(release.set)
110
111        def held(request):
112            entered.set()
113            release.wait(30)
114            return text("completed before deletion")
115
116        self.provider.script = held
117        session = self.app.session()
118        resource = f"/sessions/{session}?view=configuration"
119        with self.app.api(resource) as response:
120            revision = response.headers["ETag"]
121            response.read()
122        identity = operation_id()
123        with self.app.api(
124            f"/sessions/{session}/inputs/{identity}",
125            {"kind": "message", "text": "hold deletion"},
126            method="PUT",
127        ) as response:
128            response.read()
129        self.assertTrue(entered.wait(30))
130        with self.app.api(resource) as response:
131            revision = response.headers["ETag"]
132            response.read()
133        with self.assertRaises(urllib.error.HTTPError) as caught:
134            self.app.api(resource, method="DELETE", headers={"If-Match": revision})
135        self.assertEqual(caught.exception.code, 409)
136        release.set()
137        self.app.idle(session)
138        with self.app.api(resource) as response:
139            revision = response.headers["ETag"]
140            response.read()
141        with self.app.api(
142            resource,
143            {"name": "changed before deletion"},
144            method="PATCH",
145            headers={"If-Match": revision},
146        ) as response:
147            response.read()
148        with self.assertRaises(urllib.error.HTTPError) as caught:
149            self.app.api(resource, method="DELETE", headers={"If-Match": revision})
150        self.assertEqual(caught.exception.code, 412)
151        with self.app.api(f"/sessions/{session}") as response:
152            self.assertEqual(json.load(response)["name"], "changed before deletion")
153        self.delete_session(session)
154        with self.app.api(f"/sessions/{session}/inputs/{identity}") as response:
155            receipt = json.load(response)
156        self.assertEqual(receipt["delivery"], "committed")
157        self.assertEqual(receipt["turn"]["state"], "completed")
158        with self.assertRaises(urllib.error.HTTPError) as caught:
159            self.app.api(f"/sessions/{session}")
160        self.assertEqual(caught.exception.code, 410)
161        self.assertEqual(json.load(caught.exception)["code"], "session_deleted")
162
163    def test_rolling_and_notes_rows_go_with_the_session(self):
164        self.assert_delete_clears("rolling", "compaction_notes")
165
166    def test_lcm_rows_go_with_the_session(self):
167        self.assert_delete_clears("lcm", "lcm_compaction_node")
168
169    def test_rows_of_a_strategy_switched_off_go_too(self):
170        # Switching to rolling disables lcm; its graph still belongs to the session.
171        self.assert_delete_clears("lcm", "lcm_compaction_node", then="rolling")
172
173    def test_work_and_schedule_rows_go_with_the_session(self):
174        session = self.app.session()
175        with self.app.api(
176            "/extensions/schedule/jobs",
177            {
178                "session_id": session,
179                "kind": "once",
180                "delay_seconds": 3600,
181                "prompt": "build",
182            },
183        ) as response:
184            response.read()
185        with self.app.api(
186            "/extensions/work/items?"
187            + urllib.parse.urlencode({"workspace": str(self.app.workspace)}),
188            {"title": "assigned", "session_id": session},
189        ) as response:
190            response.read()
191        with self.database() as db:
192            for table in ("work", "schedules"):
193                self.assertGreater(self.rows(db, table, session), 0, table)
194        self.delete_session(session)
195        with self.database() as db:
196            for table in ("work", "schedules"):
197                self.assertEqual(self.rows(db, table, session), 0, table)
198
199    def test_snapcompact_archive_goes_with_the_session(self):
200        self.assert_delete_clears("snapcompact", "snapcompact_archive")
201
202    @exclusive
203    def test_snapcompact_selection_keeps_archive_across_switch_and_restart(self):
204        self.app.write_extensions({"snapcompact": {"enabled": False}})
205        self.app.restart()
206        session = self.compacted_session("snapcompact")
207
208        def selected():
209            with self.app.api(f"/sessions/{session}?tail=0") as response:
210                return json.load(response)["selection"]["effective"]["extensions"]
211
212        def archive():
213            with self.database() as db:
214                return db.execute(
215                    "SELECT * FROM snapcompact_archive WHERE session=?", (session,)
216                ).fetchall()
217
218        self.assertTrue(selected()["snapcompact"])
219        self.assertFalse(selected()["rolling"])
220        saved = archive()
221        self.assertTrue(saved)
222        self.compact(session, "rolling")
223        self.assertTrue(selected()["rolling"])
224        self.assertFalse(selected()["snapcompact"])
225        self.assertEqual(archive(), saved)
226        self.app.restart()
227        self.assertTrue(selected()["rolling"])
228        self.assertFalse(selected()["snapcompact"])
229        self.assertEqual(archive(), saved)
230
231
232if __name__ == "__main__":
233    unittest.main()