diff --git a/headroom/memory/adapters/hnsw.py b/headroom/memory/adapters/hnsw.py index 9d7f54d47..5f474bd5c 100644 --- a/headroom/memory/adapters/hnsw.py +++ b/headroom/memory/adapters/hnsw.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any import numpy as np -from ..models import Memory, ScopeLevel +from ..models import Memory, ScopeLevel, normalize_entity_refs from ..ports import VectorFilter, VectorSearchResult # hnswlib is optional - may not compile on all platforms @@ -139,7 +139,9 @@ class IndexedMemoryMetadata: valid_until=( datetime.fromisoformat(data["valid_until"]) if data.get("valid_until") else None ), - entity_refs=data.get("entity_refs", []), + # Normalized on load so rows written before #2947 was fixed heal + # themselves instead of crashing search. + entity_refs=normalize_entity_refs(data.get("entity_refs")), content=data["content"], created_at=datetime.fromisoformat(data["created_at"]), importance=data.get("importance", 0.5), diff --git a/headroom/memory/adapters/sqlite.py b/headroom/memory/adapters/sqlite.py index 3b779bd91..c5fcdf28c 100644 --- a/headroom/memory/adapters/sqlite.py +++ b/headroom/memory/adapters/sqlite.py @@ -16,7 +16,7 @@ from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING, Any -from ..models import Memory, ScopeLevel +from ..models import Memory, ScopeLevel, normalize_entity_refs from ..ports import MemoryFilter if TYPE_CHECKING: @@ -227,7 +227,11 @@ class SQLiteMemoryStore: last_accessed=datetime.fromisoformat(row["last_accessed"]) if row["last_accessed"] else None, - entity_refs=json.loads(row["entity_refs"]) if row["entity_refs"] else [], + # Normalized on load so rows written before #2947 was fixed heal + # themselves instead of crashing search. + entity_refs=normalize_entity_refs( + json.loads(row["entity_refs"]) if row["entity_refs"] else [] + ), embedding=self._deserialize_embedding(row["embedding"]), metadata=json.loads(row["metadata"]) if row["metadata"] else {}, ) diff --git a/headroom/memory/adapters/sqlite_vector.py b/headroom/memory/adapters/sqlite_vector.py index e67310320..71cd96a83 100644 --- a/headroom/memory/adapters/sqlite_vector.py +++ b/headroom/memory/adapters/sqlite_vector.py @@ -29,7 +29,7 @@ from typing import TYPE_CHECKING, Any, cast import numpy as np -from ..models import Memory, ScopeLevel +from ..models import Memory, ScopeLevel, normalize_entity_refs from ..ports import VectorFilter, VectorSearchResult if TYPE_CHECKING: @@ -135,7 +135,9 @@ class VectorMetadata: valid_until=( datetime.fromisoformat(d["valid_until"]) if d.get("valid_until") else None ), - entity_refs=d.get("entity_refs", []), + # Normalized on load so rows written before #2947 was fixed heal + # themselves instead of crashing search. + entity_refs=normalize_entity_refs(d.get("entity_refs")), content=d["content"], created_at=datetime.fromisoformat(d["created_at"]), importance=d.get("importance", 0.5), diff --git a/headroom/memory/backends/local.py b/headroom/memory/backends/local.py index 98e94b7e6..7010bc920 100644 --- a/headroom/memory/backends/local.py +++ b/headroom/memory/backends/local.py @@ -21,7 +21,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any from headroom.memory.adapters.graph_models import Entity, Relationship, Subgraph -from headroom.memory.models import Memory +from headroom.memory.models import Memory, normalize_entity_refs from headroom.memory.ports import MemorySearchResult from headroom.models.config import ML_MODEL_DEFAULTS @@ -284,8 +284,15 @@ class LocalBackend: # Determine if using pre-extraction mode has_pre_extraction = bool(facts or extracted_entities or extracted_relationships) - # Merge entity names from both simple and typed formats - all_entity_names: list[str] = list(entities) if entities else [] + # Merge entity names from both simple and typed formats. + # + # `entities` is typed list[str], but it is populated straight from + # LLM-supplied memory_save tool arguments, and callers do sometimes + # pass the typed {"entity": ..., "entity_type": ...} shape here (that + # is what extracted_entities is for). Normalizing keeps those dicts out + # of entity_refs, where they used to crash every later search that + # retrieved the row -- see issue #2947. + all_entity_names: list[str] = normalize_entity_refs(entities) entity_types: dict[str, str] = {} if extracted_entities: @@ -448,13 +455,17 @@ class LocalBackend: continue seen_memory_ids.add(vr.memory.id) - all_entity_refs.update(vr.memory.entity_refs) + # Defense in depth: the storage adapters normalize entity_refs on + # load, but a backend that hands us Memory objects some other way + # must not be able to abort the whole search with one bad row. + entity_refs = normalize_entity_refs(vr.memory.entity_refs) + all_entity_refs.update(entity_refs) results.append( MemorySearchResult( memory=vr.memory, score=vr.similarity, - related_entities=list(vr.memory.entity_refs), + related_entities=entity_refs, related_memories=[], ) ) @@ -503,15 +514,17 @@ class LocalBackend: MemorySearchResult( memory=memory, score=0.5, # Default score for graph-expanded results - related_entities=list(memory.entity_refs), + related_entities=normalize_entity_refs(memory.entity_refs), related_memories=[], ) ) seen_memory_ids.add(mem_id) - # Filter by specified entities if provided + # Filter by specified entities if provided. Like the save path, this + # argument arrives from LLM-supplied tool input, so it gets the same + # normalization rather than trusting its list[str] annotation. if entities: - entities_lower = {e.lower() for e in entities} + entities_lower = {e.lower() for e in normalize_entity_refs(entities)} results = [ r for r in results @@ -826,7 +839,7 @@ class LocalBackend: MemorySearchResult( memory=tr.memory, score=tr.score, - related_entities=list(tr.memory.entity_refs), + related_entities=normalize_entity_refs(tr.memory.entity_refs), related_memories=[], ) for tr in text_results diff --git a/headroom/memory/models.py b/headroom/memory/models.py index 6e1293c34..2ecdcabe4 100644 --- a/headroom/memory/models.py +++ b/headroom/memory/models.py @@ -14,6 +14,46 @@ except ImportError: np = None # type: ignore[assignment] +def normalize_entity_refs(values: Any) -> list[str]: + """Coerce a raw entity-reference list into the plain ``list[str]`` it claims to be. + + ``entity_refs`` (and the ``entities`` argument that feeds it) is typed + ``list[str]``, but nothing enforced that at runtime, so callers have + persisted the typed ``{"entity": ..., "entity_type": ...}`` shape -- the + format ``extracted_entities`` expects -- into it by mistake. Those dicts + then break every consumer that treats a ref as a string: ``set().update()`` + raises ``TypeError: unhashable type: 'dict'`` and ``ref.lower()`` raises + ``AttributeError``, which took down whole memory searches rather than the + one bad row (see issue #2947). + + Dicts are unwrapped to their ``entity`` name so no information is lost; + 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. + """ + if not values: + return [] + + normalized: list[str] = [] + seen: set[str] = set() + + for value in values: + if isinstance(value, str): + name = value + elif isinstance(value, dict): + # The extracted_entities shape, mistakenly used as a plain name. + candidate = value.get("entity") or value.get("name") + name = candidate if isinstance(candidate, str) else "" + else: + name = "" + + if name and name not in seen: + seen.add(name) + normalized.append(name) + + return normalized + + class ScopeLevel(Enum): """Memory scope hierarchy levels.""" @@ -132,7 +172,9 @@ class Memory: last_accessed=datetime.fromisoformat(data["last_accessed"]) if data.get("last_accessed") else None, - entity_refs=data.get("entity_refs", []), + # Normalized on load so rows written before #2947 was fixed heal + # themselves instead of crashing search. + entity_refs=normalize_entity_refs(data.get("entity_refs")), embedding=embedding, metadata=data.get("metadata", {}), ) diff --git a/tests/test_memory/test_entity_ref_sanitization.py b/tests/test_memory/test_entity_ref_sanitization.py new file mode 100644 index 000000000..70bbd6592 --- /dev/null +++ b/tests/test_memory/test_entity_ref_sanitization.py @@ -0,0 +1,302 @@ +"""Regression tests for entity_refs type safety. + +`entity_refs` is typed `list[str]` everywhere, but nothing enforced that at +runtime. A caller that mistakenly passed the typed +`{"entity": ..., "entity_type": ...}` shape (the format `extracted_entities` +expects) into the plain `entities` field of `save_memory` got those dicts +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 (not guarded per-item) it aborted the *entire* search for +any query whose top-k included one poisoned row. The proxy's memory handler +swallows the exception and returns no memories, so recall went quietly dark +rather than failing loudly. + +See https://github.com/headroomlabs-ai/headroom/issues/2947. +""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from headroom.memory.adapters.hnsw import IndexedMemoryMetadata +from headroom.memory.adapters.sqlite_vector import VectorMetadata +from headroom.memory.backends.local import LocalBackend +from headroom.memory.models import Memory, normalize_entity_refs + +# The malformed shape that started all of this: the extracted_entities format +# passed into a field that expects plain names. +DICT_REF = {"entity": "Project X", "entity_type": "project"} + + +# ============================================================================= +# The helper itself +# ============================================================================= + + +def test_normalize_entity_refs_unwraps_dicts_and_drops_junk() -> None: + """Dicts are unwrapped to their name; anything unusable is dropped.""" + assert normalize_entity_refs(["Alice", DICT_REF]) == ["Alice", "Project X"] + + # Nothing usable in these: no name to recover, so they are dropped rather + # than stringified into garbage entity names like "{'foo': 'bar'}". + assert normalize_entity_refs([{"entity_type": "project"}, {}, None, 42, ""]) == [] + + # Common no-op cases stay untouched. + assert normalize_entity_refs(["Alice", "Bob"]) == ["Alice", "Bob"] + assert normalize_entity_refs(None) == [] + assert normalize_entity_refs([]) == [] + + +def test_normalize_entity_refs_preserves_order_and_deduplicates() -> None: + """A name already present is not appended twice, and order is stable.""" + assert normalize_entity_refs(["Alice", DICT_REF, "Alice", "Project X"]) == [ + "Alice", + "Project X", + ] + + +# ============================================================================= +# Write path: stop new corruption at the door +# ============================================================================= + + +@pytest.mark.asyncio +async def test_save_memory_sanitizes_dict_shaped_entities_param() -> None: + """`entities` items that are dicts get coerced to plain names before storage.""" + backend = LocalBackend() + backend._initialized = True + + saved: list[Memory] = [] + + async def fake_add(**kwargs: object) -> Memory: + memory = Memory( + id="new-memory", + content=str(kwargs["content"]), + user_id=str(kwargs["user_id"]), + entity_refs=list(kwargs["entity_refs"]), # type: ignore[arg-type] + ) + saved.append(memory) + return memory + + backend._hierarchical_memory = SimpleNamespace(add=AsyncMock(side_effect=fake_add)) + backend._graph = SimpleNamespace( + get_entity_by_name=AsyncMock(return_value=None), + add_entity=AsyncMock(return_value=SimpleNamespace(id="entity-id")), + add_relationship=AsyncMock(), + ) + + # Without the fix this raises AttributeError: 'dict' object has no + # attribute 'lower' during graph linking. + await backend.save_memory( + content="Alice manages Project X", + user_id="alice", + entities=[DICT_REF], # type: ignore[list-item] + ) + + assert saved[0].entity_refs == ["Project X"] + + +@pytest.mark.asyncio +async def test_save_memory_merges_dict_entities_with_extracted_entities() -> None: + """A name arriving through both `entities` and `extracted_entities` is stored once.""" + backend = LocalBackend() + backend._initialized = True + + saved: list[Memory] = [] + + async def fake_add(**kwargs: object) -> Memory: + memory = Memory( + id="new-memory", + content=str(kwargs["content"]), + user_id=str(kwargs["user_id"]), + entity_refs=list(kwargs["entity_refs"]), # type: ignore[arg-type] + ) + saved.append(memory) + return memory + + backend._hierarchical_memory = SimpleNamespace(add=AsyncMock(side_effect=fake_add)) + backend._graph = SimpleNamespace( + get_entity_by_name=AsyncMock(return_value=None), + add_entity=AsyncMock(return_value=SimpleNamespace(id="entity-id")), + add_relationship=AsyncMock(), + ) + + await backend.save_memory( + content="Alice manages Project X", + user_id="alice", + entities=[DICT_REF], # type: ignore[list-item] + extracted_entities=[{"entity": "Project X", "entity_type": "project"}], + ) + + assert saved[0].entity_refs == ["Project X"] + + +# ============================================================================= +# Read path: heal rows that were already written before the fix +# ============================================================================= + + +def test_memory_from_dict_heals_stored_dict_refs() -> None: + """Rows persisted before the fix load as plain names instead of dicts.""" + now = datetime.now(timezone.utc).isoformat() + memory = Memory.from_dict( + { + "id": "poisoned-memory", + "content": "Alice manages Project X", + "user_id": "alice", + "created_at": now, + "valid_from": now, + "importance": 0.5, + "entity_refs": [DICT_REF, "Alice"], + } + ) + + assert memory.entity_refs == ["Project X", "Alice"] + + +def test_vector_metadata_from_json_heals_stored_dict_refs() -> None: + """The vector index keeps its own copy of entity_refs; heal that one too.""" + now = datetime.now(timezone.utc).isoformat() + metadata = VectorMetadata.from_json( + json.dumps( + { + "memory_id": "poisoned-memory", + "user_id": "alice", + "session_id": None, + "agent_id": None, + "valid_until": None, + "entity_refs": [DICT_REF], + "content": "Alice manages Project X", + "created_at": now, + "importance": 0.5, + "metadata": {}, + } + ) + ) + + assert metadata.entity_refs == ["Project X"] + assert metadata.to_memory().entity_refs == ["Project X"] + + +def test_indexed_memory_metadata_from_dict_heals_stored_dict_refs() -> None: + """Same for the HNSW index's metadata copy.""" + now = datetime.now(timezone.utc).isoformat() + metadata = IndexedMemoryMetadata.from_dict( + { + "memory_id": "poisoned-memory", + "user_id": "alice", + "session_id": None, + "agent_id": None, + "valid_until": None, + "entity_refs": [DICT_REF], + "content": "Alice manages Project X", + "created_at": now, + "importance": 0.5, + "metadata": {}, + } + ) + + assert metadata.entity_refs == ["Project X"] + + +# ============================================================================= +# Search: a single bad row must not take the whole query down +# ============================================================================= + + +def _backend_with_results(memories: list[Memory]) -> LocalBackend: + backend = LocalBackend() + backend._initialized = True + backend._hierarchical_memory = SimpleNamespace( + search=AsyncMock(return_value=[SimpleNamespace(memory=m, similarity=0.9) for m in memories]) + ) + backend._graph = SimpleNamespace( + get_entity_by_name=AsyncMock(return_value=None), + query_subgraph=AsyncMock(return_value=SimpleNamespace(entities=[], relationships=[])), + ) + return backend + + +@pytest.mark.asyncio +async def test_search_memories_tolerates_dict_shaped_entity_refs() -> None: + """A single legacy/corrupted row with dict entity_refs must not crash search.""" + poisoned = Memory( + id="poisoned-memory", + content="Alice manages Project X", + user_id="alice", + entity_refs=[DICT_REF], # type: ignore[list-item] + ) + clean = Memory( + id="clean-memory", + content="Bob manages Project Y", + user_id="alice", + entity_refs=["Project Y"], + ) + backend = _backend_with_results([poisoned, clean]) + + # Without the fix this raises TypeError: unhashable type: 'dict'. + results = await backend.search_memories("Alice's work", "alice", include_related=True) + + assert [r.memory.id for r in results] == ["poisoned-memory", "clean-memory"] + # The recovered name is still usable for graph expansion and is reported + # back to the caller as a plain string, not a dict. + assert results[0].related_entities == ["Project X"] + backend._graph.get_entity_by_name.assert_awaited() + + +@pytest.mark.asyncio +async def test_search_memories_entity_filter_matches_healed_refs() -> None: + """The `entities` filter lowercases each ref, which dicts also break. + + On unfixed code this never gets that far -- the unconditional + `set().update()` above raises first -- but once refs are strings again the + filter has to actually match the recovered name. + """ + poisoned = Memory( + id="poisoned-memory", + content="Alice manages Project X", + user_id="alice", + entity_refs=[DICT_REF], # type: ignore[list-item] + ) + backend = _backend_with_results([poisoned]) + + results = await backend.search_memories( + "Alice's work", + "alice", + include_related=False, + entities=["project x"], + ) + + assert [r.memory.id for r in results] == ["poisoned-memory"] + + +@pytest.mark.asyncio +async def test_search_memories_tolerates_dict_shaped_entities_filter() -> None: + """The filter argument comes from LLM tool input too, so it can be malformed.""" + clean = Memory( + id="clean-memory", + content="Alice manages Project X", + user_id="alice", + entity_refs=["Project X"], + ) + backend = _backend_with_results([clean]) + + # Without normalization this raises AttributeError: 'dict' object has no + # attribute 'lower' while building the filter set. + results = await backend.search_memories( + "Alice's work", + "alice", + include_related=False, + entities=[DICT_REF], # type: ignore[list-item] + ) + + assert [r.memory.id for r in results] == ["clean-memory"]