fix(proxy/openai): feed Codex WS traffic into the traffic learner (#2334)

## 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 <same files + test_memory/test_traffic_learner.py> 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 <mxjerrett@gmail.com>
2026-08-12 10:24:41 +05:30
parent 64cb46e24b
commit f669149769
3 changed files with 180 additions and 1 deletions
@@ -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,
                    }
                )

@@ -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,
@@ -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"}