diff --git a/tests/test_proxy.py b/tests/test_proxy.py index 08427f0..1fa36bd 100644 --- a/tests/test_proxy.py +++ b/tests/test_proxy.py @@ -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={ diff --git a/uncommon_route/proxy.py b/uncommon_route/proxy.py index fc341aa..009ef25 100644 --- a/uncommon_route/proxy.py +++ b/uncommon_route/proxy.py @@ -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: @@ -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"]