From 2ed33fb38ea9a238ff5a2c58a0cd5be1b75882f6 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Sun, 6 Sep 2026 12:01:44 -0700
Subject: [PATCH] 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.
---
 agent/agent_init.py                         |   2 +-
 agent/memory_manager.py                     |  25 ++--
 tests/agent/test_memory_provider.py         | 122 +++++---------------
 tests/agent/test_memory_recall_indicator.py |   9 +-
 tests/agent/test_memory_session_switch.py   |   6 +-
 tests/agent/test_memory_user_id.py          |   4 +-
 6 files changed, 49 insertions(+), 119 deletions(-)

diff --git a/agent/agent_init.py b/agent/agent_init.py
index 7df971bab1..562dfcc328 100644
--- a/agent/agent_init.py
+++ b/agent/agent_init.py
@@ -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 = ""
diff --git a/agent/memory_manager.py b/agent/memory_manager.py
index d3f5eb28f3..44ca6f930e 100644
--- a/agent/memory_manager.py
+++ b/agent/memory_manager.py
@@ -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
diff --git a/tests/agent/test_memory_provider.py b/tests/agent/test_memory_provider.py
index 023302f67b..b4243ce9b3 100644
--- a/tests/agent/test_memory_provider.py
+++ b/tests/agent/test_memory_provider.py
@@ -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")
diff --git a/tests/agent/test_memory_recall_indicator.py b/tests/agent/test_memory_recall_indicator.py
index 672ac5aeeb..4179ea6f2c 100644
--- a/tests/agent/test_memory_recall_indicator.py
+++ b/tests/agent/test_memory_recall_indicator.py
@@ -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"
diff --git a/tests/agent/test_memory_session_switch.py b/tests/agent/test_memory_session_switch.py
index 3d51c4eaac..2156ee475c 100644
--- a/tests/agent/test_memory_session_switch.py
+++ b/tests/agent/test_memory_session_switch.py
@@ -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
diff --git a/tests/agent/test_memory_user_id.py b/tests/agent/test_memory_user_id.py
index 6d852eb799..cecf809f2f 100644
--- a/tests/agent/test_memory_user_id.py
+++ b/tests/agent/test_memory_user_id.py
@@ -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(