Implement optoken to tool name map on client end for validation

Author: LucaButBoring

src/mcp/client/session.py

@@ -9,6 +9,7 @@
 from pydantic import AnyUrl, TypeAdapter
 
 import mcp.types as types
+from mcp.shared.async_operations import ClientAsyncOperationManager
 from mcp.shared.context import RequestContext
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
@@ -136,6 +137,7 @@ def __init__(
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
+        self._operation_manager = ClientAsyncOperationManager()
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
@@ -174,8 +176,15 @@ async def initialize(self) -> types.InitializeResult:
 
         await self.send_notification(types.ClientNotification(types.InitializedNotification()))
 
+        # Start cleanup task for operations
+        await self._operation_manager.start_cleanup_task()
+
         return result
 
+    async def close(self) -> None:
+        """Clean up resources."""
+        await self._operation_manager.stop_cleanup_task()
+
     async def send_ping(self) -> types.EmptyResult:
         """Send a ping request."""
         return await self.send_request(
@@ -305,7 +314,14 @@ async def call_tool(
         )
 
         if not result.isError:
-            await self._validate_tool_result(name, result)
+            # Track operation for async operations
+            if result.operation is not None:
+                self._operation_manager.track_operation(
+                    result.operation.token, name, result.operation.keepAlive or 3600
+                )
+                logger.debug(f"Tracking operation for token: {result.operation.token}")
+            else:
+                await self._validate_tool_result(name, result)
 
         return result
 
@@ -336,7 +352,7 @@ async def get_operation_result(self, token: str) -> types.GetOperationPayloadRes
         Returns:
             The final tool result
         """
-        return await self.send_request(
+        result = await self.send_request(
             types.ClientRequest(
                 types.GetOperationPayloadRequest(
                     params=types.GetOperationPayloadParams(token=token),
@@ -345,7 +361,18 @@ async def get_operation_result(self, token: str) -> types.GetOperationPayloadRes
             types.GetOperationPayloadResult,
         )
 
-    async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
+        # Validate using the stored tool name
+        if hasattr(result, "result") and result.result:
+            # Clean up expired operations first
+            self._operation_manager.cleanup_expired()
+
+            tool_name = self._operation_manager.get_tool_name(token)
+            await self._validate_tool_result(tool_name, result.result)
+            # Keep the operation for potential future retrievals
+
+        return result
+
+    async def _validate_tool_result(self, name: str | None, result: types.CallToolResult) -> None:
         """Validate the structured content of a tool result against its output schema."""
         if name not in self._tool_output_schemas:
             # refresh output schema cache
@@ -358,6 +385,7 @@ async def _validate_tool_result(self, name: str, result: types.CallToolResult) -
             logger.warning(f"Tool {name} not listed by server, cannot validate any structured content")
 
         if output_schema is not None:
+            logger.debug(f"Validating structured content for tool: {name}")
             if result.structuredContent is None:
                 raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
             try:

src/mcp/server/fastmcp/server.py

@@ -33,7 +33,6 @@
 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.async_operations import AsyncOperationManager
 from mcp.server.lowlevel.helper_types import ReadResourceContents
 from mcp.server.lowlevel.server import LifespanResultT
 from mcp.server.lowlevel.server import Server as MCPServer
@@ -44,6 +43,7 @@
 from mcp.server.streamable_http import EventStore
 from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
+from mcp.shared.async_operations import ServerAsyncOperationManager
 from mcp.shared.context import LifespanContextT, RequestContext, RequestT
 from mcp.types import (
     NEXT_PROTOCOL_VERSION,
@@ -138,7 +138,7 @@ def __init__(
         token_verifier: TokenVerifier | None = None,
         event_store: EventStore | None = None,
         *,
-        async_operations: AsyncOperationManager | None = None,
+        async_operations: ServerAsyncOperationManager | None = None,
         tools: list[Tool] | None = None,
         debug: bool = False,
         log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO",
@@ -178,7 +178,7 @@ def __init__(
             transport_security=transport_security,
         )
 
-        self._async_operations = async_operations or AsyncOperationManager()
+        self._async_operations = async_operations or ServerAsyncOperationManager()
 
         self._mcp_server = MCPServer(
             name=name or "FastMCP",

src/mcp/server/lowlevel/server.py

@@ -83,10 +83,10 @@ async def main():
 from typing_extensions import TypeVar
 
 import mcp.types as types
-from mcp.server.lowlevel.async_operations import AsyncOperation, AsyncOperationManager
 from mcp.server.lowlevel.helper_types import ReadResourceContents
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
+from mcp.shared.async_operations import ServerAsyncOperation, ServerAsyncOperationManager
 from mcp.shared.context import RequestContext
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
@@ -138,7 +138,7 @@ def __init__(
         name: str,
         version: str | None = None,
         instructions: str | None = None,
-        async_operations: AsyncOperationManager | None = None,
+        async_operations: ServerAsyncOperationManager | None = None,
         lifespan: Callable[
             [Server[LifespanResultT, RequestT]],
             AbstractAsyncContextManager[LifespanResultT],
@@ -148,7 +148,7 @@ def __init__(
         self.version = version
         self.instructions = instructions
         self.lifespan = lifespan
-        self.async_operations = async_operations or AsyncOperationManager()
+        self.async_operations = async_operations or ServerAsyncOperationManager()
         # Track request ID to operation token mapping for cancellation
         self._request_to_operation: dict[RequestId, str] = {}
         self.request_handlers: dict[type, Callable[..., Awaitable[types.ServerResult]]] = {
@@ -469,11 +469,9 @@ async def handler(req: types.CallToolRequest):
                     # Check for async execution
                     if tool and self.async_operations and self._should_execute_async(tool):
                         # Create async operation
-                        session_id = f"session_{id(self.request_context.session)}"
                         operation = self.async_operations.create_operation(
                             tool_name=tool_name,
                             arguments=arguments,
-                            session_id=session_id,
                         )
                         logger.debug(f"Created async operation with token: {operation.token}")
 
@@ -627,7 +625,7 @@ async def handler(req: types.CompleteRequest):
 
         return decorator
 
-    def _validate_operation_token(self, token: str) -> AsyncOperation:
+    def _validate_operation_token(self, token: str) -> ServerAsyncOperation:
         """Validate operation token and return operation if valid."""
         operation = self.async_operations.get_operation(token)
         if not operation: