Spaces:
Running
Running
File size: 5,073 Bytes
b140934 d80180f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | import asyncio
import pytest
from anyio import create_task_group
from mcp.types import LoggingLevel
from fastmcp import Client, Context, FastMCP
from fastmcp.client.logging import LogMessage
from fastmcp.client.transports import FastMCPTransport
from fastmcp.exceptions import ToolError
from fastmcp.server.proxy import FastMCPProxy, StatefulProxyClient
from fastmcp.utilities.tests import find_available_port
@pytest.fixture
def fastmcp_server():
mcp = FastMCP("TestServer")
states: dict[int, int] = {}
@mcp.tool
async def log(
message: str, level: LoggingLevel, logger: str, context: Context
) -> None:
await context.log(message=message, level=level, logger_name=logger)
@mcp.tool
async def stateful_put(value: int, context: Context) -> None:
"""put a value associated with the server session"""
key = id(context.session)
states[key] = value
@mcp.tool
async def stateful_get(context: Context) -> int:
"""get the value associated with the server session"""
key = id(context.session)
try:
return states[key]
except KeyError:
raise ToolError("Value not found")
return mcp
@pytest.fixture
async def stateful_proxy_server(fastmcp_server: FastMCP):
client = StatefulProxyClient(transport=FastMCPTransport(fastmcp_server))
return FastMCPProxy(client_factory=client.new_stateful)
@pytest.fixture
async def stateless_server(stateful_proxy_server: FastMCP):
port = find_available_port()
url = f"http://127.0.0.1:{port}/mcp/"
task = asyncio.create_task(
stateful_proxy_server.run_http_async(
host="127.0.0.1", port=port, stateless_http=True
)
)
async with Client(transport=url) as client:
assert await client.ping()
yield url
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
class TestStatefulProxyClient:
async def test_concurrent_log_requests_no_mixing(
self, stateful_proxy_server: FastMCP
):
"""Test that concurrent log requests don't mix handlers (fixes #1068)."""
results: dict[str, LogMessage] = {}
async def log_handler_a(message: LogMessage) -> None:
results["logger_a"] = message
async def log_handler_b(message: LogMessage) -> None:
results["logger_b"] = message
async with (
Client(stateful_proxy_server, log_handler=log_handler_a) as client_a,
Client(stateful_proxy_server, log_handler=log_handler_b) as client_b,
):
async with create_task_group() as tg:
tg.start_soon(
client_a.call_tool,
"log",
{"message": "Hello, world!", "level": "info", "logger": "a"},
)
tg.start_soon(
client_b.call_tool,
"log",
{"message": "Hello, world!", "level": "info", "logger": "b"},
)
assert results["logger_a"].logger == "a"
assert results["logger_b"].logger == "b"
async def test_stateful_proxy(self, stateful_proxy_server: FastMCP):
"""Test that the state shared across multiple calls for the same client (fixes #959)."""
async with Client(stateful_proxy_server) as client:
with pytest.raises(ToolError, match="Value not found"):
await client.call_tool("stateful_get", {})
await client.call_tool("stateful_put", {"value": 1})
result = await client.call_tool("stateful_get", {})
assert result.data == 1
async def test_stateless_proxy(self, stateless_server: str):
"""Test that the state will not be shared across different calls,
even if they are from the same client."""
async with Client(stateless_server) as client:
await client.call_tool("stateful_put", {"value": 1})
with pytest.raises(ToolError, match="Value not found"):
await client.call_tool("stateful_get", {})
async def test_multi_proxies_no_mixing(self):
"""Test that the stateful proxy client won't be mixed in multi-proxies sessions."""
mcp_a, mcp_b = FastMCP(), FastMCP()
@mcp_a.tool
def tool_a() -> str:
return "a"
@mcp_b.tool
def tool_b() -> str:
return "b"
proxy_mcp_a = FastMCPProxy(
client_factory=StatefulProxyClient(mcp_a).new_stateful
)
proxy_mcp_b = FastMCPProxy(
client_factory=StatefulProxyClient(mcp_b).new_stateful
)
multi_proxy_mcp = FastMCP()
multi_proxy_mcp.mount(proxy_mcp_a, prefix="a")
multi_proxy_mcp.mount(proxy_mcp_b, prefix="b")
async with Client(multi_proxy_mcp) as client:
result_a = await client.call_tool("a_tool_a", {})
result_b = await client.call_tool("b_tool_b", {})
assert result_a.data == "a"
assert result_b.data == "b"
|