code-puppy 0.0.302__py3-none-any.whl → 0.0.323__py3-none-any.whl

code_puppy/messaging/bus.py

@@ -86,6 +86,9 @@ class MessageBus:
         # Request/Response correlation: prompt_id → Future (for async usage)
         self._pending_requests: Dict[str, asyncio.Future[Any]] = {}
 
+        # Session context for multi-agent tracking
+        self._current_session_id: Optional[str] = None
+
     # =========================================================================
     # Outgoing Messages (Agent → UI)
     # =========================================================================
@@ -95,11 +98,16 @@ class MessageBus:
 
         Thread-safe. Can be called from sync or async context.
         If no renderer is active, messages are buffered for later.
+        Auto-tags message with current session_id if not already set.
 
         Args:
             message: The message to emit.
         """
+        # Auto-tag message with current session if not already set
         with self._lock:
+            if message.session_id is None and self._current_session_id is not None:
+                message.session_id = self._current_session_id
+
             if not self._has_active_renderer:
                 self._startup_buffer.append(message)
                 return
@@ -151,6 +159,43 @@ class MessageBus:
         """Emit a DEBUG level text message."""
         self.emit_text(MessageLevel.DEBUG, text)
 
+    def emit_shell_line(self, line: str, stream: str = "stdout") -> None:
+        """Emit a shell output line with ANSI preservation.
+
+        Args:
+            line: The output line (may contain ANSI codes).
+            stream: Which stream this came from ("stdout" or "stderr").
+        """
+        from .messages import ShellLineMessage
+
+        message = ShellLineMessage(line=line, stream=stream)  # type: ignore[arg-type]
+        self.emit(message)
+
+    # =========================================================================
+    # Session Context (Multi-Agent Tracking)
+    # =========================================================================
+
+    def set_session_context(self, session_id: Optional[str]) -> None:
+        """Set the current session context for auto-tagging messages.
+
+        When set, all messages emitted via emit() will be automatically tagged
+        with this session_id unless they already have one set.
+
+        Args:
+            session_id: The session ID to tag messages with, or None to clear.
+        """
+        with self._lock:
+            self._current_session_id = session_id
+
+    def get_session_context(self) -> Optional[str]:
+        """Get the current session context.
+
+        Returns:
+            The current session_id, or None if not set.
+        """
+        with self._lock:
+            return self._current_session_id
+
     # =========================================================================
     # User Input Requests (Agent waits for UI response)
     # =========================================================================
@@ -526,6 +571,21 @@ def emit_debug(text: str) -> None:
     get_message_bus().emit_debug(text)
 
 
+def emit_shell_line(line: str, stream: str = "stdout") -> None:
+    """Emit a shell output line with ANSI preservation."""
+    get_message_bus().emit_shell_line(line, stream)
+
+
+def set_session_context(session_id: Optional[str]) -> None:
+    """Set the session context on the global bus."""
+    get_message_bus().set_session_context(session_id)
+
+
+def get_session_context() -> Optional[str]:
+    """Get the session context from the global bus."""
+    return get_message_bus().get_session_context()
+
+
 # =============================================================================
 # Export all public symbols
 # =============================================================================
@@ -543,4 +603,8 @@ __all__ = [
     "emit_error",
     "emit_success",
     "emit_debug",
+    "emit_shell_line",
+    # Session context
+    "set_session_context",
+    "get_session_context",
 ]

code_puppy/messaging/markdown_patches.py

@@ -0,0 +1,57 @@
+"""Patches for Rich's Markdown rendering.
+
+This module provides customizations to Rich's default Markdown rendering,
+particularly for header justification which is hardcoded to center in Rich.
+"""
+
+from rich import box
+from rich.markdown import Heading, Markdown
+from rich.panel import Panel
+from rich.text import Text
+
+
+class LeftJustifiedHeading(Heading):
+    """A heading that left-justifies text instead of centering.
+
+    Rich's default Heading class hardcodes `text.justify = 'center'`,
+    which can look odd in a CLI context. This subclass overrides that
+    to use left justification instead.
+    """
+
+    def __rich_console__(self, console, options):
+        """Render the heading with left justification."""
+        text = self.text
+        text.justify = "left"  # Override Rich's default 'center'
+
+        if self.tag == "h1":
+            # Draw a border around h1s (same as Rich default)
+            yield Panel(
+                text,
+                box=box.HEAVY,
+                style="markdown.h1.border",
+            )
+        else:
+            # Styled text for h2 and beyond (same as Rich default)
+            if self.tag == "h2":
+                yield Text("")
+            yield text
+
+
+_patched = False
+
+
+def patch_markdown_headings():
+    """Patch Rich's Markdown to use left-justified headings.
+
+    This function is idempotent - calling it multiple times has no effect
+    after the first call.
+    """
+    global _patched
+    if _patched:
+        return
+
+    Markdown.elements["heading_open"] = LeftJustifiedHeading
+    _patched = True
+
+
+__all__ = ["patch_markdown_headings", "LeftJustifiedHeading"]

code_puppy/messaging/messages.py

@@ -56,6 +56,10 @@ class BaseMessage(BaseModel):
     category: MessageCategory = Field(
         description="Category for routing and rendering decisions"
     )
+    session_id: Optional[str] = Field(
+        default=None,
+        description="Session ID of the agent that emitted this message (for multi-agent tracking)",
+    )
 
     model_config = {"frozen": False, "extra": "forbid"}
 
@@ -207,6 +211,16 @@ class ShellStartMessage(BaseMessage):
     timeout: int = Field(default=60, description="Timeout in seconds")
 
 
+class ShellLineMessage(BaseMessage):
+    """A single line of shell command output with ANSI preservation."""
+
+    category: MessageCategory = MessageCategory.TOOL_OUTPUT
+    line: str = Field(description="The output line (may contain ANSI codes)")
+    stream: Literal["stdout", "stderr"] = Field(
+        default="stdout", description="Which output stream this line came from"
+    )
+
+
 class ShellOutputMessage(BaseMessage):
     """Output from a shell command execution with stdout, stderr, and timing."""
 
@@ -394,6 +408,7 @@ AnyMessage = Union[
     GrepResultMessage,
     DiffMessage,
     ShellStartMessage,
+    ShellLineMessage,
     ShellOutputMessage,
     AgentReasoningMessage,
     AgentResponseMessage,
@@ -433,6 +448,7 @@ __all__ = [
     "DiffMessage",
     # Shell
     "ShellStartMessage",
+    "ShellLineMessage",
     "ShellOutputMessage",
     # Agent
     "AgentReasoningMessage",

code_puppy/messaging/renderers.py

@@ -12,6 +12,7 @@ from typing import Optional
 
 from rich.console import Console
 from rich.markdown import Markdown
+from rich.markup import escape as escape_rich_markup
 
 from .message_queue import MessageQueue, MessageType, UIMessage
 
@@ -126,11 +127,15 @@ class InteractiveRenderer(MessageRenderer):
                     self.console.print(markdown)
                 except Exception:
                     # Fallback to plain text if markdown parsing fails
-                    self.console.print(message.content)
+                    safe_content = escape_rich_markup(message.content)
+                    self.console.print(safe_content)
             elif style:
-                self.console.print(message.content, style=style)
+                # Escape Rich markup to prevent crashes from malformed tags
+                safe_content = escape_rich_markup(message.content)
+                self.console.print(safe_content, style=style)
             else:
-                self.console.print(message.content)
+                safe_content = escape_rich_markup(message.content)
+                self.console.print(safe_content)
         else:
             # For complex Rich objects (Tables, Markdown, Text, etc.)
             self.console.print(message.content)
@@ -145,7 +150,8 @@ class InteractiveRenderer(MessageRenderer):
         # This renderer is not currently used in practice, but if it were:
         # We would need async input handling here
         # For now, just render as a system message
-        self.console.print(f"[bold cyan]INPUT REQUESTED:[/bold cyan] {message.content}")
+        safe_content = escape_rich_markup(str(message.content))
+        self.console.print(f"[bold cyan]INPUT REQUESTED:[/bold cyan] {safe_content}")
         if hasattr(self.console.file, "flush"):
             self.console.file.flush()
 
@@ -253,11 +259,16 @@ class SynchronousInteractiveRenderer:
                     self.console.print(markdown)
                 except Exception:
                     # Fallback to plain text if markdown parsing fails
-                    self.console.print(message.content)
+                    safe_content = escape_rich_markup(message.content)
+                    self.console.print(safe_content)
             elif style:
-                self.console.print(message.content, style=style)
+                # Escape Rich markup to prevent crashes from malformed tags
+                # in shell output or other user-provided content
+                safe_content = escape_rich_markup(message.content)
+                self.console.print(safe_content, style=style)
             else:
-                self.console.print(message.content)
+                safe_content = escape_rich_markup(message.content)
+                self.console.print(safe_content)
         else:
             # For complex Rich objects (Tables, Markdown, Text, etc.)
             self.console.print(message.content)
@@ -276,8 +287,9 @@ class SynchronousInteractiveRenderer:
             )
             return
 
-        # Display the prompt
-        self.console.print(f"[bold cyan]{message.content}[/bold cyan]")
+        # Display the prompt - escape to prevent markup injection
+        safe_content = escape_rich_markup(str(message.content))
+        self.console.print(f"[bold cyan]{safe_content}[/bold cyan]")
         if hasattr(self.console.file, "flush"):
             self.console.file.flush()
 

code_puppy/messaging/rich_renderer.py

@@ -11,6 +11,7 @@ from typing import Dict, Optional, Protocol, runtime_checkable
 
 from rich.console import Console
 from rich.markdown import Markdown
+from rich.markup import escape as escape_rich_markup
 from rich.panel import Panel
 from rich.rule import Rule
 
@@ -37,6 +38,7 @@ from .messages import (
     GrepResultMessage,
     MessageLevel,
     SelectionRequest,
+    ShellLineMessage,
     ShellOutputMessage,
     ShellStartMessage,
     SpinnerControl,
@@ -131,6 +133,32 @@ class RichConsoleRenderer:
         """Get the Rich console."""
         return self._console
 
+    def _get_banner_color(self, banner_name: str) -> str:
+        """Get the configured color for a banner.
+
+        Args:
+            banner_name: The banner identifier (e.g., 'thinking', 'shell_command')
+
+        Returns:
+            Rich color name for the banner background
+        """
+        from code_puppy.config import get_banner_color
+
+        return get_banner_color(banner_name)
+
+    def _format_banner(self, banner_name: str, text: str) -> str:
+        """Format a banner with its configured color.
+
+        Args:
+            banner_name: The banner identifier
+            text: The banner text
+
+        Returns:
+            Rich markup string for the banner
+        """
+        color = self._get_banner_color(banner_name)
+        return f"[bold white on {color}] {text} [/bold white on {color}]"
+
     # =========================================================================
     # Lifecycle (Synchronous - for compatibility with main.py)
     # =========================================================================
@@ -187,7 +215,9 @@ class RichConsoleRenderer:
             self._do_render(message)
         except Exception as e:
             # Don't let rendering errors crash the loop
-            self._console.print(f"[dim red]Render error: {e}[/dim red]")
+            # Escape the error message to prevent nested markup errors
+            safe_error = escape_rich_markup(str(e))
+            self._console.print(f"[dim red]Render error: {safe_error}[/dim red]")
 
     # =========================================================================
     # Async Lifecycle (for future async-first usage)
@@ -233,12 +263,15 @@ class RichConsoleRenderer:
             self._render_diff(message)
         elif isinstance(message, ShellStartMessage):
             self._render_shell_start(message)
+        elif isinstance(message, ShellLineMessage):
+            self._render_shell_line(message)
         elif isinstance(message, ShellOutputMessage):
             self._render_shell_output(message)
         elif isinstance(message, AgentReasoningMessage):
             self._render_agent_reasoning(message)
         elif isinstance(message, AgentResponseMessage):
-            self._render_agent_response(message)
+            # Skip rendering - we now stream agent responses via event_stream_handler
+            pass
         elif isinstance(message, SubAgentInvocationMessage):
             self._render_subagent_invocation(message)
         elif isinstance(message, SubAgentResponseMessage):
@@ -282,7 +315,12 @@ class RichConsoleRenderer:
     # =========================================================================
 
     def _render_text(self, msg: TextMessage) -> None:
-        """Render a text message with appropriate styling."""
+        """Render a text message with appropriate styling.
+
+        Text is escaped to prevent Rich markup injection which could crash
+        the renderer if malformed tags are present in shell output or other
+        user-provided content.
+        """
         style = self._styles.get(msg.level, "white")
 
         # Make version messages dim
@@ -290,7 +328,9 @@ class RichConsoleRenderer:
             style = "dim"
 
         prefix = self._get_level_prefix(msg.level)
-        self._console.print(f"{prefix}{msg.text}", style=style)
+        # Escape Rich markup to prevent crashes from malformed tags
+        safe_text = escape_rich_markup(msg.text)
+        self._console.print(f"{prefix}{safe_text}", style=style)
 
     def _get_level_prefix(self, level: MessageLevel) -> str:
         """Get a prefix icon for the message level."""
@@ -311,8 +351,9 @@ class RichConsoleRenderer:
         """Render a directory listing matching the old Rich-formatted output."""
         # Header on single line
         rec_flag = f"(recursive={msg.recursive})"
+        banner = self._format_banner("directory_listing", "DIRECTORY LISTING")
         self._console.print(
-            f"\n[bold white on blue] DIRECTORY LISTING [/bold white on blue] "
+            f"\n{banner} "
             f"📂 [bold cyan]{msg.directory}[/bold cyan] [dim]{rec_flag}[/dim]\n"
         )
 
@@ -362,26 +403,25 @@ class RichConsoleRenderer:
             line_info = f" [dim](lines {msg.start_line}-{end_line})[/dim]"
 
         # Just print the header - content is for LLM only
+        banner = self._format_banner("read_file", "READ FILE")
         self._console.print(
-            f"\n[bold white on blue] READ FILE [/bold white on blue] "
-            f"📂 [bold cyan]{msg.path}[/bold cyan]{line_info}"
+            f"\n{banner} 📂 [bold cyan]{msg.path}[/bold cyan]{line_info}"
         )
 
     def _render_grep_result(self, msg: GrepResultMessage) -> None:
         """Render grep results grouped by file matching old format."""
         import re
 
-        # Header matching old format
+        # Header
+        banner = self._format_banner("grep", "GREP")
         self._console.print(
-            f"\n[bold white on blue] GREP [/bold white on blue] "
-            f"📂 [bold cyan]{msg.directory}[/bold cyan] "
-            f"[dim]for '{msg.search_term}'[/dim]"
+            f"\n{banner} 📂 [dim]{msg.directory} for '{msg.search_term}'[/dim]"
         )
 
         if not msg.matches:
             self._console.print(
-                f"[yellow]⚠ No matches found for '{msg.search_term}' "
-                f"in {msg.directory}[/yellow]"
+                f"[dim]No matches found for '{msg.search_term}' "
+                f"in {msg.directory}[/dim]"
             )
             return
 
@@ -397,8 +437,7 @@ class RichConsoleRenderer:
                 file_matches = by_file[file_path]
                 match_word = "match" if len(file_matches) == 1 else "matches"
                 self._console.print(
-                    f"\n[bold white]📄 {file_path}[/bold white] "
-                    f"[dim]({len(file_matches)} {match_word})[/dim]"
+                    f"\n[dim]📄 {file_path} ({len(file_matches)} {match_word})[/dim]"
                 )
 
                 # Show each match with line number and content
@@ -414,7 +453,7 @@ class RichConsoleRenderer:
                     if search_term and not search_term.startswith("-"):
                         highlighted_line = re.sub(
                             f"({re.escape(search_term)})",
-                            r"[bold yellow on black]\1[/bold yellow on black]",
+                            r"[bold yellow]\1[/bold yellow]",
                             line,
                             flags=re.IGNORECASE,
                         )
@@ -422,9 +461,7 @@ class RichConsoleRenderer:
                         highlighted_line = line
 
                     ln = match.line_number
-                    self._console.print(
-                        f"  [bold cyan]{ln:4d}[/bold cyan] │ {highlighted_line}"
-                    )
+                    self._console.print(f"  [dim]{ln:4d}[/dim] │ {highlighted_line}")
         else:
             # Concise mode (default): Show only file summaries
             self._console.print("")
@@ -435,15 +472,18 @@ class RichConsoleRenderer:
                     f"[dim]📄 {file_path} ({len(file_matches)} {match_word})[/dim]"
                 )
 
-        # Summary
+        # Summary - subtle
         match_word = "match" if msg.total_matches == 1 else "matches"
         file_word = "file" if len(by_file) == 1 else "files"
         num_files = len(by_file)
         self._console.print(
-            f"[green]✓ Found [bold]{msg.total_matches}[/bold] {match_word} "
-            f"across [bold]{num_files}[/bold] {file_word}[/green]"
+            f"[dim]Found {msg.total_matches} {match_word} "
+            f"across {num_files} {file_word}[/dim]"
         )
 
+        # Trailing newline for spinner separation
+        self._console.print()
+
     # =========================================================================
     # Diff
     # =========================================================================
@@ -457,8 +497,9 @@ class RichConsoleRenderer:
         op_color = op_colors.get(msg.operation, "white")
 
         # Header on single line
+        banner = self._format_banner("edit_file", "EDIT FILE")
         self._console.print(
-            f"\n[bold white on blue] EDIT FILE [/bold white on blue] "
+            f"\n{banner} "
             f"{icon} [{op_color}]{msg.operation.upper()}[/{op_color}] "
             f"[bold cyan]{msg.path}[/bold cyan]"
         )
@@ -474,7 +515,12 @@ class RichConsoleRenderer:
             elif line.type == "remove":
                 diff_text_lines.append(f"-{line.content}")
             else:  # context
-                diff_text_lines.append(f" {line.content}")
+                # Don't add space prefix to diff headers - they need to be preserved
+                # exactly for syntax highlighting to detect the file extension
+                if line.content.startswith(("---", "+++", "@@", "diff ", "index ")):
+                    diff_text_lines.append(line.content)
+                else:
+                    diff_text_lines.append(f" {line.content}")
 
         diff_text = "\n".join(diff_text_lines)
 
@@ -488,43 +534,39 @@ class RichConsoleRenderer:
 
     def _render_shell_start(self, msg: ShellStartMessage) -> None:
         """Render shell command start notification."""
+        # Escape command to prevent Rich markup injection
+        safe_command = escape_rich_markup(msg.command)
         # Header showing command is starting
-        self._console.print(
-            f"\n[bold white on blue] SHELL COMMAND [/bold white on blue] "
-            f"🚀 [bold green]$ {msg.command}[/bold green]"
-        )
+        banner = self._format_banner("shell_command", "SHELL COMMAND")
+        self._console.print(f"\n{banner} 🚀 [dim]$ {safe_command}[/dim]")
 
         # Show working directory if specified
         if msg.cwd:
-            self._console.print(f"[dim]📂 Working directory: {msg.cwd}[/dim]")
+            safe_cwd = escape_rich_markup(msg.cwd)
+            self._console.print(f"[dim]📂 Working directory: {safe_cwd}[/dim]")
 
         # Show timeout
         self._console.print(f"[dim]⏱ Timeout: {msg.timeout}s[/dim]")
 
-    def _render_shell_output(self, msg: ShellOutputMessage) -> None:
-        """Render shell command output matching old format."""
-        # Header matching old format
-        self._console.print(
-            f"\n[bold white on blue] SHELL COMMAND [/bold white on blue] "
-            f"📂 [bold green]$ {msg.command}[/bold green]"
-        )
+    def _render_shell_line(self, msg: ShellLineMessage) -> None:
+        """Render shell output line preserving ANSI codes."""
+        from rich.text import Text
 
-        # stdout
-        if msg.stdout:
-            self._console.print(msg.stdout)
+        # Use Text.from_ansi() to parse ANSI codes into Rich styling
+        # This preserves colors while still being safe
+        text = Text.from_ansi(msg.line)
 
-        # stderr (if any)
-        if msg.stderr:
-            self._console.print(f"[red]{msg.stderr}[/red]")
+        # Make all shell output dim to reduce visual noise
+        self._console.print(text, style="dim")
 
-        # Footer with timing and exit code
-        if msg.exit_code == 0:
-            self._console.print(f"[dim]Completed in {msg.duration_seconds:.2f}s[/dim]")
-        else:
-            dur = msg.duration_seconds
-            self._console.print(
-                f"[red]Exit code: {msg.exit_code}[/red] [dim]({dur:.2f}s)[/dim]"
-            )
+    def _render_shell_output(self, msg: ShellOutputMessage) -> None:
+        """Render shell command output - just a trailing newline for spinner separation.
+
+        Shell command results are already returned to the LLM via tool responses,
+        so we don't need to clutter the UI with redundant output.
+        """
+        # Just print trailing newline for spinner separation
+        self._console.print()
 
     # =========================================================================
     # Agent Messages
@@ -533,9 +575,8 @@ class RichConsoleRenderer:
     def _render_agent_reasoning(self, msg: AgentReasoningMessage) -> None:
         """Render agent reasoning matching old format."""
         # Header matching old format
-        self._console.print(
-            "\n[bold white on purple] AGENT REASONING [/bold white on purple]"
-        )
+        banner = self._format_banner("agent_reasoning", "AGENT REASONING")
+        self._console.print(f"\n{banner}")
 
         # Current reasoning
         self._console.print("[bold cyan]Current reasoning:[/bold cyan]")
@@ -549,12 +590,14 @@ class RichConsoleRenderer:
             md_steps = Markdown(msg.next_steps)
             self._console.print(md_steps)
 
+        # Trailing newline for spinner separation
+        self._console.print()
+
     def _render_agent_response(self, msg: AgentResponseMessage) -> None:
         """Render agent response with header and markdown formatting."""
         # Header
-        self._console.print(
-            "\n[bold white on purple] AGENT RESPONSE [/bold white on purple]\n"
-        )
+        banner = self._format_banner("agent_response", "AGENT RESPONSE")
+        self._console.print(f"\n{banner}\n")
 
         # Content (markdown or plain)
         if msg.is_markdown:
@@ -571,8 +614,9 @@ class RichConsoleRenderer:
             if msg.is_new_session
             else f"Continuing ({msg.message_count} messages)"
         )
+        banner = self._format_banner("invoke_agent", "🤖 INVOKE AGENT")
         self._console.print(
-            f"\n[bold white on purple] 🤖 INVOKE AGENT [/bold white on purple] "
+            f"\n{banner} "
             f"[bold cyan]{msg.agent_name}[/bold cyan] "
             f"[dim]({session_type})[/dim]"
         )
@@ -591,10 +635,8 @@ class RichConsoleRenderer:
     def _render_subagent_response(self, msg: SubAgentResponseMessage) -> None:
         """Render sub-agent response with markdown formatting."""
         # Response header
-        self._console.print(
-            f"\n[bold white on green] ✓ AGENT RESPONSE [/bold white on green] "
-            f"[bold cyan]{msg.agent_name}[/bold cyan]"
-        )
+        banner = self._format_banner("subagent_response", "✓ AGENT RESPONSE")
+        self._console.print(f"\n{banner} [bold cyan]{msg.agent_name}[/bold cyan]")
 
         # Render response as markdown
         md = Markdown(msg.response)
@@ -633,9 +675,11 @@ class RichConsoleRenderer:
 
     async def _render_confirmation_request(self, msg: ConfirmationRequest) -> None:
         """Render confirmation dialog and send response back."""
-        # Show title and description
-        self._console.print(f"\n[bold yellow]{msg.title}[/bold yellow]")
-        self._console.print(msg.description)
+        # Show title and description - escape to prevent markup injection
+        safe_title = escape_rich_markup(msg.title)
+        safe_description = escape_rich_markup(msg.description)
+        self._console.print(f"\n[bold yellow]{safe_title}[/bold yellow]")
+        self._console.print(safe_description)
 
         # Show options
         options_str = "/".join(msg.options)
@@ -669,11 +713,13 @@ class RichConsoleRenderer:
 
     async def _render_selection_request(self, msg: SelectionRequest) -> None:
         """Render selection menu and send response back."""
-        self._console.print(f"\n[bold]{msg.prompt_text}[/bold]")
+        safe_prompt = escape_rich_markup(msg.prompt_text)
+        self._console.print(f"\n[bold]{safe_prompt}[/bold]")
 
-        # Show numbered options
+        # Show numbered options - escape to prevent markup injection
         for i, opt in enumerate(msg.options):
-            self._console.print(f"  [cyan]{i + 1}[/cyan]. {opt}")
+            safe_opt = escape_rich_markup(opt)
+            self._console.print(f"  [cyan]{i + 1}[/cyan]. {safe_opt}")
 
         if msg.allow_cancel:
             self._console.print("  [dim]0. Cancel[/dim]")