code-puppy 0.0.374__py3-none-any.whl → 0.0.376__py3-none-any.whl

code_puppy/agents/agent_manager.py

@@ -13,7 +13,7 @@ from pydantic_ai.messages import ModelMessage
 
 from code_puppy.agents.base_agent import BaseAgent
 from code_puppy.agents.json_agent import JSONAgent, discover_json_agents
-from code_puppy.callbacks import on_agent_reload
+from code_puppy.callbacks import on_agent_reload, on_register_agents
 from code_puppy.messaging import emit_success, emit_warning
 
 # Registry of available agents (Python classes and JSON file paths)
@@ -289,6 +289,38 @@ def _discover_agents(message_group_id: Optional[str] = None):
             message_group=message_group_id,
         )
 
+    # 3. Discover agents registered by plugins
+    try:
+        results = on_register_agents()
+        for result in results:
+            if result is None:
+                continue
+            # Each result should be a list of agent definitions
+            agents_list = result if isinstance(result, list) else [result]
+            for agent_def in agents_list:
+                if not isinstance(agent_def, dict) or "name" not in agent_def:
+                    continue
+
+                agent_name = agent_def["name"]
+
+                # Support both class-based and JSON path-based registration
+                if "class" in agent_def:
+                    agent_class = agent_def["class"]
+                    if isinstance(agent_class, type) and issubclass(
+                        agent_class, BaseAgent
+                    ):
+                        _AGENT_REGISTRY[agent_name] = agent_class
+                elif "json_path" in agent_def:
+                    json_path = agent_def["json_path"]
+                    if isinstance(json_path, str):
+                        _AGENT_REGISTRY[agent_name] = json_path
+
+    except Exception as e:
+        emit_warning(
+            f"Warning: Could not load plugin agents: {e}",
+            message_group=message_group_id,
+        )
+
 
 def get_available_agents() -> Dict[str, str]:
     """Get a dictionary of available agents with their display names.
@@ -612,7 +644,7 @@ def clone_agent(agent_name: str) -> Optional[str]:
                     agent_instance.display_name, clone_index
                 ),
                 "description": agent_instance.description,
-                "system_prompt": agent_instance.get_system_prompt(),
+                "system_prompt": agent_instance.get_full_system_prompt(),
                 "tools": _filter_available_tools(agent_instance.get_available_tools()),
             }
 

code_puppy/agents/base_agent.py

@@ -47,6 +47,10 @@ from pydantic_ai.messages import (
 from rich.text import Text
 
 from code_puppy.agents.event_stream_handler import event_stream_handler
+from code_puppy.callbacks import (
+    on_agent_run_end,
+    on_agent_run_start,
+)
 
 # Consolidated relative imports
 from code_puppy.config import (
@@ -101,6 +105,37 @@ class BaseAgent(ABC):
         # This is populated after the first successful run when MCP tools are retrieved
         self._mcp_tool_definitions_cache: List[Dict[str, Any]] = []
 
+    def get_identity(self) -> str:
+        """Get a unique identity for this agent instance.
+
+        Returns:
+            A string like 'python-programmer-a3f2b1' combining name + short UUID.
+        """
+        return f"{self.name}-{self.id[:6]}"
+
+    def get_identity_prompt(self) -> str:
+        """Get the identity prompt suffix to embed in system prompts.
+
+        Returns:
+            A string instructing the agent about its identity for task ownership.
+        """
+        return (
+            f"\n\nYour ID is `{self.get_identity()}`. "
+            "Use this for any tasks which require identifying yourself "
+            "such as claiming task ownership or coordination with other agents."
+        )
+
+    def get_full_system_prompt(self) -> str:
+        """Get the complete system prompt with identity automatically appended.
+
+        This wraps get_system_prompt() and appends the agent's identity,
+        so subclasses don't need to worry about it.
+
+        Returns:
+            The full system prompt including identity information.
+        """
+        return self.get_system_prompt() + self.get_identity_prompt()
+
     @property
     @abstractmethod
     def name(self) -> str:
@@ -372,35 +407,27 @@ class BaseAgent(ABC):
         total_tokens = 0
 
         # 1. Estimate tokens for system prompt / instructions
-        # For Claude Code models, the full system prompt is prepended to the first
-        # user message (already in message history), so we only count the short
-        # fixed instructions. For other models, count the full system prompt.
+        # Use prepare_prompt_for_model() to get the correct instructions for token counting.
+        # For models that prepend system prompt to user message (claude-code, antigravity),
+        # this returns the short fixed instructions. For other models, returns full prompt.
         try:
-            from code_puppy.model_utils import (
-                get_antigravity_instructions,
-                get_claude_code_instructions,
-                is_antigravity_model,
-                is_claude_code_model,
-            )
+            from code_puppy.model_utils import prepare_prompt_for_model
 
             model_name = (
                 self.get_model_name() if hasattr(self, "get_model_name") else ""
             )
-            if is_claude_code_model(model_name):
-                # For Claude Code models, only count the short fixed instructions
-                # The full system prompt is already in the message history
-                instructions = get_claude_code_instructions()
-                total_tokens += self.estimate_token_count(instructions)
-            elif is_antigravity_model(model_name):
-                # For Antigravity models, only count the short fixed instructions
-                # The full system prompt is already in the message history
-                instructions = get_antigravity_instructions()
-                total_tokens += self.estimate_token_count(instructions)
-            else:
-                # For other models, count the full system prompt
-                system_prompt = self.get_system_prompt()
-                if system_prompt:
-                    total_tokens += self.estimate_token_count(system_prompt)
+            system_prompt = self.get_full_system_prompt()
+
+            # Get the instructions that will be used (handles model-specific logic via hooks)
+            prepared = prepare_prompt_for_model(
+                model_name=model_name,
+                system_prompt=system_prompt,
+                user_prompt="",  # Empty - we just need the instructions
+                prepend_system_to_user=False,  # Don't modify prompt, just get instructions
+            )
+
+            if prepared.instructions:
+                total_tokens += self.estimate_token_count(prepared.instructions)
         except Exception:
             pass  # If we can't get system prompt, skip it
 
@@ -1122,7 +1149,7 @@ class BaseAgent(ABC):
             message_group,
         )
 
-        instructions = self.get_system_prompt()
+        instructions = self.get_full_system_prompt()
         puppy_rules = self.load_puppy_rules()
         if puppy_rules:
             instructions += f"\n{puppy_rules}"
@@ -1286,7 +1313,7 @@ class BaseAgent(ABC):
             model_name, models_config, str(uuid.uuid4())
         )
 
-        instructions = self.get_system_prompt()
+        instructions = self.get_full_system_prompt()
         puppy_rules = self.load_puppy_rules()
         if puppy_rules:
             instructions += f"\n{puppy_rules}"
@@ -1558,21 +1585,25 @@ class BaseAgent(ABC):
         if output_type is not None:
             pydantic_agent = self._create_agent_with_output_type(output_type)
 
-        # Handle claude-code, chatgpt-codex, and antigravity models: prepend system prompt to first user message
-        from code_puppy.model_utils import (
-            is_antigravity_model,
-            is_claude_code_model,
-        )
+        # Handle model-specific prompt transformations via prepare_prompt_for_model()
+        # This uses the get_model_system_prompt hook, so plugins can register their own handlers
+        from code_puppy.model_utils import prepare_prompt_for_model
 
-        if is_claude_code_model(self.get_model_name()) or is_antigravity_model(
-            self.get_model_name()
-        ):
-            if len(self.get_message_history()) == 0:
-                system_prompt = self.get_system_prompt()
-                puppy_rules = self.load_puppy_rules()
-                if puppy_rules:
-                    system_prompt += f"\n{puppy_rules}"
-                prompt = system_prompt + "\n\n" + prompt
+        # Only prepend system prompt on first message (empty history)
+        should_prepend = len(self.get_message_history()) == 0
+        if should_prepend:
+            system_prompt = self.get_full_system_prompt()
+            puppy_rules = self.load_puppy_rules()
+            if puppy_rules:
+                system_prompt += f"\n{puppy_rules}"
+
+            prepared = prepare_prompt_for_model(
+                model_name=self.get_model_name(),
+                system_prompt=system_prompt,
+                user_prompt=prompt,
+                prepend_system_to_user=True,
+            )
+            prompt = prepared.user_prompt
 
         # Build combined prompt payload when attachments are provided.
         attachment_parts: List[Any] = []
@@ -1719,6 +1750,17 @@ class BaseAgent(ABC):
         # Create the task FIRST
         agent_task = asyncio.create_task(run_agent_task())
 
+        # Fire agent_run_start hook - plugins can use this to start background tasks
+        # (e.g., token refresh heartbeats for OAuth models)
+        try:
+            await on_agent_run_start(
+                agent_name=self.name,
+                model_name=self.get_model_name(),
+                session_id=group_id,
+            )
+        except Exception:
+            pass  # Don't fail agent run if hook fails
+
         # Import shell process status helper
 
         loop = asyncio.get_running_loop()
@@ -1800,14 +1842,53 @@ class BaseAgent(ABC):
                 except Exception:
                     pass  # Don't fail the run if cache update fails
 
+            # Extract response text for the callback
+            _run_response_text = ""
+            if result is not None:
+                if hasattr(result, "data"):
+                    _run_response_text = str(result.data) if result.data else ""
+                elif hasattr(result, "output"):
+                    _run_response_text = str(result.output) if result.output else ""
+                else:
+                    _run_response_text = str(result)
+
+            _run_success = True
+            _run_error = None
             return result
         except asyncio.CancelledError:
+            _run_success = False
+            _run_error = None  # Cancellation is not an error
+            _run_response_text = ""
             agent_task.cancel()
         except KeyboardInterrupt:
-            # Handle direct keyboard interrupt during await
+            _run_success = False
+            _run_error = None  # User interrupt is not an error
+            _run_response_text = ""
             if not agent_task.done():
                 agent_task.cancel()
+        except Exception as e:
+            _run_success = False
+            _run_error = e
+            _run_response_text = ""
+            raise
         finally:
+            # Fire agent_run_end hook - plugins can use this for:
+            # - Stopping background tasks (token refresh heartbeats)
+            # - Workflow orchestration (Ralph's autonomous loop)
+            # - Logging/analytics
+            try:
+                await on_agent_run_end(
+                    agent_name=self.name,
+                    model_name=self.get_model_name(),
+                    session_id=group_id,
+                    success=_run_success,
+                    error=_run_error,
+                    response_text=_run_response_text,
+                    metadata={"model": self.get_model_name()},
+                )
+            except Exception:
+                pass  # Don't fail cleanup if hook fails
+
             # Stop keyboard listener if it was started
             if key_listener_stop_event is not None:
                 key_listener_stop_event.set()

code_puppy/callbacks.py

@@ -21,6 +21,12 @@ PhaseType = Literal[
     "pre_tool_call",
     "post_tool_call",
     "stream_event",
+    "register_tools",
+    "register_agents",
+    "register_model_type",
+    "get_model_system_prompt",
+    "agent_run_start",
+    "agent_run_end",
 ]
 CallbackFunc = Callable[..., Any]
 
@@ -42,6 +48,12 @@ _callbacks: Dict[PhaseType, List[CallbackFunc]] = {
     "pre_tool_call": [],
     "post_tool_call": [],
     "stream_event": [],
+    "register_tools": [],
+    "register_agents": [],
+    "register_model_type": [],
+    "get_model_system_prompt": [],
+    "agent_run_start": [],
+    "agent_run_end": [],
 }
 
 logger = logging.getLogger(__name__)
@@ -344,3 +356,164 @@ async def on_stream_event(
     return await _trigger_callbacks(
         "stream_event", event_type, event_data, agent_session_id
     )
+
+
+def on_register_tools() -> List[Dict[str, Any]]:
+    """Collect custom tool registrations from plugins.
+
+    Each callback should return a list of dicts with:
+    - "name": str - the tool name
+    - "register_func": callable - function that takes an agent and registers the tool
+
+    Example return: [{"name": "my_tool", "register_func": register_my_tool}]
+    """
+    return _trigger_callbacks_sync("register_tools")
+
+
+def on_register_agents() -> List[Dict[str, Any]]:
+    """Collect custom agent registrations from plugins.
+
+    Each callback should return a list of dicts with either:
+    - "name": str, "class": Type[BaseAgent] - for Python agent classes
+    - "name": str, "json_path": str - for JSON agent files
+
+    Example return: [{"name": "my-agent", "class": MyAgentClass}]
+    """
+    return _trigger_callbacks_sync("register_agents")
+
+
+def on_register_model_types() -> List[Dict[str, Any]]:
+    """Collect custom model type registrations from plugins.
+
+    This hook allows plugins to register custom model types that can be used
+    in model configurations. Each callback should return a list of dicts with:
+    - "type": str - the model type name (e.g., "antigravity", "claude_code")
+    - "handler": callable - function(model_name, model_config, config) -> model instance
+
+    The handler function receives:
+    - model_name: str - the name of the model being created
+    - model_config: dict - the model's configuration from models.json
+    - config: dict - the full models configuration
+
+    The handler should return a model instance or None if creation fails.
+
+    Example callback:
+        def register_my_model_types():
+            return [{
+                "type": "my_custom_type",
+                "handler": create_my_custom_model,
+            }]
+
+    Example return: [{"type": "antigravity", "handler": create_antigravity_model}]
+    """
+    return _trigger_callbacks_sync("register_model_type")
+
+
+def on_get_model_system_prompt(
+    model_name: str, default_system_prompt: str, user_prompt: str
+) -> List[Dict[str, Any]]:
+    """Allow plugins to provide custom system prompts for specific model types.
+
+    This hook allows plugins to override the system prompt handling for custom
+    model types (like claude_code or antigravity models). Each callback receives
+    the model name and should return a dict if it handles that model type, or None.
+
+    Args:
+        model_name: The name of the model being used (e.g., "claude-code-sonnet")
+        default_system_prompt: The default system prompt from the agent
+        user_prompt: The user's prompt/message
+
+    Each callback should return a dict with:
+    - "instructions": str - the system prompt/instructions to use
+    - "user_prompt": str - the (possibly modified) user prompt
+    - "handled": bool - True if this callback handled the model
+
+    Or return None if the callback doesn't handle this model type.
+
+    Example callback:
+        def get_my_model_system_prompt(model_name, default_system_prompt, user_prompt):
+            if model_name.startswith("my-custom-"):
+                return {
+                    "instructions": "You are MyCustomBot.",
+                    "user_prompt": f"{default_system_prompt}\n\n{user_prompt}",
+                    "handled": True,
+                }
+            return None  # Not handled by this callback
+
+    Returns:
+        List of results from registered callbacks (dicts or None values).
+    """
+    return _trigger_callbacks_sync(
+        "get_model_system_prompt", model_name, default_system_prompt, user_prompt
+    )
+
+
+async def on_agent_run_start(
+    agent_name: str,
+    model_name: str,
+    session_id: str | None = None,
+) -> List[Any]:
+    """Trigger callbacks when an agent run starts.
+
+    This fires at the beginning of run_with_mcp, before the agent task is created.
+    Useful for:
+    - Starting background tasks (like token refresh heartbeats)
+    - Logging/analytics
+    - Resource allocation
+
+    Args:
+        agent_name: Name of the agent starting
+        model_name: Name of the model being used
+        session_id: Optional session identifier
+
+    Returns:
+        List of results from registered callbacks.
+    """
+    return await _trigger_callbacks(
+        "agent_run_start", agent_name, model_name, session_id
+    )
+
+
+async def on_agent_run_end(
+    agent_name: str,
+    model_name: str,
+    session_id: str | None = None,
+    success: bool = True,
+    error: Exception | None = None,
+    response_text: str | None = None,
+    metadata: dict | None = None,
+) -> List[Any]:
+    """Trigger callbacks when an agent run ends.
+
+    This fires at the end of run_with_mcp, in the finally block.
+    Always fires regardless of success/failure/cancellation.
+
+    Useful for:
+    - Stopping background tasks (like token refresh heartbeats)
+    - Workflow orchestration (like Ralph's autonomous loop)
+    - Logging/analytics
+    - Resource cleanup
+    - Detecting completion signals in responses
+
+    Args:
+        agent_name: Name of the agent that finished
+        model_name: Name of the model that was used
+        session_id: Optional session identifier
+        success: Whether the run completed successfully
+        error: Exception if the run failed, None otherwise
+        response_text: The final text response from the agent (if successful)
+        metadata: Optional dict with additional context (tokens used, etc.)
+
+    Returns:
+        List of results from registered callbacks.
+    """
+    return await _trigger_callbacks(
+        "agent_run_end",
+        agent_name,
+        model_name,
+        session_id,
+        success,
+        error,
+        response_text,
+        metadata,
+    )

code_puppy/messaging/rich_renderer.py

@@ -675,15 +675,21 @@ class RichConsoleRenderer:
             self._console.print(f"[dim]⏱ Timeout: {msg.timeout}s[/dim]")
 
     def _render_shell_line(self, msg: ShellLineMessage) -> None:
-        """Render shell output line preserving ANSI codes."""
-        from rich.text import Text
+        """Render shell output line preserving ANSI codes and carriage returns."""
+        import sys
 
-        # Use Text.from_ansi() to parse ANSI codes into Rich styling
-        # This preserves colors while still being safe
-        text = Text.from_ansi(msg.line)
+        from rich.text import Text
 
-        # Make all shell output dim to reduce visual noise
-        self._console.print(text, style="dim")
+        # Check if line contains carriage return (progress bar style output)
+        if "\r" in msg.line:
+            # Bypass Rich entirely - write directly to stdout so terminal interprets \r
+            # Apply dim styling manually via ANSI codes
+            sys.stdout.write(f"\033[2m{msg.line}\033[0m")
+            sys.stdout.flush()
+        else:
+            # Normal line: use Rich for nice formatting
+            text = Text.from_ansi(msg.line)
+            self._console.print(text, style="dim")
 
     def _render_shell_output(self, msg: ShellOutputMessage) -> None:
         """Render shell command output - just a trailing newline for spinner separation.