codepp 0.0.437__py3-none-any.whl

code_puppy/plugins/chatgpt_oauth/register_callbacks.py

@@ -0,0 +1,176 @@
+"""ChatGPT OAuth plugin callbacks aligned with ChatMock flow.
+
+Provides OAuth authentication for ChatGPT models and registers
+the 'chatgpt_oauth' model type handler.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from code_puppy.callbacks import register_callback
+from code_puppy.messaging import emit_info, emit_success, emit_warning
+from code_puppy.model_switching import set_model_and_reload_agent
+
+from .config import CHATGPT_OAUTH_CONFIG, get_token_storage_path
+from .oauth_flow import run_oauth_flow
+from .utils import (
+    get_valid_access_token,
+    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_and_reload_agent("chatgpt-gpt-5.3-codex")
+        return True
+
+    if name == "chatgpt-status":
+        _handle_chatgpt_status()
+        return True
+
+    if name == "chatgpt-logout":
+        _handle_chatgpt_logout()
+        return True
+
+    return None
+
+
+def _create_chatgpt_oauth_model(
+    model_name: str, model_config: Dict, config: Dict
+) -> Any:
+    """Create a ChatGPT OAuth model instance.
+
+    This handler is registered via the 'register_model_type' callback to handle
+    models with type='chatgpt_oauth'.
+    """
+    from pydantic_ai.models.openai import OpenAIResponsesModel
+    from pydantic_ai.providers.openai import OpenAIProvider
+
+    from code_puppy.chatgpt_codex_client import create_codex_async_client
+    from code_puppy.http_utils import get_cert_bundle_path
+
+    # Get a valid access token (refreshing if needed)
+    access_token = get_valid_access_token()
+    if not access_token:
+        emit_warning(
+            f"Failed to get valid ChatGPT OAuth token; skipping model '{model_config.get('name')}'. "
+            "Run /chatgpt-auth to authenticate."
+        )
+        return None
+
+    # Get account_id from stored tokens (required for ChatGPT-Account-Id header)
+    tokens = load_stored_tokens()
+    account_id = tokens.get("account_id", "") if tokens else ""
+    if not account_id:
+        emit_warning(
+            f"No account_id found in ChatGPT OAuth tokens; skipping model '{model_config.get('name')}'. "
+            "Run /chatgpt-auth to re-authenticate."
+        )
+        return None
+
+    # Build headers for ChatGPT Codex API
+    originator = CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs")
+    client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.72.0")
+
+    headers = {
+        "ChatGPT-Account-Id": account_id,
+        "originator": originator,
+        "User-Agent": f"{originator}/{client_version}",
+    }
+    # Merge with any headers from model config
+    config_headers = model_config.get("custom_endpoint", {}).get("headers", {})
+    headers.update(config_headers)
+
+    # Get base URL - Codex API uses chatgpt.com, not api.openai.com
+    base_url = model_config.get("custom_endpoint", {}).get(
+        "url", CHATGPT_OAUTH_CONFIG["api_base_url"]
+    )
+
+    # Create HTTP client with Codex interceptor for store=false injection
+    verify = get_cert_bundle_path()
+    client = create_codex_async_client(headers=headers, verify=verify)
+
+    provider = OpenAIProvider(
+        api_key=access_token,
+        base_url=base_url,
+        http_client=client,
+    )
+
+    # ChatGPT Codex API only supports Responses format
+    model = OpenAIResponsesModel(model_name=model_config["name"], provider=provider)
+    model.provider = provider
+    return model
+
+
+def _register_model_types() -> List[Dict[str, Any]]:
+    """Register the chatgpt_oauth model type handler."""
+    return [{"type": "chatgpt_oauth", "handler": _create_chatgpt_oauth_model}]
+
+
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+register_callback("register_model_type", _register_model_types)

code_puppy/plugins/chatgpt_oauth/test_plugin.py

@@ -0,0 +1,301 @@
+"""
+Basic tests for ChatGPT OAuth plugin.
+"""
+
+import json
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from code_puppy.plugins.chatgpt_oauth import config, utils
+
+
+def test_config_paths():
+    """Test configuration path helpers."""
+    token_path = config.get_token_storage_path()
+    assert token_path.name == "chatgpt_oauth.json"
+    # XDG paths use "code_puppy" (without dot) in ~/.local/share or ~/.config
+    assert "code_puppy" in str(token_path)
+
+    config_dir = config.get_config_dir()
+    # Default is ~/.code_puppy; XDG paths only used when XDG env vars are set
+    assert config_dir.name in ("code_puppy", ".code_puppy")
+
+    chatgpt_models = config.get_chatgpt_models_path()
+    assert chatgpt_models.name == "chatgpt_models.json"
+
+
+def test_oauth_config():
+    """Test OAuth configuration values."""
+    assert config.CHATGPT_OAUTH_CONFIG["issuer"] == "https://auth.openai.com"
+    assert config.CHATGPT_OAUTH_CONFIG["client_id"] == "app_EMoamEEZ73f0CkXaXp7hrann"
+    assert config.CHATGPT_OAUTH_CONFIG["prefix"] == "chatgpt-"
+
+
+def test_jwt_parsing_with_nested_org():
+    """Test JWT parsing with nested organization structure like the user's payload."""
+    # This simulates the user's JWT payload structure
+    mock_claims = {
+        "aud": ["app_EMoamEEZ73f0CkXaXp7hrann"],
+        "auth_provider": "google",
+        "email": "mike.pfaf fenberger@gmail.com",
+        "https://api.openai.com/auth": {
+            "chatgpt_account_id": "d1844a91-9aac-419b-903e-f6a99c76f163",
+            "organizations": [
+                {
+                    "id": "org-iydWjnSxSr51VuYhDVMDte5",
+                    "is_default": True,
+                    "role": "owner",
+                    "title": "Personal",
+                }
+            ],
+            "groups": ["api-data-sharing-incentives-program", "verified-organization"],
+        },
+        "sub": "google-oauth2|107692466937587138174",
+    }
+
+    # Test the org extraction logic
+    auth_claims = mock_claims.get("https://api.openai.com/auth", {})
+    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")
+
+    assert org_id == "org-iydWjnSxSr51VuYhDVMDte5"
+
+    # Test fallback to top-level org_id (should not happen in this case)
+    if not org_id:
+        org_id = mock_claims.get("organization_id")
+
+    assert org_id == "org-iydWjnSxSr51VuYhDVMDte5"
+    assert config.CHATGPT_OAUTH_CONFIG["required_port"] == 1455
+
+
+def test_code_verifier_generation():
+    """Test PKCE code verifier generation."""
+    verifier = utils._generate_code_verifier()
+    assert isinstance(verifier, str)
+    assert len(verifier) > 50  # Should be long
+
+
+def test_code_challenge_computation():
+    """Test PKCE code challenge computation."""
+    verifier = "test_verifier_string"
+    challenge = utils._compute_code_challenge(verifier)
+    assert isinstance(challenge, str)
+    assert len(challenge) > 0
+    # Should be URL-safe base64
+    assert all(
+        c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
+        for c in challenge
+    )
+
+
+def test_prepare_oauth_context():
+    """Test OAuth context preparation."""
+    context = utils.prepare_oauth_context()
+    assert context.state
+    assert context.code_verifier
+    assert context.code_challenge
+    assert context.created_at > 0
+    assert context.redirect_uri is None
+
+
+def test_assign_redirect_uri():
+    """Test redirect URI assignment."""
+    context = utils.prepare_oauth_context()
+    redirect_uri = utils.assign_redirect_uri(context, 1455)
+    assert redirect_uri == "http://localhost:1455/auth/callback"
+    assert context.redirect_uri == redirect_uri
+
+
+def test_build_authorization_url():
+    """Test authorization URL building."""
+    context = utils.prepare_oauth_context()
+    utils.assign_redirect_uri(context, 1455)
+    auth_url = utils.build_authorization_url(context)
+
+    assert auth_url.startswith("https://auth.openai.com/oauth/authorize?")
+    assert "response_type=code" in auth_url
+    assert "client_id=" in auth_url
+    assert "redirect_uri=" in auth_url
+    assert "code_challenge=" in auth_url
+    assert "code_challenge_method=S256" in auth_url
+    assert f"state={context.state}" in auth_url
+
+
+def test_parse_jwt_claims():
+    """Test JWT claims parsing."""
+    # Valid JWT structure (header.payload.signature)
+    import base64
+
+    payload = base64.urlsafe_b64encode(json.dumps({"sub": "user123"}).encode()).decode()
+    token = f"header.{payload}.signature"
+
+    claims = utils.parse_jwt_claims(token)
+    assert claims is not None
+    assert claims["sub"] == "user123"
+
+    # Invalid token
+    assert utils.parse_jwt_claims("") is None
+    assert utils.parse_jwt_claims("invalid") is None
+
+
+def test_save_and_load_tokens(tmp_path):
+    """Test token storage and retrieval."""
+    with patch.object(
+        config, "get_token_storage_path", return_value=tmp_path / "tokens.json"
+    ):
+        tokens = {
+            "access_token": "test_access",
+            "refresh_token": "test_refresh",
+            "api_key": "sk-test",
+        }
+
+        # Save tokens
+        assert utils.save_tokens(tokens)
+
+        # Load tokens
+        loaded = utils.load_stored_tokens()
+        assert loaded == tokens
+
+
+def test_save_and_load_chatgpt_models(tmp_path):
+    """Test ChatGPT models configuration."""
+    with patch.object(
+        config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
+    ):
+        models = {
+            "chatgpt-gpt-4o": {
+                "type": "openai",
+                "name": "gpt-4o",
+                "oauth_source": "chatgpt-oauth-plugin",
+            }
+        }
+
+        # Save models
+        assert utils.save_chatgpt_models(models)
+
+        # Load models
+        loaded = utils.load_chatgpt_models()
+        assert loaded == models
+
+
+def test_remove_chatgpt_models(tmp_path):
+    """Test removal of ChatGPT models from config."""
+    with patch.object(
+        config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
+    ):
+        models = {
+            "chatgpt-gpt-4o": {
+                "type": "openai",
+                "oauth_source": "chatgpt-oauth-plugin",
+            },
+            "claude-3-opus": {
+                "type": "anthropic",
+                "oauth_source": "other",
+            },
+        }
+        utils.save_chatgpt_models(models)
+
+        # Remove only ChatGPT models
+        removed_count = utils.remove_chatgpt_models()
+        assert removed_count == 1
+
+        # Verify only ChatGPT model was removed
+        remaining = utils.load_chatgpt_models()
+        assert "chatgpt-gpt-4o" not in remaining
+        assert "claude-3-opus" in remaining
+
+
+@patch("code_puppy.plugins.chatgpt_oauth.utils.requests.post")
+def test_exchange_code_for_tokens(mock_post):
+    """Test authorization code exchange."""
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {
+        "access_token": "test_access",
+        "refresh_token": "test_refresh",
+        "id_token": "test_id",
+    }
+    mock_post.return_value = mock_response
+
+    context = utils.prepare_oauth_context()
+    utils.assign_redirect_uri(context, 1455)
+
+    tokens = utils.exchange_code_for_tokens("test_code", context)
+    assert tokens is not None
+    assert tokens["access_token"] == "test_access"
+    assert "last_refresh" in tokens
+
+
+@patch("code_puppy.plugins.chatgpt_oauth.utils.requests.get")
+def test_fetch_chatgpt_models(mock_get):
+    """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 = {
+        "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_access_token", "test_account_id")
+    assert models is not None
+    # Required models always injected
+    assert "gpt-5.4" in models
+    assert "gpt-5.3-instant" in models
+    # API-returned models present too
+    assert "gpt-4o" in models
+    assert "gpt-3.5-turbo" in models
+    assert "o1-preview" in models
+    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 (including new required ones)
+    assert "gpt-5.4" in models
+    assert "gpt-5.3-instant" in models
+    assert "gpt-5.3-codex-spark" in models
+    assert "gpt-5.3-codex" in models
+    assert "gpt-5.2-codex" in models
+    assert "gpt-5.2" in models
+
+
+def test_add_models_to_chatgpt_config(tmp_path):
+    """Test adding models to chatgpt_models.json."""
+    with patch.object(
+        config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
+    ):
+        models = ["gpt-4o", "gpt-3.5-turbo"]
+
+        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"] == "chatgpt_oauth"
+        assert loaded["chatgpt-gpt-4o"]["name"] == "gpt-4o"
+        assert loaded["chatgpt-gpt-4o"]["oauth_source"] == "chatgpt-oauth-plugin"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])