headroomlabs-ai
diff --git a/‎headroom/proxy/handlers/openai.py‎
Lines changed: 39 additions & 5 deletions b/‎headroom/proxy/handlers/openai.py‎
Lines changed: 39 additions & 5 deletions
diff --git a/‎tests/test_ws_http_fallback.py‎
Lines changed: 47 additions & 0 deletions b/‎tests/test_ws_http_fallback.py‎
Lines changed: 47 additions & 0 deletions
@@ -8602,9 +8602,23 @@ async def _record_ws_response_metrics() -> None:
                     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 @@ async def _ws_http_fallback(
         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 @@ async def _ws_http_fallback(
                                 },
                             }
                             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 @@ async def _ws_http_fallback(
                                     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 @@ async def _ws_http_fallback(
                         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 @@ async def _ws_http_fallback(
         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.
 
@@ -310,6 +310,53 @@ def stream(self, method, url, **kwargs):
 
         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.