From 7335f71c5aaf71d3d5339be9deab23edca5e9e46 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 19 Aug 2026 12:01:36 +0800 Subject: [PATCH] perf(memory): reuse session context for idle compaction --- nanobot/agent/memory.py | 114 ++++++++++++++++-- .../agent/consolidator_idle_archive.md | 5 + tests/agent/test_consolidator.py | 101 +++++++++++++--- 3 files changed, 192 insertions(+), 28 deletions(-) create mode 100644 nanobot/templates/agent/consolidator_idle_archive.md diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index d0d257a12..0d2d72a69 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -1009,29 +1009,49 @@ class Consolidator: "agent/consolidator_archive.md", strip=True, ) + return await self._archive_request( + request_messages=[ + { + "role": "system", + "content": system_prompt, + }, + {"role": "user", "content": formatted}, + ], + fallback_messages=messages, + runtime=runtime, + session_key=session_key, + tools=None, + tool_choice=None, + ) + + async def _archive_request( + self, + *, + request_messages: list[dict[str, Any]], + fallback_messages: list[dict[str, Any]], + runtime: LLMRuntime, + session_key: str | None, + tools: list[dict[str, Any]] | None, + tool_choice: str | dict[str, Any] | None, + ) -> str | None: + """Run one archive request and persist either its summary or a raw fallback.""" try: response = await runtime.provider.chat_with_retry( model=runtime.model, - messages=[ - { - "role": "system", - "content": system_prompt, - }, - {"role": "user", "content": formatted}, - ], - tools=None, - tool_choice=None, + messages=request_messages, + tools=tools, + tool_choice=tool_choice, temperature=runtime.generation.temperature, max_tokens=runtime.generation.max_tokens, reasoning_effort=runtime.generation.reasoning_effort, ) except Exception: logger.warning("Consolidation provider call failed, raw-dumping to history") - self.store.raw_archive(messages, session_key=session_key) + self.store.raw_archive(fallback_messages, session_key=session_key) return None if response.finish_reason == "error": logger.warning("Consolidation provider returned an error, raw-dumping to history") - self.store.raw_archive(messages, session_key=session_key) + self.store.raw_archive(fallback_messages, session_key=session_key) return None summary = response.content or "[no summary]" self.store.append_history( @@ -1041,6 +1061,74 @@ class Consolidator: ) return summary + @staticmethod + def _session_summary_for_prompt(session: Session) -> str | None: + """Rebuild the summary text used by normal turns after an idle archive.""" + meta = session.metadata.get("_last_summary") + if not isinstance(meta, dict): + return None + summary_meta = cast(dict[str, Any], meta) + text = summary_meta.get("text") + if not isinstance(text, str) or not text: + return None + last_active = summary_meta.get("last_active") + timestamp = last_active if isinstance(last_active, str) else session.updated_at.isoformat() + return f"Previous conversation summary (last active {timestamp}):\n{text}" + + def _build_idle_archive_messages( + self, + session: Session, + *, + archive_count: int, + ) -> list[dict[str, Any]]: + """Append a temporary archive turn to the session's model-facing prefix.""" + history = self._full_replay_history(session) + prompt = render_template( + "agent/consolidator_idle_archive.md", + strip=True, + archive_count=archive_count, + ) + channel = session.key.split(":", 1)[0] if ":" in session.key else None + built = self._build_messages( + history=history, + current_message=prompt, + channel=channel, + session_summary=self._session_summary_for_prompt(session), + session_key=session.key, + unified_session=self.unified_session, + ) + system_prefix = ( + [built[0]] + if built and built[0].get("role") == "system" + else [] + ) + return [ + *system_prefix, + *history, + {"role": "user", "content": prompt}, + ] + + async def _archive_idle_tail( + self, + session: Session, + messages: list[dict[str, Any]], + *, + runtime: LLMRuntime, + ) -> str | None: + """Archive an idle tail by extending the ordinary model-facing messages.""" + request_messages = self._build_idle_archive_messages( + session, + archive_count=len(messages), + ) + return await self._archive_request( + request_messages=request_messages, + fallback_messages=messages, + runtime=runtime, + session_key=session.key, + tools=self._get_tool_definitions(), + tool_choice="none", + ) + async def maybe_consolidate_by_tokens( self, session: Session, @@ -1183,10 +1271,10 @@ class Consolidator: last_active = session.updated_at archive_end = archive_start + len(messages_to_archive) - summary = await self.archive( + summary = await self._archive_idle_tail( + session, messages_to_archive, runtime=runtime, - session_key=session_key, ) if summary and summary != "(nothing)": diff --git a/nanobot/templates/agent/consolidator_idle_archive.md b/nanobot/templates/agent/consolidator_idle_archive.md new file mode 100644 index 000000000..7ca46d345 --- /dev/null +++ b/nanobot/templates/agent/consolidator_idle_archive.md @@ -0,0 +1,5 @@ +Create a memory overview for only the final {{ archive_count }} conversation messages immediately before this instruction. Earlier messages are context for resolving references; do not summarize them again. + +Do not call tools. Return only the overview, following these memory rules: + +{% include 'agent/consolidator_archive.md' %} diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index 120d84926..bbd84007b 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -88,6 +88,14 @@ def _provider_state() -> ProviderConversationState: ) +def _build_test_messages(**kwargs): + return [ + {"role": "system", "content": "system prompt"}, + *kwargs["history"], + {"role": "user", "content": kwargs["current_message"]}, + ] + + class TestConsolidatorSummarize: async def test_archive_prompt_includes_media_breadcrumb( self, consolidator, mock_provider, store, runtime @@ -634,7 +642,7 @@ class TestCompactIdleSession: return Consolidator( store=store, sessions=sessions, - build_messages=MagicMock(return_value=[]), + build_messages=MagicMock(side_effect=_build_test_messages), get_tool_definitions=MagicMock(return_value=[]), ) @@ -721,11 +729,14 @@ class TestCompactIdleSession: await real_consolidator.compact_idle_session("cli:incremental", runtime=runtime) assert mock_provider.chat_with_retry.await_count == 2 - latest_prompt = mock_provider.chat_with_retry.await_args_list[-1].kwargs["messages"][1][ - "content" + latest_messages = mock_provider.chat_with_retry.await_args_list[-1].kwargs["messages"] + assert [message["content"] for message in latest_messages[1:5]] == [ + "first user", + "first assistant", + "second user", + "second assistant", ] - assert "second user" in latest_prompt - assert "first user" not in latest_prompt + assert "final 2 conversation messages" in latest_messages[-1]["content"] assert sessions.get_or_create("cli:incremental").last_consolidated == 4 @pytest.mark.asyncio @@ -777,8 +788,11 @@ class TestCompactIdleSession: "cli:correction", runtime=runtime, max_suffix=8 ) - summarized = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"] - assert "CORRECTED_FINAL_RESULT_alpha" in summarized + sent_messages = mock_provider.chat_with_retry.call_args.kwargs["messages"] + assert any( + message.get("content") == "CORRECTED_FINAL_RESULT_alpha" + for message in sent_messages + ) @pytest.mark.asyncio async def test_raw_dumps_full_archive_batch_on_llm_failure( @@ -939,10 +953,13 @@ class TestCompactIdleSession: # Verify only the unconsolidated tail was processed: # All 10 unconsolidated messages (50-59) are archived exactly once. archived_call = mock_provider.chat_with_retry.call_args - user_content = archived_call.kwargs["messages"][1]["content"] - # Should contain only tail messages, not early ones - assert "u0" not in user_content - assert "u25" in user_content or "a25" in user_content + sent_messages = archived_call.kwargs["messages"] + sent_content = [message.get("content") for message in sent_messages] + # The ordinary replay prefix contributes recent context, while the + # temporary instruction limits the new overview to the unarchived tail. + assert "u0" not in sent_content + assert "u26" in sent_content + assert "final 10 conversation messages" in sent_messages[-1]["content"] @pytest.mark.asyncio async def test_full_archive_keeps_extended_legal_replay_suffix( @@ -988,10 +1005,64 @@ class TestCompactIdleSession: # the dropped head (user-00) and retained suffix (user-14 through # assistant-09) are all summarized. archived_call = mock_provider.chat_with_retry.call_args - user_content = archived_call.kwargs["messages"][1]["content"] - assert "user-00" in user_content - assert "assistant-09" in user_content - assert "user-14" in user_content + sent_content = [message.get("content") for message in archived_call.kwargs["messages"]] + assert "user-00" in sent_content + assert "assistant-09" in sent_content + assert "user-14" in sent_content + + @pytest.mark.asyncio + async def test_reuses_model_prefix_without_persisting_temporary_turn( + self, + real_consolidator, + mock_provider, + store, + runtime, + ): + tools = [{"type": "function", "function": {"name": "lookup"}}] + real_consolidator._get_tool_definitions.return_value = tools + mock_provider.chat_with_retry.return_value = LLMResponse( + content="Overview from the temporary turn.", + finish_reason="stop", + ) + sessions = real_consolidator.sessions + session = sessions.get_or_create("cli:tool-history") + session.add_message("user", "look this up") + session.messages.extend(_tool_round("call-1")) + session.add_message("assistant", "final answer") + sessions.save(session) + + result = await real_consolidator.compact_idle_session( + "cli:tool-history", + runtime=runtime, + ) + + assert result == "Overview from the temporary turn." + call = mock_provider.chat_with_retry.call_args.kwargs + sent_messages = call["messages"] + assert [message["role"] for message in sent_messages] == [ + "system", + "user", + "assistant", + "tool", + "assistant", + "user", + ] + assert sent_messages[2]["tool_calls"][0]["id"] == "call-1" + assert "final 4 conversation messages" in sent_messages[-1]["content"] + assert call["tools"] == tools + assert call["tool_choice"] == "none" + + reloaded = sessions.get_or_create("cli:tool-history") + assert len(reloaded.messages) == 4 + assert reloaded.messages[-1]["content"] == "final answer" + assert all( + "memory overview" not in str(message.get("content", "")).lower() + for message in reloaded.messages + ) + entries = store.read_unprocessed_history(since_cursor=0) + assert [entry["content"] for entry in entries] == [ + "Overview from the temporary turn." + ] @pytest.mark.asyncio async def test_acquires_consolidation_lock(