Support configuring async tool keepalives

Author: LucaButBoring

src/mcp/server/fastmcp/server.py

@@ -130,6 +130,8 @@ async def wrap(_: MCPServer[LifespanResultT, Request]) -> AsyncIterator[Lifespan
 
 
 class FastMCP(Generic[LifespanResultT]):
+    _tool_manager: ToolManager
+
     def __init__(
         self,
         name: str | None = None,
@@ -361,6 +363,7 @@ async def list_tools(self) -> list[MCPTool]:
                 outputSchema=info.output_schema,
                 annotations=info.annotations,
                 invocationMode=self._get_invocation_mode(info, client_supports_async),
+                _meta=info.meta,
             )
             for info in tools
             if client_supports_async or info.invocation_modes != ["async"]
@@ -434,6 +437,7 @@ def add_tool(
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
         invocation_modes: list[InvocationMode] | None = None,
+        keep_alive: int | None = None,
     ) -> None:
         """Add a tool to the server.
 
@@ -452,6 +456,8 @@ def add_tool(
                 - 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
+            keep_alive: How long (in seconds) async operation results should be kept available.
+                Only applies to async tools.
         """
         self._tool_manager.add_tool(
             fn,
@@ -461,6 +467,7 @@ def add_tool(
             annotations=annotations,
             structured_output=structured_output,
             invocation_modes=invocation_modes,
+            keep_alive=keep_alive,
         )
 
     def tool(
@@ -471,6 +478,7 @@ def tool(
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
         invocation_modes: list[InvocationMode] | None = None,
+        keep_alive: int | None = None,
     ) -> Callable[[AnyFunction], AnyFunction]:
         """Decorator to register a tool.
 
@@ -491,6 +499,8 @@ def tool(
                 - 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
+            keep_alive: How long (in seconds) async operation results should be kept available.
+                Only applies to async tools.
 
         Example:
             @server.tool()
@@ -533,6 +543,7 @@ def decorator(fn: AnyFunction) -> AnyFunction:
                 annotations=annotations,
                 structured_output=structured_output,
                 invocation_modes=invocation_modes,
+                keep_alive=keep_alive,
             )
             return fn
 

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

@@ -38,6 +38,7 @@ class Tool(BaseModel):
     invocation_modes: list[InvocationMode] = Field(
         default=["sync"], description="Supported invocation modes (sync/async)"
     )
+    meta: dict[str, Any] | None = Field(description="Optional additional tool information.", default=None)
 
     @cached_property
     def output_schema(self) -> dict[str, Any] | None:
@@ -54,6 +55,7 @@ def from_function(
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
         invocation_modes: list[InvocationMode] | None = None,
+        meta: dict[str, Any] | None = None,
     ) -> Tool:
         """Create a Tool from a function."""
         func_name = name or fn.__name__
@@ -89,6 +91,7 @@ def from_function(
             context_kwarg=context_kwarg,
             annotations=annotations,
             invocation_modes=invocation_modes,
+            meta=meta,
         )
 
     async def run(

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

@@ -51,12 +51,38 @@ def add_tool(
         annotations: ToolAnnotations | None = None,
         structured_output: bool | None = None,
         invocation_modes: list[InvocationMode] | None = None,
+        keep_alive: int | None = None,
+        meta: dict[str, Any] | 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"]
 
+        # Set appropriate default keep_alive based on async compatibility
+        # if user didn't specify custom keep_alive
+        if keep_alive is None and "async" in invocation_modes:
+            keep_alive = 3600  # Default for async-compatible tools
+
+        # Validate keep_alive is only used with async-compatible tools
+        if keep_alive is not None and "async" not in invocation_modes:
+            raise ValueError(
+                f"keep_alive parameter can only be used with async-compatible tools. "
+                f"Tool '{name or fn.__name__}' has invocation_modes={invocation_modes} "
+                f"but specifies keep_alive={keep_alive}. "
+                f"Add 'async' to invocation_modes to use keep_alive."
+            )
+
+        meta = meta or {}
+        if keep_alive is not None:
+            meta.update(
+                {
+                    # default keepalive value is stashed in _meta to pass it to the lowlevel Server
+                    # without adding it to the actual protocol-level tool definition
+                    "_keep_alive": keep_alive
+                }
+            )
+
         tool = Tool.from_function(
             fn,
             name=name,
@@ -65,6 +91,7 @@ def add_tool(
             annotations=annotations,
             structured_output=structured_output,
             invocation_modes=invocation_modes,
+            meta=meta,
         )
         existing = self._tools.get(tool.name)
         if existing:

src/mcp/server/lowlevel/server.py

@@ -468,10 +468,13 @@ async def handler(req: types.CallToolRequest):
 
                     # Check for async execution
                     if tool and self.async_operations and self._should_execute_async(tool):
+                        keep_alive = self._get_tool_keep_alive(tool)
+
                         # Create async operation
                         operation = self.async_operations.create_operation(
                             tool_name=tool_name,
                             arguments=arguments,
+                            keep_alive=keep_alive,
                         )
                         logger.debug(f"Created async operation with token: {operation.token}")
 
@@ -499,7 +502,7 @@ async def execute_async():
                                 content=[],
                                 operation=types.AsyncResultProperties(
                                     token=operation.token,
-                                    keepAlive=3600,
+                                    keepAlive=operation.keep_alive,
                                 ),
                             )
                         )
@@ -576,6 +579,12 @@ def _should_execute_async(self, tool: types.Tool) -> bool:
         invocation_mode = getattr(tool, "invocationMode", None)
         return invocation_mode == "async"
 
+    def _get_tool_keep_alive(self, tool: types.Tool) -> int:
+        """Get the keepalive value for an async tool."""
+        if not tool.meta or "_keep_alive" not in tool.meta:
+            raise ValueError(f"_keep_alive not defined for tool {tool.name}")
+        return cast(int, tool.meta["_keep_alive"])
+
     def progress_notification(self):
         def decorator(
             func: Callable[[str | int, float, float | None, str | None], Awaitable[None]],

tests/server/fastmcp/test_server.py

@@ -749,6 +749,125 @@ async def async_invalid_tool() -> list[int]:
                     pytest.fail("Operation should have failed due to validation error")
                 await asyncio.sleep(0.01)
 
+    @pytest.mark.anyio
+    async def test_tool_keep_alive_validation_no_sync_only(self):
+        """Test that keep_alive validation prevents use on sync-only tools."""
+        mcp = FastMCP()
+
+        # Should raise error when keep_alive is used on sync-only tool
+        with pytest.raises(ValueError, match="keep_alive parameter can only be used with async-compatible tools"):
+
+            @mcp.tool(keep_alive=1800)  # Custom keep_alive on sync-only tool
+            def sync_only_tool(x: int) -> str:
+                return str(x)
+
+    @pytest.mark.anyio
+    async def test_tool_keep_alive_default_async_tools(self):
+        """Test that async tools get correct default keep_alive."""
+        mcp = FastMCP()
+
+        # Async tools should get default keep_alive of 3600
+        @mcp.tool(invocation_modes=["async"])  # No keep_alive specified
+        def async_tool_default(x: int) -> str:
+            return str(x)
+
+        tools = mcp._tool_manager.list_tools()
+        tool = next(t for t in tools if t.name == "async_tool_default")
+        assert tool.meta is not None
+        assert tool.meta["_keep_alive"] == 3600
+
+    @pytest.mark.anyio
+    async def test_async_tool_keep_alive_expiry(self):
+        """Test that async operations expire after keep_alive duration."""
+        mcp = FastMCP("AsyncKeepAliveTest")
+
+        @mcp.tool(invocation_modes=["async"], keep_alive=1)  # 1 second keep_alive
+        def short_lived_tool(data: str) -> str:
+            return f"Processed: {data}"
+
+        # Check that the tool has correct keep_alive
+        tools = mcp._tool_manager.list_tools()
+        tool = next(t for t in tools if t.name == "short_lived_tool")
+        assert tool.meta is not None
+        assert tool.meta["_keep_alive"] == 1
+
+        async with client_session(mcp._mcp_server, protocol_version="next") as client:
+            # First list tools to populate keep_alive mapping
+            await client.list_tools()
+
+            # Call the async tool
+            result = await client.call_tool("short_lived_tool", {"data": "test"})
+
+            # Should get operation token
+            assert result.operation is not None
+            token = result.operation.token
+            assert result.operation.keepAlive == 1
+
+            # Wait for operation to complete
+            while True:
+                status = await client.get_operation_status(token)
+                if status.status == "completed":
+                    break
+
+            # Get result while still alive
+            operation_result = await client.get_operation_result(token)
+            assert operation_result.result is not None
+
+            # Wait for keep_alive to expire (1 second + buffer)
+            await asyncio.sleep(1.2)
+
+            # Operation should now be expired/unavailable
+            with pytest.raises(Exception):  # Should raise error for expired operation
+                await client.get_operation_result(token)
+
+    @pytest.mark.anyio
+    async def test_async_tool_keep_alive_expiry_structured_content(self):
+        """Test that async operations with structured content expire correctly."""
+        mcp = FastMCP("AsyncKeepAliveStructuredTest")
+
+        class ProcessResult(BaseModel):
+            status: str
+            data: str
+            count: int
+
+        @mcp.tool(invocation_modes=["async"], keep_alive=1)  # 1 second keep_alive
+        def structured_tool(input_data: str) -> ProcessResult:
+            return ProcessResult(status="success", data=f"Processed: {input_data}", count=42)
+
+        async with client_session(mcp._mcp_server, protocol_version="next") as client:
+            # First list tools to populate keep_alive mapping
+            await client.list_tools()
+
+            # Call the async tool
+            result = await client.call_tool("structured_tool", {"input_data": "test"})
+
+            # Should get operation token
+            assert result.operation is not None
+            token = result.operation.token
+            assert result.operation.keepAlive == 1
+
+            # Wait for operation to complete
+            while True:
+                status = await client.get_operation_status(token)
+                if status.status == "completed":
+                    break
+
+            # Get structured result while still alive
+            operation_result = await client.get_operation_result(token)
+            assert operation_result.result is not None
+            assert operation_result.result.structuredContent is not None
+            structured_data = operation_result.result.structuredContent
+            assert structured_data["status"] == "success"
+            assert structured_data["data"] == "Processed: test"
+            assert structured_data["count"] == 42
+
+            # Wait for keep_alive to expire (1 second + buffer)
+            await asyncio.sleep(1.2)
+
+            # Operation should now be expired/unavailable - validation should fail gracefully
+            with pytest.raises(Exception):  # Should raise error for expired operation
+                await client.get_operation_result(token)
+
 
 class TestServerResources:
     @pytest.mark.anyio

tests/server/fastmcp/test_tool_manager.py

@@ -633,6 +633,84 @@ def get_user() -> UserOutput:
         }
         assert tool.output_schema == expected_schema
 
+    def test_tool_meta_property(self):
+        """Test that Tool.meta property works correctly."""
+
+        def double_number(n: int) -> int:
+            """Double a number."""
+            return 10
+
+        manager = ToolManager()
+        tool = manager.add_tool(double_number, meta={"foo": "bar"})
+
+        # Test that meta is populated
+        expected_meta = {
+            "foo": "bar",
+        }
+        assert tool.meta == expected_meta
+
+    def test_tool_keep_alive_property_sync(self):
+        """Test that keep_alive property works correctly with sync-only tools."""
+
+        def double_number(n: int) -> int:
+            """Double a number."""
+            return 10
+
+        manager = ToolManager()
+
+        # Should raise error when keep_alive is used on sync-only tool
+        with pytest.raises(ValueError, match="keep_alive parameter can only be used with async-compatible tools"):
+            manager.add_tool(double_number, invocation_modes=["sync"], keep_alive=1)
+
+    def test_tool_keep_alive_property_async(self):
+        """Test that keep_alive property works correctly with async-only tools."""
+
+        def double_number(n: int) -> int:
+            """Double a number."""
+            return 10
+
+        manager = ToolManager()
+        tool = manager.add_tool(double_number, invocation_modes=["async"], keep_alive=1)
+
+        # Test that meta is populated and has the keepalive stashed in it
+        expected_meta = {
+            "_keep_alive": 1,
+        }
+        assert tool.meta == expected_meta
+
+    def test_tool_keep_alive_property_hybrid(self):
+        """Test that keep_alive property works correctly with hybrid sync/async tools."""
+
+        def double_number(n: int) -> int:
+            """Double a number."""
+            return 10
+
+        manager = ToolManager()
+        tool = manager.add_tool(double_number, invocation_modes=["sync", "async"], keep_alive=1)
+
+        # Test that meta is populated and has the keepalive stashed in it
+        expected_meta = {
+            "_keep_alive": 1,
+        }
+        assert tool.meta == expected_meta
+
+    def test_tool_keep_alive_property_meta(self):
+        """Test that keep_alive property works correctly with existing metadata defined."""
+
+        def double_number(n: int) -> int:
+            """Double a number."""
+            return 10
+
+        manager = ToolManager()
+        tool = manager.add_tool(double_number, invocation_modes=["async"], keep_alive=1, meta={"foo": "bar"})
+
+        # Test that meta is populated and has the keepalive stashed in it
+        expected_meta = {
+            "foo": "bar",
+            "_keep_alive": 1,
+        }
+        assert tool.meta == expected_meta
+
     @pytest.mark.anyio
     async def test_tool_with_dict_str_any_output(self):
         """Test tool with dict[str, Any] return type."""