mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-31 16:21:50 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(memory): preserve idle consolidation boundaries

Browse Source
This commit is contained in:
chengyongru
2026-08-19 18:40:20 +08:00
committed by chengyongru
parent 7335f71c5a
commit 12a4a8c04c
3 changed files with 206 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+23 -1
View File
@@ -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."""
nanobot/agent/memory.py
+69 -7
View File
@@ -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",
)
tests/agent/test_consolidator.py
+114
View File
@@ -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