codepp 0.0.437__py3-none-any.whl

code_puppy/plugins/chatgpt_oauth/utils.py

@@ -0,0 +1,523 @@
+"""Utility helpers for the ChatGPT OAuth plugin."""
+
+from __future__ import annotations
+
+import base64
+import datetime
+import hashlib
+import json
+import logging
+import secrets
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+from urllib.parse import parse_qs as urllib_parse_qs
+from urllib.parse import urlencode, urlparse
+
+import requests
+
+from .config import (
+    CHATGPT_OAUTH_CONFIG,
+    get_chatgpt_models_path,
+    get_token_storage_path,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OAuthContext:
+    """Runtime state for an in-progress OAuth flow."""
+
+    state: str
+    code_verifier: str
+    code_challenge: str
+    created_at: float
+    redirect_uri: Optional[str] = None
+    expires_at: Optional[float] = None  # Add expiration time
+
+    def is_expired(self) -> bool:
+        """Check if this OAuth context has expired."""
+        if self.expires_at is None:
+            # Default 5 minute expiration if not set
+            return time.time() - self.created_at > 300
+        return time.time() > self.expires_at
+
+
+def _urlsafe_b64encode(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")
+
+
+def _generate_code_verifier() -> str:
+    return secrets.token_hex(64)
+
+
+def _compute_code_challenge(code_verifier: str) -> str:
+    digest = hashlib.sha256(code_verifier.encode("utf-8")).digest()
+    return _urlsafe_b64encode(digest)
+
+
+def prepare_oauth_context() -> OAuthContext:
+    """Create a fresh OAuth PKCE context."""
+    state = secrets.token_hex(32)
+    code_verifier = _generate_code_verifier()
+    code_challenge = _compute_code_challenge(code_verifier)
+
+    # Set expiration 4 minutes from now (OpenAI sessions are short)
+    expires_at = time.time() + 240
+
+    return OAuthContext(
+        state=state,
+        code_verifier=code_verifier,
+        code_challenge=code_challenge,
+        created_at=time.time(),
+        expires_at=expires_at,
+    )
+
+
+def assign_redirect_uri(context: OAuthContext, port: int) -> str:
+    """Assign redirect URI for the given OAuth context."""
+    if context is None:
+        raise RuntimeError("OAuth context cannot be None")
+    host = CHATGPT_OAUTH_CONFIG["redirect_host"].rstrip("/")
+    path = CHATGPT_OAUTH_CONFIG["redirect_path"].lstrip("/")
+    required_port = CHATGPT_OAUTH_CONFIG.get("required_port")
+    if required_port and port != required_port:
+        raise RuntimeError(
+            f"OAuth flow must use port {required_port}; attempted to assign port {port}"
+        )
+    redirect_uri = f"{host}:{port}/{path}"
+    context.redirect_uri = redirect_uri
+    return redirect_uri
+
+
+def build_authorization_url(context: OAuthContext) -> str:
+    """Return the OpenAI authorization URL with PKCE parameters."""
+    if not context.redirect_uri:
+        raise RuntimeError("Redirect URI has not been assigned for this OAuth context")
+
+    params = {
+        "response_type": "code",
+        "client_id": CHATGPT_OAUTH_CONFIG["client_id"],
+        "redirect_uri": context.redirect_uri,
+        "scope": CHATGPT_OAUTH_CONFIG["scope"],
+        "code_challenge": context.code_challenge,
+        "code_challenge_method": "S256",
+        "id_token_add_organizations": "true",
+        "codex_cli_simplified_flow": "true",
+        "state": context.state,
+    }
+    return f"{CHATGPT_OAUTH_CONFIG['auth_url']}?{urlencode(params)}"
+
+
+def parse_authorization_error(url: str) -> Optional[str]:
+    """Parse error from OAuth callback URL."""
+    try:
+        parsed = urlparse(url)
+        params = urllib_parse_qs(parsed.query)
+        error = params.get("error", [None])[0]
+        error_description = params.get("error_description", [None])[0]
+        if error:
+            return f"{error}: {error_description or 'Unknown error'}"
+    except Exception as exc:
+        logger.error("Failed to parse OAuth error: %s", exc)
+    return None
+
+
+def parse_jwt_claims(token: str) -> Optional[Dict[str, Any]]:
+    """Parse JWT token to extract claims."""
+    if not token or token.count(".") != 2:
+        return None
+    try:
+        _, payload, _ = token.split(".")
+        padded = payload + "=" * (-len(payload) % 4)
+        data = base64.urlsafe_b64decode(padded.encode())
+        return json.loads(data.decode())
+    except Exception as exc:
+        logger.error("Failed to parse JWT: %s", exc)
+    return None
+
+
+def load_stored_tokens() -> Optional[Dict[str, Any]]:
+    try:
+        token_path = get_token_storage_path()
+        if token_path.exists():
+            with open(token_path, "r", encoding="utf-8") as handle:
+                return json.load(handle)
+    except Exception as exc:
+        logger.error("Failed to load tokens: %s", exc)
+    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")
+    try:
+        token_path = get_token_storage_path()
+        with open(token_path, "w", encoding="utf-8") as handle:
+            json.dump(tokens, handle, indent=2)
+        token_path.chmod(0o600)
+        return True
+    except Exception as exc:
+        logger.error("Failed to save tokens: %s", exc)
+    return False
+
+
+def load_chatgpt_models() -> Dict[str, Any]:
+    try:
+        models_path = get_chatgpt_models_path()
+        if models_path.exists():
+            with open(models_path, "r", encoding="utf-8") as handle:
+                return json.load(handle)
+    except Exception as exc:
+        logger.error("Failed to load ChatGPT models: %s", exc)
+    return {}
+
+
+def save_chatgpt_models(models: Dict[str, Any]) -> bool:
+    try:
+        models_path = get_chatgpt_models_path()
+        with open(models_path, "w", encoding="utf-8") as handle:
+            json.dump(models, handle, indent=2)
+        return True
+    except Exception as exc:
+        logger.error("Failed to save ChatGPT models: %s", exc)
+    return False
+
+
+def exchange_code_for_tokens(
+    auth_code: str, context: OAuthContext
+) -> Optional[Dict[str, Any]]:
+    """Exchange authorization code for access tokens."""
+    if not context.redirect_uri:
+        raise RuntimeError("Redirect URI missing from OAuth context")
+
+    if context.is_expired():
+        logger.error("OAuth context expired, cannot exchange code")
+        return None
+
+    payload = {
+        "grant_type": "authorization_code",
+        "code": auth_code,
+        "redirect_uri": context.redirect_uri,
+        "client_id": CHATGPT_OAUTH_CONFIG["client_id"],
+        "code_verifier": context.code_verifier,
+    }
+
+    headers = {
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+
+    logger.info("Exchanging code for tokens: %s", CHATGPT_OAUTH_CONFIG["token_url"])
+    try:
+        response = requests.post(
+            CHATGPT_OAUTH_CONFIG["token_url"],
+            data=payload,
+            headers=headers,
+            timeout=30,
+        )
+        logger.info("Token exchange response: %s", response.status_code)
+        if response.status_code == 200:
+            token_data = response.json()
+            # Add timestamp
+            token_data["last_refresh"] = (
+                datetime.datetime.now(datetime.timezone.utc)
+                .isoformat()
+                .replace("+00:00", "Z")
+            )
+            return token_data
+        else:
+            logger.error(
+                "Token exchange failed: %s - %s",
+                response.status_code,
+                response.text,
+            )
+            # Try to parse OAuth error
+            if response.headers.get("content-type", "").startswith("application/json"):
+                try:
+                    error_data = response.json()
+                    if "error" in error_data:
+                        logger.error(
+                            "OAuth error: %s",
+                            error_data.get("error_description", error_data["error"]),
+                        )
+                except Exception:
+                    pass
+    except Exception as exc:
+        logger.error("Token exchange error: %s", exc)
+    return None
+
+
+# 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.4",
+    "gpt-5.3-instant",
+    "gpt-5.3-codex-spark",
+    "gpt-5.3-codex",
+    "gpt-5.2-codex",
+    "gpt-5.2",
+]
+
+# Models that MUST always be registered, even if the /models endpoint
+# doesn't return them (e.g. newly launched, not yet in the API catalogue).
+# These are merged into whatever the endpoint returns.
+REQUIRED_CODEX_MODELS = [
+    "gpt-5.4",
+    "gpt-5.3-instant",
+    "gpt-5.3-codex",
+]
+
+# Per-model context length overrides (tokens).
+# Models not listed here use CHATGPT_OAUTH_CONFIG["default_context_length"] (272,000).
+CODEX_MODEL_CONTEXT_LENGTHS = {
+    "gpt-5.3-codex-spark": 131000,
+    "gpt-5.3-instant": 192000,
+}
+
+
+def _ensure_required_models(models: List[str]) -> List[str]:
+    """Merge REQUIRED_CODEX_MODELS into the given list, preserving order.
+
+    Any required model not already present is prepended so it appears first.
+    """
+    existing = set(models)
+    missing = [m for m in REQUIRED_CODEX_MODELS if m not in existing]
+    if missing:
+        logger.info("Injecting required models not returned by API: %s", missing)
+    return missing + models
+
+
+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:
+        access_token: OAuth access token for authentication
+        account_id: ChatGPT account ID (required for the API)
+
+    Returns:
+        List of model IDs, or default list if API fails
+    """
+    import platform
+
+    # 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 {access_token}",
+        "ChatGPT-Account-Id": account_id,
+        "User-Agent": user_agent,
+        "originator": originator,
+        "Accept": "application/json",
+    }
+
+    # Query params
+    params = {"client_version": client_version}
+
+    try:
+        response = requests.get(models_url, headers=headers, params=params, timeout=30)
+
+        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 _ensure_required_models(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.warning("Timeout fetching models, using default list")
+    except requests.exceptions.RequestException as exc:
+        logger.warning("Network error fetching models: %s, using default list", exc)
+    except Exception as exc:
+        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]) -> 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": "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"],
+                },
+                "context_length": CODEX_MODEL_CONTEXT_LENGTHS.get(
+                    model_name, 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):
+            logger.info("Added %s ChatGPT models", added)
+            return True
+    except Exception as exc:
+        logger.error("Error adding models to config: %s", exc)
+    return False
+
+
+def remove_chatgpt_models() -> int:
+    """Remove ChatGPT OAuth models from chatgpt_models.json."""
+    try:
+        chatgpt_models = load_chatgpt_models()
+        to_remove = [
+            name
+            for name, config in chatgpt_models.items()
+            if config.get("oauth_source") == "chatgpt-oauth-plugin"
+        ]
+        for model_name in to_remove:
+            chatgpt_models.pop(model_name, None)
+        # Always save, even if no models were removed (to match test expectations)
+        if save_chatgpt_models(chatgpt_models):
+            return len(to_remove)
+    except Exception as exc:
+        logger.error("Error removing ChatGPT models: %s", exc)
+    return 0

code_puppy/plugins/claude_code_hooks/__init__.py

@@ -0,0 +1 @@
+"""Claude Code hooks plugin."""

code_puppy/plugins/claude_code_hooks/config.py

@@ -0,0 +1,137 @@
+"""
+Configuration loader for Claude Code hooks.
+
+Loads and merges hooks from multiple locations:
+1. ~/.code_puppy/hooks.json (global level) - always loaded if exists
+2. .claude/settings.json (project-level) - merged with global
+
+Both configurations are loaded and merged so that hooks from both levels
+coexist and execute together.
+"""
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+PROJECT_HOOKS_FILE = ".claude/settings.json"
+GLOBAL_HOOKS_FILE = os.path.expanduser("~/.code_puppy/hooks.json")
+
+
+def _deep_merge_hooks(base: Dict[str, Any], overlay: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Merge hook configurations, combining event types and hook groups.
+
+    When the same event type exists in both base and overlay, their hook groups
+    are concatenated (overlay hooks appear after base hooks).
+
+    Args:
+        base: Base configuration dictionary
+        overlay: Configuration to merge on top
+
+    Returns:
+        Merged configuration with all hooks from both sources
+    """
+    merged = dict(base)
+
+    for event_type, hook_groups in overlay.items():
+        if event_type.startswith("_"):
+            # Skip comment keys
+            merged[event_type] = hook_groups
+            continue
+
+        if event_type not in merged:
+            # New event type, just add it
+            merged[event_type] = hook_groups
+        elif isinstance(merged[event_type], list) and isinstance(hook_groups, list):
+            # Both are lists, concatenate them (overlay hooks come after)
+            merged[event_type] = merged[event_type] + hook_groups
+            logger.debug(
+                f"Merged {len(hook_groups)} hook group(s) for event '{event_type}'"
+            )
+        else:
+            # Type mismatch or unexpected structure, keep base
+            logger.warning(
+                f"Cannot merge event type '{event_type}': type mismatch or unexpected structure"
+            )
+
+    return merged
+
+
+def load_hooks_config() -> Optional[Dict[str, Any]]:
+    """
+    Load and merge hooks configuration from available sources.
+
+    Priority order:
+    1. ~/.code_puppy/hooks.json (global level) - always loaded if exists
+    2. .claude/settings.json (project-level) - merged with global
+
+    Returns:
+        Configuration dictionary or None if no config found
+    """
+    merged_config: Dict[str, Any] = {}
+
+    # Load global hooks first
+    global_config_path = Path(GLOBAL_HOOKS_FILE)
+
+    if global_config_path.exists():
+        try:
+            with open(global_config_path, "r", encoding="utf-8") as f:
+                config = json.load(f)
+            if "hooks" in config and isinstance(config["hooks"], dict):
+                logger.info(
+                    f"Loaded hooks configuration (wrapped format) from {GLOBAL_HOOKS_FILE}"
+                )
+                merged_config = _deep_merge_hooks(merged_config, config["hooks"])
+            elif isinstance(config, dict):
+                logger.info(f"Loaded hooks configuration from {GLOBAL_HOOKS_FILE}")
+                merged_config = _deep_merge_hooks(merged_config, config)
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON in {GLOBAL_HOOKS_FILE}: {e}")
+        except Exception as e:
+            logger.error(f"Failed to load {GLOBAL_HOOKS_FILE}: {e}", exc_info=True)
+
+    # Load and merge project-level hooks
+    project_config_path = Path(os.getcwd()) / PROJECT_HOOKS_FILE
+
+    if project_config_path.exists():
+        try:
+            with open(project_config_path, "r", encoding="utf-8") as f:
+                config = json.load(f)
+            hooks_config = config.get("hooks")
+            if hooks_config:
+                logger.info(f"Merging hooks configuration from {project_config_path}")
+                merged_config = _deep_merge_hooks(merged_config, hooks_config)
+            else:
+                logger.debug(f"No 'hooks' section found in {project_config_path}")
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON in {project_config_path}: {e}")
+        except Exception as e:
+            logger.error(f"Failed to load {project_config_path}: {e}", exc_info=True)
+
+    if not merged_config:
+        logger.debug("No hooks configuration found")
+        return None
+
+    event_count = len(
+        [event for event in merged_config.keys() if not event.startswith("_")]
+    )
+    logger.info(f"Hooks configuration ready ({event_count} event type(s))")
+    return merged_config
+
+
+def get_hooks_config_paths() -> list:
+    """
+    Return list of hook configuration paths.
+
+    Returns paths in order of precedence (project-level first, then global).
+    Note: internally, hooks are loaded in reverse order (global first, then project)
+    so that project-level hooks can extend/append to global hooks.
+    """
+    return [
+        str(Path(os.getcwd()) / PROJECT_HOOKS_FILE),
+        GLOBAL_HOOKS_FILE,
+    ]