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
This commit is contained in:
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user