diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 1ced3014b..a73ac0ff7 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -8602,9 +8602,23 @@ class OpenAIHandlerMixin: f"[{request_id}] WS upstream failed ({_ws_detail}), " f"falling back to HTTP POST streaming" ) - await self._ws_http_fallback( + ( + fb_input_tokens, + fb_output_tokens, + fb_cache_read_tokens, + fb_cache_write_tokens, + fb_uncached_tokens, + ) = await self._ws_http_fallback( websocket, body, first_msg_raw, upstream_headers, request_id ) + # Fold the fallback's provider usage into the session totals so + # the WS session-end outcome records the authoritative wire-token + # count instead of 0 (#2957). + ws_input_tokens_total += fb_input_tokens + ws_output_tokens_total += fb_output_tokens + ws_cache_read_tokens_total += fb_cache_read_tokens + ws_cache_write_tokens_total += fb_cache_write_tokens + ws_uncached_input_tokens_total += fb_uncached_tokens # ── WS session-end metric + RequestLog ────────────────── # @@ -8850,14 +8864,31 @@ class OpenAIHandlerMixin: first_msg_raw: str, upstream_headers: dict[str, str], request_id: str, - ) -> None: + ) -> tuple[int, int, int, int, int]: """Fall back to HTTP POST streaming when upstream WS fails. Converts the WS ``response.create`` message to an HTTP POST to ``/v1/responses?stream=true``, reads SSE events, and relays each ``data:`` line as a WS text message to the client. This makes Codex work immediately instead of exhausting its WS retry budget. + + Returns ``(input, output, cache_read, cache_write, uncached)`` provider + usage parsed from the ``response.completed`` SSE event. The caller folds + it into the session totals so the WS session-end outcome uses the + authoritative wire-token count; otherwise a fallback recorded + ``input_tokens=0`` and savings percentages blew past 100 (#2957). """ + fallback_usage = [0, 0, 0, 0, 0] + + def _accumulate_usage(data_str: str) -> None: + try: + event = json.loads(data_str) + except (json.JSONDecodeError, TypeError): + return + if isinstance(event, dict) and event.get("type") == "response.completed": + for i, value in enumerate(_extract_responses_usage(event)): + fallback_usage[i] += value + # Route to correct endpoint based on auth mode is_chatgpt_fallback = has_chatgpt_account_header(upstream_headers) if is_chatgpt_fallback: @@ -8963,7 +8994,7 @@ class OpenAIHandlerMixin: }, } await websocket.send_text(json.dumps(error_event)) - return + return tuple(fallback_usage) # type: ignore[return-value] # Refresh Codex /stats from the fallback response # headers. We can't forward them onto the client 101 @@ -8989,10 +9020,11 @@ class OpenAIHandlerMixin: data = line[6:] if data == "[DONE]": continue + _accumulate_usage(data) try: await websocket.send_text(data) except Exception: - return + return tuple(fallback_usage) # type: ignore[return-value] elif line.startswith("event: "): # SSE event type — skip, the data line contains the type continue @@ -9001,9 +9033,10 @@ class OpenAIHandlerMixin: for line in buffer.strip().splitlines(): line = line.strip() if line.startswith("data: ") and line[6:] != "[DONE]": + _accumulate_usage(line[6:]) with contextlib.suppress(Exception): await websocket.send_text(line[6:]) - return + return tuple(fallback_usage) # type: ignore[return-value] except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as http_err: if http_attempt >= retry_attempts - 1: raise @@ -9034,6 +9067,7 @@ class OpenAIHandlerMixin: finally: with contextlib.suppress(Exception): await websocket.close() + return tuple(fallback_usage) # type: ignore[return-value] def _derived_compress_pipeline(self, key: str, **overrides: Any) -> Any: """Cached ``/v1/compress`` pipeline derived from the live OpenAI router. diff --git a/tests/test_ws_http_fallback.py b/tests/test_ws_http_fallback.py index 459246fb3..d0205c877 100644 --- a/tests/test_ws_http_fallback.py +++ b/tests/test_ws_http_fallback.py @@ -310,6 +310,53 @@ class TestWsHttpFallback: assert "api.openai.com" in captured_url["url"] + def test_fallback_returns_provider_usage_from_completed_event(self): + """The fallback must surface the provider's input usage (#2957). + + Otherwise the WS session-end outcome records input_tokens=0 for a large + request and savings percentages blow past 100. + """ + handler = _make_handler() + ws = FakeWebSocket() + completed = { + "type": "response.completed", + "response": { + "usage": { + "input_tokens": 31055, + "output_tokens": 246, + "input_tokens_details": {"cached_tokens": 20000}, + } + }, + } + sse_lines = [ + 'data: {"type":"response.created","response":{"id":"r1"}}\n\n', + f"data: {json.dumps(completed)}\n\n", + "data: [DONE]\n\n", + ] + handler.http_client = FakeHttpClient(FakeStreamResponse(200, sse_lines)) + + body = {"model": "gpt-5.4", "input": "big context"} + usage = asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), {}, "req_usage")) + + input_tokens, output_tokens, cache_read, _cache_write, uncached = usage + assert input_tokens == 31055 + assert output_tokens == 246 + assert cache_read == 20000 + assert uncached == 31055 - 20000 + + def test_fallback_returns_zero_usage_without_completed_event(self): + handler = _make_handler() + ws = FakeWebSocket() + handler.http_client = FakeHttpClient( + FakeStreamResponse(200, ['data: {"type":"response.created"}\n\n', "data: [DONE]\n\n"]) + ) + usage = asyncio.run( + handler._ws_http_fallback( + ws, {"model": "gpt-5.4", "input": "hi"}, json.dumps({"input": "hi"}), {}, "req_none" + ) + ) + assert usage == (0, 0, 0, 0, 0) + def test_fallback_refreshes_codex_rate_limit_state(self, monkeypatch): """A successful fallback refreshes Codex /stats from response headers.