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
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
from types import SimpleNamespace
|
||||||
from typing import Any
|
from typing import Any
|
||||||
from unittest.mock import MagicMock, patch, AsyncMock
|
from unittest.mock import MagicMock, patch, AsyncMock
|
||||||
|
|
||||||
@@ -82,3 +83,80 @@ class TestStopPreservesContext:
|
|||||||
assert len(session.messages) > 1
|
assert len(session.messages) > 1
|
||||||
# The checkpoint should be cleared
|
# The checkpoint should be cleared
|
||||||
assert "runtime_checkpoint" not in session.metadata
|
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