From f6691497692869b7067438597421ff12aace6bf4 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 12 Aug 2026 10:24:41 +0530 Subject: [PATCH] fix(proxy/openai): feed Codex WS traffic into the traffic learner (#2334) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Follow-up to the chat/completions ingestion work — this wires the Codex `/v1/responses` **WebSocket** path into the traffic learner, the remaining gap in #2060. `handle_openai_responses_ws` (the transport newer Codex versions default to) had no traffic-learner ingestion, so Codex subscription traffic produced no learned patterns even with Learn enabled. Unlike the one-shot HTTP path, a long-lived Codex WebSocket: - resends the **full transcript** on every `response.create` frame, and - replays it **wholesale on reconnect/resume**. So naive per-turn ingestion would count the same tool result as evidence over and over, and every reconnect would re-ingest the whole history. ## Fix Add `_observe_openai_ws_response_create`, which dedups per connection by tool-call id: - A per-connection `ws_learner_seen_call_ids: set[str]` tracks which tool-call ids have been observed on this WebSocket. - The **first** `response.create` frame is a **baseline**: its already-present transcript is recorded as seen but **not learned**, and preference extraction is skipped. This is the replayed/initial history, which may already have been learned on a prior connection. - **Later** frames learn only the tool results whose call id first appears after the baseline, then mark them seen. Preference extraction (`on_messages`) runs on these frames (it already looks only at the most recent messages). On reconnect the client opens a fresh WebSocket and replays the transcript in its first frame, which is baselined again, so it adds no spurious evidence. It hooks both frame paths: the first-frame handler seeds the baseline from the original client frame (parsed before memory injection / compression), and `_maybe_compress_response_create_frame` observes each subsequent frame. To dedup by identity, `TrafficLearner.extract_tool_results_from_messages` now also returns the `call_id` (the `tool_use`/`tool_result` id, which `_responses_input_to_learner_messages` already sets from the Responses `call_id`). This is additive — existing callers that don't read it are unaffected. Relationship to the chat path: the `/v1/chat/completions` ingestion is a separate change; together they cover HTTP chat, HTTP Responses (already wired), and Codex WS. This PR is independent and branches off `main`. ## 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 - `headroom/memory/traffic_learner.py`: `extract_tool_results_from_messages` now returns `call_id` for per-turn dedup (additive). - `headroom/proxy/handlers/openai.py`: add `_observe_openai_ws_response_create` (per-connection dedup + baseline); initialise `ws_learner_seen_call_ids`; observe the first frame as a baseline and each subsequent `response.create` frame. - `tests/test_openai_responses_traffic_learner.py`: add WS dedup/baseline coverage (baseline records-not-learns, later frames learn only new results, reconnect replay adds no evidence); update the existing extractor-equality assertion to include `call_id`. - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/memory/traffic_learner.py headroom/proxy/handlers/openai.py tests/test_openai_responses_traffic_learner.py All checks passed! $ uvx ruff@0.15.17 format --check all files already formatted $ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/openai.py # no errors in the changed files (the one reported error is a pre-existing # headroom/_subprocess.py:18 no-any-return, present on main with these edits stashed) ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx mypy@1.20.2`. A full `pytest` OOMs this box (ML-stack import), so I reproduced the dedup/baseline loop with a dependency-free asyncio script and left the full pytest to CI. - Exact command / steps: simulated a connection where the baseline frame carries tool-call ids A,B; later frames replay A,B and append C, then D; plus a reconnect whose first frame replays A,B,C,D. - Observed result: baseline recorded A,B without learning; frame 2 learned only C; frame 3 learned only D (A/B/C never re-counted); the reconnect's replayed transcript was baselined and learned nothing. The added unit tests assert the same through the real handler method with a recording learner. - Not tested: a live Codex WebSocket session end to end; the added tests drive `_observe_openai_ws_response_create` directly with a recording learner and the real `_responses_input_to_learner_messages` + extractor. ## 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 or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The "unit tests pass locally" box is unchecked because a local pytest run imports the ML stack and OOMs this box; the added tests reuse the existing `_RecordingLearner` harness (no real backend) and run under the normal CI pytest job, and the dedup/baseline behavior is corroborated by the standalone proof above. Design note: baselining the first frame means a brand-new conversation's first-turn tool results are not learned on that connection (subsequent turns are); this is the deliberate trade-off the issue calls for to keep reconnect/resume from inflating evidence. --------- Co-authored-by: JerrettDavis --- headroom/memory/traffic_learner.py | 5 + headroom/proxy/handlers/openai.py | 96 +++++++++++++++++++ .../test_openai_responses_traffic_learner.py | 80 +++++++++++++++- 3 files changed, 180 insertions(+), 1 deletion(-) diff --git a/headroom/memory/traffic_learner.py b/headroom/memory/traffic_learner.py index 387813911..612af996c 100644 --- a/headroom/memory/traffic_learner.py +++ b/headroom/memory/traffic_learner.py @@ -1476,6 +1476,11 @@ class TrafficLearner: "input": tool_use.get("input", {}), "output": str(result_content), "is_error": block.get("is_error", False) or _is_error(str(result_content)), + # Stable per-turn identity (the tool_use/tool_result id). + # Lets a caller dedup a replayed transcript so the same + # result is not counted as evidence twice — used by the + # Codex WebSocket ingestion path. + "call_id": tool_use_id, } ) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index b415d6e25..218ef7ec6 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1590,6 +1590,72 @@ class OpenAIHandlerMixin: except Exception as exc: logger.debug("[%s] Traffic learner (chat): %s", request_id, exc) + async def _observe_openai_ws_response_create( + self, + inner_payload: dict[str, Any], + *, + seen_call_ids: set[str], + baseline: bool, + request_id: str, + ) -> None: + """Feed one Codex WS ``response.create`` turn into the traffic learner. + + A long-lived Codex WebSocket resends the full transcript on every + ``response.create`` (and replays it wholesale on reconnect/resume), so + the HTTP one-shot ingestion would count the same tool results as new + evidence over and over. This dedups per connection by tool-call id: + + * ``baseline=True`` (the first frame) records the already-present + transcript as seen WITHOUT learning from it, so a reconnect that + replays old turns adds no spurious evidence. + * later frames learn only the tool results whose call id first appears + after the baseline, then mark them seen. + + Preference extraction (:meth:`on_messages`) is skipped on the baseline + frame and runs on later frames, where it already looks only at the most + recent messages (the new turn). + """ + traffic_learner = getattr(self, "traffic_learner", None) + if traffic_learner is None: + return + try: + memory_handler = getattr(self, "memory_handler", None) + if ( + traffic_learner._backend is None + and memory_handler + and memory_handler.initialized + and memory_handler.backend + ): + traffic_learner.set_backend(memory_handler.backend) + + learner_messages = _responses_input_to_learner_messages( + inner_payload.get("instructions"), + inner_payload.get("input", ""), + ) + tool_results = traffic_learner.extract_tool_results_from_messages(learner_messages) + for tool_result in tool_results: + call_id = tool_result.get("call_id") or "" + # A result already seen on this connection (or baselined) is not + # re-learned. Results without an id fall back to learn-once-per + # -frame on non-baseline frames (still deduped by the learner's + # own pattern accumulation). + if call_id: + if call_id in seen_call_ids: + continue + seen_call_ids.add(call_id) + if baseline: + continue + await traffic_learner.on_tool_result( + tool_name=tool_result["tool_name"], + tool_input=tool_result["input"], + tool_output=tool_result["output"], + is_error=tool_result["is_error"], + ) + if not baseline: + await traffic_learner.on_messages(learner_messages) + except Exception as exc: + logger.debug("[%s] Traffic learner (ws): %s", request_id, exc) + @staticmethod def _headroom_bypass_enabled(headers: Any) -> bool: """Return True when inbound headers request full passthrough.""" @@ -6594,6 +6660,27 @@ class OpenAIHandlerMixin: ws_recorded_tokens_saved_total = 0 ws_recorded_attempted_input_tokens_total = 0 ws_response_create_frames = 1 + # Per-connection traffic-learner dedup: tool-call ids already + # observed on this WS, so a replayed transcript (each turn resends + # the full history; reconnect replays it wholesale) is not counted + # as new evidence. Seeded from the first frame as a baseline. + ws_learner_seen_call_ids: set[str] = set() + # Baseline the first frame's transcript into the learner: record its + # tool-call ids as seen WITHOUT learning, so a reconnect that replays + # this history adds no spurious evidence. Later frames learn only the + # results appended after this point. `body` here is the original + # client frame (parsed before memory injection / compression). + if isinstance(body, dict) and body: + _ws_first_inner = ( + body["response"] if isinstance(body.get("response"), dict) else body + ) + if isinstance(_ws_first_inner, dict): + await self._observe_openai_ws_response_create( + _ws_first_inner, + seen_call_ids=ws_learner_seen_call_ids, + baseline=True, + request_id=request_id, + ) ws_client_frames_total = 1 ws_upstream_frames_total = 0 ws_cancel_frames = 0 @@ -7267,6 +7354,15 @@ class OpenAIHandlerMixin: frame_type="response.create", ) return raw_msg, False, "invalid_inner_payload" + # Learn from this turn's newly appended tool results. + # Dedup against the per-connection baseline so the + # replayed transcript prefix is not re-counted. + await self._observe_openai_ws_response_create( + inner_payload, + seen_call_ids=ws_learner_seen_call_ids, + baseline=False, + request_id=request_id, + ) store_forced = _ensure_chatgpt_responses_store_false( inner_payload, is_chatgpt_auth=is_chatgpt_auth, diff --git a/tests/test_openai_responses_traffic_learner.py b/tests/test_openai_responses_traffic_learner.py index 91e710390..9f1f7397c 100644 --- a/tests/test_openai_responses_traffic_learner.py +++ b/tests/test_openai_responses_traffic_learner.py @@ -1,12 +1,16 @@ from __future__ import annotations +import asyncio from typing import Any import httpx from fastapi.testclient import TestClient from headroom.memory.traffic_learner import TrafficLearner -from headroom.proxy.handlers.openai import _responses_input_to_learner_messages +from headroom.proxy.handlers.openai import ( + OpenAIHandlerMixin, + _responses_input_to_learner_messages, +) from headroom.proxy.server import ProxyConfig, create_app @@ -80,6 +84,7 @@ def test_responses_input_normalizes_messages_and_tool_results() -> None: "input": {"cmd": "missing-command"}, "output": "command not found", "is_error": True, + "call_id": "call_1", } ] @@ -141,3 +146,76 @@ def test_responses_http_request_reaches_traffic_learner() -> None: "is_error": True, } ] + + +def _ws_frame(call_ids: list[str]) -> dict[str, Any]: + """A response.create inner payload whose input carries one shell tool + round-trip per call id.""" + input_items: list[dict[str, Any]] = [] + for cid in call_ids: + input_items.append( + {"type": "function_call", "call_id": cid, "name": "shell", "arguments": "{}"} + ) + input_items.append( + { + "type": "function_call_output", + "call_id": cid, + "output": "ok", + "status": "completed", + } + ) + return {"input": input_items} + + +def test_ws_response_create_baselines_and_dedups_replayed_transcript() -> None: + handler = OpenAIHandlerMixin() + learner = _RecordingLearner() + handler.traffic_learner = learner + seen: set[str] = set() + + # First frame is the baseline: A and B are recorded as seen but NOT learned, + # and preference extraction is skipped. + asyncio.run( + handler._observe_openai_ws_response_create( + _ws_frame(["A", "B"]), seen_call_ids=seen, baseline=True, request_id="r" + ) + ) + assert learner.tool_results == [] + assert learner.message_batches == [] + assert seen == {"A", "B"} + + # Second frame replays A, B and appends C -> only C is learned. + asyncio.run( + handler._observe_openai_ws_response_create( + _ws_frame(["A", "B", "C"]), seen_call_ids=seen, baseline=False, request_id="r" + ) + ) + assert len(learner.tool_results) == 1 + assert seen == {"A", "B", "C"} + assert len(learner.message_batches) == 1 + + # Third frame replays A, B, C and appends D -> only D is learned. + asyncio.run( + handler._observe_openai_ws_response_create( + _ws_frame(["A", "B", "C", "D"]), seen_call_ids=seen, baseline=False, request_id="r" + ) + ) + assert len(learner.tool_results) == 2 # C then D, never A/B again + assert seen == {"A", "B", "C", "D"} + + +def test_ws_reconnect_replay_adds_no_evidence() -> None: + # A reconnect is a fresh connection: its first frame replays the whole + # transcript, which is baselined, so nothing is re-learned. + handler = OpenAIHandlerMixin() + learner = _RecordingLearner() + handler.traffic_learner = learner + seen: set[str] = set() + + asyncio.run( + handler._observe_openai_ws_response_create( + _ws_frame(["A", "B", "C", "D"]), seen_call_ids=seen, baseline=True, request_id="r" + ) + ) + assert learner.tool_results == [] + assert seen == {"A", "B", "C", "D"}