Spaces:
Running
Running
File size: 15,626 Bytes
fd0c13f c7708b2 fd0c13f 9d7c3c2 fd0c13f 6cdad57 fd0c13f 9247f8d fd0c13f 6cdad57 957b451 6cdad57 fd0c13f c7708b2 9d7c3c2 c7708b2 | 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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 | from dataclasses import dataclass
from typing import cast
import pytest
from anyio import create_task_group
from mcp.types import LoggingLevel, ModelHint, ModelPreferences, TextContent
from pydantic import BaseModel, Field
from fastmcp import Client, Context, FastMCP
from fastmcp.client.elicitation import ElicitRequestParams, ElicitResult
from fastmcp.client.logging import LogMessage
from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
from fastmcp.exceptions import ToolError
from fastmcp.server.proxy import ProxyClient
@pytest.fixture
def fastmcp_server():
mcp = FastMCP("TestServer")
@mcp.tool(tags={"echo"})
def echo(message: str) -> str:
return f"echo: {message}"
@mcp.tool
async def list_roots(context: Context) -> list[str]:
roots = await context.list_roots()
return [str(r.uri) for r in roots]
@mcp.tool
async def sampling(
context: Context,
) -> str:
result = await context.sample(
"Hello, world!",
system_prompt="You love FastMCP",
include_context="thisServer",
temperature=0.5,
max_tokens=100,
model_preferences="gpt-4o",
)
return cast(TextContent, result).text
@dataclass
class Person:
name: str
@mcp.tool
async def elicit(context: Context) -> str:
result = await context.elicit(
message="What is your name?",
response_type=Person,
)
if result.action == "accept":
return f"Hello, {result.data.name}!"
else:
return "No name provided."
@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 report_progress(context: Context) -> int:
for i in range(3):
await context.report_progress(
progress=i + 1,
total=3,
message=f"{(i + 1) / 3 * 100:.2f}% complete",
)
return 100
return mcp
@pytest.fixture
async def proxy_server(fastmcp_server: FastMCP):
"""
A proxy server that forwards interactions with the proxy client to the given fastmcp server.
"""
return FastMCP.as_proxy(ProxyClient(fastmcp_server))
class TestProxyClient:
async def test_forward_tool_meta(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `echo` tool meta.
"""
async with Client(proxy_server) as client:
tools = await client.list_tools()
echo_tool = next(t for t in tools if t.name == "echo")
assert echo_tool.meta == {"_fastmcp": {"tags": ["echo"]}}
async def test_forward_error_response(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards an error response.
"""
async with Client(proxy_server) as client:
with pytest.raises(ToolError, match="Elicitation not supported"):
await client.call_tool("elicit", {})
async def test_forward_list_roots_request(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `list_roots` request.
"""
roots_handler_called = False
async def roots_handler(ctx: RequestContext):
nonlocal roots_handler_called
roots_handler_called = True
return []
async with Client(proxy_server, roots=roots_handler) as client:
await client.call_tool("list_roots", {})
assert roots_handler_called
async def test_forward_list_roots_response(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `list_roots` response.
"""
async with Client(proxy_server, roots=["file://x/y/z"]) as client:
result = await client.call_tool("list_roots", {})
assert result.data == ["file://x/y/z"]
async def test_forward_sampling_request(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `sampling` request.
"""
sampling_handler_called = False
def sampling_handler(
messages: list[SamplingMessage],
params: SamplingParams,
ctx: RequestContext,
) -> str:
nonlocal sampling_handler_called
sampling_handler_called = True
assert messages == [
SamplingMessage(
role="user",
content=TextContent(type="text", text="Hello, world!"),
)
]
assert params.systemPrompt == "You love FastMCP"
assert params.temperature == 0.5
assert params.maxTokens == 100
assert params.modelPreferences == ModelPreferences(
hints=[ModelHint(name="gpt-4o")]
)
return ""
async with Client(proxy_server, sampling_handler=sampling_handler) as client:
await client.call_tool("sampling", {})
assert sampling_handler_called
async def test_forward_sampling_response(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `sampling` response.
"""
async with Client(
proxy_server, sampling_handler=lambda *args: "I love FastMCP"
) as client:
result = await client.call_tool("sampling", {})
assert result.data == "I love FastMCP"
async def test_elicit_request(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `elicit` request.
"""
elicitation_handler_called = False
async def elicitation_handler(
message, response_type, params: ElicitRequestParams, ctx
):
nonlocal elicitation_handler_called
elicitation_handler_called = True
assert message == "What is your name?"
assert "Person" in str(response_type)
assert params.requestedSchema == {
"title": "Person",
"type": "object",
"properties": {"name": {"title": "Name", "type": "string"}},
"required": ["name"],
}
return ElicitResult(action="accept", content=response_type(name="Alice"))
async with Client(
proxy_server, elicitation_handler=elicitation_handler
) as client:
await client.call_tool("elicit", {})
assert elicitation_handler_called
async def test_elicit_accept_response(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `elicit` accept response.
"""
async def elicitation_handler(
message, response_type, params: ElicitRequestParams, ctx
):
return ElicitResult(action="accept", content=response_type(name="Alice"))
async with Client(
proxy_server,
elicitation_handler=elicitation_handler,
) as client:
result = await client.call_tool("elicit", {})
assert result.data == "Hello, Alice!"
async def test_elicit_decline_response(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `elicit` decline response.
"""
async def elicitation_handler(
message, response_type, params: ElicitRequestParams, ctx
):
return ElicitResult(action="decline")
async with Client(
proxy_server, elicitation_handler=elicitation_handler
) as client:
result = await client.call_tool("elicit", {})
assert result.data == "No name provided."
async def test_log_request(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `log` request.
"""
log_handler_called = False
async def log_handler(message: LogMessage) -> None:
nonlocal log_handler_called
log_handler_called = True
assert message.data == "Hello, world!"
assert message.level == "info"
assert message.logger == "test"
async with Client(proxy_server, log_handler=log_handler) as client:
await client.call_tool(
"log", {"message": "Hello, world!", "level": "info", "logger": "test"}
)
assert log_handler_called
async def test_report_progress_request(self, proxy_server: FastMCP):
"""
Test that the proxy client correctly forwards the `report_progress` request.
"""
EXPECTED_PROGRESS_MESSAGES = [
dict(progress=1, total=3, message="33.33% complete"),
dict(progress=2, total=3, message="66.67% complete"),
dict(progress=3, total=3, message="100.00% complete"),
]
PROGRESS_MESSAGES = []
async def progress_handler(
progress: float, total: float | None, message: str | None
) -> None:
PROGRESS_MESSAGES.append(
dict(progress=progress, total=total, message=message)
)
async with Client(proxy_server, progress_handler=progress_handler) as client:
await client.call_tool("report_progress", {})
assert PROGRESS_MESSAGES == EXPECTED_PROGRESS_MESSAGES
async def test_concurrent_log_requests_no_mixing(self, 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(proxy_server, log_handler=log_handler_a) as client_a,
Client(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_concurrent_elicitation_no_mixing(self, proxy_server: FastMCP):
"""Test that concurrent elicitation requests don't mix handlers (fixes #1068)."""
results = {}
async def elicitation_handler_a(
message: str,
response_type: type,
params: ElicitRequestParams,
ctx: RequestContext,
) -> ElicitResult:
return ElicitResult(action="accept", content=response_type(name="Alice"))
async def elicitation_handler_b(
message: str,
response_type: type,
params: ElicitRequestParams,
ctx: RequestContext,
) -> ElicitResult:
return ElicitResult(action="accept", content=response_type(name="Bob"))
async def get_and_store(name, coro):
result = await coro
results[name] = result.data
async with (
Client(proxy_server, elicitation_handler=elicitation_handler_a) as client_a,
Client(proxy_server, elicitation_handler=elicitation_handler_b) as client_b,
):
async with create_task_group() as tg:
tg.start_soon(
get_and_store,
"elicitation_a",
client_a.call_tool("elicit", {}),
)
tg.start_soon(
get_and_store,
"elicitation_b",
client_b.call_tool("elicit", {}),
)
assert results["elicitation_a"] == "Hello, Alice!"
assert results["elicitation_b"] == "Hello, Bob!"
async def test_elicit_with_default_values(self, fastmcp_server: FastMCP):
"""
Test that the proxy client correctly handles elicitation with default values (fixes #1167).
"""
@fastmcp_server.tool
async def elicit_with_defaults(context: Context) -> str:
class TestModel(BaseModel):
content: str = Field(description="Your reply content")
acknowledge: bool = Field(
default=False, description="Send immediately or save as draft"
)
result = await context.elicit(
"Please provide input:", response_type=TestModel
)
if result.action == "accept":
return f"Content: {result.data.content}, Acknowledge: {result.data.acknowledge}"
else:
return f"Elicitation {result.action}"
proxy_server = FastMCP.as_proxy(ProxyClient(fastmcp_server))
# Test that elicitation works correctly through the proxy
async def elicitation_handler(
message: str,
response_type: type,
params: ElicitRequestParams,
ctx: RequestContext,
):
# Verify the schema is correct - acknowledge should have default=False, not be nullable
schema = params.requestedSchema
assert schema["properties"]["acknowledge"]["type"] == "boolean"
assert schema["properties"]["acknowledge"]["default"] is False
return {"content": "Test content", "acknowledge": True}
async with Client(
proxy_server, elicitation_handler=elicitation_handler
) as client:
result = await client.call_tool("elicit_with_defaults", {})
assert result.data == "Content: Test content, Acknowledge: True"
async def test_client_factory_creates_fresh_sessions(self, fastmcp_server: FastMCP):
"""Test that the client factory pattern creates fresh sessions for each request."""
from fastmcp.server.proxy import FastMCPProxy
# Create a disconnected client (should use fresh sessions per request)
base_client = Client(fastmcp_server)
# Test both as_proxy convenience method and direct client_factory usage
proxy_via_as_proxy = FastMCP.as_proxy(base_client)
proxy_via_factory = FastMCPProxy(client_factory=base_client.new)
# Verify the proxies are created successfully - this tests the client factory pattern
assert proxy_via_as_proxy is not None
assert proxy_via_factory is not None
# Verify they have the expected client factory behavior
assert hasattr(proxy_via_as_proxy, "_tool_manager")
assert hasattr(proxy_via_factory, "_tool_manager")
async def test_connected_client_reuses_sessions(self, fastmcp_server: FastMCP):
"""Test that connected clients passed to as_proxy reuse sessions (preserves #959 behavior)."""
# Create a connected client (should reuse sessions)
async with Client(fastmcp_server) as connected_client:
proxy = FastMCP.as_proxy(connected_client)
# Verify the proxy is created successfully and uses session reuse
assert proxy is not None
assert hasattr(proxy, "_tool_manager")
|