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

fix(session): preserve complete transcripts

Browse Source
This commit is contained in:
chengyongru
2026-08-19 18:40:20 +08:00
committed by chengyongru
parent 3c41d5e7f3
commit 9ef1e292ea
15 changed files with 160 additions and 770 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/python-sdk.md
+1 -1
View File
@@ -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
nanobot/agent/loop.py
+1 -19
View File
@@ -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)
nanobot/agent/memory.py
+4 -82
View File
@@ -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,
nanobot/sdk/clients.py
+1 -5
View File
@@ -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))
nanobot/session/manager.py
+15 -79
View File
@@ -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)
tests/agent/test_auto_compact.py
+1 -53
View File
@@ -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:
tests/agent/test_consolidator.py
+1 -102
View File
@@ -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,
tests/agent/test_history_replay.py
+120
View File
@@ -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
tests/agent/test_loop_consolidation_tokens.py
-2
View File
@@ -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(
tests/agent/test_max_messages_config.py
-221
View File
@@ -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
tests/agent/test_runtime_refresh.py
-1
View File
@@ -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:
tests/agent/test_session_manager_history.py
-124
View File
@@ -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."""
tests/agent/test_session_retention.py
-34
View File
@@ -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."""
tests/session/test_session_store.py
+10 -38
View File
@@ -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)
tests/test_nanobot_facade.py
+6 -9
View File
@@ -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