mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-04 02:01:48 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent): let runner own context compaction (#5568)

Browse Source 
* refactor(agent): consolidate accepted history under pressure

* fix(agent): align provider and session compaction

* refactor(agent): simplify runner context compaction

* refactor(agent): remove background token consolidation

* fix(agent): keep injected transcript messages distinct

* refactor(agent): unify native compaction summaries

* fix(agent): preserve native compaction boundary

* fix(agent): unify context compaction paths

* fix(agent): preserve exact compaction request boundaries
This commit is contained in:
chengyongru
2026-09-02 18:05:54 +08:00
committed by GitHub
parent da96c5c6eb
commit d81aa5a4ab
44 changed files with 1914 additions and 1666 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/test_auto_compact.py
-31
View File
@@ -8,7 +8,6 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.agent.runner import AgentRunResult
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
@@ -229,36 +228,6 @@ class TestAgentLoopTTLParam:
loop = _make_loop(tmp_path, session_ttl_minutes=0)
assert loop.auto_compact._ttl == 0
@pytest.mark.asyncio
async def test_process_message_reads_history_with_token_budget(self, tmp_path):
"""_process_message should pass an auto-derived token budget to get_history."""
loop = _make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:direct")
session.get_history = MagicMock(return_value=[])
loop.context.build_messages = MagicMock(return_value=[])
loop._run_agent_loop = AsyncMock(
return_value=AgentRunResult(
final_content="ok",
messages=[],
stop_reason="stop",
)
)
loop._save_turn = MagicMock()
msg = InboundMessage(
channel="cli",
sender_id="u1",
chat_id="direct",
content="hello",
)
await loop._process_message(msg)
session.get_history.assert_called_once()
kwargs = session.get_history.call_args.kwargs
assert isinstance(kwargs.get("max_tokens"), int)
assert kwargs["max_tokens"] > 0
assert set(kwargs) == {"max_tokens", "extend_to_user"}
class TestAutoCompact:
"""Test the _archive method."""
tests/agent/test_consolidator.py
+85 -263
View File
@@ -55,9 +55,7 @@ def runtime(mock_provider):
def consolidator(store):
sessions = MagicMock()
sessions.save = MagicMock()
# When maybe_consolidate_by_tokens refreshes the session reference via
# get_or_create(session.key), it should get back the same object the test
# passed in. Store sessions by key so the lookup is transparent.
# Store sessions by key so refreshes observe the same test object.
_session_cache: dict[str, MagicMock] = {}
sessions.get_or_create = MagicMock(side_effect=lambda key: _session_cache.get(key, MagicMock()))
sessions._session_cache = _session_cache
@@ -93,11 +91,17 @@ def _provider_state() -> ProviderConversationState:
def _build_test_messages(**kwargs):
return [
{"role": "system", "content": "system prompt"},
system = "system prompt"
session_summary = kwargs.get("session_summary")
if session_summary:
system += f"\n\n[Archived Context Summary]\n{session_summary['text']}"
messages = [
{"role": "system", "content": system},
*kwargs["history"],
{"role": "user", "content": kwargs["current_message"]},
]
if kwargs["current_message"] is not None:
messages.append({"role": "user", "content": kwargs["current_message"]})
return messages
async def _archive(
@@ -112,15 +116,85 @@ async def _archive(
messages,
runtime=runtime,
session_key=session_key,
request_messages=_build_test_messages(
history=messages,
current_message="consolidate",
),
history=[
{"role": "system", "content": "system prompt"},
*messages,
],
request_tools=[],
previous_summary=previous_summary,
)
class TestTurnTranscriptSummary:
async def test_uses_exact_accepted_prefix_and_existing_archiver(
self,
consolidator,
mock_provider,
runtime,
):
accepted = [
{"role": "system", "content": "stable system"},
{"role": "user", "content": "accepted history"},
]
tools = [{"type": "function", "function": {"name": "inspect"}}]
mock_provider.chat_with_retry.return_value = LLMResponse(
content="replacement checkpoint",
)
summary = await consolidator.summarize_transcript(
accepted,
"previous checkpoint",
runtime=runtime,
session_key="test:turn",
tools=tools,
)
assert summary == "replacement checkpoint"
call = mock_provider.chat_with_retry.await_args.kwargs
assert call["messages"][:-1] == accepted
assert call["messages"][-1]["role"] == "user"
assert "SNIP" in call["messages"][-1]["content"]
assert call["tools"] == tools
async def test_native_compaction_appends_only_archive_prompt(
self,
consolidator,
mock_provider,
runtime,
):
accepted = [
{"role": "system", "content": "stable system"},
{"role": "user", "content": "raw history must not be replayed"},
]
state = _provider_state()
mock_provider.can_resume_conversation_state.return_value = True
mock_provider.chat_with_retry.return_value = LLMResponse(
content="replacement checkpoint",
)
summary = await consolidator.summarize_provider_compaction(
state,
accepted,
"previous checkpoint",
runtime=runtime,
session_key="test:turn",
tools=[{"type": "function", "function": {"name": "inspect"}}],
)
assert summary == "replacement checkpoint"
call = mock_provider.chat_with_retry.await_args.kwargs
assert call["messages"][0] == accepted[0]
assert call["messages"][-1]["content"] == _ARCHIVE_PROMPT
assert accepted[1] not in call["messages"]
assert call["tools"] == []
provider_context = call["provider_context"]
assert provider_context.conversation_state is not None
assert provider_context.conversation_state.payload == state.payload
assert provider_context.conversation_state.pending_messages == [
call["messages"][-1],
]
class TestConsolidatorSummarize:
def test_format_messages_keeps_media_only_user_turn(self):
path = "/home/user/.nanobot/media/websocket/clip.mp4"
@@ -379,32 +453,7 @@ class TestConsolidatorArchiveErrorHandling:
consolidator.store.raw_archive.assert_not_called()
class TestConsolidatorTokenBudget:
async def test_prompt_below_threshold_does_not_consolidate(
self, consolidator, runtime
):
"""No consolidation when tokens are within budget."""
session = MagicMock()
session.last_archived = 0
session.messages = [{"role": "user", "content": "hi"}]
session.key = "test:key"
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
consolidator.archive_session = AsyncMock(return_value=True)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_not_called()
async def test_token_estimation_failure_propagates(self, consolidator, runtime):
session = Session(key="test:estimate-failure")
session.add_message("user", "hello")
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=RuntimeError("counter failed")
)
with pytest.raises(RuntimeError, match="counter failed"):
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
class TestConsolidatorPromptEstimate:
async def test_estimate_uses_full_unarchived_tail(self, consolidator, runtime):
"""Consolidation pressure must account for the full unarchived tail."""
session = Session(key="test:full-tail")
@@ -443,129 +492,6 @@ class TestConsolidatorTokenBudget:
assert len(captured["history"]) == 8
assert captured["history"][0]["content"] == "msg-2"
async def test_token_overflow_appends_prompt_to_replay_prefix(
self,
consolidator,
mock_provider,
runtime,
):
consolidator._SAFETY_BUFFER = 0
session = Session(key="test:token-prefix")
session.provider_state = _provider_state()
session.messages = [
{
"role": "user" if i in {0, 50, 61} else "assistant",
"content": f"m{i}",
}
for i in range(70)
]
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
consolidator.pick_consolidation_boundary = MagicMock(return_value=50)
consolidator.archiver._build_messages = MagicMock(side_effect=_build_test_messages)
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)
request = mock_provider.chat_with_retry.await_args.kwargs
assert [message["content"] for message in request["messages"][1:-1]] == [
f"m{i}" for i in range(50)
]
assert request["messages"][-1]["content"] == _ARCHIVE_PROMPT
assert request["tools"] == []
assert "tool_choice" not in request
assert session.last_archived == 50
assert session.provider_state == _provider_state()
async def test_raw_archive_fallback_advances_archive_watermark(
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
duplicate [RAW] entries into history.jsonl."""
consolidator._SAFETY_BUFFER = 0
session = Session(key="test:key")
session.provider_state = _provider_state()
session.messages = [
{"role": "user" if i in {0, 50} else "assistant", "content": f"m{i}"}
for i in range(70)
]
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
consolidator.archive_session = AsyncMock(return_value="[RAW] checkpoint")
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_awaited_once()
# The chunk is considered "materialized" (as a raw-archive breadcrumb),
# so the archive watermark must have moved past it without touching
# the provider-owned continuation state.
assert session.last_archived == 50
assert session.provider_state == _provider_state()
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
session = MagicMock()
session.last_archived = 0
session.key = "test:key"
session.messages = [
{"role": "user" if i in {0, 20, 40, 60} else "assistant", "content": f"m{i}"}
for i in range(70)
]
session.metadata = {}
consolidator.sessions._session_cache[session.key] = session
# Keep estimates high so the loop would otherwise run multiple rounds.
consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(1200, "tiktoken")
)
consolidator.archive_session = AsyncMock(return_value="[RAW] checkpoint")
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
# The fixed policy archives at most one prefix per call.
assert consolidator.archive_session.await_count == 1
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()
session.last_archived = 0
session.key = "test:key"
session.messages = [
{
"role": "user" if i in {0, 61} else "assistant",
"content": f"m{i}",
}
for i in range(70)
]
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
consolidator.archive_session = AsyncMock(return_value=True)
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_awaited_once()
# The fixed recent tail expands backward to the user at idx=61.
assert session.last_archived == 61
class TestCompactIdleSession:
"""Idle compaction tests."""
@@ -1347,110 +1273,6 @@ class TestCompactIdleSession:
assert not lock.locked()
class TestConsolidatorSessionRefresh:
"""Background consolidation must detect stale session references."""
@pytest.mark.asyncio
async def test_reloads_before_empty_session_guard(self, tmp_path):
"""A stale empty reference must not skip a non-empty cached session."""
from nanobot.agent.memory import Consolidator, MemoryStore
from nanobot.session.manager import Session, SessionManager
store = MemoryStore(tmp_path)
provider = MagicMock()
provider.chat_with_retry = AsyncMock(
return_value=MagicMock(content="summary", finish_reason="stop")
)
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,
sessions=sessions,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
fresh = sessions.get_or_create("cli:test")
fresh.add_message("user", "fresh message")
sessions.save(fresh)
stale_empty = Session(key="cli:test")
seen: dict[str, 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,
runtime=runtime,
)
assert seen["session"] is fresh
@pytest.mark.asyncio
async def test_reloads_stale_session_after_compact(self, tmp_path):
"""After compact_idle_session replaces the session, a concurrent
maybe_consolidate_by_tokens with the old reference should use the
fresh session from cache instead of overwriting."""
from nanobot.agent.memory import Consolidator, MemoryStore
from nanobot.session.manager import SessionManager
store = MemoryStore(tmp_path)
provider = MagicMock()
provider.chat_with_retry = AsyncMock(
return_value=MagicMock(content="summary", finish_reason="stop")
)
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,
sessions=sessions,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
# Populate session with many messages
session = sessions.get_or_create("cli:test")
for i in range(20):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
sessions.save(session)
# Simulate: background consolidation captures old reference
old_ref = session
await consolidator.compact_idle_session(
"cli:test",
runtime=runtime,
max_suffix=8,
)
await consolidator.maybe_consolidate_by_tokens(
old_ref,
runtime=runtime,
)
session_after = sessions.get_or_create("cli:test")
assert len(session_after.messages) == 40
assert session_after.last_archived == 40
assert len(session_after.get_history(max_messages=40)) == 8
class TestRawArchiveTruncation:
"""raw_archive() must cap entry size to avoid bloating history.jsonl."""
tests/agent/test_dream.py
+2 -18
View File
@@ -404,10 +404,9 @@ class TestEphemeralDirect:
with (
patch("nanobot.agent.loop.SessionManager"),
patch("nanobot.agent.loop.SubagentManager") as mock_sub,
patch("nanobot.agent.loop.Consolidator") as mock_consolidator_cls,
patch("nanobot.agent.loop.Consolidator"),
):
mock_sub.return_value.cancel_by_session = AsyncMock(return_value=0)
mock_consolidator_cls.return_value.maybe_consolidate_by_tokens = AsyncMock()
loop = AgentLoop(
bus=bus,
provider=provider,
@@ -493,20 +492,6 @@ class TestEphemeralDirect:
assert captured.get("ephemeral") is False
async def test_ephemeral_skips_consolidator(self, tmp_path, _make_loop):
"""When ephemeral=True, consolidator.maybe_consolidate_by_tokens is not called."""
from unittest.mock import patch
loop, store = _make_loop
with patch.object(
loop.consolidator, "maybe_consolidate_by_tokens",
) as mock_consolidate:
await loop.process_direct(
"test", session_key="dream:consolidate-test", ephemeral=True,
)
mock_consolidate.assert_not_called()
async def test_ephemeral_response_reports_stop_reason(self, tmp_path, _make_loop):
loop, store = _make_loop
loop.provider.chat_with_retry.return_value = LLMResponse(
@@ -701,10 +686,9 @@ class TestEphemeralHooks:
with (
patch("nanobot.agent.loop.SessionManager"),
patch("nanobot.agent.loop.SubagentManager") as mock_sub,
patch("nanobot.agent.loop.Consolidator") as mock_consolidator_cls,
patch("nanobot.agent.loop.Consolidator"),
):
mock_sub.return_value.cancel_by_session = AsyncMock(return_value=0)
mock_consolidator_cls.return_value.maybe_consolidate_by_tokens = AsyncMock()
loop = AgentLoop(
bus=bus,
provider=provider,
tests/agent/test_history_replay.py
+7 -9
View File
@@ -12,6 +12,7 @@ from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse
from nanobot.session.manager import Session
from nanobot.session.summary import SUMMARY_CONTINUATION_TEXT
def _make_loop(tmp_path: Path, context_window_tokens: int = 200_000) -> AgentLoop:
@@ -66,13 +67,12 @@ def test_explicit_message_limit_still_starts_at_user_turn() -> None:
@pytest.mark.asyncio
async def test_process_message_replays_with_token_budget_only(tmp_path: Path) -> None:
async def test_process_message_hands_complete_replay_to_runner(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=32_768)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
with patch.object(session, "get_history", wraps=session.get_history) as get_history:
@@ -81,20 +81,16 @@ async def test_process_message_replays_with_token_budget_only(tmp_path: Path) ->
)
assert result is not None
assert get_history.call_args.kwargs == {
"max_tokens": loop._replay_token_budget(loop.llm_runtime()),
"extend_to_user": False,
}
assert get_history.call_args.kwargs == {"extend_to_user": False}
@pytest.mark.asyncio
async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None:
async def test_runner_checkpoint_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=8_000)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old")
@@ -117,4 +113,6 @@ async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path
sent_messages = loop.provider.chat_with_retry.await_args.kwargs["messages"]
sent_text = "\n".join(str(message.get("content")) for message in sent_messages)
assert "new question" in sent_text
assert "long older turn" not in sent_text
assert [message["role"] for message in sent_messages] == ["system", "user", "user"]
assert sent_messages[1]["content"] == SUMMARY_CONTINUATION_TEXT
assert any(message.get("content") == "long older turn" for message in session.messages)
tests/agent/test_loop_consolidation_tokens.py
+90 -164
View File
@@ -4,7 +4,12 @@ import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse
from nanobot.providers.base import (
GenerationSettings,
LLMResponse,
ProviderConversationState,
)
from nanobot.session.summary import SUMMARY_CONTINUATION_TEXT
def _make_loop(
@@ -14,7 +19,6 @@ def _make_loop(
context_window_tokens: int,
max_tokens: int = 0,
) -> AgentLoop:
from nanobot.providers.base import GenerationSettings
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings(max_tokens=max_tokens)
@@ -39,186 +43,108 @@ def _make_loop(
@pytest.mark.asyncio
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.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
await loop.process_direct("hello", session_key="cli:test")
loop.consolidator.archive_session.assert_not_awaited()
@pytest.mark.asyncio
async def test_prompt_above_threshold_triggers_consolidation(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
await loop.process_direct("hello", session_key="cli:test")
assert loop.consolidator.archive_session.await_count >= 1
@pytest.mark.asyncio
async def test_token_consolidation_refreshes_summary_for_current_request(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock( # type: ignore[method-assign]
return_value="FRESH_CHECKPOINT"
)
loop.consolidator.estimate_session_prompt_tokens = MagicMock( # type: ignore[method-assign]
return_value=(1000, "test")
)
async def test_runner_pressure_commits_summary_and_current_delta(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=2_000)
loop.context_block_limit = 500
loop.provider.generation = GenerationSettings(max_tokens=100)
loop.provider.can_resume_conversation_state.return_value = False
loop.schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
{"role": role, "content": f"old-{role}-{turn}"}
for turn in range(6)
for role in ("user", "assistant")
]
loop.sessions.save(session)
await loop.process_direct("hello", session_key="cli:test")
def estimate(messages, _tools, _model):
contents = [str(message.get("content")) for message in messages]
if contents and "SNIP" in contents[-1]:
return 300, "test-counter"
if any(content.startswith("old-") for content in contents):
return 600, "test-counter"
return 100, "test-counter"
request_messages = loop.provider.chat_with_retry.await_args.kwargs["messages"]
system_prompt = request_messages[0]["content"]
assert "FRESH_CHECKPOINT" in system_prompt
assert all(message.get("content") != "u0" for message in request_messages)
assert loop.sessions.get_or_create("cli:test").last_archived == 12
loop.provider.estimate_prompt_tokens.side_effect = estimate
loop.provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="Current checkpoint.", tool_calls=[]),
LLMResponse(content="done", tool_calls=[]),
])
result = await loop.process_direct("continue the task", session_key="cli:test")
@pytest.mark.asyncio
async def test_prompt_above_threshold_uses_fixed_recent_tail(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
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] == [
"u0", "a0", "u1", "a1", "u2", "a2", "u3", "a3", "u4", "a4", "u5", "a5",
]
assert session.last_archived == 12
@pytest.mark.asyncio
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(5)
for role in ("user", "assistant")
]
loop.sessions.save(session)
def mock_estimate(_session, *, runtime):
return (500, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert result.content == "done"
assert loop.provider.chat_with_retry.await_count == 2
model_request = loop.provider.chat_with_retry.await_args_list[1].kwargs["messages"]
assert "Current checkpoint." in model_request[0]["content"]
assert model_request[1]["content"] == SUMMARY_CONTINUATION_TEXT
assert model_request[2]["content"] == "continue the task"
reloaded = loop.sessions.get_or_create("cli:test")
meta = reloaded.metadata.get("_last_summary")
assert meta is not None
assert meta["text"] == "User discussed project status."
reloaded, pending = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert pending is not None
assert pending["text"] == "User discussed project status."
# _last_summary persists for restart survival.
assert "_last_summary" in reloaded.metadata
assert reloaded.messages[0]["content"] == "old-user-0"
assert reloaded.metadata["_last_summary"]["text"] == "Current checkpoint."
assert reloaded.messages[reloaded.last_archived]["content"] == (
SUMMARY_CONTINUATION_TEXT
)
assert [message["content"] for message in reloaded.get_history()] == [
SUMMARY_CONTINUATION_TEXT,
"continue the task",
"done",
]
@pytest.mark.asyncio
async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200)
session = loop.sessions.get_or_create("cli:test")
loop.auto_compact.prepare_session = MagicMock(
return_value=(
session,
{"text": "earlier context", "last_active": session.updated_at.isoformat()},
)
) # 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]
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,
)
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
async def test_preflight_consolidation_before_llm_call(tmp_path) -> None:
"""Verify preflight consolidation runs before the LLM call in process_direct."""
order: list[str] = []
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
archived_session_keys: list[str | None] = []
async def track_consolidate(session, *, archive_end, runtime):
order.append("consolidate")
archived_session_keys.append(session.key)
return True
loop.consolidator.archive_session = track_consolidate # type: ignore[method-assign]
async def track_llm(*args, **kwargs):
order.append("llm")
return LLMResponse(content="ok", tool_calls=[])
loop.provider.chat_with_retry = track_llm
loop.provider.chat_stream_with_retry = track_llm
loop.schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
async def test_native_provider_compaction_commits_portable_terminal_checkpoint(
tmp_path,
) -> None:
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=2_000)
session = loop.sessions.get_or_create("cli:native")
session.messages = [
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
{"role": "user", "content": "accepted history"},
{"role": "assistant", "content": "accepted answer"},
]
loop.sessions.save(session)
call_count = [0]
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]
compacted_state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="test-model",
version=1,
payload={"items": [{"type": "compaction", "encrypted_content": "opaque"}]},
)
loop.provider.can_resume_conversation_state.return_value = True
loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="done",
provider_state=compacted_state,
provider_compaction_applied=True,
provider_compaction_state=compacted_state,
provider_compaction_scope="current_request",
))
loop.consolidator.summarize_provider_compaction = AsyncMock(
return_value="portable terminal checkpoint",
)
await loop.process_direct("hello", session_key="cli:test")
result = await loop.process_direct("continue", session_key="cli:native")
assert "consolidate" in order
assert "llm" in order
assert order.index("consolidate") < order.index("llm")
assert archived_session_keys == ["cli:test"]
assert result.content == "done"
summarize = loop.consolidator.summarize_provider_compaction
summarize.assert_awaited_once()
assert summarize.await_args.args[0] == compacted_state
accepted = summarize.await_args.args[1]
accepted_contents = [message.get("content") for message in accepted]
assert "accepted history" in accepted_contents
assert "accepted answer" in accepted_contents
assert "continue" in accepted_contents
assert "done" not in accepted_contents
reloaded = loop.sessions.get_or_create("cli:native")
assert reloaded.provider_state is None
assert reloaded.metadata["_last_summary"]["text"] == (
"portable terminal checkpoint"
)
assert reloaded.messages[reloaded.last_archived]["content"] == (
SUMMARY_CONTINUATION_TEXT
)
assert [message["content"] for message in reloaded.get_history()] == [
SUMMARY_CONTINUATION_TEXT,
"done",
]
tests/agent/test_loop_image_generation_media.py
-1
View File
@@ -69,7 +69,6 @@ async def test_outbound_no_longer_carries_generated_media(
),
image_generation_provider_config=ProviderConfig(api_key="sk-or-test"),
)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
result = await loop._process_message(
InboundMessage(
tests/agent/test_loop_progress.py
-15
View File
@@ -425,7 +425,6 @@ class TestToolEventProgress:
None,
),
)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -473,7 +472,6 @@ class TestToolEventProgress:
provider.chat_stream_with_retry = AsyncMock()
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="openai-codex/gpt-5.5")
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="whatsapp",
@@ -512,7 +510,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="openai-codex/gpt-5.5")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -566,7 +563,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -611,7 +607,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -655,7 +650,6 @@ class TestToolEventProgress:
_attach_webui_runtime_events(loop, bus)
loop.max_iterations = 1
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -747,7 +741,6 @@ class TestToolEventProgress:
)
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -815,9 +808,6 @@ class TestToolEventProgress:
return "ok"
loop.tools.execute = execute_tool
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock( # type: ignore[method-assign]
return_value=False
)
session_key = "websocket:chat-a"
session = loop.sessions.get_or_create(session_key)
@@ -949,7 +939,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="openai-codex/gpt-5.5")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -1048,7 +1037,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="websocket",
@@ -1132,7 +1120,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await asyncio.wait_for(loop._dispatch(InboundMessage(
channel="websocket",
@@ -1181,7 +1168,6 @@ class TestToolEventProgress:
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
_attach_webui_runtime_events(loop, bus)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
captured: dict[str, object] = {}
@@ -1268,7 +1254,6 @@ class TestToolEventProgress:
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="Done", tool_calls=[]))
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
await loop._dispatch(InboundMessage(
channel="slack",
tests/agent/test_loop_runner_integration.py
-6
View File
@@ -112,7 +112,6 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
LLMResponse(content="done", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("cli:direct")
session.add_message("user", "Let's agree on the migration implementation.")
session.add_message("assistant", "Use the staged migration plan and run integration tests.")
@@ -166,7 +165,6 @@ async def test_runtime_context_is_persisted_as_next_turn_prompt_prefix(tmp_path)
LLMResponse(content="second answer", usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("cli:direct")
provider_calls: list[str | None] = []
@@ -220,7 +218,6 @@ async def test_webui_quote_reaches_model_without_leaking_into_public_history(tmp
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage=None))
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("websocket:chat")
quote = webui_quote_runtime_context({
WEBUI_QUOTE_METADATA: "the selected answer excerpt",
@@ -265,7 +262,6 @@ async def test_runtime_context_provider_runs_once_across_tool_iterations(tmp_pat
LLMResponse(content="done", usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
provider_calls = 0
async def provide_context(_request):
@@ -310,7 +306,6 @@ async def test_non_goal_direct_turn_cannot_reuse_prior_goal_command(tmp_path):
LLMResponse(content="handled as a one-time task", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("api:default")
session.add_message("user", "/goal old completed request")
session.add_message("assistant", "The old request is complete.")
@@ -589,7 +584,6 @@ async def test_next_turn_after_llm_error_keeps_turn_boundary(tmp_path):
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
first = await loop._process_message(
InboundMessage(channel="cli", sender_id="user", chat_id="test", content="first question")
tests/agent/test_loop_save_turn.py
+58 -32
View File
@@ -45,6 +45,10 @@ from nanobot.session.recovery import (
RUNTIME_CHECKPOINT_KEY,
restore_runtime_checkpoint,
)
from nanobot.session.summary import (
SUMMARY_CONTINUATION_TEXT,
SessionSummaryCheckpoint,
)
from nanobot.session.turn_continuation import (
INTERNAL_CONTINUATION_META,
INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
@@ -506,6 +510,60 @@ def test_save_turn_keeps_multimodal_runtime_context_for_model_replay() -> None:
assert public_history_message(session.messages[0])["content"] == []
def test_save_turn_commits_summary_boundary_without_rewriting_raw_history() -> None:
loop = _mk_loop()
session = Session(key="test:summary-checkpoint")
session.add_message("user", "inspect the project")
messages = [
{"role": "system", "content": "system"},
{"role": "user", "content": "inspect the project"},
{
"role": "assistant",
"content": "",
"tool_calls": [{
"id": "call-1",
"type": "function",
"function": {"name": "inspect", "arguments": "{}"},
}],
},
{
"role": "tool",
"tool_call_id": "call-1",
"name": "inspect",
"content": "full current result",
},
{"role": "assistant", "content": "done"},
]
loop._save_turn(
session,
messages,
skip=2,
summary_checkpoint=SessionSummaryCheckpoint(
summary="Current working-memory checkpoint.",
transcript_boundary=2,
),
input_persisted_early=True,
)
assert [message["role"] for message in session.messages] == [
"user", "user", "assistant", "tool", "assistant",
]
assert session.messages[0]["content"] == "inspect the project"
assert session.messages[1]["content"] == SUMMARY_CONTINUATION_TEXT
assert session.messages[1]["_hidden_history"] is True
assert session.last_archived == 1
assert session.metadata["_last_summary"]["text"] == (
"Current working-memory checkpoint."
)
assert [message["content"] for message in session.get_history()] == [
SUMMARY_CONTINUATION_TEXT,
"",
"full current result",
"done",
]
def test_save_turn_acknowledges_every_merged_recovery_followup() -> None:
"""Persisting a merged injected row retires every durable follow-up ID."""
loop = _mk_loop()
@@ -966,7 +1024,6 @@ async def test_runtime_checkpoint_keeps_provider_state_out_of_public_metadata(
@pytest.mark.asyncio
async def test_process_message_persists_user_message_before_turn_completes(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign]
msg = InboundMessage(channel="feishu", sender_id="u1", chat_id="c1", content="persist me")
@@ -986,7 +1043,6 @@ async def test_subagent_followup_stages_provider_state_before_turn_runs(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign]
loop.provider.can_resume_conversation_state.return_value = True
session = loop.sessions.get_or_create("cli:subagent-crash")
@@ -1016,7 +1072,6 @@ async def test_subagent_followup_state_is_durable_before_prompt_assembly(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.provider.can_resume_conversation_state.return_value = True
loop.context.build_system_prompt = MagicMock( # type: ignore[method-assign]
side_effect=RuntimeError("prompt boom"),
@@ -1049,7 +1104,6 @@ async def test_subagent_redelivery_does_not_duplicate_staged_provider_input(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.provider.can_resume_conversation_state.return_value = True
build_system_prompt = loop.context.build_system_prompt
loop.context.build_system_prompt = MagicMock( # type: ignore[method-assign]
@@ -1101,7 +1155,6 @@ async def test_subagent_followup_clears_state_before_compatibility_failure(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.provider.can_resume_conversation_state.side_effect = RuntimeError(
"compatibility boom"
)
@@ -1129,7 +1182,6 @@ async def test_subagent_followup_clears_state_before_compatibility_failure(
async def test_process_message_persists_unified_session_delivery_route(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop._unified_session = True
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign]
msg = InboundMessage(
@@ -1230,7 +1282,6 @@ async def test_process_message_persists_media_paths_on_user_turn(tmp_path: Path)
img_b.write_bytes(_PNG_1X1)
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("interrupt")) # type: ignore[method-assign]
msg = InboundMessage(
@@ -1262,7 +1313,6 @@ async def test_process_message_persists_media_only_turn_without_text(tmp_path: P
img.write_bytes(_PNG_1X1)
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign]
msg = InboundMessage(
@@ -1286,7 +1336,6 @@ async def test_process_message_persists_media_only_turn_without_text(tmp_path: P
@pytest.mark.asyncio
async def test_process_message_does_not_duplicate_early_persisted_user_message(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(return_value=_agent_run_result(
"done",
[
@@ -1319,7 +1368,6 @@ async def test_internal_continuation_queues_turn_without_fake_user_history(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("feishu:c-auto")
session.metadata[GOAL_STATE_KEY] = {
"status": "active",
@@ -1388,7 +1436,6 @@ async def test_internal_continuation_preserves_streaming_route_metadata(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("feishu:c-stream")
session.metadata[GOAL_STATE_KEY] = {
"status": "active",
@@ -1462,7 +1509,6 @@ async def test_websocket_internal_continuation_keeps_single_visible_run(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("websocket:c-auto")
session.metadata[GOAL_STATE_KEY] = {
"status": "active",
@@ -1526,7 +1572,6 @@ async def test_websocket_internal_continuation_keeps_single_visible_run(
@pytest.mark.asyncio
async def test_process_message_keeps_delivery_chat_for_thread_session(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.context.build_messages = MagicMock( # type: ignore[method-assign]
return_value=[
{"role": "system", "content": "system"},
@@ -1565,7 +1610,6 @@ async def test_process_message_uses_explicit_session_for_goal_context(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
chat_session = loop.sessions.get_or_create("websocket:chat-with-goal")
chat_session.metadata[GOAL_STATE_KEY] = {
"status": "active",
@@ -1713,7 +1757,6 @@ async def test_request_context_uses_effective_key_for_spawn_tool(tmp_path: Path)
@pytest.mark.asyncio
async def test_next_turn_after_crash_closes_pending_user_turn_before_new_input(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.provider.chat_with_retry = AsyncMock(return_value=MagicMock()) # unused because _run_agent_loop is stubbed
session = loop.sessions.get_or_create("feishu:c3")
@@ -1762,7 +1805,6 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) -
from nanobot.command.router import CommandContext
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
checkpoint_saved = asyncio.Event()
@@ -1866,7 +1908,6 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) -
@pytest.mark.asyncio
async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "question")
@@ -1913,11 +1954,6 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_
assert request.metadata == {"subagent_task_id": "sub-1"}
assert request.turn_id
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"]
@@ -1952,7 +1988,6 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_
@pytest.mark.asyncio
async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
turn_usage = LLMUsage.reported(input_tokens=64, output_tokens=9)
async def fake_run_agent_loop(transcript_input, **_kwargs):
@@ -1978,9 +2013,6 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
@pytest.mark.asyncio
async def test_system_subagent_followup_does_not_log_content(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock( # type: ignore[method-assign]
return_value=False
)
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
@@ -2017,9 +2049,6 @@ async def test_system_subagent_followup_does_not_log_content(tmp_path: Path) ->
@pytest.mark.asyncio
async def test_system_subagent_followup_uses_common_turn_lifecycle(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock( # type: ignore[method-assign]
return_value=False
)
visited: list[str] = []
for name in (
@@ -2081,7 +2110,6 @@ async def test_system_subagent_followup_uses_common_turn_lifecycle(tmp_path: Pat
@pytest.mark.asyncio
async def test_multiple_subagent_followups_all_persist_as_standalone_history(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
@@ -2207,7 +2235,6 @@ async def test_request_context_passes_thread_session_key_to_spawn(tmp_path: Path
@pytest.mark.asyncio
async def test_system_subagent_followup_uses_thread_session_and_slack_metadata(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
thread_session = loop.sessions.get_or_create("slack:C123:1700.42")
thread_session.add_message("user", "thread question")
@@ -2266,7 +2293,6 @@ async def test_system_subagent_followup_uses_thread_session_and_slack_metadata(t
@pytest.mark.asyncio
async def test_turn_after_unanswered_user_keeps_tool_call_pairing(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("feishu:c-merge")
session.add_message("user", "earlier question that never got an answer")
tests/agent/test_loop_session_policy.py
-2
View File
@@ -46,7 +46,6 @@ def _loop(tmp_path, responses: list[str], **kwargs) -> AgentLoop:
async def test_transient_session_keeps_history_without_persisting_or_durable_tools(tmp_path) -> None:
loop = _loop(tmp_path, ["first answer", "second answer"])
loop.context.memory.write_memory("private durable memory")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock()
key = "websocket:transient-test"
loop.sessions.get_or_create_transient(
key,
@@ -71,7 +70,6 @@ async def test_transient_session_keeps_history_without_persisting_or_durable_too
"assistant",
]
assert loop.sessions.read_session_file(key) is None
loop.consolidator.maybe_consolidate_by_tokens.assert_not_awaited()
@pytest.mark.asyncio
tests/agent/test_runner_core.py
+4 -1
View File
@@ -61,7 +61,10 @@ def test_initial_transcript_is_built_from_structured_turn_input() -> None:
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
assert AgentRunner._initial_transcript(spec) == expected
messages, compaction = AgentRunner._initial_transcript_and_compaction(spec)
assert messages == expected
assert compaction is None
transcript_builder.assert_called_once_with(transcript_input)
tests/agent/test_runner_governance.py
+488 -7
View File
@@ -7,13 +7,14 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.context import TranscriptInput
from nanobot.agent.context_governance import (
BACKFILL_CONTENT,
ContextGovernanceConfig,
ContextGovernor,
ContextWindowExceededError,
)
from nanobot.agent.runner import AgentRunSpec
from nanobot.agent.runner import AgentRunner, AgentRunSpec
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import (
LLMProvider,
@@ -22,10 +23,23 @@ from nanobot.providers.base import (
ProviderConversationState,
ToolCallRequest,
)
from nanobot.session.summary import SUMMARY_CONTINUATION_TEXT
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
def _build_transcript(transcript: TranscriptInput) -> list[dict]:
system = (
transcript.session_summary["text"]
if transcript.session_summary is not None
else "system"
)
messages = [{"role": "system", "content": system}, *transcript.history]
if transcript.current_message is not None:
messages.append({"role": transcript.current_role, "content": transcript.current_message})
return messages
def _governance_config(
provider,
tools,
@@ -97,13 +111,16 @@ async def test_runner_locally_fits_oversized_initial_transcript(monkeypatch):
tools = MagicMock()
tools.get_definitions.return_value = []
old_content = "x" * 20_000
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
estimate = MagicMock(
side_effect=lambda _provider, _model, messages, _tools: (
(600, "test-counter")
if any(message.get("content") == old_content for message in messages)
else (100, "test-counter")
),
)
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
result = await AgentRunner().run(make_run_spec(
@@ -127,9 +144,441 @@ async def test_runner_locally_fits_oversized_initial_transcript(monkeypatch):
{"role": "system", "content": "system"},
{"role": "user", "content": "continue"},
]
estimated_messages = [call.args[2] for call in estimate.call_args_list]
assert sum(
any(message.get("content") == old_content for message in messages)
for messages in estimated_messages
) == 1
assert len(estimated_messages) == 3
assert any(message.get("content") == old_content for message in result.messages)
async def test_runner_summarizes_history_and_preserves_current_input(monkeypatch):
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True
prior_state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="test-model",
version=1,
payload={"items": [{"type": "message", "role": "assistant"}]},
)
candidate_state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="test-model",
version=1,
payload={"items": [{"type": "message", "role": "assistant", "fresh": True}]},
)
requests: list[tuple[list[dict], object]] = []
async def request(*, messages, provider_context, **_kwargs):
requests.append((messages, provider_context))
return LLMResponse(content="done", provider_state=candidate_state)
provider.chat_with_retry = request
tools = MagicMock()
tools.get_definitions.return_value = []
old_answer = "old answer " * 2_000
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
(600, "test-counter")
if any(message.get("content") == old_answer for message in messages)
else (100, "test-counter")
),
)
consolidate = AsyncMock(return_value="fresh checkpoint")
previous = {"text": "existing checkpoint", "last_active": "2026-08-30T00:00:00"}
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(
history=[
{"role": "user", "content": "old question"},
{"role": "assistant", "content": old_answer},
],
current_message="continue the current task",
session_summary=previous,
),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
provider_state=prior_state,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
consolidate.assert_awaited_once_with(
[
{"role": "system", "content": "existing checkpoint"},
{"role": "user", "content": "old question"},
{"role": "assistant", "content": old_answer},
],
"existing checkpoint",
)
assert requests[0][0] == [
{"role": "system", "content": "fresh checkpoint"},
{"role": "user", "content": SUMMARY_CONTINUATION_TEXT},
{"role": "user", "content": "continue the current task"},
]
assert requests[0][1].conversation_state is None
assert result.provider_state == candidate_state
assert result.summary_checkpoint is not None
assert result.summary_checkpoint.summary == "fresh checkpoint"
assert result.summary_checkpoint.transcript_boundary == 3
assert any(message.get("content") == old_answer for message in result.messages)
async def test_runner_rejects_oversized_delta_without_summarizable_history(monkeypatch):
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock()
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (600, "test-counter"),
)
consolidate = AsyncMock(return_value=None)
with pytest.raises(ContextWindowExceededError):
await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(
history=[],
current_message="current input is the entire oversized delta",
),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
consolidate.assert_awaited_once_with(
[{"role": "system", "content": "system"}],
None,
)
provider.chat_with_retry.assert_not_awaited()
async def test_runner_governs_history_before_summarizing_it(monkeypatch):
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = False
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="done"))
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
(100, "test-counter")
if messages[0].get("content") == "fresh checkpoint"
else (600, "test-counter")
),
)
consolidate = AsyncMock(return_value="fresh checkpoint")
await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(
history=[
{"role": "user", "content": "inspect"},
{
"role": "assistant",
"content": "",
"tool_calls": [{
"id": "call-missing",
"type": "function",
"function": {"name": "inspect", "arguments": "{}"},
}],
},
],
current_message="continue",
),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
summarized = consolidate.await_args.args[0]
assert [message["role"] for message in summarized] == [
"system", "user", "assistant", "tool",
]
assert summarized[-1]["tool_call_id"] == "call-missing"
assert summarized[-1]["content"] == BACKFILL_CONTENT
@pytest.mark.parametrize(
("scope", "expected_contents", "expected_boundary"),
[
("prior_context", ["system", "accepted question", "accepted answer"], 3),
(
"current_request",
["system", "accepted question", "accepted answer", "inspect the project"],
4,
),
],
)
async def test_native_compaction_uses_provider_request_boundary(
monkeypatch,
scope,
expected_contents,
expected_boundary,
):
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = False
compacted_state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="test-model",
version=1,
payload={"items": [{"type": "compaction", "encrypted_content": "opaque"}]},
)
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="call-1", name="inspect", arguments={})],
provider_compaction_applied=True,
provider_compaction_state=compacted_state,
provider_compaction_scope=scope,
),
LLMResponse(content="done"),
])
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(return_value="complete tool result")
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args: (100, "test-counter"),
)
consolidate = AsyncMock(return_value="portable checkpoint")
consolidate_native = AsyncMock(return_value="portable checkpoint")
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(
history=[
{"role": "user", "content": "accepted question"},
{"role": "assistant", "content": "accepted answer"},
],
current_message="inspect the project",
),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
consolidate_provider_compaction=consolidate_native,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
consolidate.assert_not_awaited()
consolidate_native.assert_awaited_once()
assert consolidate_native.await_args.args[0] == compacted_state
assert [
message["content"] for message in consolidate_native.await_args.args[1]
] == expected_contents
assert consolidate_native.await_args.args[2] is None
assert result.summary_checkpoint is not None
assert result.summary_checkpoint.transcript_boundary == expected_boundary
assert result.provider_compaction_applied is True
assert any(message.get("content") == "inspect the project" for message in result.messages)
assert any(message.get("content") == "complete tool result" for message in result.messages)
async def test_runner_keeps_current_tool_exchange_outside_summary(monkeypatch):
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = False
responses = [
LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="call-1", name="inspect", arguments={})],
),
LLMResponse(content="done"),
]
requests: list[list[dict]] = []
async def request(*, messages, **_kwargs):
requests.append(messages)
return responses.pop(0)
provider.chat_with_retry = request
tools = MagicMock()
tools.get_definitions.return_value = []
full_result = "tool-result:" + ("x" * 4_000)
tools.execute = AsyncMock(return_value=full_result)
def estimate(_provider, _model, messages, _tools):
has_tool_result = any(message.get("role") == "tool" for message in messages)
has_old_system = any(
message.get("role") == "system" and message.get("content") == "system"
for message in messages
)
return (600 if has_tool_result and has_old_system else 100, "test-counter")
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
consolidate = AsyncMock(return_value="fresh checkpoint")
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(history=[], current_message="inspect the project"),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
consolidate.assert_awaited_once_with(
[
{"role": "system", "content": "system"},
{"role": "user", "content": "inspect the project"},
],
None,
)
assert [message["role"] for message in requests[1]] == [
"system", "user", "assistant", "tool",
]
assert requests[1][1]["content"] == SUMMARY_CONTINUATION_TEXT
assert requests[1][-1]["content"] == full_result
assert result.summary_checkpoint is not None
assert result.summary_checkpoint.transcript_boundary == 2
assert any(message.get("content") == full_result for message in result.messages)
async def test_repeated_pressure_advances_summary_boundary(monkeypatch):
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = False
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="call-1", name="inspect", arguments={})],
),
LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="call-2", name="inspect", arguments={})],
),
LLMResponse(content="done"),
])
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(side_effect=["result-1", "result-2"])
def estimate(_provider, _model, messages, _tools):
system = messages[0].get("content")
contents = {message.get("content") for message in messages}
if "result-2" in contents:
return (100 if system == "checkpoint-2" else 600, "test-counter")
if "result-1" in contents:
return (100 if system == "checkpoint-1" else 600, "test-counter")
return 100, "test-counter"
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
consolidate = AsyncMock(side_effect=["checkpoint-1", "checkpoint-2"])
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(history=[], current_message="inspect"),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=3,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert consolidate.await_count == 2
assert consolidate.await_args_list[0].args[1] is None
assert consolidate.await_args_list[1].args[1] == "checkpoint-1"
second_prefix = consolidate.await_args_list[1].args[0]
assert second_prefix[0]["content"] == "checkpoint-1"
assert any(message.get("content") == "result-1" for message in second_prefix)
assert result.final_content == "done"
assert result.summary_checkpoint is not None
assert result.summary_checkpoint.summary == "checkpoint-2"
assert result.summary_checkpoint.transcript_boundary == 4
async def test_runner_refuses_checkpoint_that_cannot_fit_with_delta(monkeypatch):
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="unexpected"))
tools = MagicMock()
tools.get_definitions.return_value = []
old_answer = "old answer"
current_input = "current input must remain intact"
def estimate(_provider, _model, messages, _tools):
contents = {message.get("content") for message in messages}
if old_answer in contents or current_input in contents:
return 600, "test-counter"
return 100, "test-counter"
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
consolidate = AsyncMock(return_value="small checkpoint")
with pytest.raises(ContextWindowExceededError):
await AgentRunner().run(make_run_spec(
provider,
initial_messages=None,
transcript_input=TranscriptInput(
history=[{"role": "assistant", "content": old_answer}],
current_message=current_input,
),
transcript_builder=_build_transcript,
consolidate_history=consolidate,
tools=tools,
model="test-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
summarized = consolidate.await_args.args[0]
assert all(message.get("content") != current_input for message in summarized)
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_governs_messages_added_by_before_iteration_hook(monkeypatch):
from nanobot.agent.hook import AgentHook, AgentHookContext
@@ -145,7 +594,7 @@ async def test_runner_governs_messages_added_by_before_iteration_hook(monkeypatc
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
(2_000, "test-counter")
if any(message.get("content") == oversized for message in messages)
if any(oversized in str(message.get("content")) for message in messages)
else (100, "test-counter")
),
)
@@ -491,6 +940,39 @@ def test_changed_messages_use_local_estimate_after_reported_usage(monkeypatch):
estimate.assert_called_once()
def test_resumed_provider_context_avoids_full_transcript_estimate(monkeypatch):
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
spec = make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "pending delta"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (_ for _ in ()).throw(
AssertionError("resumed provider context must be authoritative")
),
)
pressure = ContextGovernor().request_pressure(
_governance_config(provider, tools, spec),
spec.initial_messages,
LLMUsage.reported(input_tokens=900, output_tokens=10),
usage_matches_messages=False,
tool_definitions=tools.get_definitions(),
request_context_tokens=100,
)
assert pressure is None
@pytest.mark.asyncio
async def test_runner_counts_resumed_provider_state_before_dispatch(monkeypatch):
from nanobot.agent.runner import AgentRunner
@@ -832,7 +1314,6 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar
model="test-model",
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
tests/agent/test_runner_injections.py
+56 -29
View File
@@ -750,25 +750,27 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
),
]
injected = [message for message in result.messages if message.get("role") == "user"][-1]
assert "follow-up from the second speaker" in str(injected["content"])
injected = [message for message in result.messages if message.get("role") == "user"][-2:]
assert str(injected[0]["content"]).startswith("follow-up from the second speaker\n\n")
assert str(injected[1]["content"]).startswith("another follow-up\n\n")
model_messages = provider.chat_with_retry.await_args_list[-1].kwargs["messages"]
assert "telegram | group-1 | user-b | message-2" in str(model_messages)
assert "Bob | topic-7" in str(model_messages)
assert "telegram | group-1 | user-c | message-3" in str(model_messages)
assert "Carol | topic-7" in str(model_messages)
assert injected["_meta"][RUNTIME_CONTEXT_MESSAGE_META]["sources"] == [
"identity",
"identity",
]
assert all(
message["_meta"][RUNTIME_CONTEXT_MESSAGE_META]["sources"] == ["identity"]
for message in injected
)
loop._save_turn(session, result.messages, skip=1)
persisted = [message for message in session.messages if message.get("role") == "user"][-1]
assert "telegram | group-1 | user-b | message-2" in str(persisted["content"])
assert "telegram | group-1 | user-c | message-3" in str(persisted["content"])
assert public_history_message(persisted)["content"] == (
"follow-up from the second speaker\n\nanother follow-up"
)
persisted = [message for message in session.messages if message.get("role") == "user"][-2:]
assert "telegram | group-1 | user-b | message-2" in str(persisted[0]["content"])
assert "telegram | group-1 | user-c | message-3" in str(persisted[1]["content"])
assert [public_history_message(message)["content"] for message in persisted] == [
"follow-up from the second speaker",
"another follow-up",
]
@pytest.mark.asyncio
@@ -835,8 +837,8 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p
@pytest.mark.asyncio
async def test_runner_merges_multiple_injected_user_messages_without_losing_media():
"""Multiple injected follow-ups should not create lossy consecutive user messages."""
async def test_model_request_merges_injected_user_messages_without_losing_media():
"""The model copy may merge follow-ups while the raw transcript keeps each event."""
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
@@ -895,10 +897,17 @@ async def test_runner_merges_multiple_injected_user_messages_without_losing_medi
for block in injected["content"]
if isinstance(block, dict)
)
assert [message["content"] for message in result.messages[-3:-1]] == [
[
{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}},
{"type": "text", "text": "look at this"},
],
"and answer briefly",
]
def test_runner_merge_keeps_all_recovery_followup_ids() -> None:
"""Merged follow-ups stay acknowledged together after a later save."""
def test_runner_append_keeps_recovery_followups_separate() -> None:
"""Each raw follow-up keeps its own recovery identity."""
from nanobot.agent.runner import AgentRunner
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY
@@ -908,10 +917,12 @@ def test_runner_merge_keeps_all_recovery_followup_ids() -> None:
[{"role": "user", "content": "second", PENDING_FOLLOWUP_ID_KEY: "two"}],
)
assert messages[-1][PENDING_FOLLOWUP_ID_KEY] == ["one", "two"]
assert [message["content"] for message in messages] == ["first", "second"]
assert [message[PENDING_FOLLOWUP_ID_KEY] for message in messages] == ["one", "two"]
def test_runner_merge_preserves_runtime_markers_with_media() -> None:
def test_model_request_merge_preserves_runtime_markers_with_media() -> None:
from nanobot.agent.context_governance import ContextGovernor
from nanobot.agent.runner import AgentRunner
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
@@ -948,8 +959,9 @@ def test_runner_merge_preserves_runtime_markers_with_media() -> None:
},
])
assert len(messages) == 1
merged = messages[0]
assert len(messages) == 2
merged = ContextGovernor._merge_adjacent_user_messages_for_model(messages)[0]
assert len(messages) == 2
assert "private first" in str(merged["content"])
assert "private second" in str(merged["content"])
persisted = {
@@ -1681,15 +1693,17 @@ async def test_drain_injections_after_recoverable_tool_error():
@pytest.mark.asyncio
async def test_drain_injections_on_llm_error():
"""Pending injections should be drained when the LLM returns an error finish_reason."""
"""A follow-up after an error stays raw and reaches the next model request."""
from nanobot.agent.runner import AgentRunner
from nanobot.bus.events import InboundMessage
provider = MagicMock()
call_count = {"n": 0}
requests: list[list[dict]] = []
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
requests.append(messages)
if call_count["n"] == 1:
return LLMResponse(
content=None,
@@ -1713,11 +1727,20 @@ async def test_drain_injections_on_llm_error():
runner = AgentRunner()
result = await runner.run(make_run_spec(provider,
initial_messages=[
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "previous response"},
{"role": "user", "content": "trigger error"},
initial_messages=None,
transcript_input=TranscriptInput(
history=[
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "previous response"},
{"role": "user", "content": "trigger error"},
],
current_message=None,
),
transcript_builder=lambda transcript: [
{"role": "system", "content": "system"},
*transcript.history,
],
consolidate_history=AsyncMock(return_value=None),
tools=tools,
model="test-model",
max_iterations=5,
@@ -1727,11 +1750,15 @@ async def test_drain_injections_on_llm_error():
assert result.had_injections is True
assert result.final_content == "recovered answer"
injected = [
m for m in result.messages
if m.get("role") == "user" and "follow-up after LLM error" in str(m.get("content", ""))
assert "follow-up after LLM error" in str(requests[1])
assert [
message["content"]
for message in result.messages
if message.get("role") == "user"
][-2:] == [
"trigger error",
"follow-up after LLM error",
]
assert len(injected) == 1
@pytest.mark.asyncio
tests/agent/test_session_manager_history.py
+36 -170
View File
@@ -1,10 +1,11 @@
from nanobot.providers.base import ProviderConversationState
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
RuntimeContextBlock,
append_runtime_context,
)
from nanobot.session.history_visibility import HIDDEN_HISTORY_META
from nanobot.session.manager import Session, SessionManager
from nanobot.session.summary import SUMMARY_CONTINUATION_TEXT
def _assert_no_orphans(history: list[dict]) -> None:
@@ -136,58 +137,6 @@ def test_legitimate_tool_pairs_preserved_after_trim():
assert history[0]["role"] == "user"
def test_retain_recent_legal_suffix_keeps_recent_messages():
session = Session(key="test:trim")
for i in range(10):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.retain_recent_legal_suffix(4)
assert len(session.messages) == 4
assert session.messages[0]["content"] == "msg6"
assert session.messages[-1]["content"] == "msg9"
def test_retain_recent_legal_suffix_adjusts_last_archived():
session = Session(key="test:trim-cons")
for i in range(10):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.last_archived = 7
session.retain_recent_legal_suffix(4)
assert len(session.messages) == 4
assert session.last_archived == 1
def test_retain_recent_legal_suffix_zero_clears_session():
session = Session(key="test:trim-zero")
for i in range(10):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.last_archived = 5
session.retain_recent_legal_suffix(0)
assert session.messages == []
assert session.last_archived == 0
def test_retain_recent_legal_suffix_keeps_legal_tool_boundary():
session = Session(key="test:trim-tools")
session.messages.append({"role": "user", "content": "old"})
session.messages.extend(_tool_turn("old", 0))
session.messages.append({"role": "user", "content": "keep"})
session.messages.extend(_tool_turn("keep", 0))
session.messages.append({"role": "assistant", "content": "done"})
session.retain_recent_legal_suffix(4)
history = session.get_history(max_messages=500)
_assert_no_orphans(history)
assert history[0]["role"] == "user"
assert history[0]["content"] == "keep"
# --- last_archived > 0 ---
def test_orphan_trim_with_last_archived():
@@ -635,6 +584,40 @@ def test_fork_session_allows_index_equal_to_user_count(tmp_path):
assert [m["content"] for m in forked.messages] == ["round1", "answer1"]
def test_fork_session_user_index_ignores_hidden_checkpoint_anchor(tmp_path):
manager = SessionManager(tmp_path)
source = manager.get_or_create("websocket:source")
source.add_message("user", "round1")
source.add_message("assistant", "answer1")
source.add_message("user", "round2")
source.add_message(
"user",
SUMMARY_CONTINUATION_TEXT,
**{HIDDEN_HISTORY_META: True},
)
source.add_message("assistant", "answer2")
source.last_archived = 3
source.metadata["_last_summary"] = {"text": "round1 and round2"}
manager.save(source)
forked = manager.fork_session_before_user_index(
"websocket:source",
"websocket:fork",
2,
)
assert forked is not None
assert [message["content"] for message in forked.messages] == [
"round1",
"answer1",
"round2",
SUMMARY_CONTINUATION_TEXT,
"answer2",
]
assert forked.last_archived == 3
assert forked.metadata["_last_summary"]["text"] == "round1 and round2"
def test_fork_session_drops_summary_when_fork_point_is_inside_archived_prefix(tmp_path):
manager = SessionManager(tmp_path)
source = manager.get_or_create("websocket:source")
@@ -756,44 +739,6 @@ def test_get_history_recovers_user_when_token_slice_would_be_assistant_only(monk
assert [m["content"] for m in history] == ["u2", "a2"]
def test_retain_recent_legal_suffix_hard_cap_with_long_non_user_chain():
session = Session(key="test:hard-cap-chain")
session.messages.append({"role": "user", "content": "u0"})
session.messages.append(
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "c1", "type": "function", "function": {"name": "x", "arguments": "{}"}}
],
}
)
for i in range(12):
session.messages.append({"role": "assistant", "content": f"a{i}"})
session.retain_recent_legal_suffix(6)
assert len(session.messages) <= 6
def test_retain_recent_legal_suffix_can_extend_to_user_for_long_recent_turn():
session = Session(key="test:extend-to-user")
session.messages.append({"role": "user", "content": "old"})
session.messages.append({"role": "assistant", "content": "old answer"})
session.messages.append({"role": "user", "content": "record this"})
for i in range(4):
session.messages.extend(_tool_turn("recent", i))
session.messages.append({"role": "assistant", "content": "done"})
session.retain_recent_legal_suffix(8, extend_to_user=True)
assert len(session.messages) > 8
assert session.messages[0]["content"] == "record this"
assert session.messages[-1]["content"] == "done"
history = session.get_history(max_messages=500)
_assert_no_orphans(history)
def test_get_history_can_extend_to_user_for_long_recent_turn():
session = Session(key="test:history-extend-to-user")
session.messages.append({"role": "user", "content": "old"})
@@ -828,82 +773,3 @@ def test_get_history_extend_to_user_keeps_newer_user_inside_window():
assert [m["content"] for m in history] == ["new question", "new answer"]
_assert_no_orphans(history)
def test_retain_recent_legal_suffix_returns_dropped_messages():
"""retain_recent_legal_suffix returns the actually-dropped messages."""
session = Session(
key="test:return-dropped",
provider_state=ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="test-model",
version=1,
payload={"items": []},
),
)
for i in range(10):
session.messages.append({"role": "user", "content": f"msg{i}"})
result = session.retain_recent_legal_suffix(4)
assert len(result.dropped) == 6
assert [m["content"] for m in result.dropped] == [f"msg{i}" for i in range(6)]
assert len(session.messages) == 4
assert result.already_consolidated_count == 0
assert session.provider_state is None
def test_retain_recent_legal_suffix_returns_empty_when_no_drop():
"""No messages dropped → empty list returned."""
state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="test-model",
version=1,
payload={"items": []},
)
session = Session(key="test:no-drop", provider_state=state)
for i in range(3):
session.messages.append({"role": "user", "content": f"msg{i}"})
result = session.retain_recent_legal_suffix(4)
assert result.dropped == []
assert result.already_consolidated_count == 0
assert len(session.messages) == 3
assert session.provider_state is state
def test_retain_recent_legal_suffix_returns_all_on_zero():
"""max_messages=0 clears session and returns all messages."""
session = Session(key="test:zero-return")
for i in range(5):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.last_archived = 3
result = session.retain_recent_legal_suffix(0)
assert len(result.dropped) == 5
assert result.already_consolidated_count == 3
assert session.messages == []
def test_retain_recent_legal_suffix_last_archived_correct_in_else_branch():
"""last_archived should count retained messages from the old archived prefix."""
session = Session(key="test:else-lc-correct")
# 20 messages: u0..u9, a0..a9
for i in range(10):
session.messages.append({"role": "user", "content": f"u{i}"})
for i in range(10):
session.messages.append({"role": "assistant", "content": f"a{i}"})
session.last_archived = 12 # u0..u9, a0, a1 archived
result = session.retain_recent_legal_suffix(4)
# Retained messages start from latest user (u9) + max_messages forward
# so retained = [u9, a0..a9][:4] → but these are from original indices 9..12
# Of those, indices 9,10,11 are < 12 (before_lc), so new_lc = 3
assert session.last_archived == 3
# already_cons should count dropped messages with original index < 12
assert result.already_consolidated_count == 9
tests/agent/test_session_retention.py
-225
View File
@@ -1,225 +0,0 @@
from nanobot.session.manager import Session
def _assert_no_orphans(history: list[dict]) -> None:
declared = {
tc["id"]
for m in history
if m.get("role") == "assistant"
for tc in (m.get("tool_calls") or [])
}
orphans = [
m.get("tool_call_id")
for m in history
if m.get("role") == "tool" and m.get("tool_call_id") not in declared
]
assert orphans == [], f"orphan tool_call_ids: {orphans}"
def _delivery(content: str) -> dict:
return {"role": "assistant", "content": content, "_channel_delivery": True}
def _tool_turn(prefix: str, idx: int) -> list[dict]:
return [
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": f"{prefix}_{idx}_a",
"type": "function",
"function": {"name": "x", "arguments": "{}"},
},
{
"id": f"{prefix}_{idx}_b",
"type": "function",
"function": {"name": "y", "arguments": "{}"},
},
],
},
{"role": "tool", "tool_call_id": f"{prefix}_{idx}_a", "name": "x", "content": "ok"},
{"role": "tool", "tool_call_id": f"{prefix}_{idx}_b", "name": "y", "content": "ok"},
]
def _contents(messages: list[dict]) -> list[str]:
return [m.get("content") for m in messages]
def _has_delivery(messages: list[dict]) -> bool:
return any(m.get("_channel_delivery") for m in messages)
# --- Hard-cap trimming must preserve a proactive delivery the user replied to ---
def test_retain_hard_cap_keeps_delivery_before_user():
session = Session(key="test:cap-delivery")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append(_delivery("Remember to drink water"))
session.messages.append({"role": "user", "content": "ok"})
session.messages.append({"role": "assistant", "content": "great"})
session.retain_recent_legal_suffix(3)
assert _has_delivery(session.messages), "delivery dropped by hard-cap trim"
assert _contents(session.messages) == [
"Remember to drink water",
"ok",
"great",
]
def test_retain_hard_cap_matches_get_history_boundary():
"""The trimmed suffix must start on the same message as get_history()."""
session = Session(key="test:cap-boundary")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append(_delivery("You have 3 pending tasks"))
session.messages.append({"role": "user", "content": "show them"})
session.messages.append({"role": "assistant", "content": "done"})
expected = session.get_history(max_messages=3)
session.retain_recent_legal_suffix(3)
assert _contents(session.messages) == _contents(expected)
def test_retain_extend_to_user_keeps_delivery_before_recovered_user():
session = Session(key="test:extend-delivery")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append({"role": "assistant", "content": "work"})
session.messages.append(_delivery("Reminder: deploy at 17:00"))
session.messages.append({"role": "user", "content": "ok"})
session.messages.append({"role": "assistant", "content": "a1"})
session.messages.append({"role": "assistant", "content": "a2"})
session.messages.append({"role": "assistant", "content": "a3"})
session.retain_recent_legal_suffix(3, extend_to_user=True)
assert _has_delivery(session.messages), "delivery dropped by extend_to_user trim"
assert session.messages[0]["content"] == "Reminder: deploy at 17:00"
assert session.messages[-1]["content"] == "a3"
def test_retain_extend_to_user_matches_get_history_boundary():
session = Session(key="test:extend-boundary")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append({"role": "assistant", "content": "work"})
session.messages.append(_delivery("Reminder: review the draft"))
session.messages.append({"role": "user", "content": "ok"})
session.messages.append({"role": "assistant", "content": "a1"})
session.messages.append({"role": "assistant", "content": "a2"})
session.messages.append({"role": "assistant", "content": "a3"})
expected = session.get_history(max_messages=3, extend_to_user=True)
session.retain_recent_legal_suffix(3, extend_to_user=True)
assert _contents(session.messages) == _contents(expected)
def test_retain_extend_to_user_does_not_extend_delivery_only_tail():
session = Session(key="test:extend-no-user")
for i in range(4):
session.messages.append(_delivery(f"notification {i}"))
session.retain_recent_legal_suffix(3, extend_to_user=True)
assert _contents(session.messages) == [
"notification 1",
"notification 2",
"notification 3",
]
# --- Only the immediately-preceding delivery is part of the anchor ---
def test_retain_keeps_only_immediate_delivery():
session = Session(key="test:multi-delivery")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append(_delivery("old scheduled note"))
session.messages.append(_delivery("new scheduled note"))
session.messages.append({"role": "user", "content": "ok"})
session.messages.append({"role": "assistant", "content": "great"})
session.retain_recent_legal_suffix(3)
kept = _contents(session.messages)
assert kept == ["new scheduled note", "ok", "great"], kept
def test_retain_drops_delivery_not_adjacent_to_anchor_user():
"""A delivery that does not immediately precede the retained user turn is
not part of the anchor and should not be force-retained."""
session = Session(key="test:nonadjacent")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append(_delivery("unrelated scheduled note"))
session.messages.append({"role": "assistant", "content": "reply"})
session.messages.append({"role": "user", "content": "ok"})
session.messages.append({"role": "assistant", "content": "great"})
session.retain_recent_legal_suffix(2)
assert not _has_delivery(session.messages)
assert _contents(session.messages) == ["ok", "great"]
def test_compact_probe_keeps_delivery_in_visible_suffix():
"""compact_idle_session() trims a probe copy with extend_to_user=True; the
visible suffix it keeps must still contain the delivery message."""
tail = [
{"role": "user", "content": "setup"},
{"role": "assistant", "content": "work"},
_delivery("Reminder: deploy at 17:00"),
{"role": "user", "content": "ok"},
{"role": "assistant", "content": "a1"},
{"role": "assistant", "content": "a2"},
{"role": "assistant", "content": "a3"},
]
probe = Session(key="test:probe", messages=tail)
probe.retain_recent_legal_suffix(3, extend_to_user=True)
assert _has_delivery(probe.messages)
assert probe.messages[0]["content"] == "Reminder: deploy at 17:00"
# --- Trimming must stay coherent with the rest of replay ---
def test_retain_then_replay_keeps_delivery_and_no_orphans():
session = Session(key="test:replay-after-trim")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append(_delivery("You have 3 pending tasks"))
session.messages.append({"role": "user", "content": "show them"})
session.messages.extend(_tool_turn("cur", 0))
session.messages.append({"role": "assistant", "content": "done"})
session.retain_recent_legal_suffix(6)
assert _has_delivery(session.messages)
history = session.get_history(max_messages=500)
_assert_no_orphans(history)
assert any(m.get("content") == "You have 3 pending tasks" for m in history)
def test_retain_keeps_delivery_when_user_inside_window():
"""When the capped window already contains a user, its immediately
preceding delivery must stay attached to it."""
session = Session(key="test:window-user")
session.messages.append({"role": "user", "content": "setup"})
session.messages.append({"role": "assistant", "content": "a0"})
session.messages.append(_delivery("Reminder"))
session.messages.append({"role": "user", "content": "ok"})
session.messages.append({"role": "assistant", "content": "a1"})
session.messages.append({"role": "assistant", "content": "a2"})
expected = session.get_history(max_messages=4)
session.retain_recent_legal_suffix(4)
assert _has_delivery(session.messages)
assert _contents(session.messages) == _contents(expected)
tests/agent/test_unified_session.py
+1 -111
View File
@@ -28,7 +28,7 @@ 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.session.manager import SessionManager
from nanobot.utils.llm_runtime import LLMRuntime
# ---------------------------------------------------------------------------
@@ -334,116 +334,6 @@ class TestCmdNewUnifiedSession:
assert len(sessions.get_or_create("discord:999").messages) == 1
# ---------------------------------------------------------------------------
# TestConsolidationUnaffectedByUnifiedSession — consolidation is key-agnostic
# ---------------------------------------------------------------------------
class TestConsolidationUnaffectedByUnifiedSession:
"""maybe_consolidate_by_tokens() behaviour is identical regardless of session key."""
@pytest.mark.asyncio
async def test_consolidation_skips_empty_session_for_unified_key(self):
"""Empty unified:default session → consolidation exits immediately, archive not called."""
from nanobot.agent.memory import Consolidator, MemoryStore
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,
sessions=sessions,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
consolidator.archive_session = AsyncMock()
session = Session(key="unified:default")
session.messages = []
sessions.get_or_create.return_value = session
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_not_called()
@pytest.mark.asyncio
async def test_consolidation_behaviour_identical_for_any_key(self):
"""Archive call count is the same for 'telegram:123' and 'unified:default'
under identical token conditions."""
from nanobot.agent.memory import Consolidator, MemoryStore
archive_calls: dict[str, int] = {}
for key in ("telegram:123", "unified:default"):
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,
sessions=sessions,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
session = Session(key=key)
session.messages = [] # empty → exits immediately for both keys
sessions.get_or_create.return_value = session
consolidator.archive_session = AsyncMock()
await consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
)
archive_calls[key] = consolidator.archive_session.call_count
assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0
@pytest.mark.asyncio
async def test_consolidation_triggers_when_over_budget_unified_key(self):
"""When tokens exceed budget, consolidation attempts to find a boundary —
behaviour is identical to any other session key."""
from nanobot.agent.memory import Consolidator, MemoryStore
store = MagicMock(spec=MemoryStore)
mock_provider = MagicMock()
runtime = _runtime(mock_provider)
sessions = MagicMock(spec=SessionManager)
consolidator = Consolidator(
store=store,
sessions=sessions,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
session = Session(key="unified:default")
session.messages = [{"role": "user", "content": "msg"}]
sessions.get_or_create.return_value = session
# Simulate over-budget: estimated > budget
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(950, "tiktoken"))
# No valid boundary found → returns gracefully without archiving
consolidator.pick_consolidation_boundary = MagicMock(return_value=None)
consolidator.archive_session = AsyncMock()
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_session.assert_not_called()
# ---------------------------------------------------------------------------
# TestStopCommandWithUnifiedSession — /stop command integration
# ---------------------------------------------------------------------------