code-puppy 0.0.325__py3-none-any.whl → 0.0.336__py3-none-any.whl

code_puppy/plugins/antigravity_oauth/utils.py

@@ -0,0 +1,169 @@
+"""Utility helpers for the Antigravity OAuth plugin."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from .config import (
+    ANTIGRAVITY_OAUTH_CONFIG,
+    get_antigravity_models_path,
+    get_token_storage_path,
+)
+from .constants import ANTIGRAVITY_ENDPOINT, ANTIGRAVITY_HEADERS, ANTIGRAVITY_MODELS
+
+logger = logging.getLogger(__name__)
+
+
+def load_stored_tokens() -> Optional[Dict[str, Any]]:
+    """Load stored OAuth tokens from disk."""
+    try:
+        token_path = get_token_storage_path()
+        if token_path.exists():
+            with open(token_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+    except Exception as e:
+        logger.error("Failed to load tokens: %s", e)
+    return None
+
+
+def save_tokens(tokens: Dict[str, Any]) -> bool:
+    """Save OAuth tokens to disk."""
+    try:
+        token_path = get_token_storage_path()
+        with open(token_path, "w", encoding="utf-8") as f:
+            json.dump(tokens, f, indent=2)
+        token_path.chmod(0o600)
+        return True
+    except Exception as e:
+        logger.error("Failed to save tokens: %s", e)
+        return False
+
+
+def load_antigravity_models() -> Dict[str, Any]:
+    """Load configured Antigravity models from disk."""
+    try:
+        models_path = get_antigravity_models_path()
+        if models_path.exists():
+            with open(models_path, "r", encoding="utf-8") as f:
+                return json.load(f)
+    except Exception as e:
+        logger.error("Failed to load Antigravity models: %s", e)
+    return {}
+
+
+def save_antigravity_models(models: Dict[str, Any]) -> bool:
+    """Save Antigravity models configuration to disk."""
+    try:
+        models_path = get_antigravity_models_path()
+        with open(models_path, "w", encoding="utf-8") as f:
+            json.dump(models, f, indent=2)
+        return True
+    except Exception as e:
+        logger.error("Failed to save Antigravity models: %s", e)
+        return False
+
+
+def add_models_to_config(access_token: str, project_id: str = "") -> bool:
+    """Add all available Antigravity models to the configuration."""
+    try:
+        models_config: Dict[str, Any] = {}
+        prefix = ANTIGRAVITY_OAUTH_CONFIG["prefix"]
+
+        for model_id, model_info in ANTIGRAVITY_MODELS.items():
+            prefixed_name = f"{prefix}{model_id}"
+
+            # Build custom headers
+            headers = dict(ANTIGRAVITY_HEADERS)
+
+            # Use custom_gemini type with Antigravity transport
+            models_config[prefixed_name] = {
+                "type": "custom_gemini",
+                "name": model_id,
+                "custom_endpoint": {
+                    "url": ANTIGRAVITY_ENDPOINT,
+                    "api_key": access_token,
+                    "headers": headers,
+                },
+                "project_id": project_id,
+                "context_length": model_info.get("context_length", 200000),
+                "family": model_info.get("family", "other"),
+                "oauth_source": "antigravity-plugin",
+                "antigravity": True,  # Flag to use Antigravity transport
+            }
+
+            # Add thinking budget if present
+            if model_info.get("thinking_budget"):
+                models_config[prefixed_name]["thinking_budget"] = model_info[
+                    "thinking_budget"
+                ]
+
+        if save_antigravity_models(models_config):
+            logger.info("Added %d Antigravity models", len(models_config))
+            return True
+
+    except Exception as e:
+        logger.error("Error adding models to config: %s", e)
+    return False
+
+
+def remove_antigravity_models() -> int:
+    """Remove all Antigravity models from configuration."""
+    try:
+        models = load_antigravity_models()
+        to_remove = [
+            name
+            for name, config in models.items()
+            if config.get("oauth_source") == "antigravity-plugin"
+        ]
+
+        if not to_remove:
+            return 0
+
+        for model_name in to_remove:
+            models.pop(model_name, None)
+
+        if save_antigravity_models(models):
+            return len(to_remove)
+    except Exception as e:
+        logger.error("Error removing Antigravity models: %s", e)
+    return 0
+
+
+def get_model_families_summary() -> Dict[str, List[str]]:
+    """Get a summary of available models by family."""
+    families: Dict[str, List[str]] = {
+        "gemini": [],
+        "claude": [],
+        "other": [],
+    }
+
+    for model_id, info in ANTIGRAVITY_MODELS.items():
+        family = info.get("family", "other")
+        if family in families:
+            families[family].append(model_id)
+
+    return families
+
+
+def reload_current_agent() -> None:
+    """Reload the current agent so new auth tokens are picked up immediately."""
+    try:
+        from code_puppy.agents import get_current_agent
+
+        current_agent = get_current_agent()
+        if current_agent is None:
+            logger.debug("No current agent to reload")
+            return
+
+        if hasattr(current_agent, "refresh_config"):
+            try:
+                current_agent.refresh_config()
+            except Exception:
+                pass
+
+        current_agent.reload_code_generation_agent()
+        logger.info("Active agent reloaded with new authentication")
+    except Exception as e:
+        logger.warning("Agent reload failed: %s", e)

code_puppy/plugins/chatgpt_oauth/register_callbacks.py

@@ -6,6 +6,7 @@ import os
 from typing import List, Optional, Tuple
 
 from code_puppy.callbacks import register_callback
+from code_puppy.config import set_model_name
 from code_puppy.messaging import emit_info, emit_success, emit_warning
 
 from .config import CHATGPT_OAUTH_CONFIG, get_token_storage_path
@@ -75,6 +76,7 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
 
     if name == "chatgpt-auth":
         run_oauth_flow()
+        set_model_name("chatgpt-gpt-5.2-codex")
         return True
 
     if name == "chatgpt-status":

code_puppy/plugins/claude_code_oauth/register_callbacks.py

@@ -12,6 +12,7 @@ from typing import Any, Dict, List, Optional, Tuple
 from urllib.parse import parse_qs, urlparse
 
 from code_puppy.callbacks import register_callback
+from code_puppy.config import set_model_name
 from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
 
 from ..oauth_puppy_html import oauth_failure_html, oauth_success_html
@@ -260,6 +261,7 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
                 "Existing Claude Code tokens found. Continuing will overwrite them."
             )
         _perform_authentication()
+        set_model_name("claude-code-claude-opus-4-5-20251101")
         return True
 
     if name == "claude-code-status":

code_puppy/reopenable_async_client.py

@@ -54,13 +54,15 @@ class ReopenableAsyncClient:
             if self._stream_context:
                 return await self._stream_context.__aexit__(exc_type, exc_val, exc_tb)
 
-    def __init__(self, **kwargs):
+    def __init__(self, client_class=None, **kwargs):
         """
         Initialize the ReopenableAsyncClient.
 
         Args:
-            **kwargs: All arguments that would be passed to httpx.AsyncClient()
+            client_class: Class to use for creating the internal client (defaults to httpx.AsyncClient)
+            **kwargs: All arguments that would be passed to the client constructor
         """
+        self._client_class = client_class or httpx.AsyncClient
         self._client_kwargs = kwargs.copy()
         self._client: Optional[httpx.AsyncClient] = None
         self._is_closed = True
@@ -70,7 +72,7 @@ class ReopenableAsyncClient:
         Ensure the underlying client is open and ready to use.
 
         Returns:
-            The active httpx.AsyncClient instance
+            The active client instance
 
         Raises:
             RuntimeError: If client cannot be opened
@@ -80,12 +82,12 @@ class ReopenableAsyncClient:
         return self._client
 
     async def _create_client(self) -> None:
-        """Create a new httpx.AsyncClient with the stored configuration."""
+        """Create a new client with the stored configuration."""
         if self._client is not None and not self._is_closed:
             # Close existing client first
             await self._client.aclose()
 
-        self._client = httpx.AsyncClient(**self._client_kwargs)
+        self._client = self._client_class(**self._client_kwargs)
         self._is_closed = False
 
     async def reopen(self) -> None:
@@ -171,14 +173,12 @@ class ReopenableAsyncClient:
         """
         if self._client is None or self._is_closed:
             # Create a temporary client just for building the request
-            temp_client = httpx.AsyncClient(**self._client_kwargs)
+            temp_client = self._client_class(**self._client_kwargs)
             try:
                 request = temp_client.build_request(method, url, **kwargs)
                 return request
             finally:
                 # Clean up the temporary client synchronously if possible
-                # Note: This might leave a connection open, but it's better than
-                # making this method async just for building requests
                 pass
         return self._client.build_request(method, url, **kwargs)
 

code_puppy/terminal_utils.py

@@ -6,6 +6,10 @@ Handles Windows console mode resets and Unix terminal sanity restoration.
 import platform
 import subprocess
 import sys
+from typing import Callable, Optional
+
+# Store the original console ctrl handler so we can restore it if needed
+_original_ctrl_handler: Optional[Callable] = None
 
 
 def reset_windows_terminal_ansi() -> None:
@@ -86,17 +90,36 @@ def reset_windows_console_mode() -> None:
         pass  # Silently ignore errors - best effort reset
 
 
+def flush_windows_keyboard_buffer() -> None:
+    """Flush the Windows keyboard buffer.
+
+    Clears any pending keyboard input that could interfere with
+    subsequent input operations after an interrupt.
+    """
+    if platform.system() != "Windows":
+        return
+
+    try:
+        import msvcrt
+
+        while msvcrt.kbhit():
+            msvcrt.getch()
+    except Exception:
+        pass  # Silently ignore errors - best effort flush
+
+
 def reset_windows_terminal_full() -> None:
-    """Perform a full Windows terminal reset (ANSI + console mode).
+    """Perform a full Windows terminal reset (ANSI + console mode + keyboard buffer).
 
-    Combines both ANSI reset and console mode reset for complete
-    terminal state restoration after interrupts.
+    Combines ANSI reset, console mode reset, and keyboard buffer flush
+    for complete terminal state restoration after interrupts.
     """
     if platform.system() != "Windows":
         return
 
     reset_windows_terminal_ansi()
     reset_windows_console_mode()
+    flush_windows_keyboard_buffer()
 
 
 def reset_unix_terminal() -> None:
@@ -124,3 +147,145 @@ def reset_terminal() -> None:
         reset_windows_terminal_full()
     else:
         reset_unix_terminal()
+
+
+def disable_windows_ctrl_c() -> bool:
+    """Disable Ctrl+C processing at the Windows console input level.
+
+    This removes ENABLE_PROCESSED_INPUT from stdin, which prevents
+    Ctrl+C from being interpreted as a signal at all. Instead, it
+    becomes just a regular character (^C) that gets ignored.
+
+    This is more reliable than SetConsoleCtrlHandler because it
+    prevents Ctrl+C from being processed before it reaches any handler.
+
+    Returns:
+        True if successfully disabled, False otherwise.
+    """
+    global _original_ctrl_handler
+
+    if platform.system() != "Windows":
+        return False
+
+    try:
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+
+        # Get stdin handle
+        STD_INPUT_HANDLE = -10
+        stdin_handle = kernel32.GetStdHandle(STD_INPUT_HANDLE)
+
+        # Get current console mode
+        mode = ctypes.c_ulong()
+        if not kernel32.GetConsoleMode(stdin_handle, ctypes.byref(mode)):
+            return False
+
+        # Save original mode for potential restoration
+        _original_ctrl_handler = mode.value
+
+        # Console mode flags
+        ENABLE_PROCESSED_INPUT = 0x0001  # This makes Ctrl+C generate signals
+
+        # Remove ENABLE_PROCESSED_INPUT to disable Ctrl+C signal generation
+        new_mode = mode.value & ~ENABLE_PROCESSED_INPUT
+
+        if kernel32.SetConsoleMode(stdin_handle, new_mode):
+            return True
+        return False
+
+    except Exception:
+        return False
+
+
+def enable_windows_ctrl_c() -> bool:
+    """Re-enable Ctrl+C at the Windows console level.
+
+    Restores the original console mode saved by disable_windows_ctrl_c().
+
+    Returns:
+        True if successfully re-enabled, False otherwise.
+    """
+    global _original_ctrl_handler
+
+    if platform.system() != "Windows":
+        return False
+
+    if _original_ctrl_handler is None:
+        return True  # Nothing to restore
+
+    try:
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+
+        # Get stdin handle
+        STD_INPUT_HANDLE = -10
+        stdin_handle = kernel32.GetStdHandle(STD_INPUT_HANDLE)
+
+        # Restore original mode
+        if kernel32.SetConsoleMode(stdin_handle, _original_ctrl_handler):
+            _original_ctrl_handler = None
+            return True
+        return False
+
+    except Exception:
+        return False
+
+
+# Flag to track if we should keep Ctrl+C disabled
+_keep_ctrl_c_disabled: bool = False
+
+
+def set_keep_ctrl_c_disabled(value: bool) -> None:
+    """Set whether Ctrl+C should be kept disabled.
+
+    When True, ensure_ctrl_c_disabled() will re-disable Ctrl+C
+    even if something else (like prompt_toolkit) re-enables it.
+    """
+    global _keep_ctrl_c_disabled
+    _keep_ctrl_c_disabled = value
+
+
+def ensure_ctrl_c_disabled() -> bool:
+    """Ensure Ctrl+C is disabled if it should be.
+
+    Call this after operations that might restore console mode
+    (like prompt_toolkit input).
+
+    Returns:
+        True if Ctrl+C is now disabled (or wasn't needed), False on error.
+    """
+    if not _keep_ctrl_c_disabled:
+        return True
+
+    if platform.system() != "Windows":
+        return True
+
+    try:
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+
+        # Get stdin handle
+        STD_INPUT_HANDLE = -10
+        stdin_handle = kernel32.GetStdHandle(STD_INPUT_HANDLE)
+
+        # Get current console mode
+        mode = ctypes.c_ulong()
+        if not kernel32.GetConsoleMode(stdin_handle, ctypes.byref(mode)):
+            return False
+
+        # Console mode flags
+        ENABLE_PROCESSED_INPUT = 0x0001
+
+        # Check if Ctrl+C processing is enabled
+        if mode.value & ENABLE_PROCESSED_INPUT:
+            # Disable it
+            new_mode = mode.value & ~ENABLE_PROCESSED_INPUT
+            return bool(kernel32.SetConsoleMode(stdin_handle, new_mode))
+
+        return True  # Already disabled
+
+    except Exception:
+        return False

code_puppy/tools/command_runner.py

@@ -9,7 +9,7 @@ import threading
 import time
 import traceback
 from contextlib import contextmanager
-from typing import Callable, Literal, Optional, Set
+from typing import Callable, List, Literal, Optional, Set
 
 from pydantic import BaseModel
 from pydantic_ai import RunContext
@@ -192,11 +192,6 @@ def kill_all_running_shell_processes() -> int:
     """Kill all currently tracked running shell processes and stop reader threads.
 
     Returns the number of processes signaled.
-
-    Implementation notes:
-    - Atomically snapshot and clear the registry to prevent race conditions
-    - Deduplicate by PID to ensure each process is killed at most once
-    - Let exceptions from _kill_process_group propagate (tests expect this)
     """
     global _READER_STOP_EVENT
 
@@ -204,52 +199,30 @@ def kill_all_running_shell_processes() -> int:
     if _READER_STOP_EVENT:
         _READER_STOP_EVENT.set()
 
-    # Atomically take snapshot and clear registry
-    # This prevents other threads from seeing/processing the same processes
+    procs: list[subprocess.Popen]
     with _RUNNING_PROCESSES_LOCK:
-        procs_snapshot = list(_RUNNING_PROCESSES)
-        _RUNNING_PROCESSES.clear()
-
-    # Deduplicate by pid to ensure at-most-one kill per process
-    seen_pids: set = set()
-    killed_count = 0
-
-    for proc in procs_snapshot:
-        if proc is None:
-            continue
-
-        pid = getattr(proc, "pid", None)
-        key = pid if pid is not None else id(proc)
-
-        if key in seen_pids:
-            continue
-        seen_pids.add(key)
-
-        # Close pipes first to unblock readline()
+        procs = list(_RUNNING_PROCESSES)
+    count = 0
+    for p in procs:
         try:
-            if proc.stdout and not proc.stdout.closed:
-                proc.stdout.close()
-            if proc.stderr and not proc.stderr.closed:
-                proc.stderr.close()
-            if proc.stdin and not proc.stdin.closed:
-                proc.stdin.close()
-        except (OSError, ValueError):
-            pass
-
-        # Only attempt to kill processes that are still running
-        if proc.poll() is None:
-            # Let exceptions bubble up (tests expect this behavior)
-            _kill_process_group(proc)
-            killed_count += 1
-
-            # Track user-killed PIDs
-            if pid is not None:
-                try:
-                    _USER_KILLED_PROCESSES.add(pid)
-                except Exception:
-                    pass  # Non-fatal bookkeeping
+            # Close pipes first to unblock readline()
+            try:
+                if p.stdout and not p.stdout.closed:
+                    p.stdout.close()
+                if p.stderr and not p.stderr.closed:
+                    p.stderr.close()
+                if p.stdin and not p.stdin.closed:
+                    p.stdin.close()
+            except (OSError, ValueError):
+                pass
 
-    return killed_count
+            if p.poll() is None:
+                _kill_process_group(p)
+                count += 1
+                _USER_KILLED_PROCESSES.add(p.pid)
+        finally:
+            _unregister_process(p)
+    return count
 
 
 def get_running_shell_process_count() -> int:
@@ -1104,12 +1077,21 @@ class ReasoningOutput(BaseModel):
 
 
 def share_your_reasoning(
-    context: RunContext, reasoning: str, next_steps: str | None = None
+    context: RunContext, reasoning: str, next_steps: str | List[str] | None = None
 ) -> ReasoningOutput:
+    # Handle list of next steps by formatting them
+    formatted_next_steps = next_steps
+    if isinstance(next_steps, list):
+        formatted_next_steps = "\n".join(
+            [f"{i + 1}. {step}" for i, step in enumerate(next_steps)]
+        )
+
     # Emit structured AgentReasoningMessage for the UI
     reasoning_msg = AgentReasoningMessage(
         reasoning=reasoning,
-        next_steps=next_steps if next_steps and next_steps.strip() else None,
+        next_steps=formatted_next_steps
+        if formatted_next_steps and formatted_next_steps.strip()
+        else None,
     )
     get_message_bus().emit(reasoning_msg)
 
@@ -1197,7 +1179,9 @@ def register_agent_share_your_reasoning(agent):
 
     @agent.tool
     def agent_share_your_reasoning(
-        context: RunContext, reasoning: str = "", next_steps: str | None = None
+        context: RunContext,
+        reasoning: str = "",
+        next_steps: str | List[str] | None = None,
     ) -> ReasoningOutput:
         """Share the agent's current reasoning and planned next steps with the user.
 
@@ -1211,8 +1195,8 @@ def register_agent_share_your_reasoning(agent):
                 reasoning for the current situation. This should be clear,
                 comprehensive, and explain the 'why' behind decisions.
             next_steps: Planned upcoming actions or steps
-                the agent intends to take. Can be None if no specific next steps
-                are determined. Defaults to None.
+                the agent intends to take. Can be a string or a list of strings.
+                Can be None if no specific next steps are determined. Defaults to None.
 
         Returns:
             ReasoningOutput: A simple response object containing:
@@ -1223,6 +1207,10 @@ def register_agent_share_your_reasoning(agent):
             >>> next_steps = "First, I'll list the directory contents, then read key files"
             >>> result = agent_share_your_reasoning(ctx, reasoning, next_steps)
 
+            >>> # Using a list for next steps
+            >>> next_steps_list = ["List files", "Read README.md", "Run tests"]
+            >>> result = agent_share_your_reasoning(ctx, reasoning, next_steps_list)
+
         Best Practice:
             Use this tool frequently to maintain transparency. Call it:
             - Before starting complex operations