okami.mom/albedo

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

Aly Raffauffix(search): resolve providers from active extension selectionsb82006c

main
branchesmainall refs →
10.2 KiB266 linesraw
1"""The model's web_search must use the user's preferred provider and fall
2back down the order only when one fails, and the /web-search page must
3reorder providers and turn them off for every later search.
4
5Exa is the one provider a test can stand in for: its endpoint is a setting,
6so a local fake answers it. The subscription providers have no sign-in on the
7test daemon, so each fails at once without reaching the network.
8
9The order and the Exa key are global settings, so each scenario has its own
10daemon."""
11
12import ast
13import json
14import threading
15import unittest
16import urllib.error
17from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
18
19from harness import Albedo, Provider, exclusive, python, text
20
21COLLECTION = "/extensions/web-search/providers"
22
23RESULTS = [
24    {
25        "url": "https://gleam.run/news/one",
26        "title": "Gleam one",
27        "highlights": ["first highlight", "second"],
28        "publishedDate": "2026-08-01",
29    },
30    {"url": "https://gleam.run/news/one", "title": "duplicate"},
31    {"url": "https://example.com/two", "title": None, "highlights": None},
32]
33
34SEARCH = (
35    "r = await web_search('latest gleam release', limit=2)\n"
36    "print(repr([r.answer, [(s.title, s.url, s.snippet, s.published)"
37    " for s in r.sources]]))\n"
38    "print(r)"
39)
40
41FAILING = (
42    "try:\n"
43    "    await web_search('latest gleam release')\n"
44    "except Exception as error:\n"
45    "    print('raised:', error)"
46)
47
48
49class FakeExa:
50    """Answers /search with RESULTS and records what each request sent."""
51
52    def __init__(self):
53        self.requests = []
54        owner = self
55
56        class Handler(BaseHTTPRequestHandler):
57            def do_POST(self):
58                body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
59                owner.requests.append((self.path, self.headers["x-api-key"], body))
60                payload = json.dumps({"results": RESULTS}).encode()
61                self.send_response(200)
62                self.send_header("Content-Type", "application/json")
63                self.send_header("Content-Length", str(len(payload)))
64                self.end_headers()
65                self.wfile.write(payload)
66
67            def log_message(self, format, *args):
68                pass
69
70        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
71        self.url = f"http://127.0.0.1:{self.server.server_address[1]}"
72        threading.Thread(target=self.server.serve_forever, daemon=True).start()
73
74    def close(self):
75        self.server.shutdown()
76        self.server.server_close()
77
78
79def scripting(cells):
80    """Answers each user turn with the next scripted python cell."""
81
82    def reply(request):
83        if request["messages"][-1].get("role") == "user" and cells:
84            return python(cells.pop(0))
85        return text("done")
86
87    return Provider(reply)
88
89
90def last_output(app, session):
91    results = [
92        json.loads(part["value"])
93        for entry in app.history(session)["items"]
94        if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python"
95        for part in entry["content"]
96        if part["kind"] == "json" and part["field"] == "result"
97    ]
98    return results[-1]["output"]
99
100
101@exclusive
102class WebSearchTests(unittest.TestCase):
103    def setUp(self):
104        self.exa = FakeExa()
105        self.addCleanup(self.exa.close)
106
107    def start(self, cells):
108        provider = scripting(cells)
109        self.addCleanup(provider.close)
110
111        def prepare(app):
112            app.write_extensions({"exa": {"endpoint": self.exa.url}})
113            app.store_secrets("exa", {"apiKey": "exa-key"})
114
115        app = Albedo(provider, prepare=prepare)
116        app.__enter__()
117        self.addCleanup(app.__exit__, None, None, None)
118        return app
119
120    def ask(self, app, prompt):
121        session = app.session()
122        app.prompt(session, prompt).close()
123        app.idle(session)
124        return session, last_output(app, session)
125
126    def providers(self, app):
127        with app.api(COLLECTION) as response:
128            page = json.load(response)
129        return [item["value"] for item in page["items"]], page["page"]
130
131    def change(self, app, name, change):
132        with app.api(f"{COLLECTION}/{name}", {"change": change}) as response:
133            return json.load(response)["resource"]["value"]
134
135    def patch(self, app, route, body):
136        with app.api(route) as response:
137            etag = response.headers["ETag"]
138        with app.api(
139            route, body, method="PATCH", headers={"If-Match": etag}
140        ) as response:
141            return json.load(response)
142
143    def test_session_provider_selection_is_loaded_and_isolated(self):
144        app = self.start([FAILING, SEARCH, SEARCH])
145        disabled = app.session()
146        self.patch(
147            app,
148            f"/sessions/{disabled}?view=configuration",
149            {"selection": {"extensions": {"exa": False}}},
150        )
151        with app.api(f"/sessions/{disabled}/reload", {"target": "session"}) as response:
152            self.assertEqual(json.load(response)["session"]["state"], "applied")
153        app.prompt(disabled, "search without exa").close()
154        app.idle(disabled)
155        self.assertIn("raised:", last_output(app, disabled))
156        self.assertEqual(self.exa.requests, [])
157        self.ask(app, "another session can use exa")
158        self.assertEqual(len(self.exa.requests), 1)
159        self.patch(
160            app,
161            f"/sessions/{disabled}?view=configuration",
162            {"selection": {"extensions": {"exa": True}}},
163        )
164        with app.api(f"/sessions/{disabled}/reload", {"target": "session"}) as response:
165            self.assertEqual(json.load(response)["session"]["state"], "applied")
166        app.prompt(disabled, "search after applying exa").close()
167        app.idle(disabled)
168        self.assertEqual(len(self.exa.requests), 2)
169
170    def test_global_eligibility_preserves_hidden_provider_preferences(self):
171        app = self.start([])
172        self.change(app, "exa", "up")
173        self.change(app, "exa", "toggle")
174        before = json.loads((app.home / "extensions.json").read_text())["web-search"]
175        self.patch(app, "/settings?group=extensions", {"defaults": {"exa": False}})
176        self.assertNotIn("exa", [entry["name"] for entry in self.providers(app)[0]])
177        self.change(app, "codex", "toggle")
178        saved = json.loads((app.home / "extensions.json").read_text())["web-search"]
179        self.assertEqual(saved["order"], before["order"])
180        self.assertIn("exa", saved["off"])
181        self.patch(app, "/settings?group=extensions", {"defaults": {"exa": True}})
182        exa = next(entry for entry in self.providers(app)[0] if entry["name"] == "exa")
183        self.assertFalse(exa["enabled"])
184        self.assertEqual(exa["position"], before["order"].index("exa") + 1)
185
186    def test_search_passes_over_signed_out_providers_to_exa(self):
187        app = self.start([SEARCH])
188        session, output = self.ask(app, "search the web")
189        answer, sources = ast.literal_eval(output.splitlines()[0])
190        self.assertEqual(answer, "")
191        # One source per url, cut to the limit, highlights as the snippet.
192        self.assertEqual(
193            sources,
194            [
195                (
196                    "Gleam one",
197                    "https://gleam.run/news/one",
198                    "first highlight … second",
199                    "2026-08-01",
200                ),
201                ("https://example.com/two", "https://example.com/two", "", None),
202            ],
203        )
204        self.assertIn("[1] Gleam one (2026-08-01)", output)
205        # The transcript shows the search's result beside the cell.
206        (activity,) = [
207            activity
208            for entry in app.history(session)["items"]
209            if entry["kind"] == "tool_result"
210            for part in entry["content"]
211            if part["kind"] == "trace"
212            for activity in part["trace"]["activities"]
213            if activity["kind"] == "web"
214        ]
215        self.assertEqual(activity["target"], "latest gleam release")
216        self.assertIn("- [Gleam one](https://gleam.run/news/one)", activity["detail"])
217        ((path, key, body),) = self.exa.requests
218        self.assertEqual((path, key), ("/search", "exa-key"))
219        self.assertEqual(
220            (body["query"], body["numResults"]), ("latest gleam release", 2)
221        )
222
223    def test_page_reorders_and_turns_off_providers(self):
224        app = self.start([FAILING])
225        listed, page = self.providers(app)
226        names = [entry["name"] for entry in listed]
227        self.assertEqual(sorted(names), ["antigravity", "claude", "codex", "exa"])
228        self.assertEqual([entry["position"] for entry in listed], [1, 2, 3, 4])
229        self.assertTrue(all(entry["enabled"] for entry in listed))
230        self.assertEqual(
231            [row["id"] for row in page["rows"]], names, "the page lists the same order"
232        )
233
234        index = names.index("exa")
235        moved = self.change(app, "exa", "up")
236        self.assertEqual(moved["position"], index)
237        expected = names[: index - 1] + ["exa", names[index - 1]] + names[index + 1 :]
238        self.assertEqual([entry["name"] for entry in self.providers(app)[0]], expected)
239        settings = json.loads((app.home / "extensions.json").read_text())
240        self.assertEqual(settings["web-search"]["order"], expected)
241
242        # A move past the top changes nothing.
243        top = expected[0]
244        self.assertEqual(self.change(app, top, "up")["position"], 1)
245
246        self.assertFalse(self.change(app, "exa", "toggle")["enabled"])
247        _, output = self.ask(app, "search the web")
248        self.assertIn("raised: every web search provider failed", output)
249        for name in ("codex", "claude", "antigravity"):
250            self.assertIn(name + ":", output)
251        self.assertNotIn("exa:", output)
252        self.assertEqual(self.exa.requests, [], "a provider turned off is never tried")
253
254        for name, body, status in (
255            ("nobody", {"change": "up"}, 404),
256            ("exa", {"change": "sideways"}, 400),
257        ):
258            with self.subTest(name=name, body=body):
259                with self.assertRaises(urllib.error.HTTPError) as refused:
260                    app.api(f"{COLLECTION}/{name}", body).close()
261                refused.exception.close()
262                self.assertEqual(refused.exception.code, status)
263
264
265if __name__ == "__main__":
266    unittest.main()