code-puppy 0.0.325__py3-none-any.whl → 0.0.341__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/plugins/claude_code_oauth/utils.py

@@ -21,6 +21,8 @@ from .config import (
     get_token_storage_path,
 )
 
+TOKEN_REFRESH_BUFFER_SECONDS = 60
+
 logger = logging.getLogger(__name__)
 
 
@@ -132,6 +134,124 @@ def load_stored_tokens() -> Optional[Dict[str, Any]]:
     return None
 
 
+def _calculate_expires_at(expires_in: Optional[float]) -> Optional[float]:
+    if expires_in is None:
+        return None
+    try:
+        return time.time() + float(expires_in)
+    except (TypeError, ValueError):
+        return None
+
+
+def is_token_expired(tokens: Dict[str, Any]) -> bool:
+    expires_at = tokens.get("expires_at")
+    if expires_at is None:
+        return False
+    try:
+        expires_at_value = float(expires_at)
+    except (TypeError, ValueError):
+        return False
+    return time.time() >= expires_at_value - TOKEN_REFRESH_BUFFER_SECONDS
+
+
+def update_claude_code_model_tokens(access_token: str) -> bool:
+    try:
+        claude_models = load_claude_models()
+        if not claude_models:
+            return False
+
+        updated = False
+        for config in claude_models.values():
+            if config.get("oauth_source") != "claude-code-plugin":
+                continue
+            custom_endpoint = config.get("custom_endpoint")
+            if not isinstance(custom_endpoint, dict):
+                continue
+            custom_endpoint["api_key"] = access_token
+            updated = True
+
+        if updated:
+            return save_claude_models(claude_models)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.error("Failed to update Claude model tokens: %s", exc)
+    return False
+
+
+def refresh_access_token(force: bool = False) -> Optional[str]:
+    tokens = load_stored_tokens()
+    if not tokens:
+        return None
+
+    if not force and not is_token_expired(tokens):
+        return tokens.get("access_token")
+
+    refresh_token = tokens.get("refresh_token")
+    if not refresh_token:
+        logger.debug("No refresh_token available")
+        return None
+
+    payload = {
+        "grant_type": "refresh_token",
+        "client_id": CLAUDE_CODE_OAUTH_CONFIG["client_id"],
+        "refresh_token": refresh_token,
+    }
+
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "anthropic-beta": "oauth-2025-04-20",
+    }
+
+    try:
+        response = requests.post(
+            CLAUDE_CODE_OAUTH_CONFIG["token_url"],
+            json=payload,
+            headers=headers,
+            timeout=30,
+        )
+        if response.status_code == 200:
+            new_tokens = response.json()
+            tokens["access_token"] = new_tokens.get("access_token")
+            tokens["refresh_token"] = new_tokens.get("refresh_token", refresh_token)
+            if "expires_in" in new_tokens:
+                tokens["expires_in"] = new_tokens["expires_in"]
+                tokens["expires_at"] = _calculate_expires_at(
+                    new_tokens.get("expires_in")
+                )
+            if save_tokens(tokens):
+                update_claude_code_model_tokens(tokens["access_token"])
+                return tokens["access_token"]
+        else:
+            logger.error(
+                "Token refresh failed: %s - %s", response.status_code, response.text
+            )
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.error("Token refresh error: %s", exc)
+    return None
+
+
+def get_valid_access_token() -> Optional[str]:
+    tokens = load_stored_tokens()
+    if not tokens:
+        logger.debug("No stored Claude Code OAuth tokens found")
+        return None
+
+    access_token = tokens.get("access_token")
+    if not access_token:
+        logger.debug("No access_token in stored tokens")
+        return None
+
+    if is_token_expired(tokens):
+        logger.info("Claude Code OAuth token expired, attempting refresh")
+        refreshed = refresh_access_token()
+        if refreshed:
+            return refreshed
+        logger.warning("Claude Code token refresh failed")
+        return None
+
+    return access_token
+
+
 def save_tokens(tokens: Dict[str, Any]) -> bool:
     try:
         token_path = get_token_storage_path()
@@ -243,7 +363,11 @@ def exchange_code_for_tokens(
         logger.info("Token exchange response: %s", response.status_code)
         logger.debug("Response body: %s", response.text)
         if response.status_code == 200:
-            return response.json()
+            token_data = response.json()
+            token_data["expires_at"] = _calculate_expires_at(
+                token_data.get("expires_in")
+            )
+            return token_data
         logger.error(
             "Token exchange failed: %s - %s",
             response.status_code,
@@ -341,12 +465,7 @@ def add_models_to_extra_config(models: List[str]) -> bool:
         # Start fresh - overwrite the file on every auth instead of loading existing
         claude_models = {}
         added = 0
-        tokens = load_stored_tokens()
-
-        # Handle case where tokens are None or empty
-        access_token = ""
-        if tokens and "access_token" in tokens:
-            access_token = tokens["access_token"]
+        access_token = get_valid_access_token() or ""
 
         for model_name in filtered_models:
             prefixed = f"{CLAUDE_CODE_OAUTH_CONFIG['prefix']}{model_name}"

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)