From 6cdfd3f64d2f64d50ed47644126df71872a21050 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Sat, 18 Jul 2026 22:39:01 +0530 Subject: [PATCH] fix(proxy/openai): feed chat/completions traffic into the traffic learner (#2333) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Addresses the chat/completions portion of #2060. The live traffic learner is wired into the Anthropic `/v1/messages` handler and, since then, the OpenAI Responses HTTP handler (`_observe_openai_responses_traffic`, called from `handle_openai_responses`). But `handle_openai_chat` has **no** ingestion call site: ```text headroom/proxy/handlers/openai.py handle_openai_responses -> _observe_openai_responses_traffic (wired) handle_openai_chat -> (no traffic_learner call) (gap) ``` So OpenAI-compatible clients that route through `/v1/chat/completions` — GitHub Copilot CLI, opencode, OpenAI SDKs — run through an apparently healthy proxy with Learn enabled while producing no learned patterns: the learner starts, but it never receives their tool results or user messages. ## Fix Observe the original client payload (before memory/compression mutates it) at the top of `handle_openai_chat`, mirroring the Responses and Anthropic ingestion paths: ```python await self._observe_openai_chat_traffic(original_client_messages, request_id=request_id) ``` `_observe_openai_chat_traffic` is the chat counterpart of `_observe_openai_responses_traffic`: same lazy backend wiring, same `on_tool_result` / `on_messages` lifecycle, same fail-soft `try/except`. The one format-specific piece is tool-result extraction. chat/completions encodes tool calls differently from Anthropic — the call is on an assistant message's `tool_calls` array (`id` -> function `name` + `arguments`) and each result is a separate `role: "tool"` message keyed by `tool_call_id`, so the existing `extract_tool_results_from_messages` (which scans for Anthropic `tool_use`/`tool_result` blocks) finds nothing. A new `TrafficLearner.extract_tool_results_from_openai_messages`: - builds the `tool_call_id -> function` map from assistant `tool_calls`; - for each `role: "tool"` message, resolves the tool name and joins string-or-list content; - parses the OpenAI `arguments` JSON string into a dict, so the downstream environment/recovery extractors (which call `input.get("command")`, `input.get("file_path")`, ...) see the same shape as an Anthropic `tool_use.input` instead of a raw string; - sniffs `is_error` from the output (chat tool messages carry no error flag). It returns the same `{tool_name, input, output, is_error}` shape as the Anthropic extractor, so `on_tool_result` stays format-agnostic. User-message preference extraction (`on_messages`) already reads plain `role`/`content`, so it consumes chat messages unchanged. Scope: this wires the **chat/completions** path. Codex WebSocket ingestion (`handle_openai_responses_ws`) additionally needs per-`response.create` evaluation plus transcript-replay baselining on reconnect, so it is intentionally left as a follow-up rather than half-implemented here. ## 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`: add `extract_tool_results_from_openai_messages` (OpenAI chat tool-result extraction with `arguments` JSON parsed to a dict). - `headroom/proxy/handlers/openai.py`: add `_observe_openai_chat_traffic` and call it from `handle_openai_chat` on the original client payload. - `tests/test_memory/test_traffic_learner.py`: cover the OpenAI extractor (name resolution, arguments parsing, list content, error sniff, malformed/orphan handling, empty case). - `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_memory/test_traffic_learner.py All checks passed! $ uvx ruff@0.15.17 format --check 3 files already formatted $ uvx mypy@1.20.2 --ignore-missing-imports headroom/memory/traffic_learner.py # clean for this file (the one reported error is a pre-existing # headroom/_subprocess.py:18 no-any-return, unrelated to this change and # 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 extractor with a dependency-free script and left the full pytest to CI. - Exact command / steps: replicated `extract_tool_results_from_openai_messages` and ran it over a typical chat round-trip (assistant `tool_calls` for `bash` + `read_file`, then two `role: "tool"` results, one erroring and one with list content), plus malformed-`arguments`, orphan-`tool_call_id`, and no-tool cases. - Observed result: tool names resolved from the call-id map; `arguments` parsed to a dict so `input.get("command")` works; list content joined; `is_error` sniffed from output; malformed arguments degrade to `{}` and an orphan id yields `unknown` without raising. The added unit tests assert the same through a real `TrafficLearner`. - Not tested: a live Copilot CLI session end to end; the added tests drive `TrafficLearner.extract_tool_results_from_openai_messages` directly, matching the existing `test_extract_tool_results_from_messages` pattern. ## 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 `TrafficLearner(backend=None, ...)` harness in `test_traffic_learner.py` (no real backend) and run under the normal CI pytest job, and the extractor behavior is corroborated by the standalone proof above. This PR is deliberately scoped to `/v1/chat/completions`; I'm happy to follow up with the Codex WebSocket ingestion path (which needs the transcript-replay baselining discussed in the issue) as a separate change if useful. --------- Co-authored-by: JerrettDavis --- headroom/memory/traffic_learner.py | 78 +++++++++++++++++++++++ headroom/proxy/handlers/openai.py | 47 ++++++++++++++ tests/test_memory/test_traffic_learner.py | 76 ++++++++++++++++++++++ 3 files changed, 201 insertions(+) diff --git a/headroom/memory/traffic_learner.py b/headroom/memory/traffic_learner.py index 60b5b1e68..1f80a4ae9 100644 --- a/headroom/memory/traffic_learner.py +++ b/headroom/memory/traffic_learner.py @@ -1399,6 +1399,84 @@ class TrafficLearner: return results + def extract_tool_results_from_openai_messages( + self, + messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Extract tool results from OpenAI chat/completions-format messages. + + The OpenAI counterpart of :meth:`extract_tool_results_from_messages`. + Chat/completions represents tool calls and their results differently + from Anthropic: the call lives on an assistant message's ``tool_calls`` + array (``id`` -> function ``name`` + ``arguments``), and each result is + a separate ``role: "tool"`` message keyed by ``tool_call_id``. + + Returns the same ``{tool_name, input, output, is_error}`` shape as the + Anthropic extractor so :meth:`on_tool_result` stays format-agnostic. The + OpenAI ``arguments`` JSON string is parsed into a dict so the downstream + environment/recovery extractors (which call ``input.get(...)``) see the + same shape as an Anthropic ``tool_use.input``. + """ + results: list[dict[str, Any]] = [] + + # Build tool_call_id -> function (name, arguments) from assistant turns. + tool_calls: dict[str, dict[str, Any]] = {} + for msg in messages: + if not isinstance(msg, dict) or msg.get("role") != "assistant": + continue + calls = msg.get("tool_calls") + if not isinstance(calls, list): + continue + for call in calls: + if not isinstance(call, dict): + continue + call_id = call.get("id", "") + function = call.get("function") + if isinstance(function, dict) and call_id: + tool_calls[call_id] = function + + for msg in messages: + if not isinstance(msg, dict) or msg.get("role") != "tool": + continue + function = tool_calls.get(msg.get("tool_call_id", ""), {}) + + # Tool-message content is usually a string, but the spec also allows + # a list of content parts. + result_content = msg.get("content", "") + if isinstance(result_content, list): + result_content = " ".join( + b.get("text", "") + for b in result_content + if isinstance(b, dict) and b.get("type") == "text" + ) + output = str(result_content) + + # Normalize the OpenAI ``arguments`` JSON string into a dict so the + # downstream extractors that call ``input.get(...)`` don't blow up. + raw_args = function.get("arguments", {}) + if isinstance(raw_args, dict): + tool_input: dict[str, Any] = raw_args + elif isinstance(raw_args, str) and raw_args: + try: + parsed = json.loads(raw_args) + except (ValueError, TypeError): + parsed = None + tool_input = parsed if isinstance(parsed, dict) else {} + else: + tool_input = {} + + # OpenAI tool messages carry no is_error flag; sniff the output. + results.append( + { + "tool_name": function.get("name", "unknown"), + "input": tool_input, + "output": output, + "is_error": _is_error(output), + } + ) + + return results + # ============================================================================= # Module helpers: project routing, memory.db loading, recommendation build diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 701a67261..48fa5e03c 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1319,6 +1319,47 @@ class OpenAIHandlerMixin: except Exception as exc: logger.debug("[%s] Traffic learner (responses): %s", request_id, exc) + async def _observe_openai_chat_traffic( + self, + messages: list[dict[str, Any]], + *, + request_id: str, + ) -> None: + """Feed one chat/completions request into the live traffic learner. + + The chat counterpart of :meth:`_observe_openai_responses_traffic`. + Chat/completions clients (GitHub Copilot CLI, opencode, OpenAI SDKs) + route here rather than through ``/v1/responses``, so without this call + their tool results and user preferences never reached the learner even + with Learn enabled (part of #2060). Chat messages are already + ``role``/``content`` shaped, so ``on_messages`` consumes them directly; + tool results use the OpenAI-format extractor. + """ + 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) + + tool_results = traffic_learner.extract_tool_results_from_openai_messages(messages) + for tool_result in tool_results[-5:]: + 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"], + ) + await traffic_learner.on_messages(messages) + except Exception as exc: + logger.debug("[%s] Traffic learner (chat): %s", request_id, exc) + @staticmethod def _headroom_bypass_enabled(headers: Any) -> bool: """Return True when inbound headers request full passthrough.""" @@ -2558,6 +2599,12 @@ class OpenAIHandlerMixin: stream = body.get("stream", False) + # Learn from the original client payload before memory context or + # compression mutates it, mirroring the Responses and Anthropic + # ingestion paths. Without this, chat/completions traffic (Copilot CLI, + # opencode, OpenAI SDKs) fed nothing to the learner (part of #2060). + await self._observe_openai_chat_traffic(original_client_messages, request_id=request_id) + # Bypass: skip ALL compression for explicit opt-out _bypass = self._headroom_bypass_enabled(request.headers) if _bypass: diff --git a/tests/test_memory/test_traffic_learner.py b/tests/test_memory/test_traffic_learner.py index cca523f39..5591053dc 100644 --- a/tests/test_memory/test_traffic_learner.py +++ b/tests/test_memory/test_traffic_learner.py @@ -416,6 +416,82 @@ class TestTrafficLearner: assert "file1.py" in results[0]["output"] assert not results[0]["is_error"] + def test_extract_tool_results_from_openai_messages(self, learner: TrafficLearner): + """OpenAI chat/completions tool results: assistant tool_calls + role:tool. + + Regression for the chat-path portion of #2060 — the extractor must + resolve the tool name from the assistant ``tool_calls`` id map, parse the + ``arguments`` JSON string into a dict (so downstream ``input.get(...)`` + works), join list content, and sniff errors from the output. + """ + messages = [ + {"role": "user", "content": "run the tests"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command": "pytest -q"}'}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "read_file", "arguments": '{"file_path": "/a/b.py"}'}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "Traceback (most recent call last):\nModuleNotFoundError: No module named x", + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": [{"type": "text", "text": "file body"}], + }, + ] + + results = learner.extract_tool_results_from_openai_messages(messages) + assert len(results) == 2 + + by_name = {r["tool_name"]: r for r in results} + assert by_name["bash"]["input"] == {"command": "pytest -q"} # parsed to dict + assert by_name["bash"]["input"].get("command") == "pytest -q" # downstream .get works + assert by_name["bash"]["is_error"] is True + + assert by_name["read_file"]["input"] == {"file_path": "/a/b.py"} + assert by_name["read_file"]["output"] == "file body" # list content joined + assert by_name["read_file"]["is_error"] is False + + def test_extract_openai_tool_results_handles_malformed_and_orphans( + self, learner: TrafficLearner + ): + """Malformed arguments become an empty dict; an unmatched tool_call_id + yields ``unknown`` — neither raises, so on_tool_result stays safe.""" + messages = [ + { + "role": "assistant", + "tool_calls": [{"id": "c1", "function": {"name": "grep", "arguments": "not json"}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": "ok"}, + {"role": "tool", "tool_call_id": "missing", "content": "orphan"}, + ] + + results = learner.extract_tool_results_from_openai_messages(messages) + assert results[0]["tool_name"] == "grep" + assert results[0]["input"] == {} + assert results[1]["tool_name"] == "unknown" + assert results[1]["input"] == {} + + def test_extract_openai_tool_results_empty_without_tool_messages(self, learner: TrafficLearner): + assert ( + learner.extract_tool_results_from_openai_messages([{"role": "user", "content": "hi"}]) + == [] + ) + @pytest.mark.asyncio async def test_tool_history_bounded(self, learner: TrafficLearner): """Test that tool history stays within max_history."""