refactor(memory): keep the spill, drop registration-trust rework

The prefetch spill is the fix; the is_builtin registration flag and duplicate-
instance rejection defended against a provider naming itself "builtin", which no
live path can do (only a test fixture does). Restore the name-based check and the
four test files that only changed for the new kwarg; keep two spill invariants.
This commit is contained in:
Teknium
2026-09-06 12:01:44 -07:00
parent d932fa5929
commit 2ed33fb38e
6 changed files with 49 additions and 119 deletions
+1 -1
View File
@@ -1286,7 +1286,7 @@ def _init_memory(agent, _agent_cfg, skip_memory, platform):
agent._memory_manager = _MemoryManager()
_mp = _load_mem(_mem_provider_name)
if _mp and _mp.is_available():
agent._memory_manager.add_provider(_mp, is_builtin=False)
agent._memory_manager.add_provider(_mp)
elif _mp is not None and _mem_provider_name not in _warned_unavailable_providers:
# unavailable_reason() reads config/probes importlib — skip it once warned.
_unavailable_reason = ""
+8 -17
View File
@@ -296,7 +296,6 @@ class MemoryManager:
def __init__(self, *, external_prefetch_timeout: Optional[float] = None) -> None:
self._providers: List[MemoryProvider] = []
self._tool_to_provider: Dict[str, MemoryProvider] = {}
self._builtin_provider_ids: set[int] = set()
self._external_prefetch_spill_config: Optional[Dict[str, Any]] = None
self._has_external: bool = False
timeout = external_prefetch_timeout
@@ -330,19 +329,11 @@ class MemoryManager:
logger.log(level, "Memory provider '%s' %s: %s", provider.name, label, e, exc_info=exc_info)
return results
def _is_builtin_provider(self, provider: MemoryProvider) -> bool:
return id(provider) in self._builtin_provider_ids
def add_provider(self, provider: MemoryProvider, *, is_builtin: bool = False) -> None:
"""Trust comes from registration, never the provider-controlled name.
Reject duplicate instances so their trust cannot change retroactively.
"""
if any(existing is provider for existing in self._providers):
logger.warning("Rejected duplicate memory provider instance '%s'", provider.name)
return
if not is_builtin:
def add_provider(self, provider: MemoryProvider) -> None:
"""Register a provider; builtin always accepted, only ONE external allowed."""
if provider.name != "builtin":
if self._has_external:
existing = next((p.name for p in self._providers if not self._is_builtin_provider(p)), "unknown")
existing = next((p.name for p in self._providers if p.name != "builtin"), "unknown")
logger.warning(
"Rejected memory provider '%s' — external provider '%s' is "
"already registered. Only one external memory provider is "
@@ -354,8 +345,6 @@ class MemoryManager:
self._external_prefetch_spill_config = get_spill_config()
self._providers.append(provider)
if is_builtin:
self._builtin_provider_ids.add(id(provider))
# Core tool names are reserved: built-ins always win at agent init, so a shadowing
# provider tool would linger in ``_tool_to_provider`` and hijack dispatch.
@@ -415,7 +404,7 @@ class MemoryManager:
def _prefetch_provider(self, provider: MemoryProvider, query: str, *, session_id: str = "") -> str:
"""Run one provider's prefetch; external providers are bounded by a timeout. A stuck external
call keeps running on its daemon thread and the provider is skipped on later turns until it returns."""
if self._is_builtin_provider(provider):
if provider.name == "builtin":
return provider.prefetch(query, session_id=session_id)
result_box: Dict[str, Any] = {}
@@ -450,6 +439,8 @@ class MemoryManager:
raise result_box["error"]
result = result_box.get("value", "")
if result and result.strip():
# Prefetch is stamped into the user turn's api_content and replayed every later turn;
# spill oversized results like plugin hook output so one provider can't inflate the prefix.
result = spill_if_oversized(
result, session_id=session_id, source=f"{provider.name} memory prefetch",
config=self._external_prefetch_spill_config,
@@ -729,7 +720,7 @@ class MemoryManager:
else:
provider.on_memory_write(action, target, content, metadata=dict(metadata or {}))
external = [p for p in self._providers if not self._is_builtin_provider(p)]
external = [p for p in self._providers if p.name != "builtin"]
self._each_provider("on_memory_write failed", _notify, providers=external)
# Actions mirrored to external providers; non-mutating results (errors, staged) are
+31 -91
View File
@@ -3,11 +3,10 @@
import json
import threading
import time
import pytest
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from agent.memory_provider import MemoryProvider
from agent.memory_manager import MemoryManager, inject_memory_provider_tools
@@ -146,23 +145,6 @@ class TestMemoryProviderABC:
class TestMemoryManager:
@staticmethod
def _set_spill_config(monkeypatch, tmp_path, *, max_chars):
load_config = MagicMock(
return_value={
"enabled": True,
"max_chars": max_chars,
"preview_head": 12,
"preview_tail": 12,
"directory": str(tmp_path),
}
)
monkeypatch.setattr(
"agent.memory_manager.get_spill_config",
load_config,
)
return load_config
def test_empty_manager(self):
mgr = MemoryManager()
assert mgr.providers == []
@@ -189,20 +171,13 @@ class TestMemoryManager:
def test_prefetch_merges_results(self):
mgr = MemoryManager()
p1 = FakeMemoryProvider("builtin")
p1._prefetch_result = "Memory from builtin"
p2 = FakeMemoryProvider("external")
p2._prefetch_result = "Memory from external"
mgr.add_provider(p1, is_builtin=True)
mgr.add_provider(p2)
result = mgr.prefetch_all("what do you know?")
assert "Memory from builtin" in result
assert "Memory from external" in result
assert p1.prefetch_queries == ["what do you know?"]
assert p2.prefetch_queries == ["what do you know?"]
@staticmethod
def _set_spill_config(monkeypatch, tmp_path, *, max_chars):
monkeypatch.setattr(
"agent.memory_manager.get_spill_config",
lambda: {"enabled": True, "max_chars": max_chars, "preview_head": 12,
"preview_tail": 12, "directory": str(tmp_path)},
)
def test_oversized_external_prefetch_is_spilled(self, tmp_path, monkeypatch):
self._set_spill_config(monkeypatch, tmp_path, max_chars=40)
@@ -218,72 +193,37 @@ class TestMemoryManager:
assert len(spill_files) == 1
assert spill_files[0].read_text() == provider._prefetch_result + "\n"
def test_external_prefetch_under_limit_is_unchanged(self, tmp_path, monkeypatch):
load_config = self._set_spill_config(monkeypatch, tmp_path, max_chars=100)
mgr = MemoryManager()
provider = FakeMemoryProvider("external")
provider._prefetch_result = "remember this exactly"
mgr.add_provider(provider)
result = mgr.prefetch_all("what do you remember?", session_id="session-2")
repeated = mgr.prefetch_all("what else?", session_id="session-2")
assert result == repeated == provider._prefetch_result
load_config.assert_called_once_with()
assert not list(tmp_path.rglob("*.txt"))
def test_builtin_prefetch_keeps_its_existing_budget(self, tmp_path, monkeypatch):
def test_builtin_prefetch_is_not_spilled(self, tmp_path, monkeypatch):
self._set_spill_config(monkeypatch, tmp_path, max_chars=10)
mgr = MemoryManager()
provider = FakeMemoryProvider("builtin")
provider._prefetch_result = "built-in memory has its own configured limit"
mgr.add_provider(provider, is_builtin=True)
mgr.add_provider(provider)
result = mgr.prefetch_all("what do you remember?", session_id="session-3")
assert result == provider._prefetch_result
assert mgr.prefetch_all("what do you remember?", session_id="s") == provider._prefetch_result
assert not list(tmp_path.rglob("*.txt"))
def test_external_provider_cannot_claim_builtin_spill_exemption(
self, tmp_path, monkeypatch
):
self._set_spill_config(monkeypatch, tmp_path, max_chars=40)
def test_prefetch_merges_results(self):
mgr = MemoryManager()
provider = FakeMemoryProvider("builtin")
provider._prefetch_result = "external recall " * 20
mgr.add_provider(provider)
p1 = FakeMemoryProvider("builtin")
p1._prefetch_result = "Memory from builtin"
p2 = FakeMemoryProvider("external")
p2._prefetch_result = "Memory from external"
mgr.add_provider(p1)
mgr.add_provider(p2)
result = mgr.prefetch_all("what do you remember?", session_id="session-4")
assert "builtin memory prefetch output truncated" in result
spill_files = list((tmp_path / "session-4").glob("*.txt"))
assert len(spill_files) == 1
assert spill_files[0].read_text() == provider._prefetch_result + "\n"
def test_duplicate_registration_cannot_promote_external_provider(
self, tmp_path, monkeypatch
):
self._set_spill_config(monkeypatch, tmp_path, max_chars=40)
mgr = MemoryManager()
provider = FakeMemoryProvider("external")
provider._prefetch_result = "external recall " * 20
mgr.add_provider(provider)
mgr.add_provider(provider, is_builtin=True)
result = mgr.prefetch_all("what do you remember?", session_id="session-5")
assert mgr.providers == [provider]
assert "external memory prefetch output truncated" in result
spill_files = list((tmp_path / "session-5").glob("*.txt"))
assert len(spill_files) == 1
assert spill_files[0].read_text() == provider._prefetch_result + "\n"
result = mgr.prefetch_all("what do you know?")
assert "Memory from builtin" in result
assert "Memory from external" in result
assert p1.prefetch_queries == ["what do you know?"]
assert p2.prefetch_queries == ["what do you know?"]
def test_queue_prefetch_all(self):
mgr = MemoryManager()
p1 = FakeMemoryProvider("builtin")
p2 = FakeMemoryProvider("external")
mgr.add_provider(p1, is_builtin=True)
mgr.add_provider(p1)
mgr.add_provider(p2)
mgr.queue_prefetch_all("next turn")
@@ -300,7 +240,7 @@ class TestMemoryManager:
p1 = FakeMemoryProvider("builtin")
p1.sync_turn = MagicMock(side_effect=RuntimeError("boom"))
p2 = FakeMemoryProvider("external")
mgr.add_provider(p1, is_builtin=True)
mgr.add_provider(p1)
mgr.add_provider(p2)
mgr.sync_all("user", "assistant")
@@ -342,7 +282,7 @@ class TestMemoryManager:
p2 = FakeMemoryProvider("external", tools=[
{"name": "ext_tool", "description": "External", "parameters": {}}
])
mgr.add_provider(p1, is_builtin=True)
mgr.add_provider(p1)
mgr.add_provider(p2)
r1 = json.loads(mgr.handle_tool_call("builtin_tool", {"a": 1}))
@@ -366,7 +306,7 @@ class TestMemoryManager:
builtin._prefetch_result = "builtin memory"
external = BlockingPrefetchProvider("hy-memory")
external._prefetch_result = "late external memory"
mgr.add_provider(builtin, is_builtin=True)
mgr.add_provider(builtin)
mgr.add_provider(external)
started = time.monotonic()
@@ -828,7 +768,7 @@ class TestSequentialDispatchRouting:
{"name": "ext_recall", "description": "E1", "parameters": {}},
{"name": "ext_retain", "description": "E2", "parameters": {}},
])
mgr.add_provider(builtin, is_builtin=True)
mgr.add_provider(builtin)
mgr.add_provider(external)
names = mgr.get_all_tool_names()
@@ -1013,7 +953,7 @@ class TestCommitMemorySessionRouting:
mgr = MemoryManager()
builtin = _CommitRecorder("builtin")
external = _CommitRecorder("openviking")
mgr.add_provider(builtin, is_builtin=True)
mgr.add_provider(builtin)
mgr.add_provider(external)
msgs = [{"role": "user", "content": "hi"}]
@@ -1027,7 +967,7 @@ class TestCommitMemorySessionRouting:
builtin = FakeMemoryProvider("builtin")
bad = _CommitRecorder("bad-provider")
bad.on_session_end = lambda m: (_ for _ in ()).throw(RuntimeError("boom"))
mgr.add_provider(builtin, is_builtin=True)
mgr.add_provider(builtin)
mgr.add_provider(bad)
mgr.on_session_end([]) # must not raise
@@ -1100,7 +1040,7 @@ class TestOnMemoryWriteBridge:
bad = FakeMemoryProvider("builtin")
bad.on_memory_write = MagicMock(side_effect=RuntimeError("boom"))
good = FakeMemoryProvider("good")
mgr.add_provider(bad, is_builtin=True)
mgr.add_provider(bad)
mgr.add_provider(good)
mgr.on_memory_write("add", "user", "test")
+4 -5
View File
@@ -72,11 +72,10 @@ def test_zero_count_renders_generic():
def test_aggregates_multiple_providers():
# Use the trusted builtin slot plus one external to exercise the join path.
# builtin is always accepted first; a second external is rejected, so use
# builtin + one external to exercise the join path.
mgr = MemoryManager()
mgr.add_provider(
_FakeProvider("builtin", RecallStatus("Notes", 2)), is_builtin=True
)
mgr.add_provider(_FakeProvider("builtin", RecallStatus("Notes", 2)))
mgr.add_provider(_FakeProvider("hindsight", RecallStatus("Hindsight", 5)))
result = mgr.describe_recall()
assert "🧠 Notes — recalled 2 memories" in result
@@ -85,7 +84,7 @@ def test_aggregates_multiple_providers():
def test_failing_provider_is_skipped_not_fatal():
mgr = MemoryManager()
mgr.add_provider(_FakeProvider("builtin", None, raises=True), is_builtin=True)
mgr.add_provider(_FakeProvider("builtin", None, raises=True))
mgr.add_provider(_FakeProvider("hindsight", RecallStatus("Hindsight", 1)))
# The raising provider is swallowed; the healthy one still surfaces.
assert mgr.describe_recall() == "🧠 Hindsight — recalled 1 memory"
+3 -3
View File
@@ -94,10 +94,10 @@ class _MinimalProvider(MemoryProvider):
def test_manager_fans_out_to_all_providers():
mm = MemoryManager()
# Only one external provider is allowed; use the trusted builtin slot for p1.
# Only one external provider is allowed; use the builtin slot for p1.
p1 = _RecordingProvider(name="builtin")
p2 = _RecordingProvider(name="hindsight")
mm.add_provider(p1, is_builtin=True)
mm.add_provider(p1)
mm.add_provider(p2)
mm.on_session_switch("new-sid", parent_session_id="old-sid", reset=False, reason="resume")
@@ -125,7 +125,7 @@ def test_manager_isolates_provider_failures():
# (builtin slot) with a good external one.
broken = _Broken(name="builtin")
good = _RecordingProvider(name="good")
mm.add_provider(broken, is_builtin=True)
mm.add_provider(broken)
mm.add_provider(good)
# Must not raise — exceptions in one provider are swallowed + logged
+2 -2
View File
@@ -82,10 +82,10 @@ class TestMemoryManagerUserIdThreading:
def test_multiple_providers_all_receive_user_id(self):
mgr = MemoryManager()
# Register one trusted built-in provider and one external provider.
# Use one provider named "builtin" (always accepted) and one external
p1 = RecordingProvider("builtin")
p2 = RecordingProvider("external")
mgr.add_provider(p1, is_builtin=True)
mgr.add_provider(p1)
mgr.add_provider(p2)
mgr.initialize_all(