From cbe0ec164acc5f63b1b78939374385058aad91d6 Mon Sep 17 00:00:00 2001 From: asemabdallah Date: Sun, 11 Oct 2026 02:28:21 +0200 Subject: [PATCH] fix(agent): free all thread data when BoundedMemorySaver evicts --- apps/agent/src/bounded_memory_saver.py | 24 +++++++---- apps/agent/tests/test_bounded_memory_saver.py | 42 +++++++++++++++++++ 2 files changed, 57 insertions(+), 9 deletions(-) create mode 100644 apps/agent/tests/test_bounded_memory_saver.py diff --git a/apps/agent/src/bounded_memory_saver.py b/apps/agent/src/bounded_memory_saver.py index aed473f..0521bb8 100644 --- a/apps/agent/src/bounded_memory_saver.py +++ b/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 diff --git a/apps/agent/tests/test_bounded_memory_saver.py b/apps/agent/tests/test_bounded_memory_saver.py new file mode 100644 index 0000000..0240b46 --- /dev/null +++ b/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']