fix(agent): free all thread data when BoundedMemorySaver evicts

apps/agent/src/bounded_memory_saver.py

@@ -10,7 +10,8 @@
 rather than sorting keys, so eviction order is correct even when thread IDs
 are UUIDs or other non-chronological strings.
 
-NOTE: This class relies on MemorySaver.storage (an internal attribute).
+NOTE: This class relies on MemorySaver.delete_thread and on how MemorySaver
+      stores threads internally.
       The langgraph version is pinned in pyproject.toml to guard against
       breaking changes.
 
@@ -43,13 +44,18 @@ def put(self, config, checkpoint, metadata, new_versions):
 
         result = super().put(config, checkpoint, metadata, new_versions)
 
-        while len(self.storage) > self.max_threads and self._insertion_order:
+        # Count tracked threads rather than self.storage: reading an unknown
+        # thread (get_tuple/list) adds an empty entry to the storage defaultdict
+        # that put() never tracks, which would otherwise count toward the limit
+        # and force the thread just written to be evicted.
+        while len(self._insertion_order) > self.max_threads:
             oldest_thread, _ = self._insertion_order.popitem(last=False)
-            if oldest_thread in self.storage:
-                logger.info(
-                    "BoundedMemorySaver: evicting thread %s (%d threads stored)",
-                    oldest_thread,
-                    len(self.storage),
-                )
-                del self.storage[oldest_thread]
+            logger.info(
+                "BoundedMemorySaver: evicting thread %s (%d threads tracked)",
+                oldest_thread,
+                len(self._insertion_order),
+            )
+            # delete_thread also drops the thread's channel values (blobs) and
+            # pending writes, which hold the message data.
+            self.delete_thread(oldest_thread)
         return result

apps/agent/tests/test_bounded_memory_saver.py

@@ -0,0 +1,42 @@
+"""BoundedMemorySaver must cap stored threads without dropping live ones."""
+from langchain_core.messages import AIMessage, HumanMessage
+from langgraph.graph import END, START, MessagesState, StateGraph
+
+from src.bounded_memory_saver import BoundedMemorySaver
+
+
+def _graph(saver):
+    builder = StateGraph(MessagesState)
+    builder.add_node('answer', lambda state: {'messages': [AIMessage(content='OK')]})
+    builder.add_edge(START, 'answer')
+    builder.add_edge('answer', END)
+    return builder.compile(checkpointer=saver)
+
+
+def _config(thread_id):
+    return {'configurable': {'thread_id': thread_id}}
+
+
+def test_eviction_drops_all_data_for_evicted_threads():
+    saver = BoundedMemorySaver(max_threads=2)
+    graph = _graph(saver)
+    for index in range(5):
+        graph.invoke({'messages': [HumanMessage(content='Hello')]}, _config(f't{index}'))
+
+    assert set(saver.storage) == {'t3', 't4'}
+    assert {key[0] for key in saver.blobs} <= {'t3', 't4'}
+    assert {key[0] for key in saver.writes} <= {'t3', 't4'}
+
+
+def test_reading_unknown_threads_does_not_evict_the_thread_just_written():
+    saver = BoundedMemorySaver(max_threads=2)
+    graph = _graph(saver)
+    # Reading a thread that has no checkpoint still creates an empty entry in
+    # MemorySaver.storage, as the AG-UI adapter does before every run.
+    for index in range(3):
+        assert graph.get_state(_config(f'unknown-{index}')).values == {}
+
+    graph.invoke({'messages': [HumanMessage(content='Hello')]}, _config('real'))
+
+    messages = graph.get_state(_config('real')).values['messages']
+    assert [message.content for message in messages] == ['Hello', 'OK']