Skip to content

Forward client bearer keys upstream #46

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vfeitoza wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Forward client bearer keys upstream #46

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 102 additions & 0 deletions tests/test_proxy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1668,6 +1668,108 @@ def test_virtual_model_routes(self, client: TestClient) -> None:
assert data["model"] == "uncommon-route/debug"
assert "UncommonRoute Debug" in data["choices"][0]["message"]["content"]

def test_request_bearer_key_overrides_primary_connection(self, monkeypatch: pytest.MonkeyPatch) -> None:
captured: dict[str, object] = {}

def handler(request: httpx.Request) -> httpx.Response:
captured["headers"] = dict(request.headers)
return httpx.Response(
200,
json={
"id": "chatcmpl-request-key",
"object": "chat.completion",
"created": 1,
"model": "openai/gpt-4o-mini",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
headers={"content-type": "application/json"},
)

store = ConnectionsStore(storage=InMemoryConnectionsStorage())
store.set_primary(
upstream="https://api.example.test/v1",
api_key="sk-connection-key",
)
async_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
monkeypatch.setattr("uncommon_route.proxy._get_client", lambda: async_client)

try:
app = create_app(
connections_store=store,
model_mapper=_build_test_mapper("openai/gpt-4o-mini"),
spend_control=SpendControl(storage=InMemorySpendControlStorage()),
)
client = TestClient(app, raise_server_exceptions=False)
response = client.post(
"/v1/messages",
headers={"authorization": "Bearer sk-claude-cli-key"},
json={
"model": "openai/gpt-4o-mini",
"max_tokens": 64,
"messages": [{"role": "user", "content": "hello"}],
},
)

assert response.status_code == 200
headers = captured["headers"]
assert isinstance(headers, dict)
assert headers["authorization"] == "Bearer sk-claude-cli-key"
assert store.primary().api_key == "sk-connection-key"
finally:
asyncio.run(async_client.aclose())

def test_connection_key_is_used_when_request_has_no_bearer(self, monkeypatch: pytest.MonkeyPatch) -> None:
captured: dict[str, object] = {}

def handler(request: httpx.Request) -> httpx.Response:
captured["headers"] = dict(request.headers)
return httpx.Response(
200,
json={
"id": "chatcmpl-connection-key",
"object": "chat.completion",
"created": 1,
"model": "openai/gpt-4o-mini",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
headers={"content-type": "application/json"},
)

store = ConnectionsStore(storage=InMemoryConnectionsStorage())
store.set_primary(
upstream="https://api.example.test/v1",
api_key="sk-connection-key",
)
async_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
monkeypatch.setattr("uncommon_route.proxy._get_client", lambda: async_client)

try:
app = create_app(
connections_store=store,
model_mapper=_build_test_mapper("openai/gpt-4o-mini"),
spend_control=SpendControl(storage=InMemorySpendControlStorage()),
)
client = TestClient(app, raise_server_exceptions=False)
response = client.post("/v1/chat/completions", json={
"model": "openai/gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
})

assert response.status_code == 200
assert captured["headers"]["authorization"] == "Bearer sk-connection-key"
finally:
asyncio.run(async_client.aclose())

def test_routing_headers_present(self, client: TestClient) -> None:
"""Non-debug requests forward to upstream; headers are set even if upstream fails."""
resp = client.post("/v1/chat/completions", json={
Expand Down
18 changes: 12 additions & 6 deletions uncommon_route/proxy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4622,6 +4622,10 @@ def _prepare_attempt(model_name: str) -> dict[str, Any]:
val = request.headers.get(key)
if val:
attempt_headers[key] = val
request_auth = str(request.headers.get("authorization", "") or "").strip()
request_bearer = ""
if request_auth.lower().startswith("bearer "):
request_bearer = request_auth[7:].strip()
if api_format == "anthropic" and "authorization" not in attempt_headers:
x_api_key = request.headers.get("x-api-key")
if x_api_key:
Expand Down Expand Up @@ -4721,12 +4725,14 @@ def _prepare_attempt(model_name: str) -> dict[str, Any]:
attempt_headers["x-api-key"] = attempt_provider_entry.api_key
else:
attempt_headers["authorization"] = f"Bearer {attempt_provider_entry.api_key}"
elif primary_key:
if attempt_native_anthropic_transport:
attempt_headers.pop("authorization", None)
attempt_headers["x-api-key"] = primary_key
else:
attempt_headers["authorization"] = f"Bearer {primary_key}"
else:
upstream_api_key = request_bearer or primary_key
if upstream_api_key:
if attempt_native_anthropic_transport:
attempt_headers.pop("authorization", None)
attempt_headers["x-api-key"] = upstream_api_key
else:
attempt_headers["authorization"] = f"Bearer {upstream_api_key}"
if attempt_native_anthropic_transport:
if "x-api-key" not in attempt_headers and "authorization" in attempt_headers:
bearer = attempt_headers["authorization"]
Expand Down