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