Files
turnstone/tests/test_structured_memory.py
Patrick Buckley 2fd0c29a92 fix(memory): query-aware candidate selection + OR-of-terms search (#468)
* fix(memory): query-aware candidate selection + OR-of-terms search

The system-message memory composition path used a recency-ordered
candidate set (`_list_visible_memories(limit=fetch_limit)`).  On
deployments with more than `fetch_limit` (default 50) visible
memories, BM25 only ever ranked the 50 most-recently-touched memories
— a relevant memory written months ago was silently invisible
regardless of how well it matched the recent context.  Multi-word
search at the SQL layer used AND-of-terms, killing recall on any
multi-word query without an exact field overlap.

## Functional changes

- `_init_system_messages` (`turnstone/core/session.py`): extract
  recent context first, then `_search_visible_memories(context)` to
  pull query-aware candidates.  Search hits below `fetch_limit` union
  with the recency list (deduped by memory_id) so the BM25 candidate
  pool is always a SUPERSET of the prior recency-only pool — even on
  noisy queries where the cap fills with stopwords, the recency-50
  the original bug surfaced still reaches BM25.  Empty context skips
  search entirely.  Candidate-selection logic extracted into
  `_select_memory_candidates`.

- `search_structured_memories` (PostgreSQL + SQLite): per-term
  clauses join with OR instead of AND.  A row matches if ANY term
  matches ANY of name/description/content.  Downstream BM25 narrows
  back down by relevance.

## Perf hardening

- Collapse the 1-3 fanned scope queries into a single SQL.  New
  backend methods `list_visible_structured_memories` /
  `search_visible_structured_memories` union the visibility scopes
  into one WHERE OR-group, so a composition rebuild now hits the DB
  at most twice (search + recency) instead of up to six times.

- Cap and normalize search terms.  Composition can hand a multi-KB
  pasted message to ILIKE-based search; without a cap, every distinct
  token would emit one unindexable predicate per scope-fanned query.
  `normalize_search_terms` (`storage/_utils.py`) de-dupes
  case-insensitively, drops <2-char tokens, and hard-caps at 16.

- Per-turn search cache.  `_init_system_messages` fires from many
  call sites within one turn (state transitions, MCP refresh, tool
  results) and the recent-context query is identical across them.
  Session-instance cache keyed by (query, mem_type, limit) absorbs
  the duplicates; invalidated in `_append_user_turn` and after
  memory save/delete tool actions.

- Stable secondary sort by `memory_id`.  `updated` is second-precision
  and `touch_structured_memories` can land a batch on identical
  timestamps; without a tie-breaker SQL returns rows in
  implementation-defined order, BM25 input shuffles, and the
  LLM-side prompt cache misses across calls.  All four backend ORDER
  BYs now break ties on `memory_id ASC`.

## Quality cleanups

- Coalesce `memory.search.term_count` + `memory.search.zero_results`
  into a single `memory.search` log carrying both `term_count` and
  `result_count`.
- New `memory.composition` log: source / candidates / injected.
- Promote a shared `make_chat_session` factory to `tests/_helpers.py`.
- Rename SQL builder local `extra` -> `scope_filters` for clarity.
- Add docstrings on `search_structured_memories` so the AND->OR flip
  survives future readers.

## Tests

Adds 20 tests across `tests/test_structured_memory.py`,
`tests/test_structured_memory_storage.py`, and
`tests/test_memory_relevance.py`: recency-ceiling regression,
empty-query fallback, sparse-match union, recency-preserved-when-
search-returns-noise (locks in the pool-superset invariant),
OR-of-terms on both backends, scope filtering preserved,
search-facade multi-word behavior, term-cap normalization, the new
visible-scope helpers (list + search + empty-scopes guard),
coord-scope composition isolation, end-to-end
`memory(action='search')` tool execution, per-turn cache hit +
invalidation, and stable ordering under tied `updated` timestamps.

Memory test sweep: 102/102.  Broader regression
(session, storage, coordinator, load_skill): 411/411.

* fix(memory): address Copilot review on PR #468

Three follow-ups from Copilot's inline review:

1. SUPERSET invariant violation (Copilot, session.py:5510).
   `(search_hits + extra)[:fetch_limit]` capped the union back down to
   fetch_limit, evicting the recency tail when search added distinct
   hits.  Recency tail is exactly where ancient-but-recently-touched
   memories live — the recall this PR is supposed to improve — so
   tail eviction recreated the bug for the narrow case where a query
   term fell off the 16-cap and the matching memory sat in
   recency[40-49].  Drop the cap; both halves are already SQL-capped
   at fetch_limit, so the union is at most 2 × fetch_limit (~100 with
   defaults).  BM25 over 100 candidates in pure Python is sub-ms;
   irrelevant recency fillers get score=0 and don't pollute ranking.
   Updates the docstring to actually be honest about the invariant.
   Adds `test_recency_tail_preserved_when_search_adds_distinct_hits`
   that locks the behavior in: 5 search hits + 10 recency = 15-item
   pool, every recency item present, source="union".

2. Unbounded `query.split()` in normalize_search_terms (Copilot,
   _utils.py:74).  `str.split()` allocates the full token list before
   the cap-after-16 break, so a 100KB pasted query did MB of throwaway
   work even though only 16 tokens entered SQL.  Switch to
   `re.finditer(r'\S+', query)` — streaming iterator, stops scanning
   at the first 16 normalized terms regardless of input size.

3. Misleading + unbounded log term_count (Copilot, session.py:8571).
   `len(item["query"].split())` had two problems: same unbounded
   split as #2, and the value reported the raw input token count
   rather than the normalized term count that actually hit the SQL
   WHERE clause — misleading metric for an operator trying to
   understand storage-side behavior.  Switch to
   `len(normalize_search_terms(item["query"]))` — accurate count, and
   bounded for free via #2.

Refuted: github-code-quality flagged `...` bodies in the new Protocol
methods as "statement has no effect."  False positive — `...` is the
canonical Protocol body convention, used 213 other times in the same
file.

Memory test sweep: 103/103.  Broader regression: 411/411.
2026-05-03 13:40:30 -07:00

479 lines
19 KiB
Python

"""Tests for turnstone.core.memory — structured memory facade functions."""
from turnstone.core.memory import (
count_structured_memories,
delete_structured_memory,
get_structured_memory_by_name,
list_structured_memories,
normalize_key,
save_structured_memory,
search_structured_memories,
)
class TestSaveStructuredMemory:
def test_save_new(self, tmp_db):
mid, old = save_structured_memory("test_key", "hello world")
assert mid != ""
assert old is None
def test_save_upsert(self, tmp_db):
save_structured_memory("test_key", "first")
mid, old = save_structured_memory("test_key", "second")
assert old == "first"
assert mid != ""
def test_save_normalizes_key(self, tmp_db):
save_structured_memory("My-Key", "value")
mems = list_structured_memories()
assert any(m["name"] == "my_key" for m in mems)
def test_save_with_type_and_scope(self, tmp_db):
save_structured_memory("k", "v", mem_type="user", scope="workstream", scope_id="ws1")
mems = list_structured_memories(scope="workstream", scope_id="ws1")
assert len(mems) == 1
assert mems[0]["type"] == "user"
class TestDeleteStructuredMemory:
def test_delete_existing(self, tmp_db):
save_structured_memory("mykey", "val")
assert delete_structured_memory("mykey")
def test_delete_nonexistent(self, tmp_db):
assert not delete_structured_memory("nope")
def test_delete_normalizes_key(self, tmp_db):
save_structured_memory("my_key", "val")
assert delete_structured_memory("My-Key")
class TestListStructuredMemories:
def test_list_empty(self, tmp_db):
assert list_structured_memories() == []
def test_list_returns_saved(self, tmp_db):
save_structured_memory("a", "alpha")
save_structured_memory("b", "beta")
mems = list_structured_memories()
assert len(mems) == 2
class TestSearchStructuredMemories:
def test_search_finds_match(self, tmp_db):
save_structured_memory("db_host", "localhost", description="database hostname")
save_structured_memory("api_url", "http://example.com")
results = search_structured_memories("database")
assert len(results) >= 1
assert any(r["name"] == "db_host" for r in results)
def test_multiword_or_matches_partial(self, tmp_db):
"""OR-of-terms: memory matching only 1 of 3 query terms is returned."""
save_structured_memory("postgres_config", "host=localhost port=5432")
save_structured_memory("redis_config", "host=redis port=6379")
save_structured_memory("unrelated", "nothing relevant here")
# "postgres missing_word_a missing_word_b": only postgres_config matches "postgres"
results = search_structured_memories("postgres missing_word_a missing_word_b")
names = {r["name"] for r in results}
assert "postgres_config" in names
assert "unrelated" not in names
def test_multiword_or_multiple_partial_matches(self, tmp_db):
"""Multiple memories each matching different terms are all returned."""
save_structured_memory("key_alpha", "alpha content here")
save_structured_memory("key_beta", "beta content here")
save_structured_memory("key_other", "completely different")
results = search_structured_memories("alpha beta")
names = {r["name"] for r in results}
assert "key_alpha" in names
assert "key_beta" in names
assert "key_other" not in names
def test_search_scope_filtering_preserved(self, tmp_db):
"""Search with scope filter only returns memories in that scope."""
save_structured_memory("ws1_fact", "alpha info", scope="workstream", scope_id="ws1")
save_structured_memory("ws2_fact", "alpha info", scope="workstream", scope_id="ws2")
save_structured_memory("global_fact", "alpha info", scope="global")
results = search_structured_memories("alpha", scope="workstream", scope_id="ws1")
names = {r["name"] for r in results}
assert "ws1_fact" in names
assert "ws2_fact" not in names
assert "global_fact" not in names
class TestGetStructuredMemoryByName:
def test_get_existing(self, tmp_db):
save_structured_memory("my_mem", "full content here that is quite long")
mem = get_structured_memory_by_name("my_mem", "global", "")
assert mem is not None
assert mem["content"] == "full content here that is quite long"
assert mem["name"] == "my_mem"
def test_get_nonexistent(self, tmp_db):
assert get_structured_memory_by_name("nope", "global", "") is None
def test_get_wrong_scope(self, tmp_db):
save_structured_memory("ws_mem", "data", scope="workstream", scope_id="ws1")
assert get_structured_memory_by_name("ws_mem", "global", "") is None
assert get_structured_memory_by_name("ws_mem", "workstream", "ws1") is not None
def test_get_normalizes_key(self, tmp_db):
save_structured_memory("My-Key", "value")
mem = get_structured_memory_by_name("My-Key", "global", "")
assert mem is not None
assert mem["name"] == "my_key"
class TestCountStructuredMemories:
def test_count_zero(self, tmp_db):
assert count_structured_memories() == 0
def test_count_after_save(self, tmp_db):
save_structured_memory("a", "1")
save_structured_memory("b", "2")
assert count_structured_memories() == 2
class TestNormalizeKey:
def test_basic(self):
assert normalize_key("My-Key Name") == "my_key_name"
class TestScopeIsolation:
"""Verify that list/search without scope only returns visible memories.
Reproduces the cross-workstream leak: unscoped list/search must not
return workstream-scoped memories from other workstreams or
user-scoped memories from other users.
"""
def _seed(self):
"""Create memories across multiple scopes."""
save_structured_memory("global_note", "visible to all", scope="global")
save_structured_memory("ws1_note", "belongs to ws1", scope="workstream", scope_id="ws1")
save_structured_memory("ws2_note", "belongs to ws2", scope="workstream", scope_id="ws2")
save_structured_memory("u1_note", "belongs to user1", scope="user", scope_id="u1")
save_structured_memory("u2_note", "belongs to user2", scope="user", scope_id="u2")
@staticmethod
def _list_visible(ws_id: str, user_id: str, mem_type: str = "", limit: int = 50):
"""Replicate the scope-filtered list logic from ChatSession."""
global_mems = list_structured_memories(mem_type=mem_type, scope="global", limit=limit)
ws_mems = list_structured_memories(
mem_type=mem_type, scope="workstream", scope_id=ws_id, limit=limit
)
user_mems = (
list_structured_memories(mem_type=mem_type, scope="user", scope_id=user_id, limit=limit)
if user_id
else []
)
combined = global_mems + ws_mems + user_mems
combined.sort(key=lambda m: m.get("updated", ""), reverse=True)
return combined[:limit]
@staticmethod
def _search_visible(query: str, ws_id: str, user_id: str, mem_type: str = "", limit: int = 20):
"""Replicate the scope-filtered search logic from ChatSession."""
global_mems = search_structured_memories(
query, mem_type=mem_type, scope="global", limit=limit
)
ws_mems = search_structured_memories(
query, mem_type=mem_type, scope="workstream", scope_id=ws_id, limit=limit
)
user_mems = (
search_structured_memories(
query, mem_type=mem_type, scope="user", scope_id=user_id, limit=limit
)
if user_id
else []
)
combined = global_mems + ws_mems + user_mems
combined.sort(key=lambda m: m.get("updated", ""), reverse=True)
return combined[:limit]
def test_unscoped_list_returns_all_scopes(self, tmp_db):
"""Demonstrate the leak: unscoped list returns everything."""
self._seed()
all_mems = list_structured_memories()
assert len(all_mems) == 5 # no scope filter → all memories
def test_visible_list_excludes_other_workstreams(self, tmp_db):
"""Scope-filtered list for ws1/u1 excludes ws2 and u2 memories."""
self._seed()
visible = self._list_visible("ws1", "u1")
names = {m["name"] for m in visible}
assert "global_note" in names
assert "ws1_note" in names
assert "u1_note" in names
assert "ws2_note" not in names
assert "u2_note" not in names
def test_visible_list_no_user(self, tmp_db):
"""Scope-filtered list with no user_id excludes all user memories."""
self._seed()
visible = self._list_visible("ws1", "")
names = {m["name"] for m in visible}
assert "global_note" in names
assert "ws1_note" in names
assert "u1_note" not in names
assert "u2_note" not in names
def test_visible_search_no_user(self, tmp_db):
"""Scope-filtered search with no user_id excludes all user memories."""
self._seed()
visible = self._search_visible("belongs", "ws1", "")
names = {m["name"] for m in visible}
assert "ws1_note" in names
assert "u1_note" not in names
assert "u2_note" not in names
def test_visible_search_excludes_other_workstreams(self, tmp_db):
"""Scope-filtered search for ws1/u1 excludes ws2 and u2 memories."""
self._seed()
visible = self._search_visible("belongs", "ws1", "u1")
names = {m["name"] for m in visible}
assert "ws1_note" in names
assert "u1_note" in names
assert "ws2_note" not in names
assert "u2_note" not in names
def test_explicit_scope_still_works(self, tmp_db):
"""Explicit scope filter continues to work as before."""
self._seed()
ws2_only = list_structured_memories(scope="workstream", scope_id="ws2")
assert len(ws2_only) == 1
assert ws2_only[0]["name"] == "ws2_note"
class TestSanitizeErrorText:
"""Verify error-text sanitisation strips credentials and caps length.
Pairs with the ``persist_last_error`` writer — every persisted
string flows through ``sanitize_error_text`` so a misconfigured
provider URL or a quoted response body can't park credentials in
storage where the coordinator LLM later inhales them via the
inspect/wait surface.
Sanitisation delegates to
:func:`turnstone.core.output_guard.redact_credentials` so the
pattern set is the same one audit logs and the post-tool guard
use. The tests below assert the *behaviour* (the secret is gone)
rather than the exact replacement marker — output_guard owns the
marker format and the regex catalog, and pinning the marker here
would force two-place edits whenever output_guard adds a new
redaction label.
"""
def test_strips_url_userinfo(self):
from turnstone.core.memory import sanitize_error_text
# Misconfigured OPENAI_BASE_URL → httpx ConnectError carries
# the userinfo verbatim in str(exc).
msg = "ConnectError: connection failed to https://user:hunter2@api.example.com/v1/chat"
out = sanitize_error_text(msg)
# The password is gone but the host (useful for triage) stays.
assert "hunter2" not in out
assert "api.example.com" in out
def test_strips_url_userinfo_http_too(self):
from turnstone.core.memory import sanitize_error_text
msg = "RequestError on http://admin:s3cret@internal.host/path"
out = sanitize_error_text(msg)
assert "s3cret" not in out
assert "internal.host" in out
def test_strips_db_connection_string(self):
"""Output_guard already covered DB connection-strings; assert
the delegation surfaces that coverage so a leaked
``DATABASE_URL`` echoed in an error doesn't slip through."""
from turnstone.core.memory import sanitize_error_text
msg = "OperationalError: postgresql://app:topsecret@db.host/main"
out = sanitize_error_text(msg)
assert "topsecret" not in out
def test_redacts_openai_keys(self):
from turnstone.core.memory import sanitize_error_text
msg = (
"AuthenticationError: invalid api key sk-proj-AbCdEfGhIjKlMnOpQrStUv "
"(echoed from request body)"
)
out = sanitize_error_text(msg)
assert "sk-proj-AbCdEfGhIjKlMnOpQrStUv" not in out
def test_redacts_bearer_tokens(self):
from turnstone.core.memory import sanitize_error_text
msg = "401 Unauthorized - Bearer eyJabcDEFghiJKLmnoPQRstuVWX rejected"
out = sanitize_error_text(msg)
assert "eyJabcDEFghiJKLmnoPQRstuVWX" not in out
def test_redacts_github_tokens(self):
from turnstone.core.memory import sanitize_error_text
# The output_guard ghp pattern requires exactly 36 chars, so
# use a realistic-shaped token.
msg = "git push failed: ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghij not authorized"
out = sanitize_error_text(msg)
assert "ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghij" not in out
def test_redacts_aws_access_keys(self):
from turnstone.core.memory import sanitize_error_text
msg = "S3 error: signature mismatch for AKIAIOSFODNN7EXAMPLE"
out = sanitize_error_text(msg)
assert "AKIAIOSFODNN7EXAMPLE" not in out
def test_caps_length(self):
from turnstone.core.memory import LAST_ERROR_MAX_LEN, sanitize_error_text
msg = "X" * (LAST_ERROR_MAX_LEN * 2)
out = sanitize_error_text(msg)
assert len(out) <= LAST_ERROR_MAX_LEN
# Truncation marker preserved.
assert out.endswith("...")
def test_passes_through_clean_text(self):
from turnstone.core.memory import sanitize_error_text
msg = "TimeoutError: provider did not respond within 60s"
assert sanitize_error_text(msg) == msg
def test_handles_empty(self):
from turnstone.core.memory import sanitize_error_text
assert sanitize_error_text("") == ""
class TestPersistLastError:
"""Direct unit tests for the writer-side helper.
The reader-side tests in test_coordinator_client.py write to storage
via the raw backend, so the writer's contract — sanitize, no-op on
empty inputs, swallow storage failures, use the published constant
key — is unexercised without these.
"""
def test_round_trip_uses_constant_key(self, tmp_db):
from turnstone.core.memory import (
LAST_ERROR_CONFIG_KEY,
load_last_error,
persist_last_error,
register_workstream,
)
# Pre-register a workstream so save_workstream_config has somewhere
# to land — workstream_config rows reference the workstreams table.
register_workstream("ws-1", user_id="u1")
persist_last_error("ws-1", "TimeoutError: provider stalled")
assert load_last_error("ws-1") == "TimeoutError: provider stalled"
# The persisted row uses the published constant key — pinning
# this catches future drift between the writer and the
# coordinator_client.py readers that import the same constant.
from turnstone.core.memory import load_workstream_config
cfg = load_workstream_config("ws-1")
assert LAST_ERROR_CONFIG_KEY in cfg
def test_sanitises_before_persist(self, tmp_db):
from turnstone.core.memory import (
load_last_error,
persist_last_error,
register_workstream,
)
register_workstream("ws-1", user_id="u1")
persist_last_error("ws-1", "ConnectError: https://user:secret@host/")
stored = load_last_error("ws-1")
# The secret is gone but the host (useful for triage) survives.
# We don't pin the redaction marker — output_guard owns the
# format and the assertion above is the behaviour we care about.
assert "secret" not in stored
assert "host/" in stored
def test_noop_on_empty_ws_id(self, tmp_db):
from turnstone.core.memory import persist_last_error
# Must not raise; must not write anywhere observable.
persist_last_error("", "anything") # no-op
def test_noop_on_empty_err_msg(self, tmp_db):
from turnstone.core.memory import (
load_last_error,
persist_last_error,
register_workstream,
)
register_workstream("ws-1", user_id="u1")
persist_last_error("ws-1", "")
# Empty err_msg is a no-op — the row stays absent rather than
# being upserted with an empty string.
assert load_last_error("ws-1") == ""
def test_swallows_storage_failure(self, tmp_db, monkeypatch):
"""A storage failure must not propagate — error surfacing is
advisory, not safety-critical. The exception path of a worker
thread already has enough trouble without this."""
from turnstone.core import memory as memory_mod
from turnstone.core.memory import persist_last_error
class _BoomStorage:
def save_workstream_config(self, *_args, **_kw):
raise RuntimeError("simulated storage failure")
monkeypatch.setattr(memory_mod, "get_storage", lambda: _BoomStorage())
# Must not raise.
persist_last_error("ws-1", "TimeoutError: x")
class TestClearLastError:
"""Verify clear_last_error wipes the row idempotently."""
def test_clears_existing(self, tmp_db):
from turnstone.core.memory import (
clear_last_error,
load_last_error,
persist_last_error,
register_workstream,
)
register_workstream("ws-1", user_id="u1")
persist_last_error("ws-1", "RuntimeError: boom")
assert load_last_error("ws-1") == "RuntimeError: boom"
clear_last_error("ws-1")
assert load_last_error("ws-1") == ""
def test_clear_preserves_other_config_keys(self, tmp_db):
"""clear_last_error must not delete sibling config rows
(close_reason, tasks). It writes an empty string to the
last_error key only — INSERT OR REPLACE per key, no row-wide
delete."""
from turnstone.core.memory import (
clear_last_error,
load_workstream_config,
persist_last_error,
register_workstream,
save_workstream_config,
)
register_workstream("ws-1", user_id="u1")
save_workstream_config("ws-1", {"close_reason": "user closed"})
persist_last_error("ws-1", "RuntimeError: boom")
clear_last_error("ws-1")
cfg = load_workstream_config("ws-1")
# close_reason untouched.
assert cfg.get("close_reason") == "user closed"
def test_noop_on_empty_ws_id(self, tmp_db):
from turnstone.core.memory import clear_last_error
clear_last_error("") # must not raise