mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
f542b70413
## Description
`extract_memory_query_sources` (`headroom/proxy/memory_query_policy.py`)
builds the text used to
retrieve relevant memories. It captures `latest_user` **only** when a
user message's `content` is
a plain `str`:
```python
if role == "user":
if isinstance(content, list):
_append_anthropic_tool_results(content, tool_outputs=..., lookback_tools=...)
elif isinstance(content, str) and not latest_user:
latest_user = content
```
But the standard Anthropic `/v1/messages` shape (used by Claude Code)
sends the user turn as a
**list of content blocks** — `content=[{"type":"text","text":"help me
refactor auth"}]`. That
routes into `_append_anthropic_tool_results`, which extracts only
`type=="tool_result"` blocks
and **never reads the `type=="text"` blocks** — so the actual user
prompt is discarded.
Downstream (`handlers/anthropic.py` → `MemoryQuery.from_messages` →
`to_embedding_input`):
- On a **first turn** (no prior assistant/tool context) the embedding
input is `""`, and the
memory handler then returns `None` — **memory injection is silently
skipped entirely**.
- With history present, the query is assembled from stale assistant/tool
context **minus the
current question**, so retrieval targets the wrong text.
Closes: no issue filed — found while auditing the memory retrieval query
policy.
## Fix
In the list-content user branch, also collect the `text` blocks into
`latest_user` (guarded by
`if not latest_user` so the latest turn wins), alongside the existing
tool-result extraction.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `headroom/proxy/memory_query_policy.py`: capture Anthropic user `text`
blocks into `latest_user`.
- `tests/test_memory_query_policy.py`: add
`test_extract_sources_captures_anthropic_user_text_blocks` and
`test_extract_sources_captures_user_text_alongside_tool_result`.
## Testing
- [x] New regression tests added (`tests/test_memory_query_policy.py`)
- [x] Linting/formatting clean — run with the CI-pinned `ruff==0.15.17`
- [ ] Full `pytest` deferred to CI (local-OOM reason below).
```text
$ uvx ruff@0.15.17 check headroom/proxy/memory_query_policy.py tests/test_memory_query_policy.py
All checks passed!
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.10, headroom from this branch.
Importing `headroom` pulls in the torch/transformers stack and a full
`pytest` gets OOM-killed on this box, so I verified the extraction logic
with a dependency-free script and left the full pytest to CI.
- Exact command / steps: ran a text-block user turn (and a mixed
text+tool_result turn, a plain-string turn, and multiple user turns)
through the old and new logic.
- Observed result: the old logic drops the user text (empty query →
injection skipped); the new logic captures it, still gathers tool
output, and keeps the plain-string / latest-turn behavior:
```text
text-block user: OLD user_text='' NEW user_text='help me refactor auth'
MEMORY QUERY TEXT-BLOCK FIX VERIFIED (old drops user text; new captures it)
```
- Not tested: a full memory retrieval round-trip through the embedder
(needs the heavy stack). The fix is confined to
`extract_memory_query_sources` and the new tests drive it directly. Full
local `pytest` deferred to CI (OOM, per above).
## 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 — ran
lint + a standalone logic check; full pytest deferred to CI (local OOM,
disclosed above)
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
- Small, contained fix in the query-source extractor; no new
dependencies. The existing
`test_extract_sources_handles_anthropic_tool_result_without_user_text`
still passes (its list turn has no text block).
- @JerrettDavis tagging you — this silently disables memory injection
for the standard Claude Code request shape on a first turn, so it seemed
worth surfacing. Thanks!
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
"""Tests for pure memory query construction policy."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from headroom.proxy.memory_query_policy import (
|
|
extract_memory_query_sources,
|
|
render_embedding_input,
|
|
)
|
|
|
|
|
|
def test_render_embedding_input_orders_sources_for_embedding() -> None:
|
|
rendered = render_embedding_input(
|
|
user_text="latest user",
|
|
recent_tool_outputs=("tool output",),
|
|
recent_assistant_turns=("assistant context",),
|
|
)
|
|
|
|
assert rendered.index("assistant context") < rendered.index("tool output")
|
|
assert rendered.index("tool output") < rendered.index("latest user")
|
|
|
|
|
|
def test_extract_sources_uses_latest_user_and_recent_context_in_order() -> None:
|
|
messages = [
|
|
{"role": "user", "content": "first"},
|
|
{"role": "assistant", "content": "a1"},
|
|
{"role": "tool", "content": "t1"},
|
|
{"role": "assistant", "content": "a2"},
|
|
{"role": "tool", "content": "t2"},
|
|
{"role": "user", "content": "second"},
|
|
]
|
|
|
|
user_text, tool_outputs, assistant_turns = extract_memory_query_sources(
|
|
messages,
|
|
lookback_assistant=2,
|
|
lookback_tools=2,
|
|
)
|
|
|
|
assert user_text == "second"
|
|
assert tool_outputs == ("t1", "t2")
|
|
assert assistant_turns == ("a1", "a2")
|
|
|
|
|
|
def test_extract_sources_handles_anthropic_tool_result_without_user_text() -> None:
|
|
messages = [
|
|
{"role": "user", "content": "real user"},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "content": [{"type": "text", "text": "nested"}]}],
|
|
},
|
|
]
|
|
|
|
user_text, tool_outputs, assistant_turns = extract_memory_query_sources(messages)
|
|
|
|
assert user_text == "real user"
|
|
assert tool_outputs == ("nested",)
|
|
assert assistant_turns == ()
|
|
|
|
|
|
def test_extract_sources_captures_anthropic_user_text_blocks() -> None:
|
|
"""Anthropic user turns carry the prompt as text blocks (the standard Claude
|
|
Code shape). The user's question must be captured — not dropped — so memory
|
|
retrieval keys on it."""
|
|
messages = [
|
|
{"role": "user", "content": [{"type": "text", "text": "help me refactor auth"}]},
|
|
]
|
|
|
|
user_text, _tool_outputs, _assistant_turns = extract_memory_query_sources(messages)
|
|
|
|
assert user_text == "help me refactor auth"
|
|
|
|
|
|
def test_extract_sources_captures_user_text_alongside_tool_result() -> None:
|
|
"""A user turn mixing a tool_result and a text block yields both: the text as
|
|
the user query and the tool output as context."""
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "content": "exit 0"},
|
|
{"type": "text", "text": "did the tests pass?"},
|
|
],
|
|
},
|
|
]
|
|
|
|
user_text, tool_outputs, _assistant_turns = extract_memory_query_sources(messages)
|
|
|
|
assert user_text == "did the tests pass?"
|
|
assert tool_outputs == ("exit 0",)
|