work in progress towards async protocol - simple call and get work

davemssavage · davemssavage · commit 9b1fd8193403 · 2025-06-01T06:15:47.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = [
     "sse-starlette>=1.6.1",
     "pydantic-settings>=2.5.2",
     "uvicorn>=0.23.1; sys_platform != 'emscripten'",
+    "cachetools==6",
 ]
 
 [project.optional-dependencies]
diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
@@ -147,7 +147,7 @@ def __init__(
         tools: list[Tool] | None = None,
         **settings: Any,
     ):
-        self.settings = Settings(**settings) # type: ignore
+        self.settings = Settings(**settings)  # type: ignore
 
         self._mcp_server = MCPServer(
             name=name or "FastMCP",
diff --git a/src/mcp/server/lowlevel/result_cache.py b/src/mcp/server/lowlevel/result_cache.py
@@ -0,0 +1,122 @@
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from time import time
+from typing import Any
+from uuid import uuid4
+
+from anyio import Lock, create_task_group, move_on_after
+from anyio.abc import TaskGroup
+from cachetools import TTLCache
+
+from mcp import types
+from mcp.shared.context import BaseSession, RequestContext, SessionT
+
+
+@dataclass
+class InProgress:
+    token: str
+    task_group: TaskGroup | None = None
+    sessions: list[BaseSession[Any, Any, Any, Any, Any]] = field(
+        default_factory=lambda: []
+    )
+
+
+class ResultCache:
+    _in_progress: dict[types.AsyncToken, InProgress]
+
+    def __init__(self, max_size: int, max_keep_alive: int):
+        self._max_size = max_size
+        self._max_keep_alive = max_keep_alive
+        self._result_cache = TTLCache[types.AsyncToken, types.CallToolResult](
+            self._max_size, self._max_keep_alive
+        )
+        self._in_progress = {}
+        self._lock = Lock()
+
+    async def add_call(
+        self,
+        call: Callable[[types.CallToolRequest], Awaitable[types.ServerResult]],
+        req: types.CallToolAsyncRequest,
+        ctx: RequestContext[SessionT, Any, Any],
+    ) -> types.CallToolAsyncResult:
+        in_progress = await self._new_in_progress()
+        timeout = min(
+            req.params.keepAlive or self._max_keep_alive, self._max_keep_alive
+        )
+
+        async def call_tool():
+            with move_on_after(timeout) as scope:
+                result = await call(
+                    types.CallToolRequest(
+                        method="tools/call",
+                        params=types.CallToolRequestParams(
+                            name=req.params.name, arguments=req.params.arguments
+                        ),
+                    )
+                )
+            if not scope.cancel_called:
+                async with self._lock:
+                    assert type(result.root) is types.CallToolResult
+                    self._result_cache[in_progress.token] = result.root
+
+        async with create_task_group() as tg:
+            tg.start_soon(call_tool)
+            in_progress.task_group = tg
+            in_progress.sessions.append(ctx.session)
+            result = types.CallToolAsyncResult(
+                token=in_progress.token,
+                recieved=round(time()),
+                keepAlive=timeout,
+                accepted=True,
+            )
+            return result
+
+    async def join_call(
+        self,
+        req: types.JoinCallToolAsyncRequest,
+        ctx: RequestContext[SessionT, Any, Any],
+    ) -> types.CallToolAsyncResult:
+        async with self._lock:
+            in_progress = self._in_progress.get(req.params.token)
+            if in_progress is None:
+                # TODO consider creating new token to allow client 
+                # to get message describing why it wasn't accepted
+                return types.CallToolAsyncResult(accepted=False)
+            else:
+                in_progress.sessions.append(ctx.session)
+                return types.CallToolAsyncResult(accepted=True)
+
+        return
+
+    async def cancel(self, notification: types.CancelToolAsyncNotification) -> None:
+        async with self._lock:
+            in_progress = self._in_progress.get(notification.params.token)
+            if in_progress is not None and in_progress.task_group is not None:
+                in_progress.task_group.cancel_scope.cancel()
+                del self._in_progress[notification.params.token]
+
+    async def get_result(self, req: types.GetToolAsyncResultRequest):
+        async with self._lock:
+            in_progress = self._in_progress.get(req.params.token)
+            if in_progress is None:
+                return types.CallToolResult(
+                    content=[
+                        types.TextContent(type="text", text="Unknown progress token")
+                    ],
+                    isError=True,
+                )
+            else:
+                result = self._result_cache.get(in_progress.token)
+                if result is None:
+                    return types.CallToolResult(content=[], isPending=True)
+                else:
+                    return result
+
+    async def _new_in_progress(self) -> InProgress:
+        async with self._lock:
+            while True:
+                token = str(uuid4())
+                if token not in self._in_progress:
+                    new_in_progress = InProgress(token)
+                    self._in_progress[token] = new_in_progress
+                    return new_in_progress
diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
@@ -81,6 +81,7 @@ async def main():
 
 import mcp.types as types
 from mcp.server.lowlevel.helper_types import ReadResourceContents
+from mcp.server.lowlevel.result_cache import ResultCache
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
 from mcp.server.stdio import stdio_server as stdio_server
@@ -135,6 +136,8 @@ def __init__(
             [Server[LifespanResultT, RequestT]],
             AbstractAsyncContextManager[LifespanResultT],
         ] = lifespan,
+        max_cache_size: int = 1000,
+        max_cache_ttl: int = 60,
     ):
         self.name = name
         self.version = version
@@ -145,6 +148,7 @@ def __init__(
         ] = {
             types.PingRequest: _ping_handler,
         }
+        self.result_cache = ResultCache(max_cache_size, max_cache_ttl)
         self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
         self.notification_options = NotificationOptions()
         logger.debug(f"Initializing server '{name}'")
@@ -426,7 +430,32 @@ async def handler(req: types.CallToolRequest):
                         )
                     )
 
+            async def async_call_handler(req: types.CallToolAsyncRequest):
+                ctx = request_ctx.get()
+                result = await self.result_cache.add_call(handler, req, ctx)
+                return types.ServerResult(result)
+
+            async def async_join_handler(req: types.JoinCallToolAsyncRequest):
+                ctx = request_ctx.get()
+                result = await self.result_cache.join_call(req, ctx)
+                return types.ServerResult(result)
+
+            async def async_cancel_handler(req: types.CancelToolAsyncNotification):
+                await self.result_cache.cancel(req)
+
+            async def async_result_handler(req: types.GetToolAsyncResultRequest):
+                result = await self.result_cache.get_result(req)
+                return types.ServerResult(result)
+
             self.request_handlers[types.CallToolRequest] = handler
+            self.request_handlers[types.CallToolAsyncRequest] = async_call_handler
+            self.request_handlers[types.JoinCallToolAsyncRequest] = async_join_handler
+            self.request_handlers[types.GetToolAsyncResultRequest] = (
+                async_result_handler
+            )
+            self.notification_handlers[types.CancelToolAsyncNotification] = (
+                async_cancel_handler
+            )
             return func
 
         return decorator
diff --git a/src/mcp/types.py b/src/mcp/types.py
@@ -799,6 +799,8 @@ class Tool(BaseModel):
     """A JSON Schema object defining the expected parameters for the tool."""
     annotations: ToolAnnotations | None = None
     """Optional additional tool information."""
+    preferAsync: bool | None = None
+    """Optional flag to suggest to client async calls should be preferred"""
     model_config = ConfigDict(extra="allow")
 
 
@@ -1227,6 +1229,7 @@ class ClientRequest(
         | CallToolRequest
         | CallToolAsyncRequest
         | JoinCallToolAsyncRequest
+        | GetToolAsyncResultRequest
         | ListToolsRequest
     ]
 ):
diff --git a/tests/shared/test_session.py b/tests/shared/test_session.py
@@ -129,6 +129,77 @@ async def make_request(client_session):
                 await ev_cancelled.wait()
 
 
+@pytest.mark.anyio
+async def test_request_async():
+    """Test that requests can be run asynchronously."""
+    # The tool is already registered in the fixture
+
+    ev_tool_called = anyio.Event()
+
+    # Start the request in a separate task so we can cancel it
+    def make_server() -> Server:
+        server = Server(name="TestSessionServer")
+
+        # Register the tool handler
+        @server.call_tool()
+        async def handle_call_tool(name: str, arguments: dict | None) -> list:
+            nonlocal ev_tool_called
+            if name == "async_tool":
+                ev_tool_called.set()
+                return [types.TextContent(type="text", text="test")]
+            raise ValueError(f"Unknown tool: {name}")
+
+        # Register the tool so it shows up in list_tools
+        @server.list_tools()
+        async def handle_list_tools() -> list[types.Tool]:
+            return [
+                types.Tool(
+                    name="async_tool",
+                    description="A tool that does things asynchronously",
+                    inputSchema={},
+                    preferAsync=True,
+                )
+            ]
+
+        return server
+
+    async def make_request(client_session: ClientSession):
+        return await client_session.send_request(
+            ClientRequest(
+                types.CallToolAsyncRequest(
+                    method="tools/async/call",
+                    params=types.CallToolAsyncRequestParams(
+                        name="async_tool", arguments={}
+                    ),
+                )
+            ),
+            types.CallToolAsyncResult,
+        )
+
+    async def get_result(client_session: ClientSession, async_token: types.AsyncToken):
+        return await client_session.send_request(
+            ClientRequest(
+                types.GetToolAsyncResultRequest(
+                    method="tools/async/get",
+                    params=types.GetToolAsyncResultRequestParams(token=async_token),
+                )
+            ),
+            types.CallToolResult,
+        )
+
+    async with create_connected_server_and_client_session(
+        make_server()
+    ) as client_session:
+        async_result = await make_request(client_session)
+        assert async_result is not None
+        assert async_result.token is not None
+        with anyio.fail_after(1):  # Timeout after 1 second
+            await ev_tool_called.wait()
+        result = await get_result(client_session, async_result.token)
+        assert type(result.content[0]) is types.TextContent
+        assert result.content[0].text == "test"
+
+
 @pytest.mark.anyio
 async def test_connection_closed():
     """
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,7 @@ dependencies = [`
`31`	`31`	`"sse-starlette>=1.6.1",`
`32`	`32`	`"pydantic-settings>=2.5.2",`
`33`	`33`	`"uvicorn>=0.23.1; sys_platform != 'emscripten'",`
	`34`	`+ "cachetools==6",`
`34`	`35`	`]`
`35`	`36`
`36`	`37`	`[project.optional-dependencies]`