server.py

"""
DCAP Disco MCP Server

This server implements an MCP (Model Context Protocol) server that connects to the DCAP intelligence stream
and provides real-time access to tool discovery data.

Features:
- Connects to DCAP WebSocket stream at ws://159.89.110.236:10191
- Maintains FIFO cache of up to 100 most recently discovered tools
- Exposes MCP tools for querying tool capabilities, performance, and execution sequences
- Auto-reconnects if connection drops
- Uses asyncio for efficient asynchronous operations.

"""

import asyncio
import websockets
import json
import logging
import sys
import os
import base64
import socket
import time
from collections import OrderedDict
from datetime import datetime
from typing import Any, Dict, List, Optional
from pathlib import Path
from dotenv import load_dotenv
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
import httpx
from httpx_sse import aconnect_sse
from x402.exact import prepare_payment_header, sign_payment_header
from x402.types import PaymentRequirements
from eth_account import Account

# Logging setup (must be before .env loading)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]
)
logger = logging.getLogger("dcap-disco")

# Load .env file from project root
env_path = Path(__file__).parent / '.env'
if env_path.exists():
    load_dotenv(env_path)
    logger.info(f"📄 Loaded environment from {env_path}")
else:
    logger.warning(f"⚠️ No .env file found at {env_path}")

# Configuration
DCAP_STREAM_URL = "ws://159.89.110.236:10191"
DCAP_HUB_IP = "159.89.110.236"
DCAP_HUB_PORT = 10191
MAX_TOOLS_CACHE = 100
RECONNECT_DELAY = 5  # seconds
SERVER_ID = "dcap-disco"  # Our server identifier

# Store connected clients (for original WebSocket server)
clients = set()


class DCAPBroadcaster:
    """Broadcasts DCAP messages via UDP"""
    
    def __init__(self, server_id: str, hub_ip: str = DCAP_HUB_IP, hub_port: int = DCAP_HUB_PORT):
        self.server_id = server_id
        self.hub_ip = hub_ip
        self.hub_port = hub_port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        logger.info(f"📡 DCAP broadcaster initialized: {server_id} -> {hub_ip}:{hub_port}")
    
    def send_semantic_discover(self, tool_name: str, description: str) -> None:
        """Broadcast semantic_discover message for a tool"""
        try:
            message = {
                "v": 2,
                "t": "semantic_discover",
                "ts": int(time.time()),
                "sid": self.server_id,
                "tool": tool_name,
                "does": description,
                "when": [],  # Will be filled by Semanticore
                "good_at": [],
                "bad_at": [],
                "connector": {
                    "transport": "stdio",  # DCAP Disco runs via stdio MCP
                    "endpoint": "python /path/to/server.py",  # Placeholder - user configures this
                    "protocol": {
                        "type": "mcp",
                        "version": "2024-11-05",
                        "methods": ["tools/list", "tools/call"]
                    },
                    "auth": {
                        "type": "none",
                        "required": False
                    }
                }
            }
            
            json_msg = json.dumps(message)
            self.sock.sendto(json_msg.encode('utf-8'), (self.hub_ip, self.hub_port))
            logger.debug(f"📤 Broadcast semantic_discover: {tool_name}")
            
        except Exception as e:
            logger.warning(f"Failed to broadcast semantic_discover for {tool_name}: {e}")
    
    def send_perf_update(self, tool_name: str, exec_ms: int, success: bool) -> None:
        """Broadcast perf_update message for a tool execution"""
        try:
            message = {
                "v": 2,
                "t": "perf_update",
                "ts": int(time.time()),
                "sid": self.server_id,
                "tool": tool_name,
                "exec_ms": exec_ms,
                "success": success
            }
            
            json_msg = json.dumps(message)
            self.sock.sendto(json_msg.encode('utf-8'), (self.hub_ip, self.hub_port))
            logger.debug(f"📊 Broadcast perf_update: {tool_name} ({exec_ms}ms, success={success})")
            
        except Exception as e:
            logger.warning(f"Failed to broadcast perf_update for {tool_name}: {e}")
    
    def broadcast_all_tools(self) -> None:
        """Broadcast semantic_discover for all our MCP tools"""
        tools = [
            ("list_discovered_tools", "Get all tools discovered from the DCAP intelligence stream with their capabilities and connection endpoints"),
            ("search_tools_by_capability", "Search for tools by capability keywords in tool descriptions and semantic fields"),
            ("get_tool_details", "Get detailed information about a specific tool including performance history and connection details"),
            ("get_recent_activity", "Get recent performance updates from all tools showing real-time execution data"),
            ("get_tool_sequences", "Get observed tool execution sequences and usage patterns"),
            ("get_stream_status", "Get status of the DCAP intelligence stream connection and cache statistics"),
            ("call_discovered_tool", "Autonomously call any discovered tool at runtime with automatic connection and x402 payment handling")
        ]
        
        for tool_name, description in tools:
            self.send_semantic_discover(tool_name, description)
        
        logger.info(f"✅ Broadcast {len(tools)} tool capabilities to DCAP network")


class DCAPCache:
    """FIFO cache for tool discoveries with deduplication"""
    
    def __init__(self, max_size: int = MAX_TOOLS_CACHE):
        self.max_size = max_size
        self.tools: OrderedDict[str, Dict[str, Any]] = OrderedDict()
        self.perf_updates: List[Dict[str, Any]] = []
        self.sequences: List[Dict[str, Any]] = []
        self.error_patterns: List[Dict[str, Any]] = []
        self.stream_stats = {
            "connected": False,
            "messages_received": 0,
            "connection_time": None,
            "last_message_time": None
        }
    
    def add_tool_discovery(self, message: Dict[str, Any]) -> None:
        """Add or upgrade tool entry with semantic_discover message"""
        tool_name = message.get("tool")
        if not tool_name:
            return
        
        # If tool exists (possibly as stub from perf_update), upgrade it
        if tool_name in self.tools:
            existing = self.tools[tool_name]
            # Check if this was a stub entry (missing connector)
            was_stub = not existing.get("connector")
            
            self.tools.move_to_end(tool_name)
            self.tools[tool_name] = message
            
            if was_stub:
                logger.info(f"⬆️  Upgraded stub tool to full entry: {tool_name}")
        else:
            # Add new tool
            self.tools[tool_name] = message
            
            # Remove oldest if we exceed max_size
            if len(self.tools) > self.max_size:
                oldest = next(iter(self.tools))
                del self.tools[oldest]
                logger.info(f"Cache full. Removed oldest tool: {oldest}")
    
    def add_perf_update(self, message: Dict[str, Any]) -> None:
        """Add perf_update to recent history and create stub tool entry if needed"""
        tool_name = message.get("tool")
        sid = message.get("sid")
        
        # Add to performance history
        self.perf_updates.append(message)
        if len(self.perf_updates) > 500:
            self.perf_updates = self.perf_updates[-500:]
        
        # Create stub tool entry if we haven't seen semantic_discover yet
        if tool_name and sid and tool_name not in self.tools:
            stub_entry = {
                "v": message.get("v", 2),
                "t": "stub",  # Mark as stub entry
                "ts": message.get("ts"),
                "sid": sid,
                "tool": tool_name,
                "does": "(Performance data only - awaiting semantic_discover)",
                "when": [],
                "good_at": [],
                "bad_at": [],
                # No connector - this is what makes it a stub
                "stub": True,  # Explicit flag for easy detection
                "first_seen": message.get("ts"),
                "perf_samples": 1
            }
            
            self.tools[tool_name] = stub_entry
            logger.info(f"📊 Created stub entry from perf_update: {tool_name}@{sid}")
            
            # Remove oldest if we exceed max_size
            if len(self.tools) > self.max_size:
                oldest = next(iter(self.tools))
                del self.tools[oldest]
                logger.info(f"Cache full. Removed oldest tool: {oldest}")
        
        # Update perf_samples counter for stub entries
        elif tool_name in self.tools and self.tools[tool_name].get("stub"):
            self.tools[tool_name]["perf_samples"] = self.tools[tool_name].get("perf_samples", 0) + 1
    
    def add_sequence(self, message: Dict[str, Any]) -> None:
        """Add observed_sequence to history (keep last 100)"""
        self.sequences.append(message)
        if len(self.sequences) > 100:
            self.sequences = self.sequences[-100:]
    
    def add_error_pattern(self, message: Dict[str, Any]) -> None:
        """Add error_pattern to history (keep last 100)"""
        self.error_patterns.append(message)
        if len(self.error_patterns) > 100:
            self.error_patterns = self.error_patterns[-100:]
    
    def get_all_tools(self) -> List[Dict[str, Any]]:
        """Get all cached tools (including stubs)"""
        return list(self.tools.values())
    
    def get_stats(self) -> Dict[str, int]:
        """Get cache statistics"""
        stubs = sum(1 for tool in self.tools.values() if tool.get("stub"))
        full = len(self.tools) - stubs
        return {
            "total_tools": len(self.tools),
            "full_entries": full,
            "stub_entries": stubs
        }
    
    def get_tool(self, tool_name: str) -> Optional[Dict[str, Any]]:
        """Get specific tool by name"""
        return self.tools.get(tool_name)
    
    def search_tools(self, query: str) -> List[Dict[str, Any]]:
        """Simple keyword search in tool capabilities"""
        query_lower = query.lower()
        results = []
        
        for tool in self.tools.values():
            # Search in: tool name, does, when, good_at
            searchable = [
                tool.get("tool", ""),
                tool.get("does", ""),
                " ".join(tool.get("when", [])),
                " ".join(tool.get("good_at", []))
            ]
            
            if any(query_lower in text.lower() for text in searchable):
                results.append(tool)
        
        return results
    
    def get_tool_performance(self, tool_name: str) -> List[Dict[str, Any]]:
        """Get performance updates for a specific tool"""
        return [
            perf for perf in self.perf_updates
            if perf.get("tool") == tool_name
        ]


class DCAPStreamClient:
    """WebSocket client for DCAP intelligence stream"""
    
    def __init__(self, url: str, cache: DCAPCache):
        self.url = url
        self.cache = cache
        self.ws = None
        self.running = False
    
    async def connect(self) -> None:
        """Connect to DCAP stream with auto-reconnect"""
        self.running = True
        
        while self.running:
            try:
                logger.info(f"Connecting to DCAP stream: {self.url}")
                async with websockets.connect(self.url) as websocket:
                    self.ws = websocket
                    self.cache.stream_stats["connected"] = True
                    self.cache.stream_stats["connection_time"] = datetime.now().isoformat()
                    logger.info("Connected to DCAP stream")
                    
                    async for message in websocket:
                        await self.handle_message(message)
                        
            except websockets.exceptions.WebSocketException as e:
                logger.error(f"WebSocket error: {e}")
                self.cache.stream_stats["connected"] = False
            except Exception as e:
                logger.error(f"Unexpected error: {e}")
                self.cache.stream_stats["connected"] = False
            
            if self.running:
                logger.info(f"Reconnecting in {RECONNECT_DELAY} seconds...")
                await asyncio.sleep(RECONNECT_DELAY)
    
    async def handle_message(self, message: str) -> None:
        """Process incoming DCAP message"""
        try:
            data = json.loads(message)
            self.cache.stream_stats["messages_received"] += 1
            self.cache.stream_stats["last_message_time"] = datetime.now().isoformat()
            
            # Handle welcome message
            if data.get("type") == "welcome":
                logger.info(f"Welcome: {data.get('message')}")
                return
            
            # Extract DCAP message from wrapper
            dcap_data = data.get("data", data)
            message_type = dcap_data.get("t")
            
            if message_type == "semantic_discover":
                self.cache.add_tool_discovery(dcap_data)
                logger.info(f"✅ Discovered tool: {dcap_data.get('tool')}")
            
            elif message_type == "perf_update":
                self.cache.add_perf_update(dcap_data)
                if self.cache.stream_stats["messages_received"] % 100 == 0:
                    logger.info(f"📊 Perf update: {dcap_data.get('tool')}")
            
            elif dcap_data.get("type") == "observed_sequence":
                self.cache.add_sequence(dcap_data)
            
            elif message_type == "error_pattern":
                self.cache.add_error_pattern(dcap_data)
                
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse message: {e}")
        except Exception as e:
            logger.error(f"Error handling message: {e}")
    
    async def disconnect(self) -> None:
        """Gracefully disconnect"""
        self.running = False
        if self.ws:
            await self.ws.close()


# Global cache and broadcaster instances
cache = DCAPCache()
broadcaster = DCAPBroadcaster(SERVER_ID)


class ToolConnectionManager:
    """Manages lazy connections to discovered MCP tools"""
    
    def __init__(self, cache: DCAPCache):
        self.cache = cache
        self.connections: Dict[str, Any] = {}  # endpoint → client
        self.connection_locks: Dict[str, asyncio.Lock] = {}  # endpoint → lock
        logger.info("ToolConnectionManager initialized")
    
    async def get_tool_info(self, tool_identifier: str) -> tuple[Optional[Dict], Optional[str], Optional[str]]:
        """
        Parse tool identifier and retrieve metadata.
        
        Args:
            tool_identifier: Either "tool_name@server_id" or just "tool_name"
        
        Returns:
            (tool_metadata, tool_name, server_id) or (None, None, None) if not found
        """
        if "@" in tool_identifier:
            tool_name, server_id = tool_identifier.split("@", 1)
        else:
            tool_name = tool_identifier
            server_id = None
        
        # Search cache for matching tool
        for cached_tool in self.cache.tools.values():
            cached_name = cached_tool.get("tool")
            cached_sid = cached_tool.get("sid")
            
            if cached_name == tool_name:
                if server_id is None or cached_sid == server_id:
                    return cached_tool, cached_name, cached_sid
        
        return None, None, None
    
    async def _create_x402_payment_from_requirements(self, payment_requirements: Dict) -> Optional[Dict]:
        """
        Create x402 payment from server-provided payment requirements.
        
        Args:
            payment_requirements: Payment requirements dict from 402 response
        
        Returns:
            Payment dict or None if failed
        """
        try:
            # Read private key from environment
            private_key = os.getenv("X402_PRIVATE_KEY")
            if not private_key:
                logger.error("X402_PRIVATE_KEY not set in environment")
                return None
            
            # Create eth account from private key
            if not private_key.startswith("0x"):
                private_key = f"0x{private_key}"
            
            account = Account.from_key(private_key)
            payer_address = account.address
            
            # Convert payment requirements to PaymentRequirements object
            payment_req = PaymentRequirements(
                scheme=payment_requirements["scheme"],
                network=payment_requirements["network"],
                maxAmountRequired=payment_requirements["maxAmountRequired"],
                resource=payment_requirements["resource"],
                description=payment_requirements["description"],
                mimeType=payment_requirements["mimeType"],
                payTo=payment_requirements["payTo"],
                maxTimeoutSeconds=payment_requirements["maxTimeoutSeconds"],
                asset=payment_requirements["asset"],
                extra=payment_requirements.get("extra", {})
            )
            
            # Prepare unsigned payment header
            unsigned_header = prepare_payment_header(
                sender_address=payer_address,
                x402_version=1,
                payment_requirements=payment_req
            )
            
            # Fix nonce format: x402 library returns bytes but sign_payment_header expects hex string
            if isinstance(unsigned_header['payload']['authorization']['nonce'], bytes):
                unsigned_header['payload']['authorization']['nonce'] = unsigned_header['payload']['authorization']['nonce'].hex()
            
            # Sign the payment header
            signature = sign_payment_header(
                account=account,
                payment_requirements=payment_req,
                header=unsigned_header
            )
            
            # The signature IS the payment - it's already a dict, not encoded
            logger.info(f"✅ Payment created for {payment_req.pay_to}")
            return signature  # Return the signed payment dict directly
            
        except Exception as e:
            logger.error(f"Failed to create x402 payment: {e}")
            logger.exception(e)
            return None
    
    
    async def call_tool_via_sse(
        self, 
        endpoint: str, 
        tool_name: str, 
        arguments: Dict,
        payment: Optional[Dict] = None
    ) -> Dict:
        """
        Call an MCP tool via SSE transport with x402 payment support.
        
        Follows x402 protocol:
        1. Call tool without payment
        2. If 402 received, extract payment requirements
        3. Create payment and retry
        
        Args:
            endpoint: MCP endpoint URL
            tool_name: Name of the tool to call
            arguments: Tool arguments
            payment: Optional pre-created payment dict
        
        Returns:
            Tool execution result
        """
        # MCP over SSE/HTTP: POST to endpoint with tool call request
        request_payload = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        
        headers = {
            "Content-Type": "application/json",
            "X-Caller-ID": "dcap-disco-agent"
        }
        
        # Add payment if provided
        if payment:
            headers["X-Payment"] = json.dumps(payment)
        
        try:
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(
                    endpoint,
                    json=request_payload,
                    headers=headers
                )
                
                # Handle successful response
                if response.status_code == 200:
                    result = response.json()
                    if "result" in result:
                        return {"success": True, "result": result["result"]}
                    elif "error" in result:
                        return {
                            "success": False,
                            "error": result["error"].get("message", "Unknown error"),
                            "code": result["error"].get("code")
                        }
                
                # Handle payment required (402)
                elif response.status_code == 402:
                    x402_response = response.json()
                    logger.info(f"💰 Payment required for {tool_name}")
                    
                    if "accepts" not in x402_response or len(x402_response["accepts"]) == 0:
                        return {
                            "success": False,
                            "error": "Payment required but no payment requirements provided"
                        }
                    
                    # Get payment requirements
                    payment_req = x402_response["accepts"][0]
                    logger.info(f"   💵 Price: ${float(payment_req['maxAmountRequired'])/1_000_000:.4f}")
                    logger.info(f"   🏦 Pay to: {payment_req['payTo']}")
                    
                    # Create payment
                    created_payment = await self._create_x402_payment_from_requirements(payment_req)
                    if not created_payment:
                        return {
                            "success": False,
                            "error": "Failed to create x402 payment"
                        }
                    
                    # Retry with payment
                    logger.info(f"🔐 Retrying {tool_name} with payment...")
                    return await self.call_tool_via_sse(endpoint, tool_name, arguments, created_payment)
                
                else:
                    return {
                        "success": False,
                        "error": f"HTTP {response.status_code}: {response.text}"
                    }
                
        except httpx.TimeoutException:
            logger.error(f"Timeout calling {tool_name} at {endpoint}")
            return {"success": False, "error": "Request timeout"}
        except Exception as e:
            logger.error(f"Error calling {tool_name} via SSE: {e}")
            logger.exception(e)
            return {"success": False, "error": str(e)}
    
    async def call_discovered_tool(
        self, 
        tool_identifier: str, 
        arguments: Dict
    ) -> Dict:
        """
        Main entry point: Call a discovered tool.
        
        Args:
            tool_identifier: "tool_name@server_id" or "tool_name"
            arguments: Tool arguments
        
        Returns:
            Tool execution result
        """
        # 1. Lookup tool in cache
        tool_metadata, tool_name, server_id = await self.get_tool_info(tool_identifier)
        
        if not tool_metadata:
            return {
                "success": False,
                "error": f"Tool '{tool_identifier}' not found in cache"
            }
        
        # Check if this is a stub entry (no connector info yet)
        if tool_metadata.get("stub"):
            return {
                "success": False,
                "error": f"Tool '{tool_identifier}' is a stub entry (only performance data available)",
                "hint": "Waiting for semantic_discover message with connector details",
                "perf_samples": tool_metadata.get("perf_samples", 0)
            }
        
        connector = tool_metadata.get("connector")
        if not connector:
            return {
                "success": False,
                "error": f"Tool '{tool_identifier}' has no connector information"
            }
        
        transport = connector.get("transport")
        endpoint = connector.get("endpoint")
        
        logger.info(f"Calling {tool_name}@{server_id} via {transport} at {endpoint}")
        
        # Note: x402 authentication is handled via the 402 flow inside call_tool_via_sse
        # No need to pre-create payment - server will tell us payment requirements
        
        # Route based on transport
        if transport in ["sse", "http"]:
            return await self.call_tool_via_sse(endpoint, tool_name, arguments)
        
        elif transport == "stdio":
            return {
                "success": False,
                "error": "stdio transport not yet supported for remote calls"
            }
        
        else:
            return {
                "success": False,
                "error": f"Unsupported transport: {transport}"
            }


# Global connection manager
connection_manager = ToolConnectionManager(cache)


def create_mcp_server() -> Server:
    """Create and configure MCP server"""
    server = Server("dcap-disco")
    
    @server.list_tools()
    async def list_tools() -> List[Tool]:
        """List available MCP tools"""
        return [
            Tool(
                name="list_discovered_tools",
                description="Get all tools discovered from the DCAP intelligence stream. Returns up to 100 most recently seen tools with their capabilities, performance stats, and connection endpoints.",
                inputSchema={
                    "type": "object",
                    "properties": {},
                }
            ),
            Tool(
                name="search_tools_by_capability",
                description="Search for tools by capability keywords. Searches in tool name, description (does), trigger contexts (when), and strengths (good_at).",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Search query (keywords to search for in tool capabilities)"
                        }
                    },
                    "required": ["query"]
                }
            ),
            Tool(
                name="get_tool_details",
                description="Get detailed information about a specific tool including its capabilities, performance history, and connection endpoint.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "tool_name": {
                            "type": "string",
                            "description": "Name of the tool to get details for"
                        }
                    },
                    "required": ["tool_name"]
                }
            ),
            Tool(
                name="get_recent_activity",
                description="Get recent performance updates from all tools. Shows real-time execution data including timing, success rates, and costs.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "limit": {
                            "type": "number",
                            "description": "Maximum number of recent updates to return (default: 20, max: 100)",
                            "default": 20
                        },
                        "tool_name": {
                            "type": "string",
                            "description": "Optional: filter by specific tool name"
                        }
                    }
                }
            ),
            Tool(
                name="get_tool_sequences",
                description="Get observed tool execution sequences (chains). Shows patterns of tools used together and their combined performance.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "limit": {
                            "type": "number",
                            "description": "Maximum number of sequences to return (default: 10)",
                            "default": 10
                        }
                    }
                }
            ),
            Tool(
                name="get_stream_status",
                description="Get status of the DCAP intelligence stream connection including uptime, message counts, and connection health.",
                inputSchema={
                    "type": "object",
                    "properties": {},
                }
            ),
            Tool(
                name="call_discovered_tool",
                description="Call a discovered tool at runtime. This enables autonomous tool acquisition - you can invoke ANY tool discovered via DCAP without manual configuration. Supports x402 payment for paid tools. Use format 'tool_name@server_id' or just 'tool_name'.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "tool_identifier": {
                            "type": "string",
                            "description": "Tool identifier in format 'tool_name@server_id' (e.g., 'validate_daml_business_logic@canton-mcp') or just 'tool_name'"
                        },
                        "arguments": {
                            "type": "object",
                            "description": "Arguments to pass to the tool (tool-specific schema)"
                        }
                    },
                    "required": ["tool_identifier", "arguments"]
                }
            )
        ]
    
    @server.call_tool()
    async def call_tool(name: str, arguments: Any) -> List[TextContent]:
        """Handle tool calls"""
        start_time = time.time()
        success = True
        
        try:
            result = await _handle_tool_call(name, arguments)
            return result
        except Exception as e:
            success = False
            logger.error(f"Error handling tool {name}: {e}")
            return [TextContent(
                type="text",
                text=json.dumps({"error": str(e)})
            )]
        finally:
            # Broadcast performance update
            exec_ms = int((time.time() - start_time) * 1000)
            broadcaster.send_perf_update(name, exec_ms, success)
    
    async def _handle_tool_call(name: str, arguments: Any) -> List[TextContent]:
        """Internal tool call handler"""
        
        if name == "list_discovered_tools":
            tools = cache.get_all_tools()
            stats = cache.get_stats()
            return [TextContent(
                type="text",
                text=json.dumps({
                    "total_tools": len(tools),
                    "full_entries": stats["full_entries"],
                    "stub_entries": stats["stub_entries"],
                    "tools": tools
                }, indent=2)
            )]
        
        elif name == "search_tools_by_capability":
            query = arguments.get("query", "")
            results = cache.search_tools(query)
            return [TextContent(
                type="text",
                text=json.dumps({
                    "query": query,
                    "matches": len(results),
                    "tools": results
                }, indent=2)
            )]
        
        elif name == "get_tool_details":
            tool_name = arguments.get("tool_name")
            tool = cache.get_tool(tool_name)
            
            if not tool:
                return [TextContent(
                    type="text",
                    text=json.dumps({
                        "error": f"Tool '{tool_name}' not found in cache"
                    })
                )]
            
            perf_history = cache.get_tool_performance(tool_name)
            
            return [TextContent(
                type="text",
                text=json.dumps({
                    "tool": tool,
                    "performance_history": perf_history[-20:]  # Last 20 updates
                }, indent=2)
            )]
        
        elif name == "get_recent_activity":
            limit = min(arguments.get("limit", 20), 100)
            tool_name = arguments.get("tool_name")
            
            updates = cache.perf_updates[-limit:]
            if tool_name:
                updates = [u for u in updates if u.get("tool") == tool_name]
            
            return [TextContent(
                type="text",
                text=json.dumps({
                    "count": len(updates),
                    "updates": updates
                }, indent=2)
            )]
        
        elif name == "get_tool_sequences":
            limit = arguments.get("limit", 10)
            sequences = cache.sequences[-limit:]
            
            return [TextContent(
                type="text",
                text=json.dumps({
                    "count": len(sequences),
                    "sequences": sequences
                }, indent=2)
            )]
        
        elif name == "get_stream_status":
            stats = cache.stream_stats.copy()
            cache_stats = cache.get_stats()
            stats["cached_tools"] = len(cache.tools)
            stats["full_entries"] = cache_stats["full_entries"]
            stats["stub_entries"] = cache_stats["stub_entries"]
            stats["recent_updates"] = len(cache.perf_updates)
            stats["observed_sequences"] = len(cache.sequences)
            
            return [TextContent(
                type="text",
                text=json.dumps(stats, indent=2)
            )]
        
        elif name == "call_discovered_tool":
            tool_identifier = arguments.get("tool_identifier")
            tool_arguments = arguments.get("arguments", {})
            
            if not tool_identifier:
                return [TextContent(
                    type="text",
                    text=json.dumps({
                        "error": "Missing required parameter: tool_identifier"
                    })
                )]
            
            # Call the discovered tool
            result = await connection_manager.call_discovered_tool(
                tool_identifier, 
                tool_arguments
            )
            
            return [TextContent(
                type="text",
                text=json.dumps(result, indent=2)
            )]
        
        else:
            return [TextContent(
                type="text",
                text=json.dumps({"error": f"Unknown tool: {name}"})
            )]
    
    return server


# WebSocket Handler for Push Updates (original functionality preserved)
async def websocket_handler(websocket, path):
    clients.add(websocket)
    try:
        async for message in websocket:
            data = json.loads(message)
            if data.get("action") == "subscribe":
                print(f"Client subscribed: {websocket.remote_address}")
    except websockets.exceptions.ConnectionClosed:
        pass
    finally:
        clients.remove(websocket)

# Function to Push Updates to Clients (original functionality preserved)
async def push_updates(data):
    if clients:
        message = json.dumps({"event": "new_data", "payload": data})
        await asyncio.gather(*(client.send(message) for client in clients))

async def periodic_broadcast():
    """Periodically broadcast our capabilities every 5 minutes"""
    while True:
        await asyncio.sleep(300)  # 5 minutes
        broadcaster.broadcast_all_tools()
        logger.info("🔄 Periodic capability broadcast sent")

# Main entry point
async def main():
    """Main entry point"""
    logger.info("Starting DCAP Disco MCP Server")
    
    # Create stream client
    stream_client = DCAPStreamClient(DCAP_STREAM_URL, cache)
    
    # Start stream connection in background
    stream_task = asyncio.create_task(stream_client.connect())
    
    # Wait a moment for stream to connect before broadcasting
    await asyncio.sleep(2)
    
    # Broadcast our capabilities to DCAP network (initial + periodic)
    broadcaster.broadcast_all_tools()
    periodic_task = asyncio.create_task(periodic_broadcast())
    
    # Create and run MCP server
    server = create_mcp_server()
    
    try:
        async with stdio_server() as (read_stream, write_stream):
            await server.run(
                read_stream,
                write_stream,
                server.create_initialization_options()
            )
    except KeyboardInterrupt:
        logger.info("Shutting down...")
    finally:
        await stream_client.disconnect()
        stream_task.cancel()
        periodic_task.cancel()

if __name__ == "__main__":
    asyncio.run(main())