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

code_puppy/pydantic_patches.py

@@ -0,0 +1,131 @@
+"""Monkey patches for pydantic-ai.
+
+This module contains all monkey patches needed to customize pydantic-ai behavior.
+These patches MUST be applied before any other pydantic-ai imports to work correctly.
+
+Usage:
+    from code_puppy.pydantic_patches import apply_all_patches
+    apply_all_patches()
+"""
+
+import importlib.metadata
+
+
+def _get_code_puppy_version() -> str:
+    """Get the current code-puppy version."""
+    try:
+        return importlib.metadata.version("code-puppy")
+    except Exception:
+        return "0.0.0-dev"
+
+
+def patch_user_agent() -> None:
+    """Patch pydantic-ai's User-Agent to use Code-Puppy's version.
+
+    pydantic-ai sets its own User-Agent ('pydantic-ai/x.x.x') via a @cache-decorated
+    function. We replace it with a dynamic function that returns:
+    - 'KimiCLI/0.63' for Kimi models
+    - 'Code-Puppy/{version}' for all other models
+
+    This MUST be called before any pydantic-ai models are created.
+    """
+    try:
+        import pydantic_ai.models as pydantic_models
+
+        version = _get_code_puppy_version()
+
+        # Clear cache if already called
+        if hasattr(pydantic_models.get_user_agent, "cache_clear"):
+            pydantic_models.get_user_agent.cache_clear()
+
+        def _get_dynamic_user_agent() -> str:
+            """Return User-Agent based on current model selection."""
+            try:
+                from code_puppy.config import get_global_model_name
+
+                model_name = get_global_model_name()
+                if model_name and "kimi" in model_name.lower():
+                    return "KimiCLI/0.63"
+            except Exception:
+                pass
+            return f"Code-Puppy/{version}"
+
+        pydantic_models.get_user_agent = _get_dynamic_user_agent
+    except Exception:
+        pass  # Don't crash on patch failure
+
+
+def patch_message_history_cleaning() -> None:
+    """Disable overly strict message history cleaning in pydantic-ai."""
+    try:
+        from pydantic_ai import _agent_graph
+
+        _agent_graph._clean_message_history = lambda messages: messages
+    except Exception:
+        pass
+
+
+def patch_process_message_history() -> None:
+    """Patch _process_message_history to skip strict ModelRequest validation.
+
+    Pydantic AI added a validation that history must end with ModelRequest,
+    but this breaks valid conversation flows. We patch it to skip that validation.
+    """
+    try:
+        from pydantic_ai import _agent_graph
+
+        async def _patched_process_message_history(messages, processors, run_context):
+            """Patched version that doesn't enforce ModelRequest at end."""
+            from pydantic_ai._agent_graph import (
+                _HistoryProcessorAsync,
+                _HistoryProcessorSync,
+                _HistoryProcessorSyncWithCtx,
+                cast,
+                exceptions,
+                is_async_callable,
+                is_takes_ctx,
+                run_in_executor,
+            )
+
+            for processor in processors:
+                takes_ctx = is_takes_ctx(processor)
+
+                if is_async_callable(processor):
+                    if takes_ctx:
+                        messages = await processor(run_context, messages)
+                    else:
+                        async_processor = cast(_HistoryProcessorAsync, processor)
+                        messages = await async_processor(messages)
+                else:
+                    if takes_ctx:
+                        sync_processor_with_ctx = cast(
+                            _HistoryProcessorSyncWithCtx, processor
+                        )
+                        messages = await run_in_executor(
+                            sync_processor_with_ctx, run_context, messages
+                        )
+                    else:
+                        sync_processor = cast(_HistoryProcessorSync, processor)
+                        messages = await run_in_executor(sync_processor, messages)
+
+            if len(messages) == 0:
+                raise exceptions.UserError("Processed history cannot be empty.")
+
+            # NOTE: We intentionally skip the "must end with ModelRequest" validation
+            # that was added in newer Pydantic AI versions.
+
+            return messages
+
+        _agent_graph._process_message_history = _patched_process_message_history
+    except Exception:
+        pass
+
+
+def apply_all_patches() -> None:
+    """Apply all pydantic-ai monkey patches.
+
+    Call this at the very top of main.py, before any other imports.
+    """
+    patch_user_agent()
+    patch_message_history_cleaning()
+    patch_process_message_history()

code_puppy/session_storage.py

@@ -146,6 +146,7 @@ async def restore_autosave_interactively(base_dir: Path) -> None:
 
     # Import locally to avoid pulling the messaging layer into storage modules
     from datetime import datetime
+
     from prompt_toolkit.formatted_text import FormattedText
 
     from code_puppy.agents.agent_manager import get_current_agent
@@ -186,7 +187,7 @@ async def restore_autosave_interactively(base_dir: Path) -> None:
         start = page * PAGE_SIZE
         end = min(start + PAGE_SIZE, total)
         page_entries = entries[start:end]
-        emit_system_message("[bold magenta]Autosave Sessions Available:[/bold magenta]")
+        emit_system_message("Autosave Sessions Available:")
         for idx, (name, timestamp, message_count) in enumerate(page_entries, start=1):
             timestamp_display = timestamp or "unknown time"
             message_display = (

code_puppy/status_display.py

@@ -7,6 +7,8 @@ from rich.panel import Panel
 from rich.spinner import Spinner
 from rich.text import Text
 
+from code_puppy.messaging import emit_info
+
 # Global variable to track current token per second rate
 CURRENT_TOKEN_RATE = 0.0
 
@@ -185,7 +187,7 @@ class StatusDisplay:
     async def _update_display(self) -> None:
         """Update the display continuously while active using Rich Live display"""
         # Add a newline to ensure we're below the blue bar
-        self.console.print("\n")
+        emit_info("")
 
         # Create a Live display that will update in-place
         with Live(
@@ -221,8 +223,8 @@ class StatusDisplay:
             # Print final stats
             elapsed = time.time() - self.start_time if self.start_time else 0
             avg_rate = self.token_count / elapsed if elapsed > 0 else 0
-            self.console.print(
-                f"[dim]Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)[/dim]"
+            emit_info(
+                f"Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)"
             )
 
             # Reset state
@@ -240,6 +242,6 @@ class StatusDisplay:
             # This is for testing purposes
             elapsed = time.time() - self.start_time if self.start_time else 0
             avg_rate = self.token_count / elapsed if elapsed > 0 else 0
-            self.console.print(
-                f"[dim]Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)[/dim]"
+            emit_info(
+                f"Completed: {self.token_count} tokens in {elapsed:.1f}s ({avg_rate:.1f} t/s avg)"
             )

code_puppy/terminal_utils.py

@@ -0,0 +1,126 @@
+"""Terminal utilities for cross-platform terminal state management.
+
+Handles Windows console mode resets and Unix terminal sanity restoration.
+"""
+
+import platform
+import subprocess
+import sys
+
+
+def reset_windows_terminal_ansi() -> None:
+    """Reset ANSI formatting on Windows stdout/stderr.
+
+    This is a lightweight reset that just clears ANSI escape sequences.
+    Use this for quick resets after output operations.
+    """
+    if platform.system() != "Windows":
+        return
+
+    try:
+        sys.stdout.write("\x1b[0m")  # Reset ANSI formatting
+        sys.stdout.flush()
+        sys.stderr.write("\x1b[0m")
+        sys.stderr.flush()
+    except Exception:
+        pass  # Silently ignore errors - best effort reset
+
+
+def reset_windows_console_mode() -> None:
+    """Full Windows console mode reset using ctypes.
+
+    This resets both stdout and stdin console modes to restore proper
+    terminal behavior after interrupts (Ctrl+C, Ctrl+D). Without this,
+    the terminal can become unresponsive (can't type characters).
+    """
+    if platform.system() != "Windows":
+        return
+
+    try:
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+
+        # Reset stdout
+        STD_OUTPUT_HANDLE = -11
+        handle = kernel32.GetStdHandle(STD_OUTPUT_HANDLE)
+
+        # Enable virtual terminal processing and line input
+        mode = ctypes.c_ulong()
+        kernel32.GetConsoleMode(handle, ctypes.byref(mode))
+
+        # Console mode flags for stdout
+        ENABLE_PROCESSED_OUTPUT = 0x0001
+        ENABLE_WRAP_AT_EOL_OUTPUT = 0x0002
+        ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
+
+        new_mode = (
+            mode.value
+            | ENABLE_PROCESSED_OUTPUT
+            | ENABLE_WRAP_AT_EOL_OUTPUT
+            | ENABLE_VIRTUAL_TERMINAL_PROCESSING
+        )
+        kernel32.SetConsoleMode(handle, new_mode)
+
+        # Reset stdin
+        STD_INPUT_HANDLE = -10
+        stdin_handle = kernel32.GetStdHandle(STD_INPUT_HANDLE)
+
+        # Console mode flags for stdin
+        ENABLE_LINE_INPUT = 0x0002
+        ENABLE_ECHO_INPUT = 0x0004
+        ENABLE_PROCESSED_INPUT = 0x0001
+
+        stdin_mode = ctypes.c_ulong()
+        kernel32.GetConsoleMode(stdin_handle, ctypes.byref(stdin_mode))
+
+        new_stdin_mode = (
+            stdin_mode.value
+            | ENABLE_LINE_INPUT
+            | ENABLE_ECHO_INPUT
+            | ENABLE_PROCESSED_INPUT
+        )
+        kernel32.SetConsoleMode(stdin_handle, new_stdin_mode)
+
+    except Exception:
+        pass  # Silently ignore errors - best effort reset
+
+
+def reset_windows_terminal_full() -> None:
+    """Perform a full Windows terminal reset (ANSI + console mode).
+
+    Combines both ANSI reset and console mode reset for complete
+    terminal state restoration after interrupts.
+    """
+    if platform.system() != "Windows":
+        return
+
+    reset_windows_terminal_ansi()
+    reset_windows_console_mode()
+
+
+def reset_unix_terminal() -> None:
+    """Reset Unix/Linux/macOS terminal to sane state.
+
+    Uses the `reset` command to restore terminal sanity.
+    Silently fails if the command isn't available.
+    """
+    if platform.system() == "Windows":
+        return
+
+    try:
+        subprocess.run(["reset"], check=True, capture_output=True)
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        pass  # Silently fail if reset command isn't available
+
+
+def reset_terminal() -> None:
+    """Cross-platform terminal reset.
+
+    Automatically detects the platform and performs the appropriate
+    terminal reset operation.
+    """
+    if platform.system() == "Windows":
+        reset_windows_terminal_full()
+    else:
+        reset_unix_terminal()

code_puppy/tools/agent_tools.py

@@ -1,5 +1,6 @@
 # agent_tools.py
 import asyncio
+import hashlib
 import itertools
 import json
 import pickle
@@ -17,19 +18,22 @@ from pydantic_ai import Agent, RunContext, UsageLimits
 from pydantic_ai.messages import ModelMessage
 
 from code_puppy.config import (
+    DATA_DIR,
     get_message_limit,
     get_use_dbos,
 )
 from code_puppy.messaging import (
-    emit_divider,
+    SubAgentInvocationMessage,
+    SubAgentResponseMessage,
     emit_error,
     emit_info,
-    emit_system_message,
+    get_message_bus,
+    get_session_context,
+    set_session_context,
 )
 from code_puppy.model_factory import ModelFactory, make_model_settings
 from code_puppy.tools.common import generate_group_id
 
-_temp_agent_count = 0
 # Set to track active subagent invocation tasks
 _active_subagent_tasks: Set[asyncio.Task] = set()
 
@@ -55,6 +59,16 @@ def _generate_dbos_workflow_id(base_id: str) -> str:
     return f"{base_id}-wf-{counter}"
 
 
+def _generate_session_hash_suffix() -> str:
+    """Generate a short SHA1 hash suffix based on current timestamp for uniqueness.
+
+    Returns:
+        A 6-character hex string, e.g., "a3f2b1"
+    """
+    timestamp = str(datetime.now().timestamp())
+    return hashlib.sha1(timestamp.encode()).hexdigest()[:6]
+
+
 # Regex pattern for kebab-case session IDs
 SESSION_ID_PATTERN = re.compile(r"^[a-z0-9]+(-[a-z0-9]+)*$")
 SESSION_ID_MAX_LENGTH = 128
@@ -100,10 +114,10 @@ def _get_subagent_sessions_dir() -> Path:
     """Get the directory for storing subagent session data.
 
     Returns:
-        Path to ~/.code_puppy/subagent_sessions/
+        Path to XDG data directory/subagent_sessions/
     """
-    sessions_dir = Path.home() / ".code_puppy" / "subagent_sessions"
-    sessions_dir.mkdir(parents=True, exist_ok=True)
+    sessions_dir = Path(DATA_DIR) / "subagent_sessions"
+    sessions_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
     return sessions_dir
 
 
@@ -194,6 +208,7 @@ class AgentInfo(BaseModel):
 
     name: str
     display_name: str
+    description: str
 
 
 class ListAgentsOutput(BaseModel):
@@ -208,6 +223,7 @@ class AgentInvokeOutput(BaseModel):
 
     response: str | None
     agent_name: str
+    session_id: str | None = None
     error: str | None = None
 
 
@@ -228,38 +244,50 @@ def register_list_agents(agent):
         # Generate a group ID for this tool execution
         group_id = generate_group_id("list_agents")
 
+        from rich.text import Text
+
+        from code_puppy.config import get_banner_color
+
+        list_agents_color = get_banner_color("list_agents")
         emit_info(
-            "\n[bold white on blue] LIST AGENTS [/bold white on blue]",
+            Text.from_markup(
+                f"\n[bold white on {list_agents_color}] LIST AGENTS [/bold white on {list_agents_color}]"
+            ),
             message_group=group_id,
         )
-        emit_divider(message_group=group_id)
 
         try:
-            from code_puppy.agents import get_available_agents
+            from code_puppy.agents import get_agent_descriptions, get_available_agents
 
-            # Get available agents from the agent manager
+            # Get available agents and their descriptions from the agent manager
             agents_dict = get_available_agents()
+            descriptions_dict = get_agent_descriptions()
 
             # Convert to list of AgentInfo objects
             agents = [
-                AgentInfo(name=name, display_name=display_name)
+                AgentInfo(
+                    name=name,
+                    display_name=display_name,
+                    description=descriptions_dict.get(name, "No description available"),
+                )
                 for name, display_name in agents_dict.items()
             ]
 
-            # Display the agents in the console
+            # Accumulate output into a single string and emit once
+            # Use Text.from_markup() to pass a Rich object that won't be escaped
+            lines = []
             for agent_item in agents:
-                emit_system_message(
-                    f"- [bold]{agent_item.name}[/bold]: {agent_item.display_name}",
-                    message_group=group_id,
+                lines.append(
+                    f"- [bold]{agent_item.name}[/bold]: {agent_item.display_name}\n"
+                    f"  [dim]{agent_item.description}[/dim]"
                 )
+            emit_info(Text.from_markup("\n".join(lines)), message_group=group_id)
 
-            emit_divider(message_group=group_id)
             return ListAgentsOutput(agents=agents)
 
         except Exception as e:
             error_msg = f"Error listing agents: {str(e)}"
             emit_error(error_msg, message_group=group_id)
-            emit_divider(message_group=group_id)
             return ListAgentsOutput(agents=[], error=error_msg)
 
     return list_agents
@@ -285,21 +313,25 @@ def register_invoke_agent(agent):
 
                        **Session ID Format:**
                        - Must be kebab-case (lowercase letters, numbers, hyphens only)
-                       - Should be human-readable with random suffix: e.g., "implement-oauth-abc123", "review-auth-x7k9"
-                       - Add 3-6 random characters/numbers at the end to prevent namespace collisions
+                       - Should be human-readable: e.g., "implement-oauth", "review-auth"
+                       - For NEW sessions, a SHA1 hash suffix is automatically appended for uniqueness
+                       - To CONTINUE a session, use the full session_id (with hash) from the previous invocation
                        - If None (default), auto-generates like "agent-name-session-1"
 
                        **When to use session_id:**
-                       - **REUSE** the same session_id ONLY when you need the sub-agent to remember
-                         previous conversation context (e.g., multi-turn discussions, iterative reviews)
-                       - **DO NOT REUSE** for independent, one-off tasks - let it auto-generate or use
-                         unique IDs for each invocation
+                       - **NEW SESSION**: Provide a base name like "review-auth" - we'll append a unique hash
+                       - **CONTINUE SESSION**: Use the full session_id from output (e.g., "review-auth-a3f2b1")
+                       - **ONE-OFF TASKS**: Leave as None (auto-generate)
 
                        **Most common pattern:** Leave session_id as None (auto-generate) unless you
                        specifically need conversational memory.
 
         Returns:
-            AgentInvokeOutput: The agent's response to the prompt
+            AgentInvokeOutput: Contains:
+                - response (str | None): The agent's response to the prompt
+                - agent_name (str): Name of the invoked agent
+                - session_id (str | None): The full session ID (with hash suffix) - USE THIS to continue the conversation!
+                - error (str | None): Error message if invocation failed
 
         Examples:
             # COMMON CASE: One-off invocation, no memory needed (auto-generate session)
@@ -307,46 +339,43 @@ def register_invoke_agent(agent):
                 "qa-expert",
                 "Review this function: def add(a, b): return a + b"
             )
+            # result.session_id will be something like "qa-expert-session-a3f2b1"
 
-            # MULTI-TURN: Start a conversation with explicit session ID (note random suffix)
+            # MULTI-TURN: Start a NEW conversation with a base session ID
+            # A hash suffix is auto-appended: "review-add-function" -> "review-add-function-a3f2b1"
             result1 = invoke_agent(
                 "qa-expert",
                 "Review this function: def add(a, b): return a + b",
-                session_id="review-add-function-x7k9"  # Random suffix prevents collisions
+                session_id="review-add-function"
             )
+            # result1.session_id contains the full ID like "review-add-function-a3f2b1"
 
-            # Continue the SAME conversation (reuse session_id to maintain memory)
+            # Continue the SAME conversation using session_id from the previous result
             result2 = invoke_agent(
                 "qa-expert",
                 "Can you suggest edge cases for that function?",
-                session_id="review-add-function-x7k9"  # SAME session_id = conversation memory
+                session_id=result1.session_id  # Use the session_id from previous output!
             )
 
-            # Multiple INDEPENDENT reviews (unique session IDs with random suffixes)
+            # Multiple INDEPENDENT reviews (each gets unique hash suffix)
             auth_review = invoke_agent(
                 "code-reviewer",
                 "Review my authentication code",
-                session_id="auth-review-abc123"  # Random suffix for uniqueness
+                session_id="auth-review"  # -> "auth-review-<hash1>"
             )
+            # auth_review.session_id contains the full ID to continue this review
 
             payment_review = invoke_agent(
                 "code-reviewer",
                 "Review my payment processing code",
-                session_id="payment-review-def456"  # Different session = no shared context
+                session_id="payment-review"  # -> "payment-review-<hash2>"
             )
+            # payment_review.session_id contains a different full ID
         """
-        global _temp_agent_count
-
         from code_puppy.agents.agent_manager import load_agent
 
-        # Generate or use provided session_id (kebab-case format)
-        if session_id is None:
-            # Create a new session ID in kebab-case format
-            # Example: "qa-expert-session-1", "code-reviewer-session-2"
-            _temp_agent_count += 1
-            session_id = f"{agent_name}-session-{_temp_agent_count}"
-        else:
-            # Validate user-provided session_id
+        # Validate user-provided session_id if given
+        if session_id is not None:
             try:
                 _validate_session_id(session_id)
             except ValueError as e:
@@ -360,28 +389,44 @@ def register_invoke_agent(agent):
         # Generate a group ID for this tool execution
         group_id = generate_group_id("invoke_agent", agent_name)
 
-        emit_info(
-            f"\n[bold white on blue] INVOKE AGENT [/bold white on blue] {agent_name} (session: {session_id})",
-            message_group=group_id,
-        )
-        emit_divider(message_group=group_id)
-        emit_system_message(f"Prompt: {prompt}", message_group=group_id)
-
-        # Retrieve existing message history from filesystem for this session, if any
-        message_history = _load_session_history(session_id)
-        is_new_session = len(message_history) == 0
-
-        if message_history:
-            emit_system_message(
-                f"Continuing conversation from session {session_id} ({len(message_history)} messages)",
-                message_group=group_id,
-            )
+        # Check if this is an existing session or a new one
+        # For user-provided session_id, check if it exists
+        # For None, we'll generate a new one below
+        if session_id is not None:
+            message_history = _load_session_history(session_id)
+            is_new_session = len(message_history) == 0
         else:
-            emit_system_message(
-                f"Starting new session {session_id}",
-                message_group=group_id,
+            message_history = []
+            is_new_session = True
+
+        # Generate or finalize session_id
+        if session_id is None:
+            # Auto-generate a session ID with hash suffix for uniqueness
+            # Example: "qa-expert-session-a3f2b1"
+            hash_suffix = _generate_session_hash_suffix()
+            session_id = f"{agent_name}-session-{hash_suffix}"
+        elif is_new_session:
+            # User provided a base name for a NEW session - append hash suffix
+            # Example: "review-auth" -> "review-auth-a3f2b1"
+            hash_suffix = _generate_session_hash_suffix()
+            session_id = f"{session_id}-{hash_suffix}"
+        # else: continuing existing session, use session_id as-is
+
+        # Emit structured invocation message via MessageBus
+        bus = get_message_bus()
+        bus.emit(
+            SubAgentInvocationMessage(
+                agent_name=agent_name,
+                session_id=session_id,
+                prompt=prompt,
+                is_new_session=is_new_session,
+                message_count=len(message_history),
             )
-        emit_divider(message_group=group_id)
+        )
+
+        # Save current session context and set the new one for this sub-agent
+        previous_session_id = get_session_context()
+        set_session_context(session_id)
 
         try:
             # Load the specified agent config
@@ -400,6 +445,11 @@ def register_invoke_agent(agent):
             # Create a temporary agent instance to avoid interfering with current agent state
             instructions = agent_config.get_system_prompt()
 
+            # Add AGENTS.md content to subagents
+            puppy_rules = agent_config.load_puppy_rules()
+            if puppy_rules:
+                instructions += f"\n\n{puppy_rules}"
+
             # Apply prompt additions (like file permission handling) to temporary agents
             from code_puppy import callbacks
             from code_puppy.model_utils import prepare_prompt_for_model
@@ -418,7 +468,7 @@ def register_invoke_agent(agent):
             instructions = prepared.instructions
             prompt = prepared.user_prompt
 
-            subagent_name = f"temp-invoke-agent-{_temp_agent_count}"
+            subagent_name = f"temp-invoke-agent-{session_id}"
             model_settings = make_model_settings(model_name)
 
             temp_agent = Agent(
@@ -490,21 +540,32 @@ def register_invoke_agent(agent):
                 initial_prompt=prompt if is_new_session else None,
             )
 
-            emit_system_message(f"Response: {response}", message_group=group_id)
-            emit_system_message(
-                f"Session {session_id} saved to disk ({len(updated_history)} messages)",
-                message_group=group_id,
+            # Emit structured response message via MessageBus
+            bus.emit(
+                SubAgentResponseMessage(
+                    agent_name=agent_name,
+                    session_id=session_id,
+                    response=response,
+                    message_count=len(updated_history),
+                )
             )
-            emit_divider(message_group=group_id)
 
-            return AgentInvokeOutput(response=response, agent_name=agent_name)
+            return AgentInvokeOutput(
+                response=response, agent_name=agent_name, session_id=session_id
+            )
 
         except Exception:
             error_msg = f"Error invoking agent '{agent_name}': {traceback.format_exc()}"
             emit_error(error_msg, message_group=group_id)
-            emit_divider(message_group=group_id)
             return AgentInvokeOutput(
-                response=None, agent_name=agent_name, error=error_msg
+                response=None,
+                agent_name=agent_name,
+                session_id=session_id,
+                error=error_msg,
             )
 
+        finally:
+            # Restore the previous session context
+            set_session_context(previous_session_id)
+
     return invoke_agent