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

refactor(memory): unify structured consolidation flow

Browse Source
This commit is contained in:
chengyongru
2026-08-19 18:40:20 +08:00
committed by chengyongru
parent d4de0e4e3d
commit 16be6b0f36
14 changed files with 257 additions and 340 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/autocompact.py
+3 -3
View File
@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Callable, Coroutine
from loguru import logger from loguru import logger
from nanobot.session.manager import MIN_COMPACTED_REPLAY_MESSAGES, Session, SessionManager from nanobot.session.manager import MIN_COMPACTED_REPLAY_MESSAGES, Session, SessionManager
from nanobot.session.summary import SessionSummary from nanobot.session.summary import SessionSummary, session_summary_from_metadata
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.agent.memory import Consolidator from nanobot.agent.memory import Consolidator
@@ -91,7 +91,7 @@ class AutoCompact:
) )
if summary and summary != "(nothing)": if summary and summary != "(nothing)":
session = self.sessions.get_or_create(key) session = self.sessions.get_or_create(key)
stored = SessionSummary.from_metadata( stored = session_summary_from_metadata(
session.metadata, session.metadata,
fallback_last_active=session.updated_at, fallback_last_active=session.updated_at,
) )
@@ -117,7 +117,7 @@ class AutoCompact:
# Cold path: summary persisted in session metadata (process restarted). # Cold path: summary persisted in session metadata (process restarted).
# Persisted metadata may outlive schema changes; a malformed summary must # Persisted metadata may outlive schema changes; a malformed summary must
# not abort turn preparation. # not abort turn preparation.
return session, SessionSummary.from_metadata( return session, session_summary_from_metadata(
session.metadata, session.metadata,
fallback_last_active=session.updated_at, fallback_last_active=session.updated_at,
) )
nanobot/agent/context.py
+6 -2
View File
@@ -157,7 +157,11 @@ class ContextBuilder:
parts.append("# Recent History\n\n" + history_text) parts.append("# Recent History\n\n" + history_text)
if session_summary: if session_summary:
parts.append(f"[Archived Context Summary]\n\n{session_summary.for_prompt()}") parts.append(
"[Archived Context Summary]\n\n"
f"Previous conversation summary (last active {session_summary['last_active']}):\n"
f"{session_summary['text']}"
)
return "\n\n---\n\n".join(parts) return "\n\n---\n\n".join(parts)
@@ -175,7 +179,7 @@ class ContextBuilder:
entry = entries[index] entry = entries[index]
if ( if (
entry.get("session_key") == session_key entry.get("session_key") == session_key
and entry.get("content") == session_summary.text and entry.get("content") == session_summary["text"]
): ):
return [*entries[:index], *entries[index + 1:]] return [*entries[:index], *entries[index + 1:]]
return entries return entries
nanobot/agent/memory.py
+33 -54
View File
@@ -27,7 +27,7 @@ from nanobot.session.manager import (
SessionManager, SessionManager,
replay_max_messages_for_context, replay_max_messages_for_context,
) )
from nanobot.session.summary import SessionSummary from nanobot.session.summary import session_summary_from_metadata
from nanobot.utils.gitstore import GitStore from nanobot.utils.gitstore import GitStore
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
content_with_media_breadcrumbs, content_with_media_breadcrumbs,
@@ -38,7 +38,6 @@ from nanobot.utils.helpers import (
recent_message_start_index, recent_message_start_index,
strip_think, strip_think,
truncate_text, truncate_text,
truncate_text_to_tokens,
) )
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
from nanobot.utils.workspace_prompts import ( from nanobot.utils.workspace_prompts import (
@@ -928,10 +927,10 @@ class Consolidator:
len(chunk), len(chunk),
replay_max_messages, replay_max_messages,
) )
summary = await self.archive( summary = await self.archive_session(
chunk, session,
archive_end=end_idx,
runtime=runtime, runtime=runtime,
session_key=session.key,
) )
session.last_consolidated = end_idx session.last_consolidated = end_idx
session.provider_state = None session.provider_state = None
@@ -955,7 +954,7 @@ class Consolidator:
"""Estimate prompt size from the full replayable session history.""" """Estimate prompt size from the full replayable session history."""
history = self._full_replay_history(session) history = self._full_replay_history(session)
channel = session.key.split(":", 1)[0] if ":" in session.key else None channel = session.key.split(":", 1)[0] if ":" in session.key else None
summary = SessionSummary.from_metadata( summary = session_summary_from_metadata(
session.metadata, session.metadata,
fallback_last_active=session.updated_at, fallback_last_active=session.updated_at,
) )
@@ -982,52 +981,24 @@ class Consolidator:
- self._SAFETY_BUFFER - self._SAFETY_BUFFER
) )
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(runtime)
if budget <= 0:
return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS)
return truncate_text_to_tokens(text, budget)
async def archive( async def archive(
self, self,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
*, *,
runtime: LLMRuntime, runtime: LLMRuntime,
session_key: str | None = None, session_key: str,
summary_messages: list[dict[str, Any]] | None = None, request_messages: list[dict[str, Any]],
request_messages: list[dict[str, Any]] | None = None, request_tools: list[dict[str, Any]],
request_tools: list[dict[str, Any]] | None = None,
) -> str | None: ) -> str | None:
"""Summarize messages and append the result to history.jsonl. """Execute a prepared consolidation request and persist its result."""
``summary_messages`` adds context but is excluded from raw fallback.
``request_messages`` preserves a prebuilt model-facing prefix instead
of flattening the messages; tools are included but disabled.
"""
if not messages: if not messages:
return None return None
prebuilt_request = request_messages is not None
if request_messages is None:
formatted = MemoryStore._format_messages(
public_history_messages(
summary_messages if summary_messages is not None else messages
)
)
formatted = self._truncate_to_token_budget(formatted, runtime=runtime)
request_messages = [
{
"role": "system",
"content": render_template("agent/consolidator_archive.md", strip=True),
},
{"role": "user", "content": formatted},
]
try: try:
response = await runtime.provider.chat_with_retry( response = await runtime.provider.chat_with_retry(
model=runtime.model, model=runtime.model,
messages=request_messages, messages=request_messages,
tools=request_tools if prebuilt_request else None, tools=request_tools,
tool_choice="none" if prebuilt_request else None, tool_choice="none",
temperature=runtime.generation.temperature, temperature=runtime.generation.temperature,
max_tokens=runtime.generation.max_tokens, max_tokens=runtime.generation.max_tokens,
reasoning_effort=runtime.generation.reasoning_effort, reasoning_effort=runtime.generation.reasoning_effort,
@@ -1056,23 +1027,31 @@ class Consolidator:
) )
return summary return summary
async def _archive_idle_tail( async def archive_session(
self, self,
session: Session, session: Session,
messages: list[dict[str, Any]],
*, *,
archive_end: int,
runtime: LLMRuntime, runtime: LLMRuntime,
) -> str | None: ) -> str | None:
"""Archive an idle tail by extending the ordinary model-facing messages.""" """Archive a session prefix by appending a consolidation instruction."""
messages = list(session.messages[session.last_consolidated:archive_end])
if not messages:
return None
budget = self._input_token_budget(runtime) budget = self._input_token_budget(runtime)
if budget <= 0: if budget <= 0:
logger.debug( logger.debug(
"Idle consolidation has no safe input budget for {}; raw-dumping", "Consolidation has no safe input budget for {}; raw-dumping",
session.key, session.key,
) )
self.store.raw_archive(messages, session_key=session.key) self.store.raw_archive(messages, session_key=session.key)
return None return None
history = session.get_history( prefix = Session(
key=session.key,
messages=list(session.messages[:archive_end]),
last_consolidated=session.last_consolidated,
)
history = prefix.get_history(
max_messages=replay_max_messages_for_context(runtime.context_window_tokens), max_messages=replay_max_messages_for_context(runtime.context_window_tokens),
max_tokens=budget, max_tokens=budget,
) )
@@ -1085,7 +1064,7 @@ class Consolidator:
or history[-len(archive_history):] != archive_history or history[-len(archive_history):] != archive_history
): ):
logger.debug( logger.debug(
"Idle consolidation cannot replay the full tail for {}; raw-dumping", "Consolidation cannot replay the full chunk for {}; raw-dumping",
session.key, session.key,
) )
self.store.raw_archive(messages, session_key=session.key) self.store.raw_archive(messages, session_key=session.key)
@@ -1103,7 +1082,7 @@ class Consolidator:
history=history, history=history,
current_message=prompt, current_message=prompt,
channel=channel, channel=channel,
session_summary=SessionSummary.from_metadata( session_summary=session_summary_from_metadata(
session.metadata, session.metadata,
fallback_last_active=session.updated_at, fallback_last_active=session.updated_at,
), ),
@@ -1120,7 +1099,7 @@ class Consolidator:
) )
if estimated > budget: if estimated > budget:
logger.debug( logger.debug(
"Idle consolidation prefix exceeds budget for {}; raw-dumping: {}/{} via {}", "Consolidation prefix exceeds budget for {}; raw-dumping: {}/{} via {}",
session.key, session.key,
estimated, estimated,
budget, budget,
@@ -1215,13 +1194,13 @@ class Consolidator:
source, source,
len(chunk), len(chunk),
) )
summary = await self.archive( summary = await self.archive_session(
chunk, session,
archive_end=end_idx,
runtime=runtime, 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_session() 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
# would just emit duplicate [RAW] entries. # would just emit duplicate [RAW] entries.
if summary: if summary:
@@ -1278,9 +1257,9 @@ class Consolidator:
last_active = session.updated_at last_active = session.updated_at
archive_end = archive_start + len(messages_to_archive) archive_end = archive_start + len(messages_to_archive)
summary = await self._archive_idle_tail( summary = await self.archive_session(
session, session,
messages_to_archive, archive_end=archive_end,
runtime=runtime, runtime=runtime,
) )
nanobot/command/builtin.py
+13 -5
View File
@@ -8,7 +8,7 @@ import subprocess
import sys import sys
import time import time
from contextlib import suppress from contextlib import suppress
from dataclasses import dataclass from dataclasses import dataclass, replace
from typing import TYPE_CHECKING, Any, Literal, cast from typing import TYPE_CHECKING, Any, Literal, cast
from nanobot import __version__ from nanobot import __version__
@@ -307,18 +307,26 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
loop.discard_session_file_state(ctx.key) loop.discard_session_file_state(ctx.key)
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or loop.sessions.get_or_create(ctx.key)
snapshot = session.messages[session.last_consolidated:] snapshot = session.messages[session.last_consolidated:]
archive_snapshot = None
runtime = None runtime = None
if snapshot: if snapshot:
runtime = ctx.runtime or loop.runtime_for_session(session) runtime = ctx.runtime or loop.runtime_for_session(session)
archive_snapshot = replace(
session,
messages=snapshot,
metadata=dict(session.metadata),
last_consolidated=0,
provider_state=None,
)
session.clear() session.clear()
loop.sessions.save(session) loop.sessions.save(session)
loop.sessions.invalidate(session.key) loop.sessions.invalidate(session.key)
if snapshot and runtime is not None: if archive_snapshot is not None and runtime is not None:
loop.schedule_background( loop.schedule_background(
loop.consolidator.archive( # pyright: ignore[reportUnknownMemberType] loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
snapshot, archive_snapshot,
archive_end=len(snapshot),
runtime=runtime, runtime=runtime,
session_key=ctx.key,
) )
) )
return OutboundMessage( return OutboundMessage(
nanobot/session/summary.py
+12 -23
View File
@@ -1,33 +1,22 @@
"""Structured session-summary values used while building model context.""" """Helpers for validated session-summary metadata."""
from __future__ import annotations from __future__ import annotations
from collections.abc import Mapping from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime from datetime import datetime
from typing import cast from typing import TypedDict, cast
@dataclass(frozen=True, slots=True) class SessionSummary(TypedDict):
class SessionSummary:
"""A consolidated session checkpoint before presentation formatting."""
text: str text: str
last_active: datetime last_active: str
def for_prompt(self) -> str:
return (
f"Previous conversation summary (last active {self.last_active.isoformat()}):\n"
f"{self.text}"
)
@classmethod def session_summary_from_metadata(
def from_metadata(
cls,
metadata: Mapping[str, object] | None, metadata: Mapping[str, object] | None,
*, *,
fallback_last_active: datetime, fallback_last_active: datetime,
) -> SessionSummary | None: ) -> SessionSummary | None:
raw: object = metadata.get("_last_summary") if metadata is not None else None raw: object = metadata.get("_last_summary") if metadata is not None else None
if not isinstance(raw, Mapping): if not isinstance(raw, Mapping):
return None return None
@@ -36,12 +25,12 @@ class SessionSummary:
if not isinstance(text, str) or not text: if not isinstance(text, str) or not text:
return None return None
raw_last_active = summary_data.get("last_active") raw_last_active = summary_data.get("last_active")
if isinstance(raw_last_active, str):
try: try:
last_active = (
datetime.fromisoformat(raw_last_active) datetime.fromisoformat(raw_last_active)
if isinstance(raw_last_active, str) last_active = raw_last_active
else fallback_last_active
)
except ValueError: except ValueError:
last_active = fallback_last_active last_active = fallback_last_active.isoformat()
return cls(text=text, last_active=last_active) else:
last_active = fallback_last_active.isoformat()
return {"text": text, "last_active": last_active}
tests/agent/test_auto_compact.py
+6 -7
View File
@@ -421,7 +421,7 @@ class TestAutoCompact:
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
assert entry.text == "User said hello." assert entry["text"] == "User said hello."
session_after = loop.sessions.get_or_create("cli:test") session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 12 assert len(session_after.messages) == 12
assert len(session_after.get_history(max_messages=12)) == ( assert len(session_after.get_history(max_messages=12)) == (
@@ -909,7 +909,7 @@ class TestProactiveAutoCompact:
assert len(archived_messages) == 10 assert len(archived_messages) == 10
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
assert entry.text == "User chatted about old things." assert entry["text"] == "User chatted about old things."
await loop.aclose() await loop.aclose()
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -1227,8 +1227,7 @@ class TestSummaryPersistence:
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test") _, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary is not None assert summary is not None
assert summary.text == "User said hello." assert summary["text"] == "User said hello."
assert "Previous conversation summary" in summary.for_prompt()
# _last_summary persists in metadata for restart survival. # _last_summary persists in metadata for restart survival.
assert "_last_summary" in reloaded.metadata assert "_last_summary" in reloaded.metadata
await loop.aclose() await loop.aclose()
@@ -1256,7 +1255,7 @@ class TestSummaryPersistence:
assert summary is not None assert summary is not None
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test") _, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary2 is not None assert summary2 is not None
assert summary2.text == "Summary." assert summary2["text"] == "Summary."
# _last_summary persists in metadata for restart survival. # _last_summary persists in metadata for restart survival.
assert "_last_summary" in reloaded.metadata assert "_last_summary" in reloaded.metadata
await loop.aclose() await loop.aclose()
@@ -1306,7 +1305,7 @@ class TestSummaryPersistence:
loop.sessions.get_or_create("cli:test"), "cli:test" loop.sessions.get_or_create("cli:test"), "cli:test"
) )
assert summary1 is not None assert summary1 is not None
assert summary1.text == "First summary." assert summary1["text"] == "First summary."
assert "cli:test" not in loop.auto_compact._summaries # popped by hot path assert "cli:test" not in loop.auto_compact._summaries # popped by hot path
# Add new messages and archive again (simulating a later turn) # Add new messages and archive again (simulating a later turn)
@@ -1326,7 +1325,7 @@ class TestSummaryPersistence:
reloaded = loop.sessions.get_or_create("cli:test") reloaded = loop.sessions.get_or_create("cli:test")
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test") _, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary2 is not None assert summary2 is not None
assert summary2.text == "Second summary." assert summary2["text"] == "Second summary."
await loop.aclose() await loop.aclose()
@pytest.mark.asyncio @pytest.mark.asyncio
tests/agent/test_autocompact_unit.py
+31 -32
View File
@@ -7,7 +7,6 @@ import pytest
from nanobot.agent.autocompact import AutoCompact from nanobot.agent.autocompact import AutoCompact
from nanobot.session.manager import Session, SessionManager from nanobot.session.manager import Session, SessionManager
from nanobot.session.summary import SessionSummary
def _runtime(_session: Session | None = None): def _runtime(_session: Session | None = None):
@@ -176,24 +175,6 @@ class TestIsExpired:
assert ac._is_expired(expired, now=now) is True assert ac._is_expired(expired, now=now) is True
# ---------------------------------------------------------------------------
# SessionSummary
# ---------------------------------------------------------------------------
class TestSessionSummary:
"""Test prompt rendering for the structured summary value."""
def test_formats_prompt(self):
last_active = datetime(2026, 5, 13, 14, 30, 0)
summary = SessionSummary("User discussed Python.", last_active)
assert summary.for_prompt() == (
"Previous conversation summary (last active 2026-05-13T14:30:00):\n"
"User discussed Python."
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# check_expired # check_expired
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -490,7 +471,7 @@ class TestArchiveDelegates:
entry = ac._summaries.get("cli:test") entry = ac._summaries.get("cli:test")
assert entry is not None assert entry is not None
assert entry.text == "Hello." assert entry["text"] == "Hello."
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_no_summary_when_compact_returns_empty(self): async def test_no_summary_when_compact_returns_empty(self):
@@ -569,21 +550,29 @@ class TestPrepareSession:
ac = _make_autocompact() ac = _make_autocompact()
session = _make_session() session = _make_session()
last_active = datetime(2026, 5, 13, 14, 0, 0) last_active = datetime(2026, 5, 13, 14, 0, 0)
ac._summaries["cli:test"] = SessionSummary("Hot summary.", last_active) ac._summaries["cli:test"] = {
"text": "Hot summary.",
"last_active": last_active.isoformat(),
}
result_session, summary = ac.prepare_session(session, "cli:test") result_session, summary = ac.prepare_session(session, "cli:test")
assert result_session is session assert result_session is session
assert summary is not None assert summary is not None
assert summary.text == "Hot summary." assert summary == {
assert "Previous conversation summary" in summary.for_prompt() "text": "Hot summary.",
"last_active": last_active.isoformat(),
}
def test_hot_path_pops_summary_one_shot(self): def test_hot_path_pops_summary_one_shot(self):
"""Hot path should pop the summary (one-shot; second call returns None).""" """Hot path should pop the summary (one-shot; second call returns None)."""
ac = _make_autocompact() ac = _make_autocompact()
session = _make_session() session = _make_session()
last_active = datetime(2026, 1, 1) last_active = datetime(2026, 1, 1)
ac._summaries["cli:test"] = SessionSummary("One-shot.", last_active) ac._summaries["cli:test"] = {
"text": "One-shot.",
"last_active": last_active.isoformat(),
}
_, summary1 = ac.prepare_session(session, "cli:test") _, summary1 = ac.prepare_session(session, "cli:test")
assert summary1 is not None assert summary1 is not None
@@ -606,7 +595,7 @@ class TestPrepareSession:
assert result_session is session assert result_session is session
assert summary is not None assert summary is not None
assert summary.text == "Cold summary." assert summary["text"] == "Cold summary."
def test_cold_path_tolerates_malformed_last_active(self): def test_cold_path_tolerates_malformed_last_active(self):
"""A malformed persisted last_active must not raise on the turn path. """A malformed persisted last_active must not raise on the turn path.
@@ -629,8 +618,10 @@ class TestPrepareSession:
assert result_session is session assert result_session is session
assert summary is not None assert summary is not None
assert summary.text == "Cold summary." assert summary == {
assert summary.last_active == fallback "text": "Cold summary.",
"last_active": fallback.isoformat(),
}
def test_cold_path_tolerates_missing_last_active(self): def test_cold_path_tolerates_missing_last_active(self):
"""A _last_summary dict without last_active must not raise.""" """A _last_summary dict without last_active must not raise."""
@@ -645,8 +636,10 @@ class TestPrepareSession:
assert result_session is session assert result_session is session
assert summary is not None assert summary is not None
assert summary.text == "Cold summary." assert summary == {
assert summary.last_active == fallback "text": "Cold summary.",
"last_active": fallback.isoformat(),
}
def test_cold_path_missing_text_returns_none(self): def test_cold_path_missing_text_returns_none(self):
"""A _last_summary without a non-empty string text yields no summary.""" """A _last_summary without a non-empty string text yields no summary."""
@@ -677,7 +670,10 @@ class TestPrepareSession:
ac.sessions = mock_sm ac.sessions = mock_sm
key = "dream:20260602-155256" key = "dream:20260602-155256"
ac._archiving.add(key) ac._archiving.add(key)
ac._summaries[key] = SessionSummary("Hot summary.", datetime(2026, 6, 2, 15, 52, 56)) ac._summaries[key] = {
"text": "Hot summary.",
"last_active": "2026-06-02T15:52:56",
}
session = _make_session( session = _make_session(
key=key, key=key,
updated_at=datetime.now() - timedelta(minutes=20), updated_at=datetime.now() - timedelta(minutes=20),
@@ -717,9 +713,12 @@ class TestPrepareSession:
}, },
}) })
last_active = datetime(2026, 5, 13, 14, 0, 0) last_active = datetime(2026, 5, 13, 14, 0, 0)
ac._summaries["cli:test"] = SessionSummary("Hot summary.", last_active) ac._summaries["cli:test"] = {
"text": "Hot summary.",
"last_active": last_active.isoformat(),
}
_, summary = ac.prepare_session(session, "cli:test") _, summary = ac.prepare_session(session, "cli:test")
assert summary is not None assert summary is not None
assert summary.text == "Hot summary." assert summary["text"] == "Hot summary."
# After hot path pops, cold path would kick in on next call # After hot path pops, cold path would kick in on next call
tests/agent/test_consolidate_offset.py
+19 -17
View File
@@ -520,14 +520,14 @@ class TestNewCommandArchival:
call_count = 0 call_count = 0
expected_runtime = loop.llm_runtime() expected_runtime = loop.llm_runtime()
async def _failing_summarize(_messages, *, runtime, session_key=None) -> bool: async def _failing_summarize(session, *, archive_end, runtime) -> None:
nonlocal call_count nonlocal call_count
assert runtime is expected_runtime assert runtime is expected_runtime
assert session_key == "cli:test" assert session.key == "cli:test"
assert archive_end == len(session.messages)
call_count += 1 call_count += 1
return False
loop.consolidator.archive = _failing_summarize # type: ignore[method-assign] loop.consolidator.archive_session = _failing_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new") new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime) response = await loop._process_message(new_msg, runtime=expected_runtime)
@@ -557,14 +557,14 @@ class TestNewCommandArchival:
archived_session_key = None archived_session_key = None
expected_runtime = loop.llm_runtime() expected_runtime = loop.llm_runtime()
async def _fake_summarize(messages, *, runtime, session_key=None) -> bool: async def _fake_summarize(session, *, archive_end, runtime) -> str:
nonlocal archived_count, archived_session_key nonlocal archived_count, archived_session_key
assert runtime is expected_runtime assert runtime is expected_runtime
archived_count = len(messages) archived_count = len(session.messages[:archive_end])
archived_session_key = session_key archived_session_key = session.key
return True return "Summary."
loop.consolidator.archive = _fake_summarize # type: ignore[method-assign] loop.consolidator.archive_session = _fake_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new") new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime) response = await loop._process_message(new_msg, runtime=expected_runtime)
@@ -588,12 +588,13 @@ class TestNewCommandArchival:
loop.sessions.save(session) loop.sessions.save(session)
expected_runtime = loop.llm_runtime() expected_runtime = loop.llm_runtime()
async def _ok_summarize(_messages, *, runtime, session_key=None) -> bool: async def _ok_summarize(session, *, archive_end, runtime) -> str:
assert runtime is expected_runtime assert runtime is expected_runtime
assert session_key == "cli:test" assert session.key == "cli:test"
return True assert archive_end == len(session.messages)
return "Summary."
loop.consolidator.archive = _ok_summarize # type: ignore[method-assign] loop.consolidator.archive_session = _ok_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new") new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime) response = await loop._process_message(new_msg, runtime=expected_runtime)
@@ -618,14 +619,15 @@ class TestNewCommandArchival:
release_archive = asyncio.Event() release_archive = asyncio.Event()
expected_runtime = loop.llm_runtime() expected_runtime = loop.llm_runtime()
async def _slow_summarize(_messages, *, runtime, session_key=None) -> bool: async def _slow_summarize(session, *, archive_end, runtime) -> str:
assert runtime is expected_runtime assert runtime is expected_runtime
assert session_key == "cli:test" assert session.key == "cli:test"
assert archive_end == len(session.messages)
await release_archive.wait() await release_archive.wait()
archived.set() archived.set()
return True return "Summary."
loop.consolidator.archive = _slow_summarize # type: ignore[method-assign] loop.consolidator.archive_session = _slow_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new") new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
await loop._process_message(new_msg, runtime=expected_runtime) await loop._process_message(new_msg, runtime=expected_runtime)
tests/agent/test_consolidation_ratio.py
+2 -2
View File
@@ -72,7 +72,7 @@ async def test_consolidation_ratio_controls_target(
context_window_tokens=context_window_tokens, context_window_tokens=context_window_tokens,
consolidation_ratio=ratio, consolidation_ratio=ratio,
) )
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = _session_with_turns(loop, turns=10) session = _session_with_turns(loop, turns=10)
remaining_estimates = list(estimates) remaining_estimates = list(estimates)
@@ -90,7 +90,7 @@ async def test_consolidation_ratio_controls_target(
runtime=runtime, runtime=runtime,
) )
assert loop.consolidator.archive.await_count == expected_archives assert loop.consolidator.archive_session.await_count == expected_archives
def test_ratio_propagated_from_config_schema() -> None: def test_ratio_propagated_from_config_schema() -> None:
tests/agent/test_consolidator.py
+75 -145
View File
@@ -98,28 +98,20 @@ def _build_test_messages(**kwargs):
] ]
class TestConsolidatorSummarize: async def _archive(consolidator, messages, runtime, *, session_key="test:session"):
async def test_archive_prompt_includes_media_breadcrumb( return await consolidator.archive(
self, consolidator, mock_provider, store, runtime messages,
):
path = "/home/user/.nanobot/media/websocket/upload_photo.png"
summary = "User uploaded a photo."
mock_provider.chat_with_retry.return_value = MagicMock(
content=summary,
finish_reason="stop",
)
result = await consolidator.archive(
[{"role": "user", "content": "please inspect this", "media": [path]}],
runtime=runtime, runtime=runtime,
session_key=session_key,
request_messages=_build_test_messages(
history=messages,
current_message="consolidate",
),
request_tools=[],
) )
prompt = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
entries = store.read_unprocessed_history(since_cursor=0)
assert f"[image: {path}]" in prompt
assert result == summary
assert [entry["content"] for entry in entries] == [summary]
class TestConsolidatorSummarize:
def test_format_messages_keeps_media_only_user_turn(self): def test_format_messages_keeps_media_only_user_turn(self):
path = "/home/user/.nanobot/media/websocket/clip.mp4" path = "/home/user/.nanobot/media/websocket/clip.mp4"
@@ -134,31 +126,6 @@ class TestConsolidatorSummarize:
assert formatted == f"[2026-07-27] USER: [image: {path}]" assert formatted == f"[2026-07-27] USER: [image: {path}]"
async def test_archive_excludes_model_only_runtime_context(
self, consolidator, mock_provider, runtime
):
content, marker = append_runtime_context(
"ship the feature",
[RuntimeContextBlock(source="goal", content="host-only goal guidance")],
)
mock_provider.chat_with_retry.return_value = MagicMock(
content="User wants to ship the feature.",
finish_reason="stop",
)
await consolidator.archive(
[{
"role": "user",
"content": content,
RUNTIME_CONTEXT_HISTORY_META: marker,
}],
runtime=runtime,
)
prompt = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
assert "ship the feature" in prompt
assert "host-only goal guidance" not in prompt
async def test_archive_uses_captured_generation( async def test_archive_uses_captured_generation(
self, consolidator, mock_provider, runtime self, consolidator, mock_provider, runtime
): ):
@@ -180,10 +147,7 @@ class TestConsolidatorSummarize:
finish_reason="stop", finish_reason="stop",
) )
await consolidator.archive( await _archive(consolidator, [{"role": "user", "content": "hello"}], admitted)
[{"role": "user", "content": "hello"}],
runtime=admitted,
)
call = mock_provider.chat_with_retry.call_args.kwargs call = mock_provider.chat_with_retry.call_args.kwargs
assert call["model"] == admitted.model assert call["model"] == admitted.model
@@ -202,7 +166,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, runtime=runtime) result = await _archive(consolidator, messages, 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
@@ -220,9 +184,10 @@ class TestConsolidatorSummarize:
) )
messages = [{"role": "user", "content": "fix the auth bug"}] messages = [{"role": "user", "content": "fix the auth bug"}]
await consolidator.archive( await _archive(
consolidator,
messages, messages,
runtime=runtime, runtime,
session_key="telegram:chat-1", session_key="telegram:chat-1",
) )
@@ -235,7 +200,7 @@ class TestConsolidatorSummarize:
"""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, runtime=runtime) result = await _archive(consolidator, messages, 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
@@ -251,9 +216,10 @@ class TestConsolidatorSummarize:
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( await _archive(
consolidator,
messages, messages,
runtime=runtime, runtime,
session_key="slack:chat-2", session_key="slack:chat-2",
) )
@@ -261,7 +227,7 @@ class TestConsolidatorSummarize:
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, runtime): async def test_summarize_skips_empty_messages(self, consolidator, runtime):
result = await consolidator.archive([], runtime=runtime) result = await _archive(consolidator, [], runtime)
assert result is None assert result is None
@@ -293,7 +259,7 @@ 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, runtime=runtime) result = await _archive(consolidator, messages, 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
@@ -312,7 +278,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, runtime=runtime) result = await _archive(consolidator, messages, 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
@@ -329,9 +295,10 @@ class TestConsolidatorArchiveErrorHandling:
consolidator.store.raw_archive = MagicMock() consolidator.store.raw_archive = MagicMock()
with pytest.raises(OSError, match="disk full"): with pytest.raises(OSError, match="disk full"):
await consolidator.archive( await _archive(
consolidator,
[{"role": "user", "content": "important"}], [{"role": "user", "content": "important"}],
runtime=runtime, runtime,
) )
consolidator.store.raw_archive.assert_not_called() consolidator.store.raw_archive.assert_not_called()
@@ -339,15 +306,19 @@ class TestConsolidatorArchiveErrorHandling:
async def test_archive_propagates_template_failure_without_raw_archive( async def test_archive_propagates_template_failure_without_raw_archive(
self, consolidator, mock_provider, runtime, monkeypatch self, consolidator, mock_provider, runtime, monkeypatch
): ):
runtime = replace(runtime, context_window_tokens=128_000)
consolidator.store.raw_archive = MagicMock() consolidator.store.raw_archive = MagicMock()
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.agent.memory.render_template", "nanobot.agent.memory.render_template",
MagicMock(side_effect=RuntimeError("template failed")), MagicMock(side_effect=RuntimeError("template failed")),
) )
session = Session(key="test:template")
session.add_message("user", "important")
with pytest.raises(RuntimeError, match="template failed"): with pytest.raises(RuntimeError, match="template failed"):
await consolidator.archive( await consolidator.archive_session(
[{"role": "user", "content": "important"}], session,
archive_end=len(session.messages),
runtime=runtime, runtime=runtime,
) )
@@ -366,9 +337,9 @@ class TestConsolidatorTokenBudget:
session.key = "test:key" session.key = "test:key"
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_session = AsyncMock(return_value=True)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive.assert_not_called() consolidator.archive_session.assert_not_called()
async def test_token_estimation_failure_propagates(self, consolidator, runtime): async def test_token_estimation_failure_propagates(self, consolidator, runtime):
session = Session(key="test:estimate-failure") session = Session(key="test:estimate-failure")
@@ -434,7 +405,7 @@ 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="old conversation summary") consolidator.archive_session = AsyncMock(return_value="old conversation summary")
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
@@ -442,7 +413,8 @@ class TestConsolidatorTokenBudget:
replay_max_messages=6, replay_max_messages=6,
) )
archived_chunk = consolidator.archive.await_args.args[0] archive_end = consolidator.archive_session.await_args.kwargs["archive_end"]
archived_chunk = session.messages[:archive_end]
assert archived_chunk[0]["content"] == "u0" assert archived_chunk[0]["content"] == "u0"
assert archived_chunk[-1]["content"] == "a6" assert archived_chunk[-1]["content"] == "a6"
assert session.last_consolidated == 14 assert session.last_consolidated == 14
@@ -466,7 +438,7 @@ 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="tool turn summary") consolidator.archive_session = AsyncMock(return_value="tool turn summary")
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
@@ -474,7 +446,8 @@ class TestConsolidatorTokenBudget:
replay_max_messages=4, replay_max_messages=4,
) )
archived_chunk = consolidator.archive.await_args.args[0] archive_end = consolidator.archive_session.await_args.kwargs["archive_end"]
archived_chunk = session.messages[:archive_end]
assert [m["content"] for m in archived_chunk] == ["old", "old answer"] assert [m["content"] for m in archived_chunk] == ["old", "old answer"]
assert session.last_consolidated == 2 assert session.last_consolidated == 2
@@ -501,7 +474,7 @@ 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="older turn summary") consolidator.archive_session = AsyncMock(return_value="older turn summary")
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
@@ -509,7 +482,8 @@ class TestConsolidatorTokenBudget:
replay_max_messages=6, replay_max_messages=6,
) )
archived_chunk = consolidator.archive.await_args.args[0] archive_end = consolidator.archive_session.await_args.kwargs["archive_end"]
archived_chunk = session.messages[:archive_end]
assert archived_chunk[2]["content"] == "long older turn" assert archived_chunk[2]["content"] == "long older turn"
assert archived_chunk[-1]["content"] == "older final" assert archived_chunk[-1]["content"] == "older final"
assert session.last_consolidated == len(session.messages) - 2 assert session.last_consolidated == len(session.messages) - 2
@@ -517,12 +491,14 @@ 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, runtime): async def test_token_overflow_appends_prompt_to_replay_prefix(
"""Without chunk cap, the full range from pick_consolidation_boundary is archived.""" self,
consolidator,
mock_provider,
runtime,
):
consolidator._SAFETY_BUFFER = 0 consolidator._SAFETY_BUFFER = 0
session = MagicMock() session = Session(key="test:token-prefix")
session.last_consolidated = 0
session.key = "test:key"
session.provider_state = _provider_state() session.provider_state = _provider_state()
session.messages = [ session.messages = [
{ {
@@ -535,16 +511,24 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock( consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")] side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
) )
# Use real pick_consolidation_boundary — it will find boundary at idx=50 consolidator.pick_consolidation_boundary = MagicMock(return_value=(50, 800))
# (user message at 50, token budget met) consolidator._build_messages = MagicMock(side_effect=_build_test_messages)
consolidator.archive = AsyncMock(return_value=True) mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter")
mock_provider.chat_with_retry.return_value = LLMResponse(
content="Token overflow summary.",
finish_reason="stop",
)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
archived_chunk = consolidator.archive.await_args.args[0] request = mock_provider.chat_with_retry.await_args.kwargs
# pick_consolidation_boundary returns (50, tokens) — user turn at idx 50 assert [message["content"] for message in request["messages"][1:-1]] == [
assert archived_chunk[0]["content"] == "m0" f"m{i}" for i in range(50)
assert session.last_consolidated > 0 ]
assert "final 50 conversation messages" in request["messages"][-1]["content"]
assert request["tools"] == []
assert request["tool_choice"] == "none"
assert session.last_consolidated == 50
assert session.provider_state is None assert session.provider_state is None
async def test_raw_archive_fallback_advances_last_consolidated( async def test_raw_archive_fallback_advances_last_consolidated(
@@ -567,12 +551,12 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock( consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")] side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
) )
# LLM consolidation fails — archive() returns None (raw_archive fired). # LLM consolidation fails after raw_archive fires.
consolidator.archive = AsyncMock(return_value=None) consolidator.archive_session = AsyncMock(return_value=None)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive.assert_awaited_once() consolidator.archive_session.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
@@ -596,12 +580,12 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock( consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(1200, "tiktoken") return_value=(1200, "tiktoken")
) )
consolidator.archive = AsyncMock(return_value=None) consolidator.archive_session = AsyncMock(return_value=None)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) 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_session.await_count == 1
async def test_boundary_respected_when_no_intermediate_user_turn( async def test_boundary_respected_when_no_intermediate_user_turn(
self, consolidator, runtime self, consolidator, runtime
@@ -622,11 +606,11 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock( consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")] side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
) )
consolidator.archive = AsyncMock(return_value=True) consolidator.archive_session = AsyncMock(return_value=True)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive.assert_awaited_once() consolidator.archive_session.assert_awaited_once()
# pick_consolidation_boundary finds the only boundary at idx=61 # pick_consolidation_boundary finds the only boundary at idx=61
assert session.last_consolidated == 61 assert session.last_consolidated == 61
@@ -1439,43 +1423,7 @@ class TestRawArchiveTruncation:
assert len(entries[0]["content"]) < 200 assert len(entries[0]["content"]) < 200
class TestArchiveTruncation: class TestArchivePersistence:
"""archive() must truncate formatted text before sending to consolidation LLM."""
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)
big_messages = [{"role": "user", "content": "x" * 100_000}]
mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary of large input.", finish_reason="stop"
)
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, runtime
):
"""Small context window: truncation uses actual tokenizer count."""
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, runtime=runtime)
sent_messages = mock_provider.chat_with_retry.call_args.kwargs["messages"]
user_content = sent_messages[1]["content"]
# budget = 500 - 100 - 1024 = negative, fallback char-based
# Should be truncated
assert len(user_content) < 250_000
async def test_oversized_summary_is_capped_before_append( async def test_oversized_summary_is_capped_before_append(
self, consolidator, mock_provider, store, runtime self, consolidator, mock_provider, store, runtime
): ):
@@ -1486,29 +1434,11 @@ 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( await _archive(
consolidator,
[{"role": "user", "content": "hi"}], [{"role": "user", "content": "hi"}],
runtime=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, runtime
):
"""Positive token budget should use tiktoken for precise truncation."""
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, runtime=runtime)
import tiktoken
enc = tiktoken.get_encoding("cl100k_base")
sent_content = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
token_count = len(enc.encode(sent_content))
assert token_count <= 9_900
tests/agent/test_context_builder.py
+5 -4
View File
@@ -1,13 +1,11 @@
"""Tests for ContextBuilder — system prompt and message assembly.""" """Tests for ContextBuilder — system prompt and message assembly."""
from datetime import datetime
from pathlib import Path from pathlib import Path
import pytest import pytest
from nanobot.agent.context import ContextBuilder from nanobot.agent.context import ContextBuilder
from nanobot.runtime_context import RuntimeContextBlock from nanobot.runtime_context import RuntimeContextBlock
from nanobot.session.summary import SessionSummary
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Helpers # Helpers
@@ -332,7 +330,10 @@ class TestBuildSystemPrompt:
def test_includes_session_summary(self, tmp_path): def test_includes_session_summary(self, tmp_path):
builder = _builder(tmp_path) builder = _builder(tmp_path)
summary = SessionSummary("Previous chat about Python.", datetime(2026, 8, 19, 10, 0)) summary = {
"text": "Previous chat about Python.",
"last_active": "2026-08-19T10:00:00",
}
result = builder.build_system_prompt(session_summary=summary) result = builder.build_system_prompt(session_summary=summary)
assert "Previous chat about Python." in result assert "Previous chat about Python." in result
assert "[Archived Context Summary]" in result assert "[Archived Context Summary]" in result
@@ -340,7 +341,7 @@ class TestBuildSystemPrompt:
def test_sections_separated_by_separator(self, tmp_path): def test_sections_separated_by_separator(self, tmp_path):
(tmp_path / "AGENTS.md").write_text("Rules.", encoding="utf-8") (tmp_path / "AGENTS.md").write_text("Rules.", encoding="utf-8")
builder = _builder(tmp_path) builder = _builder(tmp_path)
summary = SessionSummary("Summary.", datetime(2026, 8, 19, 10, 0)) summary = {"text": "Summary.", "last_active": "2026-08-19T10:00:00"}
result = builder.build_system_prompt(session_summary=summary) result = builder.build_system_prompt(session_summary=summary)
assert "\n\n---\n\n" in result assert "\n\n---\n\n" in result
tests/agent/test_context_prompt_cache.py
+1 -2
View File
@@ -10,7 +10,6 @@ from pathlib import Path
from nanobot.agent.context import ContextBuilder from nanobot.agent.context import ContextBuilder
from nanobot.runtime_context import RuntimeContextBlock from nanobot.runtime_context import RuntimeContextBlock
from nanobot.session.summary import SessionSummary
class _FakeDatetime(real_datetime): class _FakeDatetime(real_datetime):
@@ -125,7 +124,7 @@ def test_session_summary_replaces_interleaved_recent_history_entry(tmp_path) ->
"later telegram event", "later telegram event",
session_key="telegram:chat-1", session_key="telegram:chat-1",
) )
summary = SessionSummary(overview, real_datetime(2026, 8, 19, 10, 0)) summary = {"text": overview, "last_active": "2026-08-19T10:00:00"}
prompt = builder.build_system_prompt( prompt = builder.build_system_prompt(
session_key=session_key, session_key=session_key,
tests/agent/test_loop_consolidation_tokens.py
+20 -17
View File
@@ -7,7 +7,6 @@ from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse from nanobot.providers.base import LLMResponse
from nanobot.session.manager import replay_max_messages_for_context from nanobot.session.manager import replay_max_messages_for_context
from nanobot.session.summary import SessionSummary
def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -> AgentLoop: def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -> AgentLoop:
@@ -35,17 +34,17 @@ def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None: async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200)
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
await loop.process_direct("hello", session_key="cli:test") await loop.process_direct("hello", session_key="cli:test")
loop.consolidator.archive.assert_not_awaited() loop.consolidator.archive_session.assert_not_awaited()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None: async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"}, {"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
@@ -57,13 +56,13 @@ async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypat
await loop.process_direct("hello", session_key="cli:test") await loop.process_direct("hello", session_key="cli:test")
assert loop.consolidator.archive.await_count >= 1 assert loop.consolidator.archive_session.await_count >= 1
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None: async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
@@ -83,7 +82,8 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
) )
archived_chunk = loop.consolidator.archive.await_args.args[0] archive_end = loop.consolidator.archive_session.await_args.kwargs["archive_end"]
archived_chunk = session.messages[:archive_end]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"] assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"]
assert session.last_consolidated == 4 assert session.last_consolidated == 4
@@ -92,7 +92,7 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> None: async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> None:
"""Verify maybe_consolidate_by_tokens keeps looping until under threshold.""" """Verify maybe_consolidate_by_tokens keeps looping until under threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
@@ -123,7 +123,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
) )
assert loop.consolidator.archive.await_count == 2 assert loop.consolidator.archive_session.await_count == 2
assert session.last_consolidated == 6 assert session.last_consolidated == 6
@@ -131,7 +131,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> None: async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> None:
"""Once triggered, consolidation should continue until it drops below half threshold.""" """Once triggered, consolidation should continue until it drops below half threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
@@ -163,14 +163,14 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path,
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
) )
assert loop.consolidator.archive.await_count == 2 assert loop.consolidator.archive_session.await_count == 2
assert session.last_consolidated == 6 assert session.last_consolidated == 6
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, monkeypatch) -> None: async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, monkeypatch) -> None:
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
@@ -203,7 +203,7 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path,
reloaded, pending = loop.auto_compact.prepare_session(reloaded, "cli:test") reloaded, pending = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert pending is not None assert pending is not None
assert pending.text == "User discussed project status." assert pending["text"] == "User discussed project status."
# _last_summary persists for restart survival. # _last_summary persists for restart survival.
assert "_last_summary" in reloaded.metadata assert "_last_summary" in reloaded.metadata
@@ -213,7 +213,10 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200)
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
loop.auto_compact.prepare_session = MagicMock( loop.auto_compact.prepare_session = MagicMock(
return_value=(session, SessionSummary("earlier context", session.updated_at)) return_value=(
session,
{"text": "earlier context", "last_active": session.updated_at.isoformat()},
)
) # type: ignore[method-assign] ) # type: ignore[method-assign]
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]
@@ -242,11 +245,11 @@ 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, *, runtime, session_key=None): async def track_consolidate(session, *, archive_end, runtime):
order.append("consolidate") order.append("consolidate")
archived_session_keys.append(session_key) archived_session_keys.append(session.key)
return True return True
loop.consolidator.archive = track_consolidate # type: ignore[method-assign] loop.consolidator.archive_session = track_consolidate # type: ignore[method-assign]
async def track_llm(*args, **kwargs): async def track_llm(*args, **kwargs):
order.append("llm") order.append("llm")
tests/agent/test_unified_session.py
+16 -12
View File
@@ -258,7 +258,7 @@ class TestCmdNewUnifiedSession:
previous_file_state.record_read(tracked_file) previous_file_state.record_read(tracked_file)
loop = SimpleNamespace( loop = SimpleNamespace(
sessions=sessions, sessions=sessions,
consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), consolidator=SimpleNamespace(archive_session=AsyncMock(return_value=True)),
_cancel_active_tasks=AsyncMock(return_value=0), _cancel_active_tasks=AsyncMock(return_value=0),
discard_session_file_state=file_state_store.discard, discard_session_file_state=file_state_store.discard,
llm_runtime=MagicMock(return_value=MagicMock()), llm_runtime=MagicMock(return_value=MagicMock()),
@@ -288,10 +288,14 @@ class TestCmdNewUnifiedSession:
reset_file_state = file_state_store.for_session("unified:default") reset_file_state = file_state_store.for_session("unified:default")
assert reset_file_state is not previous_file_state assert reset_file_state is not previous_file_state
assert reset_file_state.is_unchanged(tracked_file) is False assert reset_file_state.is_unchanged(tracked_file) is False
loop.consolidator.archive.assert_called_once_with( archived = loop.consolidator.archive_session.call_args.args[0]
expected_snapshot, assert archived.key == "unified:default"
assert archived.messages == expected_snapshot
assert archived.last_consolidated == 0
loop.consolidator.archive_session.assert_called_once_with(
archived,
archive_end=len(expected_snapshot),
runtime=admitted_runtime, runtime=admitted_runtime,
session_key="unified:default",
) )
loop.llm_runtime.assert_not_called() loop.llm_runtime.assert_not_called()
@@ -310,7 +314,7 @@ class TestCmdNewUnifiedSession:
loop = SimpleNamespace( loop = SimpleNamespace(
sessions=sessions, sessions=sessions,
consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)), consolidator=SimpleNamespace(archive_session=AsyncMock(return_value=True)),
_cancel_active_tasks=AsyncMock(return_value=0), _cancel_active_tasks=AsyncMock(return_value=0),
discard_session_file_state=MagicMock(), discard_session_file_state=MagicMock(),
runtime_for_session=MagicMock(return_value=MagicMock()), runtime_for_session=MagicMock(return_value=MagicMock()),
@@ -356,7 +360,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
build_messages=MagicMock(return_value=[]), build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]), get_tool_definitions=MagicMock(return_value=[]),
) )
consolidator.archive = AsyncMock() consolidator.archive_session = AsyncMock()
session = Session(key="unified:default") session = Session(key="unified:default")
session.messages = [] session.messages = []
@@ -364,11 +368,11 @@ class TestConsolidationUnaffectedByUnifiedSession:
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive.assert_not_called() consolidator.archive_session.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_consolidation_behaviour_identical_for_any_key(self): async def test_consolidation_behaviour_identical_for_any_key(self):
"""archive call count is the same for 'telegram:123' and 'unified:default' """Archive call count is the same for 'telegram:123' and 'unified:default'
under identical token conditions.""" under identical token conditions."""
from nanobot.agent.memory import Consolidator, MemoryStore from nanobot.agent.memory import Consolidator, MemoryStore
@@ -392,12 +396,12 @@ class TestConsolidationUnaffectedByUnifiedSession:
session.messages = [] # empty → exits immediately for both keys session.messages = [] # empty → exits immediately for both keys
sessions.get_or_create.return_value = session sessions.get_or_create.return_value = session
consolidator.archive = AsyncMock() consolidator.archive_session = AsyncMock()
await consolidator.maybe_consolidate_by_tokens( await consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime, runtime=runtime,
) )
archive_calls[key] = consolidator.archive.call_count archive_calls[key] = consolidator.archive_session.call_count
assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0 assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0
@@ -427,7 +431,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(950, "tiktoken")) consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(950, "tiktoken"))
# No valid boundary found → returns gracefully without archiving # No valid boundary found → returns gracefully without archiving
consolidator.pick_consolidation_boundary = MagicMock(return_value=None) consolidator.pick_consolidation_boundary = MagicMock(return_value=None)
consolidator.archive = AsyncMock() consolidator.archive_session = AsyncMock()
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
@@ -437,7 +441,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
runtime=runtime, 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_session.assert_not_called()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------