code-puppy 0.0.127__py3-none-any.whl → 0.0.129__py3-none-any.whl

code_puppy/__init__.py

@@ -1,3 +1,4 @@
 import importlib.metadata
 
+# Biscuit was here! 🐶
 __version__ = importlib.metadata.version("code-puppy")

code_puppy/agent.py

@@ -1,3 +1,4 @@
+import uuid
 from pathlib import Path
 from typing import Dict, Optional
 
@@ -42,7 +43,9 @@ _code_generation_agent = None
 
 
 def _load_mcp_servers(extra_headers: Optional[Dict[str, str]] = None):
+    """Load MCP servers using the new manager while maintaining backward compatibility."""
     from code_puppy.config import get_value, load_mcp_server_configs
+    from code_puppy.mcp import get_mcp_manager, ServerConfig
 
     # Check if MCP servers are disabled
     mcp_disabled = get_value("disable_mcp_servers")
@@ -50,86 +53,76 @@ def _load_mcp_servers(extra_headers: Optional[Dict[str, str]] = None):
         emit_system_message("[dim]MCP servers disabled via config[/dim]")
         return []
 
+    # Get the MCP manager singleton
+    manager = get_mcp_manager()
+    
+    # Load configurations from legacy file for backward compatibility
     configs = load_mcp_server_configs()
     if not configs:
-        emit_system_message("[dim]No MCP servers configured[/dim]")
-        return []
-    servers = []
-    for name, conf in configs.items():
-        server_type = conf.get("type", "sse")
-        url = conf.get("url")
-        timeout = conf.get("timeout", 30)
-        server_headers = {}
-        if extra_headers:
-            server_headers.update(extra_headers)
-        user_headers = conf.get("headers") or {}
-        if isinstance(user_headers, dict) and user_headers:
+        # Check if manager already has servers (could be from new system)
+        existing_servers = manager.list_servers()
+        if not existing_servers:
+            emit_system_message("[dim]No MCP servers configured[/dim]")
+            return []
+    else:
+        # Register servers from legacy config with manager
+        for name, conf in configs.items():
             try:
-                user_headers = resolve_env_var_in_header(user_headers)
-            except Exception:
-                pass
-            server_headers.update(user_headers)
-        http_client = None
-
-        try:
-            if server_type == "http" and url:
-                emit_system_message(
-                    f"Registering MCP Server (HTTP) - {url} (timeout: {timeout}s, headers: {bool(server_headers)})"
-                )
-                http_client = create_reopenable_async_client(
-                    timeout=timeout, headers=server_headers or None, verify=False
+                # Convert legacy format to new ServerConfig
+                server_config = ServerConfig(
+                    id=conf.get("id", f"{name}_{hash(name)}"),
+                    name=name,
+                    type=conf.get("type", "sse"),
+                    enabled=conf.get("enabled", True),
+                    config=conf
                 )
-                servers.append(
-                    MCPServerStreamableHTTP(url=url, http_client=http_client)
-                )
-            elif (
-                server_type == "stdio"
-            ):  # Fixed: was "stdios" (plural), should be "stdio" (singular)
-                command = conf.get("command")
-                args = conf.get("args", [])
-                timeout = conf.get(
-                    "timeout", 30
-                )  # Default 30 seconds for stdio servers (npm downloads can be slow)
-                if command:
-                    emit_system_message(
-                        f"Registering MCP Server (Stdio) - {command} {args} (timeout: {timeout}s)"
-                    )
-                    servers.append(MCPServerStdio(command, args=args, timeout=timeout))
+                
+                # Check if server already registered
+                existing = manager.get_server_by_name(name)
+                if not existing:
+                    # Register new server
+                    manager.register_server(server_config)
+                    emit_system_message(f"[dim]Registered MCP server: {name}[/dim]")
                 else:
-                    emit_error(f"MCP Server '{name}' missing required 'command' field")
-            elif server_type == "sse" and url:
-                emit_system_message(
-                    f"Registering MCP Server (SSE) - {url} (timeout: {timeout}s, headers: {bool(server_headers)})"
-                )
-                # For SSE, allow long reads; only bound connect timeout
-                http_client = create_reopenable_async_client(
-                    timeout=30, headers=server_headers or None, verify=False
-                )
-                servers.append(MCPServerSSE(url=url, http_client=http_client))
-            else:
-                emit_error(
-                    f"Invalid type '{server_type}' or missing URL for MCP server '{name}'"
-                )
-        except Exception as e:
-            emit_error(f"Failed to register MCP server '{name}': {str(e)}")
-            emit_info(f"Skipping server '{name}' and continuing with other servers...")
-            # Continue with other servers instead of crashing
-            continue
-
+                    # Update existing server config if needed
+                    if existing.config != server_config.config:
+                        manager.update_server(existing.id, server_config)
+                        emit_system_message(f"[dim]Updated MCP server: {name}[/dim]")
+                        
+            except Exception as e:
+                emit_error(f"Failed to register MCP server '{name}': {str(e)}")
+                continue
+    
+    # Get pydantic-ai compatible servers from manager
+    servers = manager.get_servers_for_agent()
+    
     if servers:
         emit_system_message(
-            f"[green]Successfully registered {len(servers)} MCP server(s)[/green]"
+            f"[green]Successfully loaded {len(servers)} MCP server(s)[/green]"
         )
     else:
         emit_system_message(
-            "[yellow]No MCP servers were successfully registered[/yellow]"
+            "[yellow]No MCP servers available (check if servers are enabled)[/yellow]"
         )
-
+    
     return servers
 
 
-def reload_code_generation_agent():
+def reload_mcp_servers():
+    """Reload MCP servers without restarting the agent."""
+    from code_puppy.mcp import get_mcp_manager
+    
+    manager = get_mcp_manager()
+    # Reload configurations
+    _load_mcp_servers()
+    # Return updated servers
+    return manager.get_servers_for_agent()
+
+
+def reload_code_generation_agent(message_group: str | None):
     """Force-reload the agent, usually after a model change."""
+    if message_group is None:
+        message_group = str(uuid.uuid4())
     global _code_generation_agent, _LAST_MODEL_NAME
     from code_puppy.config import clear_model_cache, get_model_name
     from code_puppy.agents import clear_agent_cache
@@ -139,14 +132,15 @@ def reload_code_generation_agent():
     clear_agent_cache()
 
     model_name = get_model_name()
-    emit_info(f"[bold cyan]Loading Model: {model_name}[/bold cyan]")
+    emit_info(f"[bold cyan]Loading Model: {model_name}[/bold cyan]", message_group=message_group)
     models_config = ModelFactory.load_config()
     model = ModelFactory.get_model(model_name, models_config)
 
     # Get agent-specific system prompt
     agent_config = get_current_agent_config()
     emit_info(
-        f"[bold magenta]Loading Agent: {agent_config.display_name}[/bold magenta]"
+        f"[bold magenta]Loading Agent: {agent_config.display_name}[/bold magenta]",
+        message_group=message_group
     )
 
     instructions = agent_config.get_system_prompt()
@@ -183,17 +177,19 @@ def reload_code_generation_agent():
     return _code_generation_agent
 
 
-def get_code_generation_agent(force_reload=False):
+def get_code_generation_agent(force_reload=False, message_group: str | None = None):
     """
     Retrieve the agent with the currently configured model.
     Forces a reload if the model has changed, or if force_reload is passed.
     """
     global _code_generation_agent, _LAST_MODEL_NAME
+    if message_group is None:
+        message_group = str(uuid.uuid4())
     from code_puppy.config import get_model_name
 
     model_name = get_model_name()
     if _code_generation_agent is None or _LAST_MODEL_NAME != model_name or force_reload:
-        return reload_code_generation_agent()
+        return reload_code_generation_agent(message_group)
     return _code_generation_agent
 
 

code_puppy/agents/agent_code_puppy.py

@@ -140,9 +140,6 @@ Important rules:
 Your solutions should be production-ready, maintainable, and follow best practices for the chosen language.
 
 Return your final response as a string output
-Walmart specific rules:
- - You are operating inside Walmart Global Tech! Yay!
- - Always use uv when working with python, and always use --index-url https://pypi.ci.artifacts.walmart.com/artifactory/api/pypi/external-pypi/simple
 """
 
         prompt_additions = callbacks.on_load_prompt()

code_puppy/agents/runtime_manager.py

@@ -0,0 +1,231 @@
+"""
+Runtime agent manager that ensures proper agent instance updates.
+
+This module provides a wrapper around the agent singleton that ensures
+all references to the agent are properly updated when it's reloaded.
+"""
+
+import asyncio
+import signal
+import uuid
+from typing import Optional, Any
+
+import mcp
+from pydantic_ai import Agent
+from pydantic_ai.usage import UsageLimits
+
+from code_puppy.messaging.message_queue import emit_info, emit_warning
+
+
+class RuntimeAgentManager:
+    """
+    Manages the runtime agent instance and ensures proper updates.
+    
+    This class acts as a proxy that always returns the current agent instance,
+    ensuring that when the agent is reloaded, all code using this manager
+    automatically gets the updated instance.
+    """
+    
+    def __init__(self):
+        """Initialize the runtime agent manager."""
+        self._agent: Optional[Agent] = None
+        self._last_model_name: Optional[str] = None
+        
+    def get_agent(self, force_reload: bool = False, message_group: str = "") -> Agent:
+        """
+        Get the current agent instance.
+        
+        This method always returns the most recent agent instance,
+        automatically handling reloads when the model changes.
+        
+        Args:
+            force_reload: If True, force a reload of the agent
+            
+        Returns:
+            The current agent instance
+        """
+        from code_puppy.agent import get_code_generation_agent
+        
+        # Always get the current singleton - this ensures we have the latest
+        current_agent = get_code_generation_agent(force_reload=force_reload, message_group=message_group)
+        self._agent = current_agent
+        
+        return self._agent
+    
+    def reload_agent(self) -> Agent:
+        """
+        Force reload the agent.
+        
+        This is typically called after MCP servers are started/stopped.
+        
+        Returns:
+            The newly loaded agent instance
+        """
+        message_group = uuid.uuid4()
+        emit_info("[bold cyan]Reloading agent with updated configuration...[/bold cyan]", message_group=message_group)
+        return self.get_agent(force_reload=True, message_group=message_group)
+    
+    async def run_with_mcp(self, prompt: str, usage_limits: Optional[UsageLimits] = None, **kwargs) -> Any:
+        """
+        Run the agent with MCP servers and full cancellation support.
+        
+        This method ensures we're always using the current agent instance
+        and handles Ctrl+C interruption properly by creating a cancellable task.
+        
+        Args:
+            prompt: The user prompt to process
+            usage_limits: Optional usage limits for the agent
+            **kwargs: Additional arguments to pass to agent.run (e.g., message_history)
+            
+        Returns:
+            The agent's response
+            
+        Raises:
+            asyncio.CancelledError: When execution is cancelled by user
+        """
+        agent = self.get_agent()
+        group_id = str(uuid.uuid4())
+        # Function to run agent with MCP
+        async def run_agent_task():
+            try:
+                async with agent:
+                    return await agent.run(prompt, usage_limits=usage_limits, **kwargs)
+            except* mcp.shared.exceptions.McpError as mcp_error:
+                emit_warning(f"MCP server error: {str(mcp_error)}", group_id=group_id)
+                emit_warning(f"{str(mcp_error)}", group_id=group_id)
+                emit_warning(f"Try disabling any malfunctioning MCP servers", group_id=group_id)
+            except* InterruptedError as ie:
+                emit_warning(f"Interrupted: {str(ie)}")
+            except* Exception as other_error:
+                # Filter out CancelledError from the exception group - let it propagate
+                remaining_exceptions = []
+                def collect_non_cancelled_exceptions(exc):
+                    if isinstance(exc, ExceptionGroup):
+                        for sub_exc in exc.exceptions:
+                            collect_non_cancelled_exceptions(sub_exc)
+                    elif not isinstance(exc, asyncio.CancelledError):
+                        remaining_exceptions.append(exc)
+                        emit_warning(f"Unexpected error: {str(exc)}", group_id=group_id)
+                        emit_warning(f"{str(exc.args)}", group_id=group_id)
+                
+                collect_non_cancelled_exceptions(other_error)
+                
+                # If there are CancelledError exceptions in the group, re-raise them
+                cancelled_exceptions = []
+                def collect_cancelled_exceptions(exc):
+                    if isinstance(exc, ExceptionGroup):
+                        for sub_exc in exc.exceptions:
+                            collect_cancelled_exceptions(sub_exc)
+                    elif isinstance(exc, asyncio.CancelledError):
+                        cancelled_exceptions.append(exc)
+                
+                collect_cancelled_exceptions(other_error)
+                
+                if cancelled_exceptions:
+                    # Re-raise the first CancelledError to propagate cancellation
+                    raise cancelled_exceptions[0]
+        
+        # Create the task FIRST
+        agent_task = asyncio.create_task(run_agent_task())
+        
+        # Import shell process killer
+        from code_puppy.tools.command_runner import kill_all_running_shell_processes
+        
+        # Ensure the interrupt handler only acts once per task
+        handled = False
+        
+        def keyboard_interrupt_handler(sig, frame):
+            """Signal handler for Ctrl+C - replicating exact original logic"""
+            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:
+                    emit_warning(f"Cancelled {killed} running shell process(es).")
+                else:
+                    # Only cancel the agent task if no shell processes were killed
+                    if not agent_task.done():
+                        agent_task.cancel()
+            except Exception as e:
+                emit_warning(f"Shell kill error: {e}")
+                # If shell kill failed, still try to cancel the agent task
+                if not agent_task.done():
+                    agent_task.cancel()
+            # Don't call the original handler
+            # This prevents the application from exiting
+        
+        try:
+            # Save original handler and set our custom one AFTER task is created
+            original_handler = signal.signal(signal.SIGINT, keyboard_interrupt_handler)
+            
+            # Wait for the task to complete or be cancelled
+            result = await agent_task
+            return result
+        except asyncio.CancelledError:
+            # Task was cancelled by our handler
+            raise
+        except KeyboardInterrupt:
+            # Handle direct keyboard interrupt during await
+            if not agent_task.done():
+                agent_task.cancel()
+            try:
+                await agent_task
+            except asyncio.CancelledError:
+                pass
+            raise asyncio.CancelledError()
+        finally:
+            # Restore original signal handler
+            if original_handler:
+                signal.signal(signal.SIGINT, original_handler)
+    
+    async def run(self, prompt: str, usage_limits: Optional[UsageLimits] = None, **kwargs) -> Any:
+        """
+        Run the agent without explicitly managing MCP servers.
+        
+        Args:
+            prompt: The user prompt to process
+            usage_limits: Optional usage limits for the agent
+            **kwargs: Additional arguments to pass to agent.run (e.g., message_history)
+            
+        Returns:
+            The agent's response
+        """
+        agent = self.get_agent()
+        return await agent.run(prompt, usage_limits=usage_limits, **kwargs)
+    
+    def __getattr__(self, name: str) -> Any:
+        """
+        Proxy all other attribute access to the current agent.
+        
+        This allows the manager to be used as a drop-in replacement
+        for direct agent access.
+        
+        Args:
+            name: The attribute name to access
+            
+        Returns:
+            The attribute from the current agent
+        """
+        agent = self.get_agent()
+        return getattr(agent, name)
+
+
+# Global singleton instance
+_runtime_manager: Optional[RuntimeAgentManager] = None
+
+
+def get_runtime_agent_manager() -> RuntimeAgentManager:
+    """
+    Get the global runtime agent manager instance.
+    
+    Returns:
+        The singleton RuntimeAgentManager instance
+    """
+    global _runtime_manager
+    if _runtime_manager is None:
+        _runtime_manager = RuntimeAgentManager()
+    return _runtime_manager

code_puppy/command_line/command_handler.py

@@ -9,23 +9,42 @@ from code_puppy.command_line.utils import make_directory_table
 from code_puppy.config import get_config_keys
 from code_puppy.tools.tools_content import tools_content
 
-COMMANDS_HELP = """
-[bold magenta]Commands Help[/bold magenta]
-/help, /h             Show this help message
-/cd <dir>             Change directory or show directories
-/agent <name>         Switch to a different agent or show available agents
-/exit, /quit          Exit interactive mode
-/generate-pr-description [@dir]  Generate comprehensive PR description
-/m <model>            Set active model
-/motd                 Show the latest message of the day (MOTD)
-/show                 Show puppy config key-values
-/compact              Summarize and compact current chat history
-/dump_context <name>  Save current message history to file
-/load_context <name>  Load message history from file
-/set                  Set puppy config key-values (e.g., /set yolo_mode true, /set compaction_strategy truncation)
-/tools                Show available tools and capabilities
-/<unknown>            Show unknown command warning
-"""
+def get_commands_help():
+    """Generate commands help using Rich Text objects to avoid markup conflicts."""
+    from rich.text import Text
+    
+    # Build help text programmatically
+    help_lines = []
+    
+    # Title
+    help_lines.append(Text("Commands Help", style="bold magenta"))
+    
+    # Commands - build each line programmatically  
+    help_lines.append(Text("/help, /h", style="cyan") + Text("             Show this help message"))
+    help_lines.append(Text("/cd", style="cyan") + Text(" <dir>             Change directory or show directories"))
+    help_lines.append(Text("/agent", style="cyan") + Text(" <name>         Switch to a different agent or show available agents"))
+    help_lines.append(Text("/exit, /quit", style="cyan") + Text("          Exit interactive mode"))
+    help_lines.append(Text("/generate-pr-description", style="cyan") + Text(" [@dir]  Generate comprehensive PR description"))
+    help_lines.append(Text("/model", style="cyan") + Text(" <model>        Set active model"))
+    help_lines.append(Text("/mcp", style="cyan") + Text("                  Manage MCP servers (list, start, stop, status, etc.)"))
+    help_lines.append(Text("/motd", style="cyan") + Text("                 Show the latest message of the day (MOTD)"))
+    help_lines.append(Text("/show", style="cyan") + Text("                 Show puppy config key-values"))
+    help_lines.append(Text("/compact", style="cyan") + Text("              Summarize and compact current chat history"))
+    help_lines.append(Text("/dump_context", style="cyan") + Text(" <name>  Save current message history to file"))
+    help_lines.append(Text("/load_context", style="cyan") + Text(" <name>  Load message history from file"))
+    help_lines.append(Text("/set", style="cyan") + Text("                  Set puppy config key-values (e.g., /set yolo_mode true, /set compaction_strategy truncation)"))
+    help_lines.append(Text("/tools", style="cyan") + Text("                Show available tools and capabilities"))
+    help_lines.append(Text("/<unknown>", style="cyan") + Text("            Show unknown command warning"))
+    
+    
+    # Combine all lines
+    final_text = Text()
+    for i, line in enumerate(help_lines):
+        if i > 0:
+            final_text.append("\n")
+        final_text.append_text(line)
+    
+    return final_text
 
 
 def handle_command(command: str):
@@ -220,7 +239,7 @@ def handle_command(command: str):
             set_current_agent,
             get_agent_descriptions,
         )
-        from code_puppy.agent import get_code_generation_agent
+        from code_puppy.agents.runtime_manager import get_runtime_agent_manager
 
         tokens = command.split()
 
@@ -272,7 +291,8 @@ def handle_command(command: str):
 
             if set_current_agent(agent_name):
                 # Reload the agent with new configuration
-                get_code_generation_agent(force_reload=True)
+                manager = get_runtime_agent_manager()
+                manager.reload_agent()
                 new_agent = get_current_agent_config()
                 emit_success(
                     f"Switched to agent: {new_agent.display_name}",
@@ -297,25 +317,36 @@ def handle_command(command: str):
             emit_warning("Usage: /agent [agent-name]")
             return True
 
-    if command.startswith("/m"):
+    if command.startswith("/model"):
         # Try setting model and show confirmation
-        new_input = update_model_in_input(command)
+        # Handle both /model and /m for backward compatibility
+        model_command = command.replace("/model", "/m") if command.startswith("/model") else command
+        new_input = update_model_in_input(model_command)
         if new_input is not None:
-            from code_puppy.agent import get_code_generation_agent
+            from code_puppy.agents.runtime_manager import get_runtime_agent_manager
             from code_puppy.command_line.model_picker_completion import get_active_model
 
             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()
             emit_success(f"Active model set and loaded: {model}")
             return True
         # If no model matched, show available models
         model_names = load_model_names()
-        emit_warning("Usage: /m <model-name>")
+        emit_warning("Usage: /model <model-name>")
         emit_warning(f"Available models: {', '.join(model_names)}")
         return True
+    
+    if command.startswith("/mcp"):
+        from code_puppy.command_line.mcp_commands import MCPCommandHandler
+        handler = MCPCommandHandler()
+        return handler.handle_mcp_command(command)
     if command in ("/help", "/h"):
-        emit_info(COMMANDS_HELP)
+        import uuid
+        group_id = str(uuid.uuid4())
+        help_text = get_commands_help()
+        emit_info(help_text, message_group_id=group_id)
         return True
 
     if command.startswith("/generate-pr-description"):