mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
2d1e96b85c
## Description Closes #2947 `entity_refs` is annotated `list[str]` everywhere, but nothing enforced that at runtime. `LocalBackend.save_memory`'s `entities` argument is filled straight from LLM-supplied `memory_save` tool input (`headroom/memory/system.py:575` into `memory_handler.py:1242`), so a caller can pass the typed `{"entity": ..., "entity_type": ...}` shape, which is the format `extracted_entities` expects, into it by mistake. Those dicts were then persisted verbatim into `entity_refs`, both in the `memories` table and in the duplicated copy the vector index keeps for post-filtering. Every later `search_memories` call does `set().update(memory.entity_refs)` while collecting entities for graph expansion. Hashing a dict raises `TypeError: unhashable type: 'dict'`, and because that happens inside the vector-result loop rather than per-item, **one** poisoned row aborted the **entire** search. The proxy's memory handler catches the exception and returns no memories, so recall went quietly dark rather than failing loudly, and the bad row kept re-appearing in top-k for related queries, so it stayed dark. The issue reporter hit this in production: 4 bad rows disabled memory search for a whole project for a day, with nothing visible to the end user beyond a swallowed warning in `proxy.log`. The same root cause has two more crash modes, both confirmed below: `AttributeError: 'dict' object has no attribute 'lower'` during graph linking on the save path, and the same error in the `entities` search filter (`ref.lower()`). The fix adds one helper and applies it at both ends of the data flow. Dicts are **unwrapped to their `entity` name** rather than dropped, so rows that are already corrupted keep contributing to graph expansion instead of silently losing their entities. ## 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 - **New helper `normalize_entity_refs()` in `headroom/memory/models.py`.** Coerces a raw entity-reference list into the `list[str]` it claims to be: strings pass through, dicts are unwrapped via their `entity` (or `name`) key, and anything with no recoverable name is dropped rather than stringified, since a ref like `"{'entity_type': 'project'}"` would only pollute the graph. Order is preserved and duplicate names are collapsed. - **Write path, to stop new corruption at the door.** `LocalBackend.save_memory` normalizes `entities` before it reaches `entity_refs` and graph linking. `LocalBackend.search_memories` normalizes the `entities` *filter* argument too, since it arrives from the same untrusted tool input (`memory_handler.py:1320`). - **Read path, to heal rows that were written before this fix.** Applied at the three deserialization boundaries, so no data migration is needed and corrupted rows normalize themselves the next time they are loaded: `Memory.from_dict` (`headroom/memory/models.py`), `SQLiteMemoryStore._row_to_memory` (`headroom/memory/adapters/sqlite.py`), and the vector indexes' own `entity_refs` copies used for post-filtering, `VectorMetadata.from_json` (`headroom/memory/adapters/sqlite_vector.py`) and `IndexedMemoryMetadata.from_dict` (`headroom/memory/adapters/hnsw.py`). - **Defensive normalization on emitted results.** `search_memories` and `text_search` normalize the refs they return as `related_entities`, so a backend that produces `Memory` objects by some path not covered above still cannot take a whole query down, and callers never receive a dict where they expect an entity name. **Note on scope versus the patch proposed in the issue.** The issue proposed normalizing in two places (`save_memory` plus the `set().update()` line). I widened it slightly because that pair leaves three related failures live: the `entities` filter still crashes on `ref.lower()`, `related_entities` still hands dicts back to the caller, and, most importantly, already-poisoned rows stay poisoned in storage. Normalizing at the deserialization boundaries fixes all three at once and is what makes existing corrupted databases recover on their own. **Behavior change worth flagging.** `entity_refs` is now de-duplicated (case-sensitively) on both save and load. Refs were already treated as a set for graph expansion, so this is semantically a no-op, but it is a visible difference if anything asserts on exact list contents. ## 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 New file `tests/test_memory/test_entity_ref_sanitization.py` adds 10 tests covering the helper, both write paths, all three deserialization boundaries, and the three crash modes. ### Test Output ```text $ python -m pytest tests/test_memory/test_entity_ref_sanitization.py -q .......... [100%] 10 passed, 17 warnings in 0.34s ``` Full memory suite, plus a before/after comparison of the failure set to prove no regressions: ```text $ python -m pytest tests/test_memory/ -q 13 failed, 576 passed, 3 skipped, 1072 warnings, 25 errors in 44.74s # the same run with the source changes stashed (baseline on upstream/main @ |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| conftest.py | ||
| test_budget.py | ||
| test_core_metadata_index_sync.py | ||
| test_core_operations.py | ||
| test_direct_mem0.py | ||
| test_easy.py | ||
| test_embedder_mps_serialization.py | ||
| test_embedder_thread_cap.py | ||
| test_entity_ref_sanitization.py | ||
| test_extraction.py | ||
| test_factory.py | ||
| test_factory_embedder_cache.py | ||
| test_factory_external.py | ||
| test_hierarchical.py | ||
| test_hnsw_batch_capacity.py | ||
| test_learn_flag.py | ||
| test_local_backend_search.py | ||
| test_mcp_server.py | ||
| test_qdrant_env.py | ||
| test_query_conditions.py | ||
| test_skip_helpers.py | ||
| test_supersession_repair.py | ||
| test_traffic_learner.py | ||
| test_writers.py | ||