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

traia_iatp/server/templates/__main__.py.j2

@@ -24,14 +24,22 @@ from a2a.types import AgentCard, AgentSkill, AgentCapabilities
 from a2a.server.tasks import InMemoryTaskStore
 from a2a.server.request_handlers import DefaultRequestHandler
 from a2a.server.events.event_queue import EventQueue
-from hypercorn.asyncio import serve
-from hypercorn.config import Config
-from starlette.responses import StreamingResponse
+import uvicorn
+from starlette.responses import StreamingResponse, JSONResponse
 from starlette.requests import Request
 from starlette.routing import Route
 from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import Response
 
+# Import D402 payment middleware and helpers (generalized version)
+from traia_iatp.d402.servers import D402PaymentMiddleware, build_payment_config
+from traia_iatp.d402.asgi_wrapper import D402ASGIWrapper
+from traia_iatp.d402 import D402PriceBuilder
+
+# Import httpx for MCP server verification
+import httpx
+
 # Import AgentOps for monitoring and observability
 try:
     import agentops
@@ -48,13 +56,22 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
 from traia_iatp.mcp import MCPServerConfig
 from .agent_executor import {{ class_name }}AgentExecutor
 
-# Configure logging
+# Configure logging FIRST (before any logger usage)
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
+# D402 payment support is configured via environment variables and handled by middleware
+# The D402PaymentMiddleware is added to the Starlette app based on .env configuration
+# No additional imports needed here - middleware handles all D402 logic
+D402_AVAILABLE = os.getenv("D402_ENABLED", "false").lower() == "true"
+if D402_AVAILABLE:
+    logger.info("✅ D402 payment support enabled (configured via environment)")
+else:
+    logger.info("ℹ️  D402 payment support disabled (set D402_ENABLED=true to enable)")
+
 # Enable debug logging for HTTP/2 and connection events if requested
 if os.environ.get("DEBUG_PROTOCOL", "false").lower() == "true":
     logging.getLogger("hypercorn.access").setLevel(logging.DEBUG)
@@ -109,6 +126,119 @@ class ProtocolLoggingMiddleware(BaseHTTPMiddleware):
         return response
 
 
+# Global state for MCP connection tracking
+mcp_connection_status = {
+    "verified": False,
+    "last_checked": None,
+    "error": None
+}
+
+async def verify_mcp_connection(mcp_url: str, mcp_name: str) -> bool:
+    """
+    Verify connection to the MCP server by testing endpoints.
+    Returns True if connection is working, False otherwise.
+    """
+    global mcp_connection_status
+    
+    try:
+        logger.info(f"🔍 Verifying MCP server connection to: {mcp_url}")
+        
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            # Try to test the MCP server is reachable
+            # First try .well-known/mcp.json (if available)
+            card_url = f"{mcp_url}/.well-known/mcp.json"
+            logger.info(f"   Trying MCP card at: {card_url}")
+            
+            try:
+                response = await client.get(card_url)
+                response.raise_for_status()
+                
+                mcp_card = response.json()
+                logger.info(f"   ✅ MCP card fetched successfully")
+                logger.info(f"   MCP Server: {mcp_card.get('name', 'Unknown')}")
+                logger.info(f"   Capabilities: {len(mcp_card.get('capabilities', []))} tools")
+            except Exception as e:
+                # Well-known endpoint not available - that's OK for some MCP servers
+                logger.info(f"   ℹ️  MCP card not available (some servers don't expose it): {e}")
+            
+            # Test the MCP endpoint properly (like the tests do)
+            logger.info(f"   Testing MCP JSON-RPC endpoint: {mcp_url}")
+            
+            # Step 1: Initialize MCP session (required before tools/list)
+            init_request = {
+                "jsonrpc": "2.0",
+                "method": "initialize",
+                "id": "verify-init",
+                "params": {
+                    "protocolVersion": "2024-11-05",
+                    "capabilities": {},
+                    "clientInfo": {"name": "utility-agent-verifier", "version": "1.0"}
+                }
+            }
+            
+            init_response = await client.post(
+                mcp_url,
+                json=init_request,
+                headers={
+                    "Content-Type": "application/json",
+                    "Accept": "application/json, text/event-stream"
+                }
+            )
+            init_response.raise_for_status()
+            
+            session_id = init_response.headers.get("mcp-session-id")
+            if session_id:
+                logger.info(f"   ✅ MCP session initialized: {session_id[:20]}...")
+            else:
+                logger.warning(f"   ⚠️  No session ID returned (may be OK for some servers)")
+            
+            # Step 2: Test tools/list with session
+            list_request = {
+                "jsonrpc": "2.0",
+                "method": "tools/list",
+                "id": "verify-list"
+            }
+            
+            list_headers = {
+                "Content-Type": "application/json",
+                "Accept": "application/json, text/event-stream"
+            }
+            if session_id:
+                list_headers["mcp-session-id"] = session_id
+            
+            list_response = await client.post(
+                mcp_url,
+                json=list_request,
+                headers=list_headers
+            )
+            list_response.raise_for_status()
+            
+            logger.info(f"   ✅ MCP server responding correctly (HTTP {list_response.status_code})")
+            
+            # Update connection status
+            mcp_connection_status["verified"] = True
+            mcp_connection_status["last_checked"] = datetime.now()
+            mcp_connection_status["error"] = None
+            
+            logger.info(f"✅ MCP server connection verified successfully!")
+            return True
+            
+    except httpx.HTTPError as e:
+        error_msg = f"HTTP error connecting to MCP server: {e}"
+        logger.error(f"❌ {error_msg}")
+        mcp_connection_status["verified"] = False
+        mcp_connection_status["last_checked"] = datetime.now()
+        mcp_connection_status["error"] = error_msg
+        return False
+    except Exception as e:
+        error_msg = f"Unexpected error verifying MCP connection: {e}"
+        logger.error(f"❌ {error_msg}")
+        mcp_connection_status["verified"] = False
+        mcp_connection_status["last_checked"] = datetime.now()
+        mcp_connection_status["error"] = error_msg
+        return False
+
+
 class StreamingRequestHandler(DefaultRequestHandler):
     """Extended request handler with SSE streaming support."""
     
@@ -239,6 +369,48 @@ def create_app():
     supports_streaming = mcp_data.get("server_type") == "streamable-http" or \
                         "stream" in mcp_data.get("capabilities", [])
     
+    # D402 Configuration (if enabled) - for A2A endpoint protection
+    d402_enabled = D402_AVAILABLE and os.getenv("D402_ENABLED", "false").lower() == "true"
+    d402_token_amount = None
+    
+    if d402_enabled:
+        logger.info("✅ D402 payments enabled")
+        
+        contract_address = os.getenv("UTILITY_AGENT_CONTRACT_ADDRESS")
+        token_address = os.getenv("D402_TOKEN_ADDRESS")
+        token_symbol = os.getenv("D402_TOKEN_SYMBOL", "USDC")
+        token_decimals = int(os.getenv("D402_TOKEN_DECIMALS", "6"))
+        network = os.getenv("D402_NETWORK", "sepolia")
+        facilitator_url = os.getenv("D402_FACILITATOR_URL", "http://localhost:7070")
+        operator_key = os.getenv("UTILITY_AGENT_OPERATOR_PRIVATE_KEY")
+        testing_mode = os.getenv("D402_TESTING_MODE", "false").lower() == "true"
+        
+        if not contract_address:
+            logger.error("❌ D402 enabled but UTILITY_AGENT_CONTRACT_ADDRESS not set!")
+            d402_enabled = False
+        elif not token_address:
+            logger.error("❌ D402 enabled but D402_TOKEN_ADDRESS not set!")
+            d402_enabled = False
+        else:
+            # Use D402PriceBuilder for clean price creation
+            usd_price = float(os.getenv("D402_PRICE_USD", "0.01"))
+            
+            price_builder = D402PriceBuilder(
+                token_address=token_address,
+                token_decimals=token_decimals,
+                network=network,
+                token_symbol=token_symbol
+            )
+            
+            d402_token_amount = price_builder.create_price(usd_price)
+            
+            logger.info(f"💰 D402 Price: ${usd_price} USD ({d402_token_amount.amount} wei with {token_decimals} decimals)")
+            logger.info(f"📍 Pay to: {contract_address}")
+            logger.info(f"🌐 Network: {network}")
+            logger.info(f"🪙  Token: {token_symbol} ({token_address})")
+            logger.info(f"🔗 Facilitator: {facilitator_url}")
+            logger.info(f"🧪 Testing mode: {testing_mode}")
+    
     # Create agent skills based on MCP capabilities
     skills = []
     
@@ -294,17 +466,30 @@ def create_app():
     # Authentication can be added here if needed in the future
     # Currently the A2A protocol handles authentication at a different layer
     
-    # Create agent card
-    agent_card = AgentCard(
-        name="{{ agent_id }}",
-        description="{{ agent_description }}",
-        url=f"http://0.0.0.0:{os.environ.get('PORT', 8000)}",
-        version="{{ agent_version }}",
-        capabilities=capabilities,
-        skills=skills,
-        defaultInputModes=["text", "text/plain"],
-        defaultOutputModes=["text", "text/plain", "text/event-stream"] if supports_streaming else ["text", "text/plain"]
-    )
+    # Create agent card with proper URL
+    # For Cloud Run: use PUBLIC_URL environment variable
+    # For local: use http://0.0.0.0:{port}/a2a
+    public_url = os.environ.get("PUBLIC_URL") or os.environ.get("SERVICE_URL")
+    if public_url:
+        # Cloud Run - use public URL with /a2a endpoint
+        agent_url = f"{public_url}/a2a"
+    else:
+        # Local - use 0.0.0.0:{port}
+        agent_url = f"http://0.0.0.0:{os.environ.get('PORT', 8000)}/a2a"
+    
+    agent_card_dict = {
+        "name": "{{ agent_id }}",
+        "description": "{{ agent_description }}",
+        "url": agent_url,
+        "version": "{{ agent_version }}",
+        "capabilities": capabilities.model_dump(),
+        "skills": [s.model_dump() for s in skills],
+        "defaultInputModes": ["text", "text/plain"],
+        "defaultOutputModes": ["text", "text/plain", "text/event-stream"] if supports_streaming else ["text", "text/plain"]
+    }
+    
+    # D402 payment info is automatically added to agent card by middleware
+    agent_card = AgentCard(**agent_card_dict)
     
     # Create executor with MCP config
     executor = {{ class_name }}AgentExecutor(mcp_config, supports_streaming=supports_streaming)
@@ -322,8 +507,74 @@ def create_app():
         http_handler=request_handler
     )
     
-    # Build the Starlette app - this should add all necessary routes including JSON-RPC endpoint
-    starlette_app = app.build()
+    # Build the Starlette app with A2A JSON-RPC endpoint at /a2a
+    # Add D402 middleware using Starlette's raw ASGI wrapper if enabled
+    if d402_enabled and d402_token_amount:
+        logger.info("✅ Configuring D402 payment middleware for A2A endpoints")
+        
+        # Get D402 configuration from environment
+        contract_address = os.getenv("UTILITY_AGENT_CONTRACT_ADDRESS")
+        facilitator_url = os.getenv("D402_FACILITATOR_URL", "http://localhost:7070")
+        operator_key = os.getenv("UTILITY_AGENT_OPERATOR_PRIVATE_KEY")
+        testing_mode = os.getenv("D402_TESTING_MODE", "false").lower() == "true"
+        
+        # Build payment configs for A2A endpoint
+        endpoint_payment_configs = {
+            "/a2a": build_payment_config(
+                price=d402_token_amount,
+                server_address=contract_address,
+                description="{{ agent_description }}"
+            )
+        }
+        
+        logger.info(f"   💳 Middleware configured for {len(endpoint_payment_configs)} endpoint(s)")
+        logger.info(f"   🔒 Protected paths: {list(endpoint_payment_configs.keys())}")
+        logger.info(f"   🔐 Requires auth: False (payment-only)")
+        logger.info(f"   🧪 Testing mode: {testing_mode}")
+        
+        # Build Starlette app first (without middleware)
+        starlette_app = app.build(rpc_url='/a2a')
+        
+        # Add CORS middleware to allow all origins
+        starlette_app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],  # Allow all origins
+            allow_credentials=True,
+            allow_methods=["*"],  # Allow all methods
+            allow_headers=["*"],  # Allow all headers
+        )
+        logger.info("✅ Added CORS middleware (allow all origins)")
+        
+        # Wrap with raw ASGI wrapper (bypasses Starlette middleware issues)
+        # This intercepts at the lowest ASGI level before any framework processing
+        d402_wrapped_app = D402ASGIWrapper(
+            app=starlette_app,
+            server_address=contract_address,
+            endpoint_payment_configs=endpoint_payment_configs,
+            requires_auth=False,
+            internal_api_key=None,
+            testing_mode=testing_mode,
+            facilitator_url=facilitator_url,
+            facilitator_api_key=os.getenv("D402_FACILITATOR_API_KEY"),
+            server_name="{{ agent_id }}"
+        )
+        
+        logger.info("   ✅ D402 raw ASGI wrapper applied (guaranteed interception)")
+    else:
+        # No D402 - just build normally
+        starlette_app = app.build(rpc_url='/a2a')
+        
+        # Add CORS middleware even without D402
+        starlette_app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],  # Allow all origins
+            allow_credentials=True,
+            allow_methods=["*"],  # Allow all methods
+            allow_headers=["*"],  # Allow all headers
+        )
+        logger.info("✅ Added CORS middleware (allow all origins)")
+        
+        d402_wrapped_app = None
     
     # Add protocol logging middleware if debug mode is enabled
     if os.environ.get("DEBUG_PROTOCOL", "false").lower() == "true":
@@ -360,6 +611,54 @@ def create_app():
         Route("/a2a/tasks/resubscribe", handle_resubscribe_endpoint, methods=["POST"])
     )
     
+    # Add health endpoint with MCP connection verification
+    async def health_endpoint(request: Request):
+        """Health check endpoint that verifies MCP connection every 10 minutes."""
+        global mcp_connection_status
+        
+        # Check if we need to reverify (more than 10 minutes since last check)
+        now = datetime.now()
+        should_reverify = (
+            mcp_connection_status["last_checked"] is None or
+            (now - mcp_connection_status["last_checked"]).total_seconds() > 600  # 10 minutes
+        )
+        
+        if should_reverify:
+            logger.info("⏰ 10 minutes elapsed since last MCP verification, rechecking...")
+            await verify_mcp_connection(mcp_config.url, mcp_config.name)
+        
+        # Return health status
+        if mcp_connection_status["verified"]:
+            return JSONResponse({
+                "status": "healthy",
+                "mcp_connection": "verified",
+                "mcp_server": mcp_config.name,
+                "mcp_url": mcp_config.url,
+                "last_verified": mcp_connection_status["last_checked"].isoformat() if mcp_connection_status["last_checked"] else None
+            })
+        else:
+            return JSONResponse(
+                {
+                    "status": "unhealthy",
+                    "mcp_connection": "failed",
+                    "mcp_server": mcp_config.name,
+                    "mcp_url": mcp_config.url,
+                    "error": mcp_connection_status["error"],
+                    "last_checked": mcp_connection_status["last_checked"].isoformat() if mcp_connection_status["last_checked"] else None
+                },
+                status_code=503
+            )
+    
+    starlette_app.routes.append(
+        Route("/health", health_endpoint, methods=["GET"])
+    )
+    
+    # Return the D402-wrapped app if D402 is enabled, otherwise return the unwrapped app
+    if d402_wrapped_app:
+        logger.info(f"🔒 Returning D402-wrapped app (type: {type(d402_wrapped_app).__name__})")
+        return d402_wrapped_app
+    else:
+        logger.info(f"⚠️  Returning unwrapped app (D402 disabled, type: {type(starlette_app).__name__})")
     return starlette_app
 
 
@@ -391,8 +690,38 @@ def generate_self_signed_cert(cert_path: str = "cert.pem", key_path: str = "key.
         raise
 
 
-async def main():
-    """Main function to start the A2A server with HTTP/2 support."""
+async def verify_mcp_before_start():
+    """Verify MCP connection before starting server."""
+    # Get MCP configuration
+    mcp_url = os.getenv("MCP_SERVER_URL")
+    mcp_name = "{{ mcp_server_name }}"
+    
+    if not mcp_url:
+        logger.error("MCP_SERVER_URL not set!")
+        return False
+    
+    logger.info("="*80)
+    logger.info("🔍 Verifying MCP server connection before startup...")
+    logger.info("="*80)
+    
+    mcp_ok = await verify_mcp_connection(mcp_url, mcp_name)
+    
+    if not mcp_ok:
+        logger.error("="*80)
+        logger.error("❌ STARTUP FAILED: Cannot connect to MCP server!")
+        logger.error("   The utility agent will NOT start without MCP server access.")
+        logger.error("   Please check the MCP server is running and accessible.")
+        logger.error("="*80)
+        return False
+    
+    logger.info("="*80)
+    logger.info("✅ MCP server verification complete - proceeding with startup")
+    logger.info("="*80)
+    return True
+
+
+def main():
+    """Main function to start the A2A server."""
     
     # Initialize AgentOps for monitoring and observability
     agentops_session_id = None
@@ -423,65 +752,37 @@ async def main():
     host = os.environ.get("HOST", "0.0.0.0")
     port = int(os.environ.get("PORT", 8000))
     
+    # MCP was already verified by verify_mcp_before_start()
     # Create the application
     app = create_app()
     
-    # Configure Hypercorn for HTTP/2 support
-    config = Config()
-    config.bind = [f"{host}:{port}"]
-    config.alpn_protocols = ["h2", "http/1.1"]  # Support both HTTP/2 and HTTP/1.1
-    
-    # Enable access logging if debug mode is on
-    if os.environ.get("DEBUG_PROTOCOL", "false").lower() == "true":
-        config.accesslog = "-"  # Log to stdout
-        config.errorlog = "-"   # Log errors to stdout
-        config.loglevel = "DEBUG"
-        logger.info("Hypercorn access logging enabled")
-    
-    # Enable HTTP/2
-    config.h2_max_concurrent_streams = 100
-    config.h2_max_header_list_size = 8192
-    config.h2_max_inbound_frame_size = 16384
-    config.h2_initial_connection_window_size = 65536
-    
-    # Connection settings for high performance
-    config.keep_alive_timeout = 300  # 5 minutes
-    config.max_requests = 10000  # Max requests per connection
-    config.max_requests_jitter = 1000  # Add jitter to prevent thundering herd
-    
-    # SSL/TLS configuration for production (optional)
-    if os.environ.get("USE_TLS", "false").lower() == "true":
-        cert_path = os.environ.get("TLS_CERT_PATH", "cert.pem")
-        key_path = os.environ.get("TLS_KEY_PATH", "key.pem")
-        
-        # Generate self-signed certificates for local development if needed
-        if os.environ.get("GENERATE_CERTS", "true").lower() == "true":
-            generate_self_signed_cert(cert_path, key_path)
-        
-        if Path(cert_path).exists() and Path(key_path).exists():
-            config.certfile = cert_path
-            config.keyfile = key_path
-            config.alpn_protocols = ["h2", "http/1.1"]
-            logger.info("TLS enabled with HTTP/2 support")
-            logger.info("Using HTTPS - connect to https://localhost:8000")
-        else:
-            logger.warning("TLS requested but certificates not found")
-    
     # Log startup information
-    logger.info(f"Starting {{ agent_name }} A2A Server with HTTP/2 support")
+    logger.info(f"Starting {{ agent_name }} A2A Server")
     logger.info(f"MCP Server: {{ mcp_server_name }}")
     logger.info(f"Listening on {host}:{port}")
-    logger.info(f"Agent Card available at: http://{host}:{port}/.well-known/agent.json")
+    logger.info(f"Agent Card available at: http://{host}:{port}/.well-known/agent-card.json")
+    logger.info(f"A2A Endpoint: http://{host}:{port}/a2a")
     logger.info(f"SSE endpoints available at: /a2a/tasks/subscribe and /a2a/tasks/resubscribe")
-    logger.info(f"HTTP/2 multiplexing enabled with max {config.h2_max_concurrent_streams} concurrent streams")
     
     if agentops_session_id:
         logger.info(f"📊 AgentOps Session ID: {agentops_session_id}")
         logger.info("📊 Monitor agent performance at: https://app.agentops.ai")
     
+    # Configure Uvicorn (same as MCP servers for D402 compatibility)
+    uvicorn_config = {
+        "host": host,
+        "port": port,
+        "log_level": "info"
+    }
+    
+    # Enable access logging if debug mode is on
+    if os.environ.get("DEBUG_PROTOCOL", "false").lower() == "true":
+        uvicorn_config["log_level"] = "debug"
+        logger.info("Debug logging enabled")
+    
     try:
-        # Run the server with Hypercorn
-        await serve(app, config)
+        # Run the server with Uvicorn (for D402 middleware compatibility)
+        uvicorn.run(app, **uvicorn_config)
     except KeyboardInterrupt:
         logger.info("🛑 Server shutdown requested")
         if AGENTOPS_AVAILABLE and agentops_session_id:
@@ -502,5 +803,12 @@ async def main():
 
 
 if __name__ == "__main__":
+    # Run async MCP verification first
     import asyncio
-    asyncio.run(main()) 
+    mcp_verified = asyncio.run(verify_mcp_before_start())
+    
+    if not mcp_verified:
+        sys.exit(1)
+    
+    # Then run synchronous main (Uvicorn handles its own event loop)
+    main() 

traia_iatp/server/templates/agent.py.j2

@@ -6,6 +6,7 @@ Auto-generated for the {{ agent_name }} utility agent.
 """
 
 import logging
+import datetime
 from typing import List, Dict, Any, Optional
 from crewai import Agent, Task, Crew, LLM
 import os
@@ -14,19 +15,18 @@ import os
 from traia_iatp.mcp import MCPServerConfig, MCPAgentBuilder, run_with_mcp_tools, MCPServerInfo
 
 # Import AgentOps for operation tracking
-try:
-    import agentops
-    from agentops.sdk.decorators import operation
-    AGENTOPS_AVAILABLE = True
-except ImportError:
-    AGENTOPS_AVAILABLE = False
-    agentops = None
-    # Create a no-op decorator if AgentOps is not available
-    def operation(func):
-        return func
+import agentops
+
+DEFAULT_LLM = LLM(
+    model=os.getenv("LLM_MODEL", "gpt-4.1-nano"),  # Using environment variable with fallback
+    temperature=float(os.getenv("LLM_MODEL_TEMPERATURE", "0.1")),
+    api_key=os.getenv("OPENAI_API_KEY")
+)
+current_time = datetime.datetime.utcnow()
 
 logger = logging.getLogger(__name__)
 
+logger.info(f"Current LLM model used: {os.getenv("LLM_MODEL", "gpt-4.1-nano")}")
 
 class {{ class_name }}Agent:
     """{{ agent_name }} agent that processes requests using {{ mcp_server_name }}."""
@@ -61,10 +61,11 @@ class {{ class_name }}Agent:
             ),
             verbose=True,
             allow_delegation=False,
+            llm=DEFAULT_LLM,
             tools_subset=tools_subset
         )
     
-    @operation
+    agentops.init(trace_name=f"Utility Agent for MCP: {self.mcp_config.name}", default_tags=[f"current time: {current_time}"])
     def process_request(self, request: str, context: Dict[str, Any] = None) -> str:
         """Process a request using the MCP server capabilities."""
         try:

traia_iatp/server/templates/agent_config.json.j2

@@ -3,9 +3,9 @@
     "description": "{{ agent_description }}",
     "version": "{{ agent_version }}",
     "agent_id": "{{ agent_id }}",
-    "mcp_server": {
-        "name": "{{ mcp_server_name }}",
-        "url": "{{ mcp_server_url }}",
+  "mcp_server": {
+    "name": "{{ mcp_server_name }}",
+    "url": "{{ mcp_server_url | replace('localhost', 'host.docker.internal') }}",
         "description": "{{ mcp_server_description }}",
         "server_type": "{{ mcp_server_type }}",
         "capabilities": {{ mcp_server_capabilities | tojson }},