Implement session functions for async tools

Author: LucaButBoring

src/mcp/client/session.py

@@ -2,6 +2,7 @@
 from datetime import timedelta
 from typing import Any, Protocol
 
+import anyio
 import anyio.lowlevel
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from jsonschema import SchemaError, ValidationError, validate
@@ -273,15 +274,26 @@ async def call_tool(
         arguments: dict[str, Any] | None = None,
         read_timeout_seconds: timedelta | None = None,
         progress_callback: ProgressFnT | None = None,
+        *,
+        async_properties: types.AsyncRequestProperties | None = None,
     ) -> types.CallToolResult:
-        """Send a tools/call request with optional progress callback support."""
+        """Send a tools/call request with optional progress callback support.
+
+        Args:
+            name: Name of the tool to call
+            arguments: Arguments to pass to the tool
+            read_timeout_seconds: Read timeout for the request
+            progress_callback: Optional progress callback
+            async_properties: Optional async parameters for async tool execution
+        """
 
         result = await self.send_request(
             types.ClientRequest(
                 types.CallToolRequest(
                     params=types.CallToolRequestParams(
                         name=name,
                         arguments=arguments,
+                        async_properties=async_properties,
                     ),
                 )
             ),
@@ -295,6 +307,42 @@ async def call_tool(
 
         return result
 
+    async def check_tool_async_status(self, token: str) -> types.CheckToolAsyncStatusResult:
+        """Check the status of an async tool operation.
+
+        Args:
+            token: Token returned from async call_tool
+
+        Returns:
+            Status result with current operation state
+        """
+        return await self.send_request(
+            types.ClientRequest(
+                types.CheckToolAsyncStatusRequest(
+                    params=types.CheckToolAsyncStatusParams(token=token),
+                )
+            ),
+            types.CheckToolAsyncStatusResult,
+        )
+
+    async def get_tool_async_result(self, token: str) -> types.GetToolAsyncPayloadResult:
+        """Get the result of a completed async tool operation.
+
+        Args:
+            token: Token returned from async call_tool
+
+        Returns:
+            The final tool result
+        """
+        return await self.send_request(
+            types.ClientRequest(
+                types.GetToolAsyncPayloadRequest(
+                    params=types.GetToolAsyncPayloadParams(token=token),
+                )
+            ),
+            types.GetToolAsyncPayloadResult,
+        )
+
     async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
         """Validate the structured content of a tool result against its output schema."""
         if name not in self._tool_output_schemas:

src/mcp/server/fastmcp/server.py

@@ -30,6 +30,7 @@
 from mcp.server.fastmcp.prompts import Prompt, PromptManager
 from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager
 from mcp.server.fastmcp.tools import Tool, ToolManager
+from mcp.server.fastmcp.tools.base import InvocationMode
 from mcp.server.fastmcp.utilities.context_injection import find_context_parameter
 from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger
 from mcp.server.lowlevel.helper_types import ReadResourceContents
@@ -43,7 +44,7 @@
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared.context import LifespanContextT, RequestContext, RequestT
-from mcp.types import AnyFunction, ContentBlock, GetPromptResult, ToolAnnotations
+from mcp.types import NEXT_PROTOCOL_VERSION, AnyFunction, ContentBlock, GetPromptResult, ToolAnnotations
 from mcp.types import Prompt as MCPPrompt
 from mcp.types import PromptArgument as MCPPromptArgument
 from mcp.types import Resource as MCPResource
@@ -266,9 +267,39 @@ def _setup_handlers(self) -> None:
         self._mcp_server.get_prompt()(self.get_prompt)
         self._mcp_server.list_resource_templates()(self.list_resource_templates)
 
+    def _client_supports_async(self) -> bool:
+        """Check if the current client supports async tools based on protocol version."""
+        try:
+            context = self.get_context()
+            if context.request_context and context.request_context.session.client_params:
+                client_version = str(context.request_context.session.client_params.protocolVersion)
+                # Only "next" version supports async tools for now
+                return client_version == NEXT_PROTOCOL_VERSION
+        except ValueError:
+            # Context not available (outside of request), assume no async support
+            pass
+        return False
+
+    def _get_invocation_mode(self, info: Tool, client_supports_async: bool) -> Literal["sync", "async"] | None:
+        """Determine invocationMode field based on client support."""
+        if not client_supports_async:
+            return None  # Old clients don't see invocationMode field
+
+        # New clients see the invocationMode field
+        if "async" in info.invocation_modes and len(info.invocation_modes) == 1:
+            return "async"  # Async-only
+        elif len(info.invocation_modes) > 1 or info.invocation_modes == ["sync"]:
+            return "sync"  # Hybrid or explicit sync
+        return None
+
     async def list_tools(self) -> list[MCPTool]:
         """List all available tools."""
         tools = self._tool_manager.list_tools()
+
+        # Check if client supports async tools based on protocol version
+        client_supports_async = self._client_supports_async()
+
+        # Filter out async-only tools for old clients and set invocationMode based on client support
         return [
             MCPTool(
                 name=info.name,
@@ -277,8 +308,10 @@ async def list_tools(self) -> list[MCPTool]:
                 inputSchema=info.parameters,
                 outputSchema=info.output_schema,
                 annotations=info.annotations,
+                invocationMode=self._get_invocation_mode(info, client_supports_async),
             )
             for info in tools
+            if client_supports_async or info.invocation_modes != ["async"]
         ]
 
     def get_context(self) -> Context[ServerSession, LifespanResultT, Request]:
@@ -348,6 +381,7 @@ def add_tool(
         description: str | None = None,
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
+        invocation_modes: list[InvocationMode] | None = None,
     ) -> None:
         """Add a tool to the server.
 
@@ -364,6 +398,8 @@ def add_tool(
                 - If None, auto-detects based on the function's return type annotation
                 - If True, unconditionally creates a structured tool (return type annotation permitting)
                 - If False, unconditionally creates an unstructured tool
+            invocation_modes: List of supported invocation modes (e.g., ["sync", "async"])
+                - If None, defaults to ["sync"] for backwards compatibility
         """
         self._tool_manager.add_tool(
             fn,
@@ -372,6 +408,7 @@ def add_tool(
             description=description,
             annotations=annotations,
             structured_output=structured_output,
+            invocation_modes=invocation_modes,
         )
 
     def tool(
@@ -381,6 +418,7 @@ def tool(
         description: str | None = None,
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
+        invocation_modes: list[InvocationMode] | None = None,
     ) -> Callable[[AnyFunction], AnyFunction]:
         """Decorator to register a tool.
 
@@ -397,6 +435,10 @@ def tool(
                 - If None, auto-detects based on the function's return type annotation
                 - If True, unconditionally creates a structured tool (return type annotation permitting)
                 - If False, unconditionally creates an unstructured tool
+            invocation_modes: List of supported invocation modes (e.g., ["sync", "async"])
+                - If None, defaults to ["sync"] for backwards compatibility
+                - Supports "sync" for synchronous execution and "async" for asynchronous execution
+                - Tools with "async" mode will be hidden from clients that don't support async execution
 
         Example:
             @server.tool()
@@ -412,6 +454,17 @@ def tool_with_context(x: int, ctx: Context) -> str:
             async def async_tool(x: int, context: Context) -> str:
                 await context.report_progress(50, 100)
                 return str(x)
+
+            @server.tool(invocation_modes=["async"])
+            async def async_only_tool(data: str, ctx: Context) -> str:
+                # This tool only supports async execution
+                await ctx.info("Starting long-running analysis...")
+                return await analyze_data(data)
+
+            @server.tool(invocation_modes=["sync", "async"])
+            def hybrid_tool(x: int) -> str:
+                # This tool supports both sync and async execution
+                return str(x)
         """
         # Check if user passed function directly instead of calling decorator
         if callable(name):
@@ -427,6 +480,7 @@ def decorator(fn: AnyFunction) -> AnyFunction:
                 description=description,
                 annotations=annotations,
                 structured_output=structured_output,
+                invocation_modes=invocation_modes,
             )
             return fn
 

src/mcp/server/fastmcp/tools/base.py

@@ -4,7 +4,7 @@
 import inspect
 from collections.abc import Callable
 from functools import cached_property
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -18,6 +18,8 @@
     from mcp.server.session import ServerSessionT
     from mcp.shared.context import LifespanContextT, RequestT
 
+InvocationMode = Literal["sync", "async"]
+
 
 class Tool(BaseModel):
     """Internal tool registration info."""
@@ -33,6 +35,9 @@ class Tool(BaseModel):
     is_async: bool = Field(description="Whether the tool is async")
     context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context")
     annotations: ToolAnnotations | None = Field(None, description="Optional annotations for the tool")
+    invocation_modes: list[InvocationMode] = Field(
+        default=["sync"], description="Supported invocation modes (sync/async)"
+    )
 
     @cached_property
     def output_schema(self) -> dict[str, Any] | None:
@@ -48,6 +53,7 @@ def from_function(
         context_kwarg: str | None = None,
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
+        invocation_modes: list[InvocationMode] | None = None,
     ) -> Tool:
         """Create a Tool from a function."""
         func_name = name or fn.__name__
@@ -68,6 +74,10 @@ def from_function(
         )
         parameters = func_arg_metadata.arg_model.model_json_schema(by_alias=True)
 
+        # Default to sync mode if no invocation modes specified
+        if invocation_modes is None:
+            invocation_modes = ["sync"]
+
         return cls(
             fn=fn,
             name=func_name,
@@ -78,6 +88,7 @@ def from_function(
             is_async=is_async,
             context_kwarg=context_kwarg,
             annotations=annotations,
+            invocation_modes=invocation_modes,
         )
 
     async def run(

src/mcp/server/fastmcp/tools/tool_manager.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 from mcp.server.fastmcp.exceptions import ToolError
-from mcp.server.fastmcp.tools.base import Tool
+from mcp.server.fastmcp.tools.base import InvocationMode, Tool
 from mcp.server.fastmcp.utilities.logging import get_logger
 from mcp.shared.context import LifespanContextT, RequestT
 from mcp.types import ToolAnnotations
@@ -50,15 +50,21 @@ def add_tool(
         description: str | None = None,
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
+        invocation_modes: list[InvocationMode] | None = None,
     ) -> Tool:
         """Add a tool to the server."""
+        # Default to sync mode if no invocation modes specified
+        if invocation_modes is None:
+            invocation_modes = ["sync"]
+
         tool = Tool.from_function(
             fn,
             name=name,
             title=title,
             description=description,
             annotations=annotations,
             structured_output=structured_output,
+            invocation_modes=invocation_modes,
         )
         existing = self._tools.get(tool.name)
         if existing:

tests/server/fastmcp/test_server.py

@@ -603,6 +603,54 @@ def get_settings() -> dict[str, str]:
             assert result.isError is False
             assert result.structuredContent == {"theme": "dark", "language": "en", "timezone": "UTC"}
 
+    @pytest.mark.anyio
+    async def test_list_tools_invocation_mode_sync(self):
+        """Test that sync tools have proper invocationMode field."""
+        mcp = FastMCP()
+
+        @mcp.tool()
+        def sync_tool(x: int) -> int:
+            """A sync tool."""
+            return x * 2
+
+        async with client_session(mcp._mcp_server) as client:
+            tools = await client.list_tools()
+            tool = next(t for t in tools.tools if t.name == "sync_tool")
+            # Sync tools should not have invocationMode field (None) for old clients
+            assert tool.invocationMode is None
+
+    @pytest.mark.anyio
+    async def test_list_tools_invocation_mode_async_only(self):
+        """Test that async-only tools have proper invocationMode field."""
+        mcp = FastMCP()
+
+        @mcp.tool(invocation_modes=["async"])
+        async def async_only_tool(x: int) -> int:
+            """An async-only tool."""
+            return x * 2
+
+        async with client_session(mcp._mcp_server) as client:
+            tools = await client.list_tools()
+            # Async-only tools should be filtered out for old clients
+            async_tools = [t for t in tools.tools if t.name == "async_only_tool"]
+            assert len(async_tools) == 0
+
+    @pytest.mark.anyio
+    async def test_list_tools_invocation_mode_hybrid(self):
+        """Test that hybrid tools have proper invocationMode field."""
+        mcp = FastMCP()
+
+        @mcp.tool(invocation_modes=["sync", "async"])
+        def hybrid_tool(x: int) -> int:
+            """A hybrid tool."""
+            return x * 2
+
+        async with client_session(mcp._mcp_server) as client:
+            tools = await client.list_tools()
+            tool = next(t for t in tools.tools if t.name == "hybrid_tool")
+            # Hybrid tools should not have invocationMode field (None) for old clients
+            assert tool.invocationMode is None
+
 
 class TestServerResources:
     @pytest.mark.anyio