From ea13dc114fcd41f69007b41c55aa4784d60a1cb9 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 12 Apr 2025 12:16:07 -0400 Subject: [PATCH] Improve client return types --- src/fastmcp/client/client.py | 109 +++++++++++++----- src/fastmcp/server/proxy.py | 48 ++++---- ...st_fastmcp_transport.py => test_client.py} | 22 ++-- tests/client/test_roots.py | 4 +- tests/client/test_sampling.py | 6 +- 5 files changed, 119 insertions(+), 70 deletions(-) rename tests/client/{test_fastmcp_transport.py => test_client.py} (88%) diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index ea4f81cd4..dbae1b965 100644 --- a/src/fastmcp/client/client.py +++ b/src/fastmcp/client/client.py @@ -1,7 +1,7 @@ import datetime from contextlib import AbstractAsyncContextManager from pathlib import Path -from typing import Any +from typing import Any, Literal, cast, overload import mcp.types from mcp import ClientSession @@ -24,6 +24,10 @@ from .transports import ClientTransport, SessionKwargs, infer_transport __all__ = ["Client", "RootsHandler", "RootsList"] +class ClientError(ValueError): + """Base class for errors raised by the client.""" + + class Client: """ MCP client that delegates connection management to a Transport instance. @@ -122,60 +126,101 @@ class Client: """Send a logging/setLevel request.""" await self.session.set_logging_level(level) - async def list_resources(self) -> mcp.types.ListResourcesResult: + async def send_roots_list_changed(self) -> None: + """Send a roots/list_changed notification.""" + await self.session.send_roots_list_changed() + + async def list_resources(self) -> list[mcp.types.Resource]: """Send a resources/list request.""" - return await self.session.list_resources() + result = await self.session.list_resources() + return result.resources - async def list_resource_templates(self) -> mcp.types.ListResourceTemplatesResult: + async def list_resource_templates(self) -> list[mcp.types.ResourceTemplate]: """Send a resources/listResourceTemplates request.""" - return await self.session.list_resource_templates() + result = await self.session.list_resource_templates() + return result.resourceTemplates - async def read_resource(self, uri: AnyUrl | str) -> mcp.types.ReadResourceResult: + async def read_resource( + self, uri: AnyUrl | str + ) -> list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents]: """Send a resources/read request.""" if isinstance(uri, str): uri = AnyUrl(uri) # Ensure AnyUrl - return await self.session.read_resource(uri) + result = await self.session.read_resource(uri) + return result.contents - async def subscribe_resource(self, uri: AnyUrl | str) -> None: - """Send a resources/subscribe request.""" - if isinstance(uri, str): - uri = AnyUrl(uri) - await self.session.subscribe_resource(uri) + # async def subscribe_resource(self, uri: AnyUrl | str) -> None: + # """Send a resources/subscribe request.""" + # if isinstance(uri, str): + # uri = AnyUrl(uri) + # await self.session.subscribe_resource(uri) - async def unsubscribe_resource(self, uri: AnyUrl | str) -> None: - """Send a resources/unsubscribe request.""" - if isinstance(uri, str): - uri = AnyUrl(uri) - await self.session.unsubscribe_resource(uri) + # async def unsubscribe_resource(self, uri: AnyUrl | str) -> None: + # """Send a resources/unsubscribe request.""" + # if isinstance(uri, str): + # uri = AnyUrl(uri) + # await self.session.unsubscribe_resource(uri) - async def list_prompts(self) -> mcp.types.ListPromptsResult: + async def list_prompts(self) -> list[mcp.types.Prompt]: """Send a prompts/list request.""" - return await self.session.list_prompts() + result = await self.session.list_prompts() + return result.prompts async def get_prompt( self, name: str, arguments: dict[str, str] | None = None ) -> mcp.types.GetPromptResult: """Send a prompts/get request.""" - return await self.session.get_prompt(name, arguments) + result = await self.session.get_prompt(name, arguments) + return result async def complete( self, ref: mcp.types.ResourceReference | mcp.types.PromptReference, argument: dict[str, str], - ) -> mcp.types.CompleteResult: - """Send a completion/complete request.""" - return await self.session.complete(ref, argument) + ) -> mcp.types.Completion: + """Send a completion request.""" + result = await self.session.complete(ref, argument) + return result.completion - async def list_tools(self) -> mcp.types.ListToolsResult: + async def list_tools(self) -> list[mcp.types.Tool]: """Send a tools/list request.""" - return await self.session.list_tools() + result = await self.session.list_tools() + return result.tools + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + _return_raw_result: Literal[False] = False, + ) -> list[ + mcp.types.TextContent | mcp.types.ImageContent | mcp.types.EmbeddedResource + ]: ... + + @overload + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + _return_raw_result: Literal[True] = True, + ) -> mcp.types.CallToolResult: ... async def call_tool( - self, name: str, arguments: dict[str, Any] | None = None - ) -> mcp.types.CallToolResult: + self, + name: str, + arguments: dict[str, Any] | None = None, + _return_raw_result: bool = False, + ) -> ( + list[ + mcp.types.TextContent | mcp.types.ImageContent | mcp.types.EmbeddedResource + ] + | mcp.types.CallToolResult + ): """Send a tools/call request.""" - return await self.session.call_tool(name, arguments) - - async def send_roots_list_changed(self) -> None: - """Send a roots/list_changed notification.""" - await self.session.send_roots_list_changed() + result = await self.session.call_tool(name, arguments) + if _return_raw_result: + return result + elif result.isError: + msg = cast(mcp.types.TextContent, result.content[0]).text + raise ClientError(msg) + return result.content diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 048d1d1e0..6a129bc2e 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -40,11 +40,15 @@ class ProxyTool(Tool): async def run( self, arguments: dict[str, Any], context: Context | None = None ) -> Any: + # the client context manager will swallow any exceptions inside a TaskGroup + # so we return the raw result and raise an exception ourselves async with self._client: - result = await self._client.call_tool(self.name, arguments) + result = await self._client.call_tool( + self.name, arguments, _return_raw_result=True + ) if result.isError: raise ValueError(cast(mcp.types.TextContent, result.content[0]).text) - return result.content[0] + return result.content class ProxyResource(Resource): @@ -73,12 +77,12 @@ class ProxyResource(Resource): async with self._client: result = await self._client.read_resource(self.uri) - if isinstance(result.contents[0], TextResourceContents): - return result.contents[0].text - elif isinstance(result.contents[0], BlobResourceContents): - return result.contents[0].blob + if isinstance(result[0], TextResourceContents): + return result[0].text + elif isinstance(result[0], BlobResourceContents): + return result[0].blob else: - raise ValueError(f"Unsupported content type: {type(result.contents[0])}") + raise ValueError(f"Unsupported content type: {type(result[0])}") class ProxyTemplate(ResourceTemplate): @@ -103,20 +107,20 @@ class ProxyTemplate(ResourceTemplate): async with self._client: result = await self._client.read_resource(uri) - if isinstance(result.contents[0], TextResourceContents): - value = result.contents[0].text - elif isinstance(result.contents[0], BlobResourceContents): - value = result.contents[0].blob + if isinstance(result[0], TextResourceContents): + value = result[0].text + elif isinstance(result[0], BlobResourceContents): + value = result[0].blob else: - raise ValueError(f"Unsupported content type: {type(result.contents[0])}") + raise ValueError(f"Unsupported content type: {type(result[0])}") return ProxyResource( client=self._client, uri=uri, name=self.name, description=self.description, - mime_type=result.contents[0].mimeType, - contents=result.contents, + mime_type=result[0].mimeType, + contents=result, _value=value, ) @@ -177,15 +181,15 @@ class FastMCPProxy(FastMCP): async with client: # Register proxies for client tools - tools_result = await client.list_tools() - for tool in tools_result.tools: + tools = await client.list_tools() + for tool in tools: tool_proxy = await ProxyTool.from_client(client, tool) server._tool_manager._tools[tool_proxy.name] = tool_proxy logger.debug(f"Created proxy for tool: {tool_proxy.name}") # Register proxies for client resources - resources_result = await client.list_resources() - for resource in resources_result.resources: + resources = await client.list_resources() + for resource in resources: resource_proxy = await ProxyResource.from_client(client, resource) server._resource_manager._resources[str(resource_proxy.uri)] = ( resource_proxy @@ -193,8 +197,8 @@ class FastMCPProxy(FastMCP): logger.debug(f"Created proxy for resource: {resource_proxy.uri}") # Register proxies for client resource templates - templates_result = await client.list_resource_templates() - for template in templates_result.resourceTemplates: + templates = await client.list_resource_templates() + for template in templates: template_proxy = await ProxyTemplate.from_client(client, template) server._resource_manager._templates[template_proxy.uri_template] = ( template_proxy @@ -204,8 +208,8 @@ class FastMCPProxy(FastMCP): ) # Register proxies for client prompts - prompts_result = await client.list_prompts() - for prompt in prompts_result.prompts: + prompts = await client.list_prompts() + for prompt in prompts: prompt_proxy = await ProxyPrompt.from_client(client, prompt) server._prompt_manager._prompts[prompt_proxy.name] = prompt_proxy logger.debug(f"Created proxy for prompt: {prompt_proxy.name}") diff --git a/tests/client/test_fastmcp_transport.py b/tests/client/test_client.py similarity index 88% rename from tests/client/test_fastmcp_transport.py rename to tests/client/test_client.py index bf097cc2b..4d933a185 100644 --- a/tests/client/test_fastmcp_transport.py +++ b/tests/client/test_client.py @@ -51,8 +51,8 @@ async def test_list_tools(fastmcp_server): result = await client.list_tools() # Check that our tools are available - assert len(result.tools) == 2 - assert set(tool.name for tool in result.tools) == {"greet", "add"} + assert len(result) == 2 + assert set(tool.name for tool in result) == {"greet", "add"} async def test_call_tool(fastmcp_server): @@ -63,7 +63,7 @@ async def test_call_tool(fastmcp_server): result = await client.call_tool("greet", {"name": "World"}) # The result content should contain our greeting - content_str = str(result.content[0]) + content_str = str(result[0]) assert "Hello, World!" in content_str @@ -75,8 +75,8 @@ async def test_list_resources(fastmcp_server): result = await client.list_resources() # Check that our resource is available - assert len(result.resources) == 1 - assert str(result.resources[0].uri) == "data://users" + assert len(result) == 1 + assert str(result[0].uri) == "data://users" async def test_list_prompts(fastmcp_server): @@ -87,8 +87,8 @@ async def test_list_prompts(fastmcp_server): result = await client.list_prompts() # Check that our prompt is available - assert len(result.prompts) == 1 - assert result.prompts[0].name == "welcome" + assert len(result) == 1 + assert result[0].name == "welcome" async def test_get_prompt(fastmcp_server): @@ -115,7 +115,7 @@ async def test_read_resource(fastmcp_server): result = await client.read_resource(uri) # The contents should include our user list - contents_str = str(result.contents[0]) + contents_str = str(result[0]) assert "Alice" in contents_str assert "Bob" in contents_str assert "Charlie" in contents_str @@ -145,15 +145,15 @@ async def test_resource_template(fastmcp_server): result = await client.list_resource_templates() # Check that our template is available - assert len(result.resourceTemplates) == 1 - assert "data://user/{user_id}" in result.resourceTemplates[0].uriTemplate + assert len(result) == 1 + assert "data://user/{user_id}" in result[0].uriTemplate # Now use the template with a specific user_id uri = cast(AnyUrl, "data://user/123") result = await client.read_resource(uri) # Check the content matches what we expect for the provided user_id - content_str = str(result.contents[0]) + content_str = str(result[0]) assert '"id": "123"' in content_str assert '"name": "User 123"' in content_str assert '"active": true' in content_str diff --git a/tests/client/test_roots.py b/tests/client/test_roots.py index 0ede20f01..74b478a6d 100644 --- a/tests/client/test_roots.py +++ b/tests/client/test_roots.py @@ -41,8 +41,8 @@ class TestClientRoots: async def test_valid_roots(self, fastmcp_server: FastMCP, roots: list[str]): async with Client(fastmcp_server, roots=roots) as client: result = await client.call_tool("list_roots", {}) - assert isinstance(result.content[0], TextContent) - assert json.loads(result.content[0].text) == [ + assert isinstance(result[0], TextContent) + assert json.loads(result[0].text) == [ "file://x/y/z", "file://x/y/z", ] diff --git a/tests/client/test_sampling.py b/tests/client/test_sampling.py index 0707f3a44..15c2a2ff1 100644 --- a/tests/client/test_sampling.py +++ b/tests/client/test_sampling.py @@ -47,7 +47,7 @@ async def test_simple_sampling(fastmcp_server: FastMCP): async with Client(fastmcp_server, sampling_handler=sampling_handler) as client: result = await client.call_tool("simple_sample", {"message": "Hello, world!"}) - reply = cast(TextContent, result.content[0]) + reply = cast(TextContent, result[0]) assert reply.text == "This is the sample message!" @@ -62,7 +62,7 @@ async def test_sampling_with_system_prompt(fastmcp_server: FastMCP): result = await client.call_tool( "sample_with_system_prompt", {"message": "Hello, world!"} ) - reply = cast(TextContent, result.content[0]) + reply = cast(TextContent, result[0]) assert reply.text == "You love FastMCP" @@ -81,5 +81,5 @@ async def test_sampling_with_messages(fastmcp_server: FastMCP): result = await client.call_tool( "sample_with_messages", {"message": "Hello, world!"} ) - reply = cast(TextContent, result.content[0]) + reply = cast(TextContent, result[0]) assert reply.text == "I need to think."