mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## Summary CVE-2026-77775 (SSRF via `x-headroom-base-url`) is **not fully fixed on current `main`**. The advisory lists 0.36.1 as the last affected version; one route still forwards to any destination a caller names. `upstream_guard.is_safe_upstream_url` was added and wired into `/v1/messages` and the catch-all passthrough. But `select_passthrough_base_url` moved from `providers/proxy_routes.py` to `providers/proxy_targets.py`, and the guard did not follow it. Its Azure branch returns the header verbatim whenever an `api-key` header is present — **both values are caller-supplied** — and `POST /v1/alpha/search` resolves its upstream through that helper without checking the header itself. ## Verified, not inferred Against the current tree, with a listener on loopback standing in for an internal service: ``` proxy status : 200 internal service hit : 1 time(s) Authorization it received : 'Bearer SECRET-CLIENT-TOKEN' internal body relayed back : True ``` The caller's credentials are forwarded to the attacker-named host and the internal response is relayed back. After this change: `400`, zero hits, nothing relayed. A sweep of all 99 routes isolates exactly one leak on unfixed code — `POST /v1/alpha/search` with `api-key` — and zero after. ## 1. The missing enforcement **Guarded at the chokepoint, not just the route.** `select_passthrough_base_url` now validates before returning, in `proxy_targets.py` and in the parallel copy in `providers/registry.py`, so a future caller that forgets the header check cannot reopen this. `/v1/alpha/search` also rejects explicitly with 400, matching its sibling routes. ## 2. A second gap in the address policy RFC 6598 shared address space (`100.64.0.0/10`) is not `is_private`, so it passed the guard — while routing to ISP and cloud-internal infrastructure. `_is_internal_address` now also rejects anything not globally routable. Verified over a 27-vector battery — 0 bypasses, public control unaffected: | Vector | Before | After | |---|---|---| | `100.64.0.0/10` shared address space | **allowed** | blocked | | `198.18/15`, TEST-NET, `240/4` | **allowed** | blocked | | 6to4 / Teredo embedding internal IPv4 | **allowed** | blocked | | NAT64 `64:ff9b::/96` embedding loopback | **allowed** | blocked | | loopback, RFC1918, link-local, metadata, IPv4-mapped, userinfo tricks | blocked | blocked | | multicast `224.0.0.1` | blocked | blocked | | public `8.8.8.8` | allowed | allowed | The category checks are **kept alongside** `is_global` rather than replaced — `is_global` is `True` for multicast, so a replacement would have regressed. NAT64 also reports as global, so its embedded IPv4 is extracted and judged on its own. ## 3. Unauthenticated stall via the resolver `socket.getaddrinfo` takes no timeout and runs on the calling thread — the event loop. Since the hostname is caller-supplied, a deliberately slow-resolving name stalled every other in-flight request; a handful of concurrent requests made the proxy unresponsive, unauthenticated. Resolution now runs in a small dedicated pool with a budget (`HEADROOM_UPSTREAM_RESOLVE_TIMEOUT_S`, default 3s) and fails closed on overrun, which bounds every caller including the synchronous chokepoint. `is_safe_upstream_url_async` runs the lookup off the loop, and the three route handlers that validate a caller-supplied upstream now await it. Caching was deliberately avoided: a TTL cache in front of a security decision invites poisoning, and would widen the rebinding window rather than narrow it. ## Why this survived The existing tests unit-tested the guard's *logic* but never asserted it was *reached*. Added enforcement tests at the sinks plus a **sweep over the whole route table** that fails if any route forwards to a loopback address — so the next unguarded upstream resolution fails in CI rather than in a CVE. All new tests were confirmed failing against the unfixed tree and passing after. ## Known residual — deliberately not addressed **DNS rebinding.** Validation and connection resolve the host separately, so a low-TTL answer can differ between them. Closing this needs connection-time pinning in the shared `http_client` transport, which carries every request in the proxy — too broad to fold into this patch. It should not be described as fixed. ## Compatibility An endpoint that does not resolve publicly (split-horizon, on-prem) is now rejected where it previously passed unvalidated. `HEADROOM_ALLOWED_BASE_URLS` is the documented opt-in, covered by test. Three existing tests used fictional hostnames and legitimately began failing; DNS is pinned in them so they keep testing target precedence rather than depending on the missing guard. Separately: `docker-compose.yml` has already been hardened since the advisory — `HEADROOM_PROXY_TOKEN` is now mandatory and ports are loopback-only — so the "exposed by default" multiplier the advisory cites no longer applies to the shipped compose. Full suite: the 3 failures outside this area (`test_learn/test_integration`, `test_release_workflows::test_no_native_tls_in_wheel_build_tree`, and a `test_graceful_shutdown` ordering flake) reproduce on clean `main` and are unrelated. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Tejas Chopra <tejas@Tejass-MacBook-Pro.local> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
1349 lines
52 KiB
Python
1349 lines
52 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.testclient import TestClient
|
|
|
|
from headroom.providers.codex.runtime import CodexRoutingDecision
|
|
from headroom.proxy import upstream_guard
|
|
from headroom.proxy.project_context import get_current_project
|
|
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
|
|
|
|
|
|
def _app() -> Any:
|
|
return create_app(
|
|
ProxyConfig(
|
|
optimize=False,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
anthropic_api_url="https://api.anthropic.test",
|
|
openai_api_url="https://api.openai.test",
|
|
gemini_api_url="https://api.gemini.test",
|
|
cloudcode_api_url="https://cloudcode.test",
|
|
vertex_api_url="https://vertex.test",
|
|
)
|
|
)
|
|
|
|
|
|
def test_provider_passthrough_routes_forward_expected_targets(monkeypatch) -> None:
|
|
# This routing test uses reserved, intentionally unresolvable hostnames.
|
|
# Explicitly allow them so the SSRF guard can remain fail-closed on DNS errors.
|
|
monkeypatch.setenv("HEADROOM_ALLOWED_BASE_URLS", "azure.example,custom.example,opencode.ai")
|
|
calls: list[tuple[str, str, str, str]] = []
|
|
gemini_calls: list[tuple[str, str, str, str]] = []
|
|
gemini_count_calls: list[tuple[str, str, str, str]] = []
|
|
anthropic_calls: list[tuple[str, str, str, str, bool]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
calls.append((request.method, request.url.path, base_url, provider_name))
|
|
return JSONResponse(
|
|
{
|
|
"method": request.method,
|
|
"path": request.url.path,
|
|
"base_url": base_url,
|
|
"sub_path": sub_path,
|
|
"provider": provider_name,
|
|
}
|
|
)
|
|
|
|
async def fake_gemini_generate(
|
|
self,
|
|
request,
|
|
model,
|
|
upstream_base_url=None,
|
|
provider_name="gemini",
|
|
): # type: ignore[no-untyped-def]
|
|
gemini_calls.append((request.url.path, model, upstream_base_url, provider_name))
|
|
return JSONResponse(
|
|
{
|
|
"handler": "handle_gemini_generate_content",
|
|
"path": request.url.path,
|
|
"model": model,
|
|
"upstream_base_url": upstream_base_url,
|
|
"provider": provider_name,
|
|
}
|
|
)
|
|
|
|
async def fake_anthropic_messages(
|
|
self,
|
|
request,
|
|
upstream_base_url=None,
|
|
provider_name="anthropic",
|
|
model_override=None,
|
|
force_stream=False,
|
|
): # type: ignore[no-untyped-def]
|
|
anthropic_calls.append(
|
|
(request.url.path, upstream_base_url, provider_name, model_override, force_stream)
|
|
)
|
|
return JSONResponse(
|
|
{
|
|
"handler": "handle_anthropic_messages",
|
|
"path": request.url.path,
|
|
"upstream_base_url": upstream_base_url,
|
|
"provider": provider_name,
|
|
"model": model_override,
|
|
"force_stream": force_stream,
|
|
}
|
|
)
|
|
|
|
async def fake_gemini_count(
|
|
self,
|
|
request,
|
|
model,
|
|
upstream_base_url=None,
|
|
provider_name="gemini",
|
|
): # type: ignore[no-untyped-def]
|
|
gemini_count_calls.append((request.url.path, model, upstream_base_url, provider_name))
|
|
return JSONResponse(
|
|
{
|
|
"handler": "handle_gemini_count_tokens",
|
|
"path": request.url.path,
|
|
"model": model,
|
|
"upstream_base_url": upstream_base_url,
|
|
"provider": provider_name,
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
|
monkeypatch.setattr(HeadroomProxy, "handle_gemini_generate_content", fake_gemini_generate)
|
|
monkeypatch.setattr(HeadroomProxy, "handle_gemini_count_tokens", fake_gemini_count)
|
|
monkeypatch.setattr(HeadroomProxy, "handle_anthropic_messages", fake_anthropic_messages)
|
|
|
|
with TestClient(_app()) as client:
|
|
assert client.post("/v1/messages/count_tokens").json()["base_url"] == (
|
|
"https://api.anthropic.test"
|
|
)
|
|
assert client.get("/v1/models", headers={"x-goog-api-key": "test"}).json()["base_url"] == (
|
|
"https://api.openai.test"
|
|
)
|
|
assert client.get("/v1/models/demo").json()["sub_path"] == "models"
|
|
assert (
|
|
client.get(
|
|
"/azure/models",
|
|
headers={
|
|
"api-key": "azure-key",
|
|
"x-headroom-base-url": "https://azure.example/openai/",
|
|
},
|
|
).json()["base_url"]
|
|
== "https://azure.example/openai"
|
|
)
|
|
assert client.post("/v1/embeddings").json()["provider"] == "openai"
|
|
assert client.post("/v1/moderations").json()["sub_path"] == "moderations"
|
|
assert client.post("/v1/images/generations").json()["sub_path"] == "images/generations"
|
|
assert client.post("/v1/images/edits").json()["sub_path"] == "images/edits"
|
|
assert client.post("/v1/audio/transcriptions").json()["sub_path"] == "audio/transcriptions"
|
|
assert client.post("/v1/audio/speech").json()["sub_path"] == "audio/speech"
|
|
assert client.get("/v1beta/models").json()["provider"] == "gemini"
|
|
assert client.get("/v1beta/models/demo").json()["sub_path"] == "models"
|
|
assert client.post("/v1beta/models/demo:embedContent").json()["sub_path"] == "embedContent"
|
|
assert client.post(
|
|
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
|
|
).json() == {
|
|
"handler": "handle_gemini_generate_content",
|
|
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
|
|
"model": "gemini-2.0-flash",
|
|
"upstream_base_url": "https://vertex.test",
|
|
"provider": "vertex:google",
|
|
}
|
|
assert client.post(
|
|
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens"
|
|
).json() == {
|
|
"handler": "handle_gemini_count_tokens",
|
|
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens",
|
|
"model": "gemini-2.0-flash",
|
|
"upstream_base_url": "https://vertex.test",
|
|
"provider": "vertex:google",
|
|
}
|
|
assert client.post(
|
|
"/v1beta1/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict"
|
|
).json() == {
|
|
"handler": "handle_anthropic_messages",
|
|
"path": "/v1beta1/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict",
|
|
"upstream_base_url": "https://vertex.test",
|
|
"provider": "vertex:anthropic",
|
|
"model": "claude-3-5-sonnet@20240620",
|
|
"force_stream": False,
|
|
}
|
|
assert client.post(
|
|
"/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict"
|
|
).json() == {
|
|
"handler": "handle_anthropic_messages",
|
|
"path": "/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict",
|
|
"upstream_base_url": "https://vertex.test/v1",
|
|
"provider": "vertex:anthropic",
|
|
"model": "claude-3-5-sonnet@20240620",
|
|
"force_stream": False,
|
|
}
|
|
assert client.post("/anthropic/v1/messages?beta=true").json() == {
|
|
"handler": "handle_anthropic_messages",
|
|
"path": "/v1/messages",
|
|
"upstream_base_url": "https://api.anthropic.test",
|
|
"provider": "anthropic",
|
|
"model": None,
|
|
"force_stream": False,
|
|
}
|
|
non_anthropic_raw = client.post(
|
|
"/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:rawPredict"
|
|
).json()
|
|
assert non_anthropic_raw.get("handler") != "handle_anthropic_messages"
|
|
non_anthropic_stream = client.post(
|
|
"/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamRawPredict"
|
|
).json()
|
|
assert non_anthropic_stream.get("handler") != "handle_anthropic_messages"
|
|
assert client.post("/v1beta/cachedContents").json()["sub_path"] == "cachedContents"
|
|
assert client.get("/v1beta/cachedContents").json()["sub_path"] == "cachedContents"
|
|
assert client.get("/v1beta/cachedContents/cache-1").json()["sub_path"] == "cachedContents"
|
|
assert client.delete("/v1beta/cachedContents/cache-1").json()["sub_path"] == (
|
|
"cachedContents"
|
|
)
|
|
custom_passthrough = client.get(
|
|
"/unhandled/path",
|
|
headers={"x-headroom-base-url": "https://custom.example/base/"},
|
|
).json()
|
|
assert custom_passthrough["base_url"] == "https://custom.example/base"
|
|
assert custom_passthrough["sub_path"] == ""
|
|
assert custom_passthrough["provider"] == ""
|
|
|
|
opencode_zen_passthrough = client.post(
|
|
"/zen/v1/chat/completions",
|
|
headers={"x-headroom-base-url": "https://opencode.ai/"},
|
|
json={"model": "zen"},
|
|
).json()
|
|
assert opencode_zen_passthrough["base_url"] == "https://opencode.ai"
|
|
assert opencode_zen_passthrough["sub_path"] == "chat/completions"
|
|
assert opencode_zen_passthrough["provider"] == "zen"
|
|
|
|
unrelated_custom_passthrough = client.post(
|
|
"/mcp",
|
|
headers={"x-headroom-base-url": "https://opencode.ai/"},
|
|
json={},
|
|
).json()
|
|
assert unrelated_custom_passthrough["sub_path"] == ""
|
|
assert unrelated_custom_passthrough["provider"] == ""
|
|
for unrelated_path in (
|
|
"/mcp/v1/chat/completions",
|
|
"/npm/v1/chat/completions",
|
|
"/context7/v1/chat/completions",
|
|
):
|
|
unrelated_custom_passthrough = client.post(
|
|
unrelated_path,
|
|
headers={"x-headroom-base-url": "https://opencode.ai/"},
|
|
json={},
|
|
).json()
|
|
assert unrelated_custom_passthrough["sub_path"] == ""
|
|
assert unrelated_custom_passthrough["provider"] == ""
|
|
get_custom_passthrough = client.get(
|
|
"/zen/v1/chat/completions",
|
|
headers={"x-headroom-base-url": "https://opencode.ai/"},
|
|
).json()
|
|
assert get_custom_passthrough["sub_path"] == ""
|
|
assert get_custom_passthrough["provider"] == ""
|
|
other_host_custom_passthrough = client.post(
|
|
"/zen/v1/chat/completions",
|
|
headers={"x-headroom-base-url": "https://custom.example/"},
|
|
json={"model": "zen"},
|
|
).json()
|
|
assert other_host_custom_passthrough["sub_path"] == ""
|
|
assert other_host_custom_passthrough["provider"] == ""
|
|
assert client.get("/another/path", headers={"x-goog-api-key": "test"}).json()[
|
|
"base_url"
|
|
] == ("https://api.gemini.test")
|
|
|
|
# Prove Code Assist routes go to the cloudcode target and normalize paths
|
|
res1 = client.post("/v1internal:loadCodeAssist").json()
|
|
assert res1["base_url"] == "https://cloudcode.test"
|
|
assert res1["path"] == "/v1internal:loadCodeAssist"
|
|
|
|
res2 = client.post("/v1/v1internal:fetchAvailableModels").json()
|
|
assert res2["base_url"] == "https://cloudcode.test"
|
|
assert res2["path"] == "/v1internal:fetchAvailableModels"
|
|
|
|
# Prove a non-Code-Assist passthrough path containing a similar substring does not get rerouted
|
|
assert (
|
|
client.get(
|
|
"/unrelated/path/containing/v1internal:someAction",
|
|
headers={"x-goog-api-key": "test"},
|
|
).json()["base_url"]
|
|
== "https://api.gemini.test"
|
|
)
|
|
|
|
assert len(calls) >= 16
|
|
assert len(gemini_calls) >= 1
|
|
assert len(gemini_count_calls) >= 1
|
|
assert len(anthropic_calls) >= 2
|
|
|
|
|
|
def test_codex_alpha_search_route_from_headroom_issue_2525() -> None:
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str, dict[str, str], bytes]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append(
|
|
(
|
|
method,
|
|
url,
|
|
dict(kwargs.get("headers", {})),
|
|
kwargs.get("content", b""),
|
|
)
|
|
)
|
|
return httpx.Response(200, json={"ok": True})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
response = client.post(
|
|
"/v1/alpha/search?query=weather",
|
|
headers={
|
|
"Authorization": "Bearer oauth-token",
|
|
"ChatGPT-Account-ID": "acct_123",
|
|
},
|
|
json={"query": "weather"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {"ok": True}
|
|
assert len(fake_http_client.calls) == 1
|
|
method, url, headers, body = fake_http_client.calls[0]
|
|
assert method == "POST"
|
|
assert url == "https://chatgpt.com/backend-api/codex/alpha/search?query=weather"
|
|
assert headers["authorization"] == "Bearer oauth-token"
|
|
assert headers["chatgpt-account-id"] == "acct_123"
|
|
assert headers["content-length"] == "19"
|
|
assert headers["content-type"] == "application/json"
|
|
assert body == b'{"query":"weather"}'
|
|
|
|
|
|
def test_non_chatgpt_alpha_search_falls_through_to_openai_upstream(monkeypatch) -> None:
|
|
calls: list[tuple[str, str, str, str, str]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
calls.append((request.method, request.url.path, base_url, sub_path, provider_name))
|
|
return JSONResponse(
|
|
{
|
|
"base_url": base_url,
|
|
"sub_path": sub_path,
|
|
"provider": provider_name,
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
|
|
|
with TestClient(_app()) as client:
|
|
response = client.post(
|
|
"/v1/alpha/search",
|
|
headers={"Authorization": "Bearer sk-proj-openai-test"},
|
|
json={"query": "weather"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {
|
|
"base_url": "https://api.openai.test",
|
|
"sub_path": "",
|
|
"provider": "",
|
|
}
|
|
assert calls == [
|
|
(
|
|
"POST",
|
|
"/v1/alpha/search",
|
|
"https://api.openai.test",
|
|
"",
|
|
"",
|
|
)
|
|
]
|
|
|
|
|
|
def test_codex_alpha_search_route_matrix(monkeypatch) -> None:
|
|
fallback_calls: list[tuple[str, str]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
fallback_calls.append((request.url.path, base_url))
|
|
return JSONResponse({"base_url": base_url, "provider": provider_name})
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.urls: list[str] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.urls.append(url)
|
|
return httpx.Response(200, json={"ok": True})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
|
|
oauth_response = client.post(
|
|
"/v1/alpha/search",
|
|
headers={
|
|
"Authorization": "Bearer oauth-token",
|
|
"ChatGPT-Account-ID": "acct_123",
|
|
},
|
|
json={"query": "weather"},
|
|
)
|
|
api_key_response = client.post(
|
|
"/v1/alpha/search",
|
|
headers={"Authorization": "Bearer sk-proj-openai-test"},
|
|
json={"query": "weather"},
|
|
)
|
|
|
|
assert oauth_response.status_code == 200
|
|
assert api_key_response.status_code == 200
|
|
assert fake_http_client.urls == ["https://chatgpt.com/backend-api/codex/alpha/search"]
|
|
assert fallback_calls == [("/v1/alpha/search", "https://api.openai.test")]
|
|
|
|
|
|
def test_proxy_route_helpers_prefer_legacy_targets_and_gemini_passthrough() -> None:
|
|
proxy_routes = importlib.import_module("headroom.providers.proxy_routes")
|
|
proxy = type(
|
|
"Proxy",
|
|
(),
|
|
{
|
|
"ANTHROPIC_API_URL": "https://legacy.anthropic.test",
|
|
"OPENAI_API_URL": "https://legacy.openai.test",
|
|
"GEMINI_API_URL": "https://legacy.gemini.test",
|
|
"VERTEX_API_URL": "https://legacy.vertex.test",
|
|
"provider_runtime": type(
|
|
"Runtime",
|
|
(),
|
|
{
|
|
"api_target": staticmethod(lambda provider: f"https://runtime.{provider}.test"),
|
|
"model_metadata_provider": staticmethod(lambda headers: "anthropic"),
|
|
},
|
|
)(),
|
|
},
|
|
)()
|
|
|
|
assert proxy_routes._api_target(proxy, "anthropic") == "https://legacy.anthropic.test"
|
|
assert proxy_routes._api_target(proxy, "vertex") == "https://legacy.vertex.test"
|
|
assert proxy_routes._select_passthrough_base_url(proxy, {"x-goog-api-key": "test"}) == (
|
|
"https://legacy.gemini.test"
|
|
)
|
|
# The azure branch honours the override, but only after the SSRF guard
|
|
# clears the destination (CVE-2026-77775). `azure.example` does not
|
|
# resolve, and the guard fails closed on resolution failure, so pin a
|
|
# public answer to keep this assertion about target *precedence*.
|
|
with patch.object(
|
|
upstream_guard.socket,
|
|
"getaddrinfo",
|
|
return_value=[(None, None, None, None, ("20.10.10.10", 443))],
|
|
):
|
|
assert (
|
|
proxy_routes._select_passthrough_base_url(
|
|
proxy, {"api-key": "azure", "x-headroom-base-url": "https://azure.example/base/"}
|
|
)
|
|
== "https://azure.example/base"
|
|
)
|
|
assert proxy_routes._select_passthrough_base_url(proxy, {"api-key": "azure"}) == (
|
|
"https://legacy.anthropic.test"
|
|
)
|
|
assert (
|
|
proxy_routes._select_passthrough_base_url(proxy, {"chatgpt-account-id": "acct"})
|
|
== "https://chatgpt.com"
|
|
)
|
|
assert proxy_routes._select_passthrough_base_url(proxy, {}) == "https://legacy.anthropic.test"
|
|
|
|
|
|
def test_provider_specific_routes_delegate_to_expected_proxy_handlers(monkeypatch) -> None:
|
|
delegated: list[tuple[str, str, tuple[str, ...]]] = []
|
|
|
|
def install(name: str) -> None:
|
|
async def fake(self, request, *args): # type: ignore[no-untyped-def]
|
|
delegated.append((name, request.url.path, tuple(str(arg) for arg in args)))
|
|
return JSONResponse({"handler": name, "path": request.url.path, "args": list(args)})
|
|
|
|
monkeypatch.setattr(HeadroomProxy, name, fake)
|
|
|
|
for handler_name in (
|
|
"handle_anthropic_messages",
|
|
"handle_anthropic_batch_create",
|
|
"handle_anthropic_batch_passthrough",
|
|
"handle_anthropic_batch_results",
|
|
"handle_openai_chat",
|
|
"handle_openai_responses",
|
|
"handle_batch_create",
|
|
"handle_batch_list",
|
|
"handle_batch_get",
|
|
"handle_batch_cancel",
|
|
"handle_gemini_generate_content",
|
|
"handle_gemini_stream_generate_content",
|
|
"handle_gemini_count_tokens",
|
|
"handle_google_cloudcode_stream",
|
|
"handle_google_batch_create",
|
|
"handle_google_batch_results",
|
|
"handle_google_batch_passthrough",
|
|
"handle_passthrough",
|
|
):
|
|
install(handler_name)
|
|
|
|
with TestClient(_app()) as client:
|
|
assert client.post("/v1/messages").json()["handler"] == "handle_anthropic_messages"
|
|
assert client.post("/anthropic/v1/messages").json() == {
|
|
"handler": "handle_anthropic_messages",
|
|
"path": "/v1/messages",
|
|
"args": ["https://api.anthropic.test"],
|
|
}
|
|
assert (
|
|
client.post("/v1/messages/batches").json()["handler"] == "handle_anthropic_batch_create"
|
|
)
|
|
assert client.get("/v1/messages/batches").json()["handler"] == (
|
|
"handle_anthropic_batch_passthrough"
|
|
)
|
|
assert client.get("/v1/messages/batches/b1").json()["args"] == ["b1"]
|
|
assert client.get("/v1/messages/batches/b1/results").json()["handler"] == (
|
|
"handle_anthropic_batch_results"
|
|
)
|
|
assert client.post("/v1/messages/batches/b1/cancel").json()["handler"] == (
|
|
"handle_anthropic_batch_passthrough"
|
|
)
|
|
assert client.post("/v1/chat/completions").json()["handler"] == "handle_openai_chat"
|
|
assert client.post("/chat/completions").json()["handler"] == "handle_openai_chat"
|
|
assert client.post("/v1/responses").json()["handler"] == "handle_openai_responses"
|
|
assert client.post("/responses").json()["handler"] == "handle_openai_responses"
|
|
assert client.post("/v1/codex/responses").json()["handler"] == "handle_openai_responses"
|
|
assert client.post("/backend-api/responses").json()["handler"] == "handle_openai_responses"
|
|
assert client.post("/backend-api/codex/responses").json()["handler"] == (
|
|
"handle_openai_responses"
|
|
)
|
|
assert client.post("/v1/batches").json()["handler"] == "handle_batch_create"
|
|
assert client.get("/v1/batches").json()["handler"] == "handle_batch_list"
|
|
assert client.get("/v1/batches/b1").json()["handler"] == "handle_batch_get"
|
|
assert client.post("/v1/batches/b1/cancel").json()["handler"] == "handle_batch_cancel"
|
|
assert client.post("/v1beta/models/demo:generateContent").json()["handler"] == (
|
|
"handle_gemini_generate_content"
|
|
)
|
|
assert client.post("/v1beta/models/demo:streamGenerateContent").json()["handler"] == (
|
|
"handle_gemini_stream_generate_content"
|
|
)
|
|
assert client.post("/v1beta/models/demo:countTokens").json()["handler"] == (
|
|
"handle_gemini_count_tokens"
|
|
)
|
|
assert client.post(
|
|
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent"
|
|
).json() == {
|
|
"handler": "handle_gemini_generate_content",
|
|
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent",
|
|
"args": [
|
|
"gemini-2.0-flash",
|
|
"https://vertex.test",
|
|
"vertex:google",
|
|
],
|
|
}
|
|
assert client.post(
|
|
"/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens"
|
|
).json() == {
|
|
"handler": "handle_gemini_count_tokens",
|
|
"path": "/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:countTokens",
|
|
"args": [
|
|
"gemini-2.0-flash",
|
|
"https://vertex.test",
|
|
"vertex:google",
|
|
],
|
|
}
|
|
assert client.post(
|
|
"/v1beta1/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:streamRawPredict"
|
|
).json()["args"] == [
|
|
"https://vertex.test",
|
|
"vertex:anthropic",
|
|
"claude-3-5-sonnet@20240620",
|
|
True,
|
|
]
|
|
assert client.post(
|
|
"/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:rawPredict"
|
|
).json()["args"] == [
|
|
"https://vertex.test/v1",
|
|
"vertex:anthropic",
|
|
"claude-3-5-sonnet@20240620",
|
|
]
|
|
assert client.post(
|
|
"/projects/p/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet@20240620:streamRawPredict"
|
|
).json()["args"] == [
|
|
"https://vertex.test/v1",
|
|
"vertex:anthropic",
|
|
"claude-3-5-sonnet@20240620",
|
|
True,
|
|
]
|
|
assert client.post("/v1internal:streamGenerateContent").json()["handler"] == (
|
|
"handle_google_cloudcode_stream"
|
|
)
|
|
assert client.post("/v1/v1internal:streamGenerateContent").json()["handler"] == (
|
|
"handle_google_cloudcode_stream"
|
|
)
|
|
assert client.post("/v1beta/models/demo:batchGenerateContent").json()["handler"] == (
|
|
"handle_google_batch_create"
|
|
)
|
|
assert client.get("/v1beta/batches/b1").json()["handler"] == "handle_google_batch_results"
|
|
assert client.post("/v1beta/batches/b1:cancel").json()["handler"] == (
|
|
"handle_google_batch_passthrough"
|
|
)
|
|
assert client.delete("/v1beta/batches/b1").json()["handler"] == (
|
|
"handle_google_batch_passthrough"
|
|
)
|
|
|
|
assert len(delegated) >= 26
|
|
|
|
|
|
def test_openai_response_websocket_aliases_delegate_to_openai_ws_handler(monkeypatch) -> None:
|
|
seen_paths: list[str] = []
|
|
|
|
async def fake_ws(self, websocket): # type: ignore[no-untyped-def]
|
|
seen_paths.append(websocket.url.path)
|
|
await websocket.accept()
|
|
await websocket.send_json({"path": websocket.url.path})
|
|
await websocket.close()
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_openai_responses_ws", fake_ws)
|
|
|
|
with TestClient(_app()) as client:
|
|
for path in (
|
|
"/v1/responses",
|
|
"/v1/codex/responses",
|
|
"/backend-api/responses",
|
|
"/backend-api/codex/responses",
|
|
):
|
|
with client.websocket_connect(path) as websocket:
|
|
assert websocket.receive_json() == {"path": path}
|
|
|
|
assert seen_paths == [
|
|
"/v1/responses",
|
|
"/v1/codex/responses",
|
|
"/backend-api/responses",
|
|
"/backend-api/codex/responses",
|
|
]
|
|
|
|
|
|
def test_project_prefixed_openai_response_websocket_delegates_to_openai_ws_handler(
|
|
monkeypatch,
|
|
) -> None:
|
|
seen_paths: list[str] = []
|
|
seen_projects: list[str | None] = []
|
|
|
|
async def fake_ws(self, websocket): # type: ignore[no-untyped-def]
|
|
seen_paths.append(websocket.url.path)
|
|
seen_projects.append(get_current_project())
|
|
await websocket.accept()
|
|
await websocket.send_json({"path": websocket.url.path})
|
|
await websocket.close()
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_openai_responses_ws", fake_ws)
|
|
|
|
with TestClient(_app()) as client:
|
|
with client.websocket_connect("/p/test-project/v1/responses") as websocket:
|
|
assert websocket.receive_json() == {"path": "/v1/responses"}
|
|
|
|
assert seen_paths == ["/v1/responses"]
|
|
# The /p/<name> prefix is bound as the project even without a header, so a
|
|
# prefix-only Codex WS client is still attributed (not just routed).
|
|
assert seen_projects == ["test-project"]
|
|
|
|
|
|
def test_openai_response_subpath_passthrough_returns_502_on_http_failure() -> None:
|
|
class FailingAsyncClient:
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
raise RuntimeError(f"boom: {method} {url}")
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
client.app.state.proxy.http_client = FailingAsyncClient()
|
|
with patch("headroom.providers.openai_responses.logger") as logger:
|
|
response = client.post("/v1/responses/compact?trace=1", json={"model": "gpt-4o"})
|
|
|
|
assert response.status_code == 502
|
|
assert response.text == "Upstream request failed."
|
|
logger.error.assert_called_once()
|
|
assert "boom: POST https://api.openai.test/v1/responses/compact?trace=1" in str(
|
|
logger.error.call_args
|
|
)
|
|
|
|
|
|
def test_openai_response_subpath_passthrough_uses_openai_target() -> None:
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str, dict[str, str]]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append((method, url, dict(kwargs.get("headers", {}))))
|
|
return httpx.Response(200, json={"url": url})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake
|
|
response = client.delete(
|
|
"/v1/responses/items/resp_123?trace=7",
|
|
headers={"Authorization": "Bearer sk-proj-test"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert len(fake.calls) == 1
|
|
method, url, headers = fake.calls[0]
|
|
assert method == "DELETE"
|
|
assert url == "https://api.openai.test/v1/responses/items/resp_123?trace=7"
|
|
assert headers["authorization"] == "Bearer sk-proj-test"
|
|
|
|
|
|
def test_openai_response_subpath_aliases_and_chatgpt_auth_use_expected_targets(monkeypatch) -> None:
|
|
monkeypatch.setattr(
|
|
"headroom.providers.codex.responses.resolve_codex_routing",
|
|
lambda headers: CodexRoutingDecision(headers=dict(headers), is_chatgpt_auth=True),
|
|
)
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append((method, url))
|
|
return httpx.Response(200, json={"url": url})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake
|
|
assert client.get("/v1/codex/responses/items/resp_1").status_code == 200
|
|
assert client.post("/backend-api/responses/items/resp_2").status_code == 200
|
|
assert client.delete("/backend-api/codex/responses/items/resp_3").status_code == 200
|
|
|
|
assert fake.calls == [
|
|
("GET", "https://chatgpt.com/backend-api/codex/responses/items/resp_1"),
|
|
("POST", "https://chatgpt.com/backend-api/codex/responses/items/resp_2"),
|
|
("DELETE", "https://chatgpt.com/backend-api/codex/responses/items/resp_3"),
|
|
]
|
|
|
|
|
|
def test_openai_image_routes_use_codex_backend_under_chatgpt_auth(monkeypatch) -> None:
|
|
monkeypatch.setattr(
|
|
"headroom.providers.codex.images.resolve_codex_routing",
|
|
lambda headers: CodexRoutingDecision(
|
|
headers={**headers, "ChatGPT-Account-ID": "acct_123"},
|
|
is_chatgpt_auth=True,
|
|
),
|
|
)
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str, dict[str, str], bytes]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append(
|
|
(
|
|
method,
|
|
url,
|
|
dict(kwargs.get("headers", {})),
|
|
kwargs.get("content", b""),
|
|
)
|
|
)
|
|
return httpx.Response(200, json={"url": url})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake
|
|
client.app.state.proxy.http_client_h1 = fake
|
|
|
|
generate_response = client.post(
|
|
"/v1/images/generations?client_version=0.142.0",
|
|
headers={
|
|
"Authorization": "Bearer oauth-token",
|
|
"Accept-Encoding": "gzip",
|
|
"X-Headroom-Bypass": "1",
|
|
},
|
|
json={"model": "gpt-image-2", "prompt": "a route probe"},
|
|
)
|
|
edit_response = client.post(
|
|
"/v1/images/edits",
|
|
headers={"Authorization": "Bearer oauth-token"},
|
|
json={"model": "gpt-image-2", "prompt": "edit route probe", "images": []},
|
|
)
|
|
|
|
assert generate_response.status_code == 200
|
|
assert edit_response.status_code == 200
|
|
assert len(fake.calls) == 2
|
|
|
|
generate_method, generate_url, generate_headers, generate_body = fake.calls[0]
|
|
assert generate_method == "POST"
|
|
assert (
|
|
generate_url
|
|
== "https://chatgpt.com/backend-api/codex/images/generations?client_version=0.142.0"
|
|
)
|
|
assert generate_headers["authorization"] == "Bearer oauth-token"
|
|
assert generate_headers["ChatGPT-Account-ID"] == "acct_123"
|
|
assert "host" not in generate_headers
|
|
assert "accept-encoding" not in generate_headers
|
|
assert "x-headroom-bypass" not in generate_headers
|
|
assert generate_body == b'{"model":"gpt-image-2","prompt":"a route probe"}'
|
|
|
|
edit_method, edit_url, edit_headers, edit_body = fake.calls[1]
|
|
assert edit_method == "POST"
|
|
assert edit_url == "https://chatgpt.com/backend-api/codex/images/edits"
|
|
assert edit_headers["authorization"] == "Bearer oauth-token"
|
|
assert edit_headers["ChatGPT-Account-ID"] == "acct_123"
|
|
assert "host" not in edit_headers
|
|
assert edit_body == b'{"model":"gpt-image-2","prompt":"edit route probe","images":[]}'
|
|
|
|
|
|
def test_openai_image_codex_response_strips_stale_compression_headers(monkeypatch) -> None:
|
|
upstream_body = b'{"ok":true}'
|
|
stale_content_length = "9999"
|
|
monkeypatch.setattr(
|
|
"headroom.providers.codex.images.resolve_codex_routing",
|
|
lambda headers: CodexRoutingDecision(
|
|
headers={**headers, "ChatGPT-Account-ID": "acct_123"},
|
|
is_chatgpt_auth=True,
|
|
),
|
|
)
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str, bytes]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append((method, url, kwargs.get("content", b"")))
|
|
return FakeUpstreamResponse(
|
|
content=upstream_body,
|
|
status_code=200,
|
|
headers={
|
|
"content-encoding": "gzip",
|
|
"content-length": stale_content_length,
|
|
"content-type": "application/json",
|
|
"server": "upstream-edge",
|
|
"x-upstream": "kept",
|
|
},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
class FakeUpstreamResponse:
|
|
def __init__(self, content: bytes, status_code: int, headers: dict[str, str]) -> None:
|
|
self.content = content
|
|
self.status_code = status_code
|
|
self.headers = headers
|
|
|
|
with TestClient(_app()) as client:
|
|
fake = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake
|
|
client.app.state.proxy.http_client_h1 = fake
|
|
|
|
response = client.post(
|
|
"/v1/images/generations",
|
|
headers={"Authorization": "Bearer oauth-token"},
|
|
json={"model": "gpt-image-2", "prompt": "compressed response"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.content == upstream_body
|
|
assert response.headers["x-upstream"] == "kept"
|
|
assert response.headers.get("server") is None
|
|
assert response.headers.get("content-encoding") is None
|
|
assert response.headers.get("content-length") == str(len(upstream_body))
|
|
|
|
assert fake.calls == [
|
|
(
|
|
"POST",
|
|
"https://chatgpt.com/backend-api/codex/images/generations",
|
|
b'{"model":"gpt-image-2","prompt":"compressed response"}',
|
|
)
|
|
]
|
|
|
|
|
|
def test_openai_image_edits_api_key_auth_falls_through_to_openai_passthrough(
|
|
monkeypatch,
|
|
) -> None:
|
|
calls: list[tuple[str, str, str, str, str]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
calls.append((request.method, request.url.path, base_url, sub_path, provider_name))
|
|
return JSONResponse(
|
|
{
|
|
"base_url": base_url,
|
|
"sub_path": sub_path,
|
|
"provider": provider_name,
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
|
|
|
with TestClient(_app()) as client:
|
|
response = client.post(
|
|
"/v1/images/edits",
|
|
headers={"Authorization": "Bearer sk-proj-openai-test"},
|
|
json={"model": "gpt-image-1", "prompt": "fall through", "image": "file-1"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {
|
|
"base_url": "https://api.openai.test",
|
|
"sub_path": "images/edits",
|
|
"provider": "openai",
|
|
}
|
|
assert calls == [
|
|
(
|
|
"POST",
|
|
"/v1/images/edits",
|
|
"https://api.openai.test",
|
|
"images/edits",
|
|
"openai",
|
|
)
|
|
]
|
|
|
|
|
|
def test_openai_image_edits_preserves_multipart_body_under_chatgpt_auth(monkeypatch) -> None:
|
|
monkeypatch.setattr(
|
|
"headroom.providers.codex.images.resolve_codex_routing",
|
|
lambda headers: CodexRoutingDecision(
|
|
headers={**headers, "ChatGPT-Account-ID": "acct_123"},
|
|
is_chatgpt_auth=True,
|
|
),
|
|
)
|
|
boundary = "----headroom-boundary"
|
|
body = (
|
|
(
|
|
f"--{boundary}\r\n"
|
|
'Content-Disposition: form-data; name="model"\r\n\r\n'
|
|
"gpt-image-2\r\n"
|
|
f"--{boundary}\r\n"
|
|
'Content-Disposition: form-data; name="prompt"\r\n\r\n'
|
|
"preserve these bytes\r\n"
|
|
f"--{boundary}\r\n"
|
|
'Content-Disposition: form-data; name="image"; filename="input.png"\r\n'
|
|
"Content-Type: image/png\r\n\r\n"
|
|
).encode()
|
|
+ b"\x89PNG\r\n\x1a\nraw-bytes\r\n"
|
|
+ f"--{boundary}--\r\n".encode()
|
|
)
|
|
content_type = f"multipart/form-data; boundary={boundary}"
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str, dict[str, str], bytes]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append(
|
|
(
|
|
method,
|
|
url,
|
|
dict(kwargs.get("headers", {})),
|
|
kwargs.get("content", b""),
|
|
)
|
|
)
|
|
return httpx.Response(200, json={"ok": True})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake
|
|
client.app.state.proxy.http_client_h1 = fake
|
|
|
|
response = client.post(
|
|
"/v1/images/edits",
|
|
headers={
|
|
"Authorization": "Bearer oauth-token",
|
|
"Content-Type": content_type,
|
|
},
|
|
content=body,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert len(fake.calls) == 1
|
|
method, url, headers, forwarded_body = fake.calls[0]
|
|
assert method == "POST"
|
|
assert url == "https://chatgpt.com/backend-api/codex/images/edits"
|
|
assert headers["authorization"] == "Bearer oauth-token"
|
|
assert headers["ChatGPT-Account-ID"] == "acct_123"
|
|
assert headers["content-type"] == content_type
|
|
assert "host" not in headers
|
|
assert forwarded_body == body
|
|
|
|
|
|
def test_gemini_batch_embed_contents_passthrough_uses_gemini_target(monkeypatch) -> None:
|
|
calls: list[tuple[str, str, str]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
calls.append((request.url.path, base_url, sub_path))
|
|
return JSONResponse({"base_url": base_url, "sub_path": sub_path, "provider": provider_name})
|
|
|
|
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
|
|
|
|
with TestClient(_app()) as client:
|
|
response = client.post("/v1beta/models/demo:batchEmbedContents")
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {
|
|
"base_url": "https://api.gemini.test",
|
|
"sub_path": "batchEmbedContents",
|
|
"provider": "gemini",
|
|
}
|
|
assert calls == [
|
|
("/v1beta/models/demo:batchEmbedContents", "https://api.gemini.test", "batchEmbedContents")
|
|
]
|
|
|
|
|
|
def test_v1_models_fetches_codex_registry_under_chatgpt_auth(monkeypatch) -> None:
|
|
model_metadata = importlib.import_module("headroom.providers.codex.model_metadata")
|
|
debug_messages: list[tuple[str, tuple[object, ...]]] = []
|
|
monkeypatch.setattr(
|
|
model_metadata.logger,
|
|
"debug",
|
|
lambda message, *args: debug_messages.append((message, args)),
|
|
)
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str, dict[str, str]]] = []
|
|
|
|
async def get(self, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append(("GET", url, dict(kwargs.get("headers", {}))))
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"models": [
|
|
{"slug": "gpt-5.5"},
|
|
{"slug": "gpt-5.3-codex-spark"},
|
|
]
|
|
},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
response = client.get(
|
|
"/v1/models?client_version=0.135.0",
|
|
headers={
|
|
"authorization": "Bearer eyJ-chatgpt-oauth-token",
|
|
"chatgpt-account-id": "test-account",
|
|
"originator": "Codex Desktop",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
assert payload["object"] == "list"
|
|
assert payload["data"] == [
|
|
{
|
|
"id": "gpt-5.5",
|
|
"object": "model",
|
|
"created": 0,
|
|
"owned_by": "openai",
|
|
},
|
|
{
|
|
"id": "gpt-5.3-codex-spark",
|
|
"object": "model",
|
|
"created": 0,
|
|
"owned_by": "openai",
|
|
},
|
|
]
|
|
assert [entry["slug"] for entry in payload["models"]] == [
|
|
"gpt-5.5",
|
|
"gpt-5.3-codex-spark",
|
|
]
|
|
assert [entry["display_name"] for entry in payload["models"]] == [
|
|
"GPT-5.5",
|
|
"GPT-5.3-Codex-Spark",
|
|
]
|
|
for entry in payload["models"]:
|
|
assert entry["default_reasoning_level"] == "medium"
|
|
assert entry["context_window"] == 272000
|
|
assert entry["supports_parallel_tool_calls"] is True
|
|
assert len(fake_http_client.calls) == 1
|
|
method, url, headers = fake_http_client.calls[0]
|
|
assert method == "GET"
|
|
assert url == "https://chatgpt.com/backend-api/codex/models?client_version=0.135.0"
|
|
assert headers["authorization"] == "Bearer eyJ-chatgpt-oauth-token"
|
|
assert headers["chatgpt-account-id"] == "test-account"
|
|
assert headers["originator"] == "Codex Desktop"
|
|
assert headers["accept"] == "application/json"
|
|
assert "Accept" not in headers
|
|
assert debug_messages == [
|
|
(
|
|
"Fetched Codex model IDs from upstream model registry: %s",
|
|
(["gpt-5.5", "gpt-5.3-codex-spark"],),
|
|
),
|
|
]
|
|
|
|
|
|
def test_v1_models_falls_back_to_synthetic_list_under_chatgpt_auth(monkeypatch) -> None:
|
|
"""Issue #478: under Codex ChatGPT-subscription OAuth, the proxy
|
|
must NOT forward `/v1/models` to chatgpt.com/backend-api/models
|
|
(which returns 403). If the Codex-specific registry also fails,
|
|
synthesize an OpenAI-compatible response with the known-supported
|
|
Codex/ChatGPT model set instead, so Codex's model-picker refresh succeeds.
|
|
"""
|
|
|
|
class FakeAsyncClient:
|
|
async def get(self, url, **kwargs): # type: ignore[no-untyped-def]
|
|
return httpx.Response(403, json={"error": "forbidden"})
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
client.app.state.proxy.http_client = FakeAsyncClient()
|
|
# ChatGPT auth detected via Bearer + ChatGPT account header
|
|
# (mirrors what Codex Desktop sends).
|
|
response = client.get(
|
|
"/v1/models",
|
|
headers={
|
|
"authorization": "Bearer eyJ-chatgpt-oauth-token",
|
|
"chatgpt-account-id": "test-account",
|
|
"originator": "Codex Desktop",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
assert payload["object"] == "list"
|
|
assert isinstance(payload["data"], list)
|
|
assert len(payload["data"]) > 0
|
|
model_ids = {entry["id"] for entry in payload["data"]}
|
|
model_slugs = {entry["slug"] for entry in payload["models"]}
|
|
# Spot-check: the model from issue #478's repro log must be present.
|
|
assert "gpt-5.5" in model_ids
|
|
assert "gpt-5.5" in model_slugs
|
|
gpt_55 = next(entry for entry in payload["models"] if entry["slug"] == "gpt-5.5")
|
|
assert gpt_55["display_name"] == "GPT-5.5"
|
|
assert gpt_55["supported_in_api"] is True
|
|
assert gpt_55["default_reasoning_level"] == "medium"
|
|
for entry in payload["data"]:
|
|
assert entry["object"] == "model"
|
|
assert entry["owned_by"] == "openai"
|
|
|
|
|
|
def test_v1_models_get_single_dynamic_under_chatgpt_auth() -> None:
|
|
"""The single-model variant (`/v1/models/{id}`) is also called by
|
|
Codex for some flows. It should use the Codex registry first so
|
|
dynamically exposed model slugs validate consistently."""
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
async def get(self, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls += 1
|
|
return httpx.Response(
|
|
200,
|
|
json={"models": [{"slug": "gpt-5.5"}, {"slug": "gpt-5.3-codex-spark"}]},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
ok = client.get(
|
|
"/v1/models/gpt-5.3-codex-spark",
|
|
headers={
|
|
"authorization": "Bearer eyJ-chatgpt-oauth-token",
|
|
"chatgpt-account-id": "test-account",
|
|
},
|
|
)
|
|
unknown = client.get(
|
|
"/v1/models/gpt-99-future",
|
|
headers={
|
|
"authorization": "Bearer eyJ-chatgpt-oauth-token",
|
|
"chatgpt-account-id": "test-account",
|
|
},
|
|
)
|
|
assert ok.status_code == 200
|
|
assert ok.json() == {
|
|
"id": "gpt-5.3-codex-spark",
|
|
"object": "model",
|
|
"created": 0,
|
|
"owned_by": "openai",
|
|
}
|
|
assert unknown.status_code == 404
|
|
assert fake_http_client.calls == 2
|
|
|
|
|
|
def test_v1_models_still_forwards_under_non_chatgpt_auth() -> None:
|
|
"""Non-ChatGPT auth (regular API key, Gemini, etc.) must still
|
|
forward to the upstream provider — only the ChatGPT-OAuth path
|
|
short-circuits to the synthetic response."""
|
|
calls: list[tuple[str, str, str]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
calls.append((request.url.path, base_url, provider_name))
|
|
return JSONResponse({"base_url": base_url, "provider": provider_name})
|
|
|
|
with patch.object(HeadroomProxy, "handle_passthrough", fake_passthrough):
|
|
with TestClient(_app()) as client:
|
|
response = client.get(
|
|
"/v1/models",
|
|
headers={"authorization": "Bearer sk-real-api-key"},
|
|
)
|
|
assert response.status_code == 200
|
|
# Forwarded — not synthesized — because no chatgpt-account-id header.
|
|
assert calls, "Non-ChatGPT-auth /v1/models must forward, not synthesize"
|
|
|
|
|
|
def test_v1_models_routes_claude_code_gateway_discovery_to_anthropic() -> None:
|
|
"""Claude Code gateway/OAuth model discovery can use a Bearer token that
|
|
does not look like an Anthropic API key. Route those `/v1/models` requests
|
|
to Anthropic so Claude's gateway model cache is not populated from OpenAI.
|
|
"""
|
|
calls: list[tuple[str, str, str]] = []
|
|
|
|
async def fake_passthrough(self, request, base_url, sub_path="", provider_name=""): # type: ignore[no-untyped-def]
|
|
calls.append((request.url.path, base_url, provider_name))
|
|
return JSONResponse({"base_url": base_url, "provider": provider_name})
|
|
|
|
with patch.object(HeadroomProxy, "handle_passthrough", fake_passthrough):
|
|
with TestClient(_app()) as client:
|
|
list_response = client.get(
|
|
"/v1/models",
|
|
headers={
|
|
"authorization": "Bearer claude-gateway-oauth-token",
|
|
"user-agent": "claude-code/1.5.0 (darwin; arm64)",
|
|
},
|
|
)
|
|
get_response = client.get(
|
|
"/v1/models/claude-opus-4-8",
|
|
headers={
|
|
"authorization": "Bearer claude-gateway-oauth-token",
|
|
"user-agent": "claude-code/1.5.0 (darwin; arm64)",
|
|
},
|
|
)
|
|
|
|
assert list_response.status_code == 200
|
|
assert get_response.status_code == 200
|
|
assert list_response.json() == {
|
|
"base_url": "https://api.anthropic.test",
|
|
"provider": "anthropic",
|
|
}
|
|
assert get_response.json() == {
|
|
"base_url": "https://api.anthropic.test",
|
|
"provider": "anthropic",
|
|
}
|
|
assert calls == [
|
|
("/v1/models", "https://api.anthropic.test", "anthropic"),
|
|
("/v1/models/claude-opus-4-8", "https://api.anthropic.test", "anthropic"),
|
|
]
|
|
|
|
|
|
def test_anthropic_model_metadata_strips_ansi_model_ids() -> None:
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append((method, url))
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"object": "list",
|
|
"data": [
|
|
{"id": "claude-opus-4-8\x1b[1m", "object": "model"},
|
|
{"id": "claude-sonnet-4-5[1m]", "object": "model"},
|
|
],
|
|
},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
response = client.get("/v1/models", headers={"x-api-key": "sk-ant-test"})
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["data"] == [
|
|
{"id": "claude-opus-4-8", "object": "model"},
|
|
{"id": "claude-sonnet-4-5", "object": "model"},
|
|
]
|
|
assert fake_http_client.calls == [("GET", "https://api.anthropic.test/v1/models")]
|
|
|
|
|
|
def test_anthropic_model_detail_path_strips_ansi_model_id() -> None:
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
async def request(self, method, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.calls.append((method, url))
|
|
return httpx.Response(
|
|
200,
|
|
json={"id": "claude-opus-4-8\x1b[1m", "object": "model"},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
response = client.get(
|
|
"/v1/models/claude-opus-4-8%1B%5B1m",
|
|
headers={"x-api-key": "sk-ant-test"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["id"] == "claude-opus-4-8"
|
|
assert fake_http_client.calls == [
|
|
("GET", "https://api.anthropic.test/v1/models/claude-opus-4-8")
|
|
]
|
|
|
|
|
|
def test_anthropic_messages_strips_ansi_model_id_before_upstream() -> None:
|
|
class FakeAsyncClient:
|
|
def __init__(self) -> None:
|
|
self.bodies: list[dict[str, Any]] = []
|
|
|
|
async def post(self, url, **kwargs): # type: ignore[no-untyped-def]
|
|
self.bodies.append(json.loads(kwargs["content"]))
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"id": "msg_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": "claude-opus-4-8",
|
|
"content": [],
|
|
"stop_reason": "end_turn",
|
|
"usage": {"input_tokens": 1, "output_tokens": 1},
|
|
},
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
with TestClient(_app()) as client:
|
|
fake_http_client = FakeAsyncClient()
|
|
client.app.state.proxy.http_client = fake_http_client
|
|
response = client.post(
|
|
"/v1/messages",
|
|
headers={"x-api-key": "sk-ant-test"},
|
|
json={
|
|
"model": "claude-opus-4-8\x1b[1m",
|
|
"max_tokens": 16,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert fake_http_client.bodies[0]["model"] == "claude-opus-4-8"
|