code-puppy 0.0.214__py3-none-any.whl → 0.0.366__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/__init__.py

@@ -0,0 +1,8 @@
+"""ChatGPT OAuth plugin package."""
+
+from __future__ import annotations
+
+from . import register_callbacks  # noqa: F401
+from .oauth_flow import run_oauth_flow
+
+__all__ = ["run_oauth_flow"]

code_puppy/plugins/chatgpt_oauth/config.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import Any, Dict
+
+from code_puppy import config
+
+# ChatGPT OAuth configuration based on OpenAI's Codex CLI flow
+CHATGPT_OAUTH_CONFIG: Dict[str, Any] = {
+    # OAuth endpoints from OpenAI auth service
+    "issuer": "https://auth.openai.com",
+    "auth_url": "https://auth.openai.com/oauth/authorize",
+    "token_url": "https://auth.openai.com/oauth/token",
+    # 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",
+    # Callback handling (we host a localhost callback to capture the redirect)
+    "redirect_host": "http://localhost",
+    "redirect_path": "auth/callback",
+    "required_port": 1455,
+    "callback_timeout": 120,
+    # Local configuration (uses XDG_DATA_HOME)
+    "token_storage": None,  # Set dynamically in get_token_storage_path()
+    # Model configuration
+    "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",
+}
+
+
+def get_token_storage_path() -> Path:
+    """Get the path for storing OAuth tokens (uses XDG_DATA_HOME)."""
+    data_dir = Path(config.DATA_DIR)
+    data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
+    return data_dir / "chatgpt_oauth.json"
+
+
+def get_config_dir() -> Path:
+    """Get the Code Puppy configuration directory (uses XDG_CONFIG_HOME)."""
+    config_dir = Path(config.CONFIG_DIR)
+    config_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
+    return config_dir
+
+
+def get_chatgpt_models_path() -> Path:
+    """Get the path to the dedicated chatgpt_models.json file (uses XDG_DATA_HOME)."""
+    data_dir = Path(config.DATA_DIR)
+    data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
+    return data_dir / "chatgpt_models.json"

code_puppy/plugins/chatgpt_oauth/oauth_flow.py

@@ -0,0 +1,328 @@
+"""ChatGPT OAuth flow closely matching the ChatMock implementation."""
+
+from __future__ import annotations
+
+import datetime
+import threading
+import time
+import urllib.parse
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any, Optional, Tuple
+
+import requests
+
+from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
+
+from ..oauth_puppy_html import oauth_failure_html, oauth_success_html
+from .config import CHATGPT_OAUTH_CONFIG
+from .utils import (
+    add_models_to_extra_config,
+    assign_redirect_uri,
+    load_stored_tokens,
+    parse_jwt_claims,
+    prepare_oauth_context,
+    save_tokens,
+)
+
+REQUIRED_PORT = CHATGPT_OAUTH_CONFIG["required_port"]
+URL_BASE = f"http://localhost:{REQUIRED_PORT}"
+
+
+@dataclass
+class TokenData:
+    id_token: str
+    access_token: str
+    refresh_token: str
+    account_id: str
+
+
+@dataclass
+class AuthBundle:
+    api_key: Optional[str]
+    token_data: TokenData
+    last_refresh: str
+
+
+class _OAuthServer(HTTPServer):
+    def __init__(
+        self,
+        *,
+        client_id: str,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(
+            ("localhost", REQUIRED_PORT), _CallbackHandler, bind_and_activate=True
+        )
+        self.exit_code = 1
+        self.verbose = verbose
+        self.client_id = client_id
+        self.issuer = CHATGPT_OAUTH_CONFIG["issuer"]
+        self.token_endpoint = CHATGPT_OAUTH_CONFIG["token_url"]
+
+        # Create fresh OAuth context for this server instance
+        context = prepare_oauth_context()
+        self.redirect_uri = assign_redirect_uri(context, REQUIRED_PORT)
+        self.context = context
+
+    def auth_url(self) -> str:
+        params = {
+            "response_type": "code",
+            "client_id": self.client_id,
+            "redirect_uri": self.redirect_uri,
+            "scope": CHATGPT_OAUTH_CONFIG["scope"],
+            "code_challenge": self.context.code_challenge,
+            "code_challenge_method": "S256",
+            "id_token_add_organizations": "true",
+            "codex_cli_simplified_flow": "true",
+            "state": self.context.state,
+        }
+        return f"{self.issuer}/oauth/authorize?" + urllib.parse.urlencode(params)
+
+    def exchange_code(self, code: str) -> Tuple[AuthBundle, str]:
+        data = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": self.redirect_uri,
+            "client_id": self.client_id,
+            "code_verifier": self.context.code_verifier,
+        }
+
+        response = requests.post(
+            self.token_endpoint,
+            data=data,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            timeout=30,
+        )
+        response.raise_for_status()
+        payload = response.json()
+
+        id_token = payload.get("id_token", "")
+        access_token = payload.get("access_token", "")
+        refresh_token = payload.get("refresh_token", "")
+
+        id_token_claims = parse_jwt_claims(id_token) or {}
+        access_token_claims = parse_jwt_claims(access_token) or {}
+
+        auth_claims = id_token_claims.get("https://api.openai.com/auth") or {}
+        chatgpt_account_id = auth_claims.get("chatgpt_account_id", "")
+        # Extract org_id from nested auth structure like ChatMock
+        organizations = auth_claims.get("organizations", [])
+        org_id = None
+        if organizations:
+            default_org = next(
+                (org for org in organizations if org.get("is_default")),
+                organizations[0],
+            )
+            org_id = default_org.get("id")
+        # Fallback to top-level org_id if still not found
+        if not org_id:
+            org_id = id_token_claims.get("organization_id")
+
+        token_data = TokenData(
+            id_token=id_token,
+            access_token=access_token,
+            refresh_token=refresh_token,
+            account_id=chatgpt_account_id,
+        )
+
+        # Instead of exchanging for an API key, just use the access_token directly
+        # This matches how ChatMock works - no token exchange, just OAuth tokens
+        api_key = token_data.access_token
+
+        last_refresh = (
+            datetime.datetime.now(datetime.timezone.utc)
+            .isoformat()
+            .replace("+00:00", "Z")
+        )
+        bundle = AuthBundle(
+            api_key=api_key, token_data=token_data, last_refresh=last_refresh
+        )
+
+        # Build success URL with all the token info
+        success_query = {
+            "id_token": token_data.id_token,
+            "access_token": token_data.access_token,
+            "refresh_token": token_data.refresh_token,
+            "org_id": org_id or "",
+            "plan_type": access_token_claims.get("chatgpt_plan_type"),
+            "platform_url": "https://platform.openai.com",
+        }
+        success_url = f"{URL_BASE}/success?{urllib.parse.urlencode(success_query)}"
+        return bundle, success_url
+
+
+class _CallbackHandler(BaseHTTPRequestHandler):
+    server: "_OAuthServer"
+
+    def do_GET(self) -> None:  # noqa: N802
+        path = urllib.parse.urlparse(self.path).path
+        if path == "/success":
+            success_html = oauth_success_html(
+                "ChatGPT",
+                "You can now close this window and return to Code Puppy.",
+            )
+            self._send_html(success_html)
+            self._shutdown_after_delay(2.0)
+            return
+
+        if path != "/auth/callback":
+            self._send_failure(404, "Callback endpoint not found for the puppy parade.")
+            self._shutdown()
+            return
+
+        query = urllib.parse.urlparse(self.path).query
+        params = urllib.parse.parse_qs(query)
+
+        code = params.get("code", [None])[0]
+        if not code:
+            self._send_failure(400, "Missing auth code — the token treat rolled away.")
+            self._shutdown()
+            return
+
+        try:
+            auth_bundle, success_url = self.server.exchange_code(code)
+        except Exception as exc:  # noqa: BLE001
+            self._send_failure(500, f"Token exchange failed: {exc}")
+            self._shutdown()
+            return
+
+        tokens = {
+            "id_token": auth_bundle.token_data.id_token,
+            "access_token": auth_bundle.token_data.access_token,
+            "refresh_token": auth_bundle.token_data.refresh_token,
+            "account_id": auth_bundle.token_data.account_id,
+            "last_refresh": auth_bundle.last_refresh,
+        }
+        if auth_bundle.api_key:
+            tokens["api_key"] = auth_bundle.api_key
+
+        if save_tokens(tokens):
+            self.server.exit_code = 0
+            # Redirect to the success URL returned by exchange_code
+            self._send_redirect(success_url)
+        else:
+            self._send_failure(
+                500, "Unable to persist auth file — a puppy probably chewed it."
+            )
+            self._shutdown()
+        self._shutdown_after_delay(2.0)
+
+    def do_POST(self) -> None:  # noqa: N802
+        self._send_failure(
+            404, "POST not supported — the pups only fetch GET requests."
+        )
+        self._shutdown()
+
+    def log_message(self, fmt: str, *args: Any) -> None:  # noqa: A003
+        if getattr(self.server, "verbose", False):
+            super().log_message(fmt, *args)
+
+    def _send_redirect(self, url: str) -> None:
+        self.send_response(302)
+        self.send_header("Location", url)
+        self.end_headers()
+
+    def _send_html(self, body: str, status: int = 200) -> None:
+        encoded = body.encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(encoded)))
+        self.end_headers()
+        self.wfile.write(encoded)
+
+    def _send_failure(self, status: int, reason: str) -> None:
+        failure_html = oauth_failure_html("ChatGPT", reason)
+        self._send_html(failure_html, status)
+
+    def _shutdown(self) -> None:
+        threading.Thread(target=self.server.shutdown, daemon=True).start()
+
+    def _shutdown_after_delay(self, seconds: float = 2.0) -> None:
+        def _later() -> None:
+            try:
+                time.sleep(seconds)
+            finally:
+                self._shutdown()
+
+        threading.Thread(target=_later, daemon=True).start()
+
+
+def run_oauth_flow() -> None:
+    existing_tokens = load_stored_tokens()
+    if existing_tokens and existing_tokens.get("access_token"):
+        emit_warning("Existing ChatGPT tokens will be overwritten.")
+
+    try:
+        server = _OAuthServer(client_id=CHATGPT_OAUTH_CONFIG["client_id"])
+    except OSError as exc:
+        emit_error(f"Could not start OAuth server on port {REQUIRED_PORT}: {exc}")
+        emit_info(f"Use `lsof -ti:{REQUIRED_PORT} | xargs kill` to free the port.")
+        return
+
+    auth_url = server.auth_url()
+    emit_info(f"Open this URL in your browser: {auth_url}")
+
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
+
+    webbrowser_opened = False
+    try:
+        import webbrowser
+
+        from code_puppy.tools.common import should_suppress_browser
+
+        if should_suppress_browser():
+            emit_info(f"[HEADLESS MODE] Would normally open: {auth_url}")
+        else:
+            webbrowser_opened = webbrowser.open(auth_url)
+    except Exception as exc:  # noqa: BLE001
+        emit_warning(f"Could not open browser automatically: {exc}")
+
+    if not webbrowser_opened and not should_suppress_browser():
+        emit_warning("Please open the URL manually if the browser did not open.")
+
+    emit_info("Waiting for authentication callback…")
+
+    elapsed = 0.0
+    timeout = CHATGPT_OAUTH_CONFIG["callback_timeout"]
+    interval = 0.25
+    while elapsed < timeout:
+        time.sleep(interval)
+        elapsed += interval
+        if server.exit_code == 0:
+            break
+
+    server.shutdown()
+    server_thread.join(timeout=5)
+
+    if server.exit_code != 0:
+        emit_error("Authentication failed or timed out.")
+        return
+
+    tokens = load_stored_tokens()
+    if not tokens:
+        emit_error("Tokens saved during OAuth flow could not be loaded.")
+        return
+
+    api_key = tokens.get("api_key")
+    if api_key:
+        emit_success("Successfully obtained OAuth access token for API access.")
+        emit_info(
+            f"Access token saved and available via {CHATGPT_OAUTH_CONFIG['api_key_env_var']}"
+        )
+    else:
+        emit_warning(
+            "No API key obtained. You may need to configure projects at platform.openai.com."
+        )
+
+    if 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):
+                emit_success(
+                    "ChatGPT models registered. Use the `chatgpt-` prefix in /model."
+                )

code_puppy/plugins/chatgpt_oauth/register_callbacks.py

@@ -0,0 +1,94 @@
+"""ChatGPT OAuth plugin callbacks aligned with ChatMock flow."""
+
+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
+from .oauth_flow import run_oauth_flow
+from .utils import load_chatgpt_models, load_stored_tokens, remove_chatgpt_models
+
+
+def _custom_help() -> List[Tuple[str, str]]:
+    return [
+        (
+            "chatgpt-auth",
+            "Authenticate with ChatGPT via OAuth and import available models",
+        ),
+        (
+            "chatgpt-status",
+            "Check ChatGPT OAuth authentication status and configured models",
+        ),
+        ("chatgpt-logout", "Remove ChatGPT OAuth tokens and imported models"),
+    ]
+
+
+def _handle_chatgpt_status() -> None:
+    tokens = load_stored_tokens()
+    if tokens and tokens.get("access_token"):
+        emit_success("🔐 ChatGPT OAuth: Authenticated")
+
+        api_key = tokens.get("api_key")
+        if api_key:
+            os.environ[CHATGPT_OAUTH_CONFIG["api_key_env_var"]] = api_key
+            emit_info("✅ OAuth access token available for API requests")
+        else:
+            emit_warning("⚠️ No access token obtained. Authentication may have failed.")
+
+        chatgpt_models = [
+            name
+            for name, cfg in load_chatgpt_models().items()
+            if cfg.get("oauth_source") == "chatgpt-oauth-plugin"
+        ]
+        if chatgpt_models:
+            emit_info(f"🎯 Configured ChatGPT models: {', '.join(chatgpt_models)}")
+        else:
+            emit_warning("⚠️ No ChatGPT models configured yet.")
+    else:
+        emit_warning("🔓 ChatGPT OAuth: Not authenticated")
+        emit_info("🌐 Run /chatgpt-auth to launch the browser sign-in flow.")
+
+
+def _handle_chatgpt_logout() -> None:
+    token_path = get_token_storage_path()
+    if token_path.exists():
+        token_path.unlink()
+        emit_info("Removed ChatGPT OAuth tokens")
+
+    if CHATGPT_OAUTH_CONFIG["api_key_env_var"] in os.environ:
+        del os.environ[CHATGPT_OAUTH_CONFIG["api_key_env_var"]]
+
+    removed = remove_chatgpt_models()
+    if removed:
+        emit_info(f"Removed {removed} ChatGPT models from configuration")
+
+    emit_success("ChatGPT logout complete")
+
+
+def _handle_custom_command(command: str, name: str) -> Optional[bool]:
+    if not name:
+        return None
+
+    if name == "chatgpt-auth":
+        run_oauth_flow()
+        set_model_name("chatgpt-gpt-5.2-codex")
+        return True
+
+    if name == "chatgpt-status":
+        _handle_chatgpt_status()
+        return True
+
+    if name == "chatgpt-logout":
+        _handle_chatgpt_logout()
+        return True
+
+    return None
+
+
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)