code-puppy 0.0.354__py3-none-any.whl → 0.0.356__py3-none-any.whl

code_puppy/plugins/frontend_emitter/register_callbacks.py

@@ -0,0 +1,261 @@
+"""Callback registration for frontend event emission.
+
+This module registers callbacks for various agent events and emits them
+to subscribed WebSocket handlers via the emitter module.
+"""
+
+import logging
+import time
+from typing import Any, Dict, Optional
+
+from code_puppy.callbacks import register_callback
+from code_puppy.plugins.frontend_emitter.emitter import emit_event
+
+logger = logging.getLogger(__name__)
+
+
+async def on_pre_tool_call(
+    tool_name: str, tool_args: Dict[str, Any], context: Any = None
+) -> None:
+    """Emit an event when a tool call starts.
+
+    Args:
+        tool_name: Name of the tool being called
+        tool_args: Arguments being passed to the tool
+        context: Optional context data for the tool call
+    """
+    try:
+        emit_event(
+            "tool_call_start",
+            {
+                "tool_name": tool_name,
+                "tool_args": _sanitize_args(tool_args),
+                "start_time": time.time(),
+            },
+        )
+        logger.debug(f"Emitted tool_call_start for {tool_name}")
+    except Exception as e:
+        logger.error(f"Failed to emit pre_tool_call event: {e}")
+
+
+async def on_post_tool_call(
+    tool_name: str,
+    tool_args: Dict[str, Any],
+    result: Any,
+    duration_ms: float,
+    context: Any = None,
+) -> None:
+    """Emit an event when a tool call completes.
+
+    Args:
+        tool_name: Name of the tool that was called
+        tool_args: Arguments that were passed to the tool
+        result: The result returned by the tool
+        duration_ms: Execution time in milliseconds
+        context: Optional context data for the tool call
+    """
+    try:
+        emit_event(
+            "tool_call_complete",
+            {
+                "tool_name": tool_name,
+                "tool_args": _sanitize_args(tool_args),
+                "duration_ms": duration_ms,
+                "success": _is_successful_result(result),
+                "result_summary": _summarize_result(result),
+            },
+        )
+        logger.debug(
+            f"Emitted tool_call_complete for {tool_name} ({duration_ms:.2f}ms)"
+        )
+    except Exception as e:
+        logger.error(f"Failed to emit post_tool_call event: {e}")
+
+
+async def on_stream_event(
+    event_type: str, event_data: Any, agent_session_id: Optional[str] = None
+) -> None:
+    """Emit streaming events from the agent.
+
+    Args:
+        event_type: Type of the streaming event
+        event_data: Data associated with the event
+        agent_session_id: Optional session ID of the agent emitting the event
+    """
+    try:
+        emit_event(
+            "stream_event",
+            {
+                "event_type": event_type,
+                "event_data": _sanitize_event_data(event_data),
+                "agent_session_id": agent_session_id,
+            },
+        )
+        logger.debug(f"Emitted stream_event: {event_type}")
+    except Exception as e:
+        logger.error(f"Failed to emit stream_event: {e}")
+
+
+async def on_invoke_agent(*args: Any, **kwargs: Any) -> None:
+    """Emit an event when an agent is invoked.
+
+    Args:
+        *args: Positional arguments from the invoke_agent callback
+        **kwargs: Keyword arguments from the invoke_agent callback
+    """
+    try:
+        # Extract relevant info from args/kwargs
+        agent_info = {
+            "agent_name": kwargs.get("agent_name") or (args[0] if args else None),
+            "session_id": kwargs.get("session_id"),
+            "prompt_preview": _truncate_string(
+                kwargs.get("prompt") or (args[1] if len(args) > 1 else None),
+                max_length=200,
+            ),
+        }
+        emit_event("agent_invoked", agent_info)
+        logger.debug(f"Emitted agent_invoked: {agent_info.get('agent_name')}")
+    except Exception as e:
+        logger.error(f"Failed to emit invoke_agent event: {e}")
+
+
+def _sanitize_args(args: Dict[str, Any]) -> Dict[str, Any]:
+    """Sanitize tool arguments for safe emission.
+
+    Truncates large values and removes potentially sensitive data.
+
+    Args:
+        args: The raw tool arguments
+
+    Returns:
+        Sanitized arguments safe for emission
+    """
+    if not isinstance(args, dict):
+        return {}
+
+    sanitized: Dict[str, Any] = {}
+    for key, value in args.items():
+        if isinstance(value, str):
+            sanitized[key] = _truncate_string(value, max_length=500)
+        elif isinstance(value, (int, float, bool, type(None))):
+            sanitized[key] = value
+        elif isinstance(value, (list, dict)):
+            # Just indicate the type and length for complex types
+            sanitized[key] = f"<{type(value).__name__}[{len(value)}]>"
+        else:
+            sanitized[key] = f"<{type(value).__name__}>"
+
+    return sanitized
+
+
+def _sanitize_event_data(data: Any) -> Any:
+    """Sanitize event data for safe emission.
+
+    Args:
+        data: The raw event data
+
+    Returns:
+        Sanitized data safe for emission
+    """
+    if data is None:
+        return None
+
+    if isinstance(data, str):
+        return _truncate_string(data, max_length=1000)
+
+    if isinstance(data, (int, float, bool)):
+        return data
+
+    if isinstance(data, dict):
+        return {k: _sanitize_event_data(v) for k, v in list(data.items())[:20]}
+
+    if isinstance(data, (list, tuple)):
+        return [_sanitize_event_data(item) for item in data[:20]]
+
+    return f"<{type(data).__name__}>"
+
+
+def _is_successful_result(result: Any) -> bool:
+    """Determine if a tool result indicates success.
+
+    Args:
+        result: The tool result
+
+    Returns:
+        True if the result appears successful
+    """
+    if result is None:
+        return True  # No result often means success
+
+    if isinstance(result, dict):
+        # Check for error indicators
+        if result.get("error"):
+            return False
+        if result.get("success") is False:
+            return False
+        return True
+
+    if isinstance(result, bool):
+        return result
+
+    return True  # Default to success
+
+
+def _summarize_result(result: Any) -> str:
+    """Create a brief summary of a tool result.
+
+    Args:
+        result: The tool result
+
+    Returns:
+        A string summary of the result
+    """
+    if result is None:
+        return "<no result>"
+
+    if isinstance(result, str):
+        return _truncate_string(result, max_length=200)
+
+    if isinstance(result, dict):
+        if "error" in result:
+            return f"Error: {_truncate_string(str(result['error']), max_length=100)}"
+        if "message" in result:
+            return _truncate_string(str(result["message"]), max_length=100)
+        return f"<dict with {len(result)} keys>"
+
+    if isinstance(result, (list, tuple)):
+        return f"<{type(result).__name__}[{len(result)}]>"
+
+    return _truncate_string(str(result), max_length=200)
+
+
+def _truncate_string(value: Any, max_length: int = 100) -> Optional[str]:
+    """Truncate a string value if it exceeds max_length.
+
+    Args:
+        value: The value to truncate (will be converted to str)
+        max_length: Maximum length before truncation
+
+    Returns:
+        Truncated string or None if value is None
+    """
+    if value is None:
+        return None
+
+    s = str(value)
+    if len(s) > max_length:
+        return s[: max_length - 3] + "..."
+    return s
+
+
+def register() -> None:
+    """Register all frontend emitter callbacks."""
+    register_callback("pre_tool_call", on_pre_tool_call)
+    register_callback("post_tool_call", on_post_tool_call)
+    register_callback("stream_event", on_stream_event)
+    register_callback("invoke_agent", on_invoke_agent)
+    logger.debug("Frontend emitter callbacks registered")
+
+
+# Auto-register callbacks when this module is imported
+register()

code_puppy/status_display.py

@@ -7,8 +7,6 @@ from rich.panel import Panel
 from rich.spinner import Spinner
 from rich.text import Text
 
-from code_puppy.messaging import emit_info
-
 # Global variable to track current token per second rate
 CURRENT_TOKEN_RATE = 0.0
 
@@ -186,6 +184,9 @@ class StatusDisplay:
 
     async def _update_display(self) -> None:
         """Update the display continuously while active using Rich Live display"""
+        # Lazy import to avoid circular dependency during module initialization
+        from code_puppy.messaging import emit_info
+
         # Add a newline to ensure we're below the blue bar
         emit_info("")
 
@@ -214,6 +215,9 @@ class StatusDisplay:
 
     def stop(self) -> None:
         """Stop the status display"""
+        # Lazy import to avoid circular dependency during module initialization
+        from code_puppy.messaging import emit_info
+
         if self.is_active:
             self.is_active = False
             if self.task:

code_puppy/tools/agent_tools.py

@@ -1,5 +1,6 @@
 # agent_tools.py
 import asyncio
+from functools import partial
 import hashlib
 import itertools
 import json
@@ -28,12 +29,13 @@ from code_puppy.messaging import (
     SubAgentResponseMessage,
     emit_error,
     emit_info,
+    emit_success,
     get_message_bus,
     get_session_context,
     set_session_context,
 )
-from code_puppy.model_utils import is_claude_code_model
 from code_puppy.tools.common import generate_group_id
+from code_puppy.tools.subagent_context import subagent_context
 
 # Set to track active subagent invocation tasks
 _active_subagent_tasks: Set[asyncio.Task] = set()
@@ -413,6 +415,9 @@ def register_invoke_agent(agent):
             session_id = f"{session_id}-{hash_suffix}"
         # else: continuing existing session, use session_id as-is
 
+        # Lazy imports to avoid circular dependency
+        from code_puppy.agents.subagent_stream_handler import subagent_stream_handler
+
         # Emit structured invocation message via MessageBus
         bus = get_message_bus()
         bus.emit(
@@ -486,9 +491,6 @@ def register_invoke_agent(agent):
                 manager = get_mcp_manager()
                 mcp_servers = manager.get_servers_for_agent()
 
-            # Import display function for non-streaming output
-            from code_puppy.tools.display import display_non_streamed_result
-
             if get_use_dbos():
                 from pydantic_ai.durable_exec.dbos import DBOSAgent
 
@@ -543,24 +545,36 @@ def register_invoke_agent(agent):
             # Pass the message_history from the session to continue the conversation
             workflow_id = None  # Track for potential cancellation
 
-            # For claude-code models, we MUST use streaming to properly handle
-            # tool name unprefixing in the HTTP transport layer
-            stream_handler = None
-            if is_claude_code_model(model_name):
-                from code_puppy.agents.event_stream_handler import event_stream_handler
-
-                stream_handler = event_stream_handler
-
-            if get_use_dbos():
-                # Generate a unique workflow ID for DBOS - ensures no collisions in back-to-back calls
-                workflow_id = _generate_dbos_workflow_id(group_id)
-
-                # Add MCP servers to the DBOS agent's toolsets
-                # (temp_agent is discarded after this invocation, so no need to restore)
-                if subagent_mcp_servers:
-                    temp_agent._toolsets = temp_agent._toolsets + subagent_mcp_servers
+            # Always use subagent_stream_handler to silence output and update console manager
+            # This ensures all sub-agent output goes through the aggregated dashboard
+            stream_handler = partial(subagent_stream_handler, session_id=session_id)
+
+            # Wrap the agent run in subagent context for tracking
+            with subagent_context(agent_name):
+                if get_use_dbos():
+                    # Generate a unique workflow ID for DBOS - ensures no collisions in back-to-back calls
+                    workflow_id = _generate_dbos_workflow_id(group_id)
+
+                    # Add MCP servers to the DBOS agent's toolsets
+                    # (temp_agent is discarded after this invocation, so no need to restore)
+                    if subagent_mcp_servers:
+                        temp_agent._toolsets = (
+                            temp_agent._toolsets + subagent_mcp_servers
+                        )
 
-                with SetWorkflowID(workflow_id):
+                    with SetWorkflowID(workflow_id):
+                        task = asyncio.create_task(
+                            temp_agent.run(
+                                prompt,
+                                message_history=message_history,
+                                usage_limits=UsageLimits(
+                                    request_limit=get_message_limit()
+                                ),
+                                event_stream_handler=stream_handler,
+                            )
+                        )
+                        _active_subagent_tasks.add(task)
+                else:
                     task = asyncio.create_task(
                         temp_agent.run(
                             prompt,
@@ -570,38 +584,18 @@ def register_invoke_agent(agent):
                         )
                     )
                     _active_subagent_tasks.add(task)
-            else:
-                task = asyncio.create_task(
-                    temp_agent.run(
-                        prompt,
-                        message_history=message_history,
-                        usage_limits=UsageLimits(request_limit=get_message_limit()),
-                        event_stream_handler=stream_handler,
-                    )
-                )
-                _active_subagent_tasks.add(task)
 
-            try:
-                result = await task
-            finally:
-                _active_subagent_tasks.discard(task)
-                if task.cancelled():
-                    if get_use_dbos() and workflow_id:
-                        DBOS.cancel_workflow(workflow_id)
+                try:
+                    result = await task
+                finally:
+                    _active_subagent_tasks.discard(task)
+                    if task.cancelled():
+                        if get_use_dbos() and workflow_id:
+                            DBOS.cancel_workflow(workflow_id)
 
             # Extract the response from the result
             response = result.output
 
-            # Display the response using non-streaming output
-            from code_puppy.agents.event_stream_handler import get_streaming_console
-
-            display_non_streamed_result(
-                content=response,
-                console=get_streaming_console(),
-                banner_text=f"\u2713 {agent_name.upper()} RESPONSE",
-                banner_name="subagent_response",
-            )
-
             # Update the session history with the new messages from this interaction
             # The result contains all_messages which includes the full conversation
             updated_history = result.all_messages()
@@ -624,13 +618,23 @@ def register_invoke_agent(agent):
                 )
             )
 
+            # Emit clean completion summary
+            emit_success(
+                f"✓ {agent_name} completed successfully", message_group=group_id
+            )
+
             return AgentInvokeOutput(
                 response=response, agent_name=agent_name, session_id=session_id
             )
 
-        except Exception:
+        except Exception as e:
+            # Emit clean failure summary
+            emit_error(f"✗ {agent_name} failed: {str(e)}", message_group=group_id)
+
+            # Full traceback for debugging
             error_msg = f"Error invoking agent '{agent_name}': {traceback.format_exc()}"
             emit_error(error_msg, message_group=group_id)
+
             return AgentInvokeOutput(
                 response=None,
                 agent_name=agent_name,