traia-iatp 0.1.29__py3-none-any.whl

traia_iatp/d402/fastmcp_middleware.py

@@ -0,0 +1,147 @@
+"""
+FastMCP-specific middleware for d402 payment protocol.
+
+This module contains middleware classes for standalone FastMCP package.
+These are NOT used with the official MCP SDK - see starlette_middleware.py instead.
+
+DEPRECATED: Use starlette_middleware.py for official MCP SDK.
+This file is kept for backwards compatibility with standalone fastmcp package.
+"""
+
+import logging
+import os
+from typing import Optional
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from .types import PaymentPayload
+from .encoding import safe_base64_decode
+from .facilitator import IATPSettlementFacilitator
+
+logger = logging.getLogger(__name__)
+
+# Import FastMCP middleware classes (only available with standalone fastmcp)
+try:
+    from fastmcp.server.middleware import Middleware, MiddlewareContext
+    FASTMCP_AVAILABLE = True
+except ImportError:
+    # FastMCP not available - this module won't work
+    logger.warning("Standalone fastmcp not available - fastmcp_middleware.py will not work")
+    logger.warning("Use starlette_middleware.py for official MCP SDK instead")
+    Middleware = None
+    MiddlewareContext = None
+    FASTMCP_AVAILABLE = False
+
+
+class D402PaymentRequiredException(Exception):
+    """Exception raised when payment is required (HTTP 402)."""
+    def __init__(self, payment_response: dict):
+        self.payment_response = payment_response
+        super().__init__("Payment required")
+
+
+if FASTMCP_AVAILABLE and Middleware is not None:
+    class D402MCPMiddleware(Middleware):
+        """
+        HTTP 402 Payment Required middleware for standalone FastMCP servers.
+        
+        DEPRECATED: Use D402PaymentMiddleware from starlette_middleware.py for official MCP SDK.
+        
+        This is kept for backwards compatibility with standalone fastmcp package.
+        """
+        
+        def __init__(
+            self,
+            server_address: str,
+            facilitator_url: Optional[str] = None,
+            facilitator_api_key: Optional[str] = None,
+            verify_signatures: bool = True,
+            testing_mode: bool = False,
+            requires_auth: bool = False,
+            internal_api_key_env_var: Optional[str] = None
+        ):
+            super().__init__()
+            
+            self.server_address = server_address
+            self.facilitator_url = facilitator_url or os.getenv("D402_FACILITATOR_URL", "https://facilitator.d402.net")
+            self.facilitator_api_key = facilitator_api_key or os.getenv("D402_FACILITATOR_API_KEY")
+            self.verify_signatures = verify_signatures
+            self.testing_mode = testing_mode or os.getenv("D402_TESTING_MODE", "false").lower() == "true"
+            self.requires_auth = requires_auth
+            self.internal_api_key_env_var = internal_api_key_env_var
+            
+            # Initialize facilitator
+            self.facilitator = None
+            if not self.testing_mode:
+                try:
+                    operator_private_key = os.getenv("MCP_OPERATOR_PRIVATE_KEY") or os.getenv("OPERATOR_PRIVATE_KEY")
+                    self.facilitator = IATPSettlementFacilitator(
+                        relayer_url=self.facilitator_url,
+                        relayer_api_key=self.facilitator_api_key,
+                        provider_operator_key=operator_private_key,
+                        facilitator_url=self.facilitator_url,
+                        facilitator_api_key=self.facilitator_api_key
+                    )
+                except Exception as e:
+                    logger.warning(f"Could not initialize facilitator: {e}")
+                    self.testing_mode = True
+            
+            logger.info(f"D402 MCP Middleware initialized (FastMCP - deprecated)")
+        
+        async def on_request(self, context: MiddlewareContext, call_next):
+            """FastMCP middleware on_request handler."""
+            # Full implementation would go here
+            # For now, just pass through
+            logger.warning("D402MCPMiddleware for standalone fastmcp is deprecated")
+            logger.warning("Use D402PaymentMiddleware from starlette_middleware.py instead")
+            return await call_next(context)
+    
+    
+    def create_d402_mcp_middleware(
+        server_address: Optional[str] = None,
+        facilitator_url: Optional[str] = None,
+        facilitator_api_key: Optional[str] = None,
+        verify_signatures: bool = True,
+        testing_mode: bool = False,
+        requires_auth: bool = False,
+        internal_api_key_env_var: Optional[str] = None
+    ) -> D402MCPMiddleware:
+        """
+        Create d402 middleware for standalone FastMCP.
+        
+        DEPRECATED: Use D402PaymentMiddleware from starlette_middleware.py for official MCP SDK.
+        """
+        server_address = (
+            server_address or 
+            os.getenv("SERVER_ADDRESS") or 
+            os.getenv("PAYMENT_ADDRESS") or 
+            os.getenv("EVM_ADDRESS")
+        )
+        if not server_address:
+            raise ValueError("server_address is required")
+        
+        return D402MCPMiddleware(
+            server_address=server_address,
+            facilitator_url=facilitator_url,
+            facilitator_api_key=facilitator_api_key,
+            verify_signatures=verify_signatures,
+            testing_mode=testing_mode,
+            requires_auth=requires_auth,
+            internal_api_key_env_var=internal_api_key_env_var
+        )
+
+else:
+    # FastMCP not available - create stub
+    class D402MCPMiddleware:
+        """Stub class when FastMCP not available."""
+        def __init__(self, *args, **kwargs):
+            raise ImportError("Standalone fastmcp not available. Use starlette_middleware.py for official MCP SDK.")
+    
+    def create_d402_mcp_middleware(*args, **kwargs):
+        """Stub function when FastMCP not available."""
+        raise ImportError("Standalone fastmcp not available. Use starlette_middleware.py for official MCP SDK.")
+
+
+__all__ = ["D402MCPMiddleware", "D402PaymentRequiredException", "create_d402_mcp_middleware"]
+

traia_iatp/d402/mcp_middleware.py

@@ -0,0 +1,434 @@
+"""D402 payment helpers for MCP servers with official SDK.
+
+This module provides decorators and helper functions for d402 payment protocol.
+Works with official MCP SDK (mcp.server.fastmcp).
+
+FastMCP Middleware classes are in fastmcp_middleware.py (deprecated).
+"""
+
+import logging
+import os
+from typing import Optional, Dict, Any, Callable
+from functools import wraps
+
+# Import Context from official SDK
+from mcp.server.fastmcp import Context
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from web3 import Web3
+
+from .types import PaymentPayload, Price, TokenAmount, TokenAsset, EIP712Domain, PaymentRequirements, d402PaymentRequiredResponse
+from .encoding import safe_base64_decode
+from .common import process_price_to_atomic_amount, d402_VERSION
+from .facilitator import IATPSettlementFacilitator
+
+
+class D402PaymentRequiredException(Exception):
+    """Exception raised when payment is required (HTTP 402)."""
+    def __init__(self, payment_response: Dict[str, Any]):
+        self.payment_response = payment_response
+        super().__init__("Payment required")
+
+logger = logging.getLogger(__name__)
+
+
+class EndpointPaymentInfo:
+    """Payment information for a specific endpoint."""
+    def __init__(
+        self,
+        settlement_token_address: str,
+        settlement_token_network: str,
+        payment_price_float: float,
+        payment_price_wei: str,
+        server_address: str
+    ):
+        self.settlement_token_address = settlement_token_address
+        self.settlement_token_network = settlement_token_network
+        self.payment_price_float = payment_price_float
+        self.payment_price_wei = payment_price_wei
+        self.server_address = server_address
+
+
+def get_active_api_key(context: Any) -> Optional[str]:
+    """
+    Get the API key to use for calling external APIs.
+    
+    Returns api_key_to_use which was set by:
+    1. D402PaymentMiddleware (in request.state)
+    2. @require_payment_for_tool decorator (copied to context.state)
+    
+    Priority:
+    1. context.state.api_key_to_use (set by decorator)
+    2. request.state.api_key_to_use (set by middleware)
+    
+    Args:
+        context: MCP context object
+        
+    Returns:
+        API key string (client's OR server's) if authorized, None otherwise
+        
+    Usage in tools:
+        api_key = get_active_api_key(context)
+        if api_key:
+            headers = {"Authorization": f"Bearer {api_key}"}
+    """
+    try:
+        # Check request.state (set by middleware)
+        # Context is a Pydantic model - we can't set arbitrary fields on it
+        # So we read directly from request.state where middleware stored it
+        logger.debug(f"get_active_api_key: Checking context type={type(context).__name__}")
+        logger.debug(f"  has request_context: {hasattr(context, 'request_context')}")
+        
+        if hasattr(context, 'request_context') and context.request_context:
+            logger.debug(f"  request_context exists")
+            if hasattr(context.request_context, 'request') and context.request_context.request:
+                request = context.request_context.request
+                logger.debug(f"  request exists, has state: {hasattr(request, 'state')}")
+                if hasattr(request, 'state'):
+                    api_key = getattr(request.state, 'api_key_to_use', None)
+                    logger.debug(f"  api_key_to_use: {api_key[:10] if api_key else None}")
+                    if api_key:
+                        return api_key
+        
+        logger.warning(f"get_active_api_key: Could not find api_key_to_use in request.state")
+            
+    except Exception as e:
+        logger.error(f"get_active_api_key error: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+    
+    return None
+
+
+async def settle_payment(
+    context: Any,
+    endpoint_info: EndpointPaymentInfo,
+    output_data: Any,
+    middleware: Optional[Any] = None  # Middleware instance (Starlette or FastMCP)
+) -> bool:
+    """
+    Settle a payment after successful API call with output hash attestation.
+    
+    Complete 402 settlement flow:
+    1. Hash the output data (result returned to client)
+    2. Provider signs over output_hash + consumer_request
+    3. Submit to facilitator with proof of service completion
+    4. Facilitator submits to IATP Settlement Layer on-chain
+    
+    This should be called AFTER the tool successfully processes the request
+    to submit the payment settlement to the facilitator/blockchain.
+    
+    Args:
+        context: MCP context (contains payment_payload)
+        endpoint_info: Endpoint payment requirements
+        output_data: The actual output/result being returned to client (will be hashed)
+        middleware: Optional D402MCPMiddleware instance (for facilitator access)
+        
+    Returns:
+        bool: True if settlement submitted successfully
+        
+    Usage in tools (for production settlement):
+        # Execute API call
+        response = requests.get(url, headers={"Authorization": f"Bearer {api_key}"})
+        result = response.json()
+        
+        # Settle payment with output hash
+        if context.state.payment_payload:
+            await settle_payment(context, endpoint_payment, output_data=result, middleware=...)
+        
+        return result
+    """
+    try:
+        payment_payload = getattr(context.state, 'payment_payload', None) if hasattr(context, 'state') else None
+        if not payment_payload:
+            logger.debug("No payment to settle (authenticated mode)")
+            return True  # Not an error - client used their own API key
+        
+        # Skip settlement in testing mode
+        if middleware and middleware.testing_mode:
+            logger.info("⚠️  Testing mode: Skipping payment settlement")
+            return True
+        
+        # Step 1: Hash the output data (proof of service completion)
+        import json
+        from web3 import Web3
+        
+        logger.info("🔐 Starting payment settlement process...")
+        
+        # Serialize output to JSON and hash it
+        output_json = json.dumps(output_data, sort_keys=True, separators=(',', ':'))
+        output_hash = Web3.keccak(text=output_json).hex()
+        logger.info(f"📊 Output data serialized: {len(output_json)} bytes")
+        logger.info(f"🔑 Output hash calculated: {output_hash}")
+        logger.info(f"   First 1000 chars of output: {output_json[:1000]}")
+        
+        # Step 2: Get payment_uuid from context (from facilitator verify response)
+        # The payment_uuid is the primary payment identifier from the facilitator
+        # It was set in verify_endpoint_payment() after facilitator.verify() returned it
+        payment_uuid = None
+        if hasattr(context, 'state') and hasattr(context.state, 'payment_uuid'):
+            payment_uuid = context.state.payment_uuid
+        
+        if not payment_uuid:
+            logger.warning("No payment_uuid found in context - payment may not have been verified via facilitator")
+        
+        # Step 3: Get facilitator fee from context (set by verify response)
+        facilitator_fee_percent = 250  # Default
+        if hasattr(context, 'state') and hasattr(context.state, 'facilitator_fee_percent'):
+            facilitator_fee_percent = context.state.facilitator_fee_percent
+        
+        # Step 4: Create PaymentRequirements for this endpoint
+        # Include output_hash, payment_uuid, and facilitatorFeePercent in extra data
+        extra_data = {
+            "output_hash": output_hash,
+            "facilitator_fee_percent": facilitator_fee_percent
+        }
+        if payment_uuid:
+            extra_data["payment_uuid"] = payment_uuid
+        
+        payment_requirements = PaymentRequirements(
+            scheme="exact",
+            network=endpoint_info.settlement_token_network,
+            pay_to=endpoint_info.server_address,
+            max_amount_required=endpoint_info.payment_price_wei,
+            max_timeout_seconds=300,
+            description=f"Service completed - output_hash: {output_hash}",
+            resource="",
+            mime_type="application/json",
+            asset=endpoint_info.settlement_token_address,
+            extra=extra_data  # Include output hash, payment_uuid, and facilitatorFeePercent
+        )
+        
+        # Step 4: Settle via facilitator
+        # The facilitator will:
+        # - Create provider attestation signing over the consumer's request + output_hash
+        # - Submit to relayer with proof of service completion
+        # - Relayer submits to IATPSettlementLayer on-chain
+        if middleware and middleware.facilitator:
+            try:
+                logger.info(f"📤 Submitting settlement to facilitator...")
+                logger.info(f"   Payment UUID: {payment_uuid if payment_uuid else 'N/A'}")
+                logger.info(f"   Output hash: {output_hash}")
+                logger.info(f"   Amount: {endpoint_info.payment_price_wei} wei")
+                
+                settle_result = await middleware.facilitator.settle(payment_payload, payment_requirements)
+                if settle_result.success:
+                    logger.info(f"✅ Payment SETTLEMENT COMPLETE:")
+                    logger.info(f"   Transaction: {settle_result.transaction}")
+                    logger.info(f"   Network: {settle_result.network}")
+                    logger.info(f"   Payer: {settle_result.payer}")
+                    logger.info(f"   Output Hash: {output_hash}")
+                    logger.info(f"   Status: SETTLED ON-CHAIN")
+                    return True
+                else:
+                    logger.error(f"❌ Payment settlement FAILED: {settle_result.error_reason}")
+                    # TODO: Queue for retry
+                    return False
+            except Exception as e:
+                logger.error(f"Error settling payment via facilitator: {e}")
+                # Don't fail the request if settlement fails
+                # Settlement can be retried later
+                logger.warning("Settlement failed but request completed - will retry later")
+                # TODO: Queue settlement for retry
+                return False
+        else:
+            logger.warning("No facilitator available for settlement")
+            # TODO: Queue settlement for later retry
+            return False
+            
+    except Exception as e:
+        import traceback
+        logger.error(f"Error in settle_payment: {e}")
+        logger.error(f"Traceback:\n{traceback.format_exc()}")
+        return False
+
+
+def require_payment_for_tool(
+    price: Price,
+    description: str = ""
+):
+    """
+    Decorator for MCP tools that require payment (like FastAPI's require_payment).
+    
+    Matches the pattern from IATPx402/e2e/servers/fastapi/main.py where payment
+    is configured OUTSIDE the endpoint function.
+    
+    Global config (from middleware):
+    - server_address: Where payments are sent (configured at middleware init)
+    - requires_auth: Whether API needs auth (configured at middleware init)
+    - internal_api_key_env_var: Which env var has API key (configured at middleware init)
+    
+    Per-endpoint config (decorator params):
+    - price: Payment price (TokenAmount with network info, or "$0.001" string) - per-endpoint
+    - description: Description of what is being purchased - per-endpoint
+    
+    Args:
+        price: Payment price (TokenAmount with asset.network, or "$0.001" string) - per-endpoint
+               When TokenAmount is used, settlement_token_address and settlement_token_network
+               are extracted from price.asset.address and price.asset.network
+        description: Description of what is being purchased - per-endpoint
+        
+    Returns:
+        Decorator function
+        
+    Usage - API requires auth (dual-mode):
+        @mcp.tool()
+        @require_payment_for_tool(
+            price=TokenAmount(
+                amount="1000",
+                asset=TokenAsset(
+                    address="0xUSDC...",
+                    decimals=6,
+                    network="base-sepolia",
+                    eip712=EIP712Domain(name="USD Coin", version="2")
+                )
+            ),
+            description="Get data"
+        )
+        async def get_data(context: Context) -> Dict[str, Any]:
+            api_key = context.state.api_key_to_use  # Client's OR server's API key
+            response = requests.get(url, headers={"Authorization": f"Bearer {api_key}"})
+            return response.json()
+    
+    Usage - API doesn't require auth (payment-only):
+        @mcp.tool()
+        @require_payment_for_tool(
+            price=TokenAmount(
+                amount="1000",
+                asset=TokenAsset(
+                    address="0xUSDC...",
+                    decimals=6,
+                    network="base-sepolia",
+                    eip712=EIP712Domain(name="USD Coin", version="2")
+                )
+            ),
+            description="Get public data"
+        )
+        async def get_public_data(context: Context) -> Dict[str, Any]:
+            # No API key needed - just call the API
+            response = requests.get(url)
+            return response.json()
+    """
+    def decorator(func: Callable):
+        @wraps(func)
+        async def wrapper(context: Context, *args, **kwargs):
+            """
+            Decorator for official MCP SDK with Starlette middleware.
+            
+            1. Copies api_key_to_use from request.state to context.state (for tool access)
+            2. Calls tool function
+            3. Schedules async settlement if client paid
+            """
+            
+            # Get request from context for settlement
+            request = None
+            if hasattr(context, 'request_context') and context.request_context:
+                if hasattr(context.request_context, 'request'):
+                    request = context.request_context.request
+            
+            # NOTE: Don't try to set context.state - Context is immutable Pydantic model
+            # Tools should use get_active_api_key(context) which reads from request.state
+            
+            # Call the actual tool function
+            logger.info("=" * 80)
+            logger.info(f"TOOL CALL STARTED: {func.__name__}")
+            logger.info("=" * 80)
+            logger.info(f"Arguments: {kwargs}")
+            #============================================================== 
+            result = await func(context, *args, **kwargs)
+            #==============================================================
+            logger.info(f"✅ Tool execution complete: {func.__name__}")
+            
+            # Log tool result for debugging
+            logger.info("-" * 80)
+            logger.info(f"📤 TOOL RESULT:")
+            result_str = str(result)
+            if len(result_str) > 1000:
+                logger.info(f"{result_str[:1000]}...")
+                logger.info(f"   (Result truncated - {len(result_str)} total chars)")
+            else:
+                logger.info(result_str)
+            logger.info("=" * 80)
+            
+            # Auto-settle payment ASYNCHRONOUSLY if client paid
+            if request and hasattr(request, 'state'):
+                payment_validated = getattr(request.state, 'payment_validated', False)
+                payment_uuid = getattr(request.state, 'payment_uuid', None)
+                
+                if payment_validated:
+                    # Client paid - schedule settlement (fire-and-forget)
+                    logger.info(f"💰 Payment detected for {func.__name__} - scheduling settlement (async)...")
+                    
+                    # Get payment and endpoint info for settlement
+                    if isinstance(price, TokenAmount):
+                        payment_price_wei = price.amount
+                        payment_price_float = float(price.amount) / (10 ** price.asset.decimals)
+                        settlement_token_address = price.asset.address
+                        settlement_token_network = price.asset.network
+                        
+                        # Get server address from env (set in template)
+                        server_address = os.getenv("SERVER_ADDRESS")
+                        
+                        if server_address:
+                            endpoint_info = EndpointPaymentInfo(
+                                settlement_token_address=settlement_token_address,
+                                settlement_token_network=settlement_token_network,
+                                payment_price_float=payment_price_float,
+                                payment_price_wei=payment_price_wei,
+                                server_address=server_address
+                            )
+                            
+                            # Schedule settlement asynchronously
+                            import asyncio
+                            
+                            # Get middleware reference (has facilitator)
+                            d402_mw = getattr(request.state, 'd402_middleware', None)
+                            
+                            # Use existing settle_payment function (no duplication!)
+                            async def do_settlement():
+                                try:
+                                    logger.info(f"🚀 Settlement task started for {func.__name__}")
+                                    logger.info(f"   Tool result size: {len(str(result))} bytes")
+                                    
+                                    # Create a simple wrapper context with payment_payload
+                                    class SettlementContext:
+                                        class State:
+                                            def __init__(self):
+                                                self.payment_payload = getattr(request.state, 'payment_payload', None)
+                                                self.payment_uuid = payment_uuid
+                                                self.facilitator_fee_percent = getattr(request.state, 'facilitator_fee_percent', 250)
+                                        
+                                        def __init__(self):
+                                            self.state = SettlementContext.State()
+                                    
+                                    settlement_ctx = SettlementContext()
+                                    
+                                    # Call existing settle_payment function
+                                    settlement_success = await settle_payment(
+                                        context=settlement_ctx,
+                                        endpoint_info=endpoint_info,
+                                        output_data=result,
+                                        middleware=d402_mw  # Pass Starlette middleware (has facilitator)
+                                    )
+                                    
+                                    if settlement_success:
+                                        logger.info(f"✅ Settlement task completed successfully for {func.__name__}")
+                                    else:
+                                        logger.warning(f"⚠️  Settlement task failed for {func.__name__}")
+                                        
+                                except Exception as e:
+                                    logger.error(f"❌ Settlement task error for {func.__name__}: {e}")
+                                    import traceback
+                                    logger.error(traceback.format_exc())
+                            
+                            asyncio.create_task(do_settlement())
+                            logger.info(f"📅 Settlement task scheduled for {func.__name__} - client gets response immediately")
+            
+            # Return immediately - don't wait for settlement!
+            return result
+        
+        return wrapper
+    return decorator
+
+