From 12a4a8c04ca8925f2571402d6585400ca61fd4db Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 19 Aug 2026 13:33:53 +0800 Subject: [PATCH] fix(memory): preserve idle consolidation boundaries --- nanobot/agent/loop.py | 24 ++++++- nanobot/agent/memory.py | 76 +++++++++++++++++++-- tests/agent/test_consolidator.py | 114 +++++++++++++++++++++++++++++++ 3 files changed, 206 insertions(+), 8 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 5f6443bf4..ed883052e 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -74,7 +74,11 @@ from nanobot.session.goal_state import ( sustained_goal_active, ) from nanobot.session.history_visibility import HIDDEN_HISTORY_META -from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel +from nanobot.session.keys import ( + UNIFIED_SESSION_KEY, + last_channel_from_metadata, + remember_last_channel, +) from nanobot.session.manager import ( SESSION_CACHE_MAX_SIZE, Session, @@ -442,6 +446,7 @@ class AgentLoop: sessions=self.sessions, build_messages=self.context.build_messages, get_tool_definitions=self.tools.get_definitions, + resolve_prompt_context=self._idle_consolidation_prompt_context, consolidation_ratio=consolidation_ratio, unified_session=unified_session, ) @@ -918,6 +923,23 @@ class AgentLoop: return remember_last_channel(session.metadata, msg.channel, msg.chat_id) + def _idle_consolidation_prompt_context( + self, + session: Session, + ) -> tuple[str | None, Path]: + """Resolve the same persisted route and workspace used by a normal turn.""" + channel = session.key.split(":", 1)[0] if ":" in session.key else None + if self._unified_session: + route = last_channel_from_metadata(session.metadata) + if route is not None: + channel = route[0] + scope = self.workspace_scopes.for_turn( + channel=channel, + message_metadata=None, + session_metadata=session.metadata, + ) + return channel, scope.project_path + @staticmethod def _replay_token_budget(runtime: LLMRuntime) -> int: """Derive a token budget for session history replay from the context window.""" diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index 0d2d72a69..c411ccca7 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -21,7 +21,12 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast from loguru import logger from nanobot.runtime_context import public_history_messages -from nanobot.session.manager import MIN_COMPACTED_REPLAY_MESSAGES, Session, SessionManager +from nanobot.session.manager import ( + MIN_COMPACTED_REPLAY_MESSAGES, + Session, + SessionManager, + replay_max_messages_for_context, +) from nanobot.utils.gitstore import GitStore from nanobot.utils.helpers import ( content_with_media_breadcrumbs, @@ -815,6 +820,7 @@ class Consolidator: sessions: SessionManager, build_messages: Callable[..., list[dict[str, Any]]], get_tool_definitions: Callable[[], list[dict[str, Any]]], + resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None, consolidation_ratio: float = 0.5, unified_session: bool = False, ): @@ -824,6 +830,7 @@ class Consolidator: self.unified_session = unified_session self._build_messages = build_messages self._get_tool_definitions = get_tool_definitions + self._resolve_prompt_context = resolve_prompt_context self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( weakref.WeakValueDictionary() ) @@ -1078,22 +1085,43 @@ class Consolidator: def _build_idle_archive_messages( self, session: Session, + messages: list[dict[str, Any]], *, - archive_count: int, - ) -> list[dict[str, Any]]: + runtime: LLMRuntime, + ) -> list[dict[str, Any]] | None: """Append a temporary archive turn to the session's model-facing prefix.""" - history = self._full_replay_history(session) + replay_budget = self._input_token_budget(runtime) + if replay_budget <= 0: + replay_budget = max(128, runtime.context_window_tokens // 2) + history = session.get_history( + max_messages=replay_max_messages_for_context(runtime.context_window_tokens), + max_tokens=replay_budget, + ) + archive_history = Session( + key=session.key, + messages=messages, + ).get_history(max_messages=max(1, len(messages))) + if ( + not archive_history + or len(history) < len(archive_history) + or history[-len(archive_history):] != archive_history + ): + return None prompt = render_template( "agent/consolidator_idle_archive.md", strip=True, - archive_count=archive_count, + archive_count=len(archive_history), ) channel = session.key.split(":", 1)[0] if ":" in session.key else None + workspace: Path | None = None + if self._resolve_prompt_context is not None: + channel, workspace = self._resolve_prompt_context(session) built = self._build_messages( history=history, current_message=prompt, channel=channel, session_summary=self._session_summary_for_prompt(session), + workspace=workspace, session_key=session.key, unified_session=self.unified_session, ) @@ -1118,14 +1146,48 @@ class Consolidator: """Archive an idle tail by extending the ordinary model-facing messages.""" request_messages = self._build_idle_archive_messages( session, - archive_count=len(messages), + messages, + runtime=runtime, ) + if request_messages is None: + return await self.archive( + messages, + runtime=runtime, + session_key=session.key, + ) + tools = self._get_tool_definitions() + budget = self._input_token_budget(runtime) + if budget <= 0: + return await self.archive( + messages, + runtime=runtime, + session_key=session.key, + ) + estimated, source = estimate_prompt_tokens_chain( + runtime.provider, + runtime.model, + request_messages, + tools, + ) + if estimated > budget: + logger.debug( + "Idle consolidation falling back to bounded archive for {}: {}/{} via {}", + session.key, + estimated, + budget, + source, + ) + return await self.archive( + messages, + runtime=runtime, + session_key=session.key, + ) return await self._archive_request( request_messages=request_messages, fallback_messages=messages, runtime=runtime, session_key=session.key, - tools=self._get_tool_definitions(), + tools=tools, tool_choice="none", ) diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index bbd84007b..80a941068 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -712,6 +712,7 @@ class TestCompactIdleSession: async def test_new_messages_advance_existing_archive_progress( self, real_consolidator, mock_provider, runtime ): + runtime = replace(runtime, context_window_tokens=128_000) mock_provider.chat_with_retry.return_value = MagicMock( content="Summary.", finish_reason="stop" ) @@ -771,6 +772,7 @@ class TestCompactIdleSession: the recent suffix it retains. Otherwise a late user correction / final result that lands in the kept suffix is excluded from the persisted summary, leaving a stale wrong conclusion in history. Regression for #4264.""" + runtime = replace(runtime, context_window_tokens=128_000) mock_provider.chat_with_retry.return_value = MagicMock( content="Summary.", finish_reason="stop" ) @@ -931,6 +933,7 @@ class TestCompactIdleSession: self, real_consolidator, mock_provider, runtime ): """30 turns with last_consolidated=50 → only unconsolidated tail considered.""" + runtime = replace(runtime, context_window_tokens=128_000) mock_provider.chat_with_retry.return_value = MagicMock( content="Tail summary.", finish_reason="stop" ) @@ -968,6 +971,7 @@ class TestCompactIdleSession: mock_provider, runtime, ): + runtime = replace(runtime, context_window_tokens=128_000) mock_provider.chat_with_retry.return_value = MagicMock( content="Tail summary.", finish_reason="stop" ) @@ -1018,6 +1022,7 @@ class TestCompactIdleSession: store, runtime, ): + runtime = replace(runtime, context_window_tokens=128_000) tools = [{"type": "function", "function": {"name": "lookup"}}] real_consolidator._get_tool_definitions.return_value = tools mock_provider.chat_with_retry.return_value = LLMResponse( @@ -1064,6 +1069,115 @@ class TestCompactIdleSession: "Overview from the temporary turn." ] + @pytest.mark.asyncio + async def test_oversized_prefix_falls_back_to_bounded_archive( + self, + real_consolidator, + mock_provider, + store, + runtime, + ): + async def bounded_chat(**kwargs): + sent_chars = sum( + len(str(message.get("content") or "")) + for message in kwargs["messages"] + ) + if sent_chars > 20_000: + return LLMResponse(content="context too long", finish_reason="error") + return LLMResponse(content="Bounded summary.", finish_reason="stop") + + mock_provider.chat_with_retry.side_effect = bounded_chat + sessions = real_consolidator.sessions + session = sessions.get_or_create("sdk:oversized") + session.add_message("user", "x" * 100_000) + sessions.save(session) + + result = await real_consolidator.compact_idle_session( + "sdk:oversized", + runtime=runtime, + ) + + assert result == "Bounded summary." + sent = mock_provider.chat_with_retry.call_args.kwargs["messages"] + assert sum(len(str(message.get("content") or "")) for message in sent) < 20_000 + assert not any( + "[RAW]" in entry["content"] + for entry in store.read_unprocessed_history(since_cursor=0) + ) + assert sessions.get_or_create("sdk:oversized").last_consolidated == 1 + + @pytest.mark.asyncio + async def test_incremental_scope_counts_only_model_visible_messages( + self, + real_consolidator, + mock_provider, + runtime, + ): + runtime = replace(runtime, context_window_tokens=128_000) + mock_provider.chat_with_retry.return_value = LLMResponse( + content="Summary.", + finish_reason="stop", + ) + sessions = real_consolidator.sessions + session = sessions.get_or_create("cli:commands") + session.add_message("user", "already archived user") + session.add_message("assistant", "already archived answer") + session.last_consolidated = 2 + session.add_message("user", "/status", _command=True) + session.add_message("assistant", "status output", _command=True) + session.add_message("user", "new user") + session.add_message("assistant", "new answer") + sessions.save(session) + + await real_consolidator.compact_idle_session( + "cli:commands", + runtime=runtime, + ) + + sent = mock_provider.chat_with_retry.call_args.kwargs["messages"] + assert [message.get("content") for message in sent[1:-1]] == [ + "already archived user", + "already archived answer", + "new user", + "new answer", + ] + assert "final 2 conversation messages" in sent[-1]["content"] + + @pytest.mark.asyncio + async def test_uses_persisted_workspace_scope_for_system_prefix( + self, + loop_factory, + mock_provider, + tmp_path, + ): + project = tmp_path / "project" + project.mkdir() + (tmp_path / "AGENTS.md").write_text("GLOBAL_WORKSPACE_MARKER", encoding="utf-8") + (project / "AGENTS.md").write_text("PROJECT_WORKSPACE_MARKER", encoding="utf-8") + loop = loop_factory(provider=mock_provider) + runtime = loop.llm_runtime() + runtime.provider.chat_with_retry.return_value = LLMResponse( + content="Summary.", + finish_reason="stop", + ) + session = loop.sessions.get_or_create("websocket:scope") + session.metadata["workspace_scope"] = { + "project_path": str(project), + "access_mode": "restricted", + } + session.add_message("user", "project question") + session.add_message("assistant", "project answer") + loop.sessions.save(session) + + await loop.consolidator.compact_idle_session( + "websocket:scope", + runtime=runtime, + ) + + system = runtime.provider.chat_with_retry.call_args.kwargs["messages"][0]["content"] + assert "PROJECT_WORKSPACE_MARKER" in system + assert "GLOBAL_WORKSPACE_MARKER" not in system + @pytest.mark.asyncio async def test_acquires_consolidation_lock( self, real_consolidator, mock_provider, runtime