code-puppy 0.0.353__py3-none-any.whl → 0.0.355__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,13 +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_factory import ModelFactory, make_model_settings
-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()
@@ -414,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(
@@ -431,6 +435,9 @@ def register_invoke_agent(agent):
         set_session_context(session_id)
 
         try:
+            # Lazy import to break circular dependency with messaging module
+            from code_puppy.model_factory import ModelFactory, make_model_settings
+
             # Load the specified agent config
             agent_config = load_agent(agent_name)
 
@@ -484,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
 
@@ -541,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,
@@ -568,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()
@@ -622,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,

code_puppy/tools/browser/vqa_agent.py

@@ -8,7 +8,6 @@ from pydantic import BaseModel, Field
 from pydantic_ai import Agent, BinaryContent
 
 from code_puppy.config import get_use_dbos, get_vqa_model_name
-from code_puppy.model_factory import ModelFactory
 
 
 class VisualAnalysisResult(BaseModel):
@@ -31,6 +30,7 @@ def _get_vqa_instructions() -> str:
 @lru_cache(maxsize=1)
 def _load_vqa_agent(model_name: str) -> Agent[None, VisualAnalysisResult]:
     """Create a cached agent instance for visual analysis."""
+    from code_puppy.model_factory import ModelFactory
     from code_puppy.model_utils import prepare_prompt_for_model
 
     models_config = ModelFactory.load_config()

code_puppy/tools/common.py

@@ -1,10 +1,13 @@
+import asyncio
 import fnmatch
+import functools
 import hashlib
+import logging
 import os
 import sys
 import time
 from pathlib import Path
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple
 
 from prompt_toolkit import Application
 from prompt_toolkit.formatted_text import HTML
@@ -1406,3 +1409,175 @@ def generate_group_id(tool_name: str, extra_context: str = "") -> str:
     short_hash = hash_obj.hexdigest()[:8]
 
     return f"{tool_name}_{short_hash}"
+
+
+# =============================================================================
+# TOOL CALLBACK WRAPPER
+# =============================================================================
+
+logger = logging.getLogger(__name__)
+
+
+def with_tool_callbacks(tool_name: str) -> Callable:
+    """Decorator that wraps tool functions with pre/post callback hooks.
+
+    This decorator enables plugins to hook into tool execution for:
+    - Logging and analytics
+    - Pre-execution validation or modification
+    - Post-execution result processing
+    - Performance monitoring
+
+    Args:
+        tool_name: The name of the tool being wrapped (e.g., 'edit_file', 'list_files')
+
+    Returns:
+        A decorator function that wraps the tool with callbacks.
+
+    Example:
+        @with_tool_callbacks('my_tool')
+        async def my_tool_impl(ctx, **kwargs):
+            return result
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs) -> Any:
+            # Extract context from args if available (usually first arg is RunContext)
+            context = None
+            tool_args = kwargs.copy()
+
+            # Try to get session context
+            try:
+                from code_puppy.messaging import get_session_context
+
+                context = get_session_context()
+            except ImportError:
+                pass
+
+            # Fire pre-tool callback (non-blocking)
+            try:
+                from code_puppy import callbacks
+
+                asyncio.create_task(
+                    callbacks.on_pre_tool_call(tool_name, tool_args, context)
+                )
+            except ImportError:
+                logger.debug("callbacks module not available for pre_tool_call")
+            except Exception as e:
+                logger.debug(f"Error in pre_tool_call callback: {e}")
+
+            # Execute the tool and measure duration
+            start_time = time.perf_counter()
+            result = None
+            error = None
+
+            try:
+                result = await func(*args, **kwargs)
+                return result
+            except Exception as e:
+                error = e
+                raise
+            finally:
+                end_time = time.perf_counter()
+                duration_ms = (end_time - start_time) * 1000
+
+                # Fire post-tool callback (non-blocking)
+                final_result = result if error is None else {"error": str(error)}
+                try:
+                    from code_puppy import callbacks
+
+                    asyncio.create_task(
+                        callbacks.on_post_tool_call(
+                            tool_name, tool_args, final_result, duration_ms, context
+                        )
+                    )
+                except ImportError:
+                    logger.debug("callbacks module not available for post_tool_call")
+                except Exception as e:
+                    logger.debug(f"Error in post_tool_call callback: {e}")
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs) -> Any:
+            """Sync wrapper for non-async tool functions."""
+            # Extract context
+            context = None
+            tool_args = kwargs.copy()
+
+            try:
+                from code_puppy.messaging import get_session_context
+
+                context = get_session_context()
+            except ImportError:
+                pass
+
+            # For sync functions, we can't use asyncio.create_task directly
+            # Instead, we'll try to schedule it if there's a running loop
+            def fire_pre_callback():
+                try:
+                    from code_puppy import callbacks
+
+                    loop = asyncio.get_running_loop()
+                    asyncio.run_coroutine_threadsafe(
+                        callbacks.on_pre_tool_call(tool_name, tool_args, context),
+                        loop,
+                    )
+                except RuntimeError:
+                    # No running loop - skip async callback
+                    pass
+                except ImportError:
+                    pass
+                except Exception as e:
+                    logger.debug(f"Error in sync pre_tool_call: {e}")
+
+            fire_pre_callback()
+
+            # Execute the tool
+            start_time = time.perf_counter()
+            result = None
+            error = None
+
+            try:
+                result = func(*args, **kwargs)
+                return result
+            except Exception as e:
+                error = e
+                raise
+            finally:
+                end_time = time.perf_counter()
+                duration_ms = (end_time - start_time) * 1000
+
+                # Fire post-tool callback
+                final_result = result if error is None else {"error": str(error)}
+
+                def fire_post_callback():
+                    try:
+                        from code_puppy import callbacks
+
+                        loop = asyncio.get_running_loop()
+                        asyncio.run_coroutine_threadsafe(
+                            callbacks.on_post_tool_call(
+                                tool_name,
+                                tool_args,
+                                final_result,
+                                duration_ms,
+                                context,
+                            ),
+                            loop,
+                        )
+                    except RuntimeError:
+                        # No running loop - skip async callback
+                        pass
+                    except ImportError:
+                        pass
+                    except Exception as e:
+                        logger.debug(f"Error in sync post_tool_call: {e}")
+
+                fire_post_callback()
+
+        # Return appropriate wrapper based on function type
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return decorator

code_puppy/tools/display.py

@@ -8,7 +8,8 @@ from typing import Optional
 
 from rich.console import Console
 
-from code_puppy.config import get_banner_color
+from code_puppy.config import get_banner_color, get_subagent_verbose
+from code_puppy.tools.subagent_context import is_subagent
 
 
 def display_non_streamed_result(
@@ -33,6 +34,10 @@ def display_non_streamed_result(
         >>> display_non_streamed_result("# Hello\n\nThis is **bold** text.")
         # Renders with AGENT RESPONSE banner and formatted markdown
     """
+    # Skip display for sub-agents unless verbose mode
+    if is_subagent() and not get_subagent_verbose():
+        return
+
     import time
 
     from rich.text import Text