mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-10-16 15:06:04 +08:00
fix: keep async tool generators in one task (#10474)
* fix: add demand-driven task ownership for async generators * fix: keep async tool generators in one task * fix: publish generator cleanup errors without blocking * style: format tool lifecycle regression tests
This commit is contained in:
エイカク
@@ -46,6 +46,7 @@ from astrbot.core.provider.modalities import (
|
||||
sanitize_contexts_by_modalities,
|
||||
)
|
||||
from astrbot.core.provider.provider import Provider
|
||||
from astrbot.core.utils.async_generator import iterate_in_task
|
||||
|
||||
from ..context.compressor import ContextCompressor
|
||||
from ..context.config import ContextConfig
|
||||
@@ -1611,49 +1612,49 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
|
||||
self,
|
||||
executor: T.AsyncGenerator[ToolExecutorResultT, None],
|
||||
) -> T.AsyncGenerator[ToolExecutorResultT, None]:
|
||||
executor = iterate_in_task(executor)
|
||||
|
||||
async def _next_executor_result() -> ToolExecutorResultT:
|
||||
return await anext(executor)
|
||||
|
||||
while True:
|
||||
if self._is_stop_requested():
|
||||
await self._close_executor(executor)
|
||||
raise _ToolExecutionInterrupted(
|
||||
"Tool execution interrupted before reading the next tool result."
|
||||
)
|
||||
|
||||
next_result_task = asyncio.create_task(_next_executor_result())
|
||||
abort_task = asyncio.create_task(self._abort_signal.wait())
|
||||
try:
|
||||
done, _ = await asyncio.wait(
|
||||
{next_result_task, abort_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
if abort_task in done:
|
||||
if not next_result_task.done():
|
||||
next_result_task.cancel()
|
||||
with suppress(asyncio.CancelledError, StopAsyncIteration):
|
||||
await next_result_task
|
||||
|
||||
await self._close_executor(executor)
|
||||
|
||||
try:
|
||||
while True:
|
||||
if self._is_stop_requested():
|
||||
raise _ToolExecutionInterrupted(
|
||||
"Tool execution interrupted by a stop request."
|
||||
"Tool execution interrupted before reading the next tool result."
|
||||
)
|
||||
|
||||
next_result_task = asyncio.create_task(_next_executor_result())
|
||||
abort_task = asyncio.create_task(self._abort_signal.wait())
|
||||
try:
|
||||
yield next_result_task.result()
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
if not next_result_task.done():
|
||||
next_result_task.cancel()
|
||||
# The reader must finish before its async generator can be closed.
|
||||
await asyncio.gather(next_result_task, return_exceptions=True)
|
||||
await self._close_executor(executor)
|
||||
raise
|
||||
finally:
|
||||
if not abort_task.done():
|
||||
abort_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await abort_task
|
||||
done, _ = await asyncio.wait(
|
||||
{next_result_task, abort_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
if abort_task in done:
|
||||
if not next_result_task.done():
|
||||
next_result_task.cancel()
|
||||
with suppress(asyncio.CancelledError, StopAsyncIteration):
|
||||
await next_result_task
|
||||
raise _ToolExecutionInterrupted(
|
||||
"Tool execution interrupted by a stop request."
|
||||
)
|
||||
|
||||
try:
|
||||
yield next_result_task.result()
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
if not next_result_task.done():
|
||||
next_result_task.cancel()
|
||||
# The reader must finish before its async generator can be closed.
|
||||
await asyncio.gather(next_result_task, return_exceptions=True)
|
||||
raise
|
||||
finally:
|
||||
if not abort_task.done():
|
||||
abort_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await abort_task
|
||||
finally:
|
||||
await executor.aclose()
|
||||
|
||||
@@ -6,6 +6,7 @@ import typing as T
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from contextlib import aclosing
|
||||
|
||||
import mcp
|
||||
|
||||
@@ -50,6 +51,7 @@ from astrbot.core.tools.computer_tools import (
|
||||
from astrbot.core.tools.message_tools import SendMessageToUserTool
|
||||
from astrbot.core.utils.active_event_registry import active_event_registry
|
||||
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
|
||||
from astrbot.core.utils.async_generator import iterate_in_task
|
||||
from astrbot.core.utils.config_number import coerce_int_config
|
||||
from astrbot.core.utils.history_saver import persist_agent_history
|
||||
from astrbot.core.utils.image_ref_utils import is_supported_image_ref
|
||||
@@ -145,18 +147,25 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
|
||||
if isinstance(tool, HandoffTool):
|
||||
is_bg = tool_args.pop("background_task", False)
|
||||
if is_bg:
|
||||
async for r in cls._execute_handoff_background(
|
||||
tool, run_context, **tool_args
|
||||
):
|
||||
yield r
|
||||
async with aclosing(
|
||||
cls._execute_handoff_background(tool, run_context, **tool_args)
|
||||
) as results:
|
||||
async for r in results:
|
||||
yield r
|
||||
return
|
||||
async for r in cls._execute_handoff(tool, run_context, **tool_args):
|
||||
yield r
|
||||
async with aclosing(
|
||||
cls._execute_handoff(tool, run_context, **tool_args)
|
||||
) as results:
|
||||
async for r in results:
|
||||
yield r
|
||||
return
|
||||
|
||||
elif isinstance(tool, MCPTool):
|
||||
async for r in cls._execute_mcp(tool, run_context, **tool_args):
|
||||
yield r
|
||||
async with aclosing(
|
||||
cls._execute_mcp(tool, run_context, **tool_args)
|
||||
) as results:
|
||||
async for r in results:
|
||||
yield r
|
||||
return
|
||||
|
||||
elif tool.is_background_task:
|
||||
@@ -189,8 +198,11 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
|
||||
|
||||
return
|
||||
else:
|
||||
async for r in cls._execute_local(tool, run_context, **tool_args):
|
||||
yield r
|
||||
async with aclosing(
|
||||
cls._execute_local(tool, run_context, **tool_args)
|
||||
) as results:
|
||||
async for r in results:
|
||||
yield r
|
||||
return
|
||||
|
||||
@classmethod
|
||||
@@ -724,46 +736,47 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]):
|
||||
method_name=method_name,
|
||||
**tool_args,
|
||||
)
|
||||
while True:
|
||||
try:
|
||||
resp = await asyncio.wait_for(
|
||||
anext(wrapper),
|
||||
timeout=effective_timeout,
|
||||
)
|
||||
if resp is not None:
|
||||
if isinstance(resp, mcp.types.CallToolResult):
|
||||
yield resp
|
||||
async with aclosing(wrapper):
|
||||
while True:
|
||||
try:
|
||||
resp = await asyncio.wait_for(
|
||||
anext(wrapper),
|
||||
timeout=effective_timeout,
|
||||
)
|
||||
if resp is not None:
|
||||
if isinstance(resp, mcp.types.CallToolResult):
|
||||
yield resp
|
||||
else:
|
||||
text_content = mcp.types.TextContent(
|
||||
type="text",
|
||||
text=str(resp),
|
||||
)
|
||||
yield mcp.types.CallToolResult(content=[text_content])
|
||||
else:
|
||||
text_content = mcp.types.TextContent(
|
||||
type="text",
|
||||
text=str(resp),
|
||||
)
|
||||
yield mcp.types.CallToolResult(content=[text_content])
|
||||
else:
|
||||
# NOTE: Tool 在这里直接请求发送消息给用户
|
||||
# TODO: 是否需要判断 event.get_result() 是否为空?
|
||||
# 如果为空,则说明没有发送消息给用户,并且返回值为空,将返回一个特殊的 TextContent,其内容如"工具没有返回内容"
|
||||
if res := run_context.context.event.get_result():
|
||||
if res.chain:
|
||||
try:
|
||||
await event.send(
|
||||
MessageChain(
|
||||
chain=res.chain,
|
||||
type="tool_direct_result",
|
||||
# NOTE: Tool 在这里直接请求发送消息给用户
|
||||
# TODO: 是否需要判断 event.get_result() 是否为空?
|
||||
# 如果为空,则说明没有发送消息给用户,并且返回值为空,将返回一个特殊的 TextContent,其内容如"工具没有返回内容"
|
||||
if res := run_context.context.event.get_result():
|
||||
if res.chain:
|
||||
try:
|
||||
await event.send(
|
||||
MessageChain(
|
||||
chain=res.chain,
|
||||
type="tool_direct_result",
|
||||
)
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Tool 直接发送消息失败: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
yield None
|
||||
except asyncio.TimeoutError:
|
||||
raise Exception(
|
||||
f"tool {tool.name} execution timeout after {effective_timeout} seconds.",
|
||||
)
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Tool 直接发送消息失败: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
yield None
|
||||
except asyncio.TimeoutError:
|
||||
raise Exception(
|
||||
f"tool {tool.name} execution timeout after {effective_timeout} seconds.",
|
||||
)
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
|
||||
@classmethod
|
||||
async def _execute_mcp(
|
||||
@@ -847,18 +860,19 @@ async def call_local_llm_tool(
|
||||
if inspect.isasyncgen(ready_to_call):
|
||||
_has_yielded = False
|
||||
try:
|
||||
async for ret in ready_to_call:
|
||||
# 这里逐步执行异步生成器, 对于每个 yield 返回的 ret, 执行下面的代码
|
||||
# 返回值只能是 MessageEventResult 或者 None（无返回值）
|
||||
_has_yielded = True
|
||||
if isinstance(ret, MessageEventResult | CommandResult):
|
||||
# 如果返回值是 MessageEventResult, 设置结果并继续
|
||||
event.set_result(ret)
|
||||
yield
|
||||
else:
|
||||
# 如果返回值是 None, 则不设置结果并继续
|
||||
# 继续执行后续阶段
|
||||
yield ret
|
||||
async with aclosing(iterate_in_task(ready_to_call)) as results:
|
||||
async for ret in results:
|
||||
# 这里逐步执行异步生成器, 对于每个 yield 返回的 ret, 执行下面的代码
|
||||
# 返回值只能是 MessageEventResult 或者 None（无返回值）
|
||||
_has_yielded = True
|
||||
if isinstance(ret, MessageEventResult | CommandResult):
|
||||
# 如果返回值是 MessageEventResult, 设置结果并继续
|
||||
event.set_result(ret)
|
||||
yield
|
||||
else:
|
||||
# 如果返回值是 None, 则不设置结果并继续
|
||||
# 继续执行后续阶段
|
||||
yield ret
|
||||
if not _has_yielded:
|
||||
# 如果这个异步生成器没有执行到 yield 分支
|
||||
yield
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import suppress
|
||||
from typing import TypeVar, cast
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
async def iterate_in_task(
|
||||
generator: AsyncGenerator[T, None],
|
||||
) -> AsyncGenerator[T, None]:
|
||||
"""Advance and close a generator in one task, only when a caller requests it.
|
||||
|
||||
Args:
|
||||
generator: Generator owned by this iterator until completion or close.
|
||||
|
||||
Yields:
|
||||
The original generator's values without prefetching.
|
||||
|
||||
Raises:
|
||||
BaseException: Original iteration or cleanup failures.
|
||||
"""
|
||||
advance = asyncio.Event()
|
||||
results: asyncio.Queue[tuple[bool, T | BaseException]] = asyncio.Queue(maxsize=1)
|
||||
finishing = False
|
||||
stop_requested = False
|
||||
terminal: BaseException | None = None
|
||||
|
||||
async def produce() -> None:
|
||||
nonlocal finishing, terminal
|
||||
try:
|
||||
while True:
|
||||
await advance.wait()
|
||||
advance.clear()
|
||||
try:
|
||||
value = await anext(generator)
|
||||
except asyncio.CancelledError as error:
|
||||
if not stop_requested:
|
||||
terminal = error
|
||||
break
|
||||
except BaseException as error:
|
||||
terminal = error
|
||||
break
|
||||
if stop_requested:
|
||||
return
|
||||
await results.put((True, value))
|
||||
finally:
|
||||
finishing = True
|
||||
try:
|
||||
await generator.aclose()
|
||||
except BaseException as error:
|
||||
terminal = error
|
||||
if terminal is not None:
|
||||
# Closing consumers may leave a value queued; terminal is also read below.
|
||||
with suppress(asyncio.QueueFull):
|
||||
results.put_nowait((False, terminal))
|
||||
|
||||
owner = asyncio.create_task(produce())
|
||||
try:
|
||||
while True:
|
||||
advance.set()
|
||||
succeeded, value = await results.get()
|
||||
if not succeeded:
|
||||
if isinstance(value, StopAsyncIteration):
|
||||
return
|
||||
raise cast(BaseException, value)
|
||||
yield cast(T, value)
|
||||
finally:
|
||||
interrupted = False
|
||||
if not finishing:
|
||||
stop_requested = True
|
||||
owner.cancel()
|
||||
while not owner.done():
|
||||
try:
|
||||
await asyncio.shield(owner)
|
||||
except asyncio.CancelledError:
|
||||
if not owner.done():
|
||||
interrupted = True
|
||||
with suppress(asyncio.CancelledError):
|
||||
await owner
|
||||
if (
|
||||
not interrupted
|
||||
and terminal is not None
|
||||
and not isinstance(terminal, StopAsyncIteration)
|
||||
):
|
||||
raise cast(BaseException, terminal)
|
||||
if not interrupted and not results.empty():
|
||||
succeeded, value = results.get_nowait()
|
||||
if not succeeded and not isinstance(value, StopAsyncIteration):
|
||||
raise cast(BaseException, value)
|
||||
if interrupted:
|
||||
raise asyncio.CancelledError
|
||||
@@ -0,0 +1,327 @@
|
||||
"""Tests for demand-driven async generator task ownership."""
|
||||
|
||||
import asyncio
|
||||
from contextvars import ContextVar
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from astrbot.core.utils.async_generator import iterate_in_task
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("exit_mode", ["close", "exhaust", "error"])
|
||||
async def test_owner_task_context_and_demand_are_preserved(exit_mode):
|
||||
tasks = set(asyncio.all_tasks())
|
||||
context = ContextVar("generator_context", default="caller")
|
||||
owners = []
|
||||
advanced = []
|
||||
error = ValueError("controlled generator failure")
|
||||
|
||||
async def source():
|
||||
owners.append(asyncio.current_task())
|
||||
token = context.set("generator")
|
||||
try:
|
||||
with anyio.CancelScope():
|
||||
yield None
|
||||
advanced.append("next")
|
||||
owners.append(asyncio.current_task())
|
||||
if exit_mode == "error":
|
||||
raise error
|
||||
yield error
|
||||
finally:
|
||||
await asyncio.sleep(0)
|
||||
owners.append(asyncio.current_task())
|
||||
context.reset(token)
|
||||
|
||||
results = iterate_in_task(source())
|
||||
try:
|
||||
assert await anext(results) is None
|
||||
await asyncio.sleep(0)
|
||||
assert advanced == []
|
||||
assert context.get() == "caller"
|
||||
if exit_mode == "exhaust":
|
||||
assert await anext(results) is error
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(results)
|
||||
elif exit_mode == "error":
|
||||
with pytest.raises(ValueError) as caught:
|
||||
await anext(results)
|
||||
assert caught.value is error
|
||||
finally:
|
||||
await results.aclose()
|
||||
assert all(owner is owners[0] for owner in owners)
|
||||
assert len(owners) == (2 if exit_mode == "close" else 3)
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_pending_read_waits_for_cleanup():
|
||||
tasks = set(asyncio.all_tasks())
|
||||
started = asyncio.Event()
|
||||
cleanup = []
|
||||
|
||||
async def source():
|
||||
try:
|
||||
started.set()
|
||||
await asyncio.Future()
|
||||
yield "unreachable"
|
||||
finally:
|
||||
await asyncio.sleep(0)
|
||||
cleanup.append("closed")
|
||||
|
||||
results = iterate_in_task(source())
|
||||
pending = asyncio.create_task(anext(results))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await pending
|
||||
assert cleanup == ["closed"]
|
||||
finally:
|
||||
pending.cancel()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_failure_is_not_hidden():
|
||||
error = OSError("controlled cleanup failure")
|
||||
|
||||
async def source():
|
||||
try:
|
||||
yield "result"
|
||||
finally:
|
||||
raise error
|
||||
|
||||
results = iterate_in_task(source())
|
||||
assert await anext(results) == "result"
|
||||
with pytest.raises(OSError) as caught:
|
||||
await results.aclose()
|
||||
assert caught.value is error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exhaustion_waits_without_cancelling_cleanup():
|
||||
cleanup_started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
closed = []
|
||||
|
||||
async def source():
|
||||
try:
|
||||
yield "result"
|
||||
finally:
|
||||
cleanup_started.set()
|
||||
await release.wait()
|
||||
closed.append("closed")
|
||||
|
||||
results = iterate_in_task(source())
|
||||
assert await anext(results) == "result"
|
||||
pending = asyncio.create_task(anext(results))
|
||||
try:
|
||||
await asyncio.wait_for(cleanup_started.wait(), 1)
|
||||
assert not pending.done()
|
||||
release.set()
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
assert closed == ["closed"]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_waits_for_async_cleanup_in_owner_task():
|
||||
owners = []
|
||||
cleanup_started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def source():
|
||||
owners.append(asyncio.current_task())
|
||||
try:
|
||||
yield "result"
|
||||
finally:
|
||||
cleanup_started.set()
|
||||
await release.wait()
|
||||
owners.append(asyncio.current_task())
|
||||
|
||||
results = iterate_in_task(source())
|
||||
assert await anext(results) == "result"
|
||||
closing = asyncio.create_task(results.aclose())
|
||||
try:
|
||||
await asyncio.wait_for(cleanup_started.wait(), 1)
|
||||
assert not closing.done()
|
||||
release.set()
|
||||
await asyncio.wait_for(closing, 1)
|
||||
assert len(owners) == 2 and owners[0] is owners[1]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(closing, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_generator_finishes_without_pending_task():
|
||||
tasks = set(asyncio.all_tasks())
|
||||
|
||||
async def source():
|
||||
if False:
|
||||
yield
|
||||
|
||||
assert [value async for value in iterate_in_task(source())] == []
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_cancel_does_not_interrupt_owner_cleanup():
|
||||
tasks = set(asyncio.all_tasks())
|
||||
started = asyncio.Event()
|
||||
cleanup_started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
closed = []
|
||||
|
||||
async def source():
|
||||
try:
|
||||
started.set()
|
||||
await asyncio.Future()
|
||||
yield "unreachable"
|
||||
finally:
|
||||
cleanup_started.set()
|
||||
await release.wait()
|
||||
closed.append("closed")
|
||||
|
||||
results = iterate_in_task(source())
|
||||
pending = asyncio.create_task(anext(results))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
pending.cancel()
|
||||
await asyncio.wait_for(cleanup_started.wait(), 1)
|
||||
pending.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not pending.done()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
assert closed == ["closed"]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generator_that_yields_after_cancel_is_still_closed():
|
||||
tasks = set(asyncio.all_tasks())
|
||||
started = asyncio.Event()
|
||||
closed = []
|
||||
|
||||
async def source():
|
||||
try:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
except asyncio.CancelledError:
|
||||
yield "late result"
|
||||
finally:
|
||||
closed.append("closed")
|
||||
|
||||
results = iterate_in_task(source())
|
||||
pending = asyncio.create_task(anext(results))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(pending, 0.1)
|
||||
assert closed == ["closed"]
|
||||
finally:
|
||||
pending.cancel()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_source_cancellation_reaches_consumer():
|
||||
error = asyncio.CancelledError("source cancelled itself")
|
||||
closed = []
|
||||
|
||||
async def source():
|
||||
try:
|
||||
raise error
|
||||
yield
|
||||
finally:
|
||||
closed.append("closed")
|
||||
|
||||
results = iterate_in_task(source())
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError) as caught:
|
||||
await asyncio.wait_for(anext(results), 0.1)
|
||||
assert caught.value is error
|
||||
assert closed == ["closed"]
|
||||
finally:
|
||||
await results.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_source_error_is_delivered_after_cleanup():
|
||||
error = ValueError("source failure")
|
||||
cleanup = []
|
||||
|
||||
async def source():
|
||||
try:
|
||||
raise error
|
||||
yield
|
||||
finally:
|
||||
await asyncio.sleep(0)
|
||||
cleanup.append("closed")
|
||||
|
||||
results = iterate_in_task(source())
|
||||
try:
|
||||
with pytest.raises(ValueError) as caught:
|
||||
await anext(results)
|
||||
assert caught.value is error
|
||||
assert cleanup == ["closed"]
|
||||
finally:
|
||||
await results.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_with_queued_value_and_cleanup_error(monkeypatch):
|
||||
original_queue = asyncio.Queue
|
||||
pending = None
|
||||
queued = asyncio.Event()
|
||||
error = OSError("queued cleanup failure")
|
||||
|
||||
class CancelAfterPut(original_queue):
|
||||
async def get(self):
|
||||
await asyncio.Future()
|
||||
|
||||
async def put(self, item):
|
||||
if item[0]:
|
||||
pending.cancel()
|
||||
queued.set()
|
||||
await super().put(item)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"astrbot.core.utils.async_generator.asyncio.Queue", CancelAfterPut
|
||||
)
|
||||
|
||||
async def source():
|
||||
try:
|
||||
yield "queued result"
|
||||
finally:
|
||||
raise error
|
||||
|
||||
results = iterate_in_task(source())
|
||||
pending = asyncio.create_task(anext(results))
|
||||
try:
|
||||
await asyncio.wait_for(queued.wait(), 1)
|
||||
with pytest.raises(OSError) as caught:
|
||||
await asyncio.wait_for(pending, 0.1)
|
||||
assert caught.value is error
|
||||
finally:
|
||||
pending.cancel()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
@@ -0,0 +1,309 @@
|
||||
"""Real local tool lifecycle across wrappers and runner stop/cancellation."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from astrbot.core.agent.run_context import ContextWrapper
|
||||
from astrbot.core.agent.runners.tool_loop_agent_runner import (
|
||||
ToolLoopAgentRunner,
|
||||
_ToolExecutionInterrupted,
|
||||
)
|
||||
from astrbot.core.agent.tool import FunctionTool
|
||||
from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor, call_local_llm_tool
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("layer", ["handler_wrapper", "local", "execute", "runner"])
|
||||
@pytest.mark.parametrize("exit_mode", ["close", "exhaust", "error"])
|
||||
async def test_real_handler_cleanup_and_output(layer, exit_mode):
|
||||
tasks = set(asyncio.all_tasks())
|
||||
cleanup = []
|
||||
advanced = []
|
||||
owners = []
|
||||
error = ValueError("controlled tool error")
|
||||
|
||||
async def implementation():
|
||||
token_task = asyncio.current_task()
|
||||
with anyio.CancelScope():
|
||||
try:
|
||||
owners.append(token_task)
|
||||
yield "first"
|
||||
advanced.append("next")
|
||||
owners.append(asyncio.current_task())
|
||||
if exit_mode == "error":
|
||||
raise error
|
||||
yield "second"
|
||||
finally:
|
||||
await asyncio.sleep(0)
|
||||
owners.append(asyncio.current_task())
|
||||
cleanup.append("closed")
|
||||
|
||||
handler_generator = implementation()
|
||||
|
||||
def handler(event):
|
||||
return handler_generator
|
||||
|
||||
run_context = ContextWrapper(context=SimpleNamespace(event=SimpleNamespace()))
|
||||
tool = FunctionTool(
|
||||
name="cleanup_probe",
|
||||
description="controlled tool lifecycle probe",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
handler=handler,
|
||||
)
|
||||
runner = ToolLoopAgentRunner()
|
||||
runner._abort_signal = asyncio.Event()
|
||||
executor = None
|
||||
if layer == "handler_wrapper":
|
||||
results = call_local_llm_tool(run_context, handler, "decorator_handler")
|
||||
elif layer == "local":
|
||||
results = FunctionToolExecutor._execute_local(tool, run_context)
|
||||
else:
|
||||
executor = FunctionToolExecutor.execute(tool, run_context)
|
||||
results = (
|
||||
runner._iter_tool_executor_results(executor)
|
||||
if layer == "runner"
|
||||
else executor
|
||||
)
|
||||
try:
|
||||
first = await anext(results)
|
||||
assert (
|
||||
first if layer == "handler_wrapper" else first.content[0].text
|
||||
) == "first"
|
||||
await asyncio.sleep(0)
|
||||
assert not advanced and not cleanup
|
||||
if exit_mode == "exhaust":
|
||||
second = await anext(results)
|
||||
assert (
|
||||
second if layer == "handler_wrapper" else second.content[0].text
|
||||
) == "second"
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(results)
|
||||
elif exit_mode == "error":
|
||||
with pytest.raises(ValueError) as caught:
|
||||
await anext(results)
|
||||
assert caught.value is error
|
||||
await results.aclose()
|
||||
assert cleanup == ["closed"]
|
||||
assert all(owner is owners[0] for owner in owners)
|
||||
finally:
|
||||
await results.aclose()
|
||||
if executor is not None:
|
||||
await executor.aclose()
|
||||
await handler_generator.aclose()
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["coroutine", "empty", "none"])
|
||||
async def test_local_result_compatibility(kind):
|
||||
sent = []
|
||||
|
||||
async def coroutine(event):
|
||||
return "coroutine result"
|
||||
|
||||
async def generator(event):
|
||||
if kind == "none":
|
||||
yield None
|
||||
|
||||
event = SimpleNamespace(get_result=lambda: None)
|
||||
run_context = ContextWrapper(context=SimpleNamespace(event=event))
|
||||
tool = FunctionTool(
|
||||
name="compatibility_probe",
|
||||
description="controlled compatibility probe",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
handler=coroutine if kind == "coroutine" else generator,
|
||||
)
|
||||
async for result in FunctionToolExecutor.execute(tool, run_context):
|
||||
sent.append(result)
|
||||
if kind == "coroutine":
|
||||
assert len(sent) == 1 and sent[0].content[0].text == "coroutine result"
|
||||
else:
|
||||
assert sent == [None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["stop_yielded", "stop_pending", "cancel", "timeout"])
|
||||
async def test_real_handler_stop_cancel_timeout_wait_for_cleanup(mode):
|
||||
tasks = set(asyncio.all_tasks())
|
||||
started = asyncio.Event()
|
||||
closed = []
|
||||
|
||||
async def handler(event):
|
||||
try:
|
||||
if mode == "stop_yielded":
|
||||
yield "first"
|
||||
started.set()
|
||||
await asyncio.Future()
|
||||
yield "unreachable"
|
||||
finally:
|
||||
await asyncio.sleep(0)
|
||||
closed.append("closed")
|
||||
|
||||
run_context = ContextWrapper(context=SimpleNamespace(event=SimpleNamespace()))
|
||||
tool = FunctionTool(
|
||||
name="blocked_probe",
|
||||
description="controlled blocking tool",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
handler=handler,
|
||||
)
|
||||
runner = ToolLoopAgentRunner()
|
||||
runner._abort_signal = asyncio.Event()
|
||||
executor = (
|
||||
FunctionToolExecutor._execute_local(tool, run_context, tool_call_timeout=0.01)
|
||||
if mode == "timeout"
|
||||
else FunctionToolExecutor.execute(tool, run_context)
|
||||
)
|
||||
results = (
|
||||
executor if mode == "timeout" else runner._iter_tool_executor_results(executor)
|
||||
)
|
||||
pending = None
|
||||
try:
|
||||
if mode == "stop_yielded":
|
||||
assert (await anext(results)).content[0].text == "first"
|
||||
runner.request_stop()
|
||||
with pytest.raises(_ToolExecutionInterrupted):
|
||||
await anext(results)
|
||||
else:
|
||||
pending = asyncio.create_task(anext(results))
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
if mode == "cancel":
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
elif mode == "timeout":
|
||||
with pytest.raises(Exception, match="execution timeout"):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
else:
|
||||
runner.request_stop()
|
||||
with pytest.raises(_ToolExecutionInterrupted):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
assert closed == ["closed"]
|
||||
finally:
|
||||
if pending is not None:
|
||||
pending.cancel()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
await executor.aclose()
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_repeated_cancel_waits_for_real_cleanup():
|
||||
tasks = set(asyncio.all_tasks())
|
||||
started = asyncio.Event()
|
||||
cleanup_started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
closed = []
|
||||
|
||||
async def handler(event):
|
||||
try:
|
||||
started.set()
|
||||
await asyncio.Future()
|
||||
yield "unreachable"
|
||||
finally:
|
||||
cleanup_started.set()
|
||||
await release.wait()
|
||||
closed.append("closed")
|
||||
|
||||
context = ContextWrapper(context=SimpleNamespace(event=SimpleNamespace()))
|
||||
tool = FunctionTool(
|
||||
name="repeat_cancel_probe",
|
||||
description="controlled repeated cancellation",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
handler=handler,
|
||||
)
|
||||
runner = ToolLoopAgentRunner()
|
||||
runner._abort_signal = asyncio.Event()
|
||||
results = runner._iter_tool_executor_results(
|
||||
FunctionToolExecutor.execute(tool, context)
|
||||
)
|
||||
pending = asyncio.create_task(anext(results))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
pending.cancel()
|
||||
await asyncio.wait_for(cleanup_started.wait(), 1)
|
||||
pending.cancel()
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
assert closed == ["closed"]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
await results.aclose()
|
||||
assert not (set(asyncio.all_tasks()) - tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_reports_real_handler_cleanup_error():
|
||||
error = OSError("controlled tool cleanup failure")
|
||||
|
||||
async def handler(event):
|
||||
try:
|
||||
yield "first"
|
||||
finally:
|
||||
raise error
|
||||
|
||||
context = ContextWrapper(context=SimpleNamespace(event=SimpleNamespace()))
|
||||
tool = FunctionTool(
|
||||
name="cleanup_error_probe",
|
||||
description="controlled cleanup failure",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
handler=handler,
|
||||
)
|
||||
runner = ToolLoopAgentRunner()
|
||||
runner._abort_signal = asyncio.Event()
|
||||
results = runner._iter_tool_executor_results(
|
||||
FunctionToolExecutor.execute(tool, context)
|
||||
)
|
||||
assert (await anext(results)).content[0].text == "first"
|
||||
with pytest.raises(OSError) as caught:
|
||||
await results.aclose()
|
||||
assert caught.value is error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_real_generator_cancel_closes_before_return(monkeypatch):
|
||||
started = asyncio.Event()
|
||||
closed = []
|
||||
wakeups = []
|
||||
|
||||
async def handler(event):
|
||||
try:
|
||||
yield "partial"
|
||||
started.set()
|
||||
await asyncio.Future()
|
||||
finally:
|
||||
await asyncio.sleep(0)
|
||||
closed.append("closed")
|
||||
|
||||
async def wake(**kwargs):
|
||||
wakeups.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(
|
||||
FunctionToolExecutor, "_wake_main_agent_for_background_result", wake
|
||||
)
|
||||
context = ContextWrapper(context=SimpleNamespace(event=SimpleNamespace()))
|
||||
tool = FunctionTool(
|
||||
name="background_cleanup_probe",
|
||||
description="controlled background cleanup",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
handler=handler,
|
||||
)
|
||||
pending = asyncio.create_task(
|
||||
FunctionToolExecutor._execute_background(tool, context, "task")
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 1)
|
||||
pending.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(pending, 1)
|
||||
assert closed == ["closed"]
|
||||
assert wakeups == []
|
||||
finally:
|
||||
pending.cancel()
|
||||
await asyncio.gather(pending, return_exceptions=True)
|
||||
Reference in New Issue
Block a user