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

refactor(agent): make checkpoint recovery ownership explicit

Browse Source
This commit is contained in:
chengyongru
2026-08-26 16:33:59 +08:00
committed by chengyongru
parent 9f5a56f1ec
commit 56aa7296f3
4 changed files with 47 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+2 -6
View File
@@ -1466,7 +1466,7 @@ class AgentLoop:
try:
key = self._effective_session_key(msg)
session = self.sessions.get_or_create(key)
if self._restore_runtime_checkpoint(session):
if restore_runtime_checkpoint(session):
self._clear_pending_user_turn(session)
self.sessions.save(session)
logger.info(
@@ -1823,7 +1823,7 @@ class AgentLoop:
if ctx.kind is TurnKind.USER:
self.workspace_scopes.persist_message_scope(session, msg)
if self._restore_runtime_checkpoint(session):
if restore_runtime_checkpoint(session):
self.sessions.save(session)
if (
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
@@ -2313,10 +2313,6 @@ class AgentLoop:
if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None)
def _restore_runtime_checkpoint(self, session: Session) -> bool:
"""Materialize an unfinished turn into session history before a new request."""
return restore_runtime_checkpoint(session)
async def process_direct(
self,
content: str,
nanobot/session/recovery.py
+3 -3
View File
@@ -1,8 +1,8 @@
"""Durable, side-effect-safe recovery for interrupted WebUI turns.
The coordinator owns restart policy. AgentLoop only exposes checkpoint
materialization and an admission hook, so transport code never has to guess
whether an interrupted tool call is safe to replay.
The coordinator owns restart policy. Checkpoint materialization is a session
operation shared with AgentLoop lifecycle boundaries, so transport code never
has to guess whether an interrupted tool call is safe to replay.
"""
from __future__ import annotations
tests/agent/test_loop_save_turn.py
+23 -25
View File
@@ -37,7 +37,14 @@ from nanobot.session.keys import (
UNIFIED_SESSION_KEY,
)
from nanobot.session.manager import Session
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY, PENDING_FOLLOWUPS_KEY
from nanobot.session.recovery import (
PENDING_FOLLOWUP_ID_KEY,
PENDING_FOLLOWUPS_KEY,
PROVIDER_STATE_CHECKPOINT_VERSION,
PROVIDER_STATE_CHECKPOINT_VERSION_KEY,
RUNTIME_CHECKPOINT_KEY,
restore_runtime_checkpoint,
)
from nanobot.session.turn_continuation import (
INTERNAL_CONTINUATION_META,
INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
@@ -682,12 +689,11 @@ def test_save_turn_stamps_latency_on_last_assistant() -> None:
def test_restore_runtime_checkpoint_rehydrates_completed_and_pending_tools() -> None:
loop = _mk_loop()
session = Session(
key="test:checkpoint",
provider_state=_provider_state(),
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"assistant_message": {
"role": "assistant",
"content": "working",
@@ -723,10 +729,10 @@ def test_restore_runtime_checkpoint_rehydrates_completed_and_pending_tools() ->
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is None
assert session.metadata.get(RUNTIME_CHECKPOINT_KEY) is None
assert session.messages[0]["role"] == "assistant"
assert session.messages[1]["tool_call_id"] == "call_done"
assert session.messages[2]["tool_call_id"] == "call_pending"
@@ -735,17 +741,14 @@ def test_restore_runtime_checkpoint_rehydrates_completed_and_pending_tools() ->
def test_restore_final_response_checkpoint_preserves_matching_provider_state() -> None:
loop = _mk_loop()
state = _provider_state()
session = Session(
key="test:final-checkpoint",
provider_state=state,
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"phase": "final_response",
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION_KEY: (
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION
),
PROVIDER_STATE_CHECKPOINT_VERSION_KEY: PROVIDER_STATE_CHECKPOINT_VERSION,
"assistant_message": {
"role": "assistant",
"content": "finished",
@@ -756,21 +759,20 @@ def test_restore_final_response_checkpoint_preserves_matching_provider_state() -
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.messages[-1]["content"] == "finished"
assert session.provider_state is state
assert session.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is None
assert session.metadata.get(RUNTIME_CHECKPOINT_KEY) is None
def test_restore_legacy_final_checkpoint_discards_unproven_provider_state() -> None:
loop = _mk_loop()
session = Session(
key="test:legacy-final-checkpoint",
provider_state=_provider_state(),
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"phase": "final_response",
"assistant_message": {
"role": "assistant",
@@ -782,7 +784,7 @@ def test_restore_legacy_final_checkpoint_discards_unproven_provider_state() -> N
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.messages[-1]["content"] == "finished"
@@ -790,7 +792,6 @@ def test_restore_legacy_final_checkpoint_discards_unproven_provider_state() -> N
def test_restore_completed_tools_checkpoint_preserves_matching_provider_state() -> None:
loop = _mk_loop()
tool_result = {
"role": "tool",
"tool_call_id": "call_done",
@@ -802,11 +803,9 @@ def test_restore_completed_tools_checkpoint_preserves_matching_provider_state()
key="test:completed-tools-checkpoint",
provider_state=state,
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"phase": "tools_completed",
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION_KEY: (
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION
),
PROVIDER_STATE_CHECKPOINT_VERSION_KEY: PROVIDER_STATE_CHECKPOINT_VERSION,
"assistant_message": {
"role": "assistant",
"content": None,
@@ -824,7 +823,7 @@ def test_restore_completed_tools_checkpoint_preserves_matching_provider_state()
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.messages[-1]["content"] == "compacted result"
@@ -832,7 +831,6 @@ def test_restore_completed_tools_checkpoint_preserves_matching_provider_state()
def test_restore_runtime_checkpoint_dedupes_overlapping_tail() -> None:
loop = _mk_loop()
session = Session(
key="test:checkpoint-overlap",
messages=[
@@ -860,7 +858,7 @@ def test_restore_runtime_checkpoint_dedupes_overlapping_tail() -> None:
},
],
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"assistant_message": {
"role": "assistant",
"content": "working",
@@ -896,10 +894,10 @@ def test_restore_runtime_checkpoint_dedupes_overlapping_tail() -> None:
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is None
assert session.metadata.get(RUNTIME_CHECKPOINT_KEY) is None
assert len(session.messages) == 3
assert session.messages[0]["role"] == "assistant"
assert session.messages[1]["tool_call_id"] == "call_done"
tests/agent/test_stop_preserves_context.py
+19 -51
View File
@@ -18,6 +18,7 @@ import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.session.recovery import RUNTIME_CHECKPOINT_KEY
def _make_provider():
@@ -41,51 +42,6 @@ def _make_loop(tmp_path: Path) -> AgentLoop:
return AgentLoop(bus=bus, provider=provider, workspace=tmp_path)
class TestStopPreservesContext:
"""Verify that /stop restores partial context via checkpoint."""
def test_restore_checkpoint_method_exists(self, tmp_path):
"""AgentLoop should have _restore_runtime_checkpoint."""
loop = _make_loop(tmp_path)
assert hasattr(loop, "_restore_runtime_checkpoint")
def test_checkpoint_key_constant(self, tmp_path):
"""The runtime checkpoint key should be defined."""
loop = _make_loop(tmp_path)
assert loop._RUNTIME_CHECKPOINT_KEY == "runtime_checkpoint"
def test_cancel_dispatch_restores_checkpoint(self, tmp_path):
"""When a task is cancelled, the checkpoint should be restored."""
loop = _make_loop(tmp_path)
session = MagicMock()
session.metadata = {
"runtime_checkpoint": {
"phase": "awaiting_tools",
"iteration": 0,
"assistant_message": {
"role": "assistant",
"content": "Let me search for that.",
"tool_calls": [{"id": "tc_1", "type": "function",
"function": {"name": "web_search", "arguments": "{}"}}],
},
"completed_tool_results": [
{"role": "tool", "tool_call_id": "tc_1",
"content": "Search results: ..."},
],
"pending_tool_calls": [],
}
}
session.messages = [
{"role": "user", "content": "Search for something"},
]
loop.sessions.get_or_create.return_value = session
restored = loop._restore_runtime_checkpoint(session)
assert restored is True
assert len(session.messages) > 1
assert "runtime_checkpoint" not in session.metadata
@pytest.mark.asyncio
async def test_dispatch_cancellation_restores_checkpoint():
"""Regression for #2966: /stop interrupting _dispatch must materialize the
@@ -93,9 +49,8 @@ async def test_dispatch_cancellation_restores_checkpoint():
unwinds, so the next turn can see the partial work.
This exercises the real _dispatch path (locks, pending queues, the
CancelledError handler) rather than poking _restore_runtime_checkpoint in
isolation, so a future refactor that drops the cancel-time restore is
caught by CI instead of silently regressing.
CancelledError handler), so a future refactor that drops the cancel-time
restore is caught by CI instead of silently regressing.
"""
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
@@ -112,7 +67,7 @@ async def test_dispatch_cancellation_restores_checkpoint():
mock_subagent_manager.return_value.cancel_by_session = AsyncMock(return_value=0)
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
checkpoint_key = loop._RUNTIME_CHECKPOINT_KEY
checkpoint_key = RUNTIME_CHECKPOINT_KEY
session = SimpleNamespace(
key="test:c1",
metadata={
@@ -168,7 +123,19 @@ async def test_dispatch_cancellation_keeps_checkpoint_for_gateway_shutdown(tmp_p
"""Gateway shutdown preserves the checkpoint; an explicit stop restores it."""
loop = _make_loop(tmp_path)
loop.preserve_inflight_turns_on_shutdown()
loop._restore_runtime_checkpoint = MagicMock() # type: ignore[method-assign]
checkpoint_key = RUNTIME_CHECKPOINT_KEY
checkpoint = {
"phase": "final_response",
"assistant_message": {"role": "assistant", "content": "finished"},
"completed_tool_results": [],
"pending_tool_calls": [],
}
session = SimpleNamespace(
metadata={checkpoint_key: checkpoint},
messages=[],
provider_state=None,
)
loop.sessions.get_or_create.return_value = session
async def _cancel(*_args: object, **_kwargs: object) -> None:
raise asyncio.CancelledError()
@@ -182,4 +149,5 @@ async def test_dispatch_cancellation_keeps_checkpoint_for_gateway_shutdown(tmp_p
InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="work")
)
loop._restore_runtime_checkpoint.assert_not_called()
assert session.metadata[checkpoint_key] == checkpoint
assert session.messages == []