mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-09-01 00:31:51 +03:00
fix(agent): settle reasoning close before cancellation
This commit is contained in:
+24
-5
@@ -949,6 +949,7 @@ class AgentRunner:
|
|||||||
|
|
||||||
active_hosted_tools: dict[str, dict[str, Any]] = {}
|
active_hosted_tools: dict[str, dict[str, Any]] = {}
|
||||||
native_reasoning_open = False
|
native_reasoning_open = False
|
||||||
|
native_reasoning_close_task: asyncio.Task[None] | None = None
|
||||||
request_started_at = 0.0
|
request_started_at = 0.0
|
||||||
first_output_at: float | None = None
|
first_output_at: float | None = None
|
||||||
generation_started_at: float | None = None
|
generation_started_at: float | None = None
|
||||||
@@ -972,11 +973,29 @@ class AgentRunner:
|
|||||||
generation_started_at = None
|
generation_started_at = None
|
||||||
|
|
||||||
async def _close_native_reasoning() -> None:
|
async def _close_native_reasoning() -> None:
|
||||||
nonlocal native_reasoning_open
|
nonlocal native_reasoning_open, native_reasoning_close_task
|
||||||
if not native_reasoning_open:
|
if native_reasoning_close_task is None:
|
||||||
return
|
if not native_reasoning_open:
|
||||||
native_reasoning_open = False
|
return
|
||||||
await hook.emit_reasoning_end()
|
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:
|
async def _provider_tool_event(event: dict[str, Any]) -> None:
|
||||||
if event.get("kind") != "hosted_tool":
|
if event.get("kind") != "hosted_tool":
|
||||||
|
|||||||
@@ -83,6 +83,18 @@ class _LifecycleRecordingHook(AgentHook):
|
|||||||
self.events.append(f"hosted_tool:{event.get('phase')}")
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_runner_preserves_reasoning_fields_in_assistant_history():
|
async def test_runner_preserves_reasoning_fields_in_assistant_history():
|
||||||
"""Reasoning fields ride along on the persisted assistant message so
|
"""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"]
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
|
async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
|
||||||
from nanobot.agent.runner import AgentRunner
|
from nanobot.agent.runner import AgentRunner
|
||||||
|
|||||||
Reference in New Issue
Block a user