code-puppy 0.0.302__py3-none-any.whl → 0.0.335__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/config.py

@@ -9,7 +9,8 @@ CHATGPT_OAUTH_CONFIG: Dict[str, Any] = {
     "issuer": "https://auth.openai.com",
     "auth_url": "https://auth.openai.com/oauth/authorize",
     "token_url": "https://auth.openai.com/oauth/token",
-    "api_base_url": "https://api.openai.com",
+    # API endpoints - Codex uses chatgpt.com backend, not api.openai.com
+    "api_base_url": "https://chatgpt.com/backend-api/codex",
     # OAuth client configuration for Code Puppy
     "client_id": "app_EMoamEEZ73f0CkXaXp7hrann",
     "scope": "openid profile email offline_access",
@@ -24,6 +25,9 @@ CHATGPT_OAUTH_CONFIG: Dict[str, Any] = {
     "prefix": "chatgpt-",
     "default_context_length": 272000,
     "api_key_env_var": "CHATGPT_OAUTH_API_KEY",
+    # Codex CLI version info (for User-Agent header)
+    "client_version": "0.72.0",
+    "originator": "codex_cli_rs",
 }
 
 

code_puppy/plugins/chatgpt_oauth/oauth_flow.py

@@ -19,7 +19,6 @@ from .config import CHATGPT_OAUTH_CONFIG
 from .utils import (
     add_models_to_extra_config,
     assign_redirect_uri,
-    fetch_chatgpt_models,
     load_stored_tokens,
     parse_jwt_claims,
     prepare_oauth_context,
@@ -318,12 +317,12 @@ def run_oauth_flow() -> None:
         )
 
     if api_key:
-        emit_info("Fetching available ChatGPT models…")
-        models = fetch_chatgpt_models(api_key)
+        emit_info("Registering ChatGPT Codex models…")
+        from .utils import DEFAULT_CODEX_MODELS
+
+        models = DEFAULT_CODEX_MODELS
         if models:
-            if add_models_to_extra_config(models, api_key):
+            if add_models_to_extra_config(models):
                 emit_success(
                     "ChatGPT models registered. Use the `chatgpt-` prefix in /model."
                 )
-        else:
-            emit_warning("API key obtained, but model list could not be fetched.")

code_puppy/plugins/chatgpt_oauth/register_callbacks.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 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
@@ -74,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":
@@ -87,6 +90,5 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
     return None
 
 
-# Temporarily disabled - chatgpt-oauth plugin not working yet
-# register_callback("custom_command_help", _custom_help)
-# register_callback("custom_command", _handle_custom_command)
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)

code_puppy/plugins/chatgpt_oauth/test_plugin.py

@@ -235,25 +235,41 @@ def test_exchange_code_for_tokens(mock_post):
 
 @patch("code_puppy.plugins.chatgpt_oauth.utils.requests.get")
 def test_fetch_chatgpt_models(mock_get):
-    """Test fetching models from OpenAI API."""
+    """Test fetching models from ChatGPT Codex API."""
     mock_response = MagicMock()
     mock_response.status_code = 200
+    # New response format uses "models" key with "slug" field
     mock_response.json.return_value = {
-        "data": [
-            {"id": "gpt-4o"},
-            {"id": "gpt-3.5-turbo"},
-            {"id": "whisper-1"},  # Should be filtered out
-            {"id": "o1-preview"},
+        "models": [
+            {"slug": "gpt-4o"},
+            {"slug": "gpt-3.5-turbo"},
+            {"slug": "o1-preview"},
+            {"slug": "codex-mini"},
         ]
     }
     mock_get.return_value = mock_response
 
-    models = utils.fetch_chatgpt_models("test_api_key")
+    models = utils.fetch_chatgpt_models("test_access_token", "test_account_id")
     assert models is not None
     assert "gpt-4o" in models
     assert "gpt-3.5-turbo" in models
     assert "o1-preview" in models
-    assert "whisper-1" not in models  # Should be filtered
+    assert "codex-mini" in models
+
+
+@patch("code_puppy.plugins.chatgpt_oauth.utils.requests.get")
+def test_fetch_chatgpt_models_fallback(mock_get):
+    """Test that fetch_chatgpt_models returns default list on API failure."""
+    mock_response = MagicMock()
+    mock_response.status_code = 404
+    mock_response.text = '{"detail":"Not Found"}'
+    mock_get.return_value = mock_response
+
+    models = utils.fetch_chatgpt_models("test_access_token", "test_account_id")
+    assert models is not None
+    # Should return default models
+    assert "gpt-5.2" in models
+    assert "gpt-4o" in models
 
 
 def test_add_models_to_chatgpt_config(tmp_path):
@@ -262,14 +278,13 @@ def test_add_models_to_chatgpt_config(tmp_path):
         config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
     ):
         models = ["gpt-4o", "gpt-3.5-turbo"]
-        api_key = "sk-test"
 
-        assert utils.add_models_to_extra_config(models, api_key)
+        assert utils.add_models_to_extra_config(models)
 
         loaded = utils.load_chatgpt_models()
         assert "chatgpt-gpt-4o" in loaded
         assert "chatgpt-gpt-3.5-turbo" in loaded
-        assert loaded["chatgpt-gpt-4o"]["type"] == "openai"
+        assert loaded["chatgpt-gpt-4o"]["type"] == "chatgpt_oauth"
         assert loaded["chatgpt-gpt-4o"]["name"] == "gpt-4o"
         assert loaded["chatgpt-gpt-4o"]["oauth_source"] == "chatgpt-oauth-plugin"
 

code_puppy/plugins/chatgpt_oauth/utils.py

@@ -149,6 +149,98 @@ def load_stored_tokens() -> Optional[Dict[str, Any]]:
     return None
 
 
+def get_valid_access_token() -> Optional[str]:
+    """Get a valid access token, refreshing if expired.
+
+    Returns:
+        Valid access token string, or None if not authenticated or refresh failed.
+    """
+    tokens = load_stored_tokens()
+    if not tokens:
+        logger.debug("No stored ChatGPT 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
+
+    # Check if token is expired by parsing JWT claims
+    claims = parse_jwt_claims(access_token)
+    if claims:
+        exp = claims.get("exp")
+        if exp and isinstance(exp, (int, float)):
+            # Add 30 second buffer before expiry
+            if time.time() > exp - 30:
+                logger.info("ChatGPT OAuth token expired, attempting refresh")
+                refreshed = refresh_access_token()
+                if refreshed:
+                    return refreshed
+                logger.warning("Token refresh failed")
+                return None
+
+    return access_token
+
+
+def refresh_access_token() -> Optional[str]:
+    """Refresh the access token using the refresh token.
+
+    Returns:
+        New access token if refresh succeeded, None otherwise.
+    """
+    tokens = load_stored_tokens()
+    if not tokens:
+        return None
+
+    refresh_token = tokens.get("refresh_token")
+    if not refresh_token:
+        logger.debug("No refresh_token available")
+        return None
+
+    payload = {
+        "grant_type": "refresh_token",
+        "refresh_token": refresh_token,
+        "client_id": CHATGPT_OAUTH_CONFIG["client_id"],
+    }
+
+    headers = {
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+
+    try:
+        response = requests.post(
+            CHATGPT_OAUTH_CONFIG["token_url"],
+            data=payload,
+            headers=headers,
+            timeout=30,
+        )
+
+        if response.status_code == 200:
+            new_tokens = response.json()
+            # Merge with existing tokens (preserve account_id, etc.)
+            tokens.update(
+                {
+                    "access_token": new_tokens.get("access_token"),
+                    "refresh_token": new_tokens.get("refresh_token", refresh_token),
+                    "id_token": new_tokens.get("id_token", tokens.get("id_token")),
+                    "last_refresh": datetime.datetime.now(datetime.timezone.utc)
+                    .isoformat()
+                    .replace("+00:00", "Z"),
+                }
+            )
+            if save_tokens(tokens):
+                logger.info("Successfully refreshed ChatGPT OAuth token")
+                return tokens["access_token"]
+        else:
+            logger.error(
+                "Token refresh failed: %s - %s", response.status_code, response.text
+            )
+    except Exception as exc:
+        logger.error("Token refresh error: %s", exc)
+
+    return None
+
+
 def save_tokens(tokens: Dict[str, Any]) -> bool:
     if tokens is None:
         raise TypeError("tokens cannot be None")
@@ -248,103 +340,126 @@ def exchange_code_for_tokens(
     return None
 
 
-def fetch_chatgpt_models(api_key: str) -> Optional[List[str]]:
-    """Fetch available models from OpenAI API.
+# Default models available via ChatGPT Codex API
+# These are the known models that work with ChatGPT OAuth tokens
+# Based on codex-rs CLI and shell-scripts/codex-call.sh
+DEFAULT_CODEX_MODELS = [
+    "gpt-5.2",
+    "gpt-5.2-codex",
+]
+
 
-    Makes a real HTTP GET request to OpenAI's models endpoint and filters
-    the results to include only GPT series models while preserving server order.
+def fetch_chatgpt_models(access_token: str, account_id: str) -> Optional[List[str]]:
+    """Fetch available models from ChatGPT Codex API.
+
+    Attempts to fetch models from the API, but falls back to a default list
+    of known Codex-compatible models if the API is unavailable.
 
     Args:
-        api_key: OpenAI API key for authentication
+        access_token: OAuth access token for authentication
+        account_id: ChatGPT account ID (required for the API)
 
     Returns:
-        List of filtered model IDs preserving server order, or None if request fails
+        List of model IDs, or default list if API fails
     """
-    # Build the models URL, ensuring it ends with /v1/models
-    base_url = CHATGPT_OAUTH_CONFIG["api_base_url"].rstrip("/")
-    models_url = f"{base_url}/v1/models"
+    import platform
 
-    # Blocklist of model IDs to exclude
-    blocklist = {"whisper-1"}
+    # Build the models URL with client version
+    client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.72.0")
+    base_url = CHATGPT_OAUTH_CONFIG["api_base_url"].rstrip("/")
+    models_url = f"{base_url}/models"
+
+    # Build User-Agent to match codex-rs CLI format
+    originator = CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs")
+    os_name = platform.system()
+    if os_name == "Darwin":
+        os_name = "Mac OS"
+    os_version = platform.release()
+    arch = platform.machine()
+    user_agent = (
+        f"{originator}/{client_version} ({os_name} {os_version}; {arch}) "
+        "Terminal_Codex_CLI"
+    )
 
     headers = {
-        "Authorization": f"Bearer {api_key}",
+        "Authorization": f"Bearer {access_token}",
+        "ChatGPT-Account-Id": account_id,
+        "User-Agent": user_agent,
+        "originator": originator,
+        "Accept": "application/json",
     }
 
-    try:
-        response = requests.get(models_url, headers=headers, timeout=30)
-
-        if response.status_code != 200:
-            logger.error(
-                "Failed to fetch models: HTTP %d - %s",
-                response.status_code,
-                response.text,
-            )
-            return None
-
-        # Parse JSON response
-        try:
-            data = response.json()
-            if "data" not in data or not isinstance(data["data"], list):
-                logger.error("Invalid response format: missing 'data' list")
-                return None
-        except (json.JSONDecodeError, ValueError) as exc:
-            logger.error("Failed to parse JSON response: %s", exc)
-            return None
-
-        # Filter models: start with "gpt-" or "o1-" and not in blocklist
-        filtered_models = []
-        seen_models = set()  # For deduplication while preserving order
-
-        for model in data["data"]:
-            # Skip None entries
-            if model is None:
-                continue
+    # Query params
+    params = {"client_version": client_version}
 
-            model_id = model.get("id")
-            if not model_id:
-                continue
-
-            # Skip if already seen (deduplication)
-            if model_id in seen_models:
-                continue
-
-            # Check if model starts with allowed prefixes and not in blocklist
-            if (
-                model_id.startswith("gpt-") or model_id.startswith("o1-")
-            ) and model_id not in blocklist:
-                filtered_models.append(model_id)
-                seen_models.add(model_id)
+    try:
+        response = requests.get(models_url, headers=headers, params=params, timeout=30)
 
-        return filtered_models
+        if response.status_code == 200:
+            # Parse JSON response
+            try:
+                data = response.json()
+                # The response has a "models" key with list of model objects
+                if "models" in data and isinstance(data["models"], list):
+                    models = []
+                    for model in data["models"]:
+                        if model is None:
+                            continue
+                        model_id = (
+                            model.get("slug") or model.get("id") or model.get("name")
+                        )
+                        if model_id:
+                            models.append(model_id)
+                    if models:
+                        return models
+            except (json.JSONDecodeError, ValueError) as exc:
+                logger.warning("Failed to parse models response: %s", exc)
+
+        # API didn't return valid models, use default list
+        logger.info(
+            "Models endpoint returned %d, using default model list",
+            response.status_code,
+        )
 
     except requests.exceptions.Timeout:
-        logger.error("Timeout while fetching models after 30 seconds")
-        return None
+        logger.warning("Timeout fetching models, using default list")
     except requests.exceptions.RequestException as exc:
-        logger.error("Network error while fetching models: %s", exc)
-        return None
+        logger.warning("Network error fetching models: %s, using default list", exc)
     except Exception as exc:
-        logger.error("Unexpected error while fetching models: %s", exc)
-        return None
+        logger.warning("Error fetching models: %s, using default list", exc)
 
+    # Return default models when API fails
+    logger.info("Using default Codex models: %s", DEFAULT_CODEX_MODELS)
+    return DEFAULT_CODEX_MODELS
 
-def add_models_to_extra_config(models: List[str], api_key: str) -> bool:
+
+def add_models_to_extra_config(models: List[str]) -> bool:
     """Add ChatGPT models to chatgpt_models.json configuration."""
     try:
         chatgpt_models = load_chatgpt_models()
         added = 0
         for model_name in models:
             prefixed = f"{CHATGPT_OAUTH_CONFIG['prefix']}{model_name}"
+
+            # Determine supported settings based on model type
+            # All GPT-5.x models support reasoning_effort and verbosity
+            supported_settings = ["reasoning_effort", "verbosity"]
+
+            # Only codex models support xhigh reasoning effort
+            # Regular gpt-5.2 is capped at "high"
+            is_codex = "codex" in model_name.lower()
+
             chatgpt_models[prefixed] = {
-                "type": "openai",
+                "type": "chatgpt_oauth",
                 "name": model_name,
                 "custom_endpoint": {
+                    # Codex API uses chatgpt.com/backend-api/codex, not api.openai.com
                     "url": CHATGPT_OAUTH_CONFIG["api_base_url"],
-                    "api_key": "${" + CHATGPT_OAUTH_CONFIG["api_key_env_var"] + "}",
                 },
                 "context_length": CHATGPT_OAUTH_CONFIG["default_context_length"],
                 "oauth_source": "chatgpt-oauth-plugin",
+                "supported_settings": supported_settings,
+                "supports_xhigh_reasoning": is_codex,
             }
             added += 1
         if save_chatgpt_models(chatgpt_models):

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
@@ -180,6 +181,31 @@ def _custom_help() -> List[Tuple[str, str]]:
     ]
 
 
+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
+
+        # JSON agents may need to refresh their config before reload
+        if hasattr(current_agent, "refresh_config"):
+            try:
+                current_agent.refresh_config()
+            except Exception:
+                # Non-fatal, continue to reload
+                pass
+
+        current_agent.reload_code_generation_agent()
+        emit_info("Active agent reloaded with new authentication")
+    except Exception as e:
+        emit_warning(f"Authentication succeeded but agent reload failed: {e}")
+        logger.exception("Failed to reload agent after authentication")
+
+
 def _perform_authentication() -> None:
     context = prepare_oauth_context()
     code = _await_callback(context)
@@ -219,6 +245,9 @@ def _perform_authentication() -> None:
             "Claude Code models added to your configuration. Use the `claude-code-` prefix!"
         )
 
+    # Reload the current agent so the new auth token is picked up immediately
+    _reload_current_agent()
+
 
 def _handle_custom_command(command: str, name: str) -> Optional[bool]:
     if not name:
@@ -232,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

@@ -368,6 +368,7 @@ def add_models_to_extra_config(models: List[str]) -> bool:
                     "temperature",
                     "extended_thinking",
                     "budget_tokens",
+                    "interleaved_thinking",
                 ],
             }
             added += 1