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:
+1
-1
@@ -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
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user