github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy/openai): propagate provider usage on the Responses WS->HTTP fallback (#2988)

Browse Source 
## Description

When Codex uses the OpenAI Responses WebSocket endpoint through Headroom
and the upstream WebSocket is rejected, Headroom falls back to HTTPS
POST/SSE. On that fallback the dashboard reported zero or tiny input
tokens for a large request, and invalid savings:

```json
{ "input_tokens_original": 3, "input_tokens_optimized": 0,
  "output_tokens": 246, "tokens_saved": 31052, "savings_percent": 33233.33 }
```

## Root cause

`_ws_http_fallback` (openai.py) relays the SSE `data:` events to the
client but never parses the terminal `response.completed` event for
usage. The non-fallback WS path accumulates
`_extract_responses_usage(event)` into the session totals on every
`response.completed` frame (openai.py ~8182); the fallback path did not.
So `ws_input_tokens_total` stayed at the small local count, and the
session-end RequestLog computed `optimized_tokens =
residual_input_tokens = 0`, leaving `tokens_saved >
input_tokens_original` and `savings_percent` far above 100%.

## Fix

`_ws_http_fallback` now parses each relayed `response.completed` line
with the existing `_extract_responses_usage` and returns the accumulated
`(input, output, cache_read, cache_write, uncached)` provider usage. The
caller folds it into the WS session totals, so the session-end outcome
uses the authoritative provider wire-token count -- bringing the
fallback to parity with the non-fallback WS path. SSE relay behaviour is
otherwise unchanged.

Fixes #2957

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature
- [ ] Breaking change
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/handlers/openai.py` (`_ws_http_fallback`): accumulate
usage from `response.completed` SSE lines (both the main relay loop and
the buffer flush) and return the `(input, output, cache_read,
cache_write, uncached)` tuple from every exit path; the WS handler
caller adds it to `ws_input_tokens_total` / `ws_output_tokens_total` /
cache / uncached totals before the session-end RequestLog.
- `tests/test_ws_http_fallback.py`: the fallback returns the provider
usage from a `response.completed` event
(input/output/cache_read/uncached), and returns all-zeros when no
completed event arrives.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`)
- [x] New tests added

### Test Output

```text
tests/test_ws_http_fallback.py  13 passed  (11 existing + 2 new)
# uvx ruff@0.15.22 check  -> All checks passed!
# uvx mypy@1.20.2 headroom/proxy/handlers/openai.py -> Success: no issues found in 1 source file
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12.11, project venv, pytest 9.1.1,
ruff 0.15.22 and mypy 1.20.2 via uvx.
- Exact command / steps: drove `_ws_http_fallback` with the existing
WS/stream mocks, feeding an SSE `response.completed` carrying
`usage.input_tokens=31055`, `output_tokens=246`,
`input_tokens_details.cached_tokens=20000`. The method now returns
`(31055, 246, 20000, ..., 11055)`; a stream with no completed event
returns all zeros. The existing 11 relay/routing/retry tests are
unchanged (they ignore the new return value).
- Observed result: the fallback surfaces the provider's real input
usage, so the WS session-end outcome records the actual input tokens
instead of 0, and savings percentages stay within a meaningful range.
- Not tested: a live Codex WS session that triggers the upstream-WS
rejection and HTTP fallback end to end (needs a real upstream refusing
the WS). The usage-propagation contract is verified at the fallback
boundary with the same mocks the existing fallback tests use.

## Runtime Rollout Safety

- Rollout-managed feature(s): none. The OpenAI Responses WS-to-HTTP
fallback is always-on transport behavior, not rollout-channel-gated.
- Minimum rollout channel: N/A (no rollout-managed behavior).
- Stable/default behavior changed: yes, as a bug fix. On the WS-to-HTTP
fallback the session-end outcome now records the provider's real
input/output/cache usage from `response.completed` instead of leaving
`ws_input_tokens_total` at 0 (which produced >100% savings). SSE relay
to the client is unchanged.
- Kill switch / disable path: N/A. This corrects accounting only; there
is no behavioral toggle and no user-facing surface beyond the recorded
outcome numbers.
- Unsafe override required: no.
- Qualification impact: fallback-path token accounting now matches the
non-fallback WS path and the HTTP Responses path (all three use
`_extract_responses_usage`); savings percentages return to a valid
range.
- Rollback path: revert this PR; the fallback returns to reporting zero
input usage on this path.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`: it is generated by
release-please from my Conventional Commit PR title

## Additional Notes

The fix reuses the already-present `_extract_responses_usage` (same
parser the non-fallback WS path and HTTP Responses path use), so
cache-read/write and uncached accounting stay consistent across all
three transports.

Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
Abhay Singh
2026-08-17 03:34:39 +05:30
committed by GitHub
parent a06a51eca6
commit 536c949a69
2 changed files with 86 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/proxy/handlers/openai.py
+39 -5
View File
@@ -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.
tests/test_ws_http_fallback.py
+47
View File
@@ -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.