From c9d3e743423903dd1bc4ad7fd779c9ddf03ab224 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Fri, 10 Jul 2026 14:52:51 +0800 Subject: [PATCH] refactor(agent): require runtime for consolidation --- nanobot/agent/autocompact.py | 18 +- nanobot/agent/loop.py | 17 +- nanobot/agent/memory.py | 79 ++--- nanobot/command/builtin.py | 23 +- nanobot/command/router.py | 2 + nanobot/sdk/clients.py | 9 +- tests/agent/test_auto_compact.py | 53 ++-- tests/agent/test_autocompact_unit.py | 63 +++- tests/agent/test_consolidation_ratio.py | 10 +- tests/agent/test_consolidator.py | 273 +++++++++++++----- tests/agent/test_loop_consolidation_tokens.py | 39 ++- tests/agent/test_loop_save_turn.py | 5 + tests/agent/test_runtime_refresh.py | 10 +- tests/agent/test_self_model_preset.py | 19 +- tests/agent/test_unified_session.py | 55 +++- tests/cli/test_restart_command.py | 14 +- tests/command/test_model_command.py | 3 +- tests/test_nanobot_facade.py | 8 +- 18 files changed, 486 insertions(+), 214 deletions(-) diff --git a/nanobot/agent/autocompact.py b/nanobot/agent/autocompact.py index 38a2a390..6c9e580c 100644 --- a/nanobot/agent/autocompact.py +++ b/nanobot/agent/autocompact.py @@ -12,6 +12,7 @@ from nanobot.session.manager import Session, SessionManager if TYPE_CHECKING: from nanobot.agent.memory import Consolidator + from nanobot.utils.llm_runtime import LLMRuntime class AutoCompact: @@ -62,8 +63,12 @@ class AutoCompact: def _is_internal_session(cls, key: str) -> bool: return key.startswith(cls._INTERNAL_SESSION_PREFIXES) - def check_expired(self, schedule_background: Callable[[Coroutine], None], - active_session_keys: Collection[str] = ()) -> None: + def check_expired( + self, + schedule_background: Callable[[Coroutine], None], + resolve_runtime: Callable[[], LLMRuntime], + active_session_keys: Collection[str] = (), + ) -> None: """Schedule archival for idle sessions, skipping those with in-flight agent tasks.""" now = datetime.now() for info in self.sessions.list_sessions(): @@ -74,16 +79,19 @@ class AutoCompact: continue updated_at = info.get("updated_at") if self._is_expired(updated_at, now) and self._has_compactable_idle_tail(key): + runtime = resolve_runtime() self._archiving.add(key) - schedule_background(self._archive(key)) + schedule_background(self._archive(key, runtime=runtime)) - async def _archive(self, key: str) -> None: + async def _archive(self, key: str, *, runtime: LLMRuntime) -> None: if self._is_internal_session(key): self._archiving.discard(key) return try: summary = await self.consolidator.compact_idle_session( - key, self._RECENT_SUFFIX_MESSAGES, + key, + runtime=runtime, + max_suffix=self._RECENT_SUFFIX_MESSAGES, ) if summary and summary != "(nothing)": session = self.sessions.get_or_create(key) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 48063efa..2919fd80 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -388,13 +388,9 @@ class AgentLoop: ) self.consolidator = Consolidator( store=self.context.memory, - provider=provider, - model=self.model, sessions=self.sessions, - context_window_tokens=self.context_window_tokens, build_messages=self.context.build_messages, get_tool_definitions=self.tools.get_definitions, - max_completion_tokens=provider.generation.max_tokens, consolidation_ratio=consolidation_ratio, unified_session=unified_session, ) @@ -492,7 +488,6 @@ class AgentLoop: self.provider = provider self.model = model self.context_window_tokens = context_window_tokens - 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: @@ -974,6 +969,7 @@ class AgentLoop: except asyncio.TimeoutError: self.auto_compact.check_expired( self._schedule_background, + self.llm_runtime, active_session_keys=self._pending_queues.keys(), ) continue @@ -1272,6 +1268,7 @@ class AgentLoop: await self.consolidator.maybe_consolidate_by_tokens( session, + runtime=runtime, replay_max_messages=replay_max_messages_for_context( runtime.context_window_tokens ), @@ -1328,6 +1325,7 @@ class AgentLoop: self._schedule_background( self.consolidator.maybe_consolidate_by_tokens( session, + runtime=runtime, replay_max_messages=replay_max_messages_for_context( runtime.context_window_tokens ), @@ -1539,7 +1537,12 @@ class AgentLoop: async def _state_command(self, ctx: TurnContext) -> str: raw = ctx.msg.content.strip() cmd_ctx = CommandContext( - msg=ctx.msg, session=ctx.session, key=ctx.session_key, raw=raw, loop=self + msg=ctx.msg, + session=ctx.session, + key=ctx.session_key, + raw=raw, + loop=self, + runtime=ctx.runtime, ) result = await self.commands.dispatch(cmd_ctx) if result is not None: @@ -1568,6 +1571,7 @@ class AgentLoop: if not ctx.ephemeral: await self.consolidator.maybe_consolidate_by_tokens( ctx.session, + runtime=ctx.runtime, replay_max_messages=replay_max_messages, ) if message_tool := self.tools.get("message"): @@ -1673,6 +1677,7 @@ class AgentLoop: self._schedule_background( self.consolidator.maybe_consolidate_by_tokens( ctx.session, + runtime=ctx.runtime, replay_max_messages=replay_max_messages_for_context( ctx.runtime.context_window_tokens ), diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index 0decc52e..753b261d 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -30,8 +30,8 @@ from nanobot.utils.helpers import ( from nanobot.utils.prompt_templates import render_template if TYPE_CHECKING: - from nanobot.providers.base import LLMProvider from nanobot.session.manager import SessionManager + from nanobot.utils.llm_runtime import LLMRuntime # --------------------------------------------------------------------------- # MemoryStore — pure file I/O layer @@ -742,22 +742,14 @@ class Consolidator: def __init__( self, store: MemoryStore, - provider: LLMProvider, - model: str, sessions: SessionManager, - context_window_tokens: int, build_messages: Callable[..., list[dict[str, Any]]], get_tool_definitions: Callable[[], list[dict[str, Any]]], - max_completion_tokens: int = 4096, consolidation_ratio: float = 0.5, unified_session: bool = False, ): self.store = store - self.provider = provider - self.model = model self.sessions = sessions - self.context_window_tokens = context_window_tokens - self.max_completion_tokens = max_completion_tokens self.consolidation_ratio = consolidation_ratio self.unified_session = unified_session self._build_messages = build_messages @@ -766,17 +758,6 @@ class Consolidator: weakref.WeakValueDictionary() ) - def set_provider( - self, - provider: LLMProvider, - model: str, - context_window_tokens: int, - ) -> None: - self.provider = provider - self.model = model - self.context_window_tokens = context_window_tokens - self.max_completion_tokens = provider.generation.max_tokens - def get_lock(self, session_key: str) -> asyncio.Lock: """Return the shared consolidation lock for one session.""" return self._locks.setdefault(session_key, asyncio.Lock()) @@ -854,6 +835,8 @@ class Consolidator: 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) @@ -868,7 +851,11 @@ class Consolidator: len(chunk), replay_max_messages, ) - summary = await self.archive(chunk, session_key=session.key) + summary = await self.archive( + chunk, + runtime=runtime, + session_key=session.key, + ) session.last_consolidated = end_idx self.sessions.save(session) return summary @@ -884,6 +871,8 @@ class Consolidator: def estimate_session_prompt_tokens( self, session: Session, + *, + runtime: LLMRuntime, ) -> tuple[int, str]: """Estimate prompt size from the full unconsolidated session tail.""" history = self._full_unconsolidated_history(session) @@ -903,20 +892,23 @@ class Consolidator: unified_session=self.unified_session, ) return estimate_prompt_tokens_chain( - self.provider, - self.model, + runtime.provider, + runtime.model, probe_messages, self._get_tool_definitions(), ) - @property - def _input_token_budget(self) -> int: + def _input_token_budget(self, runtime: LLMRuntime) -> int: """Available input token budget for consolidation LLM.""" - return self.context_window_tokens - self.max_completion_tokens - self._SAFETY_BUFFER + return ( + runtime.context_window_tokens + - runtime.generation.max_tokens + - self._SAFETY_BUFFER + ) - def _truncate_to_token_budget(self, text: str) -> str: + def _truncate_to_token_budget(self, text: str, *, runtime: LLMRuntime) -> str: """Truncate text so it fits within the consolidation LLM's token budget.""" - budget = self._input_token_budget + budget = self._input_token_budget(runtime) if budget <= 0: return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS) return truncate_text_to_tokens(text, budget) @@ -925,6 +917,7 @@ class Consolidator: self, messages: list[dict], *, + runtime: LLMRuntime, session_key: str | None = None, summary_messages: list[dict] | None = None, ) -> str | None: @@ -942,9 +935,9 @@ class Consolidator: messages_to_summarize = summary_messages if summary_messages is not None else messages try: formatted = MemoryStore._format_messages(messages_to_summarize) - formatted = self._truncate_to_token_budget(formatted) - response = await self.provider.chat_with_retry( - model=self.model, + formatted = self._truncate_to_token_budget(formatted, runtime=runtime) + response = await runtime.provider.chat_with_retry( + model=runtime.model, messages=[ { "role": "system", @@ -957,6 +950,9 @@ class Consolidator: ], tools=None, tool_choice=None, + temperature=runtime.generation.temperature, + max_tokens=runtime.generation.max_tokens, + reasoning_effort=runtime.generation.reasoning_effort, ) if response.finish_reason == "error": raise RuntimeError(f"LLM returned error: {response.content}") @@ -976,6 +972,7 @@ class Consolidator: self, session: Session, *, + runtime: LLMRuntime, replay_max_messages: int | None = None, ) -> None: """Loop: archive old messages until prompt fits within safe budget. @@ -983,7 +980,7 @@ class Consolidator: The budget reserves space for completion tokens and a safety buffer so the LLM request never exceeds the context window. """ - if self.context_window_tokens <= 0: + if runtime.context_window_tokens <= 0: return lock = self.get_lock(session.key) @@ -995,15 +992,17 @@ class Consolidator: if not session.messages: return - budget = self._input_token_budget + 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, ) try: estimated, source = self.estimate_session_prompt_tokens( session, + runtime=runtime, ) except Exception: logger.exception("Token estimation failed for {}", session.key) @@ -1017,7 +1016,7 @@ class Consolidator: "Token consolidation idle {}: {}/{} via {}, msgs={}", session.key, estimated, - self.context_window_tokens, + runtime.context_window_tokens, source, unconsolidated_count, ) @@ -1048,11 +1047,15 @@ class Consolidator: round_num, session.key, estimated, - self.context_window_tokens, + runtime.context_window_tokens, source, len(chunk), ) - summary = await self.archive(chunk, session_key=session.key) + summary = await self.archive( + chunk, + runtime=runtime, + session_key=session.key, + ) # Advance the cursor either way: on success the chunk was # summarized; on failure archive() already raw-archived it as # a breadcrumb. Re-archiving the same chunk on the next call @@ -1069,6 +1072,7 @@ class Consolidator: try: estimated, source = self.estimate_session_prompt_tokens( session, + runtime=runtime, ) except Exception: logger.exception("Token estimation failed for {}", session.key) @@ -1084,6 +1088,8 @@ class Consolidator: async def compact_idle_session( self, session_key: str, + *, + runtime: LLMRuntime, max_suffix: int = 8, ) -> str | None: """Hard-truncate an idle session under the consolidation lock. @@ -1126,6 +1132,7 @@ class Consolidator: # the messages that are no longer kept in the live session. summary = await self.archive( messages_to_remove, + runtime=runtime, session_key=session_key, summary_messages=messages_to_summarize, ) diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 9e6a516e..05011e4d 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -227,9 +227,13 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage: """Build an outbound status message for a session.""" loop = ctx.loop session = ctx.session or loop.sessions.get_or_create(ctx.key) + runtime = ctx.runtime or loop.llm_runtime() ctx_est = 0 with suppress(Exception): - ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(session) + ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens( + session, + runtime=runtime, + ) if ctx_est <= 0: ctx_est = loop._last_usage.get("prompt_tokens", 0) @@ -253,16 +257,14 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage: channel=ctx.msg.channel, chat_id=ctx.msg.chat_id, content=build_status_content( - version=__version__, model=loop.model, + version=__version__, model=runtime.model, start_time=loop._start_time, last_usage=loop._last_usage, - context_window_tokens=loop.context_window_tokens, + context_window_tokens=runtime.context_window_tokens, session_msg_count=len(session.get_history(max_messages=0)), context_tokens_estimate=ctx_est, search_usage_text=search_usage_text, active_task_count=task_count, - max_completion_tokens=getattr( - getattr(loop.provider, "generation", None), "max_tokens", 8192 - ), + max_completion_tokens=runtime.generation.max_tokens, ), metadata={**dict(ctx.msg.metadata or {}), "render_as": "text"}, ) @@ -278,7 +280,14 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage: loop.sessions.save(session) loop.sessions.invalidate(session.key) if snapshot: - loop._schedule_background(loop.consolidator.archive(snapshot, session_key=ctx.key)) + runtime = ctx.runtime or loop.llm_runtime() + loop._schedule_background( + loop.consolidator.archive( + snapshot, + runtime=runtime, + session_key=ctx.key, + ) + ) return OutboundMessage( channel=ctx.msg.channel, chat_id=ctx.msg.chat_id, content="New session started.", diff --git a/nanobot/command/router.py b/nanobot/command/router.py index fdbe1a2a..ed174e1a 100644 --- a/nanobot/command/router.py +++ b/nanobot/command/router.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable if TYPE_CHECKING: from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.session.manager import Session + from nanobot.utils.llm_runtime import LLMRuntime Handler = Callable[["CommandContext"], Awaitable["OutboundMessage | None"]] _BOT_SUFFIX_RE = re.compile(r"^[A-Za-z0-9_]+$") @@ -43,6 +44,7 @@ class CommandContext: raw: str args: str = "" loop: Any = None + runtime: LLMRuntime | None = None class CommandRouter: diff --git a/nanobot/sdk/clients.py b/nanobot/sdk/clients.py index 49a68222..602aec76 100644 --- a/nanobot/sdk/clients.py +++ b/nanobot/sdk/clients.py @@ -13,6 +13,7 @@ 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 @@ -151,15 +152,21 @@ class RuntimeClient: async def compact_session(self, session_key: str) -> SessionSnapshot: """Run token/replay-window consolidation for one session.""" session = self._loop.sessions.get_or_create(session_key) + runtime = self._loop.llm_runtime() await self._loop.consolidator.maybe_consolidate_by_tokens( session, - replay_max_messages=self._loop._max_messages, + 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)) async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None: """Run idle-session compaction for one session and return the summary.""" + runtime = self._loop.llm_runtime() return await self._loop.consolidator.compact_idle_session( session_key, + runtime=runtime, max_suffix=max_suffix, ) diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index 78672276..34bbf04d 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -85,7 +85,7 @@ def _make_fake_compact( state = {"count": 0} - async def _fake_compact(key: str, max_suffix: int = 8) -> str: + async def _fake_compact(key: str, *, runtime, max_suffix: int = 8) -> str: state["count"] += 1 session = loop.sessions.get_or_create(key) @@ -307,7 +307,7 @@ class TestAutoCompact: loop.sessions.save(s2) loop.consolidator.compact_idle_session = _make_fake_compact(loop) - loop.auto_compact.check_expired(loop._schedule_background) + loop.auto_compact.check_expired(loop._schedule_background, loop.llm_runtime) await _drain_background_tasks(loop) active_after = loop.sessions.get_or_create("cli:active") @@ -328,7 +328,7 @@ class TestAutoCompact: loop, track_archived=archived_messages, ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) assert len(archived_messages) == 4 session_after = loop.sessions.get_or_create("cli:test") @@ -348,7 +348,7 @@ class TestAutoCompact: session.add_message("assistant", "done") loop.sessions.save(session) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) > loop.auto_compact._RECENT_SUFFIX_MESSAGES @@ -378,7 +378,7 @@ class TestAutoCompact: loop, summary="User said hello.", ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) entry = loop.auto_compact._summaries.get("cli:test") assert entry is not None @@ -394,7 +394,7 @@ class TestAutoCompact: loop.consolidator.compact_idle_session = _make_fake_compact(loop) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == 0 @@ -415,7 +415,7 @@ class TestAutoCompact: loop, track_archived=archived_messages, ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) assert len(archived_messages) == 2 await loop.close_mcp() @@ -455,7 +455,7 @@ class TestAutoCompactIdleDetection: ) # Simulate proactive archive completing before message arrives - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg") await loop._process_message(msg) @@ -579,7 +579,7 @@ class TestAutoCompactSystemMessages: loop.consolidator.compact_idle_session = _make_fake_compact(loop) # Simulate proactive archive completing before system message arrives - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) msg = InboundMessage( channel="system", sender_id="subagent", chat_id="cli:test", @@ -611,7 +611,7 @@ class TestAutoCompactEdgeCases: return_value=LLMResponse(content="(nothing)", tool_calls=[]) ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES @@ -632,7 +632,7 @@ class TestAutoCompactEdgeCases: loop.provider.chat_with_retry = AsyncMock(side_effect=Exception("API down")) # Should not raise - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) session_after = loop.sessions.get_or_create("cli:test") assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES @@ -659,7 +659,7 @@ class TestAutoCompactEdgeCases: ) # Simulate proactive archive completing before message arrives - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="continue") await loop._process_message(msg) @@ -751,7 +751,7 @@ class TestAutoCompactIntegration: loop.consolidator.compact_idle_session = _make_fake_compact(loop) # Simulate proactive archive completing before message arrives - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) msg = InboundMessage( channel="cli", sender_id="user", chat_id="test", @@ -776,6 +776,7 @@ class TestProactiveAutoCompact: """Helper: run check_expired via callback and wait for background tasks.""" loop.auto_compact.check_expired( loop._schedule_background, + loop.llm_runtime, active_session_keys=active_session_keys, ) await _drain_background_tasks(loop) @@ -867,7 +868,7 @@ class TestProactiveAutoCompact: started = asyncio.Event() block_forever = asyncio.Event() - async def _slow_compact(key, max_suffix=8): + async def _slow_compact(key, *, runtime, max_suffix=8): nonlocal archive_count archive_count += 1 started.set() @@ -877,12 +878,12 @@ class TestProactiveAutoCompact: loop.consolidator.compact_idle_session = _slow_compact # First call starts archiving via callback - loop.auto_compact.check_expired(loop._schedule_background) + loop.auto_compact.check_expired(loop._schedule_background, loop.llm_runtime) await started.wait() assert archive_count == 1 # Second call should skip (key is in _archiving) - loop.auto_compact.check_expired(loop._schedule_background) + loop.auto_compact.check_expired(loop._schedule_background, loop.llm_runtime) assert archive_count == 1 # Clean up @@ -899,7 +900,7 @@ class TestProactiveAutoCompact: session.updated_at = datetime.now() - timedelta(minutes=20) loop.sessions.save(session) - async def _failing_compact(key, max_suffix=8): + async def _failing_compact(key, *, runtime, max_suffix=8): raise RuntimeError("LLM down") loop.consolidator.compact_idle_session = _failing_compact @@ -1056,7 +1057,7 @@ class TestProactiveAutoCompact: loop.consolidator.compact_idle_session = _fake_compact # First compact cycle - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) assert _fake_compact.state["count"] == 1 # User returns, sends new messages @@ -1070,7 +1071,7 @@ class TestProactiveAutoCompact: loop.sessions.save(session2) # Second compact cycle should succeed - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) assert _fake_compact.state["count"] == 2 await loop.close_mcp() @@ -1091,7 +1092,7 @@ class TestSummaryPersistence: loop, summary="User said hello.", ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # Summary should be persisted in session metadata session_after = loop.sessions.get_or_create("cli:test") @@ -1116,7 +1117,7 @@ class TestSummaryPersistence: ) # Archive - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # Simulate restart: clear in-memory state loop.auto_compact._summaries.clear() @@ -1145,7 +1146,7 @@ class TestSummaryPersistence: loop.consolidator.compact_idle_session = _make_fake_compact(loop) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # Clear in-memory to force metadata path loop.auto_compact._summaries.clear() @@ -1173,7 +1174,7 @@ class TestSummaryPersistence: loop.consolidator.compact_idle_session = _make_fake_compact(loop) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # Both _summaries and metadata have the summary assert "cli:test" in loop.auto_compact._summaries @@ -1200,7 +1201,7 @@ class TestSummaryPersistence: loop.consolidator.compact_idle_session = _make_fake_compact( loop, summary="First summary.", ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # Consume the first summary via hot path _, summary1 = loop.auto_compact.prepare_session( @@ -1218,7 +1219,7 @@ class TestSummaryPersistence: loop.consolidator.compact_idle_session = _make_fake_compact( loop, summary="Second summary.", ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # The second archive writes a new summary assert "cli:test" in loop.auto_compact._summaries @@ -1242,7 +1243,7 @@ class TestSummaryPersistence: loop.consolidator.compact_idle_session = _make_fake_compact( loop, summary="Old summary.", ) - await loop.auto_compact._archive("cli:test") + await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime()) # Verify summary exists before /new reloaded = loop.sessions.get_or_create("cli:test") diff --git a/tests/agent/test_autocompact_unit.py b/tests/agent/test_autocompact_unit.py index dd85f956..ec610aba 100644 --- a/tests/agent/test_autocompact_unit.py +++ b/tests/agent/test_autocompact_unit.py @@ -9,6 +9,10 @@ from nanobot.agent.autocompact import AutoCompact from nanobot.session.manager import Session, SessionManager +def _runtime(): + return MagicMock(name="runtime") + + def _make_session( key: str = "cli:test", messages: list | None = None, @@ -193,7 +197,7 @@ class TestCheckExpired: mock_sm.list_sessions.return_value = [] ac.sessions = mock_sm scheduler = MagicMock() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) scheduler.assert_not_called() def test_expired_session_schedules_background(self): @@ -213,10 +217,36 @@ class TestCheckExpired: scheduled.append(coro) coro.close() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) assert len(scheduled) == 1 assert "cli:old" in ac._archiving + @pytest.mark.asyncio + async def test_runtime_is_captured_before_background_starts(self): + ac = _make_autocompact(ttl=15) + old_dt = datetime.now() - timedelta(minutes=20) + session = _make_session("cli:old", updated_at=old_dt) + _add_turns(session, 5) + ac.sessions.list_sessions.return_value = [ + {"key": "cli:old", "updated_at": old_dt.isoformat()} + ] + ac.sessions.get_or_create.return_value = session + admitted = _runtime() + replacement = _runtime() + resolve_runtime = MagicMock(return_value=admitted) + scheduled = [] + + ac.check_expired(scheduled.append, resolve_runtime) + resolve_runtime.return_value = replacement + await scheduled[0] + + resolve_runtime.assert_called_once_with() + ac.consolidator.compact_idle_session.assert_awaited_once_with( + "cli:old", + runtime=admitted, + max_suffix=ac._RECENT_SUFFIX_MESSAGES, + ) + def test_active_session_key_skips(self): """Session in active_session_keys should be skipped.""" ac = _make_autocompact(ttl=15) @@ -225,7 +255,7 @@ class TestCheckExpired: mock_sm.list_sessions.return_value = [{"key": "cli:busy", "updated_at": old_ts}] ac.sessions = mock_sm scheduler = MagicMock() - ac.check_expired(scheduler, active_session_keys={"cli:busy"}) + ac.check_expired(scheduler, _runtime, active_session_keys={"cli:busy"}) scheduler.assert_not_called() def test_session_already_in_archiving_skips(self): @@ -237,7 +267,7 @@ class TestCheckExpired: ac.sessions = mock_sm ac._archiving.add("cli:dup") scheduler = MagicMock() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) scheduler.assert_not_called() def test_session_with_no_key_skips(self): @@ -247,7 +277,7 @@ class TestCheckExpired: mock_sm.list_sessions.return_value = [{"key": "", "updated_at": "old"}] ac.sessions = mock_sm scheduler = MagicMock() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) scheduler.assert_not_called() def test_session_with_missing_key_field_skips(self): @@ -257,7 +287,7 @@ class TestCheckExpired: mock_sm.list_sessions.return_value = [{"updated_at": "old"}] ac.sessions = mock_sm scheduler = MagicMock() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) scheduler.assert_not_called() def test_dream_session_skips(self): @@ -271,7 +301,7 @@ class TestCheckExpired: ac.sessions = mock_sm scheduler = MagicMock() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) scheduler.assert_not_called() assert "dream:20260602-155256" not in ac._archiving @@ -290,7 +320,7 @@ class TestCheckExpired: ac.sessions = mock_sm scheduler = MagicMock() - ac.check_expired(scheduler) + ac.check_expired(scheduler, _runtime) scheduler.assert_not_called() @@ -310,10 +340,13 @@ class TestArchiveDelegates: ac.sessions = mock_sm ac.consolidator.compact_idle_session = AsyncMock(return_value="Summary.") - await ac._archive("cli:test") + runtime = _runtime() + await ac._archive("cli:test", runtime=runtime) ac.consolidator.compact_idle_session.assert_awaited_once_with( - "cli:test", ac._RECENT_SUFFIX_MESSAGES, + "cli:test", + runtime=runtime, + max_suffix=ac._RECENT_SUFFIX_MESSAGES, ) @pytest.mark.asyncio @@ -322,7 +355,7 @@ class TestArchiveDelegates: ac.consolidator.compact_idle_session = AsyncMock(return_value="Summary.") ac._archiving.add("dream:20260602-155256") - await ac._archive("dream:20260602-155256") + await ac._archive("dream:20260602-155256", runtime=_runtime()) ac.consolidator.compact_idle_session.assert_not_awaited() assert "dream:20260602-155256" not in ac._archiving @@ -338,7 +371,7 @@ class TestArchiveDelegates: ac.sessions = mock_sm ac.consolidator.compact_idle_session = AsyncMock(return_value="Hello.") - await ac._archive("cli:test") + await ac._archive("cli:test", runtime=_runtime()) entry = ac._summaries.get("cli:test") assert entry is not None @@ -351,7 +384,7 @@ class TestArchiveDelegates: ac.sessions = mock_sm ac.consolidator.compact_idle_session = AsyncMock(return_value="") - await ac._archive("cli:test") + await ac._archive("cli:test", runtime=_runtime()) assert "cli:test" not in ac._summaries @@ -362,7 +395,7 @@ class TestArchiveDelegates: ac.sessions = mock_sm ac.consolidator.compact_idle_session = AsyncMock(return_value="(nothing)") - await ac._archive("cli:test") + await ac._archive("cli:test", runtime=_runtime()) assert "cli:test" not in ac._summaries @@ -374,7 +407,7 @@ class TestArchiveDelegates: ac.consolidator.compact_idle_session = AsyncMock(side_effect=RuntimeError("fail")) ac._archiving.add("cli:test") - await ac._archive("cli:test") + await ac._archive("cli:test", runtime=_runtime()) assert "cli:test" not in ac._archiving diff --git a/tests/agent/test_consolidation_ratio.py b/tests/agent/test_consolidation_ratio.py index b1c95ec4..5a7b8009 100644 --- a/tests/agent/test_consolidation_ratio.py +++ b/tests/agent/test_consolidation_ratio.py @@ -77,14 +77,18 @@ async def test_consolidation_ratio_controls_target( remaining_estimates = list(estimates) - def mock_estimate(_session, *, session_summary=None): - assert session_summary is None + runtime = loop.llm_runtime() + + def mock_estimate(_session, *, runtime): return (remaining_estimates.pop(0), "test") loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100) - await loop.consolidator.maybe_consolidate_by_tokens(session) + await loop.consolidator.maybe_consolidate_by_tokens( + session, + runtime=runtime, + ) assert loop.consolidator.archive.await_count == expected_archives diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index 34350095..4b949ae2 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -1,5 +1,6 @@ """Tests for the lightweight Consolidator — append-only to HISTORY.md.""" +from dataclasses import replace from unittest.mock import AsyncMock, MagicMock import pytest @@ -9,8 +10,9 @@ from nanobot.agent.memory import ( Consolidator, MemoryStore, ) -from nanobot.providers.base import LLMResponse +from nanobot.providers.base import GenerationSettings, LLMResponse from nanobot.session.manager import Session +from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.prompt_templates import render_template @@ -23,11 +25,21 @@ def store(tmp_path): def mock_provider(): p = MagicMock() p.chat_with_retry = AsyncMock() + p.generation = GenerationSettings(max_tokens=100) return p @pytest.fixture -def consolidator(store, mock_provider): +def runtime(mock_provider): + return LLMRuntime.capture( + mock_provider, + "test-model", + context_window_tokens=1000, + ) + + +@pytest.fixture +def consolidator(store): sessions = MagicMock() sessions.save = MagicMock() # When maybe_consolidate_by_tokens refreshes the session reference via @@ -38,13 +50,9 @@ def consolidator(store, mock_provider): sessions._session_cache = _session_cache return Consolidator( store=store, - provider=mock_provider, - model="test-model", sessions=sessions, - context_window_tokens=1000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), - max_completion_tokens=100, ) @@ -62,7 +70,41 @@ def _tool_round(call_id: str) -> list[dict]: class TestConsolidatorSummarize: - async def test_summarize_appends_to_history(self, consolidator, mock_provider, store): + async def test_archive_uses_captured_generation( + self, consolidator, mock_provider, runtime + ): + admitted = replace( + runtime, + generation=GenerationSettings( + temperature=0.25, + max_tokens=321, + reasoning_effort="medium", + ), + ) + mock_provider.generation = GenerationSettings( + temperature=0.9, + max_tokens=999, + reasoning_effort="high", + ) + mock_provider.chat_with_retry.return_value = MagicMock( + content="Summary.", + finish_reason="stop", + ) + + await consolidator.archive( + [{"role": "user", "content": "hello"}], + runtime=admitted, + ) + + call = mock_provider.chat_with_retry.call_args.kwargs + assert call["model"] == admitted.model + assert call["temperature"] == 0.25 + assert call["max_tokens"] == 321 + assert call["reasoning_effort"] == "medium" + + async def test_summarize_appends_to_history( + self, consolidator, mock_provider, store, runtime + ): """Consolidator should call LLM to summarize, then append to HISTORY.md.""" mock_provider.chat_with_retry.return_value = MagicMock( content="User fixed a bug in the auth module." @@ -71,7 +113,7 @@ class TestConsolidatorSummarize: {"role": "user", "content": "fix the auth bug"}, {"role": "assistant", "content": "Done, fixed the race condition."}, ] - result = await consolidator.archive(messages) + result = await consolidator.archive(messages, runtime=runtime) assert result == "User fixed a bug in the auth module." entries = store.read_unprocessed_history(since_cursor=0) assert len(entries) == 1 @@ -81,6 +123,7 @@ class TestConsolidatorSummarize: consolidator, mock_provider, store, + runtime, ): mock_provider.chat_with_retry.return_value = MagicMock( content="User fixed a bug in the auth module.", @@ -88,16 +131,22 @@ class TestConsolidatorSummarize: ) messages = [{"role": "user", "content": "fix the auth bug"}] - await consolidator.archive(messages, session_key="telegram:chat-1") + await consolidator.archive( + messages, + runtime=runtime, + session_key="telegram:chat-1", + ) entries = store.read_unprocessed_history(since_cursor=0) assert entries[0]["session_key"] == "telegram:chat-1" - async def test_summarize_raw_dumps_on_llm_failure(self, consolidator, mock_provider, store): + async def test_summarize_raw_dumps_on_llm_failure( + self, consolidator, mock_provider, store, runtime + ): """On LLM failure, raw-dump messages to HISTORY.md.""" mock_provider.chat_with_retry.side_effect = Exception("API error") messages = [{"role": "user", "content": "hello"}] - result = await consolidator.archive(messages) + result = await consolidator.archive(messages, runtime=runtime) assert result is None # no summary on raw dump fallback entries = store.read_unprocessed_history(since_cursor=0) assert len(entries) == 1 @@ -108,17 +157,22 @@ class TestConsolidatorSummarize: consolidator, mock_provider, store, + runtime, ): mock_provider.chat_with_retry.side_effect = Exception("API error") messages = [{"role": "user", "content": "hello"}] - await consolidator.archive(messages, session_key="slack:chat-2") + await consolidator.archive( + messages, + runtime=runtime, + session_key="slack:chat-2", + ) entries = store.read_unprocessed_history(since_cursor=0) assert entries[0]["session_key"] == "slack:chat-2" - async def test_summarize_skips_empty_messages(self, consolidator): - result = await consolidator.archive([]) + async def test_summarize_skips_empty_messages(self, consolidator, runtime): + result = await consolidator.archive([], runtime=runtime) assert result is None @@ -139,7 +193,9 @@ class TestConsolidatorArchiveErrorHandling: See https://github.com/HKUDS/nanobot/issues/3244 """ - async def test_archive_falls_back_on_error_finish_reason(self, consolidator, mock_provider, store): + async def test_archive_falls_back_on_error_finish_reason( + self, consolidator, mock_provider, store, runtime + ): """LLM returning finish_reason='error' should trigger raw_archive, not write error text.""" mock_provider.chat_with_retry.return_value = MagicMock( content="Error: {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'overloaded_error (529)'}}", @@ -149,14 +205,16 @@ class TestConsolidatorArchiveErrorHandling: {"role": "user", "content": "fix the auth bug"}, {"role": "assistant", "content": "Done, fixed the race condition."}, ] - result = await consolidator.archive(messages) + result = await consolidator.archive(messages, runtime=runtime) assert result is None entries = store.read_unprocessed_history(since_cursor=0) assert len(entries) == 1 assert "[RAW]" in entries[0]["content"] assert "Error:" not in entries[0]["content"] - async def test_archive_preserves_summary_on_success(self, consolidator, mock_provider, store): + async def test_archive_preserves_summary_on_success( + self, consolidator, mock_provider, store, runtime + ): """Normal LLM response should still produce a proper summary entry.""" mock_provider.chat_with_retry.return_value = MagicMock( content="User fixed a bug in the auth module.", @@ -166,7 +224,7 @@ class TestConsolidatorArchiveErrorHandling: {"role": "user", "content": "fix the auth bug"}, {"role": "assistant", "content": "Done."}, ] - result = await consolidator.archive(messages) + result = await consolidator.archive(messages, runtime=runtime) assert result == "User fixed a bug in the auth module." entries = store.read_unprocessed_history(since_cursor=0) assert len(entries) == 1 @@ -174,7 +232,9 @@ class TestConsolidatorArchiveErrorHandling: class TestConsolidatorTokenBudget: - async def test_prompt_below_threshold_does_not_consolidate(self, consolidator): + async def test_prompt_below_threshold_does_not_consolidate( + self, consolidator, runtime + ): """No consolidation when tokens are within budget.""" session = MagicMock() session.last_consolidated = 0 @@ -183,10 +243,10 @@ class TestConsolidatorTokenBudget: consolidator.sessions._session_cache[session.key] = session consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken")) consolidator.archive = AsyncMock(return_value=True) - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) consolidator.archive.assert_not_called() - async def test_estimate_uses_full_unconsolidated_tail(self, consolidator): + async def test_estimate_uses_full_unconsolidated_tail(self, consolidator, runtime): """Consolidation pressure must see messages hidden by the replay window.""" session = Session(key="test:full-tail") for i in range(160): @@ -200,7 +260,7 @@ class TestConsolidatorTokenBudget: consolidator._build_messages = build_messages - consolidator.estimate_session_prompt_tokens(session) + consolidator.estimate_session_prompt_tokens(session, runtime=runtime) assert len(captured["history"]) == 160 assert captured["history"][0]["content"].endswith("msg-0") @@ -208,6 +268,7 @@ class TestConsolidatorTokenBudget: 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 @@ -222,6 +283,7 @@ class TestConsolidatorTokenBudget: await consolidator.maybe_consolidate_by_tokens( session, + runtime=runtime, replay_max_messages=6, ) @@ -235,6 +297,7 @@ class TestConsolidatorTokenBudget: 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") @@ -251,6 +314,7 @@ class TestConsolidatorTokenBudget: await consolidator.maybe_consolidate_by_tokens( session, + runtime=runtime, replay_max_messages=4, ) @@ -266,6 +330,7 @@ class TestConsolidatorTokenBudget: 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") @@ -284,6 +349,7 @@ class TestConsolidatorTokenBudget: await consolidator.maybe_consolidate_by_tokens( session, + runtime=runtime, replay_max_messages=6, ) @@ -295,7 +361,7 @@ class TestConsolidatorTokenBudget: 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_large_chunk_archived_without_cap(self, consolidator): + async def test_large_chunk_archived_without_cap(self, consolidator, runtime): """Without chunk cap, the full range from pick_consolidation_boundary is archived.""" consolidator._SAFETY_BUFFER = 0 session = MagicMock() @@ -316,14 +382,16 @@ class TestConsolidatorTokenBudget: # (user message at 50, token budget met) consolidator.archive = AsyncMock(return_value=True) - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) archived_chunk = consolidator.archive.await_args.args[0] # pick_consolidation_boundary returns (50, tokens) — user turn at idx 50 assert archived_chunk[0]["content"] == "m0" assert session.last_consolidated > 0 - async def test_raw_archive_fallback_advances_last_consolidated(self, consolidator): + async def test_raw_archive_fallback_advances_last_consolidated( + self, consolidator, runtime + ): """When archive() falls back to raw-archive (LLM failed), the cursor must still advance. Otherwise the same chunk gets raw-archived again on every subsequent maybe_consolidate_by_tokens() call, spamming @@ -344,14 +412,16 @@ class TestConsolidatorTokenBudget: # LLM consolidation fails — archive() returns None (raw_archive fired). consolidator.archive = AsyncMock(return_value=None) - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) consolidator.archive.assert_awaited_once() # The chunk is considered "materialized" (as a raw-archive breadcrumb), # so last_consolidated must have moved past it. assert session.last_consolidated == 50 - async def test_raw_archive_fallback_breaks_round_loop(self, consolidator): + async def test_raw_archive_fallback_breaks_round_loop( + self, consolidator, runtime + ): """A degraded LLM should not trigger more archive() calls within the same maybe_consolidate_by_tokens invocation — bail after one fallback.""" consolidator._SAFETY_BUFFER = 0 @@ -370,12 +440,14 @@ class TestConsolidatorTokenBudget: ) consolidator.archive = AsyncMock(return_value=None) - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) # Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS. assert consolidator.archive.await_count == 1 - async def test_boundary_respected_when_no_intermediate_user_turn(self, consolidator): + async def test_boundary_respected_when_no_intermediate_user_turn( + self, consolidator, runtime + ): """When boundary points past a long tool chain, the full chunk is archived.""" consolidator._SAFETY_BUFFER = 0 session = MagicMock() @@ -394,7 +466,7 @@ class TestConsolidatorTokenBudget: ) consolidator.archive = AsyncMock(return_value=True) - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) consolidator.archive.assert_awaited_once() # pick_consolidation_boundary finds the only boundary at idx=61 @@ -412,17 +484,15 @@ class TestCompactIdleSession: sessions = SessionManager(store.workspace) return Consolidator( store=store, - provider=mock_provider, - model="test-model", sessions=sessions, - context_window_tokens=1000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), - max_completion_tokens=100, ) @pytest.mark.asyncio - async def test_archives_prefix_keeps_suffix(self, real_consolidator, mock_provider): + async def test_archives_prefix_keeps_suffix( + self, real_consolidator, mock_provider, runtime + ): """20 user/assistant turns → compact with max_suffix=8 → messages ≤ 8, last_consolidated=0, _last_summary stored.""" mock_provider.chat_with_retry.return_value = MagicMock( @@ -437,7 +507,9 @@ class TestCompactIdleSession: session.updated_at = old_ts sessions.save(session) - result = await real_consolidator.compact_idle_session("cli:test", max_suffix=8) + result = await real_consolidator.compact_idle_session( + "cli:test", runtime=runtime, max_suffix=8 + ) assert result == "Summary of old conversation." reloaded = sessions.get_or_create("cli:test") @@ -451,7 +523,7 @@ class TestCompactIdleSession: @pytest.mark.asyncio async def test_summarizes_retained_suffix_not_just_dropped_prefix( - self, real_consolidator, mock_provider + self, real_consolidator, mock_provider, runtime ): """idleCompact must summarize over the full unconsolidated tail, including the recent suffix it retains. Otherwise a late user correction / final @@ -470,14 +542,16 @@ class TestCompactIdleSession: session.add_message("assistant", "CORRECTED_FINAL_RESULT_alpha") sessions.save(session) - await real_consolidator.compact_idle_session("cli:correction", max_suffix=8) + await real_consolidator.compact_idle_session( + "cli:correction", runtime=runtime, max_suffix=8 + ) summarized = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"] assert "CORRECTED_FINAL_RESULT_alpha" in summarized @pytest.mark.asyncio async def test_raw_dumps_only_dropped_messages_on_llm_failure( - self, real_consolidator, mock_provider, store + self, real_consolidator, mock_provider, store, runtime ): """Summarizing over the full tail must not widen what gets raw-dumped on LLM failure: the breadcrumb should contain only the removed prefix, not @@ -492,7 +566,9 @@ class TestCompactIdleSession: session.add_message("assistant", "RETAINED_SUFFIX_marker") sessions.save(session) - await real_consolidator.compact_idle_session("cli:rawdrop", max_suffix=8) + await real_consolidator.compact_idle_session( + "cli:rawdrop", runtime=runtime, max_suffix=8 + ) raw = "\n".join(e["content"] for e in store.read_unprocessed_history(since_cursor=0)) assert "[RAW]" in raw @@ -505,6 +581,7 @@ class TestCompactIdleSession: real_consolidator, mock_provider, store, + runtime, ): mock_provider.chat_with_retry.return_value = MagicMock( content="Summary of old conversation.", finish_reason="stop" @@ -515,14 +592,16 @@ class TestCompactIdleSession: session.add_message("assistant", f"assistant msg {i}") real_consolidator.sessions.save(session) - await real_consolidator.compact_idle_session("cli:test", max_suffix=4) + await real_consolidator.compact_idle_session( + "cli:test", runtime=runtime, max_suffix=4 + ) entries = store.read_unprocessed_history(since_cursor=0) assert entries[0]["session_key"] == "cli:test" @pytest.mark.asyncio async def test_empty_session_does_not_refresh_timestamp( - self, real_consolidator + self, real_consolidator, runtime ): """Empty session with old updated_at does not look active after compaction.""" from datetime import datetime, timedelta @@ -533,7 +612,9 @@ class TestCompactIdleSession: session.updated_at = old_ts sessions.save(session) - result = await real_consolidator.compact_idle_session("cli:empty") + result = await real_consolidator.compact_idle_session( + "cli:empty", runtime=runtime + ) assert result == "" reloaded = sessions.get_or_create("cli:empty") @@ -541,7 +622,9 @@ class TestCompactIdleSession: assert reloaded.metadata == {} @pytest.mark.asyncio - async def test_nothing_summary_not_stored(self, real_consolidator, mock_provider): + async def test_nothing_summary_not_stored( + self, real_consolidator, mock_provider, runtime + ): """LLM returns '(nothing)' → _last_summary NOT in metadata.""" mock_provider.chat_with_retry.return_value = MagicMock( content="(nothing)", finish_reason="stop" @@ -553,14 +636,18 @@ class TestCompactIdleSession: session.add_message("assistant", f"a{i}") sessions.save(session) - result = await real_consolidator.compact_idle_session("cli:nothing", max_suffix=4) + result = await real_consolidator.compact_idle_session( + "cli:nothing", runtime=runtime, max_suffix=4 + ) assert result == "(nothing)" reloaded = sessions.get_or_create("cli:nothing") assert "_last_summary" not in reloaded.metadata @pytest.mark.asyncio - async def test_llm_failure_still_truncates(self, real_consolidator, mock_provider, store): + async def test_llm_failure_still_truncates( + self, real_consolidator, mock_provider, store, runtime + ): """LLM raises RuntimeError → raw_archive fires, session still truncated, returns None.""" mock_provider.chat_with_retry.side_effect = RuntimeError("LLM unavailable") sessions = real_consolidator.sessions @@ -570,7 +657,9 @@ class TestCompactIdleSession: session.add_message("assistant", f"a{i}") sessions.save(session) - result = await real_consolidator.compact_idle_session("cli:fail", max_suffix=4) + result = await real_consolidator.compact_idle_session( + "cli:fail", runtime=runtime, max_suffix=4 + ) assert result is None # raw_archive should have been called (history.jsonl gets an entry) @@ -582,7 +671,9 @@ class TestCompactIdleSession: assert len(reloaded.messages) <= 4 @pytest.mark.asyncio - async def test_respects_last_consolidated(self, real_consolidator, mock_provider): + async def test_respects_last_consolidated( + self, real_consolidator, mock_provider, runtime + ): """30 turns with last_consolidated=50 → only unconsolidated tail considered.""" mock_provider.chat_with_retry.return_value = MagicMock( content="Tail summary.", finish_reason="stop" @@ -595,7 +686,9 @@ class TestCompactIdleSession: session.last_consolidated = 50 # Only 10 messages unconsolidated sessions.save(session) - result = await real_consolidator.compact_idle_session("cli:offset", max_suffix=4) + result = await real_consolidator.compact_idle_session( + "cli:offset", runtime=runtime, max_suffix=4 + ) assert result == "Tail summary." # Verify only the unconsolidated tail was processed: @@ -611,6 +704,7 @@ class TestCompactIdleSession: self, real_consolidator, mock_provider, + runtime, ): """Assistant-only tails extend back to the latest user turn, so archive the actual dropped messages rather than a computed prefix.""" @@ -625,7 +719,9 @@ class TestCompactIdleSession: session.add_message("assistant", f"assistant-{i:02d}") sessions.save(session) - result = await real_consolidator.compact_idle_session("cli:noncontiguous", max_suffix=6) + result = await real_consolidator.compact_idle_session( + "cli:noncontiguous", runtime=runtime, max_suffix=6 + ) assert result == "Tail summary." reloaded = sessions.get_or_create("cli:noncontiguous") @@ -653,7 +749,9 @@ class TestCompactIdleSession: assert "user-14" in user_content @pytest.mark.asyncio - async def test_acquires_consolidation_lock(self, real_consolidator, mock_provider): + async def test_acquires_consolidation_lock( + self, real_consolidator, mock_provider, runtime + ): """Verify lock is held during execution.""" import asyncio @@ -679,7 +777,9 @@ class TestCompactIdleSession: assert not lock.locked() task = asyncio.ensure_future( - real_consolidator.compact_idle_session("cli:lock", max_suffix=4) + real_consolidator.compact_idle_session( + "cli:lock", runtime=runtime, max_suffix=4 + ) ) await started.wait() assert lock.locked() @@ -702,15 +802,17 @@ class TestConsolidatorSessionRefresh: provider.chat_with_retry = AsyncMock( return_value=MagicMock(content="summary", finish_reason="stop") ) - provider.generation.max_tokens = 4096 + provider.generation = GenerationSettings(max_tokens=4096) provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test")) + runtime = LLMRuntime.capture( + provider, + "test-model", + context_window_tokens=128_000, + ) sessions = SessionManager(tmp_path) consolidator = Consolidator( store=store, - provider=provider, - model="test-model", sessions=sessions, - context_window_tokens=128_000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), ) @@ -722,13 +824,16 @@ class TestConsolidatorSessionRefresh: seen: dict[str, Session] = {} - def estimate(session: Session): + def estimate(session: Session, *, runtime): seen["session"] = session return 10, "test" consolidator.estimate_session_prompt_tokens = MagicMock(side_effect=estimate) - await consolidator.maybe_consolidate_by_tokens(stale_empty) + await consolidator.maybe_consolidate_by_tokens( + stale_empty, + runtime=runtime, + ) assert seen["session"] is fresh @@ -745,15 +850,17 @@ class TestConsolidatorSessionRefresh: provider.chat_with_retry = AsyncMock( return_value=MagicMock(content="summary", finish_reason="stop") ) - provider.generation.max_tokens = 4096 + provider.generation = GenerationSettings(max_tokens=4096) provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test")) + runtime = LLMRuntime.capture( + provider, + "test-model", + context_window_tokens=128_000, + ) sessions = SessionManager(tmp_path) consolidator = Consolidator( store=store, - provider=provider, - model="test-model", sessions=sessions, - context_window_tokens=128_000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), ) @@ -769,11 +876,18 @@ class TestConsolidatorSessionRefresh: old_ref = session # AutoCompact runs first and truncates to 8 - await consolidator.compact_idle_session("cli:test", max_suffix=8) + await consolidator.compact_idle_session( + "cli:test", + runtime=runtime, + max_suffix=8, + ) # Background consolidation runs with stale reference — # should detect the session was replaced and not undo the compact. - await consolidator.maybe_consolidate_by_tokens(old_ref) + await consolidator.maybe_consolidate_by_tokens( + old_ref, + runtime=runtime, + ) session_after = sessions.get_or_create("cli:test") # Messages should still be truncated (not restored to 40) @@ -818,7 +932,9 @@ class TestRawArchiveTruncation: class TestArchiveTruncation: """archive() must truncate formatted text before sending to consolidation LLM.""" - async def test_archive_truncates_large_formatted_text(self, consolidator, mock_provider, store): + async def test_archive_truncates_large_formatted_text( + self, consolidator, mock_provider, store, runtime + ): """Large formatted text should be truncated to token budget before LLM call.""" # context_window_tokens=1000, max_completion_tokens=100, _SAFETY_BUFFER=1024 # budget = 1000 - 100 - 1024 = -124 → fallback via truncate_text(budget*4) @@ -826,21 +942,23 @@ class TestArchiveTruncation: mock_provider.chat_with_retry.return_value = MagicMock( content="Summary of large input.", finish_reason="stop" ) - await consolidator.archive(big_messages) + await consolidator.archive(big_messages, runtime=runtime) call_args = mock_provider.chat_with_retry.call_args user_content = call_args.kwargs["messages"][1]["content"] # Should be significantly shorter than 100K assert len(user_content) < 50_000 - async def test_archive_truncates_with_small_token_budget(self, consolidator, mock_provider, store): + async def test_archive_truncates_with_small_token_budget( + self, consolidator, mock_provider, store, runtime + ): """Small context window: truncation uses actual tokenizer count.""" - consolidator.context_window_tokens = 500 + runtime = replace(runtime, context_window_tokens=500) big_messages = [{"role": "user", "content": "word " * 50_000}] mock_provider.chat_with_retry.return_value = MagicMock( content="Summary.", finish_reason="stop" ) - await consolidator.archive(big_messages) + await consolidator.archive(big_messages, runtime=runtime) sent_messages = mock_provider.chat_with_retry.call_args.kwargs["messages"] user_content = sent_messages[1]["content"] @@ -848,7 +966,9 @@ class TestArchiveTruncation: # Should be truncated assert len(user_content) < 250_000 - async def test_oversized_summary_is_capped_before_append(self, consolidator, mock_provider, store): + async def test_oversized_summary_is_capped_before_append( + self, consolidator, mock_provider, store, runtime + ): """A pathologically large LLM summary must not land full-length in history.jsonl — that would re-open the #3412 bloat vector from the *success* path instead of the fallback path.""" @@ -856,21 +976,26 @@ class TestArchiveTruncation: content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10), finish_reason="stop", ) - await consolidator.archive([{"role": "user", "content": "hi"}]) + await consolidator.archive( + [{"role": "user", "content": "hi"}], + runtime=runtime, + ) entry = store.read_unprocessed_history(since_cursor=0)[0] assert len(entry["content"]) <= _ARCHIVE_SUMMARY_MAX_CHARS + 50 - async def test_archive_truncates_via_tiktoken_with_positive_budget(self, consolidator, mock_provider, store): + async def test_archive_truncates_via_tiktoken_with_positive_budget( + self, consolidator, mock_provider, store, runtime + ): """Positive token budget should use tiktoken for precise truncation.""" - consolidator.context_window_tokens = 10_000 + runtime = replace(runtime, context_window_tokens=10_000) consolidator._SAFETY_BUFFER = 0 # budget = 10000 - 100 - 0 = 9900 tokens big_messages = [{"role": "user", "content": "word " * 50_000}] mock_provider.chat_with_retry.return_value = MagicMock( content="Summary.", finish_reason="stop" ) - await consolidator.archive(big_messages) + await consolidator.archive(big_messages, runtime=runtime) import tiktoken enc = tiktoken.get_encoding("cl100k_base") diff --git a/tests/agent/test_loop_consolidation_tokens.py b/tests/agent/test_loop_consolidation_tokens.py index 3c1f6fcb..f113b14a 100644 --- a/tests/agent/test_loop_consolidation_tokens.py +++ b/tests/agent/test_loop_consolidation_tokens.py @@ -76,7 +76,10 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path token_map = {"u1": 120, "a1": 120, "u2": 120, "a2": 120, "u3": 120} monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda message: token_map[message["content"]]) - await loop.consolidator.maybe_consolidate_by_tokens(session) + await loop.consolidator.maybe_consolidate_by_tokens( + session, + runtime=loop.llm_runtime(), + ) archived_chunk = loop.consolidator.archive.await_args.args[0] assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"] @@ -102,7 +105,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No loop.sessions.save(session) call_count = [0] - def mock_estimate(_session, *, session_summary=None): + def mock_estimate(_session, *, runtime): call_count[0] += 1 if call_count[0] == 1: return (500, "test") @@ -113,7 +116,10 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100) - await loop.consolidator.maybe_consolidate_by_tokens(session) + await loop.consolidator.maybe_consolidate_by_tokens( + session, + runtime=loop.llm_runtime(), + ) assert loop.consolidator.archive.await_count == 2 assert session.last_consolidated == 6 @@ -139,7 +145,7 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, call_count = [0] - def mock_estimate(_session, *, session_summary=None): + def mock_estimate(_session, *, runtime): call_count[0] += 1 if call_count[0] == 1: return (500, "test") @@ -150,7 +156,10 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100) - await loop.consolidator.maybe_consolidate_by_tokens(session) + await loop.consolidator.maybe_consolidate_by_tokens( + session, + runtime=loop.llm_runtime(), + ) assert loop.consolidator.archive.await_count == 2 assert session.last_consolidated == 6 @@ -171,7 +180,7 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, call_count = [0] - def mock_estimate(_session, *, session_summary=None): + def mock_estimate(_session, *, runtime): call_count[0] += 1 if call_count[0] == 1: return (500, "test") @@ -180,7 +189,10 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 150) - await loop.consolidator.maybe_consolidate_by_tokens(session) + await loop.consolidator.maybe_consolidate_by_tokens( + session, + runtime=loop.llm_runtime(), + ) reloaded = loop.sessions.get_or_create("cli:test") meta = reloaded.metadata.get("_last_summary") @@ -204,12 +216,19 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) # type: ignore[method-assign] loop._schedule_background = lambda coro: coro.close() # type: ignore[method-assign] - await loop.process_direct("hello", session_key="cli:test") + runtime = loop.llm_runtime() + await loop.process_direct("hello", session_key="cli:test", runtime=runtime) loop.consolidator.maybe_consolidate_by_tokens.assert_any_await( session, + runtime=runtime, replay_max_messages=loop._max_messages, ) + assert len(loop.consolidator.maybe_consolidate_by_tokens.call_args_list) == 2 + assert all( + call.kwargs["runtime"] is runtime + for call in loop.consolidator.maybe_consolidate_by_tokens.call_args_list + ) @pytest.mark.asyncio @@ -221,7 +240,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) -> archived_session_keys: list[str | None] = [] - async def track_consolidate(messages, *, session_key=None): + async def track_consolidate(messages, *, runtime, session_key=None): order.append("consolidate") archived_session_keys.append(session_key) return True @@ -244,7 +263,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) -> monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500) call_count = [0] - def mock_estimate(_session, *, session_summary=None): + def mock_estimate(_session, *, runtime): call_count[0] += 1 return (1000 if call_count[0] <= 1 else 80, "test") loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index 4d1d4e81..0726671d 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -1308,6 +1308,11 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_ assert seen["runtime"] is runtime record_runtime.assert_called_once_with("cli:test", runtime) + assert len(loop.consolidator.maybe_consolidate_by_tokens.call_args_list) == 2 + assert all( + call.kwargs["runtime"] is runtime + for call in loop.consolidator.maybe_consolidate_by_tokens.call_args_list + ) initial_messages = seen["initial_messages"] assert isinstance(initial_messages, list) non_system = [m for m in initial_messages if m.get("role") != "system"] diff --git a/tests/agent/test_runtime_refresh.py b/tests/agent/test_runtime_refresh.py index 3cb3f572..e387530d 100644 --- a/tests/agent/test_runtime_refresh.py +++ b/tests/agent/test_runtime_refresh.py @@ -44,10 +44,10 @@ def test_provider_refresh_updates_all_model_dependents(tmp_path: Path) -> None: assert not hasattr(loop.subagents, "provider") assert not hasattr(loop.subagents, "model") assert not hasattr(loop.subagents.runner, "provider") - assert loop.consolidator.provider is new_provider - assert loop.consolidator.model == "new-model" - assert loop.consolidator.context_window_tokens == 2000 - assert loop.consolidator.max_completion_tokens == 456 + assert not hasattr(loop.consolidator, "provider") + assert not hasattr(loop.consolidator, "model") + assert not hasattr(loop.consolidator, "context_window_tokens") + assert not hasattr(loop.consolidator, "max_completion_tokens") def test_llm_runtime_refreshes_provider_snapshot(tmp_path: Path) -> None: @@ -122,4 +122,4 @@ def test_settings_context_window_refreshes_runtime_state( assert payload["requires_restart"] is False assert loop.context_window_tokens == 262_144 - assert loop.consolidator.context_window_tokens == 262_144 + assert loop.llm_runtime().context_window_tokens == 262_144 diff --git a/tests/agent/test_self_model_preset.py b/tests/agent/test_self_model_preset.py index 2af09824..b2e71625 100644 --- a/tests/agent/test_self_model_preset.py +++ b/tests/agent/test_self_model_preset.py @@ -58,9 +58,12 @@ def test_model_preset_setter_updates_state(tmp_path) -> None: assert loop.provider.generation.max_tokens == 4096 assert loop.provider.generation.reasoning_effort == "low" assert not hasattr(loop.subagents, "model") - assert loop.consolidator.model == "openai/gpt-4.1" - assert loop.consolidator.context_window_tokens == 32_768 - assert loop.consolidator.max_completion_tokens == 4096 + assert not hasattr(loop.consolidator, "model") + assert not hasattr(loop.consolidator, "context_window_tokens") + assert loop.llm_runtime().model == "openai/gpt-4.1" + assert loop.llm_runtime().context_window_tokens == 32_768 + assert not hasattr(loop.consolidator, "max_completion_tokens") + assert loop.llm_runtime().generation.max_tokens == 4096 def test_model_preset_setter_calls_runtime_model_publisher(tmp_path) -> None: @@ -110,10 +113,11 @@ def test_model_preset_setter_replaces_provider_from_snapshot(tmp_path) -> None: assert not hasattr(loop.runner, "provider") assert not hasattr(loop.subagents, "provider") assert not hasattr(loop.subagents.runner, "provider") - assert loop.consolidator.provider is new_provider + assert not hasattr(loop.consolidator, "provider") assert loop.model == "anthropic/claude-opus-4-5" assert loop.context_window_tokens == 200_000 - assert loop.consolidator.max_completion_tokens == 2048 + assert not hasattr(loop.consolidator, "max_completion_tokens") + assert loop.llm_runtime().generation.max_tokens == 2048 def test_model_preset_setter_failure_leaves_old_state(tmp_path) -> None: @@ -136,9 +140,10 @@ def test_model_preset_setter_failure_leaves_old_state(tmp_path) -> None: assert loop.model_preset is None assert loop.model == "base-model" assert not hasattr(loop.subagents, "model") - assert loop.consolidator.model == "base-model" + assert not hasattr(loop.consolidator, "model") assert loop.context_window_tokens == 1000 - assert loop.consolidator.max_completion_tokens == 123 + assert not hasattr(loop.consolidator, "max_completion_tokens") + assert loop.llm_runtime().generation.max_tokens == 123 def test_active_model_preset_survives_unchanged_config_refresh(tmp_path) -> None: diff --git a/tests/agent/test_unified_session.py b/tests/agent/test_unified_session.py index d6059f74..8f436b3b 100644 --- a/tests/agent/test_unified_session.py +++ b/tests/agent/test_unified_session.py @@ -25,8 +25,10 @@ from nanobot.bus.queue import MessageBus from nanobot.command.builtin import cmd_new, register_builtin_commands from nanobot.command.router import CommandContext, CommandRouter from nanobot.config.schema import AgentDefaults, Config +from nanobot.providers.base import GenerationSettings from nanobot.session.keys import UNIFIED_SESSION_KEY from nanobot.session.manager import Session, SessionManager +from nanobot.utils.llm_runtime import LLMRuntime # --------------------------------------------------------------------------- # Helpers @@ -50,6 +52,15 @@ def _make_loop(tmp_path: Path, unified_session: bool = False) -> AgentLoop: return loop +def _runtime(provider) -> LLMRuntime: + provider.generation = GenerationSettings(max_tokens=100) + return LLMRuntime.capture( + provider, + "test-model", + context_window_tokens=1000, + ) + + def _make_msg(channel: str = "telegram", chat_id: str = "111", session_key_override: str | None = None) -> InboundMessage: return InboundMessage( @@ -233,14 +244,17 @@ class TestCmdNewUnifiedSession: shared.add_message("assistant", "hi there") sessions.save(shared) assert len(sessions.get_or_create("unified:default").messages) == 2 + expected_snapshot = list(shared.messages) # _schedule_background is a *sync* method that schedules a coroutine via # asyncio.create_task(). Mirror that exactly so the coroutine is consumed # and no RuntimeWarning is emitted. + admitted_runtime = MagicMock(name="admitted_runtime") loop = SimpleNamespace( sessions=sessions, consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), _cancel_active_tasks=AsyncMock(return_value=0), + llm_runtime=MagicMock(return_value=MagicMock()), ) loop._schedule_background = lambda coro: asyncio.ensure_future(coro) @@ -248,7 +262,14 @@ class TestCmdNewUnifiedSession: channel="telegram", sender_id="user1", chat_id="111", content="/new", session_key_override="unified:default", # as _dispatch() would set it ) - ctx = CommandContext(msg=msg, session=None, key="unified:default", raw="/new", loop=loop) + ctx = CommandContext( + msg=msg, + session=None, + key="unified:default", + raw="/new", + loop=loop, + runtime=admitted_runtime, + ) result = await cmd_new(ctx) @@ -257,6 +278,12 @@ class TestCmdNewUnifiedSession: sessions.invalidate("unified:default") reloaded = sessions.get_or_create("unified:default") assert reloaded.messages == [] + loop.consolidator.archive.assert_called_once_with( + expected_snapshot, + runtime=admitted_runtime, + session_key="unified:default", + ) + loop.llm_runtime.assert_not_called() @pytest.mark.asyncio async def test_cmd_new_in_unified_mode_does_not_affect_other_sessions(self, tmp_path: Path): @@ -275,6 +302,7 @@ class TestCmdNewUnifiedSession: sessions=sessions, consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), _cancel_active_tasks=AsyncMock(return_value=0), + llm_runtime=MagicMock(return_value=MagicMock()), ) loop._schedule_background = lambda coro: asyncio.ensure_future(coro) @@ -306,26 +334,23 @@ class TestConsolidationUnaffectedByUnifiedSession: store = MagicMock(spec=MemoryStore) mock_provider = MagicMock() mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary")) + runtime = _runtime(mock_provider) # Use spec= so MagicMock doesn't auto-generate AsyncMock for non-async methods, # which would leave unawaited coroutines and trigger RuntimeWarning. sessions = MagicMock(spec=SessionManager) consolidator = Consolidator( store=store, - provider=mock_provider, - model="test-model", sessions=sessions, - context_window_tokens=1000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), - max_completion_tokens=100, ) consolidator.archive = AsyncMock() session = Session(key="unified:default") session.messages = [] - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) consolidator.archive.assert_not_called() @@ -341,24 +366,24 @@ class TestConsolidationUnaffectedByUnifiedSession: store = MagicMock(spec=MemoryStore) mock_provider = MagicMock() mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary")) + runtime = _runtime(mock_provider) sessions = MagicMock(spec=SessionManager) consolidator = Consolidator( store=store, - provider=mock_provider, - model="test-model", sessions=sessions, - context_window_tokens=1000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), - max_completion_tokens=100, ) session = Session(key=key) session.messages = [] # empty → exits immediately for both keys consolidator.archive = AsyncMock() - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens( + session, + runtime=runtime, + ) archive_calls[key] = consolidator.archive.call_count assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0 @@ -371,17 +396,14 @@ class TestConsolidationUnaffectedByUnifiedSession: store = MagicMock(spec=MemoryStore) mock_provider = MagicMock() + runtime = _runtime(mock_provider) sessions = MagicMock(spec=SessionManager) consolidator = Consolidator( store=store, - provider=mock_provider, - model="test-model", sessions=sessions, - context_window_tokens=1000, build_messages=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]), - max_completion_tokens=100, ) session = Session(key="unified:default") @@ -394,11 +416,12 @@ class TestConsolidationUnaffectedByUnifiedSession: consolidator.pick_consolidation_boundary = MagicMock(return_value=None) consolidator.archive = AsyncMock() - await consolidator.maybe_consolidate_by_tokens(session) + await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) # estimate was called (consolidation was attempted) consolidator.estimate_session_prompt_tokens.assert_called_once_with( session, + runtime=runtime, ) # but archive was not called (no valid boundary) consolidator.archive.assert_not_called() diff --git a/tests/cli/test_restart_command.py b/tests/cli/test_restart_command.py index 9b5012cb..8a227a86 100644 --- a/tests/cli/test_restart_command.py +++ b/tests/cli/test_restart_command.py @@ -244,8 +244,16 @@ class TestRestartCommand: loop.subagents.get_running_count_by_session.return_value = 0 msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/status") + runtime = loop.llm_runtime() + loop.model = "replacement-model" + loop.context_window_tokens = 10 + loop.provider.generation = SimpleNamespace( + temperature=1.0, + max_tokens=1, + reasoning_effort=None, + ) - response = await loop._process_message(msg) + response = await loop._process_message(msg, runtime=runtime) assert response is not None assert "Model: test-model" in response.content @@ -255,6 +263,10 @@ class TestRestartCommand: assert "Uptime: 2m 5s" in response.content assert "Tasks: 0 active" in response.content assert response.metadata == {"render_as": "text"} + loop.consolidator.estimate_session_prompt_tokens.assert_called_once_with( + session, + runtime=runtime, + ) @pytest.mark.asyncio async def test_status_counts_running_dispatch_and_subagent_tasks(self): diff --git a/tests/command/test_model_command.py b/tests/command/test_model_command.py index 9647ce65..6f9e3d2f 100644 --- a/tests/command/test_model_command.py +++ b/tests/command/test_model_command.py @@ -86,7 +86,8 @@ async def test_model_command_switches_preset(tmp_path) -> None: assert loop.model_preset == "fast" assert loop.model == "openai/gpt-4.1" assert not hasattr(loop.subagents, "model") - assert loop.consolidator.model == "openai/gpt-4.1" + assert not hasattr(loop.consolidator, "model") + assert loop.llm_runtime().model == "openai/gpt-4.1" @pytest.mark.asyncio diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index dc1b7e8f..a1a93b46 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -1213,11 +1213,16 @@ async def test_runtime_helpers_expose_model_workspace_and_compact(tmp_path): config_path = _write_config(tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path) await bot.sessions.ingest("sdk:history", [{"role": "user", "content": "hello"}]) + runtime = bot._loop.llm_runtime() + bot._loop.llm_runtime = MagicMock(return_value=runtime) # type: ignore[method-assign] bot._loop.consolidator.maybe_consolidate_by_tokens = AsyncMock() snapshot = await bot.runtime.compact_session("sdk:history") assert snapshot.key == "sdk:history" - bot._loop.consolidator.maybe_consolidate_by_tokens.assert_awaited_once() + assert ( + bot._loop.consolidator.maybe_consolidate_by_tokens.await_args.kwargs["runtime"] + is runtime + ) assert bot.runtime.model == bot._loop.model assert bot.runtime.workspace == tmp_path @@ -1226,6 +1231,7 @@ async def test_runtime_helpers_expose_model_workspace_and_compact(tmp_path): assert summary == "Summary." bot._loop.consolidator.compact_idle_session.assert_awaited_once_with( "sdk:history", + runtime=runtime, max_suffix=4, )