diff --git a/examples/fastmcp/resume_tokens.py b/examples/fastmcp/resume_tokens.py new file mode 100644 index 000000000..ee1f1f9ea --- /dev/null +++ b/examples/fastmcp/resume_tokens.py @@ -0,0 +1,165 @@ +# resume_tokens.py +"""MCP Server demonstrating resume tokens.""" + +import logging +import asyncio +import uuid +from typing import Dict, Optional, Any, List +from pydantic import BaseModel + +from mcp.server import fastmcp + +# --- Basic Setup --- +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class OrderData(BaseModel): + total_amount: float + +class OrderResult(BaseModel): + order_id: str + total_amount: float + +class OrderProgress(BaseModel): + order_id: str + status: str + progress_percentage: int + current_step: str + result: Optional[OrderResult] = None + error: Optional[str] = None + +# Global steps definition +ORDER_STEPS: List[OrderProgress] = [ + OrderProgress(order_id="", status="processing", progress_percentage=15, current_step="Validating order data"), + OrderProgress(order_id="", status="processing", progress_percentage=30, current_step="Checking inventory"), + OrderProgress(order_id="", status="processing", progress_percentage=50, current_step="Processing payment"), + OrderProgress(order_id="", status="processing", progress_percentage=70, current_step="Allocating resources"), + OrderProgress(order_id="", status="processing", progress_percentage=85, current_step="Preparing shipment"), + OrderProgress(order_id="", status="completed", progress_percentage=100, current_step="Finalizing order") +] + +# Global storage for order operations (in production, use proper database/cache) +order_operations: Dict[str, OrderProgress] = {} + +# Map resume tokens to order IDs +resume_token_to_order_id: Dict[str, str] = {} +order_id_to_latest_token: Dict[str, str] = {} +# Store original order data for validation +order_id_to_order_data: Dict[str, OrderData] = {} + +# --- MCP Server Setup --- +mcp = fastmcp.FastMCP("OrderCreation", port=8090, stateless_http=True) + + +# --- Async Order Creation Workflow --- +async def fake_order_creation_workflow(order_id: str, order_data: OrderData) -> None: + """ + Simulated async order creation workflow using global steps. + Each step takes 10 seconds. + """ + for step in ORDER_STEPS: + # Update progress + order_operations[order_id].status = step.status + order_operations[order_id].progress_percentage = step.progress_percentage + order_operations[order_id].current_step = step.current_step + + # Set result if this is the completion step + if step.status == "completed": + order_operations[order_id].result = OrderResult( + order_id=order_id, + total_amount=order_data.total_amount + ) + + logger.info(f"Order {order_id}: {step.current_step} ({step.progress_percentage}%)") + + # Wait 10 seconds for each step + await asyncio.sleep(10) + + +def start_order_creation(order_data: OrderData) -> str: + """Start a new order creation process and return the order ID.""" + order_id = str(uuid.uuid4()) + + # Initialize progress tracking + progress = OrderProgress( + order_id=order_id, + status="started", + progress_percentage=0, + current_step="Initializing order creation" + ) + + # Store in global storage + order_operations[order_id] = progress + order_id_to_order_data[order_id] = order_data + + # Start the async workflow in a new task + asyncio.create_task(fake_order_creation_workflow(order_id, order_data)) + + return order_id + + +def generate_new_resume_token(order_id: str) -> str: + """Generate a new resume token for an order and update mappings.""" + new_token = str(uuid.uuid4()) + + # Clean up old token mapping if exists + old_token = order_id_to_latest_token.get(order_id) + if old_token and old_token in resume_token_to_order_id: + del resume_token_to_order_id[old_token] + + # Set up new mapping + resume_token_to_order_id[new_token] = order_id + order_id_to_latest_token[order_id] = new_token + + return new_token + +# --- Tool Definitions --- +@mcp.tool( + name="create_order", + description="Create an order in the database. Can start a new order or check status with resume token.", +) +def create_order(order_data: Optional[dict] = None, resume_token: Optional[str] = None) -> dict: + """ + Create an order in the database or check status of existing order. + + Args: + order_data: Dictionary containing order information (required for new orders) + resume_token: UUID token to check status of existing order (optional) + + Returns: + Dictionary with order status and optional resume token + """ + if resume_token: + # Get order ID from resume token + order_id = resume_token_to_order_id.get(resume_token) + + # Validate order data if provided + if order_data: + current_order_data = OrderData(**order_data) + original_order_data = order_id_to_order_data.get(order_id) + + if original_order_data and current_order_data.model_dump_json() != original_order_data.model_dump_json(): + return { + "error": "Order data mismatch. The provided order data does not match the original request.", + "order_id": order_id, + } + + else: + # Create new order + order_data_model = OrderData(**order_data) if order_data else OrderData(total_amount=0.0, items=[]) + order_id = start_order_creation(order_data_model) + + # Get progress and convert to dict + progress = order_operations[order_id] + result = progress.model_dump() + + # Add resume token if still in progress + if progress.status in ["started", "processing"]: + result["next_resume_token"] = generate_new_resume_token(order_id) + + return result + + +if __name__ == "__main__": + mcp.run(transport="streamable-http")