code-puppy 0.0.375__tar.gz → 0.0.377__tar.gz

{code_puppy-0.0.375 → code_puppy-0.0.377}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.375
+Version: 0.0.377
 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.375 → code_puppy-0.0.377}/code_puppy/agents/base_agent.py

@@ -47,7 +47,10 @@ from pydantic_ai.messages import (
 from rich.text import Text
 
 from code_puppy.agents.event_stream_handler import event_stream_handler
-from code_puppy.callbacks import on_agent_response_complete
+from code_puppy.callbacks import (
+    on_agent_run_end,
+    on_agent_run_start,
+)
 
 # Consolidated relative imports
 from code_puppy.config import (
@@ -404,35 +407,27 @@ class BaseAgent(ABC):
         total_tokens = 0
 
         # 1. Estimate tokens for system prompt / instructions
-        # For Claude Code models, the full system prompt is prepended to the first
-        # user message (already in message history), so we only count the short
-        # fixed instructions. For other models, count the full system prompt.
+        # Use prepare_prompt_for_model() to get the correct instructions for token counting.
+        # For models that prepend system prompt to user message (claude-code, antigravity),
+        # this returns the short fixed instructions. For other models, returns full prompt.
         try:
-            from code_puppy.model_utils import (
-                get_antigravity_instructions,
-                get_claude_code_instructions,
-                is_antigravity_model,
-                is_claude_code_model,
-            )
+            from code_puppy.model_utils import prepare_prompt_for_model
 
             model_name = (
                 self.get_model_name() if hasattr(self, "get_model_name") else ""
             )
-            if is_claude_code_model(model_name):
-                # For Claude Code models, only count the short fixed instructions
-                # The full system prompt is already in the message history
-                instructions = get_claude_code_instructions()
-                total_tokens += self.estimate_token_count(instructions)
-            elif is_antigravity_model(model_name):
-                # For Antigravity models, only count the short fixed instructions
-                # The full system prompt is already in the message history
-                instructions = get_antigravity_instructions()
-                total_tokens += self.estimate_token_count(instructions)
-            else:
-                # For other models, count the full system prompt
-                system_prompt = self.get_full_system_prompt()
-                if system_prompt:
-                    total_tokens += self.estimate_token_count(system_prompt)
+            system_prompt = self.get_full_system_prompt()
+
+            # Get the instructions that will be used (handles model-specific logic via hooks)
+            prepared = prepare_prompt_for_model(
+                model_name=model_name,
+                system_prompt=system_prompt,
+                user_prompt="",  # Empty - we just need the instructions
+                prepend_system_to_user=False,  # Don't modify prompt, just get instructions
+            )
+
+            if prepared.instructions:
+                total_tokens += self.estimate_token_count(prepared.instructions)
         except Exception:
             pass  # If we can't get system prompt, skip it
 
@@ -1590,21 +1585,25 @@ class BaseAgent(ABC):
         if output_type is not None:
             pydantic_agent = self._create_agent_with_output_type(output_type)
 
-        # Handle claude-code, chatgpt-codex, and antigravity models: prepend system prompt to first user message
-        from code_puppy.model_utils import (
-            is_antigravity_model,
-            is_claude_code_model,
-        )
+        # Handle model-specific prompt transformations via prepare_prompt_for_model()
+        # This uses the get_model_system_prompt hook, so plugins can register their own handlers
+        from code_puppy.model_utils import prepare_prompt_for_model
 
-        if is_claude_code_model(self.get_model_name()) or is_antigravity_model(
-            self.get_model_name()
-        ):
-            if len(self.get_message_history()) == 0:
-                system_prompt = self.get_full_system_prompt()
-                puppy_rules = self.load_puppy_rules()
-                if puppy_rules:
-                    system_prompt += f"\n{puppy_rules}"
-                prompt = system_prompt + "\n\n" + prompt
+        # Only prepend system prompt on first message (empty history)
+        should_prepend = len(self.get_message_history()) == 0
+        if should_prepend:
+            system_prompt = self.get_full_system_prompt()
+            puppy_rules = self.load_puppy_rules()
+            if puppy_rules:
+                system_prompt += f"\n{puppy_rules}"
+
+            prepared = prepare_prompt_for_model(
+                model_name=self.get_model_name(),
+                system_prompt=system_prompt,
+                user_prompt=prompt,
+                prepend_system_to_user=True,
+            )
+            prompt = prepared.user_prompt
 
         # Build combined prompt payload when attachments are provided.
         attachment_parts: List[Any] = []
@@ -1751,6 +1750,17 @@ class BaseAgent(ABC):
         # Create the task FIRST
         agent_task = asyncio.create_task(run_agent_task())
 
+        # Fire agent_run_start hook - plugins can use this to start background tasks
+        # (e.g., token refresh heartbeats for OAuth models)
+        try:
+            await on_agent_run_start(
+                agent_name=self.name,
+                model_name=self.get_model_name(),
+                session_id=group_id,
+            )
+        except Exception:
+            pass  # Don't fail agent run if hook fails
+
         # Import shell process status helper
 
         loop = asyncio.get_running_loop()
@@ -1832,39 +1842,53 @@ class BaseAgent(ABC):
                 except Exception:
                     pass  # Don't fail the run if cache update fails
 
-            # Trigger agent_response_complete callback for workflow orchestration
-            try:
-                # Extract the response text from the result
-                response_text = ""
-                if result is not None:
-                    if hasattr(result, "data"):
-                        response_text = str(result.data) if result.data else ""
-                    elif hasattr(result, "output"):
-                        response_text = str(result.output) if result.output else ""
-                    else:
-                        response_text = str(result)
-
-                # Fire the callback - don't await to avoid blocking return
-                # Use asyncio.create_task to run it in background
-                asyncio.create_task(
-                    on_agent_response_complete(
-                        agent_name=self.name,
-                        response_text=response_text,
-                        session_id=group_id,
-                        metadata={"model": self.get_model_name()},
-                    )
-                )
-            except Exception:
-                pass  # Don't fail the run if callback fails
+            # Extract response text for the callback
+            _run_response_text = ""
+            if result is not None:
+                if hasattr(result, "data"):
+                    _run_response_text = str(result.data) if result.data else ""
+                elif hasattr(result, "output"):
+                    _run_response_text = str(result.output) if result.output else ""
+                else:
+                    _run_response_text = str(result)
 
+            _run_success = True
+            _run_error = None
             return result
         except asyncio.CancelledError:
+            _run_success = False
+            _run_error = None  # Cancellation is not an error
+            _run_response_text = ""
             agent_task.cancel()
         except KeyboardInterrupt:
-            # Handle direct keyboard interrupt during await
+            _run_success = False
+            _run_error = None  # User interrupt is not an error
+            _run_response_text = ""
             if not agent_task.done():
                 agent_task.cancel()
+        except Exception as e:
+            _run_success = False
+            _run_error = e
+            _run_response_text = ""
+            raise
         finally:
+            # Fire agent_run_end hook - plugins can use this for:
+            # - Stopping background tasks (token refresh heartbeats)
+            # - Workflow orchestration (Ralph's autonomous loop)
+            # - Logging/analytics
+            try:
+                await on_agent_run_end(
+                    agent_name=self.name,
+                    model_name=self.get_model_name(),
+                    session_id=group_id,
+                    success=_run_success,
+                    error=_run_error,
+                    response_text=_run_response_text,
+                    metadata={"model": self.get_model_name()},
+                )
+            except Exception:
+                pass  # Don't fail cleanup if hook fails
+
             # Stop keyboard listener if it was started
             if key_listener_stop_event is not None:
                 key_listener_stop_event.set()

{code_puppy-0.0.375 → code_puppy-0.0.377}/code_puppy/callbacks.py

@@ -25,7 +25,8 @@ PhaseType = Literal[
     "register_agents",
     "register_model_type",
     "get_model_system_prompt",
-    "agent_response_complete",
+    "agent_run_start",
+    "agent_run_end",
 ]
 CallbackFunc = Callable[..., Any]
 
@@ -51,7 +52,8 @@ _callbacks: Dict[PhaseType, List[CallbackFunc]] = {
     "register_agents": [],
     "register_model_type": [],
     "get_model_system_prompt": [],
-    "agent_response_complete": [],
+    "agent_run_start": [],
+    "agent_run_end": [],
 }
 
 logger = logging.getLogger(__name__)
@@ -446,26 +448,72 @@ def on_get_model_system_prompt(
     )
 
 
-async def on_agent_response_complete(
+async def on_agent_run_start(
     agent_name: str,
-    response_text: str,
+    model_name: str,
     session_id: str | None = None,
+) -> List[Any]:
+    """Trigger callbacks when an agent run starts.
+
+    This fires at the beginning of run_with_mcp, before the agent task is created.
+    Useful for:
+    - Starting background tasks (like token refresh heartbeats)
+    - Logging/analytics
+    - Resource allocation
+
+    Args:
+        agent_name: Name of the agent starting
+        model_name: Name of the model being used
+        session_id: Optional session identifier
+
+    Returns:
+        List of results from registered callbacks.
+    """
+    return await _trigger_callbacks(
+        "agent_run_start", agent_name, model_name, session_id
+    )
+
+
+async def on_agent_run_end(
+    agent_name: str,
+    model_name: str,
+    session_id: str | None = None,
+    success: bool = True,
+    error: Exception | None = None,
+    response_text: str | None = None,
     metadata: dict | None = None,
 ) -> List[Any]:
-    """Trigger callbacks after an agent completes its full response.
+    """Trigger callbacks when an agent run ends.
+
+    This fires at the end of run_with_mcp, in the finally block.
+    Always fires regardless of success/failure/cancellation.
 
-    This fires after all tool calls are resolved and the agent has finished.
     Useful for:
+    - Stopping background tasks (like token refresh heartbeats)
     - Workflow orchestration (like Ralph's autonomous loop)
     - Logging/analytics
+    - Resource cleanup
     - Detecting completion signals in responses
 
     Args:
-        agent_name: Name of the agent that completed
-        response_text: The final text response from the agent
+        agent_name: Name of the agent that finished
+        model_name: Name of the model that was used
         session_id: Optional session identifier
+        success: Whether the run completed successfully
+        error: Exception if the run failed, None otherwise
+        response_text: The final text response from the agent (if successful)
         metadata: Optional dict with additional context (tokens used, etc.)
+
+    Returns:
+        List of results from registered callbacks.
     """
     return await _trigger_callbacks(
-        "agent_response_complete", agent_name, response_text, session_id, metadata
+        "agent_run_end",
+        agent_name,
+        model_name,
+        session_id,
+        success,
+        error,
+        response_text,
+        metadata,
     )

{code_puppy-0.0.375 → code_puppy-0.0.377}/code_puppy/command_line/model_settings_menu.py

@@ -106,6 +106,19 @@ SETTING_DEFINITIONS: Dict[str, Dict] = {
         "type": "boolean",
         "default": False,
     },
+    "thinking_enabled": {
+        "name": "Thinking Enabled",
+        "description": "Enable thinking mode for Gemini 3 Pro models. When enabled, the model will show its reasoning process.",
+        "type": "boolean",
+        "default": True,
+    },
+    "thinking_level": {
+        "name": "Thinking Level",
+        "description": "Controls the depth of thinking for Gemini 3 Pro models. Low = faster responses, High = more thorough reasoning.",
+        "type": "choice",
+        "choices": ["low", "high"],
+        "default": "low",
+    },
 }
 
 

{code_puppy-0.0.375 → code_puppy-0.0.377}/code_puppy/gemini_model.py

@@ -40,6 +40,10 @@ from pydantic_ai.usage import RequestUsage
 
 logger = logging.getLogger(__name__)
 
+# Bypass thought signature for Gemini when no pending signature is available
+# This allows function calls to work with thinking models
+BYPASS_THOUGHT_SIGNATURE = "context_engineering_is_the_way_to_go"
+
 
 def generate_tool_call_id() -> str:
     """Generate a unique tool call ID."""
@@ -410,27 +414,48 @@ class GeminiModel(Model):
         return system_instruction, contents
 
     def _map_model_response(self, m: ModelResponse) -> dict[str, Any] | None:
-        """Map a ModelResponse to Gemini content format."""
+        """Map a ModelResponse to Gemini content format.
+
+        For Gemini thinking models, we need to track thought signatures from
+        ThinkingParts and apply them to subsequent function_call parts.
+        """
         parts: list[dict[str, Any]] = []
+        pending_signature: str | None = None
 
         for item in m.parts:
             if isinstance(item, ToolCallPart):
-                parts.append(
-                    {
-                        "function_call": {
-                            "name": item.tool_name,
-                            "args": item.args_as_dict(),
-                            "id": item.tool_call_id,
-                        }
+                part_dict: dict[str, Any] = {
+                    "function_call": {
+                        "name": item.tool_name,
+                        "args": item.args_as_dict(),
+                        "id": item.tool_call_id,
                     }
+                }
+                # Gemini thinking models REQUIRE thoughtSignature on function calls
+                # Use pending signature from thinking or bypass signature
+                part_dict["thoughtSignature"] = (
+                    pending_signature
+                    if pending_signature is not None
+                    else BYPASS_THOUGHT_SIGNATURE
                 )
+                parts.append(part_dict)
             elif isinstance(item, TextPart):
-                parts.append({"text": item.content})
+                part_dict = {"text": item.content}
+                # Apply pending signature to text parts too if present
+                if pending_signature is not None:
+                    part_dict["thoughtSignature"] = pending_signature
+                    pending_signature = None
+                parts.append(part_dict)
             elif isinstance(item, ThinkingPart):
                 if item.content:
-                    part_dict: dict[str, Any] = {"text": item.content, "thought": True}
+                    part_dict = {"text": item.content, "thought": True}
                     if item.signature:
                         part_dict["thoughtSignature"] = item.signature
+                        # Store signature for subsequent parts
+                        pending_signature = item.signature
+                    else:
+                        # No signature on thinking part, use bypass
+                        pending_signature = BYPASS_THOUGHT_SIGNATURE
                     parts.append(part_dict)
 
         if not parts:
@@ -462,18 +487,34 @@ class GeminiModel(Model):
         config: dict[str, Any] = {}
 
         if model_settings:
-            if (
-                hasattr(model_settings, "temperature")
-                and model_settings.temperature is not None
-            ):
-                config["temperature"] = model_settings.temperature
-            if hasattr(model_settings, "top_p") and model_settings.top_p is not None:
-                config["topP"] = model_settings.top_p
-            if (
-                hasattr(model_settings, "max_tokens")
-                and model_settings.max_tokens is not None
-            ):
-                config["maxOutputTokens"] = model_settings.max_tokens
+            # ModelSettings is a TypedDict, so use .get() for all access
+            temperature = model_settings.get("temperature")
+            if temperature is not None:
+                config["temperature"] = temperature
+
+            top_p = model_settings.get("top_p")
+            if top_p is not None:
+                config["topP"] = top_p
+
+            max_tokens = model_settings.get("max_tokens")
+            if max_tokens is not None:
+                config["maxOutputTokens"] = max_tokens
+
+            # Handle Gemini 3 Pro thinking settings
+            thinking_enabled = model_settings.get("thinking_enabled")
+            thinking_level = model_settings.get("thinking_level")
+
+            # Build thinkingConfig if thinking settings are present
+            if thinking_enabled is False:
+                # Disable thinking by not including thinkingConfig
+                pass
+            elif thinking_level is not None:
+                # Gemini 3 Pro uses thinkingLevel with values "low" or "high"
+                # includeThoughts=True is required to surface the thinking in the response
+                config["thinkingConfig"] = {
+                    "thinkingLevel": thinking_level,
+                    "includeThoughts": True,
+                }
 
         return config
 

{code_puppy-0.0.375 → code_puppy-0.0.377}/code_puppy/model_factory.py

@@ -74,6 +74,7 @@ def make_model_settings(
         get_effective_model_settings,
         get_openai_reasoning_effort,
         get_openai_verbosity,
+        model_supports_setting,
     )
 
     model_settings_dict: dict = {}
@@ -131,6 +132,18 @@ def make_model_settings(
             }
         model_settings = AnthropicModelSettings(**model_settings_dict)
 
+    # Handle Gemini thinking models (Gemini-3)
+    # Check if model supports thinking settings and apply defaults
+    if model_supports_setting(model_name, "thinking_level"):
+        # Apply defaults if not explicitly set by user
+        # Default: thinking_enabled=True, thinking_level="low"
+        if "thinking_enabled" not in model_settings_dict:
+            model_settings_dict["thinking_enabled"] = True
+        if "thinking_level" not in model_settings_dict:
+            model_settings_dict["thinking_level"] = "low"
+        # Recreate settings with Gemini thinking config
+        model_settings = ModelSettings(**model_settings_dict)
+
     return model_settings
 
 

{code_puppy-0.0.375 → code_puppy-0.0.377}/code_puppy/models.json

@@ -33,13 +33,13 @@
     "type": "gemini",
     "name": "gemini-3-pro-preview",
     "context_length": 200000,
-    "supported_settings": ["temperature", "top_p"]
+    "supported_settings": ["temperature", "top_p", "thinking_enabled", "thinking_level"]
   },
   "Gemini-3-Long-Context": {
     "type": "gemini",
     "name": "gemini-3-pro-preview",
     "context_length": 1000000,
-    "supported_settings": ["temperature", "top_p"]
+    "supported_settings": ["temperature", "top_p", "thinking_enabled", "thinking_level"]
   },
   "gpt-5.1": {
     "type": "openai",

code_puppy-0.0.377/code_puppy/plugins/claude_code_oauth/__init__.py

@@ -0,0 +1,25 @@
+"""
+Claude Code OAuth Plugin for Code Puppy
+
+This plugin provides OAuth authentication for Claude Code and automatically
+adds available models to the extra_models.json configuration.
+
+The plugin also includes a token refresh heartbeat for maintaining fresh
+tokens during long-running agentic operations.
+"""
+
+from .token_refresh_heartbeat import (
+    TokenRefreshHeartbeat,
+    force_token_refresh,
+    get_current_heartbeat,
+    is_heartbeat_running,
+    token_refresh_heartbeat_context,
+)
+
+__all__ = [
+    "TokenRefreshHeartbeat",
+    "token_refresh_heartbeat_context",
+    "is_heartbeat_running",
+    "get_current_heartbeat",
+    "force_token_refresh",
+]

{code_puppy-0.0.375 → code_puppy-0.0.377}/code_puppy/plugins/claude_code_oauth/register_callbacks.py

@@ -363,6 +363,78 @@ def _register_model_types() -> List[Dict[str, Any]]:
     return [{"type": "claude_code", "handler": _create_claude_code_model}]
 
 
+# Global storage for the token refresh heartbeat
+# Using a dict to allow multiple concurrent agent runs (keyed by session_id)
+_active_heartbeats: Dict[str, Any] = {}
+
+
+async def _on_agent_run_start(
+    agent_name: str,
+    model_name: str,
+    session_id: Optional[str] = None,
+) -> None:
+    """Start token refresh heartbeat for Claude Code OAuth models.
+
+    This callback is triggered when an agent run starts. If the model is a
+    Claude Code OAuth model, we start a background heartbeat to keep the
+    token fresh during long-running operations.
+    """
+    # Only start heartbeat for Claude Code models
+    if not model_name.startswith("claude-code"):
+        return
+
+    try:
+        from .token_refresh_heartbeat import TokenRefreshHeartbeat
+
+        heartbeat = TokenRefreshHeartbeat()
+        await heartbeat.start()
+
+        # Store heartbeat for cleanup, keyed by session_id
+        key = session_id or "default"
+        _active_heartbeats[key] = heartbeat
+        logger.debug(
+            "Started token refresh heartbeat for session %s (model: %s)",
+            key,
+            model_name,
+        )
+    except ImportError:
+        logger.debug("Token refresh heartbeat module not available")
+    except Exception as exc:
+        logger.debug("Failed to start token refresh heartbeat: %s", exc)
+
+
+async def _on_agent_run_end(
+    agent_name: str,
+    model_name: str,
+    session_id: Optional[str] = None,
+    success: bool = True,
+    error: Optional[Exception] = None,
+    response_text: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Stop token refresh heartbeat when agent run ends.
+
+    This callback is triggered when an agent run completes (success or failure).
+    We stop any heartbeat that was started for this session.
+    """
+    # We don't use response_text or metadata, just cleanup the heartbeat
+    key = session_id or "default"
+    heartbeat = _active_heartbeats.pop(key, None)
+
+    if heartbeat is not None:
+        try:
+            await heartbeat.stop()
+            logger.debug(
+                "Stopped token refresh heartbeat for session %s (refreshed %d times)",
+                key,
+                heartbeat.refresh_count,
+            )
+        except Exception as exc:
+            logger.debug("Error stopping token refresh heartbeat: %s", exc)
+
+
 register_callback("custom_command_help", _custom_help)
 register_callback("custom_command", _handle_custom_command)
 register_callback("register_model_type", _register_model_types)
+register_callback("agent_run_start", _on_agent_run_start)
+register_callback("agent_run_end", _on_agent_run_end)