refactor(agent): require runtime for consolidation

This commit is contained in:
chengyongru 2026-07-10 14:52:51 +08:00 committed by Xubin Ren
parent 5bd3d1e0af
commit c9d3e74342
18 changed files with 486 additions and 214 deletions

View File

@ -12,6 +12,7 @@ from nanobot.session.manager import Session, SessionManager
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.agent.memory import Consolidator from nanobot.agent.memory import Consolidator
from nanobot.utils.llm_runtime import LLMRuntime
class AutoCompact: class AutoCompact:
@ -62,8 +63,12 @@ class AutoCompact:
def _is_internal_session(cls, key: str) -> bool: def _is_internal_session(cls, key: str) -> bool:
return key.startswith(cls._INTERNAL_SESSION_PREFIXES) return key.startswith(cls._INTERNAL_SESSION_PREFIXES)
def check_expired(self, schedule_background: Callable[[Coroutine], None], def check_expired(
active_session_keys: Collection[str] = ()) -> None: 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.""" """Schedule archival for idle sessions, skipping those with in-flight agent tasks."""
now = datetime.now() now = datetime.now()
for info in self.sessions.list_sessions(): for info in self.sessions.list_sessions():
@ -74,16 +79,19 @@ class AutoCompact:
continue continue
updated_at = info.get("updated_at") updated_at = info.get("updated_at")
if self._is_expired(updated_at, now) and self._has_compactable_idle_tail(key): if self._is_expired(updated_at, now) and self._has_compactable_idle_tail(key):
runtime = resolve_runtime()
self._archiving.add(key) 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): if self._is_internal_session(key):
self._archiving.discard(key) self._archiving.discard(key)
return return
try: try:
summary = await self.consolidator.compact_idle_session( 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)": if summary and summary != "(nothing)":
session = self.sessions.get_or_create(key) session = self.sessions.get_or_create(key)

View File

@ -388,13 +388,9 @@ class AgentLoop:
) )
self.consolidator = Consolidator( self.consolidator = Consolidator(
store=self.context.memory, store=self.context.memory,
provider=provider,
model=self.model,
sessions=self.sessions, sessions=self.sessions,
context_window_tokens=self.context_window_tokens,
build_messages=self.context.build_messages, build_messages=self.context.build_messages,
get_tool_definitions=self.tools.get_definitions, get_tool_definitions=self.tools.get_definitions,
max_completion_tokens=provider.generation.max_tokens,
consolidation_ratio=consolidation_ratio, consolidation_ratio=consolidation_ratio,
unified_session=unified_session, unified_session=unified_session,
) )
@ -492,7 +488,6 @@ class AgentLoop:
self.provider = provider self.provider = provider
self.model = model self.model = model
self.context_window_tokens = context_window_tokens self.context_window_tokens = context_window_tokens
self.consolidator.set_provider(provider, model, context_window_tokens)
self._sync_replay_max_messages() self._sync_replay_max_messages()
self._provider_signature = snapshot.signature self._provider_signature = snapshot.signature
if publish_update and self._runtime_model_publisher is not None: if publish_update and self._runtime_model_publisher is not None:
@ -974,6 +969,7 @@ class AgentLoop:
except asyncio.TimeoutError: except asyncio.TimeoutError:
self.auto_compact.check_expired( self.auto_compact.check_expired(
self._schedule_background, self._schedule_background,
self.llm_runtime,
active_session_keys=self._pending_queues.keys(), active_session_keys=self._pending_queues.keys(),
) )
continue continue
@ -1272,6 +1268,7 @@ class AgentLoop:
await self.consolidator.maybe_consolidate_by_tokens( await self.consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime,
replay_max_messages=replay_max_messages_for_context( replay_max_messages=replay_max_messages_for_context(
runtime.context_window_tokens runtime.context_window_tokens
), ),
@ -1328,6 +1325,7 @@ class AgentLoop:
self._schedule_background( self._schedule_background(
self.consolidator.maybe_consolidate_by_tokens( self.consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime,
replay_max_messages=replay_max_messages_for_context( replay_max_messages=replay_max_messages_for_context(
runtime.context_window_tokens runtime.context_window_tokens
), ),
@ -1539,7 +1537,12 @@ class AgentLoop:
async def _state_command(self, ctx: TurnContext) -> str: async def _state_command(self, ctx: TurnContext) -> str:
raw = ctx.msg.content.strip() raw = ctx.msg.content.strip()
cmd_ctx = CommandContext( 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) result = await self.commands.dispatch(cmd_ctx)
if result is not None: if result is not None:
@ -1568,6 +1571,7 @@ class AgentLoop:
if not ctx.ephemeral: if not ctx.ephemeral:
await self.consolidator.maybe_consolidate_by_tokens( await self.consolidator.maybe_consolidate_by_tokens(
ctx.session, ctx.session,
runtime=ctx.runtime,
replay_max_messages=replay_max_messages, replay_max_messages=replay_max_messages,
) )
if message_tool := self.tools.get("message"): if message_tool := self.tools.get("message"):
@ -1673,6 +1677,7 @@ class AgentLoop:
self._schedule_background( self._schedule_background(
self.consolidator.maybe_consolidate_by_tokens( self.consolidator.maybe_consolidate_by_tokens(
ctx.session, ctx.session,
runtime=ctx.runtime,
replay_max_messages=replay_max_messages_for_context( replay_max_messages=replay_max_messages_for_context(
ctx.runtime.context_window_tokens ctx.runtime.context_window_tokens
), ),

View File

@ -30,8 +30,8 @@ from nanobot.utils.helpers import (
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.providers.base import LLMProvider
from nanobot.session.manager import SessionManager from nanobot.session.manager import SessionManager
from nanobot.utils.llm_runtime import LLMRuntime
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# MemoryStore — pure file I/O layer # MemoryStore — pure file I/O layer
@ -742,22 +742,14 @@ class Consolidator:
def __init__( def __init__(
self, self,
store: MemoryStore, store: MemoryStore,
provider: LLMProvider,
model: str,
sessions: SessionManager, sessions: SessionManager,
context_window_tokens: int,
build_messages: Callable[..., list[dict[str, Any]]], build_messages: Callable[..., list[dict[str, Any]]],
get_tool_definitions: Callable[[], list[dict[str, Any]]], get_tool_definitions: Callable[[], list[dict[str, Any]]],
max_completion_tokens: int = 4096,
consolidation_ratio: float = 0.5, consolidation_ratio: float = 0.5,
unified_session: bool = False, unified_session: bool = False,
): ):
self.store = store self.store = store
self.provider = provider
self.model = model
self.sessions = sessions self.sessions = sessions
self.context_window_tokens = context_window_tokens
self.max_completion_tokens = max_completion_tokens
self.consolidation_ratio = consolidation_ratio self.consolidation_ratio = consolidation_ratio
self.unified_session = unified_session self.unified_session = unified_session
self._build_messages = build_messages self._build_messages = build_messages
@ -766,17 +758,6 @@ class Consolidator:
weakref.WeakValueDictionary() 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: def get_lock(self, session_key: str) -> asyncio.Lock:
"""Return the shared consolidation lock for one session.""" """Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock()) return self._locks.setdefault(session_key, asyncio.Lock())
@ -854,6 +835,8 @@ class Consolidator:
self, self,
session: Session, session: Session,
replay_max_messages: int | None, replay_max_messages: int | None,
*,
runtime: LLMRuntime,
) -> str | None: ) -> str | None:
"""Archive messages that would be hidden by the replay message window.""" """Archive messages that would be hidden by the replay message window."""
end_idx = self._replay_overflow_boundary(session, replay_max_messages) end_idx = self._replay_overflow_boundary(session, replay_max_messages)
@ -868,7 +851,11 @@ class Consolidator:
len(chunk), len(chunk),
replay_max_messages, 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 session.last_consolidated = end_idx
self.sessions.save(session) self.sessions.save(session)
return summary return summary
@ -884,6 +871,8 @@ class Consolidator:
def estimate_session_prompt_tokens( def estimate_session_prompt_tokens(
self, self,
session: Session, session: Session,
*,
runtime: LLMRuntime,
) -> tuple[int, str]: ) -> tuple[int, str]:
"""Estimate prompt size from the full unconsolidated session tail.""" """Estimate prompt size from the full unconsolidated session tail."""
history = self._full_unconsolidated_history(session) history = self._full_unconsolidated_history(session)
@ -903,20 +892,23 @@ class Consolidator:
unified_session=self.unified_session, unified_session=self.unified_session,
) )
return estimate_prompt_tokens_chain( return estimate_prompt_tokens_chain(
self.provider, runtime.provider,
self.model, runtime.model,
probe_messages, probe_messages,
self._get_tool_definitions(), self._get_tool_definitions(),
) )
@property def _input_token_budget(self, runtime: LLMRuntime) -> int:
def _input_token_budget(self) -> int:
"""Available input token budget for consolidation LLM.""" """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.""" """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: if budget <= 0:
return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS) return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS)
return truncate_text_to_tokens(text, budget) return truncate_text_to_tokens(text, budget)
@ -925,6 +917,7 @@ class Consolidator:
self, self,
messages: list[dict], messages: list[dict],
*, *,
runtime: LLMRuntime,
session_key: str | None = None, session_key: str | None = None,
summary_messages: list[dict] | None = None, summary_messages: list[dict] | None = None,
) -> str | None: ) -> str | None:
@ -942,9 +935,9 @@ class Consolidator:
messages_to_summarize = summary_messages if summary_messages is not None else messages messages_to_summarize = summary_messages if summary_messages is not None else messages
try: try:
formatted = MemoryStore._format_messages(messages_to_summarize) formatted = MemoryStore._format_messages(messages_to_summarize)
formatted = self._truncate_to_token_budget(formatted) formatted = self._truncate_to_token_budget(formatted, runtime=runtime)
response = await self.provider.chat_with_retry( response = await runtime.provider.chat_with_retry(
model=self.model, model=runtime.model,
messages=[ messages=[
{ {
"role": "system", "role": "system",
@ -957,6 +950,9 @@ class Consolidator:
], ],
tools=None, tools=None,
tool_choice=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": if response.finish_reason == "error":
raise RuntimeError(f"LLM returned error: {response.content}") raise RuntimeError(f"LLM returned error: {response.content}")
@ -976,6 +972,7 @@ class Consolidator:
self, self,
session: Session, session: Session,
*, *,
runtime: LLMRuntime,
replay_max_messages: int | None = None, replay_max_messages: int | None = None,
) -> None: ) -> None:
"""Loop: archive old messages until prompt fits within safe budget. """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 The budget reserves space for completion tokens and a safety buffer
so the LLM request never exceeds the context window. so the LLM request never exceeds the context window.
""" """
if self.context_window_tokens <= 0: if runtime.context_window_tokens <= 0:
return return
lock = self.get_lock(session.key) lock = self.get_lock(session.key)
@ -995,15 +992,17 @@ class Consolidator:
if not session.messages: if not session.messages:
return return
budget = self._input_token_budget budget = self._input_token_budget(runtime)
target = int(budget * self.consolidation_ratio) target = int(budget * self.consolidation_ratio)
last_summary = await self._consolidate_replay_overflow( last_summary = await self._consolidate_replay_overflow(
session, session,
replay_max_messages, replay_max_messages,
runtime=runtime,
) )
try: try:
estimated, source = self.estimate_session_prompt_tokens( estimated, source = self.estimate_session_prompt_tokens(
session, session,
runtime=runtime,
) )
except Exception: except Exception:
logger.exception("Token estimation failed for {}", session.key) logger.exception("Token estimation failed for {}", session.key)
@ -1017,7 +1016,7 @@ class Consolidator:
"Token consolidation idle {}: {}/{} via {}, msgs={}", "Token consolidation idle {}: {}/{} via {}, msgs={}",
session.key, session.key,
estimated, estimated,
self.context_window_tokens, runtime.context_window_tokens,
source, source,
unconsolidated_count, unconsolidated_count,
) )
@ -1048,11 +1047,15 @@ class Consolidator:
round_num, round_num,
session.key, session.key,
estimated, estimated,
self.context_window_tokens, runtime.context_window_tokens,
source, source,
len(chunk), 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 # Advance the cursor either way: on success the chunk was
# summarized; on failure archive() already raw-archived it as # summarized; on failure archive() already raw-archived it as
# a breadcrumb. Re-archiving the same chunk on the next call # a breadcrumb. Re-archiving the same chunk on the next call
@ -1069,6 +1072,7 @@ class Consolidator:
try: try:
estimated, source = self.estimate_session_prompt_tokens( estimated, source = self.estimate_session_prompt_tokens(
session, session,
runtime=runtime,
) )
except Exception: except Exception:
logger.exception("Token estimation failed for {}", session.key) logger.exception("Token estimation failed for {}", session.key)
@ -1084,6 +1088,8 @@ class Consolidator:
async def compact_idle_session( async def compact_idle_session(
self, self,
session_key: str, session_key: str,
*,
runtime: LLMRuntime,
max_suffix: int = 8, max_suffix: int = 8,
) -> str | None: ) -> str | None:
"""Hard-truncate an idle session under the consolidation lock. """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. # the messages that are no longer kept in the live session.
summary = await self.archive( summary = await self.archive(
messages_to_remove, messages_to_remove,
runtime=runtime,
session_key=session_key, session_key=session_key,
summary_messages=messages_to_summarize, summary_messages=messages_to_summarize,
) )

View File

@ -227,9 +227,13 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
"""Build an outbound status message for a session.""" """Build an outbound status message for a session."""
loop = ctx.loop loop = ctx.loop
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or loop.sessions.get_or_create(ctx.key)
runtime = ctx.runtime or loop.llm_runtime()
ctx_est = 0 ctx_est = 0
with suppress(Exception): 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: if ctx_est <= 0:
ctx_est = loop._last_usage.get("prompt_tokens", 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, channel=ctx.msg.channel,
chat_id=ctx.msg.chat_id, chat_id=ctx.msg.chat_id,
content=build_status_content( content=build_status_content(
version=__version__, model=loop.model, version=__version__, model=runtime.model,
start_time=loop._start_time, last_usage=loop._last_usage, 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)), session_msg_count=len(session.get_history(max_messages=0)),
context_tokens_estimate=ctx_est, context_tokens_estimate=ctx_est,
search_usage_text=search_usage_text, search_usage_text=search_usage_text,
active_task_count=task_count, active_task_count=task_count,
max_completion_tokens=getattr( max_completion_tokens=runtime.generation.max_tokens,
getattr(loop.provider, "generation", None), "max_tokens", 8192
),
), ),
metadata={**dict(ctx.msg.metadata or {}), "render_as": "text"}, 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.save(session)
loop.sessions.invalidate(session.key) loop.sessions.invalidate(session.key)
if snapshot: 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( return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id, channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content="New session started.", content="New session started.",

View File

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.utils.llm_runtime import LLMRuntime
Handler = Callable[["CommandContext"], Awaitable["OutboundMessage | None"]] Handler = Callable[["CommandContext"], Awaitable["OutboundMessage | None"]]
_BOT_SUFFIX_RE = re.compile(r"^[A-Za-z0-9_]+$") _BOT_SUFFIX_RE = re.compile(r"^[A-Za-z0-9_]+$")
@ -43,6 +44,7 @@ class CommandContext:
raw: str raw: str
args: str = "" args: str = ""
loop: Any = None loop: Any = None
runtime: LLMRuntime | None = None
class CommandRouter: class CommandRouter:

View File

@ -13,6 +13,7 @@ from nanobot.sdk.types import (
snapshot_from_payload, snapshot_from_payload,
snapshot_from_session, snapshot_from_session,
) )
from nanobot.session.manager import replay_max_messages_for_context
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
@ -151,15 +152,21 @@ class RuntimeClient:
async def compact_session(self, session_key: str) -> SessionSnapshot: async def compact_session(self, session_key: str) -> SessionSnapshot:
"""Run token/replay-window consolidation for one session.""" """Run token/replay-window consolidation for one session."""
session = self._loop.sessions.get_or_create(session_key) session = self._loop.sessions.get_or_create(session_key)
runtime = self._loop.llm_runtime()
await self._loop.consolidator.maybe_consolidate_by_tokens( await self._loop.consolidator.maybe_consolidate_by_tokens(
session, 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)) 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: 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.""" """Run idle-session compaction for one session and return the summary."""
runtime = self._loop.llm_runtime()
return await self._loop.consolidator.compact_idle_session( return await self._loop.consolidator.compact_idle_session(
session_key, session_key,
runtime=runtime,
max_suffix=max_suffix, max_suffix=max_suffix,
) )

View File

@ -85,7 +85,7 @@ def _make_fake_compact(
state = {"count": 0} 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 state["count"] += 1
session = loop.sessions.get_or_create(key) session = loop.sessions.get_or_create(key)
@ -307,7 +307,7 @@ class TestAutoCompact:
loop.sessions.save(s2) loop.sessions.save(s2)
loop.consolidator.compact_idle_session = _make_fake_compact(loop) 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) await _drain_background_tasks(loop)
active_after = loop.sessions.get_or_create("cli:active") active_after = loop.sessions.get_or_create("cli:active")
@ -328,7 +328,7 @@ class TestAutoCompact:
loop, track_archived=archived_messages, 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 assert len(archived_messages) == 4
session_after = loop.sessions.get_or_create("cli:test") session_after = loop.sessions.get_or_create("cli:test")
@ -348,7 +348,7 @@ class TestAutoCompact:
session.add_message("assistant", "done") session.add_message("assistant", "done")
loop.sessions.save(session) 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") session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) > loop.auto_compact._RECENT_SUFFIX_MESSAGES assert len(session_after.messages) > loop.auto_compact._RECENT_SUFFIX_MESSAGES
@ -378,7 +378,7 @@ class TestAutoCompact:
loop, summary="User said hello.", 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") entry = loop.auto_compact._summaries.get("cli:test")
assert entry is not None assert entry is not None
@ -394,7 +394,7 @@ class TestAutoCompact:
loop.consolidator.compact_idle_session = _make_fake_compact(loop) 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") session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0 assert len(session_after.messages) == 0
@ -415,7 +415,7 @@ class TestAutoCompact:
loop, track_archived=archived_messages, 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 assert len(archived_messages) == 2
await loop.close_mcp() await loop.close_mcp()
@ -455,7 +455,7 @@ class TestAutoCompactIdleDetection:
) )
# Simulate proactive archive completing before message arrives # 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") msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
await loop._process_message(msg) await loop._process_message(msg)
@ -579,7 +579,7 @@ class TestAutoCompactSystemMessages:
loop.consolidator.compact_idle_session = _make_fake_compact(loop) loop.consolidator.compact_idle_session = _make_fake_compact(loop)
# Simulate proactive archive completing before system message arrives # 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( msg = InboundMessage(
channel="system", sender_id="subagent", chat_id="cli:test", channel="system", sender_id="subagent", chat_id="cli:test",
@ -611,7 +611,7 @@ class TestAutoCompactEdgeCases:
return_value=LLMResponse(content="(nothing)", tool_calls=[]) 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") session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES 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")) loop.provider.chat_with_retry = AsyncMock(side_effect=Exception("API down"))
# Should not raise # 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") session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES assert len(session_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
@ -659,7 +659,7 @@ class TestAutoCompactEdgeCases:
) )
# Simulate proactive archive completing before message arrives # 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") msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="continue")
await loop._process_message(msg) await loop._process_message(msg)
@ -751,7 +751,7 @@ class TestAutoCompactIntegration:
loop.consolidator.compact_idle_session = _make_fake_compact(loop) loop.consolidator.compact_idle_session = _make_fake_compact(loop)
# Simulate proactive archive completing before message arrives # 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( msg = InboundMessage(
channel="cli", sender_id="user", chat_id="test", 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.""" """Helper: run check_expired via callback and wait for background tasks."""
loop.auto_compact.check_expired( loop.auto_compact.check_expired(
loop._schedule_background, loop._schedule_background,
loop.llm_runtime,
active_session_keys=active_session_keys, active_session_keys=active_session_keys,
) )
await _drain_background_tasks(loop) await _drain_background_tasks(loop)
@ -867,7 +868,7 @@ class TestProactiveAutoCompact:
started = asyncio.Event() started = asyncio.Event()
block_forever = 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 nonlocal archive_count
archive_count += 1 archive_count += 1
started.set() started.set()
@ -877,12 +878,12 @@ class TestProactiveAutoCompact:
loop.consolidator.compact_idle_session = _slow_compact loop.consolidator.compact_idle_session = _slow_compact
# First call starts archiving via callback # 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() await started.wait()
assert archive_count == 1 assert archive_count == 1
# Second call should skip (key is in _archiving) # 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 assert archive_count == 1
# Clean up # Clean up
@ -899,7 +900,7 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20) session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session) 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") raise RuntimeError("LLM down")
loop.consolidator.compact_idle_session = _failing_compact loop.consolidator.compact_idle_session = _failing_compact
@ -1056,7 +1057,7 @@ class TestProactiveAutoCompact:
loop.consolidator.compact_idle_session = _fake_compact loop.consolidator.compact_idle_session = _fake_compact
# First compact cycle # 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 assert _fake_compact.state["count"] == 1
# User returns, sends new messages # User returns, sends new messages
@ -1070,7 +1071,7 @@ class TestProactiveAutoCompact:
loop.sessions.save(session2) loop.sessions.save(session2)
# Second compact cycle should succeed # 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 assert _fake_compact.state["count"] == 2
await loop.close_mcp() await loop.close_mcp()
@ -1091,7 +1092,7 @@ class TestSummaryPersistence:
loop, summary="User said hello.", 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 # Summary should be persisted in session metadata
session_after = loop.sessions.get_or_create("cli:test") session_after = loop.sessions.get_or_create("cli:test")
@ -1116,7 +1117,7 @@ class TestSummaryPersistence:
) )
# Archive # 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 # Simulate restart: clear in-memory state
loop.auto_compact._summaries.clear() loop.auto_compact._summaries.clear()
@ -1145,7 +1146,7 @@ class TestSummaryPersistence:
loop.consolidator.compact_idle_session = _make_fake_compact(loop) 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 # Clear in-memory to force metadata path
loop.auto_compact._summaries.clear() loop.auto_compact._summaries.clear()
@ -1173,7 +1174,7 @@ class TestSummaryPersistence:
loop.consolidator.compact_idle_session = _make_fake_compact(loop) 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 # Both _summaries and metadata have the summary
assert "cli:test" in loop.auto_compact._summaries assert "cli:test" in loop.auto_compact._summaries
@ -1200,7 +1201,7 @@ class TestSummaryPersistence:
loop.consolidator.compact_idle_session = _make_fake_compact( loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="First summary.", 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 # Consume the first summary via hot path
_, summary1 = loop.auto_compact.prepare_session( _, summary1 = loop.auto_compact.prepare_session(
@ -1218,7 +1219,7 @@ class TestSummaryPersistence:
loop.consolidator.compact_idle_session = _make_fake_compact( loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="Second summary.", 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 # The second archive writes a new summary
assert "cli:test" in loop.auto_compact._summaries assert "cli:test" in loop.auto_compact._summaries
@ -1242,7 +1243,7 @@ class TestSummaryPersistence:
loop.consolidator.compact_idle_session = _make_fake_compact( loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="Old summary.", 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 # Verify summary exists before /new
reloaded = loop.sessions.get_or_create("cli:test") reloaded = loop.sessions.get_or_create("cli:test")

View File

@ -9,6 +9,10 @@ from nanobot.agent.autocompact import AutoCompact
from nanobot.session.manager import Session, SessionManager from nanobot.session.manager import Session, SessionManager
def _runtime():
return MagicMock(name="runtime")
def _make_session( def _make_session(
key: str = "cli:test", key: str = "cli:test",
messages: list | None = None, messages: list | None = None,
@ -193,7 +197,7 @@ class TestCheckExpired:
mock_sm.list_sessions.return_value = [] mock_sm.list_sessions.return_value = []
ac.sessions = mock_sm ac.sessions = mock_sm
scheduler = MagicMock() scheduler = MagicMock()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
scheduler.assert_not_called() scheduler.assert_not_called()
def test_expired_session_schedules_background(self): def test_expired_session_schedules_background(self):
@ -213,10 +217,36 @@ class TestCheckExpired:
scheduled.append(coro) scheduled.append(coro)
coro.close() coro.close()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
assert len(scheduled) == 1 assert len(scheduled) == 1
assert "cli:old" in ac._archiving 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): def test_active_session_key_skips(self):
"""Session in active_session_keys should be skipped.""" """Session in active_session_keys should be skipped."""
ac = _make_autocompact(ttl=15) ac = _make_autocompact(ttl=15)
@ -225,7 +255,7 @@ class TestCheckExpired:
mock_sm.list_sessions.return_value = [{"key": "cli:busy", "updated_at": old_ts}] mock_sm.list_sessions.return_value = [{"key": "cli:busy", "updated_at": old_ts}]
ac.sessions = mock_sm ac.sessions = mock_sm
scheduler = MagicMock() 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() scheduler.assert_not_called()
def test_session_already_in_archiving_skips(self): def test_session_already_in_archiving_skips(self):
@ -237,7 +267,7 @@ class TestCheckExpired:
ac.sessions = mock_sm ac.sessions = mock_sm
ac._archiving.add("cli:dup") ac._archiving.add("cli:dup")
scheduler = MagicMock() scheduler = MagicMock()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
scheduler.assert_not_called() scheduler.assert_not_called()
def test_session_with_no_key_skips(self): def test_session_with_no_key_skips(self):
@ -247,7 +277,7 @@ class TestCheckExpired:
mock_sm.list_sessions.return_value = [{"key": "", "updated_at": "old"}] mock_sm.list_sessions.return_value = [{"key": "", "updated_at": "old"}]
ac.sessions = mock_sm ac.sessions = mock_sm
scheduler = MagicMock() scheduler = MagicMock()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
scheduler.assert_not_called() scheduler.assert_not_called()
def test_session_with_missing_key_field_skips(self): def test_session_with_missing_key_field_skips(self):
@ -257,7 +287,7 @@ class TestCheckExpired:
mock_sm.list_sessions.return_value = [{"updated_at": "old"}] mock_sm.list_sessions.return_value = [{"updated_at": "old"}]
ac.sessions = mock_sm ac.sessions = mock_sm
scheduler = MagicMock() scheduler = MagicMock()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
scheduler.assert_not_called() scheduler.assert_not_called()
def test_dream_session_skips(self): def test_dream_session_skips(self):
@ -271,7 +301,7 @@ class TestCheckExpired:
ac.sessions = mock_sm ac.sessions = mock_sm
scheduler = MagicMock() scheduler = MagicMock()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
scheduler.assert_not_called() scheduler.assert_not_called()
assert "dream:20260602-155256" not in ac._archiving assert "dream:20260602-155256" not in ac._archiving
@ -290,7 +320,7 @@ class TestCheckExpired:
ac.sessions = mock_sm ac.sessions = mock_sm
scheduler = MagicMock() scheduler = MagicMock()
ac.check_expired(scheduler) ac.check_expired(scheduler, _runtime)
scheduler.assert_not_called() scheduler.assert_not_called()
@ -310,10 +340,13 @@ class TestArchiveDelegates:
ac.sessions = mock_sm ac.sessions = mock_sm
ac.consolidator.compact_idle_session = AsyncMock(return_value="Summary.") 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( 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 @pytest.mark.asyncio
@ -322,7 +355,7 @@ class TestArchiveDelegates:
ac.consolidator.compact_idle_session = AsyncMock(return_value="Summary.") ac.consolidator.compact_idle_session = AsyncMock(return_value="Summary.")
ac._archiving.add("dream:20260602-155256") 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() ac.consolidator.compact_idle_session.assert_not_awaited()
assert "dream:20260602-155256" not in ac._archiving assert "dream:20260602-155256" not in ac._archiving
@ -338,7 +371,7 @@ class TestArchiveDelegates:
ac.sessions = mock_sm ac.sessions = mock_sm
ac.consolidator.compact_idle_session = AsyncMock(return_value="Hello.") 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") entry = ac._summaries.get("cli:test")
assert entry is not None assert entry is not None
@ -351,7 +384,7 @@ class TestArchiveDelegates:
ac.sessions = mock_sm ac.sessions = mock_sm
ac.consolidator.compact_idle_session = AsyncMock(return_value="") 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 assert "cli:test" not in ac._summaries
@ -362,7 +395,7 @@ class TestArchiveDelegates:
ac.sessions = mock_sm ac.sessions = mock_sm
ac.consolidator.compact_idle_session = AsyncMock(return_value="(nothing)") 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 assert "cli:test" not in ac._summaries
@ -374,7 +407,7 @@ class TestArchiveDelegates:
ac.consolidator.compact_idle_session = AsyncMock(side_effect=RuntimeError("fail")) ac.consolidator.compact_idle_session = AsyncMock(side_effect=RuntimeError("fail"))
ac._archiving.add("cli:test") ac._archiving.add("cli:test")
await ac._archive("cli:test") await ac._archive("cli:test", runtime=_runtime())
assert "cli:test" not in ac._archiving assert "cli:test" not in ac._archiving

View File

@ -77,14 +77,18 @@ async def test_consolidation_ratio_controls_target(
remaining_estimates = list(estimates) remaining_estimates = list(estimates)
def mock_estimate(_session, *, session_summary=None): runtime = loop.llm_runtime()
assert session_summary is None
def mock_estimate(_session, *, runtime):
return (remaining_estimates.pop(0), "test") return (remaining_estimates.pop(0), "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100) 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 assert loop.consolidator.archive.await_count == expected_archives

View File

@ -1,5 +1,6 @@
"""Tests for the lightweight Consolidator — append-only to HISTORY.md.""" """Tests for the lightweight Consolidator — append-only to HISTORY.md."""
from dataclasses import replace
from unittest.mock import AsyncMock, MagicMock from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
@ -9,8 +10,9 @@ from nanobot.agent.memory import (
Consolidator, Consolidator,
MemoryStore, MemoryStore,
) )
from nanobot.providers.base import LLMResponse from nanobot.providers.base import GenerationSettings, LLMResponse
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.utils.llm_runtime import LLMRuntime
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
@ -23,11 +25,21 @@ def store(tmp_path):
def mock_provider(): def mock_provider():
p = MagicMock() p = MagicMock()
p.chat_with_retry = AsyncMock() p.chat_with_retry = AsyncMock()
p.generation = GenerationSettings(max_tokens=100)
return p return p
@pytest.fixture @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 = MagicMock()
sessions.save = MagicMock() sessions.save = MagicMock()
# When maybe_consolidate_by_tokens refreshes the session reference via # When maybe_consolidate_by_tokens refreshes the session reference via
@ -38,13 +50,9 @@ def consolidator(store, mock_provider):
sessions._session_cache = _session_cache sessions._session_cache = _session_cache
return Consolidator( return Consolidator(
store=store, store=store,
provider=mock_provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=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: 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.""" """Consolidator should call LLM to summarize, then append to HISTORY.md."""
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="User fixed a bug in the auth module." content="User fixed a bug in the auth module."
@ -71,7 +113,7 @@ class TestConsolidatorSummarize:
{"role": "user", "content": "fix the auth bug"}, {"role": "user", "content": "fix the auth bug"},
{"role": "assistant", "content": "Done, fixed the race condition."}, {"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." assert result == "User fixed a bug in the auth module."
entries = store.read_unprocessed_history(since_cursor=0) entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1 assert len(entries) == 1
@ -81,6 +123,7 @@ class TestConsolidatorSummarize:
consolidator, consolidator,
mock_provider, mock_provider,
store, store,
runtime,
): ):
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="User fixed a bug in the auth module.", content="User fixed a bug in the auth module.",
@ -88,16 +131,22 @@ class TestConsolidatorSummarize:
) )
messages = [{"role": "user", "content": "fix the auth bug"}] 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) entries = store.read_unprocessed_history(since_cursor=0)
assert entries[0]["session_key"] == "telegram:chat-1" 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.""" """On LLM failure, raw-dump messages to HISTORY.md."""
mock_provider.chat_with_retry.side_effect = Exception("API error") mock_provider.chat_with_retry.side_effect = Exception("API error")
messages = [{"role": "user", "content": "hello"}] 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 assert result is None # no summary on raw dump fallback
entries = store.read_unprocessed_history(since_cursor=0) entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1 assert len(entries) == 1
@ -108,17 +157,22 @@ class TestConsolidatorSummarize:
consolidator, consolidator,
mock_provider, mock_provider,
store, store,
runtime,
): ):
mock_provider.chat_with_retry.side_effect = Exception("API error") mock_provider.chat_with_retry.side_effect = Exception("API error")
messages = [{"role": "user", "content": "hello"}] 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) entries = store.read_unprocessed_history(since_cursor=0)
assert entries[0]["session_key"] == "slack:chat-2" assert entries[0]["session_key"] == "slack:chat-2"
async def test_summarize_skips_empty_messages(self, consolidator): async def test_summarize_skips_empty_messages(self, consolidator, runtime):
result = await consolidator.archive([]) result = await consolidator.archive([], runtime=runtime)
assert result is None assert result is None
@ -139,7 +193,9 @@ class TestConsolidatorArchiveErrorHandling:
See https://github.com/HKUDS/nanobot/issues/3244 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.""" """LLM returning finish_reason='error' should trigger raw_archive, not write error text."""
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="Error: {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'overloaded_error (529)'}}", 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": "user", "content": "fix the auth bug"},
{"role": "assistant", "content": "Done, fixed the race condition."}, {"role": "assistant", "content": "Done, fixed the race condition."},
] ]
result = await consolidator.archive(messages) result = await consolidator.archive(messages, runtime=runtime)
assert result is None assert result is None
entries = store.read_unprocessed_history(since_cursor=0) entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1 assert len(entries) == 1
assert "[RAW]" in entries[0]["content"] assert "[RAW]" in entries[0]["content"]
assert "Error:" not 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.""" """Normal LLM response should still produce a proper summary entry."""
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="User fixed a bug in the auth module.", content="User fixed a bug in the auth module.",
@ -166,7 +224,7 @@ class TestConsolidatorArchiveErrorHandling:
{"role": "user", "content": "fix the auth bug"}, {"role": "user", "content": "fix the auth bug"},
{"role": "assistant", "content": "Done."}, {"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." assert result == "User fixed a bug in the auth module."
entries = store.read_unprocessed_history(since_cursor=0) entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1 assert len(entries) == 1
@ -174,7 +232,9 @@ class TestConsolidatorArchiveErrorHandling:
class TestConsolidatorTokenBudget: 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.""" """No consolidation when tokens are within budget."""
session = MagicMock() session = MagicMock()
session.last_consolidated = 0 session.last_consolidated = 0
@ -183,10 +243,10 @@ class TestConsolidatorTokenBudget:
consolidator.sessions._session_cache[session.key] = session consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken")) consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
consolidator.archive = AsyncMock(return_value=True) 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() 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.""" """Consolidation pressure must see messages hidden by the replay window."""
session = Session(key="test:full-tail") session = Session(key="test:full-tail")
for i in range(160): for i in range(160):
@ -200,7 +260,7 @@ class TestConsolidatorTokenBudget:
consolidator._build_messages = build_messages 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 len(captured["history"]) == 160
assert captured["history"][0]["content"].endswith("msg-0") 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( async def test_replay_window_overflow_is_archived_even_under_token_budget(
self, self,
consolidator, consolidator,
runtime,
): ):
"""Old messages that cannot be replayed should be materialized first.""" """Old messages that cannot be replayed should be materialized first."""
consolidator._SAFETY_BUFFER = 0 consolidator._SAFETY_BUFFER = 0
@ -222,6 +283,7 @@ class TestConsolidatorTokenBudget:
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime,
replay_max_messages=6, replay_max_messages=6,
) )
@ -235,6 +297,7 @@ class TestConsolidatorTokenBudget:
async def test_replay_window_overflow_extends_to_long_recent_user_turn( async def test_replay_window_overflow_extends_to_long_recent_user_turn(
self, self,
consolidator, consolidator,
runtime,
): ):
"""Replay-window consolidation must not cut into the latest user turn.""" """Replay-window consolidation must not cut into the latest user turn."""
session = Session(key="test:replay-tool-boundary") session = Session(key="test:replay-tool-boundary")
@ -251,6 +314,7 @@ class TestConsolidatorTokenBudget:
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime,
replay_max_messages=4, replay_max_messages=4,
) )
@ -266,6 +330,7 @@ class TestConsolidatorTokenBudget:
async def test_replay_window_overflow_uses_newer_user_inside_window( async def test_replay_window_overflow_uses_newer_user_inside_window(
self, self,
consolidator, consolidator,
runtime,
): ):
"""Do not extend to an older long turn when the hard window has a newer user.""" """Do not extend to an older long turn when the hard window has a newer user."""
session = Session(key="test:replay-newer-user") session = Session(key="test:replay-newer-user")
@ -284,6 +349,7 @@ class TestConsolidatorTokenBudget:
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime,
replay_max_messages=6, replay_max_messages=6,
) )
@ -295,7 +361,7 @@ class TestConsolidatorTokenBudget:
history = session.get_history(max_messages=6, extend_to_user=True) history = session.get_history(max_messages=6, extend_to_user=True)
assert [m["content"] for m in history] == ["new question", "new answer"] 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.""" """Without chunk cap, the full range from pick_consolidation_boundary is archived."""
consolidator._SAFETY_BUFFER = 0 consolidator._SAFETY_BUFFER = 0
session = MagicMock() session = MagicMock()
@ -316,14 +382,16 @@ class TestConsolidatorTokenBudget:
# (user message at 50, token budget met) # (user message at 50, token budget met)
consolidator.archive = AsyncMock(return_value=True) 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] archived_chunk = consolidator.archive.await_args.args[0]
# pick_consolidation_boundary returns (50, tokens) — user turn at idx 50 # pick_consolidation_boundary returns (50, tokens) — user turn at idx 50
assert archived_chunk[0]["content"] == "m0" assert archived_chunk[0]["content"] == "m0"
assert session.last_consolidated > 0 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 """When archive() falls back to raw-archive (LLM failed), the cursor
must still advance. Otherwise the same chunk gets raw-archived again must still advance. Otherwise the same chunk gets raw-archived again
on every subsequent maybe_consolidate_by_tokens() call, spamming on every subsequent maybe_consolidate_by_tokens() call, spamming
@ -344,14 +412,16 @@ class TestConsolidatorTokenBudget:
# LLM consolidation fails — archive() returns None (raw_archive fired). # LLM consolidation fails — archive() returns None (raw_archive fired).
consolidator.archive = AsyncMock(return_value=None) 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() consolidator.archive.assert_awaited_once()
# The chunk is considered "materialized" (as a raw-archive breadcrumb), # The chunk is considered "materialized" (as a raw-archive breadcrumb),
# so last_consolidated must have moved past it. # so last_consolidated must have moved past it.
assert session.last_consolidated == 50 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 """A degraded LLM should not trigger more archive() calls within the
same maybe_consolidate_by_tokens invocation bail after one fallback.""" same maybe_consolidate_by_tokens invocation bail after one fallback."""
consolidator._SAFETY_BUFFER = 0 consolidator._SAFETY_BUFFER = 0
@ -370,12 +440,14 @@ class TestConsolidatorTokenBudget:
) )
consolidator.archive = AsyncMock(return_value=None) 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. # Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS.
assert consolidator.archive.await_count == 1 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.""" """When boundary points past a long tool chain, the full chunk is archived."""
consolidator._SAFETY_BUFFER = 0 consolidator._SAFETY_BUFFER = 0
session = MagicMock() session = MagicMock()
@ -394,7 +466,7 @@ class TestConsolidatorTokenBudget:
) )
consolidator.archive = AsyncMock(return_value=True) 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() consolidator.archive.assert_awaited_once()
# pick_consolidation_boundary finds the only boundary at idx=61 # pick_consolidation_boundary finds the only boundary at idx=61
@ -412,17 +484,15 @@ class TestCompactIdleSession:
sessions = SessionManager(store.workspace) sessions = SessionManager(store.workspace)
return Consolidator( return Consolidator(
store=store, store=store,
provider=mock_provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
max_completion_tokens=100,
) )
@pytest.mark.asyncio @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, """20 user/assistant turns → compact with max_suffix=8 → messages ≤ 8,
last_consolidated=0, _last_summary stored.""" last_consolidated=0, _last_summary stored."""
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
@ -437,7 +507,9 @@ class TestCompactIdleSession:
session.updated_at = old_ts session.updated_at = old_ts
sessions.save(session) 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." assert result == "Summary of old conversation."
reloaded = sessions.get_or_create("cli:test") reloaded = sessions.get_or_create("cli:test")
@ -451,7 +523,7 @@ class TestCompactIdleSession:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_summarizes_retained_suffix_not_just_dropped_prefix( 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 """idleCompact must summarize over the full unconsolidated tail, including
the recent suffix it retains. Otherwise a late user correction / final 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") session.add_message("assistant", "CORRECTED_FINAL_RESULT_alpha")
sessions.save(session) 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"] summarized = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
assert "CORRECTED_FINAL_RESULT_alpha" in summarized assert "CORRECTED_FINAL_RESULT_alpha" in summarized
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_raw_dumps_only_dropped_messages_on_llm_failure( 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 """Summarizing over the full tail must not widen what gets raw-dumped on
LLM failure: the breadcrumb should contain only the removed prefix, not LLM failure: the breadcrumb should contain only the removed prefix, not
@ -492,7 +566,9 @@ class TestCompactIdleSession:
session.add_message("assistant", "RETAINED_SUFFIX_marker") session.add_message("assistant", "RETAINED_SUFFIX_marker")
sessions.save(session) 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)) raw = "\n".join(e["content"] for e in store.read_unprocessed_history(since_cursor=0))
assert "[RAW]" in raw assert "[RAW]" in raw
@ -505,6 +581,7 @@ class TestCompactIdleSession:
real_consolidator, real_consolidator,
mock_provider, mock_provider,
store, store,
runtime,
): ):
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary of old conversation.", finish_reason="stop" content="Summary of old conversation.", finish_reason="stop"
@ -515,14 +592,16 @@ class TestCompactIdleSession:
session.add_message("assistant", f"assistant msg {i}") session.add_message("assistant", f"assistant msg {i}")
real_consolidator.sessions.save(session) 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) entries = store.read_unprocessed_history(since_cursor=0)
assert entries[0]["session_key"] == "cli:test" assert entries[0]["session_key"] == "cli:test"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_empty_session_does_not_refresh_timestamp( 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.""" """Empty session with old updated_at does not look active after compaction."""
from datetime import datetime, timedelta from datetime import datetime, timedelta
@ -533,7 +612,9 @@ class TestCompactIdleSession:
session.updated_at = old_ts session.updated_at = old_ts
sessions.save(session) 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 == "" assert result == ""
reloaded = sessions.get_or_create("cli:empty") reloaded = sessions.get_or_create("cli:empty")
@ -541,7 +622,9 @@ class TestCompactIdleSession:
assert reloaded.metadata == {} assert reloaded.metadata == {}
@pytest.mark.asyncio @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.""" """LLM returns '(nothing)' → _last_summary NOT in metadata."""
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="(nothing)", finish_reason="stop" content="(nothing)", finish_reason="stop"
@ -553,14 +636,18 @@ class TestCompactIdleSession:
session.add_message("assistant", f"a{i}") session.add_message("assistant", f"a{i}")
sessions.save(session) 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)" assert result == "(nothing)"
reloaded = sessions.get_or_create("cli:nothing") reloaded = sessions.get_or_create("cli:nothing")
assert "_last_summary" not in reloaded.metadata assert "_last_summary" not in reloaded.metadata
@pytest.mark.asyncio @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.""" """LLM raises RuntimeError → raw_archive fires, session still truncated, returns None."""
mock_provider.chat_with_retry.side_effect = RuntimeError("LLM unavailable") mock_provider.chat_with_retry.side_effect = RuntimeError("LLM unavailable")
sessions = real_consolidator.sessions sessions = real_consolidator.sessions
@ -570,7 +657,9 @@ class TestCompactIdleSession:
session.add_message("assistant", f"a{i}") session.add_message("assistant", f"a{i}")
sessions.save(session) 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 assert result is None
# raw_archive should have been called (history.jsonl gets an entry) # raw_archive should have been called (history.jsonl gets an entry)
@ -582,7 +671,9 @@ class TestCompactIdleSession:
assert len(reloaded.messages) <= 4 assert len(reloaded.messages) <= 4
@pytest.mark.asyncio @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.""" """30 turns with last_consolidated=50 → only unconsolidated tail considered."""
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="Tail summary.", finish_reason="stop" content="Tail summary.", finish_reason="stop"
@ -595,7 +686,9 @@ class TestCompactIdleSession:
session.last_consolidated = 50 # Only 10 messages unconsolidated session.last_consolidated = 50 # Only 10 messages unconsolidated
sessions.save(session) 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." assert result == "Tail summary."
# Verify only the unconsolidated tail was processed: # Verify only the unconsolidated tail was processed:
@ -611,6 +704,7 @@ class TestCompactIdleSession:
self, self,
real_consolidator, real_consolidator,
mock_provider, mock_provider,
runtime,
): ):
"""Assistant-only tails extend back to the latest user turn, so archive """Assistant-only tails extend back to the latest user turn, so archive
the actual dropped messages rather than a computed prefix.""" the actual dropped messages rather than a computed prefix."""
@ -625,7 +719,9 @@ class TestCompactIdleSession:
session.add_message("assistant", f"assistant-{i:02d}") session.add_message("assistant", f"assistant-{i:02d}")
sessions.save(session) 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." assert result == "Tail summary."
reloaded = sessions.get_or_create("cli:noncontiguous") reloaded = sessions.get_or_create("cli:noncontiguous")
@ -653,7 +749,9 @@ class TestCompactIdleSession:
assert "user-14" in user_content assert "user-14" in user_content
@pytest.mark.asyncio @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.""" """Verify lock is held during execution."""
import asyncio import asyncio
@ -679,7 +777,9 @@ class TestCompactIdleSession:
assert not lock.locked() assert not lock.locked()
task = asyncio.ensure_future( 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() await started.wait()
assert lock.locked() assert lock.locked()
@ -702,15 +802,17 @@ class TestConsolidatorSessionRefresh:
provider.chat_with_retry = AsyncMock( provider.chat_with_retry = AsyncMock(
return_value=MagicMock(content="summary", finish_reason="stop") 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")) provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test"))
runtime = LLMRuntime.capture(
provider,
"test-model",
context_window_tokens=128_000,
)
sessions = SessionManager(tmp_path) sessions = SessionManager(tmp_path)
consolidator = Consolidator( consolidator = Consolidator(
store=store, store=store,
provider=provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=128_000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
) )
@ -722,13 +824,16 @@ class TestConsolidatorSessionRefresh:
seen: dict[str, Session] = {} seen: dict[str, Session] = {}
def estimate(session: Session): def estimate(session: Session, *, runtime):
seen["session"] = session seen["session"] = session
return 10, "test" return 10, "test"
consolidator.estimate_session_prompt_tokens = MagicMock(side_effect=estimate) 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 assert seen["session"] is fresh
@ -745,15 +850,17 @@ class TestConsolidatorSessionRefresh:
provider.chat_with_retry = AsyncMock( provider.chat_with_retry = AsyncMock(
return_value=MagicMock(content="summary", finish_reason="stop") 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")) provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test"))
runtime = LLMRuntime.capture(
provider,
"test-model",
context_window_tokens=128_000,
)
sessions = SessionManager(tmp_path) sessions = SessionManager(tmp_path)
consolidator = Consolidator( consolidator = Consolidator(
store=store, store=store,
provider=provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=128_000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
) )
@ -769,11 +876,18 @@ class TestConsolidatorSessionRefresh:
old_ref = session old_ref = session
# AutoCompact runs first and truncates to 8 # 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 — # Background consolidation runs with stale reference —
# should detect the session was replaced and not undo the compact. # 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") session_after = sessions.get_or_create("cli:test")
# Messages should still be truncated (not restored to 40) # Messages should still be truncated (not restored to 40)
@ -818,7 +932,9 @@ class TestRawArchiveTruncation:
class TestArchiveTruncation: class TestArchiveTruncation:
"""archive() must truncate formatted text before sending to consolidation LLM.""" """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.""" """Large formatted text should be truncated to token budget before LLM call."""
# context_window_tokens=1000, max_completion_tokens=100, _SAFETY_BUFFER=1024 # context_window_tokens=1000, max_completion_tokens=100, _SAFETY_BUFFER=1024
# budget = 1000 - 100 - 1024 = -124 → fallback via truncate_text(budget*4) # 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( mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary of large input.", finish_reason="stop" 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 call_args = mock_provider.chat_with_retry.call_args
user_content = call_args.kwargs["messages"][1]["content"] user_content = call_args.kwargs["messages"][1]["content"]
# Should be significantly shorter than 100K # Should be significantly shorter than 100K
assert len(user_content) < 50_000 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.""" """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}] big_messages = [{"role": "user", "content": "word " * 50_000}]
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary.", finish_reason="stop" 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"] sent_messages = mock_provider.chat_with_retry.call_args.kwargs["messages"]
user_content = sent_messages[1]["content"] user_content = sent_messages[1]["content"]
@ -848,7 +966,9 @@ class TestArchiveTruncation:
# Should be truncated # Should be truncated
assert len(user_content) < 250_000 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 """A pathologically large LLM summary must not land full-length in
history.jsonl that would re-open the #3412 bloat vector from the history.jsonl that would re-open the #3412 bloat vector from the
*success* path instead of the fallback path.""" *success* path instead of the fallback path."""
@ -856,21 +976,26 @@ class TestArchiveTruncation:
content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10), content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10),
finish_reason="stop", 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] entry = store.read_unprocessed_history(since_cursor=0)[0]
assert len(entry["content"]) <= _ARCHIVE_SUMMARY_MAX_CHARS + 50 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.""" """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 consolidator._SAFETY_BUFFER = 0
# budget = 10000 - 100 - 0 = 9900 tokens # budget = 10000 - 100 - 0 = 9900 tokens
big_messages = [{"role": "user", "content": "word " * 50_000}] big_messages = [{"role": "user", "content": "word " * 50_000}]
mock_provider.chat_with_retry.return_value = MagicMock( mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary.", finish_reason="stop" content="Summary.", finish_reason="stop"
) )
await consolidator.archive(big_messages) await consolidator.archive(big_messages, runtime=runtime)
import tiktoken import tiktoken
enc = tiktoken.get_encoding("cl100k_base") enc = tiktoken.get_encoding("cl100k_base")

View File

@ -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} 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"]]) 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] archived_chunk = loop.consolidator.archive.await_args.args[0]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"] 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) loop.sessions.save(session)
call_count = [0] call_count = [0]
def mock_estimate(_session, *, session_summary=None): def mock_estimate(_session, *, runtime):
call_count[0] += 1 call_count[0] += 1
if call_count[0] == 1: if call_count[0] == 1:
return (500, "test") 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] loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100) 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 loop.consolidator.archive.await_count == 2
assert session.last_consolidated == 6 assert session.last_consolidated == 6
@ -139,7 +145,7 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path,
call_count = [0] call_count = [0]
def mock_estimate(_session, *, session_summary=None): def mock_estimate(_session, *, runtime):
call_count[0] += 1 call_count[0] += 1
if call_count[0] == 1: if call_count[0] == 1:
return (500, "test") 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] loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100) 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 loop.consolidator.archive.await_count == 2
assert session.last_consolidated == 6 assert session.last_consolidated == 6
@ -171,7 +180,7 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path,
call_count = [0] call_count = [0]
def mock_estimate(_session, *, session_summary=None): def mock_estimate(_session, *, runtime):
call_count[0] += 1 call_count[0] += 1
if call_count[0] == 1: if call_count[0] == 1:
return (500, "test") 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] loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 150) 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") reloaded = loop.sessions.get_or_create("cli:test")
meta = reloaded.metadata.get("_last_summary") 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.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) # type: ignore[method-assign]
loop._schedule_background = lambda coro: coro.close() # 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( loop.consolidator.maybe_consolidate_by_tokens.assert_any_await(
session, session,
runtime=runtime,
replay_max_messages=loop._max_messages, 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 @pytest.mark.asyncio
@ -221,7 +240,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
archived_session_keys: list[str | None] = [] 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") order.append("consolidate")
archived_session_keys.append(session_key) archived_session_keys.append(session_key)
return True 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) monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500)
call_count = [0] call_count = [0]
def mock_estimate(_session, *, session_summary=None): def mock_estimate(_session, *, runtime):
call_count[0] += 1 call_count[0] += 1
return (1000 if call_count[0] <= 1 else 80, "test") return (1000 if call_count[0] <= 1 else 80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]

View File

@ -1308,6 +1308,11 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_
assert seen["runtime"] is runtime assert seen["runtime"] is runtime
record_runtime.assert_called_once_with("cli:test", 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"] initial_messages = seen["initial_messages"]
assert isinstance(initial_messages, list) assert isinstance(initial_messages, list)
non_system = [m for m in initial_messages if m.get("role") != "system"] non_system = [m for m in initial_messages if m.get("role") != "system"]

View File

@ -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, "provider")
assert not hasattr(loop.subagents, "model") assert not hasattr(loop.subagents, "model")
assert not hasattr(loop.subagents.runner, "provider") assert not hasattr(loop.subagents.runner, "provider")
assert loop.consolidator.provider is new_provider assert not hasattr(loop.consolidator, "provider")
assert loop.consolidator.model == "new-model" assert not hasattr(loop.consolidator, "model")
assert loop.consolidator.context_window_tokens == 2000 assert not hasattr(loop.consolidator, "context_window_tokens")
assert loop.consolidator.max_completion_tokens == 456 assert not hasattr(loop.consolidator, "max_completion_tokens")
def test_llm_runtime_refreshes_provider_snapshot(tmp_path: Path) -> None: 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 payload["requires_restart"] is False
assert loop.context_window_tokens == 262_144 assert loop.context_window_tokens == 262_144
assert loop.consolidator.context_window_tokens == 262_144 assert loop.llm_runtime().context_window_tokens == 262_144

View File

@ -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.max_tokens == 4096
assert loop.provider.generation.reasoning_effort == "low" assert loop.provider.generation.reasoning_effort == "low"
assert not hasattr(loop.subagents, "model") assert not hasattr(loop.subagents, "model")
assert loop.consolidator.model == "openai/gpt-4.1" assert not hasattr(loop.consolidator, "model")
assert loop.consolidator.context_window_tokens == 32_768 assert not hasattr(loop.consolidator, "context_window_tokens")
assert loop.consolidator.max_completion_tokens == 4096 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: 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.runner, "provider")
assert not hasattr(loop.subagents, "provider") assert not hasattr(loop.subagents, "provider")
assert not hasattr(loop.subagents.runner, "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.model == "anthropic/claude-opus-4-5"
assert loop.context_window_tokens == 200_000 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: 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_preset is None
assert loop.model == "base-model" assert loop.model == "base-model"
assert not hasattr(loop.subagents, "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.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: def test_active_model_preset_survives_unchanged_config_refresh(tmp_path) -> None:

View File

@ -25,8 +25,10 @@ from nanobot.bus.queue import MessageBus
from nanobot.command.builtin import cmd_new, register_builtin_commands from nanobot.command.builtin import cmd_new, register_builtin_commands
from nanobot.command.router import CommandContext, CommandRouter from nanobot.command.router import CommandContext, CommandRouter
from nanobot.config.schema import AgentDefaults, Config from nanobot.config.schema import AgentDefaults, Config
from nanobot.providers.base import GenerationSettings
from nanobot.session.keys import UNIFIED_SESSION_KEY from nanobot.session.keys import UNIFIED_SESSION_KEY
from nanobot.session.manager import Session, SessionManager from nanobot.session.manager import Session, SessionManager
from nanobot.utils.llm_runtime import LLMRuntime
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Helpers # Helpers
@ -50,6 +52,15 @@ def _make_loop(tmp_path: Path, unified_session: bool = False) -> AgentLoop:
return loop 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", def _make_msg(channel: str = "telegram", chat_id: str = "111",
session_key_override: str | None = None) -> InboundMessage: session_key_override: str | None = None) -> InboundMessage:
return InboundMessage( return InboundMessage(
@ -233,14 +244,17 @@ class TestCmdNewUnifiedSession:
shared.add_message("assistant", "hi there") shared.add_message("assistant", "hi there")
sessions.save(shared) sessions.save(shared)
assert len(sessions.get_or_create("unified:default").messages) == 2 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 # _schedule_background is a *sync* method that schedules a coroutine via
# asyncio.create_task(). Mirror that exactly so the coroutine is consumed # asyncio.create_task(). Mirror that exactly so the coroutine is consumed
# and no RuntimeWarning is emitted. # and no RuntimeWarning is emitted.
admitted_runtime = MagicMock(name="admitted_runtime")
loop = SimpleNamespace( loop = SimpleNamespace(
sessions=sessions, sessions=sessions,
consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)),
_cancel_active_tasks=AsyncMock(return_value=0), _cancel_active_tasks=AsyncMock(return_value=0),
llm_runtime=MagicMock(return_value=MagicMock()),
) )
loop._schedule_background = lambda coro: asyncio.ensure_future(coro) 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", channel="telegram", sender_id="user1", chat_id="111", content="/new",
session_key_override="unified:default", # as _dispatch() would set it 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) result = await cmd_new(ctx)
@ -257,6 +278,12 @@ class TestCmdNewUnifiedSession:
sessions.invalidate("unified:default") sessions.invalidate("unified:default")
reloaded = sessions.get_or_create("unified:default") reloaded = sessions.get_or_create("unified:default")
assert reloaded.messages == [] 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 @pytest.mark.asyncio
async def test_cmd_new_in_unified_mode_does_not_affect_other_sessions(self, tmp_path: Path): 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, sessions=sessions,
consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)),
_cancel_active_tasks=AsyncMock(return_value=0), _cancel_active_tasks=AsyncMock(return_value=0),
llm_runtime=MagicMock(return_value=MagicMock()),
) )
loop._schedule_background = lambda coro: asyncio.ensure_future(coro) loop._schedule_background = lambda coro: asyncio.ensure_future(coro)
@ -306,26 +334,23 @@ class TestConsolidationUnaffectedByUnifiedSession:
store = MagicMock(spec=MemoryStore) store = MagicMock(spec=MemoryStore)
mock_provider = MagicMock() mock_provider = MagicMock()
mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary")) 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, # Use spec= so MagicMock doesn't auto-generate AsyncMock for non-async methods,
# which would leave unawaited coroutines and trigger RuntimeWarning. # which would leave unawaited coroutines and trigger RuntimeWarning.
sessions = MagicMock(spec=SessionManager) sessions = MagicMock(spec=SessionManager)
consolidator = Consolidator( consolidator = Consolidator(
store=store, store=store,
provider=mock_provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
max_completion_tokens=100,
) )
consolidator.archive = AsyncMock() consolidator.archive = AsyncMock()
session = Session(key="unified:default") session = Session(key="unified:default")
session.messages = [] session.messages = []
await consolidator.maybe_consolidate_by_tokens(session) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive.assert_not_called() consolidator.archive.assert_not_called()
@ -341,24 +366,24 @@ class TestConsolidationUnaffectedByUnifiedSession:
store = MagicMock(spec=MemoryStore) store = MagicMock(spec=MemoryStore)
mock_provider = MagicMock() mock_provider = MagicMock()
mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary")) mock_provider.chat_with_retry = AsyncMock(return_value=MagicMock(content="summary"))
runtime = _runtime(mock_provider)
sessions = MagicMock(spec=SessionManager) sessions = MagicMock(spec=SessionManager)
consolidator = Consolidator( consolidator = Consolidator(
store=store, store=store,
provider=mock_provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
max_completion_tokens=100,
) )
session = Session(key=key) session = Session(key=key)
session.messages = [] # empty → exits immediately for both keys session.messages = [] # empty → exits immediately for both keys
consolidator.archive = AsyncMock() 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 archive_calls[key] = consolidator.archive.call_count
assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0 assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0
@ -371,17 +396,14 @@ class TestConsolidationUnaffectedByUnifiedSession:
store = MagicMock(spec=MemoryStore) store = MagicMock(spec=MemoryStore)
mock_provider = MagicMock() mock_provider = MagicMock()
runtime = _runtime(mock_provider)
sessions = MagicMock(spec=SessionManager) sessions = MagicMock(spec=SessionManager)
consolidator = Consolidator( consolidator = Consolidator(
store=store, store=store,
provider=mock_provider,
model="test-model",
sessions=sessions, sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
max_completion_tokens=100,
) )
session = Session(key="unified:default") session = Session(key="unified:default")
@ -394,11 +416,12 @@ class TestConsolidationUnaffectedByUnifiedSession:
consolidator.pick_consolidation_boundary = MagicMock(return_value=None) consolidator.pick_consolidation_boundary = MagicMock(return_value=None)
consolidator.archive = AsyncMock() 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) # estimate was called (consolidation was attempted)
consolidator.estimate_session_prompt_tokens.assert_called_once_with( consolidator.estimate_session_prompt_tokens.assert_called_once_with(
session, session,
runtime=runtime,
) )
# but archive was not called (no valid boundary) # but archive was not called (no valid boundary)
consolidator.archive.assert_not_called() consolidator.archive.assert_not_called()

View File

@ -244,8 +244,16 @@ class TestRestartCommand:
loop.subagents.get_running_count_by_session.return_value = 0 loop.subagents.get_running_count_by_session.return_value = 0
msg = InboundMessage(channel="telegram", sender_id="u1", chat_id="c1", content="/status") 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 response is not None
assert "Model: test-model" in response.content assert "Model: test-model" in response.content
@ -255,6 +263,10 @@ class TestRestartCommand:
assert "Uptime: 2m 5s" in response.content assert "Uptime: 2m 5s" in response.content
assert "Tasks: 0 active" in response.content assert "Tasks: 0 active" in response.content
assert response.metadata == {"render_as": "text"} assert response.metadata == {"render_as": "text"}
loop.consolidator.estimate_session_prompt_tokens.assert_called_once_with(
session,
runtime=runtime,
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_status_counts_running_dispatch_and_subagent_tasks(self): async def test_status_counts_running_dispatch_and_subagent_tasks(self):

View File

@ -86,7 +86,8 @@ async def test_model_command_switches_preset(tmp_path) -> None:
assert loop.model_preset == "fast" assert loop.model_preset == "fast"
assert loop.model == "openai/gpt-4.1" assert loop.model == "openai/gpt-4.1"
assert not hasattr(loop.subagents, "model") 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 @pytest.mark.asyncio

View File

@ -1213,11 +1213,16 @@ async def test_runtime_helpers_expose_model_workspace_and_compact(tmp_path):
config_path = _write_config(tmp_path) config_path = _write_config(tmp_path)
bot = Nanobot.from_config(config_path, workspace=tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path)
await bot.sessions.ingest("sdk:history", [{"role": "user", "content": "hello"}]) 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() bot._loop.consolidator.maybe_consolidate_by_tokens = AsyncMock()
snapshot = await bot.runtime.compact_session("sdk:history") snapshot = await bot.runtime.compact_session("sdk:history")
assert snapshot.key == "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.model == bot._loop.model
assert bot.runtime.workspace == tmp_path 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." assert summary == "Summary."
bot._loop.consolidator.compact_idle_session.assert_awaited_once_with( bot._loop.consolidator.compact_idle_session.assert_awaited_once_with(
"sdk:history", "sdk:history",
runtime=runtime,
max_suffix=4, max_suffix=4,
) )