okami.mom/albedo

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

Aly Raffauffeat(api)!: migrate daemon and clients to protocol 361c43ca

main
branchesmainall refs →
18.1 KiB425 linesraw
1"""The reloadable prompt-cache TTL table: layers, merge precedence, live reload.
2
3The table is daemon-wide state read from files, so this is where its behaviour
4is visible: the shipped default served with every entry, lookups resolving
5through match order, a local override replacing an entry by id and winning
6with a new more-specific one — without a restart — and a malformed file
7keeping the last good table. A turn's usage carries the steps its cached
8count fades through by that table, which the chat footer counts down from and
9which must survive a restart. It writes extensions.json and cache-ttl.json, so
10tests that change those layers are exclusive, on a home of their own.
11Read-only default lookups and authentication share the daemon without overrides.
12
13The remote layer's fetch is background (like the models catalog's), so the
14fixture provider serves the remote table on a loopback URL and the test waits
15for it to land; `/reload` fetches it synchronously. No live network is reached.
16"""
17
18from typing import Any
19import json
20from datetime import datetime
21import os
22import time
23import unittest
24import urllib.error
25import urllib.request
26
27from harness import ROOT, Albedo, Provider, exclusive, text
28
29
30def milliseconds(timestamp):
31    return round(
32        datetime.fromisoformat(timestamp.replace("Z", "+00:00")).timestamp() * 1000
33    )
34
35
36SHIPPED = json.loads((ROOT / "priv" / "cache-ttl.json").read_text())
37
38# `/reload` fetches the models catalog before the TTL table, so it needs a
39# fixture url too — the default would reach the live models.dev.
40MODELS_CATALOG = {
41    "fixture": {
42        "id": "fixture",
43        "name": "Fixture",
44        "models": {
45            "fixture-model": {"id": "fixture-model", "limit": {"context": 8000}}
46        },
47    }
48}
49
50REMOTE_TABLE: dict[str, Any] = {
51    "version": 1,
52    "entries": [
53        # Replaces the shipped deepseek entry by id, in place.
54        {
55            "id": "deepseek",
56            "match": {"host": "api.deepseek.com"},
57            "policy": "fixed",
58            "clock": "response",
59            "evidence": "measured",
60            "source": "fixture",
61            "note": "remote override",
62        },
63        # A new id, which goes before every default entry.
64        {
65            "id": "fixture-gateway",
66            "match": {"host": "gateway.fixture"},
67            "policy": "evict",
68            "survival": {"typical": 600},
69            "evidence": "measured",
70            "source": "fixture",
71            "note": "new remote entry",
72        },
73    ],
74}
75
76LOCAL_OVERRIDE: dict[str, Any] = {
77    "version": 1,
78    "entries": [
79        # Replaces the remote or default deepseek entry by id, in place.
80        {
81            "id": "deepseek",
82            "match": {"host": "api.deepseek.com"},
83            "policy": "fixed",
84            "clock": "request",
85            "evidence": "measured",
86            "source": "local",
87            "note": "local override",
88        },
89        # A new id, more specific than anything shipped, which must win.
90        {
91            "id": "laptop-gateway",
92            "match": {"host": "gateway.laptop"},
93            "policy": "refresh",
94            "tiers": [{"seconds": 600, "write": 1.5}],
95            "evidence": "measured",
96            "source": "local",
97            "note": "new local entry",
98        },
99    ],
100}
101
102
103def write_atomic(path, payload):
104    temporary = path.with_suffix(".tmp")
105    temporary.write_text(json.dumps(payload))
106    temporary.replace(path)
107
108
109# exclusive: configures daemon-wide TTL/catalog layers and writes cache files
110@exclusive
111class CacheTtlTests(unittest.TestCase):
112    def setUp(self):
113        self.provider = Provider(lambda _request: text("ok"), catalog=REMOTE_TABLE)
114        self.addCleanup(self.provider.close)
115        self.catalog = Provider(lambda _request: text("ok"), catalog=MODELS_CATALOG)
116        self.addCleanup(self.catalog.close)
117
118        def prepare(app):
119            app.write_extensions(
120                {
121                    "models": {"url": self.catalog.url + "/models.json"},
122                    "cacheTtl": {
123                        "url": self.provider.url + "/cache-ttl.json",
124                        "refreshHours": 24,
125                    },
126                }
127            )
128
129        self.app = Albedo(
130            self.provider,
131            prepare=prepare,
132            providers={
133                "fixture": {
134                    "baseUrl": self.provider.url,
135                    "apiKey": "fixture-key",
136                    "model": "fixture-model",
137                    "protocol": "chat_completions",
138                },
139            },
140        )
141        self.app.__enter__()
142        self.addCleanup(self.app.__exit__, None, None, None)
143
144    def table(self):
145        with self.app.api("/models?view=cache-policy&limit=200") as response:
146            return json.load(response)
147
148    def resolved(self, query):
149        with self.app.api("/models?view=cache-policy" + query) as response:
150            return json.load(response)["matched"]
151
152    def layers(self, table):
153        return {layer["name"]: layer for layer in table["layers"]}
154
155    def entries(self, table):
156        return {entry["id"]: entry for entry in table["entries"]}
157
158    def test_local_override_replaces_and_adds_without_restart(self):
159        write_atomic(self.app.home / "cache-ttl.json", LOCAL_OVERRIDE)
160        table = self.table()
161        local = self.layers(table)["local"]
162        self.assertTrue(local["loaded"])
163        self.assertEqual(local["path"], str(self.app.home / "cache-ttl.json"))
164        entries = self.entries(table)
165        # Replaced by id: same entry, now from the local layer.
166        self.assertEqual(entries["deepseek"]["layer"], "local")
167        self.assertEqual(entries["deepseek"]["note"], "local override")
168        self.assertEqual(entries["deepseek"]["clock"], "request")
169        # New ids go before all earlier-layer entries and win the lookup.
170        self.assertEqual(table["entries"][0]["id"], "laptop-gateway")
171        self.assertEqual(self.resolved("&host=gateway.laptop")["id"], "laptop-gateway")
172        # The replaced default keeps answering, with the override's values.
173        self.assertEqual(
174            self.resolved("&host=api.deepseek.com")["note"], "local override"
175        )
176        # Untouched defaults still resolve.
177        self.assertEqual(
178            self.resolved("&extension=claude")["id"], "claude-subscription"
179        )
180
181    def test_three_layers_keep_replacements_in_place_and_new_ids_first(self):
182        write_atomic(self.app.home / "cache-ttl-remote.json", REMOTE_TABLE)
183        gateway = {
184            **REMOTE_TABLE["entries"][1],
185            "match": {"host": "gateway.*"},
186            "note": "local gateway replacement",
187        }
188        local = {"entries": [*LOCAL_OVERRIDE["entries"], gateway]}
189        write_atomic(self.app.home / "cache-ttl.json", local)
190        table = self.table()
191        self.assertEqual(
192            [entry["id"] for entry in table["entries"]],
193            ["laptop-gateway", "fixture-gateway"]
194            + [entry["id"] for entry in SHIPPED["entries"]],
195        )
196        self.assertTrue(self.layers(table)["remote"]["loaded"])
197        self.assertEqual(self.entries(table)["fixture-gateway"]["layer"], "local")
198        self.assertEqual(self.entries(table)["deepseek"]["note"], "local override")
199        # A fresh local rule wins before the replaced, more general gateway rule.
200        self.assertEqual(self.resolved("&host=gateway.laptop")["id"], "laptop-gateway")
201        self.assertEqual(
202            self.resolved("&host=gateway.fixture")["note"],
203            "local gateway replacement",
204        )
205        self.assertEqual(self.resolved("&host=api.deepseek.com")["clock"], "request")
206
207    def test_duplicate_ids_reject_the_layer_and_retain_last_good_entries(self):
208        duplicate = {"entries": [LOCAL_OVERRIDE["entries"][0]] * 2}
209        write_atomic(self.app.home / "cache-ttl.json", duplicate)
210        table = self.table()
211        self.assertFalse(self.layers(table)["local"]["loaded"])
212        failure = self.layers(table)["local"]["error"]
213        self.assertEqual(set(failure), {"code", "detail"})
214        self.assertEqual(failure["code"], "cache_policy_layer_failed")
215        self.assertIn("duplicate id: deepseek", failure["detail"])
216        self.assertLessEqual(len(failure["detail"]), 4096)
217        self.assertFalse(any(entry["layer"] == "local" for entry in table["entries"]))
218        write_atomic(self.app.home / "cache-ttl.json", LOCAL_OVERRIDE)
219        self.assertTrue(self.layers(self.table())["local"]["loaded"])
220        write_atomic(self.app.home / "cache-ttl.json", duplicate)
221        for _ in range(3):
222            table = self.table()
223            self.assertFalse(self.layers(table)["local"]["loaded"])
224            self.assertIn(
225                "duplicate id: deepseek", self.layers(table)["local"]["error"]["detail"]
226            )
227            self.assertEqual(self.entries(table)["deepseek"]["note"], "local override")
228        write_atomic(self.app.home / "cache-ttl.json", LOCAL_OVERRIDE)
229        self.assertTrue(self.layers(self.table())["local"]["loaded"])
230
231    def test_duplicate_remote_ids_are_rejected_before_replacement(self):
232        session = self.app.session()
233        with self.app.api(f"/sessions/{session}/reload", {"target": "both"}):
234            pass
235        path = self.app.home / "cache-ttl-remote.json"
236        before = path.read_bytes()
237        self.provider.catalog = {"entries": [REMOTE_TABLE["entries"][0]] * 2}
238        with self.app.api(
239            f"/sessions/{session}/reload", {"target": "both"}
240        ) as response:
241            outcome = json.load(response)["cache_policy"]
242        self.assertEqual(outcome["state"], "failed")
243        self.assertIn("duplicate id", outcome["failure"]["detail"])
244        self.assertEqual(path.read_bytes(), before)
245        self.assertEqual(
246            self.entries(self.table())["deepseek"]["note"], "remote override"
247        )
248
249    def test_same_revision_remote_replacement_invalidates_parsed_layer(self):
250        session = self.app.session()
251        path = self.app.home / "cache-ttl-remote.json"
252        write_atomic(path, REMOTE_TABLE)
253        # A future mtime prevents background refresh. Restoring it after reload
254        # keeps the cached revision unchanged without depending on scheduling.
255        replacement = json.loads(json.dumps(REMOTE_TABLE))
256        replacement["entries"][0]["note"] = "remote revised!"
257        self.assertEqual(len(json.dumps(replacement)), len(json.dumps(REMOTE_TABLE)))
258        self.provider.catalog = replacement
259
260        timestamp = int(time.time()) + 60
261        os.utime(path, (timestamp, timestamp))
262        self.assertEqual(
263            self.entries(self.table())["deepseek"]["note"], "remote override"
264        )
265        with self.app.api(f"/sessions/{session}/reload", {"target": "both"}):
266            pass
267        os.utime(path, (timestamp, timestamp))
268        self.assertEqual(
269            self.entries(self.table())["deepseek"]["note"], "remote revised!"
270        )
271
272    def test_invalid_override_shadows_previous_entry_before_decoding(self):
273        write_atomic(
274            self.app.home / "cache-ttl.json",
275            {"entries": [{**LOCAL_OVERRIDE["entries"][0], "policy": "invalid"}]},
276        )
277        table = self.table()
278        self.assertTrue(self.layers(table)["local"]["loaded"])
279        self.assertNotIn("deepseek", self.entries(table))
280        self.assertIsNone(self.resolved("&host=api.deepseek.com"))
281        self.assertEqual(
282            self.resolved("&extension=claude")["id"], "claude-subscription"
283        )
284
285    def test_malformed_local_file_keeps_the_last_good_table(self):
286        write_atomic(self.app.home / "cache-ttl.json", LOCAL_OVERRIDE)
287        self.assertTrue(self.layers(self.table())["local"]["loaded"])
288        (self.app.home / "cache-ttl.json").write_text("{ this is not json")
289        table = self.table()
290        local = self.layers(table)["local"]
291        self.assertFalse(local["loaded"])
292        self.assertEqual(set(local["error"]), {"code", "detail"})
293        self.assertEqual(local["error"]["code"], "cache_policy_layer_failed")
294        self.assertIn("not valid JSON", local["error"]["detail"])
295        self.assertLessEqual(len(local["error"]["detail"]), 4096)
296        entries = self.entries(table)
297        self.assertEqual(entries["deepseek"]["layer"], "local")
298        self.assertEqual(entries["deepseek"]["note"], "local override")
299        self.assertEqual(self.resolved("&host=gateway.laptop")["id"], "laptop-gateway")
300        # Fixing the file is picked up live, without a restart.
301        write_atomic(self.app.home / "cache-ttl.json", LOCAL_OVERRIDE)
302        self.assertTrue(self.layers(self.table())["local"]["loaded"])
303
304    def test_remote_url_is_fetched_and_merged(self):
305        deadline = time.monotonic() + 20
306        table = self.table()
307        while (
308            not self.layers(table)["remote"]["loaded"] and time.monotonic() < deadline
309        ):
310            time.sleep(0.1)
311            table = self.table()
312        remote = self.layers(table)["remote"]
313        self.assertTrue(remote["loaded"])
314        self.assertEqual(remote["path"], str(self.app.home / "cache-ttl-remote.json"))
315        entries = self.entries(table)
316        # Replaced by id, in the default entry's place, ahead of general rules.
317        ids = [entry["id"] for entry in table["entries"]]
318        self.assertEqual(entries["deepseek"]["layer"], "remote")
319        self.assertEqual(entries["deepseek"]["note"], "remote override")
320        self.assertEqual(entries["deepseek"]["policy"], "fixed")
321        self.assertLess(ids.index("claude-subscription"), ids.index("deepseek"))
322        # The remote layer's new id sits in front of the default entries.
323        self.assertEqual(ids[0], "fixture-gateway")
324        self.assertEqual(
325            self.resolved("&host=gateway.fixture")["id"], "fixture-gateway"
326        )
327
328    def test_session_reload_refetches_the_remote_copy(self):
329        session = self.app.session()
330        self.assertFalse((self.app.home / "cache-ttl-remote.json").exists())
331        with self.app.api(
332            f"/sessions/{session}/reload", {"target": "both"}
333        ) as response:
334            outcome = json.load(response)
335        self.assertEqual(outcome["session"]["state"], "applied")
336        self.assertEqual(outcome["cache_policy"]["state"], "refreshed")
337        table = self.table()
338        self.assertTrue(self.layers(table)["remote"]["loaded"])
339        self.assertEqual(self.entries(table)["deepseek"]["note"], "remote override")
340
341    def test_a_turns_usage_says_when_its_cache_fades_across_a_restart(self):
342        # A provider that only evicts: the count is unknown past typical
343        # survival and gone past the bound, counted from the response's end.
344        write_atomic(
345            self.app.home / "cache-ttl.json",
346            {
347                "version": 1,
348                "entries": [
349                    {
350                        "id": "fixture-evict",
351                        "match": {"host": "127.0.0.1"},
352                        "policy": "evict",
353                        "survival": {"typical": 600, "max": 3600},
354                        "evidence": "measured",
355                    }
356                ],
357            },
358        )
359        session = self.app.session()
360        self.app.prompt(session, "hello").close()
361        self.app.idle(session)
362        with self.app.api(f"/sessions/{session}/context?view=requests") as response:
363            finished = milliseconds(json.load(response)["items"][0]["ended_at"])
364        fading = [(finished + 600_000, None), (finished + 3_600_000, 0)]
365
366        def latest_fade():
367            with self.app.api(f"/sessions/{session}?tail=0") as response:
368                fade = json.load(response)["usage"]["cache_fade"]
369            return [(milliseconds(step["at"]), step["cached_tokens"]) for step in fade]
370
371        self.assertEqual(latest_fade(), fading)
372        # The steps are stored with the usage, so a client attaching to a
373        # restarted daemon still counts down from the same call.
374        self.app.restart()
375        self.assertEqual(latest_fade(), fading)
376
377
378class DefaultCacheTtlTests(unittest.TestCase):
379    table = CacheTtlTests.table
380    resolved = CacheTtlTests.resolved
381    layers = CacheTtlTests.layers
382    entries = CacheTtlTests.entries
383
384    def setUp(self):
385        self.app = Albedo()
386        self.app.__enter__()
387        self.addCleanup(self.app.__exit__, None, None, None)
388
389    def test_shipped_table_is_served_without_overrides(self):
390        table = self.table()
391        self.assertTrue(self.layers(table)["default"]["loaded"])
392        self.assertIn("priv/cache-ttl.json", self.layers(table)["default"]["path"])
393        self.assertFalse(self.layers(table)["local"]["loaded"])
394        # Every shipped entry, each tagged with the layer it came from.
395        self.assertEqual(
396            set(self.entries(table)), {e["id"] for e in SHIPPED["entries"]}
397        )
398        for entry in table["entries"]:
399            self.assertEqual(entry["layer"], "default")
400        shipped = {entry["id"]: entry for entry in SHIPPED["entries"]}
401        for entry in table["entries"]:
402            for field, value in shipped[entry["id"]].items():
403                if field == "match":
404                    expected = {
405                        name: ([pattern] if isinstance(pattern, str) else pattern)
406                        for name, pattern in value.items()
407                    }
408                    for name in ("extension", "host", "model"):
409                        self.assertEqual(entry[field][name], expected.get(name))
410                elif field == "survival":
411                    for name, number in value.items():
412                        self.assertEqual(entry[field][name], number)
413                elif field == "tiers":
414                    for observed, tier in zip(entry[field], value, strict=True):
415                        for name, number in tier.items():
416                            self.assertEqual(observed[name], number)
417                else:
418                    self.assertEqual(entry[field], value)
419        self.assertIsNone(self.resolved("&extension=never-heard-of-it"))
420
421    def test_route_requires_authentication(self):
422        request = urllib.request.Request(self.app.base + "/models?view=cache-policy")
423        with self.assertRaises(urllib.error.HTTPError) as refused:
424            urllib.request.urlopen(request, timeout=20)
425        self.assertEqual(refused.exception.code, 401)