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

code_puppy/agents/base_agent.py

@@ -4,10 +4,23 @@ import asyncio
 import json
 import math
 import signal
+import sys
 import threading
 import uuid
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
+from collections.abc import AsyncIterable
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
 import mcp
 import pydantic
@@ -18,6 +31,7 @@ from pydantic_ai import (
     BinaryContent,
     DocumentUrl,
     ImageUrl,
+    PartEndEvent,
     RunContext,
     UsageLimitExceeded,
     UsageLimits,
@@ -33,6 +47,7 @@ from pydantic_ai.messages import (
     ToolReturn,
     ToolReturnPart,
 )
+from rich.text import Text
 
 # Consolidated relative imports
 from code_puppy.config import (
@@ -44,9 +59,10 @@ from code_puppy.config import (
     get_protected_token_count,
     get_use_dbos,
     get_value,
-    load_mcp_server_configs,
 )
-from code_puppy.mcp_ import ServerConfig, get_mcp_manager
+from code_puppy.error_logging import log_error
+from code_puppy.keymap import cancel_agent_uses_signal, get_cancel_agent_char_code
+from code_puppy.mcp_ import get_mcp_manager
 from code_puppy.messaging import (
     emit_error,
     emit_info,
@@ -85,6 +101,9 @@ class BaseAgent(ABC):
         # Cache for MCP tool definitions (for token estimation)
         # This is populated after the first successful run when MCP tools are retrieved
         self._mcp_tool_definitions_cache: List[Dict[str, Any]] = []
+        # Shared console for streaming output - should be set by cli_runner
+        # to avoid conflicts between spinner's Live display and response streaming
+        self._console: Optional[Any] = None
 
     @property
     @abstractmethod
@@ -362,8 +381,10 @@ class BaseAgent(ABC):
         # fixed instructions. For other models, count the full system prompt.
         try:
             from code_puppy.model_utils import (
-                is_claude_code_model,
+                get_chatgpt_codex_instructions,
                 get_claude_code_instructions,
+                is_chatgpt_codex_model,
+                is_claude_code_model,
             )
 
             model_name = (
@@ -374,6 +395,11 @@ class BaseAgent(ABC):
                 # The full system prompt is already in the message history
                 instructions = get_claude_code_instructions()
                 total_tokens += self.estimate_token_count(instructions)
+            elif is_chatgpt_codex_model(model_name):
+                # For ChatGPT Codex models, only count the short fixed instructions
+                # The full system prompt is already in the message history
+                instructions = get_chatgpt_codex_instructions()
+                total_tokens += self.estimate_token_count(instructions)
             else:
                 # For other models, count the full system prompt
                 system_prompt = self.get_system_prompt()
@@ -827,30 +853,11 @@ class BaseAgent(ABC):
         total_current_tokens = message_tokens + context_overhead
         proportion_used = total_current_tokens / model_max
 
-        # Check if we're in TUI mode and can update the status bar
-        from code_puppy.tui_state import get_tui_app_instance, is_tui_mode
-
         context_summary = SpinnerBase.format_context_info(
             total_current_tokens, model_max, proportion_used
         )
         update_spinner_context(context_summary)
 
-        if is_tui_mode():
-            tui_app = get_tui_app_instance()
-            if tui_app:
-                try:
-                    # Update the status bar instead of emitting a chat message
-                    status_bar = tui_app.query_one("StatusBar")
-                    status_bar.update_token_info(
-                        total_current_tokens, model_max, proportion_used
-                    )
-                except Exception as e:
-                    emit_error(e)
-            else:
-                emit_info(
-                    f"Final token count after processing: {total_current_tokens}",
-                    message_group="token_context_status",
-                )
         # Get the configured compaction threshold
         compaction_threshold = get_compaction_threshold()
 
@@ -889,30 +896,12 @@ class BaseAgent(ABC):
             final_token_count = sum(
                 self.estimate_tokens_for_message(msg) for msg in result_messages
             )
-            # Update status bar with final token count if in TUI mode
+            # Update spinner with final token count
             final_summary = SpinnerBase.format_context_info(
                 final_token_count, model_max, final_token_count / model_max
             )
             update_spinner_context(final_summary)
 
-            if is_tui_mode():
-                tui_app = get_tui_app_instance()
-                if tui_app:
-                    try:
-                        status_bar = tui_app.query_one("StatusBar")
-                        status_bar.update_token_info(
-                            final_token_count, model_max, final_token_count / model_max
-                        )
-                    except Exception:
-                        emit_info(
-                            f"Final token count after processing: {final_token_count}",
-                            message_group="token_context_status",
-                        )
-                else:
-                    emit_info(
-                        f"Final token count after processing: {final_token_count}",
-                        message_group="token_context_status",
-                    )
             self.set_message_history(result_messages)
             for m in summarized_messages:
                 self.add_compacted_message_hash(self.hash_message(m))
@@ -974,60 +963,71 @@ class BaseAgent(ABC):
 
     # ===== Agent wiring formerly in code_puppy/agent.py =====
     def load_puppy_rules(self) -> Optional[str]:
-        """Load AGENT(S).md if present and cache the contents."""
+        """Load AGENT(S).md from both global config and project directory.
+
+        Checks for AGENTS.md/AGENT.md/agents.md/agent.md in this order:
+        1. Global config directory (~/.code_puppy/ or XDG config)
+        2. Current working directory (project-specific)
+
+        If both exist, they are combined with global rules first, then project rules.
+        This allows project-specific rules to override or extend global rules.
+        """
         if self._puppy_rules is not None:
             return self._puppy_rules
         from pathlib import Path
 
         possible_paths = ["AGENTS.md", "AGENT.md", "agents.md", "agent.md"]
+
+        # Load global rules from CONFIG_DIR
+        global_rules = None
+        from code_puppy.config import CONFIG_DIR
+
         for path_str in possible_paths:
-            puppy_rules_path = Path(path_str)
-            if puppy_rules_path.exists():
-                with open(puppy_rules_path, "r") as f:
-                    self._puppy_rules = f.read()
-                    break
+            global_path = Path(CONFIG_DIR) / path_str
+            if global_path.exists():
+                global_rules = global_path.read_text(encoding="utf-8-sig")
+                break
+
+        # Load project-local rules from current working directory
+        project_rules = None
+        for path_str in possible_paths:
+            project_path = Path(path_str)
+            if project_path.exists():
+                project_rules = project_path.read_text(encoding="utf-8-sig")
+                break
+
+        # Combine global and project rules
+        # Global rules come first, project rules second (allowing project to override)
+        rules = [r for r in [global_rules, project_rules] if r]
+        self._puppy_rules = "\n\n".join(rules) if rules else None
         return self._puppy_rules
 
     def load_mcp_servers(self, extra_headers: Optional[Dict[str, str]] = None):
-        """Load MCP servers through the manager and return pydantic-ai compatible servers."""
+        """Load MCP servers through the manager and return pydantic-ai compatible servers.
+
+        Note: The manager automatically syncs from mcp_servers.json during initialization,
+        so we don't need to sync here. Use reload_mcp_servers() to force a re-sync.
+        """
 
         mcp_disabled = get_value("disable_mcp_servers")
         if mcp_disabled and str(mcp_disabled).lower() in ("1", "true", "yes", "on"):
             return []
 
         manager = get_mcp_manager()
-        configs = load_mcp_server_configs()
-        if not configs:
-            existing_servers = manager.list_servers()
-            if not existing_servers:
-                return []
-        else:
-            for name, conf in configs.items():
-                try:
-                    server_config = ServerConfig(
-                        id=conf.get("id", f"{name}_{hash(name)}"),
-                        name=name,
-                        type=conf.get("type", "sse"),
-                        enabled=conf.get("enabled", True),
-                        config=conf,
-                    )
-                    existing = manager.get_server_by_name(name)
-                    if not existing:
-                        manager.register_server(server_config)
-                    else:
-                        if existing.config != server_config.config:
-                            manager.update_server(existing.id, server_config)
-                except Exception:
-                    continue
-
         return manager.get_servers_for_agent()
 
     def reload_mcp_servers(self):
-        """Reload MCP servers and return updated servers."""
+        """Reload MCP servers and return updated servers.
+
+        Forces a re-sync from mcp_servers.json to pick up any configuration changes.
+        """
         # Clear the MCP tool cache when servers are reloaded
         self._mcp_tool_definitions_cache = []
-        self.load_mcp_servers()
+
+        # Force re-sync from mcp_servers.json
         manager = get_mcp_manager()
+        manager.sync_from_config()
+
         return manager.get_servers_for_agent()
 
     def _load_model_with_fallback(
@@ -1049,8 +1049,8 @@ class BaseAgent(ABC):
             )
             emit_warning(
                 (
-                    f"[yellow]Model '{requested_model_name}' not found. "
-                    f"Available models: {available_str}[/yellow]"
+                    f"Model '{requested_model_name}' not found. "
+                    f"Available models: {available_str}"
                 ),
                 message_group=message_group,
             )
@@ -1070,7 +1070,7 @@ class BaseAgent(ABC):
                 try:
                     model = ModelFactory.get_model(candidate, models_config)
                     emit_info(
-                        f"[bold cyan]Using fallback model: {candidate}[/bold cyan]",
+                        f"Using fallback model: {candidate}",
                         message_group=message_group,
                     )
                     return model, candidate
@@ -1082,7 +1082,7 @@ class BaseAgent(ABC):
                 "a valid model with `config set`."
             )
             emit_error(
-                f"[bold red]{friendly_message}[/bold red]",
+                friendly_message,
                 message_group=message_group,
             )
             raise ValueError(friendly_message) from exc
@@ -1110,13 +1110,7 @@ class BaseAgent(ABC):
 
         mcp_servers = self.load_mcp_servers()
 
-        output_tokens = max(
-            2048,
-            min(int(0.05 * self.get_model_context_length()) - 1024, 16384),
-        )
-        model_settings = make_model_settings(
-            resolved_model_name, max_tokens=output_tokens
-        )
+        model_settings = make_model_settings(resolved_model_name)
 
         # Handle claude-code models: swap instructions (prompt prepending happens in run_with_mcp)
         from code_puppy.model_utils import prepare_prompt_for_model
@@ -1189,7 +1183,9 @@ class BaseAgent(ABC):
 
         if len(filtered_mcp_servers) != len(mcp_servers):
             emit_info(
-                f"[dim]Filtered {len(mcp_servers) - len(filtered_mcp_servers)} conflicting MCP tools[/dim]"
+                Text.from_markup(
+                    f"[dim]Filtered {len(mcp_servers) - len(filtered_mcp_servers)} conflicting MCP tools[/dim]"
+                )
             )
 
         self._last_model_name = resolved_model_name
@@ -1246,6 +1242,74 @@ class BaseAgent(ABC):
             self._mcp_servers = mcp_servers
         return self._code_generation_agent
 
+    def _create_agent_with_output_type(self, output_type: Type[Any]) -> PydanticAgent:
+        """Create a temporary agent configured with a custom output_type.
+
+        This is used when structured output is requested via run_with_mcp.
+        The agent is created fresh with the same configuration as the main agent
+        but with the specified output_type instead of str.
+
+        Args:
+            output_type: The Pydantic model or type for structured output.
+
+        Returns:
+            A configured PydanticAgent (or DBOSAgent wrapper) with the custom output_type.
+        """
+        from code_puppy.model_utils import prepare_prompt_for_model
+        from code_puppy.tools import register_tools_for_agent
+
+        model_name = self.get_model_name()
+        models_config = ModelFactory.load_config()
+        model, resolved_model_name = self._load_model_with_fallback(
+            model_name, models_config, str(uuid.uuid4())
+        )
+
+        instructions = self.get_system_prompt()
+        puppy_rules = self.load_puppy_rules()
+        if puppy_rules:
+            instructions += f"\n{puppy_rules}"
+
+        mcp_servers = getattr(self, "_mcp_servers", []) or []
+        model_settings = make_model_settings(resolved_model_name)
+
+        prepared = prepare_prompt_for_model(
+            model_name, instructions, "", prepend_system_to_user=False
+        )
+        instructions = prepared.instructions
+
+        global _reload_count
+        _reload_count += 1
+
+        if get_use_dbos():
+            temp_agent = PydanticAgent(
+                model=model,
+                instructions=instructions,
+                output_type=output_type,
+                retries=3,
+                toolsets=[],
+                history_processors=[self.message_history_accumulator],
+                model_settings=model_settings,
+            )
+            agent_tools = self.get_available_tools()
+            register_tools_for_agent(temp_agent, agent_tools)
+            dbos_agent = DBOSAgent(
+                temp_agent, name=f"{self.name}-structured-{_reload_count}"
+            )
+            return dbos_agent
+        else:
+            temp_agent = PydanticAgent(
+                model=model,
+                instructions=instructions,
+                output_type=output_type,
+                retries=3,
+                toolsets=mcp_servers,
+                history_processors=[self.message_history_accumulator],
+                model_settings=model_settings,
+            )
+            agent_tools = self.get_available_tools()
+            register_tools_for_agent(temp_agent, agent_tools)
+            return temp_agent
+
     # It's okay to decorate it with DBOS.step even if not using DBOS; the decorator is a no-op in that case.
     @DBOS.step()
     def message_history_accumulator(self, ctx: RunContext, messages: List[Any]):
@@ -1271,12 +1335,204 @@ class BaseAgent(ABC):
             self.set_message_history(result_messages_filtered_empty_thinking)
         return self.get_message_history()
 
+    async def _event_stream_handler(
+        self, ctx: RunContext, events: AsyncIterable[Any]
+    ) -> None:
+        """Handle streaming events from the agent run.
+
+        This method processes streaming events and emits TextPart and ThinkingPart
+        content with styled banners as they stream in.
+
+        Args:
+            ctx: The run context.
+            events: Async iterable of streaming events (PartStartEvent, PartDeltaEvent, etc.).
+        """
+        from pydantic_ai import PartDeltaEvent, PartStartEvent
+        from pydantic_ai.messages import TextPartDelta, ThinkingPartDelta
+        from rich.console import Console
+        from rich.markdown import Markdown
+        from rich.markup import escape
+
+        from code_puppy.messaging.spinner import pause_all_spinners
+
+        # IMPORTANT: Use the shared console (set by cli_runner) to avoid conflicts
+        # with the spinner's Live display. Multiple Console instances with separate
+        # Live displays cause cursor positioning chaos and line duplication.
+        if self._console is not None:
+            console = self._console
+        else:
+            # Fallback if console not set (shouldn't happen in normal use)
+            console = Console()
+
+        # Track which part indices we're currently streaming (for Text/Thinking 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
+        banner_printed: set[int] = set()  # Track if banner was already printed
+        text_buffer: dict[int, list[str]] = {}  # Buffer text for final markdown render
+        token_count: dict[int, int] = {}  # Track token count per text part
+        did_stream_anything = False  # Track if we streamed any content
+
+        def _print_thinking_banner() -> None:
+            """Print the THINKING banner with spinner pause and line clear."""
+            nonlocal did_stream_anything
+            import sys
+            import time
+
+            from code_puppy.config import get_banner_color
+
+            pause_all_spinners()
+            time.sleep(0.1)  # Delay to let spinner fully clear
+            sys.stdout.write("\r\x1b[K")  # Clear line
+            sys.stdout.flush()
+            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]⚡ "
+                ),
+                end="",
+            )
+            sys.stdout.flush()
+            did_stream_anything = True
+
+        def _print_response_banner() -> None:
+            """Print the AGENT RESPONSE banner with spinner pause and line clear."""
+            nonlocal did_stream_anything
+            import sys
+            import time
+
+            from code_puppy.config import get_banner_color
+
+            pause_all_spinners()
+            time.sleep(0.1)  # Delay to let spinner fully clear
+            sys.stdout.write("\r\x1b[K")  # Clear line
+            sys.stdout.flush()
+            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}]"
+                )
+            )
+            sys.stdout.flush()
+            did_stream_anything = True
+
+        async for event in events:
+            # PartStartEvent - register the part but defer banner until content arrives
+            if isinstance(event, PartStartEvent):
+                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():
+                        _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)
+                    text_buffer[event.index] = []  # Initialize buffer
+                    token_count[event.index] = 0  # Initialize token counter
+                    # Buffer initial content if present
+                    if part.content and part.content.strip():
+                        text_buffer[event.index].append(part.content)
+                        token_count[event.index] += 1
+
+            # PartDeltaEvent - stream the content as it arrives
+            elif isinstance(event, PartDeltaEvent):
+                if event.index in streaming_parts:
+                    delta = event.delta
+                    if isinstance(delta, (TextPartDelta, ThinkingPartDelta)):
+                        if delta.content_delta:
+                            # For text parts, show token counter then render at end
+                            if event.index in text_parts:
+                                import sys
+
+                                # Print banner on first content
+                                if event.index not in banner_printed:
+                                    _print_response_banner()
+                                    banner_printed.add(event.index)
+                                # Accumulate text for final markdown render
+                                text_buffer[event.index].append(delta.content_delta)
+                                token_count[event.index] += 1
+                                # Update token counter in place (single line)
+                                count = token_count[event.index]
+                                sys.stdout.write(
+                                    f"\r\x1b[K  ⏳ Receiving... {count} tokens"
+                                )
+                                sys.stdout.flush()
+                            else:
+                                # For thinking parts, stream immediately (dim)
+                                if event.index not in banner_printed:
+                                    _print_thinking_banner()
+                                    banner_printed.add(event.index)
+                                escaped = escape(delta.content_delta)
+                                console.print(f"[dim]{escaped}[/dim]", end="")
+
+            # PartEndEvent - finish the streaming with a newline
+            elif isinstance(event, PartEndEvent):
+                if event.index in streaming_parts:
+                    # For text parts, clear counter line and render markdown
+                    if event.index in text_parts:
+                        import sys
+
+                        # Clear the token counter line
+                        sys.stdout.write("\r\x1b[K")
+                        sys.stdout.flush()
+                        # Render the final markdown nicely
+                        if event.index in text_buffer:
+                            try:
+                                final_content = "".join(text_buffer[event.index])
+                                if final_content.strip():
+                                    console.print(Markdown(final_content))
+                            except Exception:
+                                pass
+                            del text_buffer[event.index]
+                        # Clean up token count
+                        token_count.pop(event.index, None)
+                    # For thinking parts, just print newline
+                    elif event.index in banner_printed:
+                        console.print()  # Final newline after streaming
+                    # Clean up all tracking sets
+                    streaming_parts.discard(event.index)
+                    thinking_parts.discard(event.index)
+                    text_parts.discard(event.index)
+                    banner_printed.discard(event.index)
+
+                    # Resume spinner if next part is NOT text/thinking (avoid race condition)
+                    # If next part is a tool call or None, it's safe to resume
+                    # Note: spinner itself handles blank line before appearing
+                    from code_puppy.messaging.spinner import resume_all_spinners
+
+                    next_kind = getattr(event, "next_part_kind", None)
+                    if next_kind not in ("text", "thinking"):
+                        resume_all_spinners()
+
+        # Spinner is resumed in PartEndEvent when appropriate (based on next_part_kind)
+
     def _spawn_ctrl_x_key_listener(
         self,
         stop_event: threading.Event,
         on_escape: Callable[[], None],
+        on_cancel_agent: Optional[Callable[[], None]] = None,
     ) -> Optional[threading.Thread]:
-        """Start a Ctrl+X key listener thread for CLI sessions."""
+        """Start a keyboard listener thread for CLI sessions.
+
+        Listens for Ctrl+X (shell command cancel) and optionally the configured
+        cancel_agent_key (when not using SIGINT/Ctrl+C).
+
+        Args:
+            stop_event: Event to signal the listener to stop.
+            on_escape: Callback for Ctrl+X (shell command cancel).
+            on_cancel_agent: Optional callback for cancel_agent_key (only used
+                when cancel_agent_uses_signal() returns False).
+        """
         try:
             import sys
         except ImportError:
@@ -1294,16 +1550,20 @@ class BaseAgent(ABC):
         def listener() -> None:
             try:
                 if sys.platform.startswith("win"):
-                    self._listen_for_ctrl_x_windows(stop_event, on_escape)
+                    self._listen_for_ctrl_x_windows(
+                        stop_event, on_escape, on_cancel_agent
+                    )
                 else:
-                    self._listen_for_ctrl_x_posix(stop_event, on_escape)
+                    self._listen_for_ctrl_x_posix(
+                        stop_event, on_escape, on_cancel_agent
+                    )
             except Exception:
                 emit_warning(
-                    "Ctrl+X key listener stopped unexpectedly; press Ctrl+C to cancel."
+                    "Key listener stopped unexpectedly; press Ctrl+C to cancel."
                 )
 
         thread = threading.Thread(
-            target=listener, name="code-puppy-esc-listener", daemon=True
+            target=listener, name="code-puppy-key-listener", daemon=True
         )
         thread.start()
         return thread
@@ -1312,10 +1572,16 @@ class BaseAgent(ABC):
         self,
         stop_event: threading.Event,
         on_escape: Callable[[], None],
+        on_cancel_agent: Optional[Callable[[], None]] = None,
     ) -> None:
         import msvcrt
         import time
 
+        # Get the cancel agent char code if we're using keyboard-based cancel
+        cancel_agent_char: Optional[str] = None
+        if on_cancel_agent is not None and not cancel_agent_uses_signal():
+            cancel_agent_char = get_cancel_agent_char_code()
+
         while not stop_event.is_set():
             try:
                 if msvcrt.kbhit():
@@ -1327,9 +1593,18 @@ class BaseAgent(ABC):
                             emit_warning(
                                 "Ctrl+X handler raised unexpectedly; Ctrl+C still works."
                             )
+                    elif (
+                        cancel_agent_char
+                        and on_cancel_agent
+                        and key == cancel_agent_char
+                    ):
+                        try:
+                            on_cancel_agent()
+                        except Exception:
+                            emit_warning("Cancel agent handler raised unexpectedly.")
             except Exception:
                 emit_warning(
-                    "Windows Ctrl+X listener error; Ctrl+C is still available for cancel."
+                    "Windows key listener error; Ctrl+C is still available for cancel."
                 )
                 return
             time.sleep(0.05)
@@ -1338,12 +1613,18 @@ class BaseAgent(ABC):
         self,
         stop_event: threading.Event,
         on_escape: Callable[[], None],
+        on_cancel_agent: Optional[Callable[[], None]] = None,
     ) -> None:
         import select
         import sys
         import termios
         import tty
 
+        # Get the cancel agent char code if we're using keyboard-based cancel
+        cancel_agent_char: Optional[str] = None
+        if on_cancel_agent is not None and not cancel_agent_uses_signal():
+            cancel_agent_char = get_cancel_agent_char_code()
+
         stdin = sys.stdin
         try:
             fd = stdin.fileno()
@@ -1373,6 +1654,13 @@ class BaseAgent(ABC):
                         emit_warning(
                             "Ctrl+X handler raised unexpectedly; Ctrl+C still works."
                         )
+                elif (
+                    cancel_agent_char and on_cancel_agent and data == cancel_agent_char
+                ):
+                    try:
+                        on_cancel_agent()
+                    except Exception:
+                        emit_warning("Cancel agent handler raised unexpectedly.")
         finally:
             termios.tcsetattr(fd, termios.TCSADRAIN, original_attrs)
 
@@ -1382,6 +1670,7 @@ class BaseAgent(ABC):
         *,
         attachments: Optional[Sequence[BinaryContent]] = None,
         link_attachments: Optional[Sequence[Union[ImageUrl, DocumentUrl]]] = None,
+        output_type: Optional[Type[Any]] = None,
         **kwargs,
     ) -> Any:
         """Run the agent with MCP servers, attachments, and full cancellation support.
@@ -1390,10 +1679,13 @@ class BaseAgent(ABC):
             prompt: Primary user prompt text (may be empty when attachments present).
             attachments: Local binary payloads (e.g., dragged images) to include.
             link_attachments: Remote assets (image/document URLs) to include.
+            output_type: Optional Pydantic model or type for structured output.
+                When provided, creates a temporary agent configured to return
+                this type instead of the default string output.
             **kwargs: Additional arguments forwarded to `pydantic_ai.Agent.run`.
 
         Returns:
-            The agent's response.
+            The agent's response (typed according to output_type if specified).
 
         Raises:
             asyncio.CancelledError: When execution is cancelled by user.
@@ -1417,12 +1709,23 @@ class BaseAgent(ABC):
         pydantic_agent = (
             self._code_generation_agent or self.reload_code_generation_agent()
         )
-        # Handle claude-code models: prepend system prompt to first user message
-        from code_puppy.model_utils import is_claude_code_model
 
-        if is_claude_code_model(self.get_model_name()):
+        # If a custom output_type is specified, create a temporary agent with that type
+        if output_type is not None:
+            pydantic_agent = self._create_agent_with_output_type(output_type)
+
+        # Handle claude-code and chatgpt-codex models: prepend system prompt to first user message
+        from code_puppy.model_utils import is_chatgpt_codex_model, is_claude_code_model
+
+        if is_claude_code_model(self.get_model_name()) or is_chatgpt_codex_model(
+            self.get_model_name()
+        ):
             if len(self.get_message_history()) == 0:
-                prompt = self.get_system_prompt() + "\n\n" + prompt
+                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
 
         # Build combined prompt payload when attachments are provided.
         attachment_parts: List[Any] = []
@@ -1480,6 +1783,7 @@ class BaseAgent(ABC):
                                 prompt_payload,
                                 message_history=self.get_message_history(),
                                 usage_limits=usage_limits,
+                                event_stream_handler=self._event_stream_handler,
                                 **kwargs,
                             )
                     finally:
@@ -1492,6 +1796,7 @@ class BaseAgent(ABC):
                             prompt_payload,
                             message_history=self.get_message_history(),
                             usage_limits=usage_limits,
+                            event_stream_handler=self._event_stream_handler,
                             **kwargs,
                         )
                 else:
@@ -1500,6 +1805,7 @@ class BaseAgent(ABC):
                         prompt_payload,
                         message_history=self.get_message_history(),
                         usage_limits=usage_limits,
+                        event_stream_handler=self._event_stream_handler,
                         **kwargs,
                     )
                 return result_
@@ -1537,6 +1843,12 @@ class BaseAgent(ABC):
                         remaining_exceptions.append(exc)
                         emit_info(f"Unexpected error: {str(exc)}", group_id=group_id)
                         emit_info(f"{str(exc.args)}", group_id=group_id)
+                        # Log to file for debugging
+                        log_error(
+                            exc,
+                            context=f"Agent run (group_id={group_id})",
+                            include_traceback=True,
+                        )
 
                 collect_non_cancelled_exceptions(other_error)
 
@@ -1595,10 +1907,78 @@ class BaseAgent(ABC):
 
             schedule_agent_cancel()
 
+        def graceful_sigint_handler(_sig, _frame):
+            # When using keyboard-based cancel, SIGINT should be a no-op
+            # (just show a hint to user about the configured cancel key)
+            from code_puppy.keymap import get_cancel_agent_display_name
+            import sys
+
+            cancel_key = get_cancel_agent_display_name()
+            if sys.platform == "win32":
+                # On Windows, we use keyboard listener, so SIGINT might still fire
+                # but we handle cancellation via the key listener
+                pass  # Silent on Windows - the key listener handles it
+            else:
+                emit_info(f"Use {cancel_key} to cancel the agent task.")
+
         original_handler = None
+        key_listener_stop_event = None
+        _key_listener_thread = None
+        _windows_ctrl_handler = None  # Store reference to prevent garbage collection
+
         try:
-            # Save original handler and set our custom one AFTER task is created
-            original_handler = signal.signal(signal.SIGINT, keyboard_interrupt_handler)
+            if sys.platform == "win32":
+                # Windows: Use SetConsoleCtrlHandler for reliable Ctrl+C handling
+                import ctypes
+
+                # Define the handler function type
+                HANDLER_ROUTINE = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_ulong)
+
+                def windows_ctrl_handler(ctrl_type):
+                    """Handle Windows console control events."""
+                    CTRL_C_EVENT = 0
+                    CTRL_BREAK_EVENT = 1
+
+                    if ctrl_type in (CTRL_C_EVENT, CTRL_BREAK_EVENT):
+                        # Check if we're awaiting user input
+                        if is_awaiting_user_input():
+                            return False  # Let default handler run
+
+                        # Schedule agent cancellation
+                        schedule_agent_cancel()
+                        return True  # We handled it, don't terminate
+
+                    return False  # Let other handlers process it
+
+                # Create the callback - must keep reference alive!
+                _windows_ctrl_handler = HANDLER_ROUTINE(windows_ctrl_handler)
+
+                # Register the handler
+                kernel32 = ctypes.windll.kernel32
+                if not kernel32.SetConsoleCtrlHandler(_windows_ctrl_handler, True):
+                    emit_warning("Failed to set Windows Ctrl+C handler")
+
+                # Also spawn keyboard listener for Ctrl+X (shell cancel) and other keys
+                key_listener_stop_event = threading.Event()
+                _key_listener_thread = self._spawn_ctrl_x_key_listener(
+                    key_listener_stop_event,
+                    on_escape=lambda: None,  # Ctrl+X handled by command_runner
+                    on_cancel_agent=None,  # Ctrl+C handled by SetConsoleCtrlHandler above
+                )
+            elif cancel_agent_uses_signal():
+                # Unix with Ctrl+C: Use SIGINT-based cancellation
+                original_handler = signal.signal(
+                    signal.SIGINT, keyboard_interrupt_handler
+                )
+            else:
+                # Unix with different cancel key: Use keyboard listener
+                original_handler = signal.signal(signal.SIGINT, graceful_sigint_handler)
+                key_listener_stop_event = threading.Event()
+                _key_listener_thread = self._spawn_ctrl_x_key_listener(
+                    key_listener_stop_event,
+                    on_escape=lambda: None,
+                    on_cancel_agent=schedule_agent_cancel,
+                )
 
             # Wait for the task to complete or be cancelled
             result = await agent_task
@@ -1618,6 +1998,20 @@ class BaseAgent(ABC):
             if not agent_task.done():
                 agent_task.cancel()
         finally:
-            # Restore original signal handler
-            if original_handler:
+            # Stop keyboard listener if it was started
+            if key_listener_stop_event is not None:
+                key_listener_stop_event.set()
+
+            # Unregister Windows Ctrl handler
+            if sys.platform == "win32" and _windows_ctrl_handler is not None:
+                try:
+                    import ctypes
+
+                    kernel32 = ctypes.windll.kernel32
+                    kernel32.SetConsoleCtrlHandler(_windows_ctrl_handler, False)
+                except Exception:
+                    pass  # Best effort cleanup
+
+            # Restore original signal handler (Unix)
+            if original_handler is not None:
                 signal.signal(signal.SIGINT, original_handler)