mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-10-29 07:11:26 +01:00
116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
"""Per-server MCP config timeouts reach the sessions that send requests."""
|
|
|
|
import asyncio
|
|
from collections.abc import AsyncGenerator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from mcp import ClientSession, MCPError
|
|
|
|
from fastmcp import Client, FastMCP
|
|
from fastmcp.client.transports import (
|
|
FastMCPTransport,
|
|
SSETransport,
|
|
StdioTransport,
|
|
StreamableHttpTransport,
|
|
)
|
|
from fastmcp.exceptions import ToolError
|
|
|
|
|
|
@pytest.fixture
|
|
def backend_sessions(monkeypatch: pytest.MonkeyPatch) -> list[ClientSession]:
|
|
"""Replace external connections while retaining config parsing and routing."""
|
|
server = FastMCP("Timeout backend")
|
|
|
|
@server.tool(tags={"public"})
|
|
async def slow(seconds: float = 1) -> str:
|
|
await asyncio.sleep(seconds)
|
|
return "completed"
|
|
|
|
sessions: list[ClientSession] = []
|
|
|
|
@asynccontextmanager
|
|
async def connect_session(
|
|
self: Any, **kwargs: Any
|
|
) -> AsyncGenerator[ClientSession, None]:
|
|
async with FastMCPTransport(server).connect_session(**kwargs) as session:
|
|
sessions.append(session)
|
|
yield session
|
|
|
|
for transport_type in (StreamableHttpTransport, SSETransport, StdioTransport):
|
|
monkeypatch.setattr(transport_type, "connect_session", connect_session)
|
|
return sessions
|
|
|
|
|
|
@pytest.fixture(params=["http", "sse", "stdio"])
|
|
def server_entry(request: pytest.FixtureRequest) -> dict[str, Any]:
|
|
if request.param == "stdio":
|
|
return {"command": "unused"}
|
|
return {"url": "http://example.invalid/mcp", "transport": request.param}
|
|
|
|
|
|
@pytest.mark.parametrize("server_count", [1, 2])
|
|
@pytest.mark.parametrize("transformed", [False, True])
|
|
async def test_entry_timeout_limits_requests(
|
|
backend_sessions: list[ClientSession],
|
|
server_entry: dict[str, Any],
|
|
server_count: int,
|
|
transformed: bool,
|
|
) -> None:
|
|
entry = {**server_entry, "timeout": 500}
|
|
if transformed:
|
|
entry["include_tags"] = ["public"]
|
|
servers = {"limited": entry}
|
|
if server_count == 2:
|
|
servers["unlimited"] = server_entry
|
|
|
|
async with Client({"mcpServers": servers}, mode="legacy") as client:
|
|
name = "slow" if server_count == 1 else "limited_slow"
|
|
with pytest.raises((MCPError, ToolError), match="timed out"):
|
|
await client.call_tool(name)
|
|
|
|
if server_count == 2:
|
|
result = await client.call_tool("unlimited_slow", {"seconds": 0})
|
|
assert result.data == "completed"
|
|
|
|
assert backend_sessions[0]._session_read_timeout_seconds == 0.5
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("entry_timeout", "client_timeout", "expected"),
|
|
[
|
|
(None, None, None),
|
|
(None, 3, 3.0),
|
|
(1250, None, 1.25),
|
|
(1250, 3, 1.25),
|
|
(0, 3, None),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("server_count", [1, 2])
|
|
async def test_entry_timeout_defaults_and_units(
|
|
backend_sessions: list[ClientSession],
|
|
server_entry: dict[str, Any],
|
|
entry_timeout: int | None,
|
|
client_timeout: int | None,
|
|
expected: float | None,
|
|
server_count: int,
|
|
) -> None:
|
|
servers = {"first": {**server_entry, "timeout": entry_timeout}}
|
|
if server_count == 2:
|
|
servers["second"] = server_entry
|
|
|
|
async with Client({"mcpServers": servers}, timeout=client_timeout):
|
|
assert backend_sessions[0]._session_read_timeout_seconds == expected
|
|
if server_count == 2:
|
|
assert backend_sessions[1]._session_read_timeout_seconds == client_timeout
|
|
|
|
|
|
async def test_single_server_per_call_timeout_overrides_entry(
|
|
backend_sessions: list[ClientSession], server_entry: dict[str, Any]
|
|
) -> None:
|
|
config = {"mcpServers": {"only": {**server_entry, "timeout": 500}}}
|
|
async with Client(config) as client:
|
|
result = await client.call_tool("slow", {"seconds": 0.6}, timeout=2)
|
|
assert result.data == "completed"
|
|
assert backend_sessions[0]._session_read_timeout_seconds == 0.5
|