code-puppy 0.0.314__tar.gz → 0.0.315__tar.gz

{code_puppy-0.0.314 → code_puppy-0.0.315}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.314
+Version: 0.0.315
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy

{code_puppy-0.0.314 → code_puppy-0.0.315}/code_puppy/agents/base_agent.py

@@ -7,6 +7,7 @@ import signal
 import threading
 import uuid
 from abc import ABC, abstractmethod
+from collections.abc import AsyncIterable
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import mcp
@@ -18,6 +19,7 @@ from pydantic_ai import (
     BinaryContent,
     DocumentUrl,
     ImageUrl,
+    PartEndEvent,
     RunContext,
     UsageLimitExceeded,
     UsageLimits,
@@ -1265,6 +1267,224 @@ 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.).
+        """
+        import os
+        import time as time_module
+
+        from pydantic_ai import PartDeltaEvent, PartStartEvent
+        from pydantic_ai.messages import TextPartDelta, ThinkingPartDelta
+        from rich.console import Console
+        from rich.live import Live
+        from rich.markdown import Markdown
+        from rich.markup import escape
+
+        from code_puppy.messaging.spinner import pause_all_spinners
+
+        console = Console()
+
+        # Disable Live display in test mode or non-interactive environments
+        # This fixes issues with pexpect PTY where Live() hangs
+        use_live_display = (
+            console.is_terminal
+            and os.environ.get("CODE_PUPPY_TEST_FAST", "").lower() not in ("1", "true")
+            and os.environ.get("CI", "").lower() not in ("1", "true")
+        )
+
+        # 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 markdown
+        live_displays: dict[int, Live] = {}  # Live displays for streaming markdown
+        did_stream_anything = False  # Track if we streamed any content
+        last_render_time: dict[int, float] = {}  # Track last render time per part
+        render_interval = 0.1  # Only re-render markdown every 100ms (throttle)
+
+        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
+                    # Buffer initial content if present
+                    if part.content and part.content.strip():
+                        text_buffer[event.index].append(part.content)
+
+            # 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, stream markdown with Live display
+                            if event.index in text_parts:
+                                # Print banner and start Live on first content
+                                if event.index not in banner_printed:
+                                    _print_response_banner()
+                                    banner_printed.add(event.index)
+                                    # Only use Live display if enabled (disabled in test/CI)
+                                    if use_live_display:
+                                        live = Live(
+                                            Markdown(""),
+                                            console=console,
+                                            refresh_per_second=10,
+                                            vertical_overflow="visible",  # Allow scrolling for long content
+                                        )
+                                        live.start()
+                                        live_displays[event.index] = live
+                                # Accumulate text and throttle markdown rendering
+                                # (Markdown parsing is O(n), doing it on every token = O(n²) death)
+                                text_buffer[event.index].append(delta.content_delta)
+                                now = time_module.monotonic()
+                                last_render = last_render_time.get(event.index, 0)
+
+                                # Only re-render if enough time has passed (throttle)
+                                # Skip Live updates when not using live display
+                                if (
+                                    use_live_display
+                                    and now - last_render >= render_interval
+                                ):
+                                    content = "".join(text_buffer[event.index])
+                                    if event.index in live_displays:
+                                        try:
+                                            live_displays[event.index].update(
+                                                Markdown(content)
+                                            )
+                                            last_render_time[event.index] = now
+                                        except Exception:
+                                            pass
+                            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, do final render then stop the Live display
+                    if event.index in text_parts:
+                        # Final render to ensure we show complete content
+                        # (throttling may have skipped the last few tokens)
+                        if event.index in live_displays and event.index in text_buffer:
+                            try:
+                                final_content = "".join(text_buffer[event.index])
+                                live_displays[event.index].update(
+                                    Markdown(final_content)
+                                )
+                            except Exception:
+                                pass
+                        if event.index in live_displays:
+                            try:
+                                live_displays[event.index].stop()
+                            except Exception:
+                                pass
+                            del live_displays[event.index]
+                        # When not using Live display, print the final content as markdown
+                        elif 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
+                        if event.index in text_buffer:
+                            del text_buffer[event.index]
+                        # Clean up render time tracking
+                        last_render_time.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,
@@ -1523,6 +1743,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:
@@ -1535,6 +1756,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:
@@ -1543,6 +1765,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_

{code_puppy-0.0.314 → code_puppy-0.0.315}/code_puppy/chatgpt_codex_client.py

@@ -140,9 +140,11 @@ class ChatGPTCodexAsyncClient(httpx.AsyncClient):
             modified = True
 
         # CRITICAL: ChatGPT Codex backend requires stream=true
+        # If stream is already true (e.g., pydantic-ai with event_stream_handler),
+        # don't force conversion - let streaming events flow through naturally
         if data.get("stream") is not True:
             data["stream"] = True
-            forced_stream = True
+            forced_stream = True  # Only convert if WE forced streaming
             modified = True
 
         # Add reasoning settings for reasoning models (gpt-5.2, o-series, etc.)

code_puppy-0.0.314/code_puppy/main.py → code_puppy-0.0.315/code_puppy/cli_runner.py

@@ -1,3 +1,8 @@
+"""CLI runner for Code Puppy.
+
+Contains the main application logic, interactive mode, and entry point.
+"""
+
 # Apply pydantic-ai patches BEFORE any pydantic-ai imports
 from code_puppy.pydantic_patches import apply_all_patches
 
@@ -6,8 +11,6 @@ apply_all_patches()
 import argparse
 import asyncio
 import os
-import platform
-import subprocess
 import sys
 import time
 import traceback
@@ -39,15 +42,19 @@ from code_puppy.keymap import (
     validate_cancel_agent_key,
 )
 from code_puppy.messaging import emit_info
+from code_puppy.terminal_utils import (
+    reset_unix_terminal,
+    reset_windows_terminal_ansi,
+    reset_windows_terminal_full,
+)
 from code_puppy.tools.common import console
-
-# message_history_accumulator and prune_interrupted_tool_calls have been moved to BaseAgent class
 from code_puppy.version_checker import default_version_mismatch_behavior
 
 plugins.load_plugin_callbacks()
 
 
 async def main():
+    """Main async entry point for Code Puppy CLI."""
     parser = argparse.ArgumentParser(description="Code Puppy - A code generation agent")
     parser.add_argument(
         "--version",
@@ -295,11 +302,9 @@ async def main():
             DBOS.destroy()
 
 
-# Add the file handling functionality for interactive mode
 async def interactive_mode(message_renderer, initial_command: str = None) -> None:
-    from code_puppy.command_line.command_handler import handle_command
-
     """Run the agent in interactive mode."""
+    from code_puppy.command_line.command_handler import handle_command
 
     display_console = message_renderer.console
     from code_puppy.messaging import emit_info, emit_system_message
@@ -429,12 +434,7 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
             # Use prompt_toolkit for enhanced input with path completion
             try:
                 # Windows-specific: Reset terminal state before prompting
-                if platform.system() == "Windows":
-                    try:
-                        sys.stdout.write("\x1b[0m")  # Reset ANSI formatting
-                        sys.stdout.flush()
-                    except Exception:
-                        pass
+                reset_windows_terminal_ansi()
 
                 # Use the async version of get_input_with_combined_completion
                 task = await get_input_with_combined_completion(
@@ -446,6 +446,9 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
 
         except (KeyboardInterrupt, EOFError):
             # Handle Ctrl+C or Ctrl+D
+            # Windows-specific: Reset terminal state after interrupt to prevent
+            # the terminal from becoming unresponsive (can't type characters)
+            reset_windows_terminal_full()
             from code_puppy.messaging import emit_warning
 
             emit_warning("\nInput cancelled")
@@ -601,14 +604,7 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
                 # Check if the task was cancelled (but don't show message if we just killed processes)
                 if result is None:
                     # Windows-specific: Reset terminal state after cancellation
-                    if platform.system() == "Windows":
-                        try:
-                            sys.stdout.write("\x1b[0m")  # Reset ANSI formatting
-                            sys.stdout.flush()
-                            sys.stderr.write("\x1b[0m")
-                            sys.stderr.flush()
-                        except Exception:
-                            pass
+                    reset_windows_terminal_ansi()
                     continue
                 # Get the structured response
                 agent_response = result.output
@@ -651,6 +647,8 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non
 
 
 def prettier_code_blocks():
+    """Configure Rich to use prettier code block rendering."""
+
     class SimpleCodeBlock(CodeBlock):
         def __rich_console__(
             self, console: Console, options: ConsoleOptions
@@ -787,14 +785,4 @@ def main_entry():
         return 0
     finally:
         # Reset terminal on Unix-like systems (not Windows)
-        if platform.system() != "Windows":
-            try:
-                # Reset terminal to sanity state
-                subprocess.run(["reset"], check=True, capture_output=True)
-            except (subprocess.CalledProcessError, FileNotFoundError):
-                # Silently fail if reset command isn't available
-                pass
-
-
-if __name__ == "__main__":
-    main_entry()
+        reset_unix_terminal()