code-puppy 0.0.374__py3-none-any.whl → 0.0.376__py3-none-any.whl

code_puppy/plugins/claude_code_oauth/register_callbacks.py

@@ -1,5 +1,8 @@
 """
 Claude Code OAuth Plugin for Code Puppy.
+
+Provides OAuth authentication for Claude Code models and registers
+the 'claude_code' model type handler.
 """
 
 from __future__ import annotations
@@ -24,6 +27,7 @@ from .utils import (
     build_authorization_url,
     exchange_code_for_tokens,
     fetch_claude_code_models,
+    get_valid_access_token,
     load_claude_models_filtered,
     load_stored_tokens,
     prepare_oauth_context,
@@ -276,5 +280,161 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
     return None
 
 
+def _create_claude_code_model(model_name: str, model_config: Dict, config: Dict) -> Any:
+    """Create a Claude Code model instance.
+
+    This handler is registered via the 'register_model_type' callback to handle
+    models with type='claude_code'.
+    """
+    from anthropic import AsyncAnthropic
+    from pydantic_ai.models.anthropic import AnthropicModel
+    from pydantic_ai.providers.anthropic import AnthropicProvider
+
+    from code_puppy.claude_cache_client import (
+        ClaudeCacheAsyncClient,
+        patch_anthropic_client_messages,
+    )
+    from code_puppy.config import get_effective_model_settings
+    from code_puppy.http_utils import get_cert_bundle_path, get_http2
+    from code_puppy.model_factory import get_custom_config
+
+    url, headers, verify, api_key = get_custom_config(model_config)
+
+    # Refresh token if this is from the plugin
+    if model_config.get("oauth_source") == "claude-code-plugin":
+        refreshed_token = get_valid_access_token()
+        if refreshed_token:
+            api_key = refreshed_token
+            custom_endpoint = model_config.get("custom_endpoint")
+            if isinstance(custom_endpoint, dict):
+                custom_endpoint["api_key"] = refreshed_token
+
+    if not api_key:
+        emit_warning(
+            f"API key is not set for Claude Code endpoint; skipping model '{model_config.get('name')}'."
+        )
+        return None
+
+    # Check if interleaved thinking is enabled (defaults to True for OAuth models)
+    effective_settings = get_effective_model_settings(model_name)
+    interleaved_thinking = effective_settings.get("interleaved_thinking", True)
+
+    # Handle anthropic-beta header based on interleaved_thinking setting
+    if "anthropic-beta" in headers:
+        beta_parts = [p.strip() for p in headers["anthropic-beta"].split(",")]
+        if interleaved_thinking:
+            if "interleaved-thinking-2025-05-14" not in beta_parts:
+                beta_parts.append("interleaved-thinking-2025-05-14")
+        else:
+            beta_parts = [p for p in beta_parts if "interleaved-thinking" not in p]
+        headers["anthropic-beta"] = ",".join(beta_parts) if beta_parts else None
+        if headers.get("anthropic-beta") is None:
+            del headers["anthropic-beta"]
+    elif interleaved_thinking:
+        headers["anthropic-beta"] = "interleaved-thinking-2025-05-14"
+
+    # Use a dedicated client wrapper that injects cache_control on /v1/messages
+    if verify is None:
+        verify = get_cert_bundle_path()
+
+    http2_enabled = get_http2()
+
+    client = ClaudeCacheAsyncClient(
+        headers=headers,
+        verify=verify,
+        timeout=180,
+        http2=http2_enabled,
+    )
+
+    anthropic_client = AsyncAnthropic(
+        base_url=url,
+        http_client=client,
+        auth_token=api_key,
+    )
+    patch_anthropic_client_messages(anthropic_client)
+    anthropic_client.api_key = None
+    anthropic_client.auth_token = api_key
+    provider = AnthropicProvider(anthropic_client=anthropic_client)
+    return AnthropicModel(model_name=model_config["name"], provider=provider)
+
+
+def _register_model_types() -> List[Dict[str, Any]]:
+    """Register the claude_code model type handler."""
+    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)

code_puppy/plugins/claude_code_oauth/token_refresh_heartbeat.py

@@ -0,0 +1,242 @@
+"""Token refresh heartbeat for long-running Claude Code OAuth sessions.
+
+This module provides a background task that periodically checks and refreshes
+Claude Code OAuth tokens during long-running agentic operations. This ensures
+that tokens don't expire during extended streaming responses or tool processing.
+
+Usage:
+    async with token_refresh_heartbeat_context():
+        # Long running agent operation
+        await agent.run(...)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from contextlib import asynccontextmanager
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Heartbeat interval in seconds - check token every 2 minutes
+# This is frequent enough to catch expiring tokens before they cause issues
+# but not so frequent as to spam the token endpoint
+HEARTBEAT_INTERVAL_SECONDS = 120
+
+# Minimum time between refresh attempts to avoid hammering the endpoint
+MIN_REFRESH_INTERVAL_SECONDS = 60
+
+# Global tracking of last refresh time to coordinate across heartbeats
+_last_refresh_time: float = 0.0
+_heartbeat_lock = asyncio.Lock()
+
+
+class TokenRefreshHeartbeat:
+    """Background task that periodically refreshes Claude Code OAuth tokens.
+
+    This runs as an asyncio task during agent operations and checks if the
+    token needs refreshing at regular intervals.
+    """
+
+    def __init__(
+        self,
+        interval: float = HEARTBEAT_INTERVAL_SECONDS,
+        min_refresh_interval: float = MIN_REFRESH_INTERVAL_SECONDS,
+    ):
+        self._interval = interval
+        self._min_refresh_interval = min_refresh_interval
+        self._task: Optional[asyncio.Task] = None
+        self._stop_event = asyncio.Event()
+        self._refresh_count = 0
+
+    async def start(self) -> None:
+        """Start the heartbeat background task."""
+        if self._task is not None:
+            logger.debug("Heartbeat already running")
+            return
+
+        self._stop_event.clear()
+        self._task = asyncio.create_task(self._heartbeat_loop())
+        logger.debug("Token refresh heartbeat started")
+
+    async def stop(self) -> None:
+        """Stop the heartbeat background task."""
+        if self._task is None:
+            return
+
+        self._stop_event.set()
+        self._task.cancel()
+
+        try:
+            await self._task
+        except asyncio.CancelledError:
+            pass
+
+        self._task = None
+        logger.debug(
+            "Token refresh heartbeat stopped (refreshed %d times)",
+            self._refresh_count,
+        )
+
+    async def _heartbeat_loop(self) -> None:
+        """Main heartbeat loop that periodically checks token status."""
+        global _last_refresh_time
+
+        while not self._stop_event.is_set():
+            try:
+                # Wait for the interval or until stopped
+                try:
+                    await asyncio.wait_for(
+                        self._stop_event.wait(), timeout=self._interval
+                    )
+                    # If we got here, stop event was set
+                    break
+                except asyncio.TimeoutError:
+                    # Normal timeout - time to check token
+                    pass
+
+                # Check if we should attempt refresh
+                async with _heartbeat_lock:
+                    now = time.time()
+                    if now - _last_refresh_time < self._min_refresh_interval:
+                        logger.debug(
+                            "Skipping refresh - last refresh was %.1f seconds ago",
+                            now - _last_refresh_time,
+                        )
+                        continue
+
+                    # Attempt the refresh
+                    refreshed = await self._attempt_refresh()
+                    if refreshed:
+                        _last_refresh_time = now
+                        self._refresh_count += 1
+
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                logger.debug("Error in heartbeat loop: %s", exc)
+                # Continue running - don't let errors kill the heartbeat
+                await asyncio.sleep(5)  # Brief pause before retrying
+
+    async def _attempt_refresh(self) -> bool:
+        """Attempt to refresh the token if needed.
+
+        Returns:
+            True if a refresh was performed, False otherwise.
+        """
+        try:
+            # Import here to avoid circular imports
+            from .utils import (
+                is_token_expired,
+                load_stored_tokens,
+                refresh_access_token,
+            )
+
+            tokens = load_stored_tokens()
+            if not tokens:
+                logger.debug("No stored tokens found")
+                return False
+
+            if not is_token_expired(tokens):
+                logger.debug("Token not yet expired, skipping refresh")
+                return False
+
+            # Token is expiring soon, refresh it
+            logger.info("Heartbeat: Token expiring soon, refreshing proactively")
+            refreshed_token = refresh_access_token(force=False)
+
+            if refreshed_token:
+                logger.info("Heartbeat: Successfully refreshed token")
+                return True
+            else:
+                logger.warning("Heartbeat: Token refresh returned None")
+                return False
+
+        except Exception as exc:
+            logger.error("Heartbeat: Error during token refresh: %s", exc)
+            return False
+
+    @property
+    def refresh_count(self) -> int:
+        """Get the number of successful refreshes performed by this heartbeat."""
+        return self._refresh_count
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the heartbeat is currently running."""
+        return self._task is not None and not self._task.done()
+
+
+# Global heartbeat instance for the current session
+_current_heartbeat: Optional[TokenRefreshHeartbeat] = None
+
+
+@asynccontextmanager
+async def token_refresh_heartbeat_context(
+    interval: float = HEARTBEAT_INTERVAL_SECONDS,
+):
+    """Context manager that runs token refresh heartbeat during its scope.
+
+    Use this around long-running agent operations to ensure tokens stay fresh.
+
+    Args:
+        interval: Seconds between heartbeat checks. Default is 2 minutes.
+
+    Example:
+        async with token_refresh_heartbeat_context():
+            result = await agent.run(prompt)
+    """
+    global _current_heartbeat
+
+    heartbeat = TokenRefreshHeartbeat(interval=interval)
+
+    try:
+        await heartbeat.start()
+        _current_heartbeat = heartbeat
+        yield heartbeat
+    finally:
+        await heartbeat.stop()
+        _current_heartbeat = None
+
+
+def is_heartbeat_running() -> bool:
+    """Check if a token refresh heartbeat is currently active."""
+    return _current_heartbeat is not None and _current_heartbeat.is_running
+
+
+def get_current_heartbeat() -> Optional[TokenRefreshHeartbeat]:
+    """Get the currently running heartbeat instance, if any."""
+    return _current_heartbeat
+
+
+async def force_token_refresh() -> bool:
+    """Force an immediate token refresh.
+
+    This can be called from anywhere to trigger a token refresh,
+    regardless of whether a heartbeat is running.
+
+    Returns:
+        True if refresh was successful, False otherwise.
+    """
+    global _last_refresh_time
+
+    try:
+        from .utils import refresh_access_token
+
+        logger.info("Forcing token refresh")
+        refreshed_token = refresh_access_token(force=True)
+
+        if refreshed_token:
+            async with _heartbeat_lock:
+                _last_refresh_time = time.time()
+            logger.info("Force refresh successful")
+            return True
+        else:
+            logger.warning("Force refresh returned None")
+            return False
+
+    except Exception as exc:
+        logger.error("Force refresh error: %s", exc)
+        return False

code_puppy/plugins/ralph/__init__.py

@@ -0,0 +1,13 @@
+"""Ralph Plugin - Autonomous AI agent loop for completing PRDs.
+
+Based on Geoffrey Huntley's Ralph pattern: https://ghuntley.com/ralph/
+
+This plugin provides:
+- PRD Generator agent for creating detailed requirements
+- Ralph Converter agent for converting PRDs to JSON format
+- Ralph Orchestrator agent for autonomous execution
+- Tools for managing prd.json and progress.txt
+- /ralph commands for controlling the workflow
+"""
+
+__version__ = "0.1.0"