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 →
16.6 KiB422 linesraw
1"""OpenAI-compatible proxy enablement, model routing, streaming, and state projection."""
2
3import json
4import unittest
5import urllib.error
6import urllib.request
7
8from daemon_test import delayed_request, refused_request
9from harness import Albedo, Provider, exclusive, Reply
10
11TOOLS = [
12    {"type": "function", "function": {"name": "bash", "parameters": {"type": "object"}}}
13]
14OPENING = [
15    {"role": "system", "content": "be brief"},
16    {"role": "user", "content": "list files"},
17]
18CALL = {
19    "id": "call-1",
20    "type": "function",
21    "function": {"name": "bash", "arguments": '{"command":"ls"}'},
22}
23HISTORY = OPENING + [
24    {"role": "assistant", "content": "hel", "tool_calls": [CALL]},
25    {"role": "tool", "tool_call_id": "call-1", "content": "a.txt"},
26]
27
28
29def fixture(request):
30    if "messages" in request:
31        if request["messages"][-1]["role"] == "tool":
32            deltas = [({"role": "assistant", "content": "ran it"}, None), ({}, "stop")]
33        else:
34            deltas = [
35                ({"role": "assistant", "reasoning_content": "hmm"}, None),
36                ({"content": "hel"}, None),
37                (
38                    {
39                        "tool_calls": [
40                            {
41                                "index": 0,
42                                **CALL,
43                                "function": {
44                                    "name": "bash",
45                                    "arguments": '{"command":',
46                                },
47                            }
48                        ]
49                    },
50                    None,
51                ),
52                (
53                    {"tool_calls": [{"index": 0, "function": {"arguments": '"ls"}'}}]},
54                    None,
55                ),
56                ({}, "tool_calls"),
57            ]
58        events = [
59            {
60                "id": "c",
61                "choices": [{"index": 0, "delta": delta, "finish_reason": reason}],
62                **(
63                    {"usage": {"prompt_tokens": 7, "completion_tokens": 3}}
64                    if reason == "tool_calls"
65                    else {}
66                ),
67            }
68            for delta, reason in deltas
69        ]
70    else:
71        last = request["input"][-1]
72        if last.get("role") == "user" and last.get("content") == "call a tool":
73            output = [
74                {
75                    "type": "reasoning",
76                    "id": "rs_1",
77                    "summary": [],
78                    "encrypted_content": "sealed-reasoning",
79                },
80                {
81                    "type": "function_call",
82                    "id": "fc_1",
83                    "call_id": "call_r1",
84                    "name": "bash",
85                    "arguments": '{"command":"ls"}',
86                    "status": "completed",
87                },
88            ]
89        else:
90            output = [
91                {
92                    "type": "message",
93                    "role": "assistant",
94                    "status": "completed",
95                    "content": [
96                        {"type": "output_text", "text": "done", "annotations": []}
97                    ],
98                }
99            ]
100        events = [
101            {
102                "type": "response.completed",
103                "response": {
104                    "id": "r",
105                    "status": "completed",
106                    "output": output,
107                    "usage": {"input_tokens": 5, "output_tokens": 1},
108                },
109            }
110        ]
111    return Reply("raw", events=events)
112
113
114class ProxyTests(unittest.TestCase):
115    def setUp(self):
116        self.provider = Provider(fixture)
117        self.addCleanup(self.provider.close)
118        endpoint = self.provider.url + "/v1"
119        self.chat = f"chat-{self.provider.route}"
120        self.resp = f"resp-{self.provider.route}"
121        self.cx = f"cx-{self.provider.route}"
122        self.broken = f"broken-{self.provider.route}"
123        profiles = {
124            self.chat: {
125                "baseUrl": endpoint,
126                "apiKey": "k",
127                "model": "m-chat",
128                "protocol": "chat_completions",
129            },
130            self.resp: {
131                "baseUrl": endpoint,
132                "apiKey": "k",
133                "model": "m-resp",
134                "protocol": "responses",
135            },
136            self.cx: {
137                "extension": "codex",
138                "model": "gpt-saved",
139                "protocol": "responses",
140            },
141            self.broken: {"model": 3},
142        }
143
144        self.app = Albedo(self.provider, providers=profiles)
145        self.app.__enter__()
146        self.addCleanup(self.app.__exit__, None, None, None)
147        self.base = self.app.base + "/extensions/proxy/v1"
148
149    def get_models(self):
150        return urllib.request.urlopen(
151            urllib.request.Request(
152                self.base + "/models",
153                headers={"Authorization": "Bearer " + self.app.connection["token"]},
154            ),
155            timeout=10,
156        )
157
158    def enable(self):
159        self.app.write_extensions({"enabled": {"proxy": True}})
160
161    def post(self, body, *, headers=None):
162        request = urllib.request.Request(
163            self.base + "/chat/completions",
164            json.dumps(body).encode(),
165            {
166                "Content-Type": "application/json",
167                "Authorization": "Bearer " + self.app.connection["token"],
168                **(headers or {}),
169            },
170        )
171        return urllib.request.urlopen(request, timeout=30)
172
173    def sent(self):
174        record = self.provider.requests[-1]
175        return record["path"], record["request"]
176
177    # exclusive: toggles the global proxy gateway enablement
178    @exclusive
179    def test_enablement_listing_and_profile_errors(self):
180        self.app.write_extensions({"enabled": {"proxy": False}})
181        with self.assertRaises(urllib.error.HTTPError) as rejected:
182            self.get_models()
183        self.assertEqual(rejected.exception.code, 404)
184        self.enable()
185        with self.get_models() as response:
186            listing = json.load(response)
187        self.assertTrue(
188            {f"{self.chat}/m-chat", f"{self.resp}/m-resp", f"{self.cx}/gpt-saved"}
189            <= {model["id"] for model in listing["data"]}
190        )
191        self.assertIn(self.broken, [error["profile"] for error in listing["errors"]])
192        with self.post(
193            {
194                "model": f"{self.chat}/unlisted",
195                "messages": [{"role": "user", "content": "hi"}],
196            }
197        ) as response:
198            self.assertEqual(json.load(response)["model"], f"{self.chat}/unlisted")
199        self.assertEqual(self.sent()[1]["model"], "unlisted")
200        self.assertEqual(self.provider.requests[-1]["authorization"], "Bearer k")
201        for model in (self.broken, f"{self.cx}/gpt-saved"):
202            with (
203                self.subTest(model=model),
204                self.assertRaises(urllib.error.HTTPError) as rejected,
205            ):
206                self.post(
207                    {"model": model, "messages": [{"role": "user", "content": "hi"}]}
208                )
209            self.assertEqual(rejected.exception.code, 400)
210            self.assertIn("/login", json.load(rejected.exception)["error"]["message"])
211        with self.assertRaises(urllib.error.HTTPError) as rejected:
212            self.post({}, headers={"Origin": "https://example.com"})
213        self.assertEqual(rejected.exception.code, 403)
214        with self.assertRaises(urllib.error.HTTPError) as rejected:
215            self.post({"model": "missing/x", "messages": OPENING})
216        self.assertEqual(rejected.exception.code, 400)
217
218    def test_ingress_limits_and_origin_refuse_before_provider_work(self):
219        path = "/extensions/proxy/v1/chat/completions"
220        for framing, expected_status, code in (
221            ({"Transfer-Encoding": "identity"}, 400, "unsupported_transfer_encoding"),
222            ({"Content-Length": "invalid"}, 400, "invalid_request"),
223            ({"Content-Length": "32000001"}, 413, "request_body_too_large"),
224        ):
225            with self.subTest(framing=framing):
226                status, headers, body = refused_request(
227                    self.app,
228                    path,
229                    {
230                        **framing,
231                        "Authorization": "Bearer " + self.app.connection["token"],
232                    },
233                    method="POST",
234                )
235                self.assertEqual(status, expected_status)
236                self.assertEqual(json.loads(body)["error"]["code"], code)
237                status, _, _ = refused_request(
238                    self.app,
239                    path,
240                    {
241                        **framing,
242                        "Origin": "https://attacker.example",
243                        "Authorization": "Bearer " + self.app.connection["token"],
244                    },
245                    method="POST",
246                )
247                self.assertEqual(status, 403)
248        self.assertFalse(self.provider.requests)
249
250    def test_authentication_precedes_body_and_authenticated_unknown_routes_keep_connection_usable(
251        self,
252    ):
253        for method, path, expected_status in (
254            ("GET", "/extensions/proxy/v1/models", 401),
255            ("GET", "/extensions/proxy/v1/missing", 401),
256            ("POST", "/extensions/proxy/v1/missing", 401),
257        ):
258            with self.subTest(method=method, path=path):
259                status, headers, _ = refused_request(
260                    self.app, path, {"Content-Length": "2"}, method=method
261                )
262                self.assertEqual(status, expected_status)
263                self.assertEqual(headers["www-authenticate"], "Bearer")
264                if expected_status == 401:
265                    self.assertEqual(
266                        delayed_request(
267                            self.app,
268                            method,
269                            path,
270                            b"{}",
271                            {"Authorization": "Bearer " + self.app.connection["token"]},
272                        ),
273                        (200 if path.endswith("/models") else 404, 200),
274                    )
275
276    @exclusive
277    def test_anonymous_proxy_requires_explicit_operator_opt_in(self):
278        self.enable()
279        with self.assertRaises(urllib.error.HTTPError) as rejected:
280            urllib.request.urlopen(self.base + "/models", timeout=10)
281        self.assertEqual(rejected.exception.code, 401)
282        with self.get_models() as response:
283            self.assertEqual(response.status, 200)
284        previous = self.app.env.get("ALBEDO_PROXY_ALLOW_ANONYMOUS")
285        try:
286            self.app.env["ALBEDO_PROXY_ALLOW_ANONYMOUS"] = "1"
287            self.app.restart()
288            self.base = self.app.base + "/extensions/proxy/v1"
289            with urllib.request.urlopen(self.base + "/models", timeout=10) as response:
290                self.assertEqual(response.status, 200)
291            for route in (
292                "/server",
293                "/extensions/work/items",
294                "/extensions/proxy/v1/missing",
295            ):
296                with (
297                    self.subTest(route=route),
298                    self.assertRaises(urllib.error.HTTPError) as rejected,
299                ):
300                    urllib.request.urlopen(self.app.base + route, timeout=10)
301                self.assertEqual(rejected.exception.code, 401)
302        finally:
303            if previous is None:
304                self.app.env.pop("ALBEDO_PROXY_ALLOW_ANONYMOUS", None)
305            else:
306                self.app.env["ALBEDO_PROXY_ALLOW_ANONYMOUS"] = previous
307            self.app.restart()
308            self.base = self.app.base + "/extensions/proxy/v1"
309
310    def test_streamed_chat_deltas_usage_and_provider_routing(self):
311        with self.post(
312            {
313                "model": self.chat,
314                "messages": OPENING,
315                "tools": TOOLS,
316                "stream": True,
317                "stream_options": {"include_usage": True},
318            }
319        ) as response:
320            self.assertTrue(
321                response.headers["content-type"].startswith("text/event-stream")
322            )
323            lines = [
324                line[6:]
325                for line in response.read().decode().splitlines()
326                if line.startswith("data: ")
327            ]
328        self.assertEqual(lines[-1], "[DONE]")
329        stream = [json.loads(line) for line in lines[:-1]]
330        deltas = [event["choices"][0]["delta"] for event in stream if event["choices"]]
331        self.assertEqual(deltas[0]["role"], "assistant")
332        self.assertEqual("".join(delta.get("content", "") for delta in deltas), "hel")
333        self.assertEqual(
334            "".join(delta.get("reasoning_content", "") for delta in deltas), "hmm"
335        )
336        [call] = [call for delta in deltas for call in delta.get("tool_calls", [])]
337        self.assertTrue(call["id"].startswith("call-1__albedo__"))
338        self.assertEqual(call["function"], CALL["function"])
339        self.assertEqual(
340            [
341                event["choices"][0]["finish_reason"]
342                for event in stream
343                if event["choices"]
344            ][-1],
345            "tool_calls",
346        )
347        self.assertEqual(stream[-1]["usage"]["prompt_tokens"], 7)
348        self.assertEqual(stream[-1]["choices"], [])
349        self.assertTrue(all(event["model"] == self.chat for event in stream))
350        path, sent = self.sent()
351        self.assertTrue(path.endswith("/v1/chat/completions"), path)
352        self.assertEqual(sent["model"], "m-chat")
353        self.assertEqual(sent["messages"][0], OPENING[0])
354
355    def test_chat_and_responses_history_projection(self):
356        for model in (f"{self.chat}/m-chat", f"{self.resp}/m-resp"):
357            with self.subTest(model=model):
358                with self.post(
359                    {"model": model, "messages": HISTORY, "tools": TOOLS}
360                ) as response:
361                    completion = json.load(response)
362                self.assertEqual(completion["object"], "chat.completion")
363                self.assertEqual(completion["choices"][0]["finish_reason"], "stop")
364                path, sent = self.sent()
365                if model.startswith(self.chat):
366                    self.assertEqual(
367                        [message["role"] for message in sent["messages"]],
368                        ["system", "user", "assistant", "tool"],
369                    )
370                    self.assertEqual(
371                        sent["messages"][2]["tool_calls"][0]["id"], "call-1"
372                    )
373                else:
374                    self.assertTrue(path.endswith("/v1/responses"), path)
375                    self.assertEqual(sent["instructions"], "be brief")
376                    self.assertEqual(
377                        [item.get("type", item.get("role")) for item in sent["input"]],
378                        ["user", "assistant", "function_call", "function_call_output"],
379                    )
380                    self.assertEqual(
381                        sent["input"][3],
382                        {
383                            "type": "function_call_output",
384                            "call_id": "call-1",
385                            "output": "a.txt",
386                        },
387                    )
388
389    def test_response_state_round_trips_only_to_originating_profile(self):
390        ask = [{"role": "user", "content": "call a tool"}]
391        with self.post(
392            {"model": self.resp, "messages": ask, "tools": TOOLS}
393        ) as response:
394            first = json.load(response)["choices"][0]["message"]
395        [call] = first["tool_calls"]
396        self.assertTrue(call["id"].startswith("call_r1__albedo__"))
397        echoed = ask + [
398            {"role": "assistant", "content": None, "tool_calls": [call]},
399            {"role": "tool", "tool_call_id": call["id"], "content": "a.txt"},
400        ]
401        with self.post({"model": self.resp, "messages": echoed, "tools": TOOLS}):
402            pass
403        sent = self.sent()[1]
404        self.assertEqual(
405            sent["input"][1],
406            {
407                "type": "reasoning",
408                "id": "rs_1",
409                "summary": [],
410                "encrypted_content": "sealed-reasoning",
411            },
412        )
413        self.assertEqual(
414            [sent["input"][index]["call_id"] for index in (2, 3)], ["call_r1"] * 2
415        )
416        with self.post({"model": self.chat, "messages": echoed, "tools": TOOLS}):
417            pass
418        path, sent = self.sent()
419        self.assertTrue(path.endswith("/v1/chat/completions"), path)
420        self.assertNotIn("sealed-reasoning", json.dumps(sent))
421        self.assertEqual(sent["messages"][1]["tool_calls"][0]["id"], "call_r1")
422        self.assertEqual(sent["messages"][2]["tool_call_id"], "call_r1")