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