From 00de55072de307f992f1731d0cc129073d2eb89e Mon Sep 17 00:00:00 2001
From: Xubin Ren <xubinrencs@gmail.com>
Date: Mon, 20 Apr 2026 17:13:04 +0000
Subject: [PATCH] test(agent): exercise /stop cancellation through _dispatch

Add a regression test that actually runs the CancelledError branch of
AgentLoop._dispatch end-to-end and asserts the in-flight checkpoint is
materialized into session.messages before the cancellation unwinds.

The three existing tests call _restore_runtime_checkpoint directly, so
they pass even if the cancel-time restore is ever removed from
_dispatch. This new test is the one that actually locks the fix in
place.

Made-with: Cursor
---
 tests/agent/test_stop_preserves_context.py | 78 ++++++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/tests/agent/test_stop_preserves_context.py b/tests/agent/test_stop_preserves_context.py
index 0f4047b9..2a082850 100644
--- a/tests/agent/test_stop_preserves_context.py
+++ b/tests/agent/test_stop_preserves_context.py
@@ -10,6 +10,7 @@ See: https://github.com/HKUDS/nanobot/issues/2966
 from __future__ import annotations
 
 import asyncio
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import MagicMock, patch, AsyncMock
 
@@ -82,3 +83,80 @@ class TestStopPreservesContext:
         assert len(session.messages) > 1
         # The checkpoint should be cleared
         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
+    in-flight runtime checkpoint into session.messages before the cancellation
+    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.
+    """
+    from nanobot.bus.events import InboundMessage
+    from nanobot.bus.queue import MessageBus
+
+    bus = MessageBus()
+    provider = MagicMock()
+    provider.get_default_model.return_value = "test-model"
+    workspace = MagicMock()
+    workspace.__truediv__ = MagicMock(return_value=MagicMock())
+
+    with patch("nanobot.agent.loop.ContextBuilder"), \
+         patch("nanobot.agent.loop.SessionManager"), \
+         patch("nanobot.agent.loop.SubagentManager") as MockSubMgr:
+        MockSubMgr.return_value.cancel_by_session = AsyncMock(return_value=0)
+        loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
+
+    checkpoint_key = loop._RUNTIME_CHECKPOINT_KEY
+    session = SimpleNamespace(
+        key="test:c1",
+        metadata={
+            checkpoint_key: {
+                "phase": "awaiting_tools",
+                "iteration": 0,
+                "assistant_message": {
+                    "role": "assistant",
+                    "content": "Let me search.",
+                    "tool_calls": [
+                        {
+                            "id": "tc_1",
+                            "type": "function",
+                            "function": {"name": "web_search", "arguments": "{}"},
+                        }
+                    ],
+                },
+                "completed_tool_results": [
+                    {"role": "tool", "tool_call_id": "tc_1", "content": "Search hit."},
+                ],
+                "pending_tool_calls": [],
+            }
+        },
+        messages=[{"role": "user", "content": "Search for something"}],
+    )
+
+    loop.sessions.get_or_create = MagicMock(return_value=session)
+    loop.sessions.save = MagicMock()
+
+    async def _cancel(*_args, **_kwargs):
+        raise asyncio.CancelledError()
+
+    loop._process_message = _cancel
+
+    msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="work")
+
+    with pytest.raises(asyncio.CancelledError):
+        await loop._dispatch(msg)
+
+    roles = [m.get("role") for m in session.messages]
+    assert roles == ["user", "assistant", "tool"], (
+        "Expected the assistant message and completed tool result from the "
+        f"interrupted turn to be materialized into session.messages; got {roles}"
+    )
+    assert checkpoint_key not in session.metadata, \
+        "Checkpoint metadata should be cleared after restore"
+    assert loop.sessions.save.called, \
+        "Session should be persisted so the restored state survives process restart"