fix(context): scale replay cap with context window

This commit is contained in:
chengyongru
2026-06-29 14:23:55 +08:00
committed by Xubin Ren
parent dacc699293
commit 40282e3b74
5 changed files with 83 additions and 11 deletions
+15 -3
View File
@@ -57,7 +57,11 @@ from nanobot.session.goal_state import (
sustained_goal_active,
)
from nanobot.session.keys import UNIFIED_SESSION_KEY, session_key_for_channel
from nanobot.session.manager import DEFAULT_REPLAY_MAX_MESSAGES, Session, SessionManager
from nanobot.session.manager import (
Session,
SessionManager,
replay_max_messages_for_context,
)
from nanobot.utils.document import extract_documents, reference_non_image_attachments
from nanobot.utils.helpers import image_placeholder_text
from nanobot.utils.helpers import truncate_text as truncate_text_fn
@@ -201,7 +205,7 @@ class AgentLoop:
timezone: str | None = None,
session_ttl_minutes: int = 0,
consolidation_ratio: float = 0.5,
max_messages: int = DEFAULT_REPLAY_MAX_MESSAGES,
max_messages: int = 0,
hooks: list[AgentHook] | None = None,
unified_session: bool = False,
disabled_skills: list[str] | None = None,
@@ -292,8 +296,11 @@ class AgentLoop:
llm_wall_timeout_for_session=lambda sk: runner_wall_llm_timeout_s(self.sessions, sk),
)
self._unified_session = unified_session
self._explicit_max_messages = max_messages > 0
self._max_messages = (
max_messages if max_messages > 0 else DEFAULT_REPLAY_MAX_MESSAGES
max_messages
if self._explicit_max_messages
else replay_max_messages_for_context(self.context_window_tokens)
)
self._running = False
self._mcp_servers = mcp_servers or {}
@@ -422,6 +429,7 @@ class AgentLoop:
self.runner.provider = provider
self.subagents.set_provider(provider, model)
self.consolidator.set_provider(provider, model, context_window_tokens)
self._sync_replay_max_messages()
self._provider_signature = snapshot.signature
if publish_update and self._runtime_model_publisher is not None:
self._runtime_model_publisher(
@@ -435,6 +443,10 @@ class AgentLoop:
)
logger.info("Runtime model switched for next turn: {} -> {}", old_model, model)
def _sync_replay_max_messages(self) -> None:
if not self._explicit_max_messages:
self._max_messages = replay_max_messages_for_context(self.context_window_tokens)
def _refresh_provider_snapshot(self) -> None:
if self._provider_snapshot_loader is None:
return
+3
View File
@@ -438,6 +438,9 @@ class MyTool(Tool, ContextAware):
setattr(self._runtime_state, key, value)
if key == "model":
self._runtime_state._active_preset = None
sync_replay = getattr(self._runtime_state, "_sync_replay_max_messages", None)
if key == "context_window_tokens" and callable(sync_replay):
sync_replay()
if key == "max_iterations" and hasattr(self._runtime_state, "_sync_subagent_runtime_limits"):
self._runtime_state._sync_subagent_runtime_limits()
self._audit("modify", f"{key}: {old!r} -> {value!r}")
+12 -1
View File
@@ -27,7 +27,9 @@ from nanobot.utils.helpers import (
from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body
FILE_MAX_MESSAGES = 2000
DEFAULT_REPLAY_MAX_MESSAGES = 500
MIN_REPLAY_MAX_MESSAGES = 120
REPLAY_TOKENS_PER_MESSAGE = 100
DEFAULT_REPLAY_MAX_MESSAGES = FILE_MAX_MESSAGES
_MESSAGE_TIME_PREFIX_RE = re.compile(r"^\[Message Time: [^\]]+\]\n?")
_LOCAL_IMAGE_BREADCRUMB_RE = re.compile(r"^\[image: (?:/|~)[^\]]+\]\s*$")
_TOOL_CALL_ECHO_RE = re.compile(r'^\s*(?:generate_image|message)\([^)]*\)\s*$')
@@ -44,6 +46,15 @@ _FORK_VOLATILE_METADATA_KEYS = {
}
def replay_max_messages_for_context(context_window_tokens: int | None) -> int:
if not context_window_tokens or context_window_tokens <= 0:
return DEFAULT_REPLAY_MAX_MESSAGES
return min(
FILE_MAX_MESSAGES,
max(MIN_REPLAY_MAX_MESSAGES, context_window_tokens // REPLAY_TOKENS_PER_MESSAGE),
)
def _sanitize_assistant_replay_text(content: str) -> str:
"""Remove internal replay artifacts that the model may have copied before.
+49 -5
View File
@@ -11,20 +11,31 @@ from nanobot.agent.loop import AgentLoop
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse
from nanobot.session.manager import DEFAULT_REPLAY_MAX_MESSAGES, Session
from nanobot.providers.factory import ProviderSnapshot
from nanobot.session.manager import (
DEFAULT_REPLAY_MAX_MESSAGES,
Session,
replay_max_messages_for_context,
)
DEFAULT_MAX_MESSAGES = DEFAULT_REPLAY_MAX_MESSAGES
def _make_loop(tmp_path: Path, max_messages: int = DEFAULT_MAX_MESSAGES) -> AgentLoop:
def _make_loop(
tmp_path: Path,
max_messages: int = 0,
context_window_tokens: int = 200_000,
) -> AgentLoop:
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation.max_tokens = 4096
return AgentLoop(
bus=MessageBus(),
provider=provider,
workspace=tmp_path,
model="test-model",
max_messages=max_messages,
context_window_tokens=context_window_tokens,
)
@@ -51,17 +62,26 @@ def _tool_round(call_id: str) -> list[dict]:
class TestMaxMessagesInit:
"""Verify AgentLoop stores the config value correctly."""
"""Verify AgentLoop derives the internal replay cap correctly."""
def test_default_is_builtin_limit(self, tmp_path: Path) -> None:
def test_context_formula(self) -> None:
assert replay_max_messages_for_context(8_000) == 120
assert replay_max_messages_for_context(32_768) == 327
assert replay_max_messages_for_context(200_000) == DEFAULT_MAX_MESSAGES
def test_default_for_200k_context_reaches_file_cap(self, tmp_path: Path) -> None:
loop = _make_loop(tmp_path)
assert loop._max_messages == DEFAULT_MAX_MESSAGES
def test_default_scales_with_context_window(self, tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=32_768)
assert loop._max_messages == 327
def test_positive_value_stored(self, tmp_path: Path) -> None:
loop = _make_loop(tmp_path, max_messages=25)
assert loop._max_messages == 25
def test_zero_uses_builtin_limit(self, tmp_path: Path) -> None:
def test_zero_uses_context_derived_limit(self, tmp_path: Path) -> None:
loop = _make_loop(tmp_path, max_messages=0)
assert loop._max_messages == DEFAULT_MAX_MESSAGES
@@ -70,6 +90,30 @@ class TestMaxMessagesInit:
loop = _make_loop(tmp_path, max_messages=-5)
assert loop._max_messages == DEFAULT_MAX_MESSAGES
def test_provider_refresh_resyncs_context_derived_limit(self, tmp_path: Path) -> None:
old_provider = MagicMock()
old_provider.get_default_model.return_value = "old-model"
old_provider.generation.max_tokens = 4096
new_provider = MagicMock()
new_provider.generation.max_tokens = 4096
loop = AgentLoop(
bus=MessageBus(),
provider=old_provider,
workspace=tmp_path,
model="old-model",
context_window_tokens=32_768,
provider_snapshot_loader=lambda: ProviderSnapshot(
provider=new_provider,
model="new-model",
context_window_tokens=200_000,
signature=("new-model",),
),
)
assert loop._max_messages == 327
loop._refresh_provider_snapshot()
assert loop._max_messages == DEFAULT_MAX_MESSAGES
class TestGetHistoryWithMaxMessages:
"""Verify get_history respects max_messages parameter."""
+4 -2
View File
@@ -236,10 +236,12 @@ class TestModifyRestricted:
@pytest.mark.asyncio
async def test_modify_context_window_valid(self):
tool = _make_tool()
loop = _make_mock_loop(_sync_replay_max_messages=MagicMock())
tool = _make_tool(runtime_state=loop)
result = await tool.execute(action="set", key="context_window_tokens", value=131072)
assert "Set context_window_tokens" in result
assert tool._runtime_state.context_window_tokens == 131072
assert loop.context_window_tokens == 131072
loop._sync_replay_max_messages.assert_called_once_with()
@pytest.mark.asyncio
async def test_modify_none_value_for_restricted_int(self):