codepp 0.0.437__py3-none-any.whl

code_puppy/keymap.py

@@ -0,0 +1,128 @@
+"""Keymap configuration for code-puppy.
+
+This module handles configurable keyboard shortcuts, starting with the
+cancel_agent_key feature that allows users to override Ctrl+C with a
+different key for cancelling agent tasks.
+"""
+
+# Character codes for Ctrl+letter combinations (Ctrl+A = 0x01, Ctrl+Z = 0x1A)
+KEY_CODES: dict[str, str] = {
+    "ctrl+a": "\x01",
+    "ctrl+b": "\x02",
+    "ctrl+c": "\x03",
+    "ctrl+d": "\x04",
+    "ctrl+e": "\x05",
+    "ctrl+f": "\x06",
+    "ctrl+g": "\x07",
+    "ctrl+h": "\x08",
+    "ctrl+i": "\x09",
+    "ctrl+j": "\x0a",
+    "ctrl+k": "\x0b",
+    "ctrl+l": "\x0c",
+    "ctrl+m": "\x0d",
+    "ctrl+n": "\x0e",
+    "ctrl+o": "\x0f",
+    "ctrl+p": "\x10",
+    "ctrl+q": "\x11",
+    "ctrl+r": "\x12",
+    "ctrl+s": "\x13",
+    "ctrl+t": "\x14",
+    "ctrl+u": "\x15",
+    "ctrl+v": "\x16",
+    "ctrl+w": "\x17",
+    "ctrl+x": "\x18",
+    "ctrl+y": "\x19",
+    "ctrl+z": "\x1a",
+    "escape": "\x1b",
+}
+
+# Valid keys for cancel_agent_key configuration
+# NOTE: "escape" is excluded because it conflicts with ANSI escape sequences
+# (arrow keys, F-keys, etc. all start with \x1b)
+VALID_CANCEL_KEYS: set[str] = {
+    "ctrl+c",
+    "ctrl+k",
+    "ctrl+q",
+}
+
+DEFAULT_CANCEL_AGENT_KEY: str = "ctrl+c"
+
+
+class KeymapError(Exception):
+    """Exception raised for keymap configuration errors."""
+
+
+def get_cancel_agent_key() -> str:
+    """Get the configured cancel agent key from config.
+
+    On Windows when launched via uvx, this automatically returns "ctrl+k"
+    to work around uvx capturing Ctrl+C before it reaches Python.
+
+    Returns:
+        The key name (e.g., "ctrl+c", "ctrl+k") from config,
+        or the default if not configured.
+    """
+    from code_puppy.config import get_value
+    from code_puppy.uvx_detection import should_use_alternate_cancel_key
+
+    # On Windows + uvx, force ctrl+k to bypass uvx's SIGINT capture
+    if should_use_alternate_cancel_key():
+        return "ctrl+k"
+
+    key = get_value("cancel_agent_key")
+    if key is None or key.strip() == "":
+        return DEFAULT_CANCEL_AGENT_KEY
+    return key.strip().lower()
+
+
+def validate_cancel_agent_key() -> None:
+    """Validate the configured cancel agent key.
+
+    Raises:
+        KeymapError: If the configured key is invalid.
+    """
+    key = get_cancel_agent_key()
+    if key not in VALID_CANCEL_KEYS:
+        valid_keys_str = ", ".join(sorted(VALID_CANCEL_KEYS))
+        raise KeymapError(
+            f"Invalid cancel_agent_key '{key}' in puppy.cfg. "
+            f"Valid options are: {valid_keys_str}"
+        )
+
+
+def cancel_agent_uses_signal() -> bool:
+    """Check if the cancel agent key uses SIGINT (Ctrl+C).
+
+    Returns:
+        True if the cancel key is ctrl+c (uses SIGINT handler),
+        False if it uses keyboard listener approach.
+    """
+    return get_cancel_agent_key() == "ctrl+c"
+
+
+def get_cancel_agent_char_code() -> str:
+    """Get the character code for the cancel agent key.
+
+    Returns:
+        The character code (e.g., "\x0b" for ctrl+k).
+
+    Raises:
+        KeymapError: If the key is not found in KEY_CODES.
+    """
+    key = get_cancel_agent_key()
+    if key not in KEY_CODES:
+        raise KeymapError(f"Unknown key '{key}' - no character code mapping found.")
+    return KEY_CODES[key]
+
+
+def get_cancel_agent_display_name() -> str:
+    """Get a human-readable display name for the cancel agent key.
+
+    Returns:
+        A formatted display name like "Ctrl+K".
+    """
+    key = get_cancel_agent_key()
+    if key.startswith("ctrl+"):
+        letter = key.split("+")[1].upper()
+        return f"Ctrl+{letter}"
+    return key.upper()

code_puppy/main.py

@@ -0,0 +1,10 @@
+"""Main entry point for Code Puppy CLI.
+
+This module re-exports the main_entry function from cli_runner for backwards compatibility.
+All the actual logic lives in cli_runner.py.
+"""
+
+from code_puppy.cli_runner import main_entry
+
+if __name__ == "__main__":
+    main_entry()

code_puppy/mcp_/__init__.py

@@ -0,0 +1,66 @@
+"""MCP (Model Context Protocol) management system for Code Puppy.
+
+Note: Be careful not to create circular imports with config_wizard.py.
+config_wizard.py imports ServerConfig and get_mcp_manager directly from
+.manager to avoid circular dependencies with this package __init__.py
+"""
+
+from .circuit_breaker import CircuitBreaker, CircuitOpenError, CircuitState
+from .config_wizard import MCPConfigWizard, run_add_wizard
+from .dashboard import MCPDashboard
+from .error_isolation import (
+    ErrorCategory,
+    ErrorStats,
+    MCPErrorIsolator,
+    QuarantinedServerError,
+    get_error_isolator,
+)
+from .managed_server import ManagedMCPServer, ServerConfig, ServerState
+from .manager import MCPManager, ServerInfo, get_mcp_manager
+from .mcp_logs import (
+    clear_logs,
+    get_log_file_path,
+    get_log_stats,
+    get_mcp_logs_dir,
+    list_servers_with_logs,
+    read_logs,
+    write_log,
+)
+from .registry import ServerRegistry
+from .retry_manager import RetryManager, RetryStats, get_retry_manager, retry_mcp_call
+from .status_tracker import Event, ServerStatusTracker
+
+__all__ = [
+    "ManagedMCPServer",
+    "ServerConfig",
+    "ServerState",
+    "ServerStatusTracker",
+    "Event",
+    "MCPManager",
+    "ServerInfo",
+    "get_mcp_manager",
+    "ServerRegistry",
+    "MCPErrorIsolator",
+    "ErrorStats",
+    "ErrorCategory",
+    "QuarantinedServerError",
+    "get_error_isolator",
+    "CircuitBreaker",
+    "CircuitState",
+    "CircuitOpenError",
+    "RetryManager",
+    "RetryStats",
+    "get_retry_manager",
+    "retry_mcp_call",
+    "MCPDashboard",
+    "MCPConfigWizard",
+    "run_add_wizard",
+    # Log management
+    "get_mcp_logs_dir",
+    "get_log_file_path",
+    "read_logs",
+    "write_log",
+    "clear_logs",
+    "list_servers_with_logs",
+    "get_log_stats",
+]

code_puppy/mcp_/async_lifecycle.py

@@ -0,0 +1,286 @@
+"""
+Async server lifecycle management using pydantic-ai's context managers.
+
+This module properly manages MCP server lifecycles by maintaining async contexts
+within the same task, allowing servers to start and stay running.
+"""
+
+import asyncio
+import logging
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, Optional, Union
+
+from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ManagedServerContext:
+    """Represents a managed MCP server with its async context."""
+
+    server_id: str
+    server: Union[MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP]
+    exit_stack: AsyncExitStack
+    start_time: datetime
+    task: asyncio.Task  # The task that manages this server's lifecycle
+
+
+class AsyncServerLifecycleManager:
+    """
+    Manages MCP server lifecycles asynchronously.
+
+    This properly maintains async contexts within the same task,
+    allowing servers to start and stay running independently of agents.
+    """
+
+    def __init__(self):
+        """Initialize the async lifecycle manager."""
+        self._servers: Dict[str, ManagedServerContext] = {}
+        self._lock = asyncio.Lock()
+        logger.info("AsyncServerLifecycleManager initialized")
+
+    async def start_server(
+        self,
+        server_id: str,
+        server: Union[MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP],
+    ) -> bool:
+        """
+        Start an MCP server and maintain its context.
+
+        This creates a dedicated task that enters the server's context
+        and keeps it alive until explicitly stopped.
+
+        Args:
+            server_id: Unique identifier for the server
+            server: The pydantic-ai MCP server instance
+
+        Returns:
+            True if server started successfully, False otherwise
+        """
+        async with self._lock:
+            # Check if already running
+            if server_id in self._servers:
+                if self._servers[server_id].server.is_running:
+                    logger.info(f"Server {server_id} is already running")
+                    return True
+                else:
+                    # Server exists but not running, clean it up
+                    logger.warning(
+                        f"Server {server_id} exists but not running, cleaning up"
+                    )
+                    await self._stop_server_internal(server_id)
+
+            # Create an event so we know when the server is actually registered
+            ready_event = asyncio.Event()
+
+            # Create a task that will manage this server's lifecycle
+            task = asyncio.create_task(
+                self._server_lifecycle_task(server_id, server, ready_event),
+                name=f"mcp_server_{server_id}",
+            )
+
+        # Release the lock while waiting for the server to become ready
+        try:
+            await asyncio.wait_for(ready_event.wait(), timeout=10.0)
+        except asyncio.TimeoutError:
+            logger.error(f"Timed out waiting for server {server_id} to start")
+            if task.done():
+                try:
+                    await task
+                except Exception as e:
+                    logger.error(f"Server {server_id} task failed: {e}")
+            return False
+
+        # Check if task failed during startup
+        if task.done():
+            try:
+                await task
+            except Exception as e:
+                logger.error(f"Failed to start server {server_id}: {e}")
+                return False
+
+        logger.info(f"Server {server_id} started successfully")
+        return True
+
+    async def _server_lifecycle_task(
+        self,
+        server_id: str,
+        server: Union[MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP],
+        ready_event: asyncio.Event,
+    ) -> None:
+        """
+        Task that manages a server's lifecycle.
+
+        This task enters the server's context and keeps it alive
+        until the server is stopped or an error occurs.
+        """
+        exit_stack = AsyncExitStack()
+
+        try:
+            logger.info(f"Starting server lifecycle for {server_id}")
+            logger.info(
+                f"Server {server_id} _running_count before enter: {getattr(server, '_running_count', 'N/A')}"
+            )
+
+            # Enter the server's context
+            await exit_stack.enter_async_context(server)
+
+            logger.info(
+                f"Server {server_id} _running_count after enter: {getattr(server, '_running_count', 'N/A')}"
+            )
+
+            # Store the managed context
+            async with self._lock:
+                self._servers[server_id] = ManagedServerContext(
+                    server_id=server_id,
+                    server=server,
+                    exit_stack=exit_stack,
+                    start_time=datetime.now(),
+                    task=asyncio.current_task(),
+                )
+
+            # Signal that the server is registered and ready
+            ready_event.set()
+
+            logger.info(
+                f"Server {server_id} started successfully and stored in _servers"
+            )
+
+            # Keep the task alive until cancelled
+            loop_count = 0
+            while True:
+                await asyncio.sleep(1)
+                loop_count += 1
+
+                # Check if server is still running
+                running_count = getattr(server, "_running_count", "N/A")
+                is_running = server.is_running
+                logger.debug(
+                    f"Server {server_id} heartbeat #{loop_count}: "
+                    f"is_running={is_running}, _running_count={running_count}"
+                )
+
+                if not is_running:
+                    logger.warning(
+                        f"Server {server_id} stopped unexpectedly! "
+                        f"_running_count={running_count}"
+                    )
+                    break
+
+        except asyncio.CancelledError:
+            logger.info(f"Server {server_id} lifecycle task cancelled")
+            raise
+        except Exception as e:
+            logger.error(f"Error in server {server_id} lifecycle: {e}", exc_info=True)
+        finally:
+            running_count = getattr(server, "_running_count", "N/A")
+            logger.info(
+                f"Server {server_id} lifecycle ending, _running_count={running_count}"
+            )
+
+            # Clean up the context
+            await exit_stack.aclose()
+
+            running_count_after = getattr(server, "_running_count", "N/A")
+            logger.info(
+                f"Server {server_id} context closed, _running_count={running_count_after}"
+            )
+
+            # Remove from managed servers
+            async with self._lock:
+                if server_id in self._servers:
+                    del self._servers[server_id]
+
+            logger.info(f"Server {server_id} lifecycle ended")
+
+    async def stop_server(self, server_id: str) -> bool:
+        """
+        Stop a running MCP server.
+
+        This cancels the lifecycle task, which properly exits the context.
+
+        Args:
+            server_id: ID of the server to stop
+
+        Returns:
+            True if server was stopped, False if not found
+        """
+        async with self._lock:
+            return await self._stop_server_internal(server_id)
+
+    async def _stop_server_internal(self, server_id: str) -> bool:
+        """
+        Internal method to stop a server (must be called with lock held).
+        """
+        if server_id not in self._servers:
+            logger.warning(f"Server {server_id} not found")
+            return False
+
+        context = self._servers[server_id]
+
+        # Cancel the lifecycle task
+        # This will cause the task to exit and clean up properly
+        context.task.cancel()
+
+        try:
+            await context.task
+        except asyncio.CancelledError:
+            pass  # Expected
+
+        logger.info(f"Stopped server {server_id}")
+        return True
+
+    def is_running(self, server_id: str) -> bool:
+        """
+        Check if a server is running.
+
+        Args:
+            server_id: ID of the server
+
+        Returns:
+            True if server is running, False otherwise
+        """
+        context = self._servers.get(server_id)
+        return context.server.is_running if context else False
+
+    def list_servers(self) -> Dict[str, Dict[str, Any]]:
+        """
+        List all running servers.
+
+        Returns:
+            Dictionary of server IDs to server info
+        """
+        servers = {}
+        for server_id, context in self._servers.items():
+            uptime = (datetime.now() - context.start_time).total_seconds()
+            servers[server_id] = {
+                "type": context.server.__class__.__name__,
+                "is_running": context.server.is_running,
+                "uptime_seconds": uptime,
+                "start_time": context.start_time.isoformat(),
+            }
+        return servers
+
+    async def stop_all(self) -> None:
+        """Stop all running servers."""
+        server_ids = list(self._servers.keys())
+
+        for server_id in server_ids:
+            await self.stop_server(server_id)
+
+        logger.info("All MCP servers stopped")
+
+
+# Global singleton instance
+_lifecycle_manager: Optional[AsyncServerLifecycleManager] = None
+
+
+def get_lifecycle_manager() -> AsyncServerLifecycleManager:
+    """Get the global lifecycle manager instance."""
+    global _lifecycle_manager
+    if _lifecycle_manager is None:
+        _lifecycle_manager = AsyncServerLifecycleManager()
+    return _lifecycle_manager