Add snippet for async tool with keepalive

Author: LucaButBoring

examples/snippets/servers/async_tools.py

@@ -135,5 +135,27 @@ async def batch_operation_tool(items: list[str], ctx: Context) -> list[str]:  #
     return results
 
 
+@mcp.tool(invocation_modes=["async"], keep_alive=1800)
+async def long_running_task(task_name: str, ctx: Context) -> str:  # type: ignore[type-arg]
+    """A long-running task with custom keep_alive duration."""
+    await ctx.info(f"Starting long-running task: {task_name}")
+
+    # Simulate extended processing
+    await asyncio.sleep(2)
+    await ctx.report_progress(0.5, 1.0, "Halfway through processing")
+    await asyncio.sleep(2)
+
+    await ctx.info(f"Task '{task_name}' completed successfully")
+    return f"Long-running task '{task_name}' finished with 30-minute keep_alive"
+
+
+@mcp.tool(invocation_modes=["async"], keep_alive=2)
+async def quick_expiry_task(message: str, ctx: Context) -> str:  # type: ignore[type-arg]
+    """A task with very short keep_alive for testing expiry."""
+    await ctx.info(f"Quick task starting: {message}")
+    await asyncio.sleep(1)
+    return f"Quick task completed: {message} (expires in 2 seconds)"
+
+
 if __name__ == "__main__":
     mcp.run()

tests/server/fastmcp/test_integration.py

@@ -10,6 +10,7 @@
 # pyright: reportUnknownVariableType=false
 # pyright: reportUnknownArgumentType=false
 
+import asyncio
 import json
 import multiprocessing
 import socket
@@ -702,9 +703,6 @@ async def test_async_tools(server_transport: str, server_url: str) -> None:
             assert async_result.operation is not None
             token = async_result.operation.token
 
-            # Poll for completion
-            import asyncio
-
             while True:
                 status = await session.get_operation_status(token)
                 if status.status == "completed":
@@ -717,14 +715,51 @@ async def test_async_tools(server_transport: str, server_url: str) -> None:
                     break
                 elif status.status == "failed":
                     pytest.fail(f"Async operation failed: {status.error}")
-                await asyncio.sleep(0.01)
 
             # Test hybrid tool (should work as sync by default)
             hybrid_result = await session.call_tool("hybrid_tool", {"message": "hello"})
             assert len(hybrid_result.content) == 1
             assert isinstance(hybrid_result.content[0], TextContent)
             assert "Hybrid result: HELLO" in hybrid_result.content[0].text
 
+            # Test long-running task with custom keep_alive
+            long_task_result = await session.call_tool("long_running_task", {"task_name": "test_task"})
+            assert long_task_result.operation is not None
+            long_token = long_task_result.operation.token
+
+            while True:
+                status = await session.get_operation_status(long_token)
+                if status.status == "completed":
+                    final_result = await session.get_operation_result(long_token)
+                    assert not final_result.result.isError
+                    assert len(final_result.result.content) == 1
+                    content = final_result.result.content[0]
+                    assert isinstance(content, TextContent)
+                    assert "Long-running task 'test_task' finished with 30-minute keep_alive" in content.text
+                    break
+                elif status.status == "failed":
+                    pytest.fail(f"Long-running task failed: {status.error}")
+
+            # Test quick expiry task (should complete then expire)
+            quick_result = await session.call_tool("quick_expiry_task", {"message": "test_expiry"})
+            assert quick_result.operation is not None
+            quick_token = quick_result.operation.token
+
+            # Wait for completion
+            while True:
+                status = await session.get_operation_status(quick_token)
+                if status.status == "completed":
+                    break
+                elif status.status == "failed":
+                    pytest.fail(f"Quick task failed: {status.error}")
+
+            # Wait for expiry (2 seconds + buffer)
+            await asyncio.sleep(3)
+
+            # Should now be expired
+            with pytest.raises(Exception):  # Should raise error when trying to access expired operation
+                await session.get_operation_result(quick_token)
+
 
 # Test async tools example with legacy protocol
 @pytest.mark.anyio