code-puppy 0.0.135__py3-none-any.whl → 0.0.136__py3-none-any.whl

code_puppy/mcp/captured_stdio_server.py

@@ -6,20 +6,15 @@ 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 typing import AsyncIterator, Optional, Sequence
 
-from pydantic_ai.mcp import MCPServerStdio
-from mcp.client.stdio import StdioServerParameters, stdio_client
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.client.stdio import StdioServerParameters, stdio_client
 from mcp.shared.session import SessionMessage
+from pydantic_ai.mcp import MCPServerStdio
 
 logger = logging.getLogger(__name__)
 
@@ -28,11 +23,11 @@ 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
@@ -43,75 +38,75 @@ class StderrCapture:
         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''
-        
+        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'')
+                        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')
+                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')
+                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:
@@ -120,12 +115,12 @@ class StderrCapture:
                 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()
@@ -134,11 +129,11 @@ class StderrCapture:
 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,
@@ -146,11 +141,11 @@ class CapturedMCPServerStdio(MCPServerStdio):
         env: dict[str, str] | None = None,
         cwd: str | None = None,
         stderr_handler: Optional[callable] = None,
-        **kwargs
+        **kwargs,
     ):
         """
         Initialize captured stdio server.
-        
+
         Args:
             command: The command to run
             args: Arguments for the command
@@ -163,7 +158,7 @@ class CapturedMCPServerStdio(MCPServerStdio):
         self.stderr_handler = stderr_handler
         self._stderr_capture = None
         self._captured_lines = []
-        
+
     @asynccontextmanager
     async def client_streams(
         self,
@@ -175,40 +170,40 @@ class CapturedMCPServerStdio(MCPServerStdio):
     ]:
         """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
+            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):
+        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()
@@ -217,56 +212,55 @@ class CapturedMCPServerStdio(MCPServerStdio):
 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()
-            }
-            
+
+            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:
@@ -274,9 +268,8 @@ class StderrCollector:
                 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
+                    entry for entry in self.all_lines if entry["server"] != server_name
                 ]
         else:
             self.servers.clear()
-            self.all_lines.clear()
+            self.all_lines.clear()

code_puppy/mcp/circuit_breaker.py

@@ -9,129 +9,137 @@ failures when MCP servers become unhealthy. The circuit breaker has three states
 """
 
 import asyncio
+import logging
 import time
 from enum import Enum
 from typing import Any, Callable
-import logging
 
 logger = logging.getLogger(__name__)
 
 
 class CircuitState(Enum):
     """Circuit breaker states."""
-    CLOSED = "closed"      # Normal operation
-    OPEN = "open"          # Blocking calls
+
+    CLOSED = "closed"  # Normal operation
+    OPEN = "open"  # Blocking calls
     HALF_OPEN = "half_open"  # Testing recovery
 
 
 class CircuitOpenError(Exception):
     """Raised when circuit breaker is in OPEN state."""
+
     pass
 
 
 class CircuitBreaker:
     """
     Circuit breaker to prevent cascading failures in MCP servers.
-    
+
     The circuit breaker monitors the success/failure rate of operations and
     transitions between states to protect the system from unhealthy dependencies.
-    
+
     States:
     - CLOSED: Normal operation, all calls allowed
     - OPEN: Circuit is open, all calls fail fast with CircuitOpenError
     - HALF_OPEN: Testing recovery, limited calls allowed
-    
+
     State Transitions:
     - CLOSED → OPEN: After failure_threshold consecutive failures
     - OPEN → HALF_OPEN: After timeout seconds
     - HALF_OPEN → CLOSED: After success_threshold consecutive successes
     - HALF_OPEN → OPEN: After any failure
     """
-    
-    def __init__(self, failure_threshold: int = 5, success_threshold: int = 2, timeout: int = 60):
+
+    def __init__(
+        self, failure_threshold: int = 5, success_threshold: int = 2, timeout: int = 60
+    ):
         """
         Initialize circuit breaker.
-        
+
         Args:
             failure_threshold: Number of consecutive failures before opening circuit
             success_threshold: Number of consecutive successes needed to close circuit from half-open
             timeout: Seconds to wait before transitioning from OPEN to HALF_OPEN
         """
         self.failure_threshold = failure_threshold
-        self.success_threshold = success_threshold  
+        self.success_threshold = success_threshold
         self.timeout = timeout
-        
+
         self._state = CircuitState.CLOSED
         self._failure_count = 0
         self._success_count = 0
         self._last_failure_time = None
         self._lock = asyncio.Lock()
-        
+
         logger.info(
             f"Circuit breaker initialized: failure_threshold={failure_threshold}, "
             f"success_threshold={success_threshold}, timeout={timeout}s"
         )
-    
+
     async def call(self, func: Callable, *args, **kwargs) -> Any:
         """
         Execute a function through the circuit breaker.
-        
+
         Args:
             func: Function to execute
             *args: Positional arguments for the function
             **kwargs: Keyword arguments for the function
-            
+
         Returns:
             Result of the function call
-            
+
         Raises:
             CircuitOpenError: If circuit is in OPEN state
             Exception: Any exception raised by the wrapped function
         """
         async with self._lock:
             current_state = self._get_current_state()
-            
+
             if current_state == CircuitState.OPEN:
                 logger.warning("Circuit breaker is OPEN, failing fast")
                 raise CircuitOpenError("Circuit breaker is open")
-            
+
             if current_state == CircuitState.HALF_OPEN:
                 # In half-open state, we're testing recovery
                 logger.info("Circuit breaker is HALF_OPEN, allowing test call")
-        
+
         # Execute the function outside the lock to avoid blocking other calls
         try:
-            result = await func(*args, **kwargs) if asyncio.iscoroutinefunction(func) else func(*args, **kwargs)
+            result = (
+                await func(*args, **kwargs)
+                if asyncio.iscoroutinefunction(func)
+                else func(*args, **kwargs)
+            )
             await self._on_success()
             return result
         except Exception as e:
             await self._on_failure()
             raise e
-    
+
     def record_success(self) -> None:
         """Record a successful operation."""
         asyncio.create_task(self._on_success())
-    
+
     def record_failure(self) -> None:
         """Record a failed operation."""
         asyncio.create_task(self._on_failure())
-    
+
     def get_state(self) -> CircuitState:
         """Get current circuit breaker state."""
         return self._get_current_state()
-    
+
     def is_open(self) -> bool:
         """Check if circuit breaker is in OPEN state."""
         return self._get_current_state() == CircuitState.OPEN
-    
+
     def is_half_open(self) -> bool:
         """Check if circuit breaker is in HALF_OPEN state."""
         return self._get_current_state() == CircuitState.HALF_OPEN
-    
+
     def is_closed(self) -> bool:
         """Check if circuit breaker is in CLOSED state."""
         return self._get_current_state() == CircuitState.CLOSED
-    
+
     def reset(self) -> None:
         """Reset circuit breaker to CLOSED state and clear counters."""
         logger.info("Resetting circuit breaker to CLOSED state")
@@ -139,13 +147,13 @@ class CircuitBreaker:
         self._failure_count = 0
         self._success_count = 0
         self._last_failure_time = None
-    
+
     def force_open(self) -> None:
         """Force circuit breaker to OPEN state."""
         logger.warning("Forcing circuit breaker to OPEN state")
         self._state = CircuitState.OPEN
         self._last_failure_time = time.time()
-    
+
     def force_close(self) -> None:
         """Force circuit breaker to CLOSED state and reset counters."""
         logger.info("Forcing circuit breaker to CLOSED state")
@@ -153,11 +161,11 @@ class CircuitBreaker:
         self._failure_count = 0
         self._success_count = 0
         self._last_failure_time = None
-    
+
     def _get_current_state(self) -> CircuitState:
         """
         Get the current state, handling automatic transitions.
-        
+
         This method handles the automatic transition from OPEN to HALF_OPEN
         after the timeout period has elapsed.
         """
@@ -165,54 +173,62 @@ class CircuitBreaker:
             logger.info("Timeout reached, transitioning from OPEN to HALF_OPEN")
             self._state = CircuitState.HALF_OPEN
             self._success_count = 0  # Reset success counter for half-open testing
-        
+
         return self._state
-    
+
     def _should_attempt_reset(self) -> bool:
         """Check if enough time has passed to attempt reset from OPEN to HALF_OPEN."""
         if self._last_failure_time is None:
             return False
-        
+
         return time.time() - self._last_failure_time >= self.timeout
-    
+
     async def _on_success(self) -> None:
         """Handle successful operation."""
         async with self._lock:
             current_state = self._get_current_state()
-            
+
             if current_state == CircuitState.CLOSED:
                 # Reset failure count on success in closed state
                 if self._failure_count > 0:
                     logger.debug("Resetting failure count after success")
                     self._failure_count = 0
-            
+
             elif current_state == CircuitState.HALF_OPEN:
                 self._success_count += 1
-                logger.debug(f"Success in HALF_OPEN state: {self._success_count}/{self.success_threshold}")
-                
+                logger.debug(
+                    f"Success in HALF_OPEN state: {self._success_count}/{self.success_threshold}"
+                )
+
                 if self._success_count >= self.success_threshold:
-                    logger.info("Success threshold reached, transitioning from HALF_OPEN to CLOSED")
+                    logger.info(
+                        "Success threshold reached, transitioning from HALF_OPEN to CLOSED"
+                    )
                     self._state = CircuitState.CLOSED
                     self._failure_count = 0
                     self._success_count = 0
                     self._last_failure_time = None
-    
+
     async def _on_failure(self) -> None:
         """Handle failed operation."""
         async with self._lock:
             current_state = self._get_current_state()
-            
+
             if current_state == CircuitState.CLOSED:
                 self._failure_count += 1
-                logger.debug(f"Failure in CLOSED state: {self._failure_count}/{self.failure_threshold}")
-                
+                logger.debug(
+                    f"Failure in CLOSED state: {self._failure_count}/{self.failure_threshold}"
+                )
+
                 if self._failure_count >= self.failure_threshold:
-                    logger.warning("Failure threshold reached, transitioning from CLOSED to OPEN")
+                    logger.warning(
+                        "Failure threshold reached, transitioning from CLOSED to OPEN"
+                    )
                     self._state = CircuitState.OPEN
                     self._last_failure_time = time.time()
-            
+
             elif current_state == CircuitState.HALF_OPEN:
                 logger.warning("Failure in HALF_OPEN state, transitioning back to OPEN")
                 self._state = CircuitState.OPEN
                 self._success_count = 0
-                self._last_failure_time = time.time()
+                self._last_failure_time = time.time()