akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy/openai): feed chat/completions traffic into the traffic learner (#2333)

Browse source 
## 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 <same files>
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 <mxjerrett@gmail.com>
This commit is contained in:
Abhay Singh 2026-07-18 22:39:01 +05:30 • committed by GitHub
parent e3c7964038
commit 6cdfd3f64d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 201 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

78
headroom/memory/traffic_learner.py
View file

@ -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

47
headroom/proxy/handlers/openai.py
View file

@ -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:

76
tests/test_memory/test_traffic_learner.py
View file

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