diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index b4273e78b..63f7e7670 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -949,6 +949,7 @@ class AgentRunner: active_hosted_tools: dict[str, dict[str, Any]] = {} native_reasoning_open = False + native_reasoning_close_task: asyncio.Task[None] | None = None request_started_at = 0.0 first_output_at: float | None = None generation_started_at: float | None = None @@ -972,11 +973,29 @@ class AgentRunner: generation_started_at = None async def _close_native_reasoning() -> None: - nonlocal native_reasoning_open - if not native_reasoning_open: - return - native_reasoning_open = False - await hook.emit_reasoning_end() + nonlocal native_reasoning_open, native_reasoning_close_task + if native_reasoning_close_task is None: + if not native_reasoning_open: + return + native_reasoning_open = False + native_reasoning_close_task = asyncio.create_task( + hook.emit_reasoning_end() + ) + + close_task = native_reasoning_close_task + cancellation: asyncio.CancelledError | None = None + while not close_task.done(): + try: + await asyncio.shield(close_task) + except asyncio.CancelledError as exc: + cancellation = cancellation or exc + try: + close_task.result() + finally: + if native_reasoning_close_task is close_task: + native_reasoning_close_task = None + if cancellation is not None: + raise cancellation async def _provider_tool_event(event: dict[str, Any]) -> None: if event.get("kind") != "hosted_tool": diff --git a/tests/agent/test_runner_reasoning.py b/tests/agent/test_runner_reasoning.py index 1ecdde672..dca0abec9 100644 --- a/tests/agent/test_runner_reasoning.py +++ b/tests/agent/test_runner_reasoning.py @@ -83,6 +83,18 @@ class _LifecycleRecordingHook(AgentHook): self.events.append(f"hosted_tool:{event.get('phase')}") +class _BlockingReasoningEndHook(_LifecycleRecordingHook): + def __init__(self) -> None: + super().__init__() + self.reasoning_end_started = asyncio.Event() + self.release_reasoning_end = asyncio.Event() + + async def emit_reasoning_end(self) -> None: + self.reasoning_end_started.set() + await self.release_reasoning_end.wait() + await super().emit_reasoning_end() + + @pytest.mark.asyncio async def test_runner_preserves_reasoning_fields_in_assistant_history(): """Reasoning fields ride along on the persisted assistant message so @@ -595,6 +607,46 @@ async def test_runner_closes_native_reasoning_when_stream_is_cancelled(): assert hook.events == ["reasoning:inspect", "reasoning_end"] +@pytest.mark.asyncio +async def test_runner_settles_native_reasoning_end_before_propagating_cancellation(): + from nanobot.agent.runner import AgentRunner + + provider = MagicMock() + + async def chat_stream_with_retry( + *, on_content_delta=None, on_thinking_delta=None, **kwargs + ): + if on_thinking_delta: + await on_thinking_delta("inspect") + if on_content_delta: + await on_content_delta("done") + 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 = _BlockingReasoningEndHook() + + 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 hook.reasoning_end_started.wait() + + task.cancel() + await asyncio.sleep(0) + hook.release_reasoning_end.set() + 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