diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 258d2b802..b4273e78b 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -1051,6 +1051,10 @@ class AgentRunner: await coro if outer_timeout_s is None else await asyncio.wait_for(coro, timeout=outer_timeout_s) ) + except asyncio.CancelledError: + _pause_generation() + await _close_native_reasoning() + raise except asyncio.TimeoutError: if outer_timeout_s is None: response = LLMResponse( diff --git a/tests/agent/test_runner_reasoning.py b/tests/agent/test_runner_reasoning.py index 82026b837..1ecdde672 100644 --- a/tests/agent/test_runner_reasoning.py +++ b/tests/agent/test_runner_reasoning.py @@ -9,6 +9,7 @@ channels, gated by ``context.streamed_reasoning`` rather than from __future__ import annotations +import asyncio from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -554,6 +555,46 @@ async def test_runner_closes_native_reasoning_before_hosted_tool_event(): ] +@pytest.mark.asyncio +async def test_runner_closes_native_reasoning_when_stream_is_cancelled(): + from nanobot.agent.runner import AgentRunner + + provider = MagicMock() + reasoning_started = asyncio.Event() + release_provider = asyncio.Event() + + async def chat_stream_with_retry( + *, on_thinking_delta=None, **kwargs + ): + if on_thinking_delta: + await on_thinking_delta("inspect") + reasoning_started.set() + await release_provider.wait() + raise AssertionError("the cancelled provider call should not complete") + + provider.chat_stream_with_retry = chat_stream_with_retry + tools = MagicMock() + tools.get_definitions.return_value = [] + hook = _LifecycleRecordingHook() + + task = asyncio.create_task(AgentRunner().run(make_run_spec( + provider, + initial_messages=[{"role": "user", "content": "inspect"}], + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + hook=hook, + ))) + await reasoning_started.wait() + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert hook.events == ["reasoning:inspect", "reasoning_end"] + + @pytest.mark.asyncio async def test_runner_strips_thinking_tags_from_native_thinking_deltas(): from nanobot.agent.runner import AgentRunner