code-puppy 0.0.130__py3-none-any.whl → 0.0.132__py3-none-any.whl

code_puppy/mcp/blocking_startup.py

@@ -0,0 +1,404 @@
+"""
+MCP Server with blocking startup capability and stderr capture.
+
+This module provides MCP servers that:
+1. Capture stderr output from stdio servers
+2. Block until fully initialized before allowing operations
+3. Emit stderr to users via emit_info with message groups
+"""
+
+import asyncio
+import os
+import tempfile
+import threading
+import uuid
+from typing import Optional, Callable, List
+from contextlib import asynccontextmanager
+from pydantic_ai.mcp import MCPServerStdio
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.shared.session import SessionMessage
+from code_puppy.messaging import emit_info
+
+
+class StderrFileCapture:
+    """Captures stderr to a file and monitors it in a background thread."""
+    
+    def __init__(self, server_name: str, emit_to_user: bool = True, message_group: Optional[uuid.UUID] = None):
+        self.server_name = server_name
+        self.emit_to_user = emit_to_user
+        self.message_group = message_group or uuid.uuid4()
+        self.temp_file = None
+        self.temp_path = None
+        self.monitor_thread = None
+        self.stop_monitoring = threading.Event()
+        self.captured_lines = []
+        
+    def start(self):
+        """Start capture by creating temp file and monitor thread."""
+        # Create temp file
+        self.temp_file = tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix='.err')
+        self.temp_path = self.temp_file.name
+        
+        # Start monitoring thread
+        self.stop_monitoring.clear()
+        self.monitor_thread = threading.Thread(target=self._monitor_file)
+        self.monitor_thread.daemon = True
+        self.monitor_thread.start()
+        
+        return self.temp_file
+        
+    def _monitor_file(self):
+        """Monitor the temp file for new content."""
+        if not self.temp_path:
+            return
+            
+        last_pos = 0
+        while not self.stop_monitoring.is_set():
+            try:
+                with open(self.temp_path, 'r') as f:
+                    f.seek(last_pos)
+                    new_content = f.read()
+                    if new_content:
+                        last_pos = f.tell()
+                        # Process new lines
+                        for line in new_content.splitlines():
+                            if line.strip():
+                                self.captured_lines.append(line)
+                                if self.emit_to_user:
+                                    emit_info(
+                                        f"[bold white on blue] MCP {self.server_name} [/bold white on blue] {line}", 
+                                        style="dim cyan",
+                                        message_group=self.message_group
+                                    )
+                                    
+            except Exception:
+                pass  # File might not exist yet or be deleted
+                
+            self.stop_monitoring.wait(0.1)  # Check every 100ms
+            
+    def stop(self):
+        """Stop monitoring and clean up."""
+        self.stop_monitoring.set()
+        if self.monitor_thread:
+            self.monitor_thread.join(timeout=1)
+            
+        if self.temp_file:
+            try:
+                self.temp_file.close()
+            except:
+                pass
+                
+        if self.temp_path and os.path.exists(self.temp_path):
+            try:
+                # Read any remaining content
+                with open(self.temp_path, 'r') as f:
+                    content = f.read()
+                    for line in content.splitlines():
+                        if line.strip() and line not in self.captured_lines:
+                            self.captured_lines.append(line)
+                            if self.emit_to_user:
+                                emit_info(
+                                    f"[bold white on blue] MCP {self.server_name} [/bold white on blue] {line}", 
+                                    style="dim cyan",
+                                    message_group=self.message_group
+                                )
+                                
+                os.unlink(self.temp_path)
+            except:
+                pass
+                
+    def get_captured_lines(self) -> List[str]:
+        """Get all captured lines."""
+        return self.captured_lines.copy()
+
+
+class SimpleCapturedMCPServerStdio(MCPServerStdio):
+    """
+    MCPServerStdio that captures stderr to a file and optionally emits to user.
+    """
+    
+    def __init__(
+        self,
+        command: str,
+        args=(),
+        env=None,
+        cwd=None,
+        emit_stderr: bool = True,
+        message_group: Optional[uuid.UUID] = None,
+        **kwargs
+    ):
+        super().__init__(command=command, args=args, env=env, cwd=cwd, **kwargs)
+        self.emit_stderr = emit_stderr
+        self.message_group = message_group or uuid.uuid4()
+        self._stderr_capture = None
+        
+    @asynccontextmanager
+    async def client_streams(self):
+        """Create streams with stderr capture."""
+        server = StdioServerParameters(
+            command=self.command,
+            args=list(self.args),
+            env=self.env,
+            cwd=self.cwd
+        )
+        
+        # Create stderr capture
+        server_name = getattr(self, 'tool_prefix', self.command)
+        self._stderr_capture = StderrFileCapture(server_name, self.emit_stderr, self.message_group)
+        stderr_file = self._stderr_capture.start()
+        
+        try:
+            async with stdio_client(server=server, errlog=stderr_file) as (read_stream, write_stream):
+                yield read_stream, write_stream
+        finally:
+            self._stderr_capture.stop()
+            
+    def get_captured_stderr(self) -> List[str]:
+        """Get captured stderr lines."""
+        if self._stderr_capture:
+            return self._stderr_capture.get_captured_lines()
+        return []
+
+
+class BlockingMCPServerStdio(SimpleCapturedMCPServerStdio):
+    """
+    MCP Server that blocks until fully initialized.
+    
+    This server ensures that initialization is complete before
+    allowing any operations, preventing race conditions.
+    """
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._initialized = asyncio.Event()
+        self._init_error: Optional[Exception] = None
+        self._initialization_task = None
+        
+    async def __aenter__(self):
+        """Enter context and track initialization."""
+        try:
+            # Start initialization
+            result = await super().__aenter__()
+            
+            # Mark as initialized
+            self._initialized.set()
+            
+            # Emit success message
+            server_name = getattr(self, 'tool_prefix', self.command)
+            emit_info(
+                f"✅ MCP Server '{server_name}' initialized successfully", 
+                style="green",
+                message_group=self.message_group
+            )
+            
+            return result
+            
+        except Exception as e:
+            # Store error and mark as initialized (with error)
+            self._init_error = e
+            self._initialized.set()
+            
+            # Emit error message
+            server_name = getattr(self, 'tool_prefix', self.command)
+            emit_info(
+                f"❌ MCP Server '{server_name}' failed to initialize: {e}", 
+                style="red",
+                message_group=self.message_group
+            )
+            
+            raise
+            
+    async def wait_until_ready(self, timeout: float = 30.0) -> bool:
+        """
+        Wait until the server is ready.
+        
+        Args:
+            timeout: Maximum time to wait in seconds
+            
+        Returns:
+            True if server is ready, False if timeout or error
+            
+        Raises:
+            TimeoutError: If server doesn't initialize within timeout
+            Exception: If server initialization failed
+        """
+        try:
+            await asyncio.wait_for(self._initialized.wait(), timeout=timeout)
+            
+            # Check if there was an initialization error
+            if self._init_error:
+                raise self._init_error
+                
+            return True
+            
+        except asyncio.TimeoutError:
+            server_name = getattr(self, 'tool_prefix', self.command)
+            raise TimeoutError(f"Server '{server_name}' initialization timeout after {timeout}s")
+            
+    async def ensure_ready(self, timeout: float = 30.0):
+        """
+        Ensure server is ready before proceeding.
+        
+        This is a convenience method that raises if not ready.
+        
+        Args:
+            timeout: Maximum time to wait in seconds
+            
+        Raises:
+            TimeoutError: If server doesn't initialize within timeout
+            Exception: If server initialization failed
+        """
+        await self.wait_until_ready(timeout)
+        
+    def is_ready(self) -> bool:
+        """
+        Check if server is ready without blocking.
+        
+        Returns:
+            True if server is initialized and ready
+        """
+        return self._initialized.is_set() and self._init_error is None
+
+
+class StartupMonitor:
+    """
+    Monitor for tracking multiple server startups.
+    
+    This class helps coordinate startup of multiple MCP servers
+    and ensures all are ready before proceeding.
+    """
+    
+    def __init__(self, message_group: Optional[uuid.UUID] = None):
+        self.servers = {}
+        self.startup_times = {}
+        self.message_group = message_group or uuid.uuid4()
+        
+    def add_server(self, name: str, server: BlockingMCPServerStdio):
+        """Add a server to monitor."""
+        self.servers[name] = server
+        
+    async def wait_all_ready(self, timeout: float = 30.0) -> dict:
+        """
+        Wait for all servers to be ready.
+        
+        Args:
+            timeout: Maximum time to wait for all servers
+            
+        Returns:
+            Dictionary of server names to ready status
+        """
+        import time
+        results = {}
+        
+        # Create tasks for all servers
+        async def wait_server(name: str, server: BlockingMCPServerStdio):
+            start = time.time()
+            try:
+                await server.wait_until_ready(timeout)
+                self.startup_times[name] = time.time() - start
+                results[name] = True
+                emit_info(
+                    f"   {name}: Ready in {self.startup_times[name]:.2f}s", 
+                    style="dim green",
+                    message_group=self.message_group
+                )
+            except Exception as e:
+                self.startup_times[name] = time.time() - start
+                results[name] = False
+                emit_info(
+                    f"   {name}: Failed after {self.startup_times[name]:.2f}s - {e}", 
+                    style="dim red",
+                    message_group=self.message_group
+                )
+                
+        # Wait for all servers in parallel
+        emit_info(
+            f"⏳ Waiting for {len(self.servers)} MCP servers to initialize...", 
+            style="cyan",
+            message_group=self.message_group
+        )
+        
+        tasks = [
+            asyncio.create_task(wait_server(name, server))
+            for name, server in self.servers.items()
+        ]
+        
+        await asyncio.gather(*tasks, return_exceptions=True)
+        
+        # Report summary
+        ready_count = sum(1 for r in results.values() if r)
+        total_count = len(results)
+        
+        if ready_count == total_count:
+            emit_info(
+                f"✅ All {total_count} servers ready!", 
+                style="green bold",
+                message_group=self.message_group
+            )
+        else:
+            emit_info(
+                f"⚠️  {ready_count}/{total_count} servers ready", 
+                style="yellow",
+                message_group=self.message_group
+            )
+            
+        return results
+        
+    def get_startup_report(self) -> str:
+        """Get a report of startup times."""
+        lines = ["Server Startup Times:"]
+        for name, time_taken in self.startup_times.items():
+            status = "✅" if self.servers[name].is_ready() else "❌"
+            lines.append(f"  {status} {name}: {time_taken:.2f}s")
+        return "\n".join(lines)
+
+
+async def start_servers_with_blocking(*servers: BlockingMCPServerStdio, timeout: float = 30.0, message_group: Optional[uuid.UUID] = None):
+    """
+    Start multiple servers and wait for all to be ready.
+    
+    Args:
+        *servers: Variable number of BlockingMCPServerStdio instances
+        timeout: Maximum time to wait for all servers
+        message_group: Optional UUID for grouping log messages
+        
+    Returns:
+        List of ready servers
+        
+    Example:
+        server1 = BlockingMCPServerStdio(...)
+        server2 = BlockingMCPServerStdio(...)
+        ready = await start_servers_with_blocking(server1, server2)
+    """
+    monitor = StartupMonitor(message_group=message_group)
+    
+    for i, server in enumerate(servers):
+        name = getattr(server, 'tool_prefix', f"server-{i}")
+        monitor.add_server(name, server)
+        
+    # Start all servers
+    async def start_server(server):
+        async with server:
+            await asyncio.sleep(0.1)  # Keep context alive briefly
+            return server
+            
+    # Start servers in parallel
+    server_tasks = [
+        asyncio.create_task(start_server(server))
+        for server in servers
+    ]
+    
+    # Wait for all to be ready
+    results = await monitor.wait_all_ready(timeout)
+    
+    # Get the report
+    emit_info(monitor.get_startup_report(), message_group=monitor.message_group)
+    
+    # Return ready servers
+    ready_servers = [
+        server for name, server in monitor.servers.items()
+        if results.get(name, False)
+    ]
+    
+    return ready_servers

code_puppy/mcp/captured_stdio_server.py

@@ -0,0 +1,282 @@
+"""
+Custom MCPServerStdio that captures stderr output properly.
+
+This module provides a version of MCPServerStdio that captures subprocess
+stderr output and makes it available through proper logging channels.
+"""
+
+import asyncio
+import io
+import logging
+import os
+import sys
+import tempfile
+from contextlib import asynccontextmanager
+from typing import AsyncIterator, Sequence, Optional, Any
+from threading import Thread
+from queue import Queue, Empty
+
+from pydantic_ai.mcp import MCPServerStdio
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.shared.session import SessionMessage
+
+logger = logging.getLogger(__name__)
+
+
+class StderrCapture:
+    """
+    Captures stderr output using a pipe and background reader.
+    """
+    
+    def __init__(self, name: str, handler: Optional[callable] = None):
+        """
+        Initialize stderr capture.
+        
+        Args:
+            name: Name for this capture stream
+            handler: Optional function to call with captured lines
+        """
+        self.name = name
+        self.handler = handler or self._default_handler
+        self._captured_lines = []
+        self._reader_task = None
+        self._pipe_r = None
+        self._pipe_w = None
+        
+    def _default_handler(self, line: str):
+        """Default handler that logs to Python logging."""
+        if line.strip():
+            logger.debug(f"[MCP {self.name}] {line.rstrip()}")
+    
+    async def start_capture(self):
+        """Start capturing stderr by creating a pipe and reader task."""
+        # Create a pipe for capturing stderr
+        self._pipe_r, self._pipe_w = os.pipe()
+        
+        # Make the read end non-blocking
+        os.set_blocking(self._pipe_r, False)
+        
+        # Start background task to read from pipe
+        self._reader_task = asyncio.create_task(self._read_pipe())
+        
+        # Return the write end as the file descriptor for stderr
+        return self._pipe_w
+    
+    async def _read_pipe(self):
+        """Background task to read from the pipe."""
+        loop = asyncio.get_event_loop()
+        buffer = b''
+        
+        try:
+            while True:
+                # Use asyncio's add_reader for efficient async reading
+                future = asyncio.Future()
+                
+                def read_callback():
+                    try:
+                        data = os.read(self._pipe_r, 4096)
+                        future.set_result(data)
+                    except BlockingIOError:
+                        future.set_result(b'')
+                    except Exception as e:
+                        future.set_exception(e)
+                
+                loop.add_reader(self._pipe_r, read_callback)
+                try:
+                    data = await future
+                finally:
+                    loop.remove_reader(self._pipe_r)
+                
+                if not data:
+                    await asyncio.sleep(0.1)
+                    continue
+                
+                # Process the data
+                buffer += data
+                
+                # Look for complete lines
+                while b'\n' in buffer:
+                    line, buffer = buffer.split(b'\n', 1)
+                    line_str = line.decode('utf-8', errors='replace')
+                    if line_str:
+                        self._captured_lines.append(line_str)
+                        self.handler(line_str)
+                        
+        except asyncio.CancelledError:
+            # Process any remaining buffer
+            if buffer:
+                line_str = buffer.decode('utf-8', errors='replace')
+                if line_str:
+                    self._captured_lines.append(line_str)
+                    self.handler(line_str)
+            raise
+            
+    async def stop_capture(self):
+        """Stop capturing and clean up."""
+        if self._reader_task:
+            self._reader_task.cancel()
+            try:
+                await self._reader_task
+            except asyncio.CancelledError:
+                pass
+                
+        if self._pipe_r is not None:
+            os.close(self._pipe_r)
+        if self._pipe_w is not None:
+            os.close(self._pipe_w)
+            
+    def get_captured_lines(self) -> list[str]:
+        """Get all captured lines."""
+        return self._captured_lines.copy()
+
+
+class CapturedMCPServerStdio(MCPServerStdio):
+    """
+    Extended MCPServerStdio that captures and handles stderr output.
+    
+    This class captures stderr from the subprocess and makes it available
+    through proper logging channels instead of letting it pollute the console.
+    """
+    
+    def __init__(
+        self,
+        command: str,
+        args: Sequence[str] = (),
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+        stderr_handler: Optional[callable] = None,
+        **kwargs
+    ):
+        """
+        Initialize captured stdio server.
+        
+        Args:
+            command: The command to run
+            args: Arguments for the command
+            env: Environment variables
+            cwd: Working directory
+            stderr_handler: Optional function to handle stderr lines
+            **kwargs: Additional arguments for MCPServerStdio
+        """
+        super().__init__(command=command, args=args, env=env, cwd=cwd, **kwargs)
+        self.stderr_handler = stderr_handler
+        self._stderr_capture = None
+        self._captured_lines = []
+        
+    @asynccontextmanager
+    async def client_streams(
+        self,
+    ) -> AsyncIterator[
+        tuple[
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
+        ]
+    ]:
+        """Create the streams for the MCP server with stderr capture."""
+        server = StdioServerParameters(
+            command=self.command, 
+            args=list(self.args), 
+            env=self.env, 
+            cwd=self.cwd
+        )
+        
+        # Create stderr capture
+        def stderr_line_handler(line: str):
+            """Handle captured stderr lines."""
+            self._captured_lines.append(line)
+            
+            if self.stderr_handler:
+                self.stderr_handler(line)
+            else:
+                # Default: log at DEBUG level to avoid console spam
+                logger.debug(f"[MCP Server {self.command}] {line}")
+        
+        self._stderr_capture = StderrCapture(self.command, stderr_line_handler)
+        
+        # For now, use devnull for stderr to suppress output
+        # We'll capture it through other means if needed
+        with open(os.devnull, 'w') as devnull:
+            async with stdio_client(server=server, errlog=devnull) as (read_stream, write_stream):
+                yield read_stream, write_stream
+            
+    def get_captured_stderr(self) -> list[str]:
+        """
+        Get all captured stderr lines.
+        
+        Returns:
+            List of captured stderr lines
+        """
+        return self._captured_lines.copy()
+    
+    def clear_captured_stderr(self):
+        """Clear the captured stderr buffer."""
+        self._captured_lines.clear()
+
+
+class StderrCollector:
+    """
+    A centralized collector for stderr from multiple MCP servers.
+    
+    This can be used to aggregate stderr from all MCP servers in one place.
+    """
+    
+    def __init__(self):
+        """Initialize the collector."""
+        self.servers = {}
+        self.all_lines = []
+        
+    def create_handler(self, server_name: str, emit_to_user: bool = False):
+        """
+        Create a handler function for a specific server.
+        
+        Args:
+            server_name: Name to identify this server
+            emit_to_user: If True, emit stderr lines to user via emit_info
+            
+        Returns:
+            Handler function that can be passed to CapturedMCPServerStdio
+        """
+        def handler(line: str):
+            # Store with server identification
+            import time
+            entry = {
+                'server': server_name,
+                'line': line,
+                'timestamp': time.time()
+            }
+            
+            if server_name not in self.servers:
+                self.servers[server_name] = []
+            
+            self.servers[server_name].append(line)
+            self.all_lines.append(entry)
+            
+            # Emit to user if requested
+            if emit_to_user:
+                from code_puppy.messaging import emit_info
+                emit_info(f"[MCP {server_name}] {line}", style="dim cyan")
+            
+        return handler
+    
+    def get_server_output(self, server_name: str) -> list[str]:
+        """Get all output from a specific server."""
+        return self.servers.get(server_name, []).copy()
+    
+    def get_all_output(self) -> list[dict]:
+        """Get all output from all servers with metadata."""
+        return self.all_lines.copy()
+    
+    def clear(self, server_name: Optional[str] = None):
+        """Clear captured output."""
+        if server_name:
+            if server_name in self.servers:
+                self.servers[server_name].clear()
+                # Also clear from all_lines
+                self.all_lines = [
+                    entry for entry in self.all_lines 
+                    if entry['server'] != server_name
+                ]
+        else:
+            self.servers.clear()
+            self.all_lines.clear()