diff --git a/fastmcp_slim/fastmcp/server/providers/aggregate.py b/fastmcp_slim/fastmcp/server/providers/aggregate.py index 766d64846..53928d055 100644 --- a/fastmcp_slim/fastmcp/server/providers/aggregate.py +++ b/fastmcp_slim/fastmcp/server/providers/aggregate.py @@ -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") diff --git a/fastmcp_slim/fastmcp/server/providers/base.py b/fastmcp_slim/fastmcp/server/providers/base.py index 402dff351..959dc9de9 100644 --- a/fastmcp_slim/fastmcp/server/providers/base.py +++ b/fastmcp_slim/fastmcp/server/providers/base.py @@ -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]) diff --git a/fastmcp_slim/fastmcp/server/sampling/run.py b/fastmcp_slim/fastmcp/server/sampling/run.py index bcb73cc66..4aedd60ec 100644 --- a/fastmcp_slim/fastmcp/server/sampling/run.py +++ b/fastmcp_slim/fastmcp/server/sampling/run.py @@ -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 --- diff --git a/fastmcp_slim/fastmcp/utilities/async_utils.py b/fastmcp_slim/fastmcp/utilities/async_utils.py index 8fc24f49d..978c84a6e 100644 --- a/fastmcp_slim/fastmcp/utilities/async_utils.py +++ b/fastmcp_slim/fastmcp/utilities/async_utils.py @@ -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 diff --git a/tests/utilities/test_async_utils.py b/tests/utilities/test_async_utils.py index 6d1366613..7f9d02d5d 100644 --- a/tests/utilities/test_async_utils.py +++ b/tests/utilities/test_async_utils.py @@ -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: