mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 07:09:11 +02:00
Stop gather() from creating coroutines it may never schedule (#4559)
* Fix: gather() eagerly creates coroutines before scheduling them AggregateProvider fans out Provider.get_tool() (and sibling calls) across child providers via gather(*[p.get_tool(x) for p in providers]). The list comprehension builds every coroutine up front, then gather()'s scheduling loop hands them to an anyio task group one at a time. If that loop is interrupted partway through - e.g. by pytest-timeout's SIGALRM-based per-test timeout, which can fire between any two bytecode instructions, unlike normal async cancellation - any coroutine not yet scheduled is abandoned and silently garbage collected later, producing a "coroutine 'Provider.get_tool' was never awaited" warning attributed to whatever unrelated test happens to be running when the GC gets to it. Change gather() to take a single iterable consumed lazily, one awaitable at a time, right before each is scheduled, and close any awaitable that was just retrieved if scheduling it raises. Update call sites to pass generator expressions instead of eagerly-built lists so coroutine creation and scheduling stay tightly coupled. * Close unscheduled awaitables from eager callers; make get_tasks lazy
This commit is contained in:
Jeremiah Lowin
GitHub
parent
c33a3c3b29
commit
c8b8911226
5 changed files with 192 additions and 29 deletions
|
|
@ -190,7 +190,7 @@ class AggregateProvider(Provider):
|
|||
async def _list_tools(self) -> Sequence[Tool]:
|
||||
"""List all tools from all providers."""
|
||||
results = await gather(
|
||||
*[p.list_tools() for p in self.providers],
|
||||
(p.list_tools() for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._collect_list_results(results, "list_tools")
|
||||
|
|
@ -200,7 +200,7 @@ class AggregateProvider(Provider):
|
|||
) -> Tool | None:
|
||||
"""Get tool by name from providers."""
|
||||
results = await gather(
|
||||
*[p.get_tool(name, version) for p in self.providers],
|
||||
(p.get_tool(name, version) for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._get_highest_version_result(results, f"get_tool({name!r})") # type: ignore[return-value] # ty:ignore[invalid-argument-type, invalid-return-type]
|
||||
|
|
@ -208,7 +208,7 @@ class AggregateProvider(Provider):
|
|||
async def get_app_tool(self, app_name: str, tool_name: str) -> Tool | None:
|
||||
"""Query all child providers for an app tool."""
|
||||
results = await gather(
|
||||
*[p.get_app_tool(app_name, tool_name) for p in self.providers],
|
||||
(p.get_app_tool(app_name, tool_name) for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
for r in results:
|
||||
|
|
@ -223,7 +223,7 @@ class AggregateProvider(Provider):
|
|||
async def get_tool_by_hash(self, tool_hash: str, tool_name: str) -> Tool | None:
|
||||
"""Query all child providers for a tool matching a hash."""
|
||||
results = await gather(
|
||||
*[p.get_tool_by_hash(tool_hash, tool_name) for p in self.providers],
|
||||
(p.get_tool_by_hash(tool_hash, tool_name) for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
for r in results:
|
||||
|
|
@ -242,7 +242,7 @@ class AggregateProvider(Provider):
|
|||
async def _list_resources(self) -> Sequence[Resource]:
|
||||
"""List all resources from all providers."""
|
||||
results = await gather(
|
||||
*[p.list_resources() for p in self.providers],
|
||||
(p.list_resources() for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._collect_list_results(results, "list_resources")
|
||||
|
|
@ -252,7 +252,7 @@ class AggregateProvider(Provider):
|
|||
) -> Resource | None:
|
||||
"""Get resource by URI from providers."""
|
||||
results = await gather(
|
||||
*[p.get_resource(uri, version) for p in self.providers],
|
||||
(p.get_resource(uri, version) for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._get_highest_version_result(results, f"get_resource({uri!r})") # type: ignore[return-value] # ty:ignore[invalid-argument-type, invalid-return-type]
|
||||
|
|
@ -264,7 +264,7 @@ class AggregateProvider(Provider):
|
|||
async def _list_resource_templates(self) -> Sequence[ResourceTemplate]:
|
||||
"""List all resource templates from all providers."""
|
||||
results = await gather(
|
||||
*[p.list_resource_templates() for p in self.providers],
|
||||
(p.list_resource_templates() for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._collect_list_results(results, "list_resource_templates")
|
||||
|
|
@ -274,7 +274,7 @@ class AggregateProvider(Provider):
|
|||
) -> ResourceTemplate | None:
|
||||
"""Get resource template by URI from providers."""
|
||||
results = await gather(
|
||||
*[p.get_resource_template(uri, version) for p in self.providers],
|
||||
(p.get_resource_template(uri, version) for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._get_highest_version_result(
|
||||
|
|
@ -288,7 +288,7 @@ class AggregateProvider(Provider):
|
|||
async def _list_prompts(self) -> Sequence[Prompt]:
|
||||
"""List all prompts from all providers."""
|
||||
results = await gather(
|
||||
*[p.list_prompts() for p in self.providers],
|
||||
(p.list_prompts() for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._collect_list_results(results, "list_prompts")
|
||||
|
|
@ -298,7 +298,7 @@ class AggregateProvider(Provider):
|
|||
) -> Prompt | None:
|
||||
"""Get prompt by name from providers."""
|
||||
results = await gather(
|
||||
*[p.get_prompt(name, version) for p in self.providers],
|
||||
(p.get_prompt(name, version) for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._get_highest_version_result(results, f"get_prompt({name!r})") # type: ignore[return-value] # ty:ignore[invalid-argument-type, invalid-return-type]
|
||||
|
|
@ -310,7 +310,7 @@ class AggregateProvider(Provider):
|
|||
async def get_tasks(self) -> Sequence[FastMCPComponent]:
|
||||
"""Get all task-eligible components from all providers."""
|
||||
results = await gather(
|
||||
*[p.get_tasks() for p in self.providers],
|
||||
(p.get_tasks() for p in self.providers),
|
||||
return_exceptions=True,
|
||||
)
|
||||
return self._collect_list_results(results, "get_tasks")
|
||||
|
|
|
|||
|
|
@ -496,12 +496,19 @@ class Provider:
|
|||
|
||||
Used by the server during startup to register functions with Docket.
|
||||
"""
|
||||
# Fetch all component types in parallel
|
||||
# Fetch all component types in parallel. Iterate the bound methods
|
||||
# rather than a tuple of already-called coroutines: a parenthesized
|
||||
# comma expression is a tuple, so it would create all four coroutines
|
||||
# before `gather` starts, which is exactly what `gather` asks callers
|
||||
# to avoid.
|
||||
results = await gather(
|
||||
self._list_tools(),
|
||||
self._list_resources(),
|
||||
self._list_resource_templates(),
|
||||
self._list_prompts(),
|
||||
fetch()
|
||||
for fetch in (
|
||||
self._list_tools,
|
||||
self._list_resources,
|
||||
self._list_resource_templates,
|
||||
self._list_prompts,
|
||||
)
|
||||
)
|
||||
tools = cast("Sequence[Tool]", results[0])
|
||||
resources = cast("Sequence[Resource]", results[1])
|
||||
|
|
|
|||
|
|
@ -390,7 +390,7 @@ async def execute_tools(
|
|||
# Execute in parallel
|
||||
if tool_concurrency == 0:
|
||||
# Unlimited parallel execution
|
||||
return await gather(*[_execute_single_tool(tc) for tc in tool_calls])
|
||||
return await gather(_execute_single_tool(tc) for tc in tool_calls)
|
||||
else:
|
||||
# Bounded parallel execution with semaphore
|
||||
semaphore = anyio.Semaphore(tool_concurrency)
|
||||
|
|
@ -399,7 +399,7 @@ async def execute_tools(
|
|||
async with semaphore:
|
||||
return await _execute_single_tool(tool_use)
|
||||
|
||||
return await gather(*[bounded_execute(tc) for tc in tool_calls])
|
||||
return await gather(bounded_execute(tc) for tc in tool_calls)
|
||||
|
||||
|
||||
# --- Helper functions for sampling ---
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from typing import Any, Literal, TypeVar, overload
|
||||
|
||||
import anyio
|
||||
|
|
@ -36,35 +36,55 @@ async def call_sync_fn_in_threadpool(
|
|||
|
||||
@overload
|
||||
async def gather(
|
||||
*awaitables: Awaitable[T],
|
||||
awaitables: Iterable[Awaitable[T]],
|
||||
*,
|
||||
return_exceptions: Literal[True],
|
||||
) -> list[T | BaseException]: ...
|
||||
|
||||
|
||||
@overload
|
||||
async def gather(
|
||||
*awaitables: Awaitable[T],
|
||||
awaitables: Iterable[Awaitable[T]],
|
||||
*,
|
||||
return_exceptions: Literal[False] = ...,
|
||||
) -> list[T]: ...
|
||||
|
||||
|
||||
async def gather(
|
||||
*awaitables: Awaitable[T],
|
||||
awaitables: Iterable[Awaitable[T]],
|
||||
*,
|
||||
return_exceptions: bool = False,
|
||||
) -> list[T] | list[T | BaseException]:
|
||||
"""Run awaitables concurrently and return results in order.
|
||||
|
||||
Uses anyio TaskGroup for structured concurrency.
|
||||
|
||||
``awaitables`` is consumed lazily, one item at a time, right before each
|
||||
is handed to the task group. Callers with a dynamic number of awaitables
|
||||
should pass a generator expression (e.g. ``gather(f(x) for x in xs)``)
|
||||
rather than a list or list comprehension: a list comprehension calls
|
||||
every ``f(x)`` up front, creating a batch of coroutine objects before
|
||||
this function even starts, whereas a generator expression creates each
|
||||
coroutine only as this function's own scheduling loop asks for it. That
|
||||
matters because coroutine creation and scheduling can be interrupted
|
||||
between any two bytecode instructions by a synchronous signal handler
|
||||
(for example pytest-timeout's SIGALRM-based per-test timeout). If that
|
||||
happens while a whole batch of coroutines is sitting unscheduled, they
|
||||
are silently abandoned and eventually trigger a "coroutine was never
|
||||
awaited" warning attributed to whatever unrelated code happens to be
|
||||
running when the garbage collector gets to them. Lazy consumption keeps
|
||||
the window in which a created-but-unscheduled coroutine can exist as
|
||||
small as possible.
|
||||
|
||||
Args:
|
||||
*awaitables: Awaitables to run concurrently
|
||||
awaitables: Iterable of awaitables to run concurrently.
|
||||
return_exceptions: If True, exceptions are returned in results.
|
||||
If False, first exception cancels all and raises.
|
||||
|
||||
Returns:
|
||||
List of results in the same order as input awaitables.
|
||||
"""
|
||||
results: list[T | BaseException] = [None] * len(awaitables) # type: ignore[assignment] # ty:ignore[invalid-assignment]
|
||||
results: list[T | BaseException] = []
|
||||
|
||||
async def run_at(i: int, aw: Awaitable[T]) -> None:
|
||||
try:
|
||||
|
|
@ -75,8 +95,26 @@ async def gather(
|
|||
else:
|
||||
raise
|
||||
|
||||
pending = enumerate(awaitables)
|
||||
async with anyio.create_task_group() as tg:
|
||||
for i, aw in enumerate(awaitables):
|
||||
tg.start_soon(run_at, i, aw)
|
||||
for i, aw in pending:
|
||||
results.append(None) # type: ignore[arg-type] # ty:ignore[invalid-argument-type]
|
||||
try:
|
||||
tg.start_soon(run_at, i, aw)
|
||||
except BaseException:
|
||||
# `aw` was just created (possibly moments ago, by the
|
||||
# generator's own iteration) but never handed off - close it
|
||||
# explicitly so it isn't silently garbage collected later.
|
||||
if inspect.iscoroutine(aw):
|
||||
aw.close()
|
||||
# Lazy consumption keeps the leak window small, but a caller
|
||||
# that passed an already-built sequence has coroutines sitting
|
||||
# behind this one that were never scheduled either. Draining
|
||||
# the iterator closes them too, so `gather` cannot leak
|
||||
# regardless of how eagerly its argument was constructed.
|
||||
for _, remaining in pending:
|
||||
if inspect.iscoroutine(remaining):
|
||||
remaining.close()
|
||||
raise
|
||||
|
||||
return results
|
||||
|
|
|
|||
|
|
@ -1,7 +1,11 @@
|
|||
"""Tests for fastmcp.utilities.async_utils."""
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Awaitable, Iterator
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from exceptiongroup import BaseExceptionGroup
|
||||
|
||||
|
|
@ -9,6 +13,7 @@ from fastmcp import Client, FastMCP
|
|||
from fastmcp.prompts import prompt
|
||||
from fastmcp.resources import resource
|
||||
from fastmcp.tools import tool
|
||||
from fastmcp.utilities import async_utils
|
||||
from fastmcp.utilities.async_utils import gather, is_coroutine_function
|
||||
|
||||
|
||||
|
|
@ -55,14 +60,20 @@ class TestGather:
|
|||
async def value(result: int) -> int:
|
||||
return result
|
||||
|
||||
assert await gather(value(1), value(2), value(3)) == [1, 2, 3]
|
||||
assert await gather([value(1), value(2), value(3)]) == [1, 2, 3]
|
||||
|
||||
async def test_accepts_a_generator(self) -> None:
|
||||
async def value(result: int) -> int:
|
||||
return result
|
||||
|
||||
assert await gather(value(i) for i in [1, 2, 3]) == [1, 2, 3]
|
||||
|
||||
async def test_raises_by_default(self) -> None:
|
||||
async def fail() -> int:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with pytest.raises(BaseExceptionGroup) as exc_info:
|
||||
await gather(fail())
|
||||
await gather([fail()])
|
||||
|
||||
assert len(exc_info.value.exceptions) == 1
|
||||
assert isinstance(exc_info.value.exceptions[0], RuntimeError)
|
||||
|
|
@ -74,11 +85,118 @@ class TestGather:
|
|||
async def value() -> int:
|
||||
return 1
|
||||
|
||||
result = await gather(fail(), value(), return_exceptions=True)
|
||||
result = await gather([fail(), value()], return_exceptions=True)
|
||||
|
||||
assert isinstance(result[0], ValueError)
|
||||
assert result[1] == 1
|
||||
|
||||
async def test_does_not_leak_coroutine_when_scheduling_is_interrupted(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""If handing an already-created awaitable off to the task group
|
||||
raises partway through scheduling, that awaitable must be closed
|
||||
rather than silently garbage collected later - which is what
|
||||
produces a "coroutine was never awaited" RuntimeWarning attributed
|
||||
to whatever unrelated code happens to be running when the garbage
|
||||
collector eventually reclaims it.
|
||||
|
||||
In production this can happen when a synchronous signal handler
|
||||
(e.g. pytest-timeout's SIGALRM-based per-test timeout) fires inside
|
||||
anyio's task-spawning internals. This test reproduces the same
|
||||
shape of interruption deterministically by making the task group's
|
||||
``start_soon`` raise partway through scheduling, instead of relying
|
||||
on real signal timing.
|
||||
"""
|
||||
real_create_task_group = anyio.create_task_group
|
||||
|
||||
class _FailOnSecondStart:
|
||||
def __init__(self) -> None:
|
||||
self._real_tg = real_create_task_group()
|
||||
self._calls = 0
|
||||
|
||||
async def __aenter__(self) -> "_FailOnSecondStart":
|
||||
await self._real_tg.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: Any) -> bool | None:
|
||||
return await self._real_tg.__aexit__(*exc_info)
|
||||
|
||||
def start_soon(self, func: Any, *args: Any) -> None:
|
||||
self._calls += 1
|
||||
if self._calls == 2:
|
||||
raise RuntimeError("interrupted while scheduling")
|
||||
self._real_tg.start_soon(func, *args)
|
||||
|
||||
monkeypatch.setattr(async_utils.anyio, "create_task_group", _FailOnSecondStart)
|
||||
|
||||
created: list[Any] = []
|
||||
|
||||
async def value(result: int) -> int:
|
||||
return result
|
||||
|
||||
def awaitables() -> Iterator[Awaitable[int]]:
|
||||
for i in range(3):
|
||||
aw = value(i)
|
||||
created.append(aw)
|
||||
yield aw
|
||||
|
||||
with pytest.raises(BaseExceptionGroup) as exc_info:
|
||||
await gather(awaitables())
|
||||
|
||||
assert len(exc_info.value.exceptions) == 1
|
||||
assert isinstance(exc_info.value.exceptions[0], RuntimeError)
|
||||
assert "interrupted while scheduling" in str(exc_info.value.exceptions[0])
|
||||
|
||||
# created[1] was being handed to start_soon() when it raised - it
|
||||
# must have been closed rather than abandoned.
|
||||
assert inspect.getcoroutinestate(created[1]) == "CORO_CLOSED"
|
||||
|
||||
async def test_closes_unscheduled_coroutines_from_an_eager_caller(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Lazy consumption keeps the leak window small, but a caller that
|
||||
builds its awaitables eagerly (a list or a parenthesized tuple) has
|
||||
coroutines queued behind the failing one that were never scheduled
|
||||
either. ``gather`` drains what is left of the iterable and closes
|
||||
those too, so it cannot leak regardless of how its argument was
|
||||
constructed."""
|
||||
real_create_task_group = anyio.create_task_group
|
||||
|
||||
class _FailOnSecondStart:
|
||||
def __init__(self) -> None:
|
||||
self._real_tg = real_create_task_group()
|
||||
self._calls = 0
|
||||
|
||||
async def __aenter__(self) -> "_FailOnSecondStart":
|
||||
await self._real_tg.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: Any) -> bool | None:
|
||||
return await self._real_tg.__aexit__(*exc_info)
|
||||
|
||||
def start_soon(self, func: Any, *args: Any) -> None:
|
||||
self._calls += 1
|
||||
if self._calls == 2:
|
||||
raise RuntimeError("interrupted while scheduling")
|
||||
self._real_tg.start_soon(func, *args)
|
||||
|
||||
monkeypatch.setattr(async_utils.anyio, "create_task_group", _FailOnSecondStart)
|
||||
|
||||
async def value(result: int) -> int:
|
||||
return result
|
||||
|
||||
# Eagerly built: all four coroutines exist before gather() runs.
|
||||
eager = [value(0), value(1), value(2), value(3)]
|
||||
|
||||
with pytest.raises(BaseExceptionGroup):
|
||||
await gather(eager)
|
||||
|
||||
# The one that failed to schedule *and* the two queued behind it are
|
||||
# all closed; none is left to surface as a stray warning later.
|
||||
assert [inspect.getcoroutinestate(aw) for aw in eager[1:]] == [
|
||||
"CORO_CLOSED"
|
||||
] * 3
|
||||
|
||||
|
||||
class TestAsyncPartialIntegration:
|
||||
async def test_async_partial_tool_runs(self) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue