diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 1024238c3..d2a2e39f1 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -75,7 +75,11 @@ from headroom.proxy.outcome import RequestOutcome from headroom.proxy.passthrough import ( custom_base_passthrough_telemetry as _custom_base_passthrough_telemetry, ) -from headroom.proxy.project_context import classify_project, set_current_project +from headroom.proxy.project_context import ( + classify_project, + get_current_project, + set_current_project, +) from headroom.proxy.token_counting import gemini_output_tokens logger = logging.getLogger("headroom.proxy") @@ -5726,8 +5730,11 @@ class OpenAIHandlerMixin: # Captured in closure so per-turn RequestOutcome can stamp it. client = classify_client(ws_headers) # WS sessions bypass the HTTP middleware, so bind the project here; - # per-turn outcome emission inside this task inherits the context. - set_current_project(classify_project(ws_headers)) + # per-turn outcome emission inside this task inherits the context. An + # explicit X-Headroom-Project header wins; otherwise fall back to the + # /p/ path prefix already bound by WebSocketProjectPrefixMiddleware + # so prefix-only clients (aider, Copilot BYOK, Cursor) stay attributed. + set_current_project(classify_project(ws_headers) or get_current_project()) metrics_for_inbound_ws = getattr(self, "metrics", None) if metrics_for_inbound_ws is not None and hasattr( metrics_for_inbound_ws, "record_inbound_request" diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index a9ac18caa..c03c76ee8 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -2364,6 +2364,22 @@ _is_known_websocket_callback_failure = is_known_websocket_callback_failure _tool_schema_saved_from_tags = tool_schema_saved_from_tags +class WebSocketProjectPrefixMiddleware: + """Normalize project-prefixed WebSocket paths before route matching.""" + + def __init__(self, app: Any) -> None: + self.app = app + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + if scope["type"] == "websocket": + prefix_project = strip_project_path_prefix(scope) + headers = { + name.decode("latin-1"): value.decode("latin-1") for name, value in scope["headers"] + } + set_current_project(classify_project(headers) or prefix_project) + await self.app(scope, receive, send) + + def create_app(config: ProxyConfig | None = None) -> FastAPI: """Create FastAPI application.""" if not FASTAPI_AVAILABLE: @@ -2573,6 +2589,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: version=__version__, lifespan=lifespan, ) + app.add_middleware(WebSocketProjectPrefixMiddleware) loop_health_state: LoopHealthState = { "status": "healthy", "known_failures": 0, diff --git a/tests/test_provider_proxy_routes.py b/tests/test_provider_proxy_routes.py index 7a6ebb599..fe7f9ddd0 100644 --- a/tests/test_provider_proxy_routes.py +++ b/tests/test_provider_proxy_routes.py @@ -10,6 +10,7 @@ from fastapi.responses import JSONResponse from fastapi.testclient import TestClient from headroom.providers.codex.runtime import CodexRoutingDecision +from headroom.proxy.project_context import get_current_project from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app @@ -481,6 +482,31 @@ def test_openai_response_websocket_aliases_delegate_to_openai_ws_handler(monkeyp ] +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/ 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]