"""OpenAI-compatible proxy enablement, model routing, streaming, and state projection."""
import json
import unittest
import urllib.error
import urllib.request
from daemon_test import delayed_request, refused_request
from harness import Albedo, Provider, exclusive, Reply
TOOLS = [
{"type": "function", "function": {"name": "bash", "parameters": {"type": "object"}}}
]
OPENING = [
{"role": "system", "content": "be brief"},
{"role": "user", "content": "list files"},
]
CALL = {
"id": "call-1",
"type": "function",
"function": {"name": "bash", "arguments": '{"command":"ls"}'},
}
HISTORY = OPENING + [
{"role": "assistant", "content": "hel", "tool_calls": [CALL]},
{"role": "tool", "tool_call_id": "call-1", "content": "a.txt"},
]
def fixture(request):
if "messages" in request:
if request["messages"][-1]["role"] == "tool":
deltas = [({"role": "assistant", "content": "ran it"}, None), ({}, "stop")]
else:
deltas = [
({"role": "assistant", "reasoning_content": "hmm"}, None),
({"content": "hel"}, None),
(
{
"tool_calls": [
{
"index": 0,
**CALL,
"function": {
"name": "bash",
"arguments": '{"command":',
},
}
]
},
None,
),
(
{"tool_calls": [{"index": 0, "function": {"arguments": '"ls"}'}}]},
None,
),
({}, "tool_calls"),
]
events = [
{
"id": "c",
"choices": [{"index": 0, "delta": delta, "finish_reason": reason}],
**(
{"usage": {"prompt_tokens": 7, "completion_tokens": 3}}
if reason == "tool_calls"
else {}
),
}
for delta, reason in deltas
]
else:
last = request["input"][-1]
if last.get("role") == "user" and last.get("content") == "call a tool":
output = [
{
"type": "reasoning",
"id": "rs_1",
"summary": [],
"encrypted_content": "sealed-reasoning",
},
{
"type": "function_call",
"id": "fc_1",
"call_id": "call_r1",
"name": "bash",
"arguments": '{"command":"ls"}',
"status": "completed",
},
]
else:
output = [
{
"type": "message",
"role": "assistant",
"status": "completed",
"content": [
{"type": "output_text", "text": "done", "annotations": []}
],
}
]
events = [
{
"type": "response.completed",
"response": {
"id": "r",
"status": "completed",
"output": output,
"usage": {"input_tokens": 5, "output_tokens": 1},
},
}
]
return Reply("raw", events=events)
class ProxyTests(unittest.TestCase):
def setUp(self):
self.provider = Provider(fixture)
self.addCleanup(self.provider.close)
endpoint = self.provider.url + "/v1"
self.chat = f"chat-{self.provider.route}"
self.resp = f"resp-{self.provider.route}"
self.cx = f"cx-{self.provider.route}"
self.broken = f"broken-{self.provider.route}"
profiles = {
self.chat: {
"baseUrl": endpoint,
"apiKey": "k",
"model": "m-chat",
"protocol": "chat_completions",
},
self.resp: {
"baseUrl": endpoint,
"apiKey": "k",
"model": "m-resp",
"protocol": "responses",
},
self.cx: {
"extension": "codex",
"model": "gpt-saved",
"protocol": "responses",
},
self.broken: {"model": 3},
}
self.app = Albedo(self.provider, providers=profiles)
self.app.__enter__()
self.addCleanup(self.app.__exit__, None, None, None)
self.base = self.app.base + "/proxy/v1"
def get_models(self):
return urllib.request.urlopen(self.base + "/models", timeout=10)
def enable(self):
self.app.write_extensions({"enabled": {"proxy": True}})
def post(self, body, *, headers=None):
request = urllib.request.Request(
self.base + "/chat/completions",
json.dumps(body).encode(),
{"Content-Type": "application/json", **(headers or {})},
)
return urllib.request.urlopen(request, timeout=30)
def sent(self):
record = self.provider.requests[-1]
return record["path"], record["request"]
@exclusive
def test_enablement_listing_and_profile_errors(self):
self.app.write_extensions({"enabled": {"proxy": False}})
with self.assertRaises(urllib.error.HTTPError) as rejected:
self.get_models()
self.assertEqual(rejected.exception.code, 403)
self.enable()
with self.get_models() as response:
listing = json.load(response)
self.assertTrue(
{f"{self.chat}/m-chat", f"{self.resp}/m-resp", f"{self.cx}/gpt-saved"}
<= {model["id"] for model in listing["data"]}
)
self.assertIn(self.broken, [error["profile"] for error in listing["errors"]])
with self.post(
{
"model": f"{self.chat}/unlisted",
"messages": [{"role": "user", "content": "hi"}],
}
) as response:
self.assertEqual(json.load(response)["model"], f"{self.chat}/unlisted")
self.assertEqual(self.sent()[1]["model"], "unlisted")
self.assertEqual(self.provider.requests[-1]["authorization"], "Bearer k")
for model in (self.broken, f"{self.cx}/gpt-saved"):
with (
self.subTest(model=model),
self.assertRaises(urllib.error.HTTPError) as rejected,
):
self.post(
{"model": model, "messages": [{"role": "user", "content": "hi"}]}
)
self.assertEqual(rejected.exception.code, 400)
self.assertIn("/login", json.load(rejected.exception)["error"]["message"])
with self.assertRaises(urllib.error.HTTPError) as rejected:
self.post({}, headers={"Origin": "https://example.com"})
self.assertEqual(rejected.exception.code, 403)
with self.assertRaises(urllib.error.HTTPError) as rejected:
self.post({"model": "missing/x", "messages": OPENING})
self.assertEqual(rejected.exception.code, 400)
def test_ingress_limits_and_origin_refuse_before_provider_work(self):
path = "/proxy/v1/chat/completions"
for framing, expected_status, code in (
({"Transfer-Encoding": "chunked"}, 400, "unsupported_transfer_encoding"),
({"Content-Length": "invalid"}, 400, "invalid_content_length"),
({"Content-Length": "32000001"}, 413, "request_body_too_large"),
):
with self.subTest(framing=framing):
status, headers, body = refused_request(self.app, path, framing)
self.assertEqual(status, expected_status)
self.assertEqual(json.loads(body)["code"], code)
self.assertEqual(headers["albedo-error-code"], code)
status, _, _ = refused_request(
self.app, path, {**framing, "Origin": "https://attacker.example"}
)
self.assertEqual(status, 403)
self.assertFalse(self.provider.requests)
def test_get_and_unknown_route_bodies_keep_anonymous_connection_usable(self):
for method, path, expected_status in (
("GET", "/proxy/v1/models", 200),
("GET", "/proxy/v1/missing", 404),
("POST", "/proxy/v1/missing", 404),
):
with self.subTest(method=method, path=path):
self.assertEqual(
delayed_request(self.app, method, path, b"{}", {}),
(expected_status, 200),
)
def test_streamed_chat_deltas_usage_and_provider_routing(self):
with self.post(
{
"model": self.chat,
"messages": OPENING,
"tools": TOOLS,
"stream": True,
"stream_options": {"include_usage": True},
}
) as response:
self.assertTrue(
response.headers["content-type"].startswith("text/event-stream")
)
lines = [
line[6:]
for line in response.read().decode().splitlines()
if line.startswith("data: ")
]
self.assertEqual(lines[-1], "[DONE]")
stream = [json.loads(line) for line in lines[:-1]]
deltas = [event["choices"][0]["delta"] for event in stream if event["choices"]]
self.assertEqual(deltas[0]["role"], "assistant")
self.assertEqual("".join(delta.get("content", "") for delta in deltas), "hel")
self.assertEqual(
"".join(delta.get("reasoning_content", "") for delta in deltas), "hmm"
)
[call] = [call for delta in deltas for call in delta.get("tool_calls", [])]
self.assertTrue(call["id"].startswith("call-1__albedo__"))
self.assertEqual(call["function"], CALL["function"])
self.assertEqual(
[
event["choices"][0]["finish_reason"]
for event in stream
if event["choices"]
][-1],
"tool_calls",
)
self.assertEqual(stream[-1]["usage"]["prompt_tokens"], 7)
self.assertEqual(stream[-1]["choices"], [])
self.assertTrue(all(event["model"] == self.chat for event in stream))
path, sent = self.sent()
self.assertTrue(path.endswith("/v1/chat/completions"), path)
self.assertEqual(sent["model"], "m-chat")
self.assertEqual(sent["messages"][0], OPENING[0])
def test_chat_and_responses_history_projection(self):
for model in (f"{self.chat}/m-chat", f"{self.resp}/m-resp"):
with self.subTest(model=model):
with self.post(
{"model": model, "messages": HISTORY, "tools": TOOLS}
) as response:
completion = json.load(response)
self.assertEqual(completion["object"], "chat.completion")
self.assertEqual(completion["choices"][0]["finish_reason"], "stop")
path, sent = self.sent()
if model.startswith(self.chat):
self.assertEqual(
[message["role"] for message in sent["messages"]],
["system", "user", "assistant", "tool"],
)
self.assertEqual(
sent["messages"][2]["tool_calls"][0]["id"], "call-1"
)
else:
self.assertTrue(path.endswith("/v1/responses"), path)
self.assertEqual(sent["instructions"], "be brief")
self.assertEqual(
[item.get("type", item.get("role")) for item in sent["input"]],
["user", "assistant", "function_call", "function_call_output"],
)
self.assertEqual(
sent["input"][3],
{
"type": "function_call_output",
"call_id": "call-1",
"output": "a.txt",
},
)
def test_response_state_round_trips_only_to_originating_profile(self):
ask = [{"role": "user", "content": "call a tool"}]
with self.post(
{"model": self.resp, "messages": ask, "tools": TOOLS}
) as response:
first = json.load(response)["choices"][0]["message"]
[call] = first["tool_calls"]
self.assertTrue(call["id"].startswith("call_r1__albedo__"))
echoed = ask + [
{"role": "assistant", "content": None, "tool_calls": [call]},
{"role": "tool", "tool_call_id": call["id"], "content": "a.txt"},
]
with self.post({"model": self.resp, "messages": echoed, "tools": TOOLS}):
pass
sent = self.sent()[1]
self.assertEqual(
sent["input"][1],
{
"type": "reasoning",
"id": "rs_1",
"summary": [],
"encrypted_content": "sealed-reasoning",
},
)
self.assertEqual(
[sent["input"][index]["call_id"] for index in (2, 3)], ["call_r1"] * 2
)
with self.post({"model": self.chat, "messages": echoed, "tools": TOOLS}):
pass
path, sent = self.sent()
self.assertTrue(path.endswith("/v1/chat/completions"), path)
self.assertNotIn("sealed-reasoning", json.dumps(sent))
self.assertEqual(sent["messages"][1]["tool_calls"][0]["id"], "call_r1")
self.assertEqual(sent["messages"][2]["tool_call_id"], "call_r1")