Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(agent): require runtime for consolidation

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