mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-09-04 02:01:48 +03:00
refactor(agent): let runner own context compaction (#5568)
* 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
@@ -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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user