fix(proxy): preserve Codex WebSocket model attribution (#3029)
## Description Codex can switch models during a multi-turn Responses WebSocket conversation. Headroom was not consistently attributing each completed turn to the model that handled it, which made per-model usage and savings reporting inaccurate. Closes #3027 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Attribute each completed WebSocket response to its reported model. - Keep session-end metrics consistent with the response that completed. - Add a regression test covering two different models on one WebSocket session. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added new functionality - [x] Manual testing performed ### Test Output ```text uv run pytest -q tests/test_openai_codex_ws_lifecycle.py -k session_metrics_track_model_per_response_create 1 passed, 51 deselected in 2.09s Full Codex WebSocket lifecycle module: 52 passed Adjacent Codex WebSocket suites: 77 passed, 1 skipped uv run ruff check . All checks passed uv run ruff format --check . 1411 files already formatted uv run mypy headroom Success: no issues found in 520 source files ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.3, OpenAI Codex Responses WebSocket. - Exact command / steps: From the repository root, run `uv sync --extra dev --extra proxy`, then run `uv run headroom wrap codex`; in one live Codex conversation complete one turn with model A, switch to model B, complete a second turn, and inspect the proxy dashboard or `http://localhost:8787/stats` recent requests. - Observed result: Both completed turns appeared under the models that handled them, in order. - Not tested: Production deployment and non-Codex transports. ## Runtime Rollout Safety - Rollout-managed feature(s): None. - Minimum rollout channel: Stable/default. - Stable/default behavior changed: Corrects telemetry attribution only; no public API or routing changes. - Kill switch / disable path: Revert the change or use the previous release. - Unsafe override required: No. - Qualification impact: None. - Rollback path: Revert commit `d5d8d7ca`. ## 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 - [ ] 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 or that my feature works - [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 (a CI guard enforces this) ## Screenshots (if applicable) ### Pre change In a single session, started with `5.6-sol` and then switched to `5.6-luna`. The dashboard did not reflect the model change. <img width="1274" height="207" alt="image" src="https://github.com/user-attachments/assets/697803e4-d33d-4660-b8dd-f1a8d6404517" /> ### After change Repeated the same steps: started with `5.6-sol` and switched to `5.6-luna`. The dashboard now correctly reflects the model change. <img width="1264" height="202" alt="image" src="https://github.com/user-attachments/assets/722e5fab-ad1b-4e62-9344-5f9dd312d614" /> ## Additional Notes
This commit is contained in:
Gil Korzen
@@ -7194,7 +7194,6 @@ class OpenAIHandlerMixin:
|
||||
if isinstance(first_response_body, dict)
|
||||
else None
|
||||
)
|
||||
|
||||
# Hot-fix follow-up to PR #406 — inline Rust compression on the
|
||||
# WS first frame before forwarding upstream. PR #406 enabled
|
||||
# the same call for HTTP /v1/responses; PR-C5's "WS-side
|
||||
@@ -8034,6 +8033,7 @@ class OpenAIHandlerMixin:
|
||||
response_output_items.clear()
|
||||
|
||||
response_started_ms: float | None = None
|
||||
completed_response_model = "unknown"
|
||||
|
||||
async def _record_ws_response_metrics() -> None:
|
||||
"""Record one completed Responses turn on long-lived WS sessions."""
|
||||
@@ -8090,7 +8090,7 @@ class OpenAIHandlerMixin:
|
||||
):
|
||||
return
|
||||
|
||||
model_for_metrics = str(body.get("model") or "unknown")
|
||||
model_for_metrics = completed_response_model
|
||||
latency_ms = (
|
||||
(time.perf_counter() * 1000.0 - response_started_ms)
|
||||
if response_started_ms is not None
|
||||
@@ -8243,6 +8243,13 @@ class OpenAIHandlerMixin:
|
||||
upstream_frame_index,
|
||||
ws_last_upstream_frame_type,
|
||||
)
|
||||
response = event.get("response")
|
||||
completed_response_model = (
|
||||
str(response.get("model") or "unknown")
|
||||
if isinstance(response, dict)
|
||||
else "unknown"
|
||||
)
|
||||
|
||||
if event_type == "response.created":
|
||||
response_started_ms = time.perf_counter() * 1000.0
|
||||
(
|
||||
@@ -8615,11 +8622,7 @@ class OpenAIHandlerMixin:
|
||||
)
|
||||
if not isinstance(ws_inner_for_telemetry, dict):
|
||||
ws_inner_for_telemetry = {}
|
||||
model_name = (
|
||||
ws_inner_for_telemetry.get("model")
|
||||
or (body.get("model") if isinstance(body, dict) else None)
|
||||
or "unknown"
|
||||
)
|
||||
model_name = str(current_response_template.get("model") or "unknown")
|
||||
_final_auth_mode = classify_auth_mode(ws_headers)
|
||||
residual_input_tokens = max(0, ws_input_tokens_total - ws_recorded_input_tokens_total)
|
||||
residual_output_tokens = max(
|
||||
|
||||
@@ -2243,3 +2243,56 @@ async def test_ws_memory_continuation_continues_pre_stream_and_passes_late_call(
|
||||
assert second_response[6]["item"] == function_call_two
|
||||
assert second_response[7]["response"]["id"] == "r-2"
|
||||
assert executed == [("memory_search", {}, "user-1", "openai")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_session_metrics_track_model_per_response_create():
|
||||
"""A model switch on one WS session must affect the next request outcome."""
|
||||
upstream_events = [
|
||||
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "r_1",
|
||||
"model": "model-a",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
},
|
||||
}
|
||||
),
|
||||
json.dumps({"type": "response.created", "response": {"id": "r_2"}}),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "r_2",
|
||||
"model": "model-b",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
},
|
||||
}
|
||||
),
|
||||
]
|
||||
first_frame = json.dumps(
|
||||
{
|
||||
"type": "response.create",
|
||||
"response": {"model": "model-a", "input": "first turn"},
|
||||
}
|
||||
)
|
||||
second_frame = json.dumps(
|
||||
{
|
||||
"type": "response.create",
|
||||
"response": {"model": "model-b", "input": "second turn"},
|
||||
}
|
||||
)
|
||||
upstream = _FakeUpstream(upstream_events)
|
||||
fake_ws_mod = _make_fake_websockets_module(upstream)
|
||||
client_ws = _FakeWebSocket(frames=[first_frame, second_frame])
|
||||
handler = _DummyOpenAIHandler()
|
||||
|
||||
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
|
||||
await handler.handle_openai_responses_ws(client_ws)
|
||||
|
||||
assert [request["model"] for request in handler.metrics.recorded_requests] == [
|
||||
"model-a",
|
||||
"model-b",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user