traia-iatp 0.1.2__py3-none-any.whl → 0.1.67__py3-none-any.whl

traia_iatp/d402/mcp_middleware.py

@@ -0,0 +1,296 @@
+"""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 request accepted by facilitator:")
+                    logger.info(f"   Status: PENDING_SETTLEMENT (queued for on-chain settlement)")
+                    logger.info(f"   Network: {settle_result.network}")
+                    logger.info(f"   Payer: {settle_result.payer}")
+                    logger.info(f"   Output Hash: {output_hash}")
+                    logger.info(f"   Note: Facilitator cron will batch-settle 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 - METADATA ONLY.
+    
+    This decorator ONLY stores payment configuration metadata on the function.
+    All payment processing (verify, settle) is handled by D402PaymentMiddleware.
+    
+    Usage:
+        @mcp.tool()
+        @require_payment_for_tool(
+            price=TokenAmount(
+                amount="1000",
+                asset=TokenAsset(
+                    address="0xUSDC...",
+                    decimals=6,
+                    network="base-sepolia",
+                    eip712=EIP712Domain(name="IATPWallet", version="1")
+                )
+            ),
+            description="Get cryptocurrency price data"
+        )
+        async def get_price(context: Context, coin_id: str) -> Dict[str, Any]:
+            # Payment already verified by middleware
+            api_key = get_active_api_key(context)
+            response = requests.get(url, headers={"Authorization": f"Bearer {api_key}"})
+            return response.json()
+    
+    Args:
+        price: Payment configuration (TokenAmount with network, token, etc)
+        description: Service description for settlement
+    
+    Returns:
+        Decorator that attaches metadata to the function
+    """
+    def decorator(func: Callable):
+        # Store payment metadata on function for middleware extraction
+        func._d402_payment_config = {
+            "price": price,
+            "description": description
+        }
+        
+        # Return function unchanged - middleware handles all payment logic
+        return func
+    
+    return decorator
+
+

traia_iatp/d402/models.py

@@ -0,0 +1,116 @@
+"""D402 payment models for IATP protocol."""
+
+from enum import Enum
+from typing import Optional, Dict, Any
+from pydantic import BaseModel, Field
+
+
+class PaymentScheme(str, Enum):
+    """Payment schemes supported by d402."""
+    EXACT = "exact"  # EIP-3009 exact payment
+
+
+class D402ServicePrice(BaseModel):
+    """Pricing configuration for an IATP service.
+    
+    Supports any ERC20 token payment with full token details.
+    """
+    
+    # Token details
+    token_address: str = Field(..., description="Token contract address (e.g., USDC, TRAIA, DAI)")
+    token_symbol: str = Field(..., description="Token symbol for display")
+    token_decimals: int = Field(..., description="Token decimals (6 for USDC, 18 for most)")
+    
+    # Price (stored in multiple formats for convenience)
+    price_wei: str = Field(..., description="Price in wei/atomic units")
+    price_float: float = Field(..., description="Price in token units (human-readable)")
+    
+    # Network configuration
+    network: str = Field(..., description="Network (sepolia, base-sepolia, etc.)")
+    chain_id: int = Field(..., description="Chain ID")
+    
+    # Optional USD equivalent (for display only)
+    usd_amount: Optional[float] = Field(None, description="Approximate USD value")
+    
+    # Maximum timeout for payment completion
+    max_timeout_seconds: int = Field(default=300, description="Max time to complete payment")
+
+
+class D402Config(BaseModel):
+    """Configuration for d402 payment integration in IATP.
+    
+    Supports per-path pricing with different tokens.
+    """
+    
+    # Enable/disable d402 payments
+    enabled: bool = Field(default=False, description="Enable d402 payments")
+    
+    # Payment address (utility agent contract address)
+    pay_to_address: str = Field(..., description="Ethereum address to receive payments")
+    
+    # Pricing configuration
+    # Can be per-path (e.g., {"/analyze": D402ServicePrice(...), "/extract": D402ServicePrice(...)})
+    # or default for all paths
+    path_prices: Dict[str, D402ServicePrice] = Field(
+        default_factory=dict,
+        description="Price configuration per path"
+    )
+    
+    # Default price (used if path not in path_prices)
+    default_price: Optional[D402ServicePrice] = Field(None, description="Default price for all paths")
+    
+    # Legacy: Pricing per service/skill (deprecated, use path_prices)
+    skill_prices: Dict[str, D402ServicePrice] = Field(
+        default_factory=dict, 
+        description="Custom prices per skill ID (deprecated)"
+    )
+    
+    # Facilitator configuration
+    facilitator_url: str = Field(
+        default="https://api.traia.io/d402/facilitator",
+        description="URL of the d402 facilitator service"
+    )
+    
+    # Custom facilitator authentication (if needed)
+    facilitator_api_key: Optional[str] = Field(None, description="API key for facilitator")
+    
+    # Paths to gate with payments (* for all)
+    protected_paths: list[str] = Field(
+        default_factory=lambda: ["*"],
+        description="Paths that require payment"
+    )
+    
+    # Service description for payment prompt
+    service_description: str = Field(..., description="Description shown in payment UI")
+    
+    # Metadata
+    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
+
+
+class D402PaymentInfo(BaseModel):
+    """Payment information for agent card discovery."""
+    
+    enabled: bool = Field(..., description="Whether d402 is enabled")
+    payment_schemes: list[PaymentScheme] = Field(
+        default_factory=lambda: [PaymentScheme.EXACT],
+        description="Supported payment schemes"
+    )
+    networks: list[str] = Field(..., description="Supported blockchain networks")
+    default_price: D402ServicePrice = Field(..., description="Default pricing")
+    facilitator_url: str = Field(..., description="Facilitator service URL")
+    
+    class Config:
+        use_enum_values = True
+
+
+class IATPSettlementRequest(BaseModel):
+    """Request to settle a payment through IATP settlement layer."""
+    
+    consumer: str = Field(..., description="Consumer address (client agent)")
+    provider: str = Field(..., description="Provider address (utility agent)")
+    amount: str = Field(..., description="Amount in atomic units")
+    timestamp: int = Field(..., description="Request timestamp")
+    service_description: str = Field(..., description="Description of service")
+    consumer_signature: str = Field(..., description="Consumer's EIP-712 signature")
+    provider_signature: str = Field(..., description="Provider's attestation signature")
+

traia_iatp/d402/networks.py

@@ -0,0 +1,98 @@
+"""Network configuration for d402 payments.
+
+This module defines supported networks and their token configurations.
+Customized for IATP - uses database-driven network configuration.
+"""
+
+from typing import Literal, Dict, Any
+from typing_extensions import TypedDict
+
+
+# Network type definition
+SupportedNetworks = Literal[
+    "sepolia",
+    "base-sepolia",
+    "arbitrum-sepolia",
+    "base-mainnet",
+    "arbitrum-mainnet",
+]
+
+
+class NetworkConfig(TypedDict):
+    """Network configuration."""
+    chain_id: int
+    name: str
+    rpc_url: str
+    explorer_url: str
+    usdc_address: str
+
+
+# Network configurations
+NETWORKS: Dict[str, NetworkConfig] = {
+    "sepolia": {
+        "chain_id": 11155111,
+        "name": "Ethereum Sepolia",
+        "rpc_url": "https://ethereum-sepolia-rpc.publicnode.com",
+        "explorer_url": "https://sepolia.etherscan.io",
+        "usdc_address": "0x1c7D4B196Cb0C7B01d743Fbc6116a902379C7238"
+    },
+    "base-sepolia": {
+        "chain_id": 84532,
+        "name": "Base Sepolia",
+        "rpc_url": "https://sepolia.base.org",
+        "explorer_url": "https://sepolia.basescan.org",
+        "usdc_address": "0x036CbD53842c5426634e7929541eC2318f3dCF7e"
+    },
+    "arbitrum-sepolia": {
+        "chain_id": 421614,
+        "name": "Arbitrum Sepolia",
+        "rpc_url": "https://arbitrum-sepolia-rpc.publicnode.com",
+        "explorer_url": "https://sepolia.arbiscan.io",
+        "usdc_address": "0x75faf114eafb1BDbe2F0316DF893fd58CE46AA4d"
+    },
+    "base-mainnet": {
+        "chain_id": 8453,
+        "name": "Base",
+        "rpc_url": "https://mainnet.base.org",
+        "explorer_url": "https://basescan.org",
+        "usdc_address": "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913"
+    },
+    "arbitrum-mainnet": {
+        "chain_id": 42161,
+        "name": "Arbitrum One",
+        "rpc_url": "https://arbitrum-one-rpc.publicnode.com",
+        "explorer_url": "https://arbiscan.io",
+        "usdc_address": "0xaf88d065e77c8cC2239327C5EDb3A432268e5831"
+    }
+}
+
+
+def get_network_config(network: str) -> NetworkConfig:
+    """Get configuration for a network."""
+    if network not in NETWORKS:
+        raise ValueError(f"Unsupported network: {network}")
+    return NETWORKS[network]
+
+
+def get_usdc_address(network: str) -> str:
+    """Get USDC address for a network."""
+    return get_network_config(network)["usdc_address"]
+
+
+def get_chain_id(network: str) -> int:
+    """Get chain ID for a network."""
+    return get_network_config(network)["chain_id"]
+
+
+# TODO: Load from database network table
+async def load_networks_from_db() -> Dict[str, NetworkConfig]:
+    """Load network configurations from database.
+    
+    This will query the Network table and build NETWORKS dict dynamically.
+    For now, returns static config.
+    """
+    # from db.dal.models import Network
+    # networks = await db.query(Network).filter(Network.active == True).all()
+    # return {net.shortname: {...} for net in networks}
+    return NETWORKS
+

traia_iatp/d402/path.py

@@ -0,0 +1,43 @@
+import fnmatch
+import re
+from typing import Union
+
+
+def path_is_match(path: Union[str, list[str]], request_path: str) -> bool:
+    """
+    Check if request path matches the specified path pattern(s).
+
+    Supports:
+    - Exact matching: "/api/users"
+    - Glob patterns: "/api/users/*", "/api/*/profile"
+    - Regex patterns (prefix with 'regex:'): "regex:^/api/users/\\d+$"
+    - List of any of the above
+
+    Args:
+        path: Path pattern(s) to match against. Can be a string or list of strings.
+        request_path: The actual request path to check.
+
+    Returns:
+        bool: True if the request path matches any of the patterns, False otherwise.
+    """
+
+    def single_path_match(pattern: str) -> bool:
+        # Regex pattern
+        if pattern.startswith("regex:"):
+            regex_pattern = pattern[6:]  # Remove 'regex:' prefix
+            return bool(re.match(regex_pattern, request_path))
+
+        # Glob pattern (contains * or ?)
+        elif "*" in pattern or "?" in pattern:
+            return fnmatch.fnmatch(request_path, pattern)
+
+        # Exact match
+        else:
+            return pattern == request_path
+
+    if isinstance(path, str):
+        return single_path_match(path)
+    elif isinstance(path, list):
+        return any(single_path_match(p) for p in path)
+
+    return False

traia_iatp/d402/payment_introspection.py

@@ -0,0 +1,126 @@
+"""
+Helper to extract payment configurations from @require_payment_for_tool decorators.
+
+This allows us to have a single source of truth - payment config is declared
+in the decorator, and we introspect it to build TOOL_PAYMENT_CONFIGS.
+"""
+
+import logging
+from typing import Dict, Any, Optional
+
+from .types import TokenAmount
+
+logger = logging.getLogger(__name__)
+
+
+def extract_payment_configs_from_mcp(mcp_server, server_address: str) -> Dict[str, Dict[str, Any]]:
+    """
+    Extract payment configurations from tools decorated with @require_payment_for_tool.
+    
+    This introspects the decorator closures to extract TokenAmount objects,
+    eliminating the need to duplicate payment configuration.
+    
+    Args:
+        mcp_server: FastMCP server instance
+        server_address: Server's payment address
+    
+    Returns:
+        Dict mapping tool names to payment configurations
+        Format: {"tool_name": {"price_wei": "1000", "token_address": "0x...", ...}}
+    
+    Usage:
+        mcp = FastMCP("Server")
+        
+        # Add tools with @require_payment_for_tool decorators
+        @mcp.tool()
+        @require_payment_for_tool(price=TokenAmount(...))
+        async def my_tool(context): ...
+        
+        # Extract configs dynamically
+        TOOL_PAYMENT_CONFIGS = extract_payment_configs_from_mcp(mcp, SERVER_ADDRESS)
+        
+        # Add middleware with extracted configs
+        app.add_middleware(D402PaymentMiddleware, tool_payment_configs=TOOL_PAYMENT_CONFIGS, ...)
+    """
+    tool_payment_configs = {}
+    
+    try:
+        # Get registered tools from FastMCP
+        tools = mcp_server._tool_manager.list_tools()
+        
+        for tool in tools:
+            if not hasattr(tool, 'fn'):
+                continue
+            
+            fn = tool.fn
+            tool_name = tool.name
+            
+            # Check if function has payment metadata (from @require_payment_for_tool decorator)
+            # New approach: metadata stored as attribute (no closure needed)
+            if hasattr(fn, '_d402_payment_config'):
+                payment_config = fn._d402_payment_config
+                price = payment_config.get('price')
+                description = payment_config.get('description', tool_name)
+                
+                if isinstance(price, TokenAmount):
+                    token_amount = price
+                else:
+                    logger.debug(f"Tool {tool_name}: price is not TokenAmount")
+                    continue
+            # Fallback: Check closure for backwards compatibility
+            elif hasattr(fn, '__closure__') and fn.__closure__:
+                token_amount = None
+                description = None
+                for cell in fn.__closure__:
+                    try:
+                        val = cell.cell_contents
+                        if isinstance(val, TokenAmount):
+                            token_amount = val
+                        elif isinstance(val, str):
+                            description = val
+                    except:
+                        pass
+                
+                if not token_amount:
+                    logger.debug(f"Tool {tool_name}: No TokenAmount in closure")
+                    continue
+                    
+                description = description or tool.description or tool_name
+            else:
+                logger.debug(f"Tool {tool_name}: No payment metadata found")
+                continue
+            
+            if token_amount:
+                # Extract payment config from TokenAmount (including EIP712 domain)
+                # Use description from metadata if available, otherwise from tool
+                final_description = description if 'description' in locals() else (tool.description or tool_name)
+                
+                config = {
+                    "price_wei": token_amount.amount,
+                    "token_address": token_amount.asset.address,
+                    "network": token_amount.asset.network,
+                    "server_address": server_address,
+                    "description": final_description,
+                    "eip712_domain": {
+                        "name": token_amount.asset.eip712.name if token_amount.asset.eip712 else "USD Coin",
+                        "version": token_amount.asset.eip712.version if token_amount.asset.eip712 else "2"
+                    }
+                }
+                
+                tool_payment_configs[tool_name] = config
+                logger.info(f"✅ Extracted payment config for {tool_name}: {config['price_wei']} wei on {config['network']}")
+            else:
+                logger.debug(f"Tool {tool_name}: No TokenAmount found (free tool)")
+        
+        logger.info(f"📊 Extracted {len(tool_payment_configs)} payment configs from decorators")
+        
+    except Exception as e:
+        logger.error(f"Error extracting payment configs: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+    
+    return tool_payment_configs
+
+
+__all__ = ["extract_payment_configs_from_mcp"]
+