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

code_puppy/mcp/status_tracker.py

@@ -0,0 +1,355 @@
+"""
+Server Status Tracker for monitoring MCP server runtime status.
+
+This module provides the ServerStatusTracker class that tracks the runtime 
+status of MCP servers including state, metrics, and events.
+"""
+
+import logging
+import threading
+from collections import deque, defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional
+
+from .managed_server import ServerState
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Event:
+    """Data class representing a server event."""
+    timestamp: datetime
+    event_type: str  # "started", "stopped", "error", "health_check", etc.
+    details: Dict
+    server_id: str
+
+
+class ServerStatusTracker:
+    """
+    Tracks the runtime status of MCP servers including state, metrics, and events.
+    
+    This class provides in-memory storage for server states, metadata, and events
+    with thread-safe operations using locks. Events are stored using collections.deque
+    for automatic size limiting.
+    
+    Example usage:
+        tracker = ServerStatusTracker()
+        tracker.set_status("server1", ServerState.RUNNING)
+        tracker.record_event("server1", "started", {"message": "Server started successfully"})
+        events = tracker.get_events("server1", limit=10)
+    """
+    
+    def __init__(self):
+        """Initialize the status tracker with thread-safe data structures."""
+        # Thread safety lock
+        self._lock = threading.RLock()
+        
+        # Server states (server_id -> ServerState)
+        self._server_states: Dict[str, ServerState] = {}
+        
+        # Server metadata (server_id -> key -> value)
+        self._server_metadata: Dict[str, Dict[str, Any]] = defaultdict(dict)
+        
+        # Server events (server_id -> deque of events)
+        # Using deque with maxlen for automatic size limiting
+        self._server_events: Dict[str, deque] = defaultdict(lambda: deque(maxlen=1000))
+        
+        # Server timing information
+        self._start_times: Dict[str, datetime] = {}
+        self._stop_times: Dict[str, datetime] = {}
+        
+        logger.info("ServerStatusTracker initialized")
+    
+    def set_status(self, server_id: str, state: ServerState) -> None:
+        """
+        Set the current state of a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+            state: New server state
+        """
+        with self._lock:
+            old_state = self._server_states.get(server_id)
+            self._server_states[server_id] = state
+            
+            # Record state change event
+            self.record_event(
+                server_id,
+                "state_change",
+                {
+                    "old_state": old_state.value if old_state else None,
+                    "new_state": state.value,
+                    "message": f"State changed from {old_state.value if old_state else 'unknown'} to {state.value}"
+                }
+            )
+            
+            logger.debug(f"Server {server_id} state changed: {old_state} -> {state}")
+    
+    def get_status(self, server_id: str) -> ServerState:
+        """
+        Get the current state of a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+            
+        Returns:
+            Current server state, defaults to STOPPED if not found
+        """
+        with self._lock:
+            return self._server_states.get(server_id, ServerState.STOPPED)
+    
+    def set_metadata(self, server_id: str, key: str, value: Any) -> None:
+        """
+        Set metadata value for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+            key: Metadata key
+            value: Metadata value (can be any type)
+        """
+        with self._lock:
+            if server_id not in self._server_metadata:
+                self._server_metadata[server_id] = {}
+            
+            old_value = self._server_metadata[server_id].get(key)
+            self._server_metadata[server_id][key] = value
+            
+            # Record metadata change event
+            self.record_event(
+                server_id,
+                "metadata_update",
+                {
+                    "key": key,
+                    "old_value": old_value,
+                    "new_value": value,
+                    "message": f"Metadata '{key}' updated"
+                }
+            )
+            
+            logger.debug(f"Server {server_id} metadata updated: {key} = {value}")
+    
+    def get_metadata(self, server_id: str, key: str) -> Any:
+        """
+        Get metadata value for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+            key: Metadata key
+            
+        Returns:
+            Metadata value or None if not found
+        """
+        with self._lock:
+            return self._server_metadata.get(server_id, {}).get(key)
+    
+    def record_event(self, server_id: str, event_type: str, details: Dict) -> None:
+        """
+        Record an event for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+            event_type: Type of event (e.g., "started", "stopped", "error", "health_check")
+            details: Dictionary containing event details
+        """
+        with self._lock:
+            event = Event(
+                timestamp=datetime.now(),
+                event_type=event_type,
+                details=details.copy() if details else {},  # Copy to prevent modification
+                server_id=server_id
+            )
+            
+            # Add to deque (automatically handles size limiting)
+            self._server_events[server_id].append(event)
+            
+            logger.debug(f"Event recorded for server {server_id}: {event_type}")
+    
+    def get_events(self, server_id: str, limit: int = 100) -> List[Event]:
+        """
+        Get recent events for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+            limit: Maximum number of events to return (default: 100)
+            
+        Returns:
+            List of events ordered by timestamp (most recent first)
+        """
+        with self._lock:
+            events = list(self._server_events.get(server_id, deque()))
+            
+            # Return most recent events first, limited by count
+            events.reverse()  # Most recent first
+            return events[:limit]
+    
+    def clear_events(self, server_id: str) -> None:
+        """
+        Clear all events for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+        """
+        with self._lock:
+            if server_id in self._server_events:
+                self._server_events[server_id].clear()
+                logger.info(f"Cleared all events for server: {server_id}")
+    
+    def get_uptime(self, server_id: str) -> Optional[timedelta]:
+        """
+        Calculate uptime for a server based on start/stop times.
+        
+        Args:
+            server_id: Unique identifier for the server
+            
+        Returns:
+            Server uptime as timedelta, or None if server never started
+        """
+        with self._lock:
+            start_time = self._start_times.get(server_id)
+            if start_time is None:
+                return None
+            
+            # If server is currently running, calculate from start time to now
+            current_state = self.get_status(server_id)
+            if current_state == ServerState.RUNNING:
+                return datetime.now() - start_time
+            
+            # If server is stopped, calculate from start to stop time
+            stop_time = self._stop_times.get(server_id)
+            if stop_time is not None and stop_time > start_time:
+                return stop_time - start_time
+            
+            # If we have start time but no valid stop time, assume currently running
+            return datetime.now() - start_time
+    
+    def record_start_time(self, server_id: str) -> None:
+        """
+        Record the start time for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+        """
+        with self._lock:
+            start_time = datetime.now()
+            self._start_times[server_id] = start_time
+            
+            # Record start event
+            self.record_event(
+                server_id,
+                "started",
+                {
+                    "start_time": start_time.isoformat(),
+                    "message": "Server started"
+                }
+            )
+            
+            logger.info(f"Recorded start time for server: {server_id}")
+    
+    def record_stop_time(self, server_id: str) -> None:
+        """
+        Record the stop time for a server.
+        
+        Args:
+            server_id: Unique identifier for the server
+        """
+        with self._lock:
+            stop_time = datetime.now()
+            self._stop_times[server_id] = stop_time
+            
+            # Calculate final uptime
+            start_time = self._start_times.get(server_id)
+            uptime = None
+            if start_time:
+                uptime = stop_time - start_time
+            
+            # Record stop event
+            self.record_event(
+                server_id,
+                "stopped",
+                {
+                    "stop_time": stop_time.isoformat(),
+                    "uptime_seconds": uptime.total_seconds() if uptime else None,
+                    "message": "Server stopped"
+                }
+            )
+            
+            logger.info(f"Recorded stop time for server: {server_id}")
+    
+    def get_all_server_ids(self) -> List[str]:
+        """
+        Get all server IDs that have been tracked.
+        
+        Returns:
+            List of all server IDs
+        """
+        with self._lock:
+            # Combine all sources of server IDs
+            all_ids = set()
+            all_ids.update(self._server_states.keys())
+            all_ids.update(self._server_metadata.keys())
+            all_ids.update(self._server_events.keys())
+            all_ids.update(self._start_times.keys())
+            all_ids.update(self._stop_times.keys())
+            
+            return sorted(list(all_ids))
+    
+    def get_server_summary(self, server_id: str) -> Dict[str, Any]:
+        """
+        Get comprehensive summary of server status.
+        
+        Args:
+            server_id: Unique identifier for the server
+            
+        Returns:
+            Dictionary containing current state, metadata, recent events, and uptime
+        """
+        with self._lock:
+            return {
+                "server_id": server_id,
+                "state": self.get_status(server_id).value,
+                "metadata": self._server_metadata.get(server_id, {}).copy(),
+                "recent_events_count": len(self._server_events.get(server_id, deque())),
+                "uptime": self.get_uptime(server_id),
+                "start_time": self._start_times.get(server_id),
+                "stop_time": self._stop_times.get(server_id),
+                "last_event_time": (
+                    list(self._server_events.get(server_id, deque()))[-1].timestamp
+                    if server_id in self._server_events and len(self._server_events[server_id]) > 0
+                    else None
+                )
+            }
+    
+    def cleanup_old_data(self, days_to_keep: int = 7) -> None:
+        """
+        Clean up old data to prevent memory bloat.
+        
+        Args:
+            days_to_keep: Number of days of data to keep (default: 7)
+        """
+        cutoff_time = datetime.now() - timedelta(days=days_to_keep)
+        
+        with self._lock:
+            cleaned_servers = []
+            
+            for server_id in list(self._server_events.keys()):
+                events = self._server_events[server_id]
+                if events:
+                    # Filter out old events
+                    original_count = len(events)
+                    # Convert to list, filter, then create new deque
+                    filtered_events = [
+                        event for event in events 
+                        if event.timestamp >= cutoff_time
+                    ]
+                    
+                    # Replace the deque with filtered events
+                    self._server_events[server_id] = deque(filtered_events, maxlen=1000)
+                    
+                    if len(filtered_events) < original_count:
+                        cleaned_servers.append(server_id)
+            
+            if cleaned_servers:
+                logger.info(f"Cleaned old events for {len(cleaned_servers)} servers")

code_puppy/messaging/spinner/textual_spinner.py

@@ -80,7 +80,7 @@ class TextualSpinner(Static):
             self._paused = True
             self._timer.pause()
             # Store current state but don't clear it completely
-            self._previous_state = self.text
+            self._previous_state = self.renderable
             self.update("")
 
     def resume(self):
@@ -94,4 +94,8 @@ class TextualSpinner(Static):
         if self._is_spinning and self._timer and self._paused:
             self._paused = False
             self._timer.resume()
-            self.update_frame_display()
+            # Restore previous state instead of immediately updating display
+            if self._previous_state:
+                self.update(self._previous_state)
+            else:
+                self.update_frame_display()

code_puppy/model_factory.py

@@ -9,10 +9,11 @@ from anthropic import AsyncAnthropic
 from openai import AsyncAzureOpenAI  # For Azure OpenAI client
 from pydantic_ai.models.anthropic import AnthropicModel
 from pydantic_ai.models.gemini import GeminiModel
-from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
 from pydantic_ai.providers.google_gla import GoogleGLAProvider
 from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.providers.cerebras import CerebrasProvider
 
 from . import callbacks
 from .config import EXTRA_MODELS_FILE
@@ -116,7 +117,7 @@ class ModelFactory:
         elif model_type == "openai":
             provider = OpenAIProvider(api_key=os.environ.get("OPENAI_API_KEY", ""))
 
-            model = OpenAIModel(model_name=model_config["name"], provider=provider)
+            model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
             setattr(model, "provider", provider)
             return model
 
@@ -191,7 +192,7 @@ class ModelFactory:
                 max_retries=azure_max_retries,
             )
             provider = OpenAIProvider(openai_client=azure_client)
-            model = OpenAIModel(model_name=model_config["name"], provider=provider)
+            model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
             setattr(model, "provider", provider)
             return model
 
@@ -206,7 +207,7 @@ class ModelFactory:
                 provider_args["api_key"] = api_key
             provider = OpenAIProvider(**provider_args)
 
-            model = OpenAIModel(model_name=model_config["name"], provider=provider)
+            model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
             setattr(model, "provider", provider)
             return model
 
@@ -231,5 +232,19 @@ class ModelFactory:
             google_gla = CustomGoogleGLAProvider(api_key=api_key)
             model = GeminiModel(model_name=model_config["name"], provider=google_gla)
             return model
+        elif model_type == "cerebras":
+            url, headers, verify, api_key = get_custom_config(model_config)
+            client = create_async_client(headers=headers, verify=verify)
+            provider_args = dict(
+                api_key=api_key,
+                http_client=client,
+            )
+            if api_key:
+                provider_args["api_key"] = api_key
+            provider = CerebrasProvider(**provider_args)
+
+            model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
+            setattr(model, "provider", provider)
+            return model
         else:
             raise ValueError(f"Unsupported model type: {model_type}")

code_puppy/models.json

@@ -5,7 +5,7 @@
     "context_length": 400000
   },
   "Cerebras-Qwen3-Coder-480b": {
-    "type": "custom_openai",
+    "type": "cerebras",
     "name": "qwen-3-coder-480b",
     "custom_endpoint": {
       "url": "https://api.cerebras.ai/v1",
@@ -14,7 +14,7 @@
     "context_length": 131072
   },
   "Cerebras-Qwen3-235b-a22b-instruct-2507": {
-    "type": "custom_openai",
+    "type": "cerebras",
     "name": "qwen-3-235b-a22b-instruct-2507",
     "custom_endpoint": {
       "url": "https://api.cerebras.ai/v1",
@@ -23,7 +23,7 @@
     "context_length": 64000
   },
   "Cerebras-gpt-oss-120b": {
-    "type": "custom_openai",
+    "type": "cerebras",
     "name": "gpt-oss-120b",
     "custom_endpoint": {
       "url": "https://api.cerebras.ai/v1",
@@ -32,7 +32,7 @@
     "context_length": 131072
   },
   "Cerebras-Qwen-3-32b": {
-    "type": "custom_openai",
+    "type": "cerebras",
     "name": "qwen-3-32b",
     "custom_endpoint": {
       "url": "https://api.cerebras.ai/v1",
@@ -45,6 +45,24 @@
     "name": "o3",
     "context_length": 200000
   },
+  "grok-4": {
+    "type": "custom_openai",
+    "name": "grok-4",
+    "custom_endpoint": {
+      "url": "https://api.x.ai/v1",
+      "api_key": "$XAI_API_KEY"
+    },
+    "context_length": 256000
+  },
+    "grok-code-fast-1": {
+    "type": "custom_openai",
+    "name": "grok-code-fast-1",
+    "custom_endpoint": {
+      "url": "https://api.x.ai/v1",
+      "api_key": "$XAI_API_KEY"
+    },
+    "context_length": 256000
+  },
   "gemini-2.5-flash-preview-05-20": {
     "type": "gemini",
     "name": "gemini-2.5-flash-preview-05-20",

code_puppy/tui/app.py

@@ -13,6 +13,7 @@ from textual.reactive import reactive
 from textual.widgets import Footer, ListView
 
 from code_puppy.agent import get_code_generation_agent, get_custom_usage_limits
+from code_puppy.agents.runtime_manager import get_runtime_agent_manager
 from code_puppy.command_line.command_handler import handle_command
 from code_puppy.config import (
     get_model_name,
@@ -95,7 +96,7 @@ class CodePuppyTUI(App):
 
     def __init__(self, initial_command: str = None, **kwargs):
         super().__init__(**kwargs)
-        self.agent = None
+        self.agent_manager = None
         self._current_worker = None
         self.initial_command = initial_command
 
@@ -125,7 +126,8 @@ class CodePuppyTUI(App):
         self.current_model = get_model_name()
         self.puppy_name = get_puppy_name()
 
-        self.agent = get_code_generation_agent()
+        # Use runtime manager to ensure we always have the current agent
+        self.agent_manager = get_runtime_agent_manager()
 
         # Update status bar
         status_bar = self.query_one(StatusBar)
@@ -138,6 +140,10 @@ class CodePuppyTUI(App):
             "Welcome to Code Puppy 🐶!\n💨 YOLO mode is enabled in TUI: commands will execute without confirmation."
         )
 
+        # Get current agent and display info
+        get_code_generation_agent()
+        self.add_system_message(f"🐕 Loaded agent '{self.puppy_name}' with model '{self.current_model}'")
+
         # Start the message renderer EARLY to catch startup messages
         # Using call_after_refresh to start it as soon as possible after mount
         self.call_after_refresh(self.start_message_renderer_sync)
@@ -413,8 +419,7 @@ class CodePuppyTUI(App):
                 if message.strip().startswith("/agent"):
                     # The command handler will emit messages directly to our messaging system
                     handle_command(message.strip())
-                    # Refresh our agent instance after potential change
-                    self.agent = get_code_generation_agent()
+                    # Agent manager will automatically use the latest agent
                     return
 
                 # Handle exit commands
@@ -435,31 +440,18 @@ class CodePuppyTUI(App):
                 return
 
             # Process with agent
-            if self.agent:
+            if self.agent_manager:
                 try:
                     self.update_agent_progress("Processing", 25)
 
-                    # Handle MCP servers with specific TaskGroup exception handling
+                    # Use agent_manager's run_with_mcp to handle MCP servers properly
                     try:
-                        try:
-                            async with self.agent.run_mcp_servers():
-                                self.update_agent_progress("Processing", 50)
-                                result = await self.agent.run(
-                                    message,
-                                    message_history=get_message_history(),
-                                    usage_limits=get_custom_usage_limits(),
-                                )
-                        except Exception as mcp_error:
-                            # Log MCP error and fall back to running without MCP servers
-                            self.log(f"MCP server error: {str(mcp_error)}")
-                            self.add_system_message(
-                                "⚠️ MCP server error, running without MCP servers"
-                            )
-                            result = await self.agent.run(
-                                message,
-                                message_history=get_message_history(),
-                                usage_limits=get_custom_usage_limits(),
-                            )
+                        self.update_agent_progress("Processing", 50)
+                        result = await self.agent_manager.run_with_mcp(
+                            message,
+                            message_history=get_message_history(),
+                            usage_limits=get_custom_usage_limits(),
+                        )
 
                         if not result or not hasattr(result, "output"):
                             self.add_error_message("Invalid response format from agent")
@@ -496,7 +488,7 @@ class CodePuppyTUI(App):
                         f"Agent processing failed: {str(agent_error)}"
                     )
             else:
-                self.add_error_message("Agent not initialized")
+                self.add_error_message("Agent manager not initialized")
 
         except Exception as e:
             self.add_error_message(f"Error processing message: {str(e)}")
@@ -618,7 +610,7 @@ class CodePuppyTUI(App):
                     new_model = get_model_name()
                     self.current_model = new_model
                     # Reinitialize agent with new model
-                    self.agent = get_code_generation_agent()
+                    self.agent_manager.reload_agent()
 
                 # Update status bar
                 status_bar = self.query_one(StatusBar)

code_puppy/tui/tests/test_agent_command.py

@@ -8,16 +8,18 @@ from code_puppy.tui.app import CodePuppyTUI
 class TestTUIAgentCommand:
     """Test the TUI's handling of /agent commands."""
 
-    @patch("code_puppy.tui.app.get_code_generation_agent")
+    @patch("code_puppy.tui.app.get_runtime_agent_manager")
     @patch("code_puppy.tui.app.handle_command")
-    def test_tui_handles_agent_command(self, mock_handle_command, mock_get_agent):
+    def test_tui_handles_agent_command(self, mock_handle_command, mock_get_manager):
         """Test that TUI properly delegates /agent commands to command handler."""
         # Create a TUI app instance
         app = CodePuppyTUI()
 
-        # Mock the agent
+        # Mock the agent manager and agent
         mock_agent_instance = MagicMock()
-        mock_get_agent.return_value = mock_agent_instance
+        mock_manager = MagicMock()
+        mock_manager.get_agent.return_value = mock_agent_instance
+        mock_get_manager.return_value = mock_manager
 
         # Mock handle_command to simulate successful processing
         mock_handle_command.return_value = True
@@ -44,22 +46,27 @@ class TestTUIAgentCommand:
         # Verify that handle_command was called with the correct argument
         mock_handle_command.assert_called_once_with(message)
 
-        # Verify that get_code_generation_agent was called to refresh the agent instance
-        mock_get_agent.assert_called()
+        # Verify that agent manager's get_agent was called to refresh the agent instance
+        mock_manager.get_agent.assert_called()
 
-    @patch("code_puppy.tui.app.get_code_generation_agent")
-    def test_tui_refreshes_agent_after_command(self, mock_get_agent):
+    @patch("code_puppy.tui.app.get_runtime_agent_manager")
+    def test_tui_refreshes_agent_after_command(self, mock_get_manager):
         """Test that TUI refreshes its agent instance after processing /agent command."""
         # Create a TUI app instance
         app = CodePuppyTUI()
 
-        # Set initial agent
+        # Mock the agent manager
+        mock_manager = MagicMock()
         initial_agent = MagicMock()
-        app.agent = initial_agent
-
-        # Mock get_code_generation_agent to return a new agent instance
         new_agent = MagicMock()
-        mock_get_agent.return_value = new_agent
+        
+        # Set initial agent
+        app.agent = initial_agent
+        app.agent_manager = mock_manager
+        
+        # Mock manager to return a new agent instance
+        mock_manager.get_agent.return_value = new_agent
+        mock_get_manager.return_value = mock_manager
 
         # Simulate that an /agent command was processed
         with patch("code_puppy.tui.app.handle_command"):
@@ -68,5 +75,5 @@ class TestTUIAgentCommand:
             loop = asyncio.get_event_loop()
             loop.run_until_complete(app.process_message("/agent code-puppy"))
 
-        # Verify that the agent was refreshed
-        mock_get_agent.assert_called()
+        # Verify that the agent was refreshed through the manager
+        mock_manager.get_agent.assert_called()