Operation token plumbing to support async elicitation/sampling

Author: LucaButBoring

examples/snippets/clients/async_tools_client.py

@@ -13,6 +13,7 @@
 
 from mcp import ClientSession, StdioServerParameters, types
 from mcp.client.stdio import stdio_client
+from mcp.shared.context import RequestContext
 
 # Create server parameters for stdio connection
 server_params = StdioServerParameters(
@@ -22,6 +23,22 @@
 )
 
 
+async def elicitation_callback(context: RequestContext[ClientSession, None], params: types.ElicitRequestParams):
+    """Handle elicitation requests from the server."""
+    if "data_migration" in params.message:
+        return types.ElicitResult(
+            action="accept",
+            content={"continue_processing": True, "priority_level": "normal"},
+        )
+    else:
+        return types.ElicitResult(action="decline")
+
+
+async def logging_callback(params: types.LoggingMessageNotificationParams):
+    """Handle logging messages from the server."""
+    print(f"Server log: {params.data}", file=sys.stderr)
+
+
 async def demonstrate_sync_tool(session: ClientSession):
     """Demonstrate calling a synchronous tool."""
     print("\n=== Synchronous Tool Demo ===")
@@ -174,6 +191,37 @@ async def demonstrate_data_processing(session: ClientSession):
             await asyncio.sleep(0.8)
 
 
+async def demonstrate_elicitation(session: ClientSession):
+    """Demonstrate async elicitation tool."""
+    print("\n=== Async Elicitation Demo ===")
+
+    result = await session.call_tool("async_elicitation_tool", arguments={"operation": "data_migration"})
+
+    if result.operation:
+        token = result.operation.token
+        print(f"Elicitation operation started with token: {token}")
+
+        # Poll for completion
+        while True:
+            status = await session.get_operation_status(token)
+            print(f"Status: {status.status}")
+
+            if status.status == "completed":
+                final_result = await session.get_operation_result(token)
+                for content in final_result.result.content:
+                    if isinstance(content, types.TextContent):
+                        print(f"Elicitation result: {content.text}")
+                break
+            elif status.status == "failed":
+                print(f"Elicitation failed: {status.error}")
+                break
+            elif status.status in ("canceled", "unknown"):
+                print(f"Elicitation ended with status: {status.status}")
+                break
+
+            await asyncio.sleep(0.5)
+
+
 async def run():
     """Run all async tool demonstrations."""
     # Determine protocol version from command line
@@ -189,7 +237,13 @@ async def run():
 
     async with stdio_client(server_params) as (read, write):
         # Use configured protocol version
-        async with ClientSession(read, write, protocol_version=protocol_version) as session:
+        async with ClientSession(
+            read,
+            write,
+            protocol_version=protocol_version,
+            elicitation_callback=elicitation_callback,
+            logging_callback=logging_callback,
+        ) as session:
             # Initialize the connection
             await session.initialize()
 
@@ -206,6 +260,7 @@ async def run():
             await demonstrate_async_tool(session)
             await demonstrate_batch_processing(session)
             await demonstrate_data_processing(session)
+            await demonstrate_elicitation(session)
 
             print("\n=== All demonstrations complete! ===")
 

examples/snippets/servers/async_tools.py

@@ -7,12 +7,61 @@
 
 import asyncio
 
+from pydantic import BaseModel, Field
+
 from mcp.server.fastmcp import Context, FastMCP
 
 # Create an MCP server with async operations support
 mcp = FastMCP("Async Tools Demo")
 
 
+class UserPreferences(BaseModel):
+    """Schema for collecting user preferences."""
+
+    continue_processing: bool = Field(description="Should we continue with the operation?")
+    priority_level: str = Field(
+        default="normal",
+        description="Priority level: low, normal, high",
+    )
+
+
+@mcp.tool(invocation_modes=["async"])
+async def async_elicitation_tool(operation: str, ctx: Context) -> str:  # type: ignore[type-arg]
+    """An async tool that uses elicitation to get user input."""
+    await ctx.info(f"Starting operation: {operation}")
+
+    # Simulate some initial processing
+    await asyncio.sleep(0.5)
+    await ctx.report_progress(0.3, 1.0, "Initial processing complete")
+
+    await ctx.debug("About to call elicit")
+    try:
+        # Ask user for preferences
+        result = await ctx.elicit(
+            message=f"Operation '{operation}' requires user input. How should we proceed?",
+            schema=UserPreferences,
+        )
+        await ctx.debug(f"Elicit result: {result}")
+    except Exception as e:
+        await ctx.error(f"Elicitation failed: {e}")
+        raise
+
+    if result.action == "accept" and result.data:
+        if result.data.continue_processing:
+            await ctx.info(f"Continuing with {result.data.priority_level} priority")
+            # Simulate processing based on user choice
+            processing_time = {"low": 0.5, "normal": 1.0, "high": 1.5}.get(result.data.priority_level, 1.0)
+            await asyncio.sleep(processing_time)
+            await ctx.report_progress(1.0, 1.0, "Operation complete")
+            return f"Operation '{operation}' completed successfully with {result.data.priority_level} priority"
+        else:
+            await ctx.warning("User chose not to continue")
+            return f"Operation '{operation}' cancelled by user"
+    else:
+        await ctx.error("User declined or cancelled the operation")
+        return f"Operation '{operation}' aborted"
+
+
 @mcp.tool()
 def sync_tool(x: int) -> str:
     """An implicitly-synchronous tool."""

src/mcp/client/session.py

@@ -466,6 +466,7 @@ async def send_roots_list_changed(self) -> None:
     async def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
         ctx = RequestContext[ClientSession, Any](
             request_id=responder.request_id,
+            operation_token=responder.operation.token if responder.operation is not None else None,
             meta=responder.request_meta,
             session=self,
             lifespan_context=None,
@@ -475,12 +476,36 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
             case types.CreateMessageRequest(params=params):
                 with responder:
                     response = await self._sampling_callback(ctx, params)
+                    if isinstance(response, types.CreateMessageResult):
+                        response.operation_props = (
+                            types.Operation(token=responder.operation.token)
+                            if responder.operation is not None
+                            else None
+                        )
+                    else:
+                        response.operation = (
+                            types.Operation(token=responder.operation.token)
+                            if responder.operation is not None
+                            else None
+                        )
                     client_response = ClientResponse.validate_python(response)
                     await responder.respond(client_response)
 
             case types.ElicitRequest(params=params):
                 with responder:
                     response = await self._elicitation_callback(ctx, params)
+                    if isinstance(response, types.ElicitResult):
+                        response.operation_props = (
+                            types.Operation(token=responder.operation.token)
+                            if responder.operation is not None
+                            else None
+                        )
+                    else:
+                        response.operation = (
+                            types.Operation(token=responder.operation.token)
+                            if responder.operation is not None
+                            else None
+                        )
                     client_response = ClientResponse.validate_python(response)
                     await responder.respond(client_response)
 

src/mcp/server/elicitation.py

@@ -78,6 +78,7 @@ async def elicit_with_validation(
     message: str,
     schema: type[ElicitSchemaModelT],
     related_request_id: RequestId | None = None,
+    related_operation_token: str | None = None,
 ) -> ElicitationResult[ElicitSchemaModelT]:
     """Elicit information from the client/user with schema validation.
 
@@ -96,6 +97,7 @@ async def elicit_with_validation(
         message=message,
         requestedSchema=json_schema,
         related_request_id=related_request_id,
+        related_operation_token=related_operation_token,
     )
 
     if result.action == "accept" and result.content is not None:

src/mcp/server/fastmcp/server.py

@@ -1213,6 +1213,8 @@ async def report_progress(self, progress: float, total: float | None = None, mes
             progress=progress,
             total=total,
             message=message,
+            related_request_id=self.request_id,
+            related_operation_token=self.request_context.operation_token,
         )
 
     async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]:
@@ -1255,7 +1257,11 @@ async def elicit(
         """
 
         return await elicit_with_validation(
-            session=self.request_context.session, message=message, schema=schema, related_request_id=self.request_id
+            session=self.request_context.session,
+            message=message,
+            schema=schema,
+            related_request_id=self.request_id,
+            related_operation_token=self.request_context.operation_token,
         )
 
     async def log(

src/mcp/server/lowlevel/server.py

@@ -91,7 +91,7 @@ async def main():
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.session import RequestResponder
-from mcp.types import RequestId
+from mcp.types import Operation, RequestId
 
 logger = logging.getLogger(__name__)
 
@@ -478,6 +478,10 @@ async def handler(req: types.CallToolRequest):
                         )
                         logger.debug(f"Created async operation with token: {operation.token}")
 
+                        ctx = self.request_context
+                        ctx.operation_token = operation.token
+                        request_ctx.set(ctx)
+
                         # Start async execution in background
                         async def execute_async():
                             try:
@@ -560,6 +564,9 @@ def _process_tool_result(
             content=list(unstructured_content),
             structuredContent=maybe_structured_content,
             isError=False,
+            _operation=Operation(token=self.request_context.operation_token)
+            if self.request_context and self.request_context.operation_token
+            else None,
         )
 
     def _should_execute_async(self, tool: types.Tool) -> bool:
@@ -720,9 +727,7 @@ def send_request_for_operation(self, token: str, request: types.ServerRequest) -
             # Add operation token to request
             if hasattr(request.root, "params") and request.root.params is not None:
                 if not hasattr(request.root.params, "operation") or request.root.params.operation is None:
-                    # Create operation field if it doesn't exist
-                    operation_data = types.RequestParams.Operation(token=token)
-                    request.root.params.operation = operation_data
+                    request.root.params.operation = Operation(token=token)
             logger.debug(f"Marked operation {token} as input_required and added to request")
 
     def send_notification_for_operation(self, token: str, notification: types.ServerNotification) -> None:
@@ -732,9 +737,7 @@ def send_notification_for_operation(self, token: str, notification: types.Server
             # Add operation token to notification
             if hasattr(notification.root, "params") and notification.root.params is not None:
                 if not hasattr(notification.root.params, "operation") or notification.root.params.operation is None:
-                    # Create operation field if it doesn't exist
-                    operation_data = types.NotificationParams.Operation(token=token)
-                    notification.root.params.operation = operation_data
+                    notification.root.params.operation = Operation(token=token)
             logger.debug(f"Marked operation {token} as input_required and added to notification")
 
     def complete_request_for_operation(self, token: str) -> None:
@@ -833,25 +836,16 @@ async def _handle_request(
                 # app.get_request_context()
                 context_token = request_ctx.set(
                     RequestContext(
-                        message.request_id,
-                        message.request_meta,
-                        session,
-                        lifespan_context,
+                        request_id=message.request_id,
+                        operation_token=message.operation.token if message.operation else None,
+                        meta=message.request_meta,
+                        session=session,
+                        lifespan_context=lifespan_context,
                         request=request_data,
                     )
                 )
                 response = await handler(req)
 
-                # Handle operation token in response (for input_required operations)
-                if (
-                    hasattr(req, "params")
-                    and req.params is not None
-                    and hasattr(req.params, "operation")
-                    and req.params.operation is not None
-                ):
-                    operation_token = req.params.operation.token
-                    self.complete_request_for_operation(operation_token)
-
                 # Track async operations for cancellation
                 if isinstance(req, types.CallToolRequest):
                     result = response.root