mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-10-11 23:43:20 +02:00
2dd1c319c6Pydantic validation errors raised inside a tool body are now returned as tool error results instead of JSON-RPC -32602, so server-side failures no longer look like invalid client requests. Co-authored-by: qingshungLI <201055554+qingshungLI@users.noreply.github.com> Co-authored-by: Alastair Porter <19217+alastair@users.noreply.github.com>
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
"""Generator results must finish before context-manager dependencies exit."""
|
|
|
|
import asyncio
|
|
import io
|
|
from collections.abc import AsyncIterator, Iterator
|
|
from contextlib import asynccontextmanager
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from fastmcp import FastMCP
|
|
from fastmcp.client import Client
|
|
from fastmcp.dependencies import Depends
|
|
from fastmcp.exceptions import ValidationError
|
|
from fastmcp.tools.function_tool import FunctionTool, _ToolBodyError
|
|
|
|
|
|
@pytest.mark.parametrize("run_in_thread", [False, True])
|
|
@pytest.mark.parametrize(
|
|
"kind", ["async_generator", "sync_generator", "coroutine", "sync_wrapper"]
|
|
)
|
|
async def test_generator_result_keeps_dependency_open(kind: str, run_in_thread: bool):
|
|
stream = io.StringIO("payload")
|
|
events = []
|
|
|
|
@asynccontextmanager
|
|
async def open_stream():
|
|
events.append("enter")
|
|
try:
|
|
yield stream
|
|
finally:
|
|
events.append("exit")
|
|
stream.close()
|
|
|
|
async def produce(value) -> AsyncIterator[str]:
|
|
try:
|
|
events.append("body")
|
|
yield value.read()
|
|
finally:
|
|
events.append("generator-finally")
|
|
|
|
if kind == "async_generator":
|
|
|
|
async def read(value=Depends(open_stream)) -> AsyncIterator[str]:
|
|
async for item in produce(value):
|
|
yield item
|
|
elif kind == "sync_generator":
|
|
|
|
def read(value=Depends(open_stream)) -> Iterator[str]:
|
|
try:
|
|
events.append("body")
|
|
yield value.read()
|
|
finally:
|
|
events.append("generator-finally")
|
|
elif kind == "coroutine":
|
|
|
|
async def read(value=Depends(open_stream)) -> AsyncIterator[str]:
|
|
return produce(value)
|
|
else:
|
|
|
|
def read(value=Depends(open_stream)) -> AsyncIterator[str]:
|
|
return produce(value)
|
|
|
|
mcp = FastMCP("generator-dependency")
|
|
mcp.tool(read, run_in_thread=run_in_thread)
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("read", {})
|
|
assert result.content[0].text == '["payload"]'
|
|
assert events == ["enter", "body", "generator-finally", "exit"]
|
|
assert stream.closed
|
|
|
|
|
|
@pytest.mark.parametrize("failure", ["exception", "validation", "cancel"])
|
|
async def test_generator_failure_closes_dependency_after_generator(failure: str):
|
|
events = []
|
|
started = asyncio.Event()
|
|
|
|
@asynccontextmanager
|
|
async def dependency():
|
|
events.append("enter")
|
|
try:
|
|
yield None
|
|
finally:
|
|
events.append("exit")
|
|
|
|
class Model(BaseModel):
|
|
value: int
|
|
|
|
async def read(value=Depends(dependency)) -> AsyncIterator[str]:
|
|
try:
|
|
events.append("body")
|
|
yield "first"
|
|
if failure == "exception":
|
|
raise RuntimeError("iteration failed")
|
|
if failure == "validation":
|
|
Model(value="invalid")
|
|
started.set()
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
events.append("generator-finally")
|
|
|
|
tool = FunctionTool.from_function(read)
|
|
if failure == "cancel":
|
|
task = asyncio.create_task(tool.run({}))
|
|
await started.wait()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
else:
|
|
error = RuntimeError if failure == "exception" else _ToolBodyError
|
|
with pytest.raises(error):
|
|
await tool.run({})
|
|
assert events == ["enter", "body", "generator-finally", "exit"]
|
|
|
|
|
|
async def test_invalid_generator_arguments_do_not_enter_dependency():
|
|
events = []
|
|
|
|
@asynccontextmanager
|
|
async def dependency():
|
|
events.append("enter")
|
|
yield None
|
|
events.append("exit")
|
|
|
|
async def read(number: int, value=Depends(dependency)) -> AsyncIterator[str]:
|
|
yield str(number)
|
|
|
|
tool = FunctionTool.from_function(read)
|
|
with pytest.raises(ValidationError):
|
|
await tool.run({"number": "invalid"})
|
|
assert events == []
|