diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 4b096e070..e8e4b4e18 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -7779,7 +7779,8 @@ class OpenAIHandlerMixin: # as HTTP turns. await self._record_request_outcome( RequestOutcome( - request_id=request_id, + # Per-emission ids keep dashboard request-log keys unique. + request_id=await self._next_request_id(), provider="openai", model=model_for_metrics, original_tokens=max(0, input_delta) + max(0, saved_delta), @@ -8354,7 +8355,8 @@ class OpenAIHandlerMixin: ws_messages_for_log.append({"role": "user", "content": ws_input_for_log}) await self._record_request_outcome( RequestOutcome( - request_id=request_id, + # Per-emission ids keep dashboard request-log keys unique. + request_id=await self._next_request_id(), provider="openai", model=model_name, original_tokens=residual_input_tokens + residual_tokens_saved, diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 481496b4e..b89a81a7e 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -114,6 +114,14 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): await emit_request_outcome(self, outcome) +class _CapturingLogger: + def __init__(self) -> None: + self.entries = [] + + def log(self, entry) -> None: # noqa: ANN001 + self.entries.append(entry) + + class _FakeWebSocketDisconnect(Exception): """Mirrors the ``WebSocketDisconnect`` type-name check in the handler. @@ -894,6 +902,138 @@ async def test_ws_session_metrics_include_dashboard_performance_timings(): ) +@pytest.mark.asyncio +async def test_ws_multi_turn_request_ids_are_unique(): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps( + { + "type": "response.completed", + "response": { + "id": "r_1", + "usage": { + "input_tokens": 100, + "input_tokens_details": {"cached_tokens": 75}, + "output_tokens": 12, + }, + }, + } + ), + json.dumps({"type": "response.created", "response": {"id": "r_2"}}), + json.dumps( + { + "type": "response.completed", + "response": { + "id": "r_2", + "usage": { + "input_tokens": 160, + "input_tokens_details": {"cached_tokens": 120}, + "output_tokens": 20, + }, + }, + } + ), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket(frames=[_first_frame()]) + handler = _DummyOpenAIHandler() + handler.logger = _CapturingLogger() + + counter = 0 + + async def _next_request_id() -> str: + nonlocal counter + counter += 1 + return f"req-ws-{counter}" + + handler._next_request_id = _next_request_id # type: ignore[method-assign] + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + logged = handler.logger.entries + assert len(logged) == 2 + request_ids = [entry.request_id for entry in logged] + assert len(set(request_ids)) == len(request_ids) + assert [entry.input_tokens_optimized for entry in logged] == [100, 160] + assert [entry.output_tokens for entry in logged] == [12, 20] + + +@pytest.mark.asyncio +async def test_ws_no_delta_turn_emits_no_extra_request_log(): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket(frames=[_first_frame()]) + handler = _DummyOpenAIHandler() + handler.logger = _CapturingLogger() + + counter = 0 + + async def _next_request_id() -> str: + nonlocal counter + counter += 1 + return f"req-ws-{counter}" + + handler._next_request_id = _next_request_id # type: ignore[method-assign] + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(handler.logger.entries) == 1 + assert all(entry.input_tokens_optimized == 0 for entry in handler.logger.entries) + assert all(entry.output_tokens == 0 for entry in handler.logger.entries) + + +@pytest.mark.asyncio +async def test_ws_session_log_prefix_uses_session_id(caplog: pytest.LogCaptureFixture): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps( + { + "type": "response.completed", + "response": { + "id": "r_1", + "usage": { + "input_tokens": 100, + "input_tokens_details": {"cached_tokens": 75}, + "output_tokens": 12, + }, + }, + } + ), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket(frames=[_first_frame()]) + handler = _DummyOpenAIHandler() + handler.logger = _CapturingLogger() + + counter = 0 + + async def _next_request_id() -> str: + nonlocal counter + counter += 1 + return f"req-ws-{counter}" + + handler._next_request_id = _next_request_id # type: ignore[method-assign] + caplog.set_level(logging.INFO, logger="headroom.proxy") + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert handler.logger.entries + assert handler.logger.entries[0].request_id != "req-ws-1" + assert "[req-ws-1] PERF" in caplog.text + + @pytest.mark.asyncio async def test_ws_opt_in_flattens_response_create_for_openai_compatible_upstream(monkeypatch): """Some OpenAI-compatible WS gateways expect top-level response.create payloads."""