From e90aa7f2802ea1a688df7189b490843f829c6caf Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Sat, 9 May 2026 12:28:42 -0700 Subject: [PATCH] fix(agent): notify context engine on commit_memory_session (#22764) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When session_id rotates (e.g. /new), commit_memory_session was firing MemoryManager.on_session_end but skipping ContextEngine.on_session_end. Engines that accumulate per-session state (LCM-style DAGs, summary stores) leaked that state from the rotated-out session into whatever continued under the same compressor instance. Mirror the call shutdown_memory_provider already makes — same lifecycle moment, same hook contract ("real session boundaries (CLI exit, /reset, gateway expiry)"). /new is a real boundary for the old session_id; providers keep their state but the rotated-out session_id is done. 6 regression tests covering both-hooks-fire, no-memory-manager, no-context-engine, both failure-tolerant paths. Closes #22394. --- run_agent.py | 25 +++-- ...st_commit_memory_session_context_engine.py | 102 ++++++++++++++++++ 2 files changed, 121 insertions(+), 6 deletions(-) create mode 100644 tests/run_agent/test_commit_memory_session_context_engine.py diff --git a/run_agent.py b/run_agent.py index 801678f371e..aaceb79c752 100644 --- a/run_agent.py +++ b/run_agent.py @@ -5067,12 +5067,25 @@ class AIAgent: Called when session_id rotates (e.g. /new, context compression); providers keep their state and continue running under the old session_id — they just flush pending extraction now.""" - if not self._memory_manager: - return - try: - self._memory_manager.on_session_end(messages or []) - except Exception: - pass + if self._memory_manager: + try: + self._memory_manager.on_session_end(messages or []) + except Exception: + pass + # Notify context engine of session end too — same lifecycle moment as + # the memory manager's on_session_end. Without this, engines that + # accumulate per-session state (DAGs, summaries) leak that state from + # the rotated-out session into whatever comes next under the same + # compressor instance. Mirrors the call in shutdown_memory_provider(). + # See issue #22394. + if hasattr(self, "context_compressor") and self.context_compressor: + try: + self.context_compressor.on_session_end( + self.session_id or "", + messages or [], + ) + except Exception: + pass def _sync_external_memory_for_turn( self, diff --git a/tests/run_agent/test_commit_memory_session_context_engine.py b/tests/run_agent/test_commit_memory_session_context_engine.py new file mode 100644 index 00000000000..307814891a2 --- /dev/null +++ b/tests/run_agent/test_commit_memory_session_context_engine.py @@ -0,0 +1,102 @@ +"""Regression tests for AIAgent.commit_memory_session. + +Issue #22394: commit_memory_session was calling MemoryManager.on_session_end +but never ContextEngine.on_session_end. Context engines that accumulate +per-session state (LCM-style DAGs, summary stores) leaked that state from a +rotated-out session into whatever continued under the same compressor +instance. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + + +def _make_minimal_agent(memory_manager, context_compressor, session_id="abc"): + """Build an object with just enough surface for commit_memory_session to run. + + AIAgent.__init__ is too heavy for a focused unit test — bind the method + to a SimpleNamespace-style object that has the attributes the method + actually touches. + """ + from run_agent import AIAgent + + obj = SimpleNamespace( + _memory_manager=memory_manager, + context_compressor=context_compressor, + session_id=session_id, + ) + obj.commit_memory_session = AIAgent.commit_memory_session.__get__(obj) + return obj + + +def test_commit_memory_session_notifies_context_engine(): + """Both the memory manager AND the context engine receive on_session_end.""" + mm = MagicMock() + ctx = MagicMock() + agent = _make_minimal_agent(mm, ctx, session_id="sess-42") + + msgs = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}] + agent.commit_memory_session(msgs) + + mm.on_session_end.assert_called_once_with(msgs) + ctx.on_session_end.assert_called_once_with("sess-42", msgs) + + +def test_commit_memory_session_with_no_messages_passes_empty_list(): + """Empty/None messages must still fire both hooks with an empty list.""" + mm = MagicMock() + ctx = MagicMock() + agent = _make_minimal_agent(mm, ctx, session_id="sess-7") + + agent.commit_memory_session(None) + + mm.on_session_end.assert_called_once_with([]) + ctx.on_session_end.assert_called_once_with("sess-7", []) + + +def test_commit_memory_session_no_memory_manager_still_notifies_context_engine(): + """If only the context engine is configured, it still gets the hook.""" + ctx = MagicMock() + agent = _make_minimal_agent(None, ctx, session_id="sess-9") + + agent.commit_memory_session([{"role": "user", "content": "x"}]) + + ctx.on_session_end.assert_called_once_with("sess-9", [{"role": "user", "content": "x"}]) + + +def test_commit_memory_session_no_context_engine_still_notifies_memory_manager(): + """If only the memory manager is configured, it still gets the hook.""" + mm = MagicMock() + agent = _make_minimal_agent(mm, None, session_id="sess-3") + + agent.commit_memory_session([{"role": "user", "content": "x"}]) + + mm.on_session_end.assert_called_once_with([{"role": "user", "content": "x"}]) + + +def test_commit_memory_session_tolerates_memory_manager_failure(): + """A raising memory manager must not block the context engine notification.""" + mm = MagicMock() + mm.on_session_end.side_effect = RuntimeError("boom") + ctx = MagicMock() + agent = _make_minimal_agent(mm, ctx, session_id="sess-X") + + # Must not raise + agent.commit_memory_session([{"role": "user", "content": "x"}]) + + ctx.on_session_end.assert_called_once_with("sess-X", [{"role": "user", "content": "x"}]) + + +def test_commit_memory_session_tolerates_context_engine_failure(): + """A raising context engine must not surface the exception.""" + mm = MagicMock() + ctx = MagicMock() + ctx.on_session_end.side_effect = RuntimeError("boom") + agent = _make_minimal_agent(mm, ctx, session_id="sess-Y") + + # Must not raise + agent.commit_memory_session([{"role": "user", "content": "x"}]) + + mm.on_session_end.assert_called_once()