mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(memory): don't crash inline memory extraction on a non-object <memory> block (#2470)
## Description
`parse_response_with_memory` extracts an inline `<memory>...</memory>`
block from a model response and parses its JSON:
```python
try:
data = json.loads(memory_json)
memories = data.get("memories", [])
except json.JSONDecodeError as e:
logger.warning(f"Failed to parse memory JSON: {e}")
```
The block content is fully model-controlled. `json.loads` succeeds on
any valid JSON, including a non-object such as a bare array
(`<memory>["x"]</memory>`), a string, or a number. `data.get("memories",
[])` then raises `AttributeError: 'list' object has no attribute 'get'`,
which the `except json.JSONDecodeError` does not catch, so the inline
memory path crashes on output a model can realistically produce.
## Fix
Guard the parsed value: read `memories` only when the block is a JSON
object, and accept it only when it is a list (logging and ignoring
otherwise). Malformed JSON is still handled by the existing decode
guard, and a well-formed object is unchanged. This mirrors the
non-object hardening already applied to the batch JSONL path.
## 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/inline_extractor.py`: only read `memories` from a
dict-typed parsed block, and only when the field is a list; log and
ignore other shapes.
- `tests/test_memory_wrapper.py`: regression covering a non-object
memory block, a non-list `memories` field, malformed JSON, and a
well-formed block.
## Testing
- [x] 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
$ python -m pytest tests/test_memory_wrapper.py -q
6 passed
# with the fix reverted, the new test fails with
# AttributeError: 'list' object has no attribute 'get'
$ uvx ruff@0.15.17 check headroom/memory/inline_extractor.py tests/test_memory_wrapper.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/memory/inline_extractor.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, project venv (`uv sync --extra
proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv.
- Exact command / steps: called the real `parse_response_with_memory`
with a `<memory>["x"]</memory>` block, a `{"memories": "nope"}` block, a
malformed block, and a valid block; then reverted `inline_extractor.py`
and re-ran.
- Observed result: with the fix all four return cleanly (empty memories
for the bad shapes, the parsed list for the valid one) and the memory
block is still stripped from the content; with the fix reverted the
non-object block raises `AttributeError: 'list' object has no attribute
'get'`. Ran against the actual module.
- Not tested: a live end-to-end chat where a model emits a non-object
memory block.
## 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
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable
This commit is contained in:
parent
fc5c4e239c
commit
e00c6ff81c
2 changed files with 32 additions and 1 deletions
|
|
@ -126,9 +126,21 @@ def parse_response_with_memory(response_text: str) -> ParsedResponse:
|
|||
# Parse the JSON
|
||||
try:
|
||||
data = json.loads(memory_json)
|
||||
memories = data.get("memories", [])
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Failed to parse memory JSON: {e}")
|
||||
else:
|
||||
# The model controls the block content. A valid-JSON non-object
|
||||
# (e.g. a bare array) would make `.get` raise AttributeError, which
|
||||
# the JSONDecodeError guard above does not catch; likewise a
|
||||
# non-list `memories` would break downstream iteration.
|
||||
if isinstance(data, dict):
|
||||
extracted = data.get("memories", [])
|
||||
if isinstance(extracted, list):
|
||||
memories = extracted
|
||||
else:
|
||||
logger.warning("Memory JSON 'memories' field is not a list; ignoring")
|
||||
else:
|
||||
logger.warning("Memory JSON is not an object; ignoring")
|
||||
|
||||
return ParsedResponse(
|
||||
content=content,
|
||||
|
|
|
|||
|
|
@ -205,3 +205,22 @@ def test_memory_api_methods_delegate_to_underlying_memory() -> None:
|
|||
assert api.clear() == 2
|
||||
assert api.stats() == {"total": 2}
|
||||
assert fake_memory.last_clear == {"user_id": "alice"}
|
||||
|
||||
|
||||
def test_parse_response_with_memory_tolerates_non_object_memory_block() -> None:
|
||||
from headroom.memory.inline_extractor import parse_response_with_memory
|
||||
|
||||
# The model controls the <memory> block. A valid-JSON non-object (a bare
|
||||
# array) or a non-list `memories` field must not crash the parse: the
|
||||
# `json.loads` succeeds, so the JSONDecodeError guard does not apply, and
|
||||
# `.get`/iteration on the wrong type would otherwise raise.
|
||||
assert parse_response_with_memory('hi <memory>["x"]</memory> bye').memories == []
|
||||
assert parse_response_with_memory('<memory>{"memories": "nope"}</memory>').memories == []
|
||||
|
||||
# Malformed JSON is still handled, and a well-formed block still parses.
|
||||
assert parse_response_with_memory("<memory>{not json</memory>").memories == []
|
||||
parsed = parse_response_with_memory(
|
||||
'text <memory>{"memories": [{"content": "User likes Python"}]}</memory>'
|
||||
)
|
||||
assert parsed.memories == [{"content": "User likes Python"}]
|
||||
assert parsed.content == "text"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue