code-puppy 0.0.126__py3-none-any.whl → 0.0.128__py3-none-any.whl

code_puppy/command_line/meta_command_handler.py

@@ -117,11 +117,12 @@ def handle_meta_command(command: str, console: Console) -> bool:
         new_input = update_model_in_input(command)
         if new_input is not None:
             from code_puppy.command_line.model_picker_completion import get_active_model
-            from code_puppy.agent import get_code_generation_agent
+            from code_puppy.agents.runtime_manager import get_runtime_agent_manager
 
             model = get_active_model()
             # Make sure this is called for the test
-            get_code_generation_agent(force_reload=True)
+            manager = get_runtime_agent_manager()
+            manager.reload_agent()
             console.print(
                 f"[bold green]Active model set and loaded:[/bold green] [cyan]{model}[/cyan]"
             )

code_puppy/command_line/model_picker_completion.py

@@ -40,11 +40,11 @@ def set_active_model(model_name: str):
 
 class ModelNameCompleter(Completer):
     """
-    A completer that triggers on '/m' to show available models from models.json.
-    Only '/m' (not just '/') will trigger the dropdown.
+    A completer that triggers on '/model' to show available models from models.json.
+    Only '/model' (not just '/') will trigger the dropdown.
     """
 
-    def __init__(self, trigger: str = "/m"):
+    def __init__(self, trigger: str = "/model"):
         self.trigger = trigger
         self.model_names = load_model_names()
 
@@ -70,14 +70,27 @@ class ModelNameCompleter(Completer):
 
 
 def update_model_in_input(text: str) -> Optional[str]:
-    # If input starts with /m and a model name, set model and strip it out
+    # If input starts with /model or /m and a model name, set model and strip it out
     content = text.strip()
-    if content.startswith("/m"):
-        rest = content[2:].strip()
+    
+    # Check for /model command
+    if content.startswith("/model"):
+        rest = content[6:].strip()  # Remove '/model'
         for model in load_model_names():
             if rest == model:
                 set_active_model(model)
-                # Remove /mmodel from the input
+                # Remove /model from the input
+                idx = text.find("/model" + model)
+                if idx != -1:
+                    new_text = (text[:idx] + text[idx + len("/model" + model) :]).strip()
+                    return new_text
+    # Also check for legacy /m command for backward compatibility
+    elif content.startswith("/m"):
+        rest = content[2:].strip()  # Remove '/m'
+        for model in load_model_names():
+            if rest == model:
+                set_active_model(model)
+                # Remove /m from the input
                 idx = text.find("/m" + model)
                 if idx != -1:
                     new_text = (text[:idx] + text[idx + len("/m" + model) :]).strip()
@@ -86,7 +99,7 @@ def update_model_in_input(text: str) -> Optional[str]:
 
 
 async def get_input_with_model_completion(
-    prompt_str: str = ">>> ", trigger: str = "/m", history_file: Optional[str] = None
+    prompt_str: str = ">>> ", trigger: str = "/model", history_file: Optional[str] = None
 ) -> str:
     history = FileHistory(os.path.expanduser(history_file)) if history_file else None
     session = PromptSession(

code_puppy/main.py

@@ -12,7 +12,8 @@ from rich.syntax import Syntax
 from rich.text import Text
 
 from code_puppy import __version__, callbacks, plugins, state_management
-from code_puppy.agent import get_code_generation_agent, get_custom_usage_limits
+from code_puppy.agent import get_custom_usage_limits
+from code_puppy.agents.runtime_manager import get_runtime_agent_manager
 from code_puppy.command_line.prompt_toolkit_completion import (
     get_input_with_combined_completion,
     get_prompt_with_active_model,
@@ -250,9 +251,10 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
     emit_system_message(
         "Press [bold red]Ctrl+C[/bold red] during processing to cancel the current task or inference."
     )
-    from code_puppy.command_line.command_handler import COMMANDS_HELP
+    from code_puppy.command_line.command_handler import get_commands_help
 
-    emit_system_message(COMMANDS_HELP)
+    help_text = get_commands_help()
+    emit_system_message(help_text)
     try:
         from code_puppy.command_line.motd import print_motd
 
@@ -262,9 +264,12 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
 
         emit_warning(f"MOTD error: {e}")
     from code_puppy.messaging import emit_info
+    from code_puppy.agents.runtime_manager import get_runtime_agent_manager
 
     emit_info("[bold cyan]Initializing agent...[/bold cyan]")
-    get_code_generation_agent()
+    # Initialize the runtime agent manager
+    agent_manager = get_runtime_agent_manager()
+    agent_manager.get_agent()
     if initial_command:
         from code_puppy.messaging import emit_info, emit_system_message
 
@@ -273,9 +278,6 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
         )
 
         try:
-            # Get the agent (already loaded above)
-            agent = get_code_generation_agent()
-
             # Check if any tool is waiting for user input before showing spinner
             try:
                 from code_puppy.tools.command_runner import is_awaiting_user_input
@@ -286,44 +288,22 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
 
             # Run with or without spinner based on whether we're awaiting input
             if awaiting_input:
-                # No spinner - just run the agent
-                try:
-                    async with agent.run_mcp_servers():
-                        response = await agent.run(
-                            initial_command, usage_limits=get_custom_usage_limits()
-                        )
-                except Exception as mcp_error:
-                    from code_puppy.messaging import emit_warning
-
-                    emit_warning(f"MCP server error: {str(mcp_error)}")
-                    emit_warning("Running without MCP servers...")
-                    # Run without MCP servers as fallback
-                    response = await agent.run(
-                        initial_command, usage_limits=get_custom_usage_limits()
-                    )
+                # No spinner - use agent_manager's run_with_mcp method
+                response = await agent_manager.run_with_mcp(
+                    initial_command, usage_limits=get_custom_usage_limits()
+                )
             else:
                 # Use our custom spinner for better compatibility with user input
                 from code_puppy.messaging.spinner import ConsoleSpinner
 
                 with ConsoleSpinner(console=display_console):
-                    try:
-                        async with agent.run_mcp_servers():
-                            response = await agent.run(
-                                initial_command, usage_limits=get_custom_usage_limits()
-                            )
-                    except Exception as mcp_error:
-                        from code_puppy.messaging import emit_warning
-
-                        emit_warning(f"MCP server error: {str(mcp_error)}")
-                        emit_warning("Running without MCP servers...")
-                        # Run without MCP servers as fallback
-                        response = await agent.run(
-                            initial_command, usage_limits=get_custom_usage_limits()
-                        )
-                    finally:
-                        set_message_history(
-                            prune_interrupted_tool_calls(get_message_history())
-                        )
+                    # Use agent_manager's run_with_mcp method
+                    response = await agent_manager.run_with_mcp(
+                        initial_command, usage_limits=get_custom_usage_limits()
+                    )
+                    set_message_history(
+                        prune_interrupted_tool_calls(get_message_history())
+                    )
 
             agent_response = response.output
 
@@ -438,109 +418,35 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
             try:
                 prettier_code_blocks()
 
-                # Store agent's full response
-                agent_response = None
-
-                # Get the agent (uses cached version from early initialization)
-                agent = get_code_generation_agent()
+                # No need to get agent directly - use manager's run methods
 
                 # Use our custom spinner for better compatibility with user input
                 from code_puppy.messaging import emit_warning
                 from code_puppy.messaging.spinner import ConsoleSpinner
 
-                # Create a simple flag to track cancellation locally
-                local_cancelled = False
-
-                # Run with spinner
-                with ConsoleSpinner(console=display_console):
-                    # Use a separate asyncio task that we can cancel
-                    async def run_agent_task():
-                        try:
-                            async with agent.run_mcp_servers():
-                                return await agent.run(
-                                    task,
-                                    message_history=get_message_history(),
-                                    usage_limits=get_custom_usage_limits(),
-                                )
-                        except Exception as mcp_error:
-                            # Handle MCP server errors
-                            emit_warning(f"MCP server error: {str(mcp_error)}")
-                            emit_warning("Running without MCP servers...")
-                            # Run without MCP servers as fallback
-                            return await agent.run(
-                                task,
-                                message_history=get_message_history(),
-                                usage_limits=get_custom_usage_limits(),
-                            )
-                        finally:
-                            set_message_history(
-                                prune_interrupted_tool_calls(get_message_history())
-                            )
-
-                    # Create the task
-                    agent_task = asyncio.create_task(run_agent_task())
-
-                    # Set up signal handling for Ctrl+C
-                    import signal
-
-                    from code_puppy.tools.command_runner import (
-                        kill_all_running_shell_processes,
+                # Use ConsoleSpinner for better user experience
+                try:
+                    with ConsoleSpinner(console=display_console):
+                        # The manager handles all cancellation logic internally
+                        result = await agent_manager.run_with_mcp(
+                            task,
+                            message_history=get_message_history(),
+                            usage_limits=get_custom_usage_limits(),
+                        )
+                except asyncio.CancelledError:
+                    # Agent was cancelled by user
+                    result = None
+                except KeyboardInterrupt:
+                    # Keyboard interrupt
+                    emit_warning("\n⚠️ Caught KeyboardInterrupt in main")
+                    result = None
+                finally:
+                    set_message_history(
+                        prune_interrupted_tool_calls(get_message_history())
                     )
 
-                    original_handler = None
-
-                    # Ensure the interrupt handler only acts once per task
-                    handled = False
-
-                    def keyboard_interrupt_handler(sig, frame):
-                        nonlocal local_cancelled
-                        nonlocal handled
-                        if handled:
-                            return
-                        handled = True
-                        # First, nuke any running shell processes triggered by tools
-                        try:
-                            killed = kill_all_running_shell_processes()
-                            if killed:
-                                from code_puppy.messaging import emit_warning
-
-                                emit_warning(
-                                    f"Cancelled {killed} running shell process(es)."
-                                )
-                            else:
-                                # Then cancel the agent task
-                                if not agent_task.done():
-                                    state_management._message_history = (
-                                        prune_interrupted_tool_calls(
-                                            state_management._message_history
-                                        )
-                                    )
-                                    agent_task.cancel()
-                                    local_cancelled = True
-                        except Exception as e:
-                            from code_puppy.messaging import emit_warning
-
-                            emit_warning(f"Shell kill error: {e}")
-                        # Don't call the original handler
-                        # This prevents the application from exiting
-
-                    try:
-                        # Save original handler and set our custom one
-                        original_handler = signal.getsignal(signal.SIGINT)
-                        signal.signal(signal.SIGINT, keyboard_interrupt_handler)
-
-                        # Wait for the task to complete or be cancelled
-                        result = await agent_task
-                    except asyncio.CancelledError:
-                        # Task was cancelled by our handler
-                        pass
-                    finally:
-                        # Restore original signal handler
-                        if original_handler:
-                            signal.signal(signal.SIGINT, original_handler)
-
                 # Check if the task was cancelled
-                if local_cancelled:
+                if result is None:
                     emit_warning("\n⚠️ Processing cancelled by user (Ctrl+C)")
                     # Skip the rest of this loop iteration
                     continue
@@ -605,37 +511,26 @@ async def execute_single_prompt(prompt: str, message_renderer) -> None:
     emit_info(f"[bold blue]Executing prompt:[/bold blue] {prompt}")
 
     try:
-        # Get the agent
-        agent = get_code_generation_agent()
-
-        # Use our custom spinner for better compatibility with user input
+        # Get agent through runtime manager and use its run_with_mcp method
+        agent_manager = get_runtime_agent_manager()
+        
         from code_puppy.messaging.spinner import ConsoleSpinner
-
-        display_console = message_renderer.console
-        with ConsoleSpinner(console=display_console):
-            try:
-                async with agent.run_mcp_servers():
-                    response = await agent.run(
-                        prompt, usage_limits=get_custom_usage_limits()
-                    )
-            except Exception as mcp_error:
-                from code_puppy.messaging import emit_warning
-
-                emit_warning(f"MCP server error: {str(mcp_error)}")
-                emit_warning("Running without MCP servers...")
-                # Run without MCP servers as fallback
-                response = await agent.run(
-                    prompt, usage_limits=get_custom_usage_limits()
-                )
+        with ConsoleSpinner(console=message_renderer.console):
+            response = await agent_manager.run_with_mcp(
+                prompt, 
+                usage_limits=get_custom_usage_limits()
+            )
 
         agent_response = response.output
         emit_system_message(
             f"\n[bold purple]AGENT RESPONSE: [/bold purple]\n{agent_response}"
         )
 
+    except asyncio.CancelledError:
+        from code_puppy.messaging import emit_warning
+        emit_warning("Execution cancelled by user")
     except Exception as e:
         from code_puppy.messaging import emit_error
-
         emit_error(f"Error executing prompt: {str(e)}")
 
 

code_puppy/mcp/__init__.py

@@ -0,0 +1,23 @@
+"""MCP (Model Context Protocol) management system for Code Puppy."""
+
+from .managed_server import ManagedMCPServer, ServerConfig, ServerState
+from .status_tracker import ServerStatusTracker, Event
+from .manager import MCPManager, ServerInfo, get_mcp_manager
+from .registry import ServerRegistry
+from .error_isolation import MCPErrorIsolator, ErrorStats, ErrorCategory, QuarantinedServerError, get_error_isolator
+from .circuit_breaker import CircuitBreaker, CircuitState, CircuitOpenError
+from .retry_manager import RetryManager, RetryStats, get_retry_manager, retry_mcp_call
+from .dashboard import MCPDashboard
+from .config_wizard import MCPConfigWizard, run_add_wizard
+
+__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'
+]

code_puppy/mcp/async_lifecycle.py

@@ -0,0 +1,237 @@
+"""
+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 typing import Dict, Optional, Any, Union
+from datetime import datetime
+from dataclasses import dataclass
+from contextlib import AsyncExitStack
+
+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 a task that will manage this server's lifecycle
+            task = asyncio.create_task(
+                self._server_lifecycle_task(server_id, server),
+                name=f"mcp_server_{server_id}"
+            )
+            
+            # Wait briefly for the server to start
+            await asyncio.sleep(0.1)
+            
+            # Check if task failed immediately
+            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} starting in background task")
+            return True
+    
+    async def _server_lifecycle_task(
+        self,
+        server_id: str,
+        server: Union[MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP]
+    ) -> 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}")
+            
+            # Enter the server's context
+            await exit_stack.enter_async_context(server)
+            
+            # 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()
+                )
+            
+            logger.info(f"Server {server_id} started successfully")
+            
+            # Keep the task alive until cancelled
+            while True:
+                await asyncio.sleep(1)
+                
+                # Check if server is still running
+                if not server.is_running:
+                    logger.warning(f"Server {server_id} stopped unexpectedly")
+                    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}")
+        finally:
+            # Clean up the context
+            await exit_stack.aclose()
+            
+            # 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