fix(memory): sanitize entity_refs to prevent dict-shaped entries crashing search (#2951)
## 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 @ 941c25d3):
13 failed, 566 passed, 3 skipped, 1055 warnings, 25 errors in 44.92s
# diff of the failing/erroring test IDs, before versus after:
$ diff baseline.txt after.txt && echo "NO NEW FAILURES vs baseline"
NO NEW FAILURES vs baseline
```
576 passed equals the 566 baseline plus the 10 new tests. The 13
failures and 25 errors are pre-existing on `upstream/main` and unrelated
to this change: they are Windows-only temp-directory cleanup failures in
this local environment.
```text
E PermissionError: [WinError 32] The process cannot access the file because it is being
used by another process: 'C:\Users\...\Temp\tmp_qnqtami\test.db'
```
Adjacent suites that construct `Memory` objects:
```text
$ python -m pytest tests/test_memory_system.py tests/test_memory_eval.py tests/test_critical_gaps.py -q
158 passed, 1 skipped, 514 warnings in 20.58s
```
Lint and format on the changed files:
```text
$ ruff check headroom/memory tests/test_memory/test_entity_ref_sanitization.py
All checks passed!
$ ruff format --check headroom/memory tests/test_memory/test_entity_ref_sanitization.py
48 files already formatted
```
`mypy headroom --ignore-missing-imports --python-version 3.13` reports
12 errors, all pre-existing on `upstream/main` and all in files this PR
does not touch (`headroom/ccr/mcp_server.py`,
`headroom/memory/mcp_server.py`, `headroom/release_version.py`; they
come from a local MCP SDK version mismatch). Zero errors in any changed
file.
## Real Behavior Proof
- Environment: Windows 11, Python 3.13.11, pytest 9.1.1, ruff 0.15.17,
local `headroom._core` built. Branched from `upstream/main` at
`941c25d3`, the same branch point the issue reports.
- Exact command / steps: Ran a standalone script (not a mock-only test)
driving `LocalBackend.search_memories` and `LocalBackend.save_memory`
with `entity_refs=[{"entity": "Project X", "entity_type": "project"}]`,
first against unmodified `941c25d3` and then against this branch. Three
scenarios: vector search with graph expansion, search with an `entities`
filter, and a save carrying dict-shaped `entities`.
- Observed result: on unmodified `941c25d3` all three crashed, printing
`SEARCH: TypeError: unhashable type: 'dict'`, `FILTER: TypeError:
unhashable type: 'dict'`, and `SAVE: AttributeError: 'dict' object has
no attribute 'lower'`. With this branch applied all three succeed:
search returns both the poisoned and the clean memory with
`related_entities == ["Project X"]`, the filter matches the recovered
name, and the save persists `entity_refs == ["Project X"]`. Those three
scenarios are now the regression tests in
`test_entity_ref_sanitization.py`.
- Not tested: no live end-to-end run through the MCP `memory_save` tool
against a real LLM, and no test against a real pre-existing SQLite
database containing dict-shaped rows. The healing-on-load path is
covered at the deserialization functions (`Memory.from_dict`,
`VectorMetadata.from_json`, `IndexedMemoryMetadata.from_dict`) rather
than through an actual corrupted `.db` file. The non-local backends
(`mem0`, `direct_mem0`, `qdrant-neo4j`, `cognee`) were not exercised;
this PR only changes the local backend and the shared models and
adapters.
## 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
- [x] I did **not** edit `CHANGELOG.md`, it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)
## Additional Notes
Documentation was not changed because `normalize_entity_refs()` is an
internal helper and no public API or user-facing behavior changes;
`entity_refs` still behaves exactly as its existing `list[str]` contract
always documented.
Credit for the diagnosis, the root-cause analysis, and the original
repro goes to @apacheco-RT in #2947, who could not open a PR directly
because GitHub blocks Enterprise Managed User accounts from forking
outside their enterprise.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
Parideboy
@@ -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),
|
||||
|
||||
@@ -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 {},
|
||||
)
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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", {}),
|
||||
)
|
||||
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user