code-muse 0.0.1__py3-none-any.whl

code_muse/agents/event_stream_handler.py

@@ -0,0 +1,361 @@
+"""Event stream handler for processing streaming events from agent runs."""
+
+import asyncio
+import logging
+import math
+from collections.abc import AsyncIterable
+from typing import Any
+
+from pydantic_ai import PartDeltaEvent, PartEndEvent, PartStartEvent, RunContext
+from pydantic_ai.messages import (
+    TextPart,
+    TextPartDelta,
+    ThinkingPart,
+    ThinkingPartDelta,
+    ToolCallPart,
+    ToolCallPartDelta,
+)
+from rich.console import Console
+from rich.markup import escape
+from rich.text import Text
+
+from code_muse.config import get_banner_color, get_subagent_verbose
+from code_muse.messaging.spinner import pause_all_spinners, resume_all_spinners
+from code_muse.tools.subagent_context import is_subagent
+
+logger = logging.getLogger(__name__)
+
+
+def _handle_stream_task_exception(task: asyncio.Task) -> None:
+    """Consume any exception from a fire-and-forget stream-event task."""
+    try:
+        task.result()
+    except Exception:
+        logger.debug("Stream event callback failed", exc_info=True)
+
+
+def _fire_stream_event(event_type: str, event_data: Any) -> None:
+    """Fire a stream event callback asynchronously (non-blocking).
+
+    PERF-08: Only fires for structural events (part_start, part_end).
+    High-frequency part_delta events are skipped since no current plugin
+    consumes them, and creating 100+ micro-tasks per response is wasteful.
+    """
+    if event_type == "part_delta":
+        return
+    try:
+        from code_muse import callbacks
+        from code_muse.messaging import get_session_context
+
+        agent_session_id = get_session_context()
+        task = asyncio.create_task(
+            callbacks.on_stream_event(event_type, event_data, agent_session_id)
+        )
+        task.add_done_callback(_handle_stream_task_exception)
+    except ImportError:
+        logger.debug("callbacks or messaging module not available for stream event")
+    except Exception as e:
+        logger.debug("Error firing stream event callback: %s", e)
+
+
+# Module-level console for streaming output
+# Set via set_streaming_console() to share console with spinner
+_streaming_console: Console | None = None
+
+
+def set_streaming_console(console: Console | None) -> None:
+    """Set the console used for streaming output.
+
+    This should be called with the same console used by the spinner
+    to avoid Live display conflicts that cause line duplication.
+
+    Args:
+        console: The Rich console to use, or None to use a fallback.
+    """
+    global _streaming_console
+    _streaming_console = console
+
+
+def get_streaming_console() -> Console:
+    """Get the console for streaming output.
+
+    Returns the configured console or creates a fallback Console.
+    """
+    if _streaming_console is not None:
+        return _streaming_console
+    return Console()
+
+
+def _should_suppress_output() -> bool:
+    """Check if sub-agent output should be suppressed.
+
+    Returns:
+        True if we're in a sub-agent context and verbose mode is disabled.
+    """
+    return is_subagent() and not get_subagent_verbose()
+
+
+async def event_stream_handler(
+    ctx: RunContext,
+    events: AsyncIterable[Any],
+) -> None:
+    """Handle streaming events from the agent run.
+
+    This function processes streaming events and emits TextPart, ThinkingPart,
+    and ToolCallPart content with styled banners/tokens as they stream in.
+
+    Args:
+        ctx: The run context.
+        events: Async iterable of streaming events (PartStartEvent, PartDeltaEvent, etc.).
+    """
+    # If we're in a sub-agent and verbose mode is disabled, silently consume events
+    if _should_suppress_output():
+        async for _ in events:
+            pass  # Just consume events without rendering
+        return
+
+    from termflow import Parser as TermflowParser
+    from termflow import Renderer as TermflowRenderer
+    from termflow.render.style import RenderFeatures
+
+    # Use the module-level console (set via set_streaming_console)
+    console = get_streaming_console()
+
+    # Track which part indices we're currently streaming (for Text/Thinking/Tool parts)
+    streaming_parts: set[int] = set()
+    thinking_parts: set[int] = set()  # Track which parts are thinking (for dim style)
+    text_parts: set[int] = set()  # Track which parts are text
+    tool_parts: set[int] = set()  # Track which parts are tool calls
+    banner_printed: set[int] = set()  # Track if banner was already printed
+    token_count: dict[int, int] = {}  # Track token count per text/tool part
+    tool_names: dict[int, str] = {}  # Track tool name per tool part index
+    did_stream_anything = False  # Track if we streamed any content
+
+    # Termflow streaming state for text parts
+    termflow_parsers: dict[int, TermflowParser] = {}
+    termflow_renderers: dict[int, TermflowRenderer] = {}
+    termflow_line_buffers: dict[int, str] = {}  # Buffer incomplete lines
+
+    async def _print_thinking_banner() -> None:
+        """Print the THINKING banner with spinner pause and line clear."""
+        nonlocal did_stream_anything
+
+        pause_all_spinners()
+        await asyncio.sleep(0.1)  # Delay to let spinner fully clear
+        # Clear line and print newline before banner
+        console.print(" " * 50, end="\r")
+        console.print()  # Newline before banner
+        # Bold banner with configurable color and lightning bolt
+        thinking_color = get_banner_color("thinking")
+        console.print(
+            Text.from_markup(
+                f"[bold white on {thinking_color}] THINKING [/bold white on {thinking_color}] [dim]\u26a1 "
+            ),
+            end="",
+        )
+        did_stream_anything = True
+
+    async def _print_response_banner() -> None:
+        """Print the AGENT RESPONSE banner with spinner pause and line clear."""
+        nonlocal did_stream_anything
+
+        pause_all_spinners()
+        await asyncio.sleep(0.1)  # Delay to let spinner fully clear
+        # Clear line and print newline before banner
+        console.print(" " * 50, end="\r")
+        console.print()  # Newline before banner
+        response_color = get_banner_color("agent_response")
+        console.print(
+            Text.from_markup(
+                f"[bold white on {response_color}] AGENT RESPONSE [/bold white on {response_color}]"
+            )
+        )
+        did_stream_anything = True
+
+    async for event in events:
+        # PartStartEvent - register the part but defer banner until content arrives
+        if isinstance(event, PartStartEvent):
+            # Fire stream event callback for part_start
+            _fire_stream_event(
+                "part_start",
+                {
+                    "index": event.index,
+                    "part_type": type(event.part).__name__,
+                    "part": event.part,
+                },
+            )
+
+            part = event.part
+            if isinstance(part, ThinkingPart):
+                streaming_parts.add(event.index)
+                thinking_parts.add(event.index)
+                # If there's initial content, print banner + content now
+                if part.content and part.content.strip():
+                    await _print_thinking_banner()
+                    escaped = escape(part.content)
+                    console.print(f"[dim]{escaped}[/dim]", end="")
+                    banner_printed.add(event.index)
+            elif isinstance(part, TextPart):
+                streaming_parts.add(event.index)
+                text_parts.add(event.index)
+                # Initialize termflow streaming for this text part
+                termflow_parsers[event.index] = TermflowParser()
+                termflow_renderers[event.index] = TermflowRenderer(
+                    output=console.file,
+                    width=console.width,
+                    features=RenderFeatures(clipboard=False),
+                )
+                termflow_line_buffers[event.index] = ""
+                # Handle initial content if present
+                if part.content and part.content.strip():
+                    await _print_response_banner()
+                    banner_printed.add(event.index)
+                    termflow_line_buffers[event.index] = part.content
+            elif isinstance(part, ToolCallPart):
+                streaming_parts.add(event.index)
+                tool_parts.add(event.index)
+                token_count[event.index] = 0  # Initialize token counter
+                # Capture tool name from the start event
+                tool_names[event.index] = part.tool_name or ""
+                # Track tool name for display
+                banner_printed.add(
+                    event.index
+                )  # Use banner_printed to track if we've shown tool info
+
+        # PartDeltaEvent - stream the content as it arrives
+        elif isinstance(event, PartDeltaEvent):
+            # Fire stream event callback for part_delta
+            _fire_stream_event(
+                "part_delta",
+                {
+                    "index": event.index,
+                    "delta_type": type(event.delta).__name__,
+                    "delta": event.delta,
+                },
+            )
+
+            if event.index in streaming_parts:
+                delta = event.delta
+                if isinstance(delta, (TextPartDelta, ThinkingPartDelta)):
+                    if delta.content_delta:
+                        # For text parts, stream markdown with termflow
+                        if event.index in text_parts:
+                            # Print banner on first content
+                            if event.index not in banner_printed:
+                                await _print_response_banner()
+                                banner_printed.add(event.index)
+
+                            # Add content to line buffer
+                            termflow_line_buffers[event.index] += delta.content_delta
+
+                            # Process complete lines
+                            parser = termflow_parsers[event.index]
+                            renderer = termflow_renderers[event.index]
+                            buffer = termflow_line_buffers[event.index]
+
+                            while "\n" in buffer:
+                                line, buffer = buffer.split("\n", 1)
+                                events_to_render = parser.parse_line(line)
+                                renderer.render_all(events_to_render)
+
+                            termflow_line_buffers[event.index] = buffer
+                        else:
+                            # For thinking parts, stream immediately (dim)
+                            if event.index not in banner_printed:
+                                await _print_thinking_banner()
+                                banner_printed.add(event.index)
+                            escaped = escape(delta.content_delta)
+                            console.print(f"[dim]{escaped}[/dim]", end="")
+                elif isinstance(delta, ToolCallPartDelta):
+                    # For tool calls, estimate tokens from args_delta content
+                    # args_delta contains the streaming JSON arguments
+                    args_delta = getattr(delta, "args_delta", "") or ""
+                    if args_delta:
+                        # Same 2.5 chars/token heuristic as BaseAgent and file_operations
+                        estimated_tokens = max(1, math.floor(len(args_delta) / 2.5))
+                        token_count[event.index] += estimated_tokens
+                    else:
+                        # Even empty deltas count as activity
+                        token_count[event.index] += 1
+
+                    # Update tool name if delta provides more of it
+                    tool_name_delta = getattr(delta, "tool_name_delta", "") or ""
+                    if tool_name_delta:
+                        tool_names[event.index] = (
+                            tool_names.get(event.index, "") + tool_name_delta
+                        )
+
+                    # Use stored tool name for display
+                    tool_name = tool_names.get(event.index, "")
+                    count = token_count[event.index]
+                    # Display with tool wrench icon and tool name
+                    if tool_name:
+                        console.print(
+                            f"  \U0001f527 Calling {tool_name}... {count} token(s)   ",
+                            end="\r",
+                        )
+                    else:
+                        console.print(
+                            f"  \U0001f527 Calling tool... {count} token(s)   ",
+                            end="\r",
+                        )
+
+        # PartEndEvent - finish the streaming with a newline
+        elif isinstance(event, PartEndEvent):
+            # Fire stream event callback for part_end
+            _fire_stream_event(
+                "part_end",
+                {
+                    "index": event.index,
+                    "next_part_kind": getattr(event, "next_part_kind", None),
+                },
+            )
+
+            if event.index in streaming_parts:
+                # For text parts, finalize termflow rendering
+                if event.index in text_parts:
+                    # Render any remaining buffered content
+                    if event.index in termflow_parsers:
+                        parser = termflow_parsers[event.index]
+                        renderer = termflow_renderers[event.index]
+                        remaining = termflow_line_buffers.get(event.index, "")
+
+                        # Parse and render any remaining partial line
+                        if remaining.strip():
+                            events_to_render = parser.parse_line(remaining)
+                            renderer.render_all(events_to_render)
+
+                        # Finalize the parser to close any open blocks
+                        final_events = parser.finalize()
+                        renderer.render_all(final_events)
+
+                        # Clean up termflow state
+                        del termflow_parsers[event.index]
+                        del termflow_renderers[event.index]
+                        del termflow_line_buffers[event.index]
+                # For tool parts, clear the chunk counter line
+                elif event.index in tool_parts:
+                    # Clear the chunk counter line by printing spaces and returning
+                    console.print(" " * 50, end="\r")
+                # For thinking parts, just print newline
+                elif event.index in banner_printed:
+                    console.print()  # Final newline after streaming
+
+                # Clean up token count and tool names
+                token_count.pop(event.index, None)
+                tool_names.pop(event.index, None)
+                # Clean up all tracking sets
+                streaming_parts.discard(event.index)
+                thinking_parts.discard(event.index)
+                text_parts.discard(event.index)
+                tool_parts.discard(event.index)
+                banner_printed.discard(event.index)
+
+                # Resume spinner if next part is NOT text/thinking/tool (avoid race condition)
+                # If next part is None or handled differently, it's safe to resume
+                # Note: spinner itself handles blank line before appearing
+                next_kind = getattr(event, "next_part_kind", None)
+                if next_kind not in ("text", "thinking", "tool-call"):
+                    resume_all_spinners()
+
+    # Spinner is resumed in PartEndEvent when appropriate (based on next_part_kind)

code_muse/agents/json_agent.py

@@ -0,0 +1,201 @@
+"""JSON-based agent configuration system."""
+
+import json
+import logging
+from pathlib import Path
+
+from .base_agent import BaseAgent
+
+logger = logging.getLogger(__name__)
+
+
+class JSONAgent(BaseAgent):
+    """Agent configured from a JSON file."""
+
+    def __init__(self, json_path: str):
+        """Initialize agent from JSON file.
+
+        Args:
+            json_path: Path to the JSON configuration file.
+        """
+        super().__init__()
+        self.json_path = json_path
+        self._config = self._load_config()
+        self._validate_config()
+
+    def _load_config(self) -> dict:
+        """Load configuration from JSON file."""
+        try:
+            with open(self.json_path, encoding="utf-8") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, FileNotFoundError) as e:
+            raise ValueError(
+                f"Failed to load JSON agent config from {self.json_path}: {e}"
+            ) from e
+
+    def _validate_config(self) -> None:
+        """Validate required fields in configuration."""
+        required_fields = ["name", "description", "system_prompt", "tools"]
+        for field in required_fields:
+            if field not in self._config:
+                raise ValueError(
+                    f"Missing required field '{field}' in JSON agent config: {self.json_path}"
+                )
+
+        # Validate tools is a list
+        if not isinstance(self._config["tools"], list):
+            raise ValueError(
+                f"'tools' must be a list in JSON agent config: {self.json_path}"
+            )
+
+        # Validate system_prompt is string or list
+        system_prompt = self._config["system_prompt"]
+        if not isinstance(system_prompt, (str, list)):
+            raise ValueError(
+                f"'system_prompt' must be a string or list in JSON agent config: {self.json_path}"
+            )
+
+    @property
+    def name(self) -> str:
+        """Get agent name from JSON config."""
+        return self._config["name"]
+
+    @property
+    def display_name(self) -> str:
+        """Get display name from JSON config, fallback to name with emoji."""
+        return self._config.get("display_name", f"{self.name.title()} 🤖")
+
+    @property
+    def description(self) -> str:
+        """Get description from JSON config."""
+        return self._config["description"]
+
+    def get_system_prompt(self) -> str:
+        """Get system prompt from JSON config."""
+        system_prompt = self._config["system_prompt"]
+
+        # If it's a list, join with newlines
+        if isinstance(system_prompt, list):
+            return "\n".join(system_prompt)
+
+        return system_prompt
+
+    def get_available_tools(self) -> list[str]:
+        """Get available tools from JSON config.
+
+        Supports both built-in tools and Universal Constructor (UC) tools.
+        UC tools are identified by checking the UC registry.
+        """
+        from code_muse.tools import get_available_tool_names
+
+        available_tools = get_available_tool_names()
+
+        # Also get UC tool names
+        uc_tool_names = set()
+        try:
+            from code_muse.plugins.universal_constructor.registry import get_registry
+
+            registry = get_registry()
+            for tool in registry.list_tools():
+                if tool.meta.enabled:
+                    uc_tool_names.add(tool.full_name)
+        except ImportError:
+            pass  # UC module not available
+        except Exception as e:
+            # Log unexpected errors but don't fail
+            import logging
+
+            logging.debug(f"UC registry access failed: {e}")
+
+        # Return tools that are either built-in OR UC tools
+        requested_tools = []
+        for tool in self._config["tools"]:
+            if tool in available_tools:
+                requested_tools.append(tool)
+            elif tool in uc_tool_names:
+                # UC tool - mark it specially so base_agent knows to handle it
+                requested_tools.append(f"uc:{tool}")
+
+        return requested_tools
+
+    def get_user_prompt(self) -> str | None:
+        """Get custom user prompt from JSON config."""
+        return self._config.get("user_prompt")
+
+    def get_tools_config(self) -> dict | None:
+        """Get tool configuration from JSON config."""
+        return self._config.get("tools_config")
+
+    def refresh_config(self) -> None:
+        """Reload the agent configuration from disk.
+
+        This keeps long-lived agent instances in sync after external edits.
+        """
+        self._config = self._load_config()
+        self._validate_config()
+
+    def get_model_name(self) -> str | None:
+        """Get pinned model name from JSON config, if specified.
+
+        Returns:
+            Model name to use for this agent, or None to use global default.
+        """
+        result = self._config.get("model")
+        if result is None:
+            result = super().get_model_name()
+        return result
+
+
+def discover_json_agents() -> dict[str, str]:
+    """Discover JSON agent files in the user's and project's agents directories.
+
+    Searches two locations:
+    1. User agents directory (~/.muse/agents/)
+    2. Project agents directory (<CWD>/.muse/agents/) - if it exists
+
+    Project agents take priority over user agents when names collide.
+
+    Returns:
+        dict mapping agent names to their JSON file paths.
+    """
+    from code_muse.config import (
+        get_project_agents_directory,
+        get_user_agents_directory,
+    )
+
+    agents: dict[str, str] = {}
+
+    # 1. Discover user-level agents first
+    user_agents_dir = Path(get_user_agents_directory())
+    if user_agents_dir.exists() and user_agents_dir.is_dir():
+        for json_file in user_agents_dir.glob("*.json"):
+            try:
+                agent = JSONAgent(str(json_file))
+                agents[agent.name] = str(json_file)
+            except Exception as e:
+                logger.debug(
+                    "Skipping invalid user agent file: %s (reason: %s: %s)",
+                    json_file,
+                    type(e).__name__,
+                    str(e),
+                )
+                continue
+
+    # 2. Discover project-level agents (overrides user agents on name collision)
+    project_agents_dir_str = get_project_agents_directory()
+    if project_agents_dir_str is not None:
+        project_agents_dir = Path(project_agents_dir_str)
+        for json_file in project_agents_dir.glob("*.json"):
+            try:
+                agent = JSONAgent(str(json_file))
+                agents[agent.name] = str(json_file)
+            except Exception as e:
+                logger.debug(
+                    "Skipping invalid project agent file: %s (reason: %s: %s)",
+                    json_file,
+                    type(e).__name__,
+                    str(e),
+                )
+                continue
+
+    return agents