code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/chatgpt_oauth/oauth_flow.py

@@ -0,0 +1,338 @@
+"""ChatGPT OAuth flow closely matching the ChatMock implementation."""
+
+import datetime
+import threading
+import time
+import urllib.parse
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any
+
+import httpx
+
+from code_muse.messaging import emit_error, emit_info, emit_success, emit_warning
+
+from ..oauth_muse_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://127.0.0.1:{REQUIRED_PORT}"
+
+
+@dataclass
+class TokenData:
+    id_token: str
+    access_token: str
+    refresh_token: str
+    account_id: str
+
+
+@dataclass
+class AuthBundle:
+    api_key: str | None
+    token_data: TokenData
+    last_refresh: str
+
+    @property
+    def success_url(self) -> str:
+        """The URL to redirect to after successful auth."""
+        return f"{URL_BASE}/success"
+
+    def __iter__(self):
+        """Allow unpacking as ``bundle, success_url = exchange_code(...)``."""
+        yield self
+        yield self.success_url
+
+
+class _OAuthServer(HTTPServer):
+    allow_reuse_address = True  # Reduce port-in-use flakes in tests
+
+    def __init__(
+        self,
+        *,
+        client_id: str,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(
+            ("127.0.0.1", 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) -> AuthBundle:
+        data = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": self.redirect_uri,
+            "client_id": self.client_id,
+            "code_verifier": self.context.code_verifier,
+        }
+
+        response = httpx.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 {}
+
+        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,
+        )
+
+        # Use access_token directly as the API key (ChatMock behaviour)
+        api_key = token_data.access_token
+
+        last_refresh = (
+            datetime.datetime.now(datetime.UTC).isoformat().replace("+00:00", "Z")
+        )
+        return AuthBundle(
+            api_key=api_key, token_data=token_data, last_refresh=last_refresh
+        )
+
+
+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 Muse.",
+            )
+            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 Muse 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
+
+        state = params.get("state", [None])[0]
+        if state != self.server.context.state:
+            self._send_failure(400, "Invalid state parameter — possible CSRF attack.")
+            self._shutdown()
+            return
+
+        try:
+            auth_bundle = 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 a token-free success URL
+            self._send_redirect(f"{URL_BASE}/success")
+        else:
+            self._send_failure(
+                500, "Unable to persist auth file — a Muse probably misplaced 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):
+            safe_args = []
+            for arg in args:
+                arg_str = str(arg)
+                if "?" in arg_str:
+                    parsed = urllib.parse.urlparse(arg_str)
+                    if parsed.query:
+                        arg_str = urllib.parse.urlunparse(
+                            parsed._replace(query="<redacted>")
+                        )
+                safe_args.append(arg_str)
+            super().log_message(fmt, *safe_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_muse.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 fetch_chatgpt_models
+
+        account_id = tokens.get("account_id", "")
+        models = fetch_chatgpt_models(api_key, account_id)
+        if models and add_models_to_extra_config(models):
+            emit_success(
+                "ChatGPT models registered. Use the `chatgpt-` prefix in /model."
+            )

code_muse/plugins/chatgpt_oauth/register_callbacks.py

@@ -0,0 +1,172 @@
+"""ChatGPT OAuth plugin callbacks aligned with ChatMock flow.
+
+Provides OAuth authentication for ChatGPT models and registers
+the 'chatgpt_oauth' model type handler.
+"""
+
+import os
+from typing import Any
+
+from code_muse.callbacks import register_callback
+from code_muse.messaging import emit_info, emit_success, emit_warning
+from code_muse.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) -> bool | None:
+    if not name:
+        return None
+
+    if name == "chatgpt-auth":
+        run_oauth_flow()
+        set_model_and_reload_agent("chatgpt-gpt-5.4")
+        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_muse.chatgpt_codex_client import create_codex_async_client
+    from code_muse.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
+    return OpenAIResponsesModel(model_name=model_config["name"], provider=provider)
+
+
+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)