From 9ef1e292eab797477ba71fbfec6b06fc654443a9 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 19 Aug 2026 18:03:03 +0800 Subject: [PATCH] fix(session): preserve complete transcripts --- docs/python-sdk.md | 2 +- nanobot/agent/loop.py | 20 +- nanobot/agent/memory.py | 86 +------ nanobot/sdk/clients.py | 6 +- nanobot/session/manager.py | 94 ++------ tests/agent/test_auto_compact.py | 54 +---- tests/agent/test_consolidator.py | 103 +------- tests/agent/test_history_replay.py | 120 ++++++++++ tests/agent/test_loop_consolidation_tokens.py | 2 - tests/agent/test_max_messages_config.py | 221 ------------------ tests/agent/test_runtime_refresh.py | 1 - tests/agent/test_session_manager_history.py | 124 ---------- tests/agent/test_session_retention.py | 34 --- tests/session/test_session_store.py | 48 +--- tests/test_nanobot_facade.py | 15 +- 15 files changed, 160 insertions(+), 770 deletions(-) create mode 100644 tests/agent/test_history_replay.py delete mode 100644 tests/agent/test_max_messages_config.py diff --git a/docs/python-sdk.md b/docs/python-sdk.md index 508ad5f3d..0a4271db2 100644 --- a/docs/python-sdk.md +++ b/docs/python-sdk.md @@ -634,7 +634,7 @@ Do not expose exported snapshots directly to chat users. | `workspace` | Current runtime workspace path. | | `add_context_provider(provider)` | Register an async per-turn context provider and return an unsubscribe callback. | | `on_session_turn_persisted(handler)` | Register a best-effort sync or async callback for locally persisted turns and return an unsubscribe callback. | -| `await compact_session(session_key)` | Run token/replay-window consolidation for a session. | +| `await compact_session(session_key)` | Run token-based consolidation for a session. | | `await compact_idle_session(session_key, max_suffix=8)` | Run idle-session compaction and return its summary. | ### Host integration context and persisted-turn callbacks diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 8d298e2e0..184cb4a29 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -14,7 +14,6 @@ from collections.abc import Coroutine, Iterable, Mapping from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress from dataclasses import dataclass, field from enum import Enum, auto -from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Awaitable, Callable, TypeVar, cast @@ -75,12 +74,7 @@ from nanobot.session.goal_state import ( ) from nanobot.session.history_visibility import HIDDEN_HISTORY_META from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel -from nanobot.session.manager import ( - SESSION_CACHE_MAX_SIZE, - Session, - SessionManager, - replay_max_messages_for_context, -) +from nanobot.session.manager import SESSION_CACHE_MAX_SIZE, Session, SessionManager from nanobot.session.model_selection import ( SESSION_MODEL_PRESET_METADATA_KEY, model_preset_from_metadata, @@ -386,7 +380,6 @@ class AgentLoop: # WebUI and fork rollback paths. Observe that boundary once instead of # duplicating cleanup in each consumer. self.sessions.set_delete_observer(self._file_state_store.discard) - self.sessions.set_file_cap_archiver(self.context.memory.raw_archive) self.tools = tool_registry if tool_registry is not None else ToolRegistry() self._exec_session_manager = ExecSessionManager() self.runner = AgentRunner() @@ -1819,14 +1812,10 @@ class AgentLoop: ) if ctx.on_runtime_admitted is not None: await ctx.on_runtime_admitted(runtime) - replay_max_messages = replay_max_messages_for_context( - runtime.context_window_tokens - ) if not ctx.ephemeral: await self.consolidator.maybe_consolidate_by_tokens( session, runtime=runtime, - replay_max_messages=replay_max_messages, ) is_subagent = ctx.kind is TurnKind.SYSTEM and ctx.msg.sender_id == "subagent" @@ -1835,7 +1824,6 @@ class AgentLoop: message_tool.start_turn() _hist_kwargs: dict[str, Any] = { - "max_messages": replay_max_messages, "max_tokens": self._replay_token_budget(runtime), "extend_to_user": is_subagent, } @@ -1999,16 +1987,10 @@ class AgentLoop: ) ctx.delivery.record_latency(ctx.turn_latency_ms) if not ctx.ephemeral: - session.enforce_file_cap( - on_archive=partial(self.context.memory.raw_archive, session_key=ctx.session_key) - ) self.schedule_background( self.consolidator.maybe_consolidate_by_tokens( session, runtime=runtime, - replay_max_messages=replay_max_messages_for_context( - runtime.context_window_tokens - ), ) ) self._clear_pending_user_turn(session) diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index a736335b5..056a24ca7 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -25,7 +25,6 @@ from nanobot.session.manager import ( MIN_COMPACTED_REPLAY_MESSAGES, Session, SessionManager, - replay_max_messages_for_context, ) from nanobot.session.summary import session_summary_from_metadata from nanobot.utils.gitstore import GitStore @@ -34,8 +33,6 @@ from nanobot.utils.helpers import ( ensure_dir, estimate_message_tokens, estimate_prompt_tokens_chain, - find_legal_message_start, - recent_message_start_index, strip_think, truncate_text, ) @@ -868,74 +865,7 @@ class Consolidator: """Return all messages that can reach the next model prompt.""" if not session.messages: return [] - return session.get_history(max_messages=len(session.messages)) - - @staticmethod - def _replay_overflow_boundary( - session: Session, - replay_max_messages: int | None, - ) -> int | None: - if not replay_max_messages or replay_max_messages <= 0: - return None - tail = list(enumerate(session.messages[session.last_consolidated:], session.last_consolidated)) - if len(tail) <= replay_max_messages: - return None - - tail_messages = [message for _idx, message in tail] - start_idx = recent_message_start_index( - tail_messages, - replay_max_messages, - extend_to_user=True, - ) - sliced = tail[start_idx:] - for i, (_idx, message) in enumerate(sliced): - if message.get("role") == "user": - start = i - if i > 0 and sliced[i - 1][1].get("_channel_delivery"): - start = i - 1 - sliced = sliced[start:] - break - - legal_start = find_legal_message_start([message for _idx, message in sliced]) - if legal_start: - sliced = sliced[legal_start:] - if not sliced: - return len(session.messages) - - first_visible_idx = sliced[0][0] - if first_visible_idx <= session.last_consolidated: - return None - return first_visible_idx - - async def _consolidate_replay_overflow( - self, - session: Session, - replay_max_messages: int | None, - *, - runtime: LLMRuntime, - ) -> str | None: - """Archive messages that would be hidden by the replay message window.""" - end_idx = self._replay_overflow_boundary(session, replay_max_messages) - if end_idx is None: - return None - chunk = session.messages[session.last_consolidated:end_idx] - if not chunk: - return None - logger.info( - "Replay-window consolidation for {}: chunk={} msgs, replay_max={}", - session.key, - len(chunk), - replay_max_messages, - ) - summary = await self.archive_session( - session, - archive_end=end_idx, - runtime=runtime, - ) - session.last_consolidated = end_idx - session.provider_state = None - self.sessions.save(session) - return summary + return session.get_history() def _persist_last_summary(self, session: Session, summary: str | None) -> None: if summary and summary != "(nothing)": @@ -1056,14 +986,11 @@ class Consolidator: messages=list(session.messages[:archive_end]), last_consolidated=session.last_consolidated, ) - history = prefix.get_history( - max_messages=replay_max_messages_for_context(runtime.context_window_tokens), - max_tokens=budget, - ) + history = prefix.get_history(max_tokens=budget) archive_history = Session( key=session.key, messages=messages, - ).get_history(max_messages=len(messages)) + ).get_history() if ( not archive_history or history[-len(archive_history):] != archive_history @@ -1125,7 +1052,6 @@ class Consolidator: session: Session, *, runtime: LLMRuntime, - replay_max_messages: int | None = None, ) -> None: """Loop: archive old messages until prompt fits within safe budget. @@ -1146,11 +1072,7 @@ class Consolidator: budget = self._input_token_budget(runtime) target = int(budget * self.consolidation_ratio) - last_summary = await self._consolidate_replay_overflow( - session, - replay_max_messages, - runtime=runtime, - ) + last_summary: str | None = None estimated, source = self.estimate_session_prompt_tokens( session, runtime=runtime, diff --git a/nanobot/sdk/clients.py b/nanobot/sdk/clients.py index afd975764..08ff1b9db 100644 --- a/nanobot/sdk/clients.py +++ b/nanobot/sdk/clients.py @@ -15,7 +15,6 @@ from nanobot.sdk.types import ( snapshot_from_payload, snapshot_from_session, ) -from nanobot.session.manager import replay_max_messages_for_context if TYPE_CHECKING: from nanobot.agent.loop import AgentLoop @@ -210,15 +209,12 @@ class RuntimeClient: return self._loop.runtime_events.subscribe(handler, SessionTurnPersisted) async def compact_session(self, session_key: str) -> SessionSnapshot: - """Run token/replay-window consolidation for one session.""" + """Run token consolidation for one session.""" session = self._loop.sessions.get_or_create(session_key) runtime = self._loop.runtime_for_session(session) await self._loop.consolidator.maybe_consolidate_by_tokens( session, runtime=runtime, - replay_max_messages=replay_max_messages_for_context( - runtime.context_window_tokens - ), ) return snapshot_from_session(self._loop.sessions.get_or_create(session_key)) diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index d5c7e3c4a..e7d07c433 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -39,11 +39,8 @@ from nanobot.utils.helpers import ( ) from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body -FILE_MAX_MESSAGES = 2000 SESSION_CACHE_MAX_SIZE = 128 -MIN_REPLAY_MAX_MESSAGES = 120 MIN_COMPACTED_REPLAY_MESSAGES = 8 -REPLAY_TOKENS_PER_MESSAGE = 100 _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*$') @@ -84,15 +81,6 @@ def _is_provider_state_record_line(line: str) -> bool: return _PROVIDER_STATE_RECORD_PREFIX_RE.match(line) is not None -def replay_max_messages_for_context(context_window_tokens: int | None) -> int: - if not context_window_tokens or context_window_tokens <= 0: - return FILE_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. @@ -209,7 +197,7 @@ class Session: def get_history( self, - max_messages: int = FILE_MAX_MESSAGES, + max_messages: int = 0, *, max_tokens: int = 0, extend_to_user: bool = False, @@ -217,8 +205,8 @@ class Session: ) -> list[dict[str, Any]]: """Return recent replayable messages for LLM input. - History is sliced by message count first (``max_messages``), then by - token budget from the tail (``max_tokens``) when provided. + A positive ``max_messages`` applies an explicit caller-owned count + limit. The normal model path relies on ``max_tokens`` instead. """ replay_start = self.last_consolidated if replay_start: @@ -233,18 +221,20 @@ class Session: replay_start = min(replay_start, recent_start) replayable = self.messages[replay_start:] - max_messages = max_messages if max_messages > 0 else FILE_MAX_MESSAGES - unarchived_count = len(self.messages) - self.last_consolidated - if replay_start < self.last_consolidated and unarchived_count < max_messages: - # The archived replay suffix can exceed the nominal count when one - # tool-heavy turn spans the boundary. Preserve that complete turn. + if max_messages <= 0: start_idx = 0 else: - start_idx = recent_message_start_index( - replayable, - max_messages, - extend_to_user=extend_to_user, - ) + unarchived_count = len(self.messages) - self.last_consolidated + if replay_start < self.last_consolidated and unarchived_count < max_messages: + # The archived replay suffix can exceed the nominal count when one + # tool-heavy turn spans the boundary. Preserve that complete turn. + start_idx = 0 + else: + start_idx = recent_message_start_index( + replayable, + max_messages, + extend_to_user=extend_to_user, + ) sliced = replayable[start_idx:] # Avoid starting mid-turn when possible, except for proactive @@ -467,46 +457,6 @@ class Session: already_consolidated_count=already_consolidated, ) - def enforce_file_cap( - self, - on_archive: Callable[[list[dict[str, Any]]], None] | None = None, - limit: int = FILE_MAX_MESSAGES, - ) -> None: - """Bound session message growth by archiving and trimming old prefixes.""" - if limit <= 0 or len(self.messages) <= limit: - return - - original_messages = self.messages - original_last_consolidated = self.last_consolidated - original_provider_state = self.provider_state - original_updated_at = self.updated_at - result = self.retain_recent_legal_suffix(limit) - if not result.dropped: - return - - archive_chunk = result.dropped[result.already_consolidated_count:] - if archive_chunk and on_archive: - try: - on_archive(archive_chunk) - except BaseException: - # Retention runs before the archive callback so the callback can - # receive the exact dropped prefix. Restore the in-memory session - # if archival fails; otherwise a later save would persist the - # trimmed state and make that prefix impossible to retry. - self.messages = original_messages - self.last_consolidated = original_last_consolidated - self.provider_state = original_provider_state - self.updated_at = original_updated_at - raise - logger.info( - "Session file cap hit for {}: dropped {}, raw-archived {}, kept {}", - self.key, - len(result.dropped), - len(archive_chunk), - len(self.messages), - ) - - class SessionPayload(TypedDict): key: str created_at: str | None @@ -1576,7 +1526,6 @@ class SessionManager: # Preserve identity for sessions held by active callers without retaining idle ones. self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary() self._max_cached_sessions = SESSION_CACHE_MAX_SIZE - self._file_cap_archiver: Callable[..., None] | None = None self._delete_observer: Callable[[str], None] | None = None def _remember(self, session: Session) -> None: @@ -1603,10 +1552,6 @@ class SessionManager: """Return a cached session without creating or loading one from disk.""" return self._cached(key) - def set_file_cap_archiver(self, archiver: Callable[..., None]) -> None: - """Archive unconsolidated overflow whenever a session is persisted.""" - self._file_cap_archiver = archiver - def set_delete_observer(self, observer: Callable[[str], None]) -> None: """Observe explicit session deletion for process-local state cleanup.""" self._delete_observer = observer @@ -1705,15 +1650,6 @@ class SessionManager: if not session.policy.persist: return - archiver = self._file_cap_archiver - if archiver is not None: - session.enforce_file_cap( - on_archive=lambda messages: archiver( - messages, - session_key=session.key, - ) - ) - self._store.save(session, fsync=fsync) self._remember(session) diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index 6ce29787d..9013f142f 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -171,12 +171,6 @@ class TestSessionTTLConfig: data = defaults.model_dump(mode="json", by_alias=True) assert data["idleCompactCheckIntervalSeconds"] == 10 - def test_session_file_cap_is_internal_constant(self): - """Session file cap should remain an internal constant, not a config field.""" - from nanobot.session.manager import FILE_MAX_MESSAGES - assert FILE_MAX_MESSAGES == 2000 - - class TestIdleScanThrottling: """Test scheduling of full idle-session scans.""" @@ -255,53 +249,7 @@ class TestAgentLoopTTLParam: kwargs = session.get_history.call_args.kwargs assert isinstance(kwargs.get("max_tokens"), int) assert kwargs["max_tokens"] > 0 - assert set(kwargs) == {"max_messages", "max_tokens", "extend_to_user"} - - @pytest.mark.asyncio - async def test_session_file_cap_archives_and_trims_old_messages(self, tmp_path): - loop = _make_loop(tmp_path) - loop.context.memory.raw_archive = MagicMock() - - for i in range(4): - msg = InboundMessage( - channel="cli", - sender_id="u1", - chat_id="direct", - content=f"hello {i}", - ) - await loop._process_message(msg) - - session = loop.sessions.get_or_create("cli:direct") - from nanobot.session.manager import FILE_MAX_MESSAGES - assert len(session.messages) <= FILE_MAX_MESSAGES - - def test_session_enforce_file_cap_skips_archive_when_dropped_prefix_already_consolidated(self, tmp_path): - from nanobot.session.manager import Session - archive_fn = MagicMock() - session = Session(key="cli:direct") - for i in range(8): - session.add_message("user", f"u{i}") - session.last_consolidated = 6 - - session.enforce_file_cap(on_archive=archive_fn, limit=4) - - assert len(session.messages) <= 4 - archive_fn.assert_not_called() - - def test_session_enforce_file_cap_archives_only_unconsolidated_dropped_prefix(self, tmp_path): - from nanobot.session.manager import Session - archive_fn = MagicMock() - session = Session(key="cli:direct") - for i in range(8): - session.add_message("user", f"u{i}") - session.last_consolidated = 2 - - session.enforce_file_cap(on_archive=archive_fn, limit=4) - - assert len(session.messages) <= 4 - archive_fn.assert_called_once() - archived = archive_fn.call_args.args[0] - assert [m["content"] for m in archived] == ["u2", "u3"] + assert set(kwargs) == {"max_tokens", "extend_to_user"} class TestAutoCompact: diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index 83558c09e..1b88a738a 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -363,7 +363,7 @@ class TestConsolidatorTokenBudget: await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) async def test_estimate_uses_full_unconsolidated_tail(self, consolidator, runtime): - """Consolidation pressure must see messages hidden by the replay window.""" + """Consolidation pressure must account for the full unarchived tail.""" session = Session(key="test:full-tail") for i in range(160): session.add_message("user", f"msg-{i}") @@ -400,107 +400,6 @@ class TestConsolidatorTokenBudget: assert len(captured["history"]) == 8 assert captured["history"][0]["content"] == "msg-2" - async def test_replay_window_overflow_is_archived_even_under_token_budget( - self, - consolidator, - runtime, - ): - """Old messages that cannot be replayed should be materialized first.""" - consolidator._SAFETY_BUFFER = 0 - session = Session(key="test:replay-overflow") - session.provider_state = _provider_state() - for i in range(10): - session.add_message("user", f"u{i}") - session.add_message("assistant", f"a{i}") - - consolidator.sessions._session_cache[session.key] = session - consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken")) - consolidator.archive_session = AsyncMock(return_value="old conversation summary") - - await consolidator.maybe_consolidate_by_tokens( - session, - runtime=runtime, - replay_max_messages=6, - ) - - archive_end = consolidator.archive_session.await_args.kwargs["archive_end"] - archived_chunk = session.messages[:archive_end] - assert archived_chunk[0]["content"] == "u0" - assert archived_chunk[-1]["content"] == "a6" - assert session.last_consolidated == 14 - assert session.metadata["_last_summary"]["text"] == "old conversation summary" - assert session.provider_state is None - consolidator.sessions.save.assert_called() - - async def test_replay_window_overflow_extends_to_long_recent_user_turn( - self, - consolidator, - runtime, - ): - """Replay-window consolidation must not cut into the latest user turn.""" - session = Session(key="test:replay-tool-boundary") - session.add_message("user", "old") - session.add_message("assistant", "old answer") - session.add_message("user", "record this") - for i in range(4): - session.messages.extend(_tool_round(f"call-{i}")) - session.add_message("assistant", "final answer") - - consolidator.sessions._session_cache[session.key] = session - consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken")) - consolidator.archive_session = AsyncMock(return_value="tool turn summary") - - await consolidator.maybe_consolidate_by_tokens( - session, - runtime=runtime, - replay_max_messages=4, - ) - - archive_end = consolidator.archive_session.await_args.kwargs["archive_end"] - archived_chunk = session.messages[:archive_end] - assert [m["content"] for m in archived_chunk] == ["old", "old answer"] - assert session.last_consolidated == 2 - - history = session.get_history(max_messages=4, extend_to_user=True) - assert len(history) > 4 - assert history[0]["content"] == "record this" - assert history[-1]["content"] == "final answer" - - async def test_replay_window_overflow_uses_newer_user_inside_window( - self, - consolidator, - runtime, - ): - """Do not extend to an older long turn when the hard window has a newer user.""" - session = Session(key="test:replay-newer-user") - session.add_message("user", "old") - session.add_message("assistant", "old answer") - session.add_message("user", "long older turn") - for i in range(8): - session.messages.extend(_tool_round(f"older-{i}")) - session.add_message("assistant", "older final") - session.add_message("user", "new question") - session.add_message("assistant", "new answer") - - consolidator.sessions._session_cache[session.key] = session - consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken")) - consolidator.archive_session = AsyncMock(return_value="older turn summary") - - await consolidator.maybe_consolidate_by_tokens( - session, - runtime=runtime, - replay_max_messages=6, - ) - - archive_end = consolidator.archive_session.await_args.kwargs["archive_end"] - archived_chunk = session.messages[:archive_end] - assert archived_chunk[2]["content"] == "long older turn" - assert archived_chunk[-1]["content"] == "older final" - assert session.last_consolidated == len(session.messages) - 2 - - history = session.get_history(max_messages=6, extend_to_user=True) - assert [m["content"] for m in history] == ["new question", "new answer"] - async def test_token_overflow_appends_prompt_to_replay_prefix( self, consolidator, diff --git a/tests/agent/test_history_replay.py b/tests/agent/test_history_replay.py new file mode 100644 index 000000000..04372c3ca --- /dev/null +++ b/tests/agent/test_history_replay.py @@ -0,0 +1,120 @@ +"""Tests for token-bounded session history replay.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +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 Session + + +def _make_loop(tmp_path: Path, 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", + context_window_tokens=context_window_tokens, + ) + + +def _populated_session(turns: int) -> Session: + session = Session(key="test:populated") + for index in range(turns): + session.add_message("user", f"msg-{index}") + session.add_message("assistant", f"reply-{index}") + return session + + +def _tool_round(call_id: str) -> list[dict]: + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": call_id, "type": "function", "function": {"name": "x", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": call_id, "name": "x", "content": "ok"}, + ] + + +def test_default_history_has_no_message_count_limit() -> None: + session = _populated_session(1_001) + + history = session.get_history() + + assert len(history) == 2_002 + assert history[0]["content"] == "msg-0" + assert history[-1]["content"] == "reply-1000" + + +def test_explicit_message_limit_still_starts_at_user_turn() -> None: + history = _populated_session(30).get_history(max_messages=25) + + assert len(history) <= 25 + assert history[0]["role"] == "user" + + +@pytest.mark.asyncio +async def test_process_message_replays_with_token_budget_only(tmp_path: Path) -> None: + loop = _make_loop(tmp_path, context_window_tokens=32_768) + loop.provider.chat_with_retry = AsyncMock( + return_value=LLMResponse(content="ok", tool_calls=[], usage={}) + ) + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + + session = loop.sessions.get_or_create("cli:test") + with patch.object(session, "get_history", wraps=session.get_history) as get_history: + result = await loop._process_message( + InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello") + ) + + assert result is not None + assert get_history.call_args.kwargs == { + "max_tokens": loop._replay_token_budget(loop.llm_runtime()), + "extend_to_user": False, + } + + +@pytest.mark.asyncio +async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None: + loop = _make_loop(tmp_path, context_window_tokens=8_000) + loop.provider.chat_with_retry = AsyncMock( + return_value=LLMResponse(content="ok", tool_calls=[], usage={}) + ) + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + + session = loop.sessions.get_or_create("cli:test") + session.add_message("user", "old") + session.add_message("assistant", "old answer") + session.add_message("user", "long older turn") + for index in range(70): + session.messages.extend(_tool_round(f"older-{index}")) + session.add_message("assistant", "older final") + + result = await loop._process_message( + InboundMessage( + channel="cli", + sender_id="user", + chat_id="test", + content="new question", + ) + ) + + assert result is not None + sent_messages = loop.provider.chat_with_retry.await_args.kwargs["messages"] + sent_text = "\n".join(str(message.get("content")) for message in sent_messages) + assert "new question" in sent_text + assert "long older turn" not in sent_text diff --git a/tests/agent/test_loop_consolidation_tokens.py b/tests/agent/test_loop_consolidation_tokens.py index b30957e03..6cb47a85d 100644 --- a/tests/agent/test_loop_consolidation_tokens.py +++ b/tests/agent/test_loop_consolidation_tokens.py @@ -6,7 +6,6 @@ import nanobot.agent.memory as memory_module from nanobot.agent.loop import AgentLoop from nanobot.bus.queue import MessageBus from nanobot.providers.base import LLMResponse -from nanobot.session.manager import replay_max_messages_for_context def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -> AgentLoop: @@ -227,7 +226,6 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non loop.consolidator.maybe_consolidate_by_tokens.assert_any_await( session, runtime=runtime, - replay_max_messages=replay_max_messages_for_context(runtime.context_window_tokens), ) assert len(loop.consolidator.maybe_consolidate_by_tokens.call_args_list) == 2 assert all( diff --git a/tests/agent/test_max_messages_config.py b/tests/agent/test_max_messages_config.py deleted file mode 100644 index c00051d87..000000000 --- a/tests/agent/test_max_messages_config.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Tests for the internal max_messages replay cap.""" - -from __future__ import annotations - -from dataclasses import replace -from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -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.providers.factory import ProviderSnapshot -from nanobot.session.manager import ( - FILE_MAX_MESSAGES, - Session, - replay_max_messages_for_context, -) - - -def _make_loop( - tmp_path: Path, - 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", - context_window_tokens=context_window_tokens, - ) - - -def _populated_session(n: int) -> Session: - """Create a session with *n* user/assistant turn pairs.""" - session = Session(key="test:populated") - for i in range(n): - session.add_message("user", f"msg-{i}") - session.add_message("assistant", f"reply-{i}") - return session - - -def _tool_round(call_id: str) -> list[dict]: - return [ - { - "role": "assistant", - "content": None, - "tool_calls": [ - {"id": call_id, "type": "function", "function": {"name": "x", "arguments": "{}"}} - ], - }, - {"role": "tool", "tool_call_id": call_id, "name": "x", "content": "ok"}, - ] - - -class TestMaxMessagesInit: - """Verify AgentLoop derives the internal replay cap correctly.""" - - 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) == FILE_MAX_MESSAGES - - def test_default_for_200k_context_reaches_file_cap(self, tmp_path: Path) -> None: - loop = _make_loop(tmp_path) - runtime = loop.runtime_resolver.runtime - assert replay_max_messages_for_context(runtime.context_window_tokens) == FILE_MAX_MESSAGES - - def test_default_scales_with_context_window(self, tmp_path: Path) -> None: - loop = _make_loop(tmp_path, context_window_tokens=32_768) - runtime = loop.runtime_resolver.runtime - assert replay_max_messages_for_context(runtime.context_window_tokens) == 327 - - 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",), - ), - ) - - initial = loop.runtime_resolver.runtime - assert replay_max_messages_for_context(initial.context_window_tokens) == 327 - loop.runtime_resolver.invalidate() - refreshed = loop.llm_runtime() - assert replay_max_messages_for_context(refreshed.context_window_tokens) == FILE_MAX_MESSAGES - - -class TestGetHistoryWithMaxMessages: - """Verify get_history respects max_messages parameter.""" - - def test_default_uses_builtin_limit(self) -> None: - session = _populated_session(80) - history = session.get_history() - assert len(history) <= FILE_MAX_MESSAGES - - def test_explicit_max_messages_limits_output(self) -> None: - session = _populated_session(40) # 80 messages total - history = session.get_history(max_messages=20) - assert len(history) <= 20 - - def test_max_messages_starts_at_user_turn(self) -> None: - """Sliced history should start with a user message, not mid-turn.""" - session = _populated_session(30) # 60 messages - history = session.get_history(max_messages=25) - assert history[0]["role"] == "user" - - def test_max_messages_zero_uses_builtin_limit(self) -> None: - session = _populated_session(80) # 160 messages total - history = session.get_history(max_messages=0) - assert len(history) <= FILE_MAX_MESSAGES - - def test_small_session_unaffected(self) -> None: - """When session has fewer messages than max_messages, all are returned.""" - session = _populated_session(5) # 10 messages - history = session.get_history(max_messages=25) - assert len(history) == 10 - - -class TestMaxMessagesIntegration: - """Verify AgentLoop passes the replay cap into get_history calls.""" - - @pytest.mark.asyncio - async def test_process_message_passes_limit_to_history_call(self, tmp_path: Path) -> None: - """The real message path should pass max_messages into session history replay.""" - loop = _make_loop(tmp_path) - runtime = replace(loop.llm_runtime(), context_window_tokens=32_768) - loop.provider.chat_with_retry = AsyncMock( - return_value=LLMResponse(content="ok", tool_calls=[], usage={}) - ) - loop.tools.get_definitions = MagicMock(return_value=[]) - loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] - - session = loop.sessions.get_or_create("cli:test") - with patch.object(session, "get_history", wraps=session.get_history) as mock_hist: - result = await loop._process_message( - InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello"), - runtime=runtime, - ) - - assert result is not None - assert mock_hist.call_count == 1 - assert mock_hist.call_args.kwargs["max_messages"] == 327 - assert mock_hist.call_args.kwargs["extend_to_user"] is False - - @pytest.mark.asyncio - async def test_default_limit_passes_context_derived_limit_to_history_call( - self, - tmp_path: Path, - ) -> None: - loop = _make_loop(tmp_path) - loop.provider.chat_with_retry = AsyncMock( - return_value=LLMResponse(content="ok", tool_calls=[], usage={}) - ) - loop.tools.get_definitions = MagicMock(return_value=[]) - loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] - - session = loop.sessions.get_or_create("cli:test") - with patch.object(session, "get_history", wraps=session.get_history) as mock_hist: - result = await loop._process_message( - InboundMessage(channel="cli", sender_id="user", chat_id="test", content="hello") - ) - - assert result is not None - assert mock_hist.call_args.kwargs["max_messages"] == FILE_MAX_MESSAGES - assert mock_hist.call_args.kwargs["extend_to_user"] is False - - @pytest.mark.asyncio - async def test_process_message_uses_current_user_as_replay_boundary( - self, - tmp_path: Path, - ) -> None: - """A live user turn should not extend history to an older long tool turn.""" - loop = _make_loop(tmp_path, context_window_tokens=8_000) - loop.provider.chat_with_retry = AsyncMock( - return_value=LLMResponse(content="ok", tool_calls=[], usage={}) - ) - loop.tools.get_definitions = MagicMock(return_value=[]) - loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] - - session = loop.sessions.get_or_create("cli:test") - session.add_message("user", "old") - session.add_message("assistant", "old answer") - session.add_message("user", "long older turn") - for i in range(70): - session.messages.extend(_tool_round(f"older-{i}")) - session.add_message("assistant", "older final") - - with patch.object(session, "get_history", wraps=session.get_history) as mock_hist: - result = await loop._process_message( - InboundMessage( - channel="cli", - sender_id="user", - chat_id="test", - content="new question", - ) - ) - - assert result is not None - assert mock_hist.call_args.kwargs["extend_to_user"] is False - sent_messages = loop.provider.chat_with_retry.await_args.kwargs["messages"] - sent_text = "\n".join(str(message.get("content")) for message in sent_messages) - assert "new question" in sent_text - assert "long older turn" not in sent_text diff --git a/tests/agent/test_runtime_refresh.py b/tests/agent/test_runtime_refresh.py index 3ed91e3d0..0174d3873 100644 --- a/tests/agent/test_runtime_refresh.py +++ b/tests/agent/test_runtime_refresh.py @@ -83,7 +83,6 @@ def test_loop_has_no_mutable_runtime_mirrors_or_legacy_snapshot_api(tmp_path: Pa }.isdisjoint(loop.__dict__) assert not hasattr(loop, "_apply_provider_snapshot") assert not hasattr(loop, "_build_model_preset_snapshot") - assert not hasattr(loop, "_sync_replay_max_messages") def test_llm_runtime_refreshes_provider_snapshot(tmp_path: Path) -> None: diff --git a/tests/agent/test_session_manager_history.py b/tests/agent/test_session_manager_history.py index 8cf75f791..07e7bd62d 100644 --- a/tests/agent/test_session_manager_history.py +++ b/tests/agent/test_session_manager_history.py @@ -1,5 +1,3 @@ -import pytest - from nanobot.providers.base import ProviderConversationState from nanobot.runtime_context import ( RUNTIME_CONTEXT_HISTORY_META, @@ -832,9 +830,6 @@ def test_get_history_extend_to_user_keeps_newer_user_inside_window(): _assert_no_orphans(history) -# --- enforce_file_cap archive correctness (issue #4128) --- - - def test_retain_recent_legal_suffix_returns_dropped_messages(): """retain_recent_legal_suffix returns the actually-dropped messages.""" session = Session( @@ -894,125 +889,6 @@ def test_retain_recent_legal_suffix_returns_all_on_zero(): assert session.messages == [] -def test_enforce_file_cap_no_duplicate_archive_in_else_branch(): - """When the tail is assistant-only, enforce_file_cap must not archive - messages that are also retained (the bug from issue #4128).""" - from unittest.mock import MagicMock - - session = Session(key="test:else-archive") - # Build: 15 user messages, then 10 assistant messages (no user in tail) - for i in range(15): - session.messages.append({"role": "user", "content": f"u{i}"}) - for i in range(10): - session.messages.append({"role": "assistant", "content": f"a{i}"}) - - archive_fn = MagicMock() - session.enforce_file_cap(on_archive=archive_fn, limit=6) - - assert len(session.messages) <= 6 - - # Verify archived messages have NO overlap with retained - if archive_fn.called: - archived = archive_fn.call_args.args[0] - archived_ids = set(id(m) for m in archived) - retained_ids = set(id(m) for m in session.messages) - assert not archived_ids & retained_ids, ( - f"Duplicate messages in archive and retained: " - f"overlap contents = {[m['content'] for m in archived if id(m) in retained_ids]}" - ) - - -def test_enforce_file_cap_no_message_loss_in_else_branch(): - """In the else branch, no messages should silently disappear — every - message must be either retained or archived.""" - from unittest.mock import MagicMock - - session = Session(key="test:else-no-loss") - all_messages = [] - for i in range(15): - msg = {"role": "user", "content": f"u{i}"} - session.messages.append(msg) - all_messages.append(msg) - for i in range(10): - msg = {"role": "assistant", "content": f"a{i}"} - session.messages.append(msg) - all_messages.append(msg) - - archive_fn = MagicMock() - session.enforce_file_cap(on_archive=archive_fn, limit=6) - - # Collect all messages accounted for (retained + archived) - accounted = set(id(m) for m in session.messages) - if archive_fn.called: - for m in archive_fn.call_args.args[0]: - accounted.add(id(m)) - - all_ids = set(id(m) for m in all_messages) - missing = all_ids - accounted - assert not missing, ( - f"Lost {len(missing)} message(s) — neither retained nor archived" - ) - - -def test_enforce_file_cap_correct_archive_with_last_consolidated_in_else_branch(): - """When last_consolidated > 0 and the else branch fires, only the - unconsolidated dropped messages should be raw-archived. Messages in the - consolidated prefix that are dropped do NOT need raw archiving.""" - from unittest.mock import MagicMock - - session = Session(key="test:else-lc-archive") - # 20 messages total: u0..u9 (user), a0..a9 (assistant) - for i in range(10): - session.messages.append({"role": "user", "content": f"u{i}"}) - for i in range(10): - session.messages.append({"role": "assistant", "content": f"a{i}"}) - # First 8 messages already consolidated - session.last_consolidated = 8 - - archive_fn = MagicMock() - session.enforce_file_cap(on_archive=archive_fn, limit=4) - - if archive_fn.called: - archived = archive_fn.call_args.args[0] - # Archived messages should NOT include any from the consolidated prefix - # (u0..u7). They should only be unconsolidated dropped messages. - archived_contents = [m["content"] for m in archived] - for c in archived_contents: - assert c not in [f"u{i}" for i in range(8)], ( - f"Consolidated message {c!r} should not be raw-archived" - ) - - -def test_enforce_file_cap_restores_session_when_archive_fails(): - state = ProviderConversationState( - kind="openai_responses", - provider="openai:test", - model="test-model", - version=1, - payload={"items": []}, - ) - session = Session(key="test:archive-failure", provider_state=state) - for i in range(8): - session.messages.append({"role": "user", "content": f"msg{i}"}) - original_messages = session.messages - original_updated_at = session.updated_at - session.last_consolidated = 2 - - def fail_archive(_messages): - raise RuntimeError("history unavailable") - - with pytest.raises(RuntimeError, match="history unavailable"): - session.enforce_file_cap(on_archive=fail_archive, limit=4) - - assert session.messages is original_messages - assert [message["content"] for message in session.messages] == [ - f"msg{i}" for i in range(8) - ] - assert session.last_consolidated == 2 - assert session.provider_state is state - assert session.updated_at == original_updated_at - - def test_retain_recent_legal_suffix_last_consolidated_correct_in_else_branch(): """last_consolidated after retain_recent_legal_suffix should reflect how many retained messages were inside the old consolidated prefix.""" diff --git a/tests/agent/test_session_retention.py b/tests/agent/test_session_retention.py index f35e23e41..f3dc7a531 100644 --- a/tests/agent/test_session_retention.py +++ b/tests/agent/test_session_retention.py @@ -167,40 +167,6 @@ def test_retain_drops_delivery_not_adjacent_to_anchor_user(): assert _contents(session.messages) == ["ok", "great"] -# --- Delivery preservation through the production entry points --- - - -def test_enforce_file_cap_keeps_delivery_in_session(): - session = Session(key="test:cap-delivery") - session.messages.append({"role": "user", "content": "setup"}) - session.messages.append(_delivery("Remember to drink water")) - session.messages.append({"role": "user", "content": "ok"}) - session.messages.append({"role": "assistant", "content": "great"}) - - archived: list[list[dict]] = [] - session.enforce_file_cap(on_archive=archived.append, limit=3) - - archived_flat = [m for chunk in archived for m in chunk] - assert _has_delivery(session.messages) - assert not any(m.get("_channel_delivery") for m in archived_flat) - - -def test_enforce_file_cap_archives_only_prefix(): - session = Session(key="test:cap-prefix") - session.messages.append({"role": "user", "content": "setup"}) - session.messages.append({"role": "assistant", "content": "first reply"}) - session.messages.append(_delivery("Remember to drink water")) - session.messages.append({"role": "user", "content": "ok"}) - session.messages.append({"role": "assistant", "content": "great"}) - - archived: list[list[dict]] = [] - session.enforce_file_cap(on_archive=archived.append, limit=3) - - archived_flat = [m for chunk in archived for m in chunk] - assert _has_delivery(session.messages) - assert _contents(archived_flat) == ["setup", "first reply"] - - def test_compact_probe_keeps_delivery_in_visible_suffix(): """compact_idle_session() trims a probe copy with extend_to_user=True; the visible suffix it keeps must still contain the delivery message.""" diff --git a/tests/session/test_session_store.py b/tests/session/test_session_store.py index 0d36c22d4..1c6d7c1ce 100644 --- a/tests/session/test_session_store.py +++ b/tests/session/test_session_store.py @@ -1,10 +1,8 @@ from unittest.mock import MagicMock -import pytest - import nanobot.session as session_api from nanobot.session import Session, SessionManager -from nanobot.session.manager import FILE_MAX_MESSAGES, SessionStore +from nanobot.session.manager import SessionStore from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY @@ -105,49 +103,23 @@ def test_read_session_snapshot_does_not_populate_runtime_cache(tmp_path) -> None store.load.assert_called_once_with(stored.key) -def test_manager_applies_file_cap_before_store_save(tmp_path) -> None: +def test_manager_preserves_full_session_before_store_save(tmp_path) -> None: store = MagicMock(spec=SessionStore) - archiver = MagicMock() manager = SessionManager(tmp_path, store=store) - manager.set_file_cap_archiver(archiver) session = Session( key="cli:large", messages=[ - {"role": "user", "content": str(index)} - for index in range(FILE_MAX_MESSAGES + 1) + { + "role": "user" if index % 2 == 0 else "assistant", + "content": str(index), + } + for index in range(2_001) ], ) manager.save(session) - assert len(session.messages) == FILE_MAX_MESSAGES - archiver.assert_called_once() - store.save.assert_called_once_with(session, fsync=False) - - -def test_manager_retries_file_cap_archive_after_failure(tmp_path) -> None: - store = MagicMock(spec=SessionStore) - archiver = MagicMock(side_effect=[RuntimeError("history unavailable"), None]) - manager = SessionManager(tmp_path, store=store) - manager.set_file_cap_archiver(archiver) - session = Session( - key="cli:retry-large", - messages=[ - {"role": "user", "content": str(index)} - for index in range(FILE_MAX_MESSAGES + 1) - ], - ) - - with pytest.raises(RuntimeError, match="history unavailable"): - manager.save(session) - - assert len(session.messages) == FILE_MAX_MESSAGES + 1 - store.save.assert_not_called() - - manager.save(session) - - assert len(session.messages) == FILE_MAX_MESSAGES - assert archiver.call_count == 2 - assert archiver.call_args_list[0].args[0][0]["content"] == "0" - assert archiver.call_args_list[1].args[0][0]["content"] == "0" + assert len(session.messages) == 2_001 + assert session.messages[0]["content"] == "0" + assert session.messages[-1]["content"] == "2000" store.save.assert_called_once_with(session, fsync=False) diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index d3ef4d233..d168849c0 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -35,7 +35,6 @@ from nanobot.runtime_context import ( RuntimeContextBlock, append_runtime_context, ) -from nanobot.session.manager import FILE_MAX_MESSAGES from nanobot.utils.llm_runtime import runtime_from_provider_snapshot @@ -1451,7 +1450,7 @@ async def test_sessions_ingest_imports_transcript_without_running_model(tmp_path @pytest.mark.asyncio -async def test_sessions_ingest_archives_overflow_at_persistence_boundary(tmp_path): +async def test_sessions_ingest_preserves_full_transcript(tmp_path): config_path = _write_config(tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path) @@ -1459,16 +1458,14 @@ async def test_sessions_ingest_archives_overflow_at_persistence_boundary(tmp_pat "sdk:overflow", [ {"role": "user", "content": f"message-{index}"} - for index in range(FILE_MAX_MESSAGES + 1) + for index in range(2_001) ], ) - assert len(snapshot.messages) == FILE_MAX_MESSAGES - assert snapshot.messages[0]["content"] == "message-1" - history = bot.memory.read_history(session_key="sdk:overflow") - assert len(history) == 1 - assert "[RAW] 1 messages" in history[0]["content"] - assert "message-0" in history[0]["content"] + assert len(snapshot.messages) == 2_001 + assert snapshot.messages[0]["content"] == "message-0" + assert snapshot.messages[-1]["content"] == "message-2000" + assert bot.memory.read_history(session_key="sdk:overflow") == [] @pytest.mark.asyncio