diff --git a/run_agent.py b/run_agent.py index 801678f371e5..62b0b4f0068d 100644 --- a/run_agent.py +++ b/run_agent.py @@ -5063,16 +5063,25 @@ def shutdown_memory_provider(self, messages: list = None) -> None: pass def commit_memory_session(self, messages: list = None) -> None: - """Trigger end-of-session extraction without tearing providers down. + """Trigger end-of-session hooks without tearing providers down. 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 + session_messages = messages or [] + memory_manager = getattr(self, "_memory_manager", None) + if memory_manager: + try: + memory_manager.on_session_end(session_messages) + except Exception: + pass + if hasattr(self, "context_compressor") and self.context_compressor: + try: + self.context_compressor.on_session_end( + self.session_id or "", + session_messages, + ) + except Exception: + pass def _sync_external_memory_for_turn( self, diff --git a/tests/run_agent/test_commit_memory_session.py b/tests/run_agent/test_commit_memory_session.py new file mode 100644 index 000000000000..2aa5171490d8 --- /dev/null +++ b/tests/run_agent/test_commit_memory_session.py @@ -0,0 +1,79 @@ +from unittest.mock import MagicMock + +from run_agent import AIAgent + + +def _bare_agent(): + agent = AIAgent.__new__(AIAgent) + agent.session_id = "session-123" + agent._memory_manager = MagicMock() + agent.context_compressor = MagicMock() + return agent + + +def test_commit_memory_session_notifies_memory_and_context_engine(): + agent = _bare_agent() + messages = [{"role": "user", "content": "persist this final turn"}] + + agent.commit_memory_session(messages) + + agent._memory_manager.on_session_end.assert_called_once_with(messages) + agent._memory_manager.shutdown_all.assert_not_called() + agent.context_compressor.on_session_end.assert_called_once_with( + "session-123", + messages, + ) + + +def test_commit_memory_session_notifies_context_engine_without_memory_manager(): + agent = _bare_agent() + agent._memory_manager = None + messages = [{"role": "assistant", "content": "context engines still need final flush"}] + + agent.commit_memory_session(messages) + + agent.context_compressor.on_session_end.assert_called_once_with( + "session-123", + messages, + ) + + +def test_commit_memory_session_notifies_context_engine_without_memory_attr(): + agent = _bare_agent() + del agent._memory_manager + messages = [{"role": "assistant", "content": "partially constructed agents should still flush"}] + + agent.commit_memory_session(messages) + + agent.context_compressor.on_session_end.assert_called_once_with( + "session-123", + messages, + ) + + +def test_commit_memory_session_still_notifies_context_engine_when_memory_fails(): + agent = _bare_agent() + agent._memory_manager.on_session_end.side_effect = RuntimeError("memory flush failed") + messages = [{"role": "user", "content": "context engine should still flush"}] + + agent.commit_memory_session(messages) + + agent._memory_manager.on_session_end.assert_called_once_with(messages) + agent.context_compressor.on_session_end.assert_called_once_with( + "session-123", + messages, + ) + + +def test_commit_memory_session_swallows_context_engine_errors(): + agent = _bare_agent() + agent.context_compressor.on_session_end.side_effect = RuntimeError("flush failed") + messages = [{"role": "user", "content": "do not block reset"}] + + agent.commit_memory_session(messages) + + agent._memory_manager.on_session_end.assert_called_once_with(messages) + agent.context_compressor.on_session_end.assert_called_once_with( + "session-123", + messages, + ) diff --git a/tests/run_agent/test_compression_boundary_hook.py b/tests/run_agent/test_compression_boundary_hook.py index 26bac74163b2..0a48ec23b76f 100644 --- a/tests/run_agent/test_compression_boundary_hook.py +++ b/tests/run_agent/test_compression_boundary_hook.py @@ -86,6 +86,21 @@ def test_on_session_start_called_with_compression_boundary(self): assert call.kwargs.get("old_session_id") == original_sid, \ f"Expected old_session_id={original_sid!r}, got {call.kwargs!r}" + compressor.on_session_end.assert_called_once_with(original_sid, messages) + end_index = next( + idx for idx, mock_call in enumerate(compressor.mock_calls) + if mock_call[0] == "on_session_end" + ) + start_index = next( + idx for idx, mock_call in enumerate(compressor.mock_calls) + if ( + mock_call[0] == "on_session_start" + and mock_call.kwargs.get("boundary_reason") == "compression" + ) + ) + assert end_index < start_index, \ + "context engine should receive old-session end before new compression-boundary start" + def test_no_hook_when_no_session_db(self): """Without session_db, session_id does not rotate and the hook is not fired.""" from run_agent import AIAgent