From 56aa7296f35e9a952b5c06dc91c1c5ce31972c48 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Wed, 26 Aug 2026 16:23:19 +0800 Subject: [PATCH] refactor(agent): make checkpoint recovery ownership explicit --- nanobot/agent/loop.py | 8 +-- nanobot/session/recovery.py | 6 +- tests/agent/test_loop_save_turn.py | 48 +++++++-------- tests/agent/test_stop_preserves_context.py | 70 ++++++---------------- 4 files changed, 47 insertions(+), 85 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 789508dc1..a5defd4e7 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -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, diff --git a/nanobot/session/recovery.py b/nanobot/session/recovery.py index 8d8e248b8..b2016f0f6 100644 --- a/nanobot/session/recovery.py +++ b/nanobot/session/recovery.py @@ -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 diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index 18069bedd..807291fce 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -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" diff --git a/tests/agent/test_stop_preserves_context.py b/tests/agent/test_stop_preserves_context.py index d53abab93..2034ac932 100644 --- a/tests/agent/test_stop_preserves_context.py +++ b/tests/agent/test_stop_preserves_context.py @@ -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 == []