refactor(agent): make checkpoint recovery ownership explicit

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,

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
+The coordinator owns restart policy.  Checkpoint materialization is a session
-materialization and an admission hook, so transport code never has to guess
+operation shared with AgentLoop lifecycle boundaries, so transport code never
-whether an interrupted tool call is safe to replay.
+has to guess whether an interrupted tool call is safe to replay.
 """
 
 from __future__ import annotations

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: (
+                PROVIDER_STATE_CHECKPOINT_VERSION_KEY: PROVIDER_STATE_CHECKPOINT_VERSION,
-                    AgentLoop._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: (
+                PROVIDER_STATE_CHECKPOINT_VERSION_KEY: PROVIDER_STATE_CHECKPOINT_VERSION,
-                    AgentLoop._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

@@ -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
+    CancelledError handler), so a future refactor that drops the cancel-time
-    isolation, so a future refactor that drops the cancel-time restore is
+    restore is caught by CI instead of silently regressing.
-    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 == []