From 98741dff72bdb66d32a21f42245a3c87c18fa8ba Mon Sep 17 00:00:00 2001 From: Jason Cook Date: Thu, 16 Apr 2026 20:30:10 -0400 Subject: [PATCH] feat(chatgpt, github_copilot): OAuth sign-in + token refresh in proxy UI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add "Sign in with ChatGPT" and "Sign in with GitHub Copilot" device-code OAuth flows to the Add Credential modal. Tokens persist as encrypted JSON in LiteLLM_CredentialsTable and are picked up at request time when a model's api_key is set to "oauth:"; a DBAuthenticator subclass reads from the in-memory credential cache (sync) and writes refreshed tokens back via a fire-and-forget worker thread. Per-row Refresh button rotates tokens on demand — ChatGPT via the IdP's refresh_token grant, Copilot by re-deriving the short-lived API key from the stored GitHub access token. Also ships a litellm-chatgpt-login CLI with an optional PKCE+loopback flow alongside the existing device-code path for local sign-in outside the UI. Provider-specific surface lives in new files under litellm/llms/{chatgpt,github_copilot}/db_authenticator.py and litellm/proxy/{chatgpt,copilot}_oauth_endpoints/. Shared-file delta is ~275 lines across 10 pre-existing files, the bulk being a pure append at the end of networking.tsx. Dispatch convention + operational caveats documented in the new "ChatGPT / Copilot OAuth Credentials" section of CLAUDE.md. Security: - /chatgpt/oauth/* and /copilot/oauth/* endpoints require PROXY_ADMIN (view-only admins excluded from write paths) - session_id query params are Query(...) annotated - session-cap check + slot reservation are atomic; reserved slot is cleaned up on device-code failure - PKCE loopback callback html.escape()s the IdP's error_description before rendering Tests: 165 cases covering DBAuthenticator read/write, config dispatch, endpoint admin-only, 429 cap, slot cleanup, background-worker success / auth-failure / DB-persist-failure paths, PKCE helpers + XSS regression, CLI broad-exception + KeyboardInterrupt. Co-Authored-By: Claude Opus 4.7 (1M context) --- CLAUDE.md | 9 + Makefile | 2 +- litellm/llms/chatgpt/authenticator.py | 4 + litellm/llms/chatgpt/cli.py | 65 +++ litellm/llms/chatgpt/db_authenticator.py | 208 ++++++++ litellm/llms/chatgpt/pkce.py | 259 ++++++++++ .../llms/chatgpt/responses/transformation.py | 12 +- .../github_copilot/chat/transformation.py | 11 +- .../llms/github_copilot/db_authenticator.py | 234 +++++++++ .../responses/transformation.py | 12 +- .../chatgpt_oauth_endpoints/endpoints.py | 309 ++++++++++++ .../copilot_oauth_endpoints/endpoints.py | 277 +++++++++++ litellm/proxy/proxy_server.py | 14 +- pyproject.toml | 1 + .../test_chatgpt_oauth_credential_dispatch.py | 68 +++ .../chatgpt/test_chatgpt_db_authenticator.py | 237 +++++++++ .../llms/chatgpt/test_chatgpt_pkce_login.py | 303 +++++++++++ .../test_copilot_db_authenticator.py | 281 +++++++++++ .../test_copilot_oauth_dispatch.py | 84 ++++ .../proxy/chatgpt_oauth_endpoints/__init__.py | 0 .../test_chatgpt_oauth_endpoints.py | 470 ++++++++++++++++++ .../proxy/copilot_oauth_endpoints/__init__.py | 0 .../test_copilot_oauth_endpoints.py | 324 ++++++++++++ .../model_add/AddCredentialModal.tsx | 35 +- .../model_add/ChatGPTLoginButton.tsx | 25 + .../model_add/CopilotLoginButton.tsx | 25 + .../model_add/OAuthDeviceLoginButton.tsx | 162 ++++++ .../src/components/model_add/credentials.tsx | 76 ++- .../src/components/networking.tsx | 177 +++++++ .../src/components/provider_info_helpers.tsx | 1 + 30 files changed, 3650 insertions(+), 35 deletions(-) create mode 100644 litellm/llms/chatgpt/cli.py create mode 100644 litellm/llms/chatgpt/db_authenticator.py create mode 100644 litellm/llms/chatgpt/pkce.py create mode 100644 litellm/llms/github_copilot/db_authenticator.py create mode 100644 litellm/proxy/chatgpt_oauth_endpoints/endpoints.py create mode 100644 litellm/proxy/copilot_oauth_endpoints/endpoints.py create mode 100644 tests/test_litellm/llms/chatgpt/responses/test_chatgpt_oauth_credential_dispatch.py create mode 100644 tests/test_litellm/llms/chatgpt/test_chatgpt_db_authenticator.py create mode 100644 tests/test_litellm/llms/chatgpt/test_chatgpt_pkce_login.py create mode 100644 tests/test_litellm/llms/github_copilot/test_copilot_db_authenticator.py create mode 100644 tests/test_litellm/llms/github_copilot/test_copilot_oauth_dispatch.py create mode 100644 tests/test_litellm/proxy/chatgpt_oauth_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/chatgpt_oauth_endpoints/test_chatgpt_oauth_endpoints.py create mode 100644 tests/test_litellm/proxy/copilot_oauth_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/copilot_oauth_endpoints/test_copilot_oauth_endpoints.py create mode 100644 ui/litellm-dashboard/src/components/model_add/ChatGPTLoginButton.tsx create mode 100644 ui/litellm-dashboard/src/components/model_add/CopilotLoginButton.tsx create mode 100644 ui/litellm-dashboard/src/components/model_add/OAuthDeviceLoginButton.tsx diff --git a/CLAUDE.md b/CLAUDE.md index 043055408c2..7bd8b1afeef 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -124,6 +124,15 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: - Always pass the raw `expires_at` timestamp to the client — never set it to `None` for expired credentials. Let the frontend compute the "Expired" display state from the timestamp. - Use `RecordNotFoundError` (not bare `except Exception`) when catching "already deleted" in credential delete endpoints. +### ChatGPT / Copilot OAuth Credentials +- Two providers use OAuth device-code flows with tokens stored in `LiteLLM_CredentialsTable`: `chatgpt` (ChatGPT / Codex backend) and `github_copilot`. Each has an `Authenticator` (filesystem-backed, used by CLI / direct SDK use) and a `DBAuthenticator` subclass (DB-backed, used by the proxy). +- **Request-time dispatch**: the convention is `api_key: "oauth:"` in the model's `litellm_params`. When the transformation sees this prefix it resolves via `resolve_authenticator(...)` in `llms/chatgpt/db_authenticator.py` (or `llms/github_copilot/db_authenticator.py`), which swaps in a `DBAuthenticator` bound to that credential. Any other `api_key` falls through to the filesystem authenticator. Do not invent alternative markers — reuse the `oauth:` prefix so both providers stay symmetric. +- **Copilot has two use-sites**: `_get_openai_compatible_provider_info` may rewrite `api_key` to the resolved Copilot key before `validate_environment` runs. The Copilot `resolve_authenticator` therefore takes both `api_key` and `litellm_params` and checks both; the ChatGPT version only needs `litellm_params`. +- **Credential payload shape**: store each OAuth field as a separate string key in `credential_values` (e.g. `access_token`, `refresh_token`, `id_token`, `account_id`, `expires_at`). The encryption helper (`encrypt_value_helper`) only encrypts strings — stringify ints (`expires_at`) before writing and parse back on read. Always set `credential_info = {"type": "chatgpt_oauth" | "copilot_oauth", "custom_llm_provider": "chatgpt" | "github_copilot"}` so the UI badge renders and the refresh button shows up for the right rows. +- **Refresh semantics differ**: ChatGPT's `POST /chatgpt/oauth/refresh` calls the IdP refresh-token endpoint and rotates `access_token` + `refresh_token` in the DB. Copilot's `POST /copilot/oauth/refresh` re-derives the short-lived Copilot API key from the stored GitHub access token (which itself doesn't expire) and caches the result in-process — no DB write on that path. +- **Session state is in-memory per replica**. The device-code `/start` → poll → `/status` sequence uses a module-level `_sessions` dict. For multi-replica deploys, enable sticky sessions on the UI→proxy path during the ~1-minute login window, or run the admin UI against a single-replica node. +- **The refresh path from `DBAuthenticator._write_auth_file` is fire-and-forget**. The in-memory `litellm.credential_list` is updated synchronously; DB persist runs on a worker thread that creates a fresh event loop via `asyncio.run`. If the DB write fails, the in-process cache still has fresh tokens — next request is fine; reconciliation happens on the next successful write. + ### Browser Storage Safety (UI) - Never write LiteLLM access tokens or API keys to `localStorage` — use `sessionStorage` only. `localStorage` survives browser close and is readable by any injected script (XSS). - Shared utility functions (e.g. `extractErrorMessage`) belong in `src/utils/` — never define them inline in hooks or duplicate them across files. diff --git a/Makefile b/Makefile index b6b674ff3b1..c0a8d827e98 100644 --- a/Makefile +++ b/Makefile @@ -146,7 +146,7 @@ test-unit-proxy-core: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20 test-unit-proxy-misc: install-test-deps - $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/chatgpt_oauth_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/copilot_oauth_endpoints tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 test-unit-integrations: install-test-deps $(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20 diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index e35b04a3fb3..a75399a6324 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -21,6 +21,7 @@ from .common_utils import ( GetDeviceCodeError, RefreshAccessTokenError, ) +from .pkce import login_pkce as _login_pkce_flow TOKEN_EXPIRY_SKEW_SECONDS = 60 DEVICE_CODE_TIMEOUT_SECONDS = 15 * 60 @@ -341,6 +342,9 @@ class Authenticator: self._write_auth_file(auth_data) return refreshed + def login_pkce(self, **kwargs: Any) -> Dict[str, str]: + return _login_pkce_flow(self, **kwargs) + def _build_auth_record(self, tokens: Dict[str, str]) -> Dict[str, Any]: access_token = tokens.get("access_token") id_token = tokens.get("id_token") diff --git a/litellm/llms/chatgpt/cli.py b/litellm/llms/chatgpt/cli.py new file mode 100644 index 00000000000..dc0994ce4ea --- /dev/null +++ b/litellm/llms/chatgpt/cli.py @@ -0,0 +1,65 @@ +"""CLI entrypoint for logging in to ChatGPT / Codex backend via OAuth.""" + +import argparse +import sys + +from .authenticator import Authenticator +from .common_utils import ChatGPTAuthError +from .pkce import REDIRECT_PORT as PKCE_DEFAULT_PORT + + +def cli() -> int: + parser = argparse.ArgumentParser( + prog="litellm-chatgpt-login", + description=( + "Sign in to the ChatGPT / Codex backend and store OAuth credentials " + "for use with `litellm.responses(model='chatgpt/...')`." + ), + ) + parser.add_argument( + "--method", + choices=["device", "pkce"], + default="device", + help=( + "OAuth flow to use. `device` (default) shows a code to enter in a " + "browser (works over SSH). `pkce` opens a browser to a loopback " + "redirect (one-click, requires a local browser)." + ), + ) + parser.add_argument( + "--port", + type=int, + default=PKCE_DEFAULT_PORT, + help="Loopback port for the PKCE redirect (pkce only).", + ) + parser.add_argument( + "--no-browser", + action="store_true", + help="Do not try to auto-open a browser for the PKCE flow.", + ) + args = parser.parse_args() + + auth = Authenticator() + try: + if args.method == "pkce": + auth.login_pkce(open_browser=not args.no_browser, port=args.port) + else: + auth._login_device_code() + except ChatGPTAuthError as exc: + print(f"Login failed: {exc.message}", file=sys.stderr) # noqa: T201 + return 1 + except KeyboardInterrupt: + print("\nLogin cancelled.", file=sys.stderr) # noqa: T201 + return 130 + except Exception as exc: # noqa: BLE001 + # Any other failure mode (network, filesystem, etc.) — surface a + # clean message instead of a traceback. + print(f"Login failed: {exc}", file=sys.stderr) # noqa: T201 + return 1 + + print(f"Saved credentials to {auth.auth_file}") # noqa: T201 + return 0 + + +if __name__ == "__main__": + sys.exit(cli()) diff --git a/litellm/llms/chatgpt/db_authenticator.py b/litellm/llms/chatgpt/db_authenticator.py new file mode 100644 index 00000000000..e1abed66b84 --- /dev/null +++ b/litellm/llms/chatgpt/db_authenticator.py @@ -0,0 +1,208 @@ +""" +DB-backed ChatGPT / Codex OAuth authenticator. + +Mirrors the filesystem-backed :class:`Authenticator` but reads tokens from +``litellm.credential_list`` (the in-memory decrypted cache) and writes them +back to ``LiteLLM_CredentialsTable`` via a fire-and-forget background thread. + +Kept in its own module so the upstream-facing ``authenticator.py`` (which +houses the device-code flow shared with the CLI) stays nearly untouched. +""" + +import asyncio +import threading +from typing import Any, Dict, Optional + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.types.utils import CredentialItem + +from .authenticator import Authenticator + +CREDENTIAL_TYPE = "chatgpt_oauth" + +# When ``api_key`` begins with this prefix, the suffix is a credential_name +# in ``LiteLLM_CredentialsTable`` and a DB-backed authenticator is used in +# place of the filesystem one. +OAUTH_CREDENTIAL_API_KEY_PREFIX = "oauth:" + + +class DBAuthenticator(Authenticator): + """ + Uses ``LiteLLM_CredentialsTable`` as the backing store instead of + ``~/.config/litellm/chatgpt/auth.json``. The in-memory + ``litellm.credential_list`` serves reads (sync, called from + :meth:`validate_environment`); writes fan out to both the in-memory cache + (sync) and the database (async, via a worker thread so the sync refresh + path need not await). + """ + + def __init__(self, credential_name: str) -> None: + self.credential_name = credential_name + # Parent fields that aren't used by this subclass — kept non-None so + # any accidental reference fails loudly instead of silently reading + # or writing the wrong file. + self.token_dir = "" + self.auth_file = "" + + def _ensure_token_dir(self) -> None: + return + + def _read_auth_file(self) -> Optional[Dict[str, Any]]: + values = CredentialAccessor.get_credential_values(self.credential_name) + if not values: + return None + return _unpack_auth_record(values) + + def _write_auth_file(self, data: Dict[str, Any]) -> None: + credential_values = _pack_auth_record(data) + item = CredentialItem( + credential_name=self.credential_name, + credential_values=credential_values, + credential_info={ + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "chatgpt", + }, + ) + CredentialAccessor.upsert_credentials([item]) + _schedule_db_persist(item) + + +# --------------------------------------------------------------------------- +# Auth-record (de)serialization +# --------------------------------------------------------------------------- +# +# ``credential_values`` is a flat ``dict[str, str]`` — the encryption helper +# only handles strings, so ``expires_at`` (an int) is stored as a string and +# parsed back here. + + +def _pack_auth_record(data: Dict[str, Any]) -> Dict[str, str]: + packed: Dict[str, str] = {} + for key in ("access_token", "refresh_token", "id_token", "account_id"): + value = data.get(key) + if value is not None: + packed[key] = str(value) + expires_at = data.get("expires_at") + if expires_at is not None: + packed["expires_at"] = str(int(expires_at)) + return packed + + +def _unpack_auth_record(values: Dict[str, Any]) -> Dict[str, Any]: + record: Dict[str, Any] = {} + for key in ("access_token", "refresh_token", "id_token", "account_id"): + value = values.get(key) + if value: + record[key] = value + expires_at = values.get("expires_at") + if expires_at is not None: + try: + record["expires_at"] = int(expires_at) + except (TypeError, ValueError): + pass + return record + + +# --------------------------------------------------------------------------- +# Fire-and-forget DB persistence +# --------------------------------------------------------------------------- + + +def _schedule_db_persist(item: CredentialItem) -> None: + """ + Persist the credential to ``LiteLLM_CredentialsTable`` on a worker thread. + + This is fire-and-forget: the caller (typically ``_refresh_tokens``) runs + in sync context and has already updated the in-memory cache, so the + request-in-flight can proceed with fresh tokens even if the DB write + lags or fails. Failures are logged; the cache will be reconciled on the + next successful write. + """ + thread = threading.Thread( + target=_persist_item_sync, + args=(item,), + daemon=True, + name="chatgpt-oauth-persist", + ) + thread.start() + + +def _persist_item_sync(item: CredentialItem) -> None: + try: + asyncio.run(persist_credential_to_db(item)) + except Exception as exc: + verbose_logger.error( + "Failed to persist refreshed ChatGPT OAuth credential %s: %s", + item.credential_name, + exc, + ) + + +async def persist_credential_to_db(item: CredentialItem) -> None: + """ + Encrypt and upsert the credential row. + + Intended for both the fire-and-forget refresh path and the interactive + login endpoint (which can ``await`` it directly on the request event + loop). + """ + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + verbose_logger.debug( + "prisma_client unavailable; skipping DB persist for %s", + item.credential_name, + ) + return + + encrypted_values = { + k: encrypt_value_helper(v) for k, v in item.credential_values.items() + } + credential_info = item.credential_info or { + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "chatgpt", + } + await prisma_client.db.litellm_credentialstable.upsert( + where={"credential_name": item.credential_name}, + data={ + "create": { + "credential_name": item.credential_name, + "credential_values": encrypted_values, + "credential_info": credential_info, + "created_by": "chatgpt_oauth_flow", + "updated_by": "chatgpt_oauth_flow", + }, + "update": { + "credential_values": encrypted_values, + "credential_info": credential_info, + "updated_by": "chatgpt_oauth_flow", + }, + }, + ) + + +def resolve_authenticator( + litellm_params: Any, + fallback: Authenticator, +) -> Authenticator: + """ + If ``litellm_params.api_key`` starts with ``oauth:``, returns a + :class:`DBAuthenticator` for the named credential. Otherwise returns + the given fallback (typically the filesystem :class:`Authenticator`). + """ + if litellm_params is None: + return fallback + api_key = ( + litellm_params.get("api_key") + if isinstance(litellm_params, dict) + else getattr(litellm_params, "api_key", None) + ) + if isinstance(api_key, str) and api_key.startswith( + OAUTH_CREDENTIAL_API_KEY_PREFIX + ): + return DBAuthenticator( + credential_name=api_key[len(OAUTH_CREDENTIAL_API_KEY_PREFIX) :] + ) + return fallback diff --git a/litellm/llms/chatgpt/pkce.py b/litellm/llms/chatgpt/pkce.py new file mode 100644 index 00000000000..a2b7ed656f0 --- /dev/null +++ b/litellm/llms/chatgpt/pkce.py @@ -0,0 +1,259 @@ +""" +Authorization Code + PKCE OAuth flow for the ChatGPT / Codex backend. + +Kept in its own module so the upstream-facing ``authenticator.py`` (which +already hosts the device-code flow) stays nearly untouched — minimising +merge conflicts when syncing from BerriAI/litellm. +""" + +import base64 +import hashlib +import html +import http.server +import secrets +import threading +import urllib.parse +import webbrowser +from typing import TYPE_CHECKING, Any, Dict + +import httpx + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import _get_httpx_client + +from .common_utils import ( + CHATGPT_AUTH_BASE, + CHATGPT_CLIENT_ID, + CHATGPT_OAUTH_TOKEN_URL, + GetAccessTokenError, + get_chatgpt_originator, +) + +if TYPE_CHECKING: + from .authenticator import Authenticator + + +OAUTH_AUTHORIZE_URL = f"{CHATGPT_AUTH_BASE}/oauth/authorize" +REDIRECT_HOST = "127.0.0.1" +REDIRECT_PORT = 1455 +REDIRECT_PATH = "/auth/callback" +SCOPE = "openid profile email offline_access" +LOGIN_TIMEOUT_SECONDS = 10 * 60 + +_SUCCESS_HTML = ( + "LiteLLM login complete" + '' + "

Sign-in complete

" + "

You can close this tab and return to your terminal.

" + "" +) +_ERROR_HTML = ( + "LiteLLM login failed" + '' + "

Sign-in failed

" + "

{message}

" + "

Return to your terminal for details.

" + "" +) + + +def login_pkce( + authenticator: "Authenticator", + open_browser: bool = True, + port: int = REDIRECT_PORT, + timeout_seconds: float = LOGIN_TIMEOUT_SECONDS, +) -> Dict[str, str]: + """ + Run the full PKCE + loopback OAuth flow and persist tokens via the + authenticator (reusing ``_build_auth_record``/``_write_auth_file`` so the + on-disk format matches the device-code flow). + """ + code_verifier = _generate_code_verifier() + code_challenge = _generate_code_challenge(code_verifier) + state = secrets.token_hex(16) + redirect_uri = f"http://{REDIRECT_HOST}:{port}{REDIRECT_PATH}" + authorize_url = _build_authorize_url( + redirect_uri=redirect_uri, + code_challenge=code_challenge, + state=state, + ) + + result: Dict[str, Any] = {} + completed = threading.Event() + handler_cls = _make_handler(result, completed, expected_state=state) + + try: + server = http.server.HTTPServer((REDIRECT_HOST, port), handler_cls) + except OSError as exc: + raise GetAccessTokenError( + message=( + f"Failed to bind loopback server on {REDIRECT_HOST}:{port}: {exc}. " + "Pass a different port or close the process holding it." + ), + status_code=400, + ) + + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + print( # noqa: T201 + "Sign in with ChatGPT in your browser.\n" + f"If the browser does not open, visit:\n{authorize_url}", + flush=True, + ) + if open_browser: + try: + webbrowser.open(authorize_url) + except Exception as exc: + verbose_logger.debug("webbrowser.open failed: %s", exc) + + if not completed.wait(timeout=timeout_seconds): + raise GetAccessTokenError( + message="Timed out waiting for OAuth callback", + status_code=408, + ) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + if "error" in result: + raise GetAccessTokenError( + message=f"OAuth callback error: {result['error']}", + status_code=400, + ) + code = result.get("code") + if not code: + raise GetAccessTokenError( + message="OAuth callback did not include an authorization code", + status_code=400, + ) + + tokens = _exchange_code_for_tokens( + code=code, code_verifier=code_verifier, redirect_uri=redirect_uri + ) + auth_data = authenticator._build_auth_record(tokens) + authenticator._write_auth_file(auth_data) + return tokens + + +def _exchange_code_for_tokens( + code: str, code_verifier: str, redirect_uri: str +) -> Dict[str, str]: + try: + client = _get_httpx_client() + resp = client.post( + CHATGPT_OAUTH_TOKEN_URL, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data={ + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": CHATGPT_CLIENT_ID, + "code_verifier": code_verifier, + }, + ) + resp.raise_for_status() + data = resp.json() + except httpx.HTTPStatusError as exc: + raise GetAccessTokenError( + message=f"PKCE token exchange failed: {exc}", + status_code=exc.response.status_code, + ) + except Exception as exc: + raise GetAccessTokenError( + message=f"PKCE token exchange failed: {exc}", + status_code=400, + ) + + if not all(key in data for key in ("access_token", "refresh_token", "id_token")): + raise GetAccessTokenError( + message=f"PKCE token response missing fields: {data}", + status_code=400, + ) + return { + "access_token": data["access_token"], + "refresh_token": data["refresh_token"], + "id_token": data["id_token"], + } + + +def _generate_code_verifier() -> str: + """RFC 7636 PKCE verifier: 43-128 chars from the unreserved set.""" + return secrets.token_urlsafe(64)[:64] + + +def _generate_code_challenge(verifier: str) -> str: + """S256 challenge: base64url(sha256(verifier)) without padding.""" + digest = hashlib.sha256(verifier.encode("ascii")).digest() + return base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + + +def _build_authorize_url(redirect_uri: str, code_challenge: str, state: str) -> str: + params = { + "response_type": "code", + "client_id": CHATGPT_CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": SCOPE, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + "state": state, + "id_token_add_organizations": "true", + "codex_cli_simplified_flow": "true", + "originator": get_chatgpt_originator(), + } + return f"{OAUTH_AUTHORIZE_URL}?{urllib.parse.urlencode(params)}" + + +def _make_handler( + result: Dict[str, Any], + completed: threading.Event, + expected_state: str, +) -> type: + class _Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 + parsed = urllib.parse.urlparse(self.path) + if parsed.path != REDIRECT_PATH: + self.send_response(404) + self.end_headers() + return + params = urllib.parse.parse_qs(parsed.query) + error = params.get("error", [None])[0] + if error: + result["error"] = params.get("error_description", [error])[0] + self._respond_error(result["error"]) + completed.set() + return + state = params.get("state", [None])[0] + code = params.get("code", [None])[0] + if state != expected_state: + result["error"] = "state mismatch" + self._respond_error("state mismatch") + completed.set() + return + if not code: + result["error"] = "missing code" + self._respond_error("missing code") + completed.set() + return + result["code"] = code + self._respond_html(200, _SUCCESS_HTML) + completed.set() + + def _respond_error(self, message: str) -> None: + # ``message`` can originate from the IdP's ``error_description`` + # query param, so escape before interpolating into HTML. + self._respond_html(400, _ERROR_HTML.format(message=html.escape(message))) + + def _respond_html(self, status: int, body: str) -> None: + payload = 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(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args: Any, **kwargs: Any) -> None: + return + + return _Handler diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 3c59ca16581..c26ad634423 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -25,6 +25,7 @@ from ..common_utils import ( get_chatgpt_default_headers, get_chatgpt_default_instructions, ) +from ..db_authenticator import resolve_authenticator class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): @@ -42,8 +43,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): model: str, litellm_params: Optional[GenericLiteLLMParams], ) -> dict: + authenticator = resolve_authenticator(litellm_params, self.authenticator) try: - access_token = self.authenticator.get_access_token() + access_token = authenticator.get_access_token() except GetAccessTokenError as e: raise AuthenticationError( model=model, @@ -51,7 +53,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): message=str(e), ) - account_id = self.authenticator.get_account_id() + account_id = authenticator.get_account_id() session_id = ensure_chatgpt_session_id(litellm_params) default_headers = get_chatgpt_default_headers( access_token, account_id, session_id @@ -77,9 +79,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): existing_instructions = request.get("instructions") if existing_instructions: if base_instructions not in existing_instructions: - request[ - "instructions" - ] = f"{base_instructions}\n\n{existing_instructions}" + request["instructions"] = ( + f"{base_instructions}\n\n{existing_instructions}" + ) else: request["instructions"] = base_instructions request["store"] = False diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index be8ad7d0877..4aff1025056 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -11,6 +11,7 @@ from ..common_utils import ( GetAPIKeyError, get_copilot_default_headers, ) +from ..db_authenticator import resolve_authenticator class GithubCopilotConfig(OpenAIConfig): @@ -30,9 +31,10 @@ class GithubCopilotConfig(OpenAIConfig): api_key: Optional[str], custom_llm_provider: str, ) -> Tuple[Optional[str], Optional[str], str]: - dynamic_api_base = self.authenticator.get_api_base() or GITHUB_COPILOT_API_BASE + authenticator = resolve_authenticator(api_key, None, self.authenticator) + dynamic_api_base = authenticator.get_api_base() or GITHUB_COPILOT_API_BASE try: - dynamic_api_key = self.authenticator.get_api_key() + dynamic_api_key = authenticator.get_api_key() except GetAPIKeyError as e: raise AuthenticationError( model=model, @@ -83,8 +85,11 @@ class GithubCopilotConfig(OpenAIConfig): ) # Add Copilot-specific headers (editor-version, user-agent, etc.) + authenticator = resolve_authenticator( + api_key, litellm_params, self.authenticator + ) try: - copilot_api_key = self.authenticator.get_api_key() + copilot_api_key = authenticator.get_api_key() copilot_headers = get_copilot_default_headers(copilot_api_key) validated_headers = {**copilot_headers, **validated_headers} except GetAPIKeyError: diff --git a/litellm/llms/github_copilot/db_authenticator.py b/litellm/llms/github_copilot/db_authenticator.py new file mode 100644 index 00000000000..866ec3f3d31 --- /dev/null +++ b/litellm/llms/github_copilot/db_authenticator.py @@ -0,0 +1,234 @@ +""" +DB-backed GitHub Copilot authenticator. + +Mirrors the filesystem-backed :class:`Authenticator` but reads the long-lived +GitHub OAuth access token from ``litellm.credential_list`` (the in-memory +decrypted cache) and persists it to ``LiteLLM_CredentialsTable`` via a +fire-and-forget background thread. + +The short-lived Copilot API key (obtained from +``/copilot_internal/v2/token``) is cached in-memory per-process keyed by +credential name — it rotates frequently (~30 min), is cheap to refresh, and +cross-replica coherence is unnecessary. +""" + +import asyncio +import threading +from datetime import datetime +from typing import Any, Dict, Optional + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.types.utils import CredentialItem + +from .authenticator import Authenticator +from .common_utils import GetAccessTokenError, GetAPIKeyError + +CREDENTIAL_TYPE = "copilot_oauth" + + +class DBAuthenticator(Authenticator): + """ + Uses ``LiteLLM_CredentialsTable`` as the backing store for the GitHub + access token. The Copilot API key lives only in a per-process cache + (``_api_key_cache``). + """ + + # Shared across instances so multiple requests for the same credential + # reuse the same cached API key payload without re-hitting GitHub. + _api_key_cache: Dict[str, Dict[str, Any]] = {} + _api_key_cache_lock = threading.Lock() + + def __init__(self, credential_name: str) -> None: + self.credential_name = credential_name + # Parent fields that aren't used by this subclass — kept non-None so + # accidental references fail loudly rather than silently hitting disk. + self.token_dir = "" + self.access_token_file = "" + self.api_key_file = "" + + def _ensure_token_dir(self) -> None: + return + + def get_access_token(self) -> str: + values = CredentialAccessor.get_credential_values(self.credential_name) + token = values.get("access_token") if values else None + if not token: + raise GetAccessTokenError( + message=( + f"No GitHub access token stored for credential " + f"'{self.credential_name}'. Sign in via the UI first." + ), + status_code=401, + ) + return token + + def get_api_key(self) -> str: + cached = self._get_cached_api_key() + if cached is not None: + return cached + info = self._refresh_api_key() + self._cache_api_key(info) + token = info.get("token") + if not token: + raise GetAPIKeyError( + message="API key response missing token", + status_code=401, + ) + return token + + def get_api_base(self) -> Optional[str]: + with self._api_key_cache_lock: + info = self._api_key_cache.get(self.credential_name) + if info is None: + return None + endpoints = info.get("endpoints") or {} + return endpoints.get("api") + + def force_refresh_api_key(self) -> Dict[str, Any]: + """ + Force a call to ``/copilot_internal/v2/token`` even if the cached + key is still valid. Used by the UI's "Refresh" button. + """ + info = self._refresh_api_key() + self._cache_api_key(info) + return info + + def store_access_token(self, access_token: str) -> None: + """ + Called by the OAuth login flow to persist a freshly-obtained GitHub + access token. Writes to the in-memory credential cache and schedules + a DB write. + """ + item = CredentialItem( + credential_name=self.credential_name, + credential_values={"access_token": access_token}, + credential_info={ + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "github_copilot", + }, + ) + CredentialAccessor.upsert_credentials([item]) + # Invalidate any cached API key tied to an old access token. + with self._api_key_cache_lock: + self._api_key_cache.pop(self.credential_name, None) + _schedule_db_persist(item) + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _get_cached_api_key(self) -> Optional[str]: + with self._api_key_cache_lock: + info = self._api_key_cache.get(self.credential_name) + if info is None: + return None + if info.get("expires_at", 0) <= datetime.now().timestamp(): + return None + return info.get("token") + + def _cache_api_key(self, info: Dict[str, Any]) -> None: + with self._api_key_cache_lock: + self._api_key_cache[self.credential_name] = info + + +# --------------------------------------------------------------------------- +# Fire-and-forget DB persistence +# --------------------------------------------------------------------------- + + +def _schedule_db_persist(item: CredentialItem) -> None: + thread = threading.Thread( + target=_persist_item_sync, + args=(item,), + daemon=True, + name="copilot-oauth-persist", + ) + thread.start() + + +def _persist_item_sync(item: CredentialItem) -> None: + try: + asyncio.run(persist_credential_to_db(item)) + except Exception as exc: + verbose_logger.error( + "Failed to persist Copilot OAuth credential %s: %s", + item.credential_name, + exc, + ) + + +async def persist_credential_to_db(item: CredentialItem) -> None: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + verbose_logger.debug( + "prisma_client unavailable; skipping DB persist for %s", + item.credential_name, + ) + return + + encrypted_values = { + k: encrypt_value_helper(v) for k, v in item.credential_values.items() + } + credential_info = item.credential_info or { + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "github_copilot", + } + await prisma_client.db.litellm_credentialstable.upsert( + where={"credential_name": item.credential_name}, + data={ + "create": { + "credential_name": item.credential_name, + "credential_values": encrypted_values, + "credential_info": credential_info, + "created_by": "copilot_oauth_flow", + "updated_by": "copilot_oauth_flow", + }, + "update": { + "credential_values": encrypted_values, + "credential_info": credential_info, + "updated_by": "copilot_oauth_flow", + }, + }, + ) + + +# --------------------------------------------------------------------------- +# api_key-prefix dispatch helper +# --------------------------------------------------------------------------- + +OAUTH_CREDENTIAL_API_KEY_PREFIX = "oauth:" + + +def resolve_authenticator( + api_key: Optional[str], + litellm_params: Any, + fallback: Authenticator, +) -> Authenticator: + """ + If ``api_key`` (or ``litellm_params.api_key``) starts with ``oauth:``, + the suffix names a credential in ``LiteLLM_CredentialsTable`` and this + returns a :class:`DBAuthenticator` for it. Otherwise returns the given + fallback (typically the filesystem-backed :class:`Authenticator`). + + Two sources are checked because upstream may rewrite ``api_key`` to the + resolved Copilot token before ``validate_environment`` runs, while the + raw marker survives on ``litellm_params``. + """ + candidates = [api_key] + if litellm_params is not None: + if isinstance(litellm_params, dict): + candidates.append(litellm_params.get("api_key")) + else: + candidates.append(getattr(litellm_params, "api_key", None)) + + for candidate in candidates: + if isinstance(candidate, str) and candidate.startswith( + OAUTH_CREDENTIAL_API_KEY_PREFIX + ): + return DBAuthenticator( + credential_name=candidate[len(OAUTH_CREDENTIAL_API_KEY_PREFIX) :] + ) + return fallback diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index 46efc124b1d..633f7e2ad97 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -7,6 +7,7 @@ which is required for models like gpt-5.1-codex that only support the /responses Implementation based on analysis of the copilot-api project by caozhiyuan: https://github.com/caozhiyuan/copilot-api """ + from typing import TYPE_CHECKING, Any, Dict, Optional, Union from litellm._logging import verbose_logger @@ -26,6 +27,7 @@ from ..common_utils import ( GetAPIKeyError, get_copilot_default_headers, ) +from ..db_authenticator import resolve_authenticator if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -104,7 +106,10 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ try: # Get GitHub Copilot API key via OAuth - api_key = self.authenticator.get_api_key() + authenticator = resolve_authenticator( + None, litellm_params, self.authenticator + ) + api_key = authenticator.get_api_key() if not api_key: raise AuthenticationError( @@ -165,9 +170,8 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): added in the future by detecting account type. """ # Use provided api_base or fall back to authenticator's base or default - api_base = ( - api_base or self.authenticator.get_api_base() or GITHUB_COPILOT_API_BASE - ) + authenticator = resolve_authenticator(None, litellm_params, self.authenticator) + api_base = api_base or authenticator.get_api_base() or GITHUB_COPILOT_API_BASE # Remove trailing slashes api_base = api_base.rstrip("/") diff --git a/litellm/proxy/chatgpt_oauth_endpoints/endpoints.py b/litellm/proxy/chatgpt_oauth_endpoints/endpoints.py new file mode 100644 index 00000000000..a2e5faa2f96 --- /dev/null +++ b/litellm/proxy/chatgpt_oauth_endpoints/endpoints.py @@ -0,0 +1,309 @@ +""" +Admin-only endpoints that let the LiteLLM UI sign in to the ChatGPT / Codex +backend using the OpenAI device-code OAuth flow and persist the resulting +tokens in ``LiteLLM_CredentialsTable``. + +The browser-based PKCE flow redirects to ``127.0.0.1:1455`` on the user's +machine, which does not work for a remotely-hosted UI — device code is the +correct flow here: the user's browser visits ``auth.openai.com/codex/device``, +enters a code, and the proxy polls the token endpoint in the background. + +Endpoints: + POST /chatgpt/oauth/start — start a new flow; returns user_code + verification URL + GET /chatgpt/oauth/status — poll session status (pending/success/error) + POST /chatgpt/oauth/cancel — abandon an in-flight flow + +Session state is held in-memory per proxy replica. For a multi-replica +deploy, enable sticky sessions for the UI → proxy path or use a +single-replica admin node. +""" + +import asyncio +import threading +import time +import uuid +from typing import Any, Dict, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.llms.chatgpt.authenticator import Authenticator +from litellm.llms.chatgpt.common_utils import ( + CHATGPT_DEVICE_VERIFY_URL, + ChatGPTAuthError, +) +from litellm.llms.chatgpt.db_authenticator import ( + CREDENTIAL_TYPE, + DBAuthenticator, + persist_credential_to_db, +) +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.utils import CredentialItem + +router = APIRouter(prefix="/chatgpt/oauth", tags=["chatgpt oauth"]) + +SESSION_TTL_SECONDS = 20 * 60 +SESSIONS_MAX_SIZE = 100 + +_sessions: Dict[str, Dict[str, Any]] = {} +_sessions_lock = threading.Lock() + + +class StartRequest(BaseModel): + credential_name: str + + +class StartResponse(BaseModel): + session_id: str + user_code: str + verification_url: str + interval: int + + +class StatusResponse(BaseModel): + status: str # "pending" | "success" | "error" | "cancelled" + credential_name: Optional[str] = None + message: Optional[str] = None + + +class RefreshRequest(BaseModel): + credential_name: str + + +class RefreshResponse(BaseModel): + credential_name: str + expires_at: Optional[int] = None + + +def _require_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + # These endpoints write to LiteLLM_CredentialsTable (start → insert, + # refresh → rotate). PROXY_ADMIN_VIEW_ONLY must not reach them. + role = getattr(user_api_key_dict, "user_role", None) + if role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only PROXY_ADMIN may initiate ChatGPT OAuth flows.", + ) + + +def _purge_expired_sessions() -> None: + now = time.time() + with _sessions_lock: + expired = [k for k, v in _sessions.items() if v.get("expires_at", 0) < now] + for k in expired: + del _sessions[k] + + +def _update_session(session_id: str, **fields: Any) -> None: + with _sessions_lock: + entry = _sessions.get(session_id) + if entry is None: + return + entry.update(fields) + + +def _get_session(session_id: str) -> Optional[Dict[str, Any]]: + with _sessions_lock: + entry = _sessions.get(session_id) + return dict(entry) if entry else None + + +@router.post( + "/start", + response_model=StartResponse, + dependencies=[Depends(user_api_key_auth)], +) +async def start_oauth( + body: StartRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> StartResponse: + _require_admin(user_api_key_dict) + _purge_expired_sessions() + + # Atomically reserve a slot so concurrent callers cannot all pass the + # capacity check and blow past SESSIONS_MAX_SIZE. + session_id = uuid.uuid4().hex + now = time.time() + with _sessions_lock: + if len(_sessions) >= SESSIONS_MAX_SIZE: + raise HTTPException( + status_code=429, + detail="Too many in-flight OAuth sessions. Retry shortly.", + ) + _sessions[session_id] = { + "status": "starting", + "credential_name": body.credential_name, + "started_by": user_api_key_dict.user_id, + "started_at": now, + "expires_at": now + SESSION_TTL_SECONDS, + "cancelled": False, + } + + authenticator = Authenticator() + try: + device_code = authenticator._request_device_code() + except ChatGPTAuthError as exc: + with _sessions_lock: + _sessions.pop(session_id, None) + raise HTTPException(status_code=502, detail=exc.message) + + with _sessions_lock: + entry = _sessions.get(session_id) + if entry is None: + # Cancelled or evicted while the device-code call was in flight. + raise HTTPException(status_code=410, detail="Session was cancelled") + entry["status"] = "pending" + entry["device_code"] = device_code + + thread = threading.Thread( + target=_run_device_code_flow, + args=(session_id, body.credential_name, device_code, authenticator), + daemon=True, + name=f"chatgpt-oauth-{session_id[:8]}", + ) + thread.start() + + return StartResponse( + session_id=session_id, + user_code=device_code["user_code"], + verification_url=CHATGPT_DEVICE_VERIFY_URL, + interval=int(device_code.get("interval", "5")), + ) + + +@router.get( + "/status", + response_model=StatusResponse, + dependencies=[Depends(user_api_key_auth)], +) +async def oauth_status( + session_id: str = Query( + ..., description="Session ID returned by /chatgpt/oauth/start" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> StatusResponse: + _require_admin(user_api_key_dict) + entry = _get_session(session_id) + if entry is None: + raise HTTPException(status_code=404, detail="Unknown or expired session_id") + return StatusResponse( + status=entry["status"], + credential_name=entry.get("credential_name"), + message=entry.get("message"), + ) + + +@router.post( + "/cancel", + dependencies=[Depends(user_api_key_auth)], +) +async def oauth_cancel( + session_id: str = Query( + ..., description="Session ID returned by /chatgpt/oauth/start" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Dict[str, bool]: + _require_admin(user_api_key_dict) + with _sessions_lock: + entry = _sessions.get(session_id) + if entry is None: + raise HTTPException(status_code=404, detail="Unknown or expired session_id") + entry["cancelled"] = True + if entry["status"] == "pending": + entry["status"] = "cancelled" + entry["message"] = "Cancelled by user" + return {"success": True} + + +@router.post( + "/refresh", + response_model=RefreshResponse, + dependencies=[Depends(user_api_key_auth)], +) +async def oauth_refresh( + body: RefreshRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> RefreshResponse: + """ + Force a refresh of the OAuth access token using the stored refresh + token. Rotates the refresh token too if the IdP issues a new one. + """ + _require_admin(user_api_key_dict) + db_auth = DBAuthenticator(credential_name=body.credential_name) + auth_data = db_auth._read_auth_file() + if not auth_data or not auth_data.get("refresh_token"): + raise HTTPException( + status_code=404, + detail=( + f"No refresh_token stored for credential '{body.credential_name}'. " + "Re-run the sign-in flow." + ), + ) + try: + db_auth._refresh_tokens(auth_data["refresh_token"]) + except ChatGPTAuthError as exc: + raise HTTPException(status_code=502, detail=exc.message) + # _refresh_tokens already wrote back via DBAuthenticator._write_auth_file, + # which scheduled the DB persist. Re-read to get the canonical expires_at. + fresh = db_auth._read_auth_file() or {} + return RefreshResponse( + credential_name=body.credential_name, + expires_at=fresh.get("expires_at"), + ) + + +def _run_device_code_flow( + session_id: str, + credential_name: str, + device_code: Dict[str, str], + authenticator: Authenticator, +) -> None: + """ + Background worker: polls for the authorization code, exchanges it for + tokens, then upserts the credential into the in-memory cache and DB. + """ + try: + auth_code = authenticator._poll_for_authorization_code(device_code) + session_snapshot = _get_session(session_id) + if session_snapshot is None or session_snapshot.get("cancelled"): + return + tokens = authenticator._exchange_code_for_tokens(auth_code) + except ChatGPTAuthError as exc: + _update_session(session_id, status="error", message=exc.message) + return + except Exception as exc: # pragma: no cover - defensive + verbose_proxy_logger.exception("Unexpected error in ChatGPT OAuth flow") + _update_session(session_id, status="error", message=str(exc)) + return + + auth_record = authenticator._build_auth_record(tokens) + credential_values = {k: str(v) for k, v in auth_record.items() if v is not None} + item = CredentialItem( + credential_name=credential_name, + credential_values=credential_values, + credential_info={ + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "chatgpt", + }, + ) + CredentialAccessor.upsert_credentials([item]) + + try: + asyncio.run(persist_credential_to_db(item)) + except Exception as exc: + verbose_proxy_logger.error( + "Failed to persist ChatGPT OAuth credential %s: %s", + credential_name, + exc, + ) + _update_session( + session_id, + status="error", + message=f"Tokens obtained but DB persist failed: {exc}", + ) + return + + _update_session(session_id, status="success", message=None) diff --git a/litellm/proxy/copilot_oauth_endpoints/endpoints.py b/litellm/proxy/copilot_oauth_endpoints/endpoints.py new file mode 100644 index 00000000000..74d5d4dd44d --- /dev/null +++ b/litellm/proxy/copilot_oauth_endpoints/endpoints.py @@ -0,0 +1,277 @@ +""" +Admin-only endpoints that let the LiteLLM UI sign in to GitHub Copilot using +GitHub's device-code OAuth flow, and persist the resulting ``access_token`` +to ``LiteLLM_CredentialsTable``. + +Mirrors ``litellm/proxy/chatgpt_oauth_endpoints/`` — the only user-facing +differences are Copilot-specific verification URLs and the ``refresh`` +endpoint semantics (refreshes the short-lived Copilot API key rather than +the long-lived GitHub OAuth token, since the latter does not expire). + +Endpoints: + POST /copilot/oauth/start — start a new flow; returns user_code + verification URL + GET /copilot/oauth/status — poll session status (pending/success/error) + POST /copilot/oauth/cancel — abandon an in-flight flow + POST /copilot/oauth/refresh — force a Copilot API key refresh for a stored credential +""" + +import threading +import time +import uuid +from typing import Any, Dict, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.llms.github_copilot.authenticator import Authenticator +from litellm.llms.github_copilot.common_utils import GithubCopilotError +from litellm.llms.github_copilot.db_authenticator import DBAuthenticator +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter(prefix="/copilot/oauth", tags=["copilot oauth"]) + +SESSION_TTL_SECONDS = 20 * 60 +SESSIONS_MAX_SIZE = 100 + +_sessions: Dict[str, Dict[str, Any]] = {} +_sessions_lock = threading.Lock() + + +class StartRequest(BaseModel): + credential_name: str + + +class StartResponse(BaseModel): + session_id: str + user_code: str + verification_url: str + interval: int + + +class StatusResponse(BaseModel): + status: str # "pending" | "success" | "error" | "cancelled" + credential_name: Optional[str] = None + message: Optional[str] = None + + +class RefreshRequest(BaseModel): + credential_name: str + + +class RefreshResponse(BaseModel): + credential_name: str + api_key_expires_at: Optional[int] = None + + +def _require_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + # These endpoints write to LiteLLM_CredentialsTable (start → insert, + # refresh → rotate). PROXY_ADMIN_VIEW_ONLY must not reach them. + role = getattr(user_api_key_dict, "user_role", None) + if role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only PROXY_ADMIN may initiate Copilot OAuth flows.", + ) + + +def _purge_expired_sessions() -> None: + now = time.time() + with _sessions_lock: + expired = [k for k, v in _sessions.items() if v.get("expires_at", 0) < now] + for k in expired: + del _sessions[k] + + +def _update_session(session_id: str, **fields: Any) -> None: + with _sessions_lock: + entry = _sessions.get(session_id) + if entry is None: + return + entry.update(fields) + + +def _get_session(session_id: str) -> Optional[Dict[str, Any]]: + with _sessions_lock: + entry = _sessions.get(session_id) + return dict(entry) if entry else None + + +@router.post( + "/start", + response_model=StartResponse, + dependencies=[Depends(user_api_key_auth)], +) +async def start_oauth( + body: StartRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> StartResponse: + _require_admin(user_api_key_dict) + _purge_expired_sessions() + + # Atomically reserve a slot so concurrent callers cannot all pass the + # capacity check and blow past SESSIONS_MAX_SIZE. + session_id = uuid.uuid4().hex + now = time.time() + with _sessions_lock: + if len(_sessions) >= SESSIONS_MAX_SIZE: + raise HTTPException( + status_code=429, + detail="Too many in-flight OAuth sessions. Retry shortly.", + ) + _sessions[session_id] = { + "status": "starting", + "credential_name": body.credential_name, + "started_by": user_api_key_dict.user_id, + "started_at": now, + "expires_at": now + SESSION_TTL_SECONDS, + "cancelled": False, + } + + authenticator = Authenticator() + try: + device_code_info = authenticator._get_device_code() + except GithubCopilotError as exc: + with _sessions_lock: + _sessions.pop(session_id, None) + raise HTTPException(status_code=502, detail=exc.message) + + with _sessions_lock: + entry = _sessions.get(session_id) + if entry is None: + raise HTTPException(status_code=410, detail="Session was cancelled") + entry["status"] = "pending" + entry["device_code_info"] = device_code_info + + thread = threading.Thread( + target=_run_device_code_flow, + args=(session_id, body.credential_name, device_code_info, authenticator), + daemon=True, + name=f"copilot-oauth-{session_id[:8]}", + ) + thread.start() + + return StartResponse( + session_id=session_id, + user_code=device_code_info["user_code"], + verification_url=device_code_info["verification_uri"], + interval=int(device_code_info.get("interval", 5)), + ) + + +@router.get( + "/status", + response_model=StatusResponse, + dependencies=[Depends(user_api_key_auth)], +) +async def oauth_status( + session_id: str = Query( + ..., description="Session ID returned by /copilot/oauth/start" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> StatusResponse: + _require_admin(user_api_key_dict) + entry = _get_session(session_id) + if entry is None: + raise HTTPException(status_code=404, detail="Unknown or expired session_id") + return StatusResponse( + status=entry["status"], + credential_name=entry.get("credential_name"), + message=entry.get("message"), + ) + + +@router.post( + "/cancel", + dependencies=[Depends(user_api_key_auth)], +) +async def oauth_cancel( + session_id: str = Query( + ..., description="Session ID returned by /copilot/oauth/start" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Dict[str, bool]: + _require_admin(user_api_key_dict) + with _sessions_lock: + entry = _sessions.get(session_id) + if entry is None: + raise HTTPException(status_code=404, detail="Unknown or expired session_id") + entry["cancelled"] = True + if entry["status"] == "pending": + entry["status"] = "cancelled" + entry["message"] = "Cancelled by user" + return {"success": True} + + +@router.post( + "/refresh", + response_model=RefreshResponse, + dependencies=[Depends(user_api_key_auth)], +) +async def oauth_refresh( + body: RefreshRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> RefreshResponse: + """ + Force a refresh of the Copilot API key derived from the stored GitHub + access token. The GitHub access token itself does not have a short TTL; + this endpoint is primarily a health-check / "force new session key" + action surfaced in the UI. + """ + _require_admin(user_api_key_dict) + db_auth = DBAuthenticator(credential_name=body.credential_name) + try: + info = db_auth.force_refresh_api_key() + except GithubCopilotError as exc: + raise HTTPException(status_code=502, detail=exc.message) + expires_at = info.get("expires_at") + return RefreshResponse( + credential_name=body.credential_name, + api_key_expires_at=int(expires_at) if expires_at is not None else None, + ) + + +def _run_device_code_flow( + session_id: str, + credential_name: str, + device_code_info: Dict[str, str], + authenticator: Authenticator, +) -> None: + """ + Background worker: polls GitHub for the access token, then persists it + via :class:`DBAuthenticator.store_access_token`. + """ + try: + access_token = authenticator._poll_for_access_token( + device_code_info["device_code"] + ) + session_snapshot = _get_session(session_id) + if session_snapshot is None or session_snapshot.get("cancelled"): + return + except GithubCopilotError as exc: + _update_session(session_id, status="error", message=exc.message) + return + except Exception as exc: # pragma: no cover - defensive + verbose_proxy_logger.exception("Unexpected error in Copilot OAuth flow") + _update_session(session_id, status="error", message=str(exc)) + return + + try: + DBAuthenticator(credential_name=credential_name).store_access_token( + access_token + ) + except Exception as exc: + verbose_proxy_logger.error( + "Failed to persist Copilot OAuth credential %s: %s", + credential_name, + exc, + ) + _update_session( + session_id, + status="error", + message=f"Tokens obtained but DB persist failed: {exc}", + ) + return + + _update_session(session_id, status="success", message=None) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index db9468356dd..bc07da20cef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -320,6 +320,12 @@ from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES from litellm.proxy.container_endpoints.endpoints import router as container_router +from litellm.proxy.chatgpt_oauth_endpoints.endpoints import ( + router as chatgpt_oauth_router, +) +from litellm.proxy.copilot_oauth_endpoints.endpoints import ( + router as copilot_oauth_router, +) from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -2161,9 +2167,11 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug(f""" + verbose_proxy_logger.debug( + f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """) + """ + ) def _get_process_rss_mb() -> Optional[float]: @@ -13896,6 +13904,8 @@ app.include_router(vector_store_router) app.include_router(vector_store_management_router) app.include_router(vector_store_files_router) app.include_router(credential_router) +app.include_router(chatgpt_oauth_router) +app.include_router(copilot_oauth_router) app.include_router(llm_passthrough_router) app.include_router(webrtc_router) app.include_router(mcp_management_router) diff --git a/pyproject.toml b/pyproject.toml index a896d520abd..b9722e6aeaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -112,6 +112,7 @@ proxy-runtime = [ [project.scripts] litellm = "litellm:run_server" litellm-proxy = "litellm.proxy.client.cli:cli" +litellm-chatgpt-login = "litellm.llms.chatgpt.cli:cli" [dependency-groups] dev = [ diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_oauth_credential_dispatch.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_oauth_credential_dispatch.py new file mode 100644 index 00000000000..b3e2f8432b7 --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_oauth_credential_dispatch.py @@ -0,0 +1,68 @@ +from unittest.mock import MagicMock, patch + +from litellm.llms.chatgpt.authenticator import Authenticator +from litellm.llms.chatgpt.db_authenticator import ( + OAUTH_CREDENTIAL_API_KEY_PREFIX, + DBAuthenticator, + resolve_authenticator, +) +from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams + + +class TestResolveAuthenticator: + def test_plain_api_key_returns_fallback(self): + fallback = MagicMock(spec=Authenticator) + resolved = resolve_authenticator( + GenericLiteLLMParams(api_key="sk-plain"), fallback + ) + assert resolved is fallback + + def test_none_litellm_params_returns_fallback(self): + fallback = MagicMock(spec=Authenticator) + assert resolve_authenticator(None, fallback) is fallback + + def test_oauth_prefix_returns_db_authenticator(self): + fallback = MagicMock(spec=Authenticator) + resolved = resolve_authenticator( + GenericLiteLLMParams(api_key=f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}my-creds"), + fallback, + ) + assert isinstance(resolved, DBAuthenticator) + assert resolved.credential_name == "my-creds" + + def test_oauth_prefix_with_empty_suffix(self): + fallback = MagicMock(spec=Authenticator) + resolved = resolve_authenticator( + GenericLiteLLMParams(api_key=OAUTH_CREDENTIAL_API_KEY_PREFIX), fallback + ) + assert isinstance(resolved, DBAuthenticator) + assert resolved.credential_name == "" + + +class TestValidateEnvironmentRoutesThroughDispatch: + def test_validate_environment_uses_db_authenticator_when_prefixed(self): + """ + With ``api_key=oauth:`` the config should ask the DB-backed + authenticator (not ``self.authenticator``) for tokens. + """ + config = ChatGPTResponsesAPIConfig() + fs_auth = MagicMock(spec=Authenticator) + fs_auth.get_access_token.side_effect = AssertionError( + "Filesystem authenticator must not be called" + ) + config.authenticator = fs_auth + + with ( + patch.object(DBAuthenticator, "get_access_token", return_value="db-token"), + patch.object(DBAuthenticator, "get_account_id", return_value="acct-db"), + ): + headers = config.validate_environment( + headers={}, + model="chatgpt/gpt-5.3-codex", + litellm_params=GenericLiteLLMParams( + api_key=f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}my-creds" + ), + ) + assert headers["Authorization"] == "Bearer db-token" + assert headers["ChatGPT-Account-Id"] == "acct-db" diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_db_authenticator.py b/tests/test_litellm/llms/chatgpt/test_chatgpt_db_authenticator.py new file mode 100644 index 00000000000..aa950abd68d --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/test_chatgpt_db_authenticator.py @@ -0,0 +1,237 @@ +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.llms.chatgpt.db_authenticator import ( + CREDENTIAL_TYPE, + DBAuthenticator, + _pack_auth_record, + _unpack_auth_record, + persist_credential_to_db, +) +from litellm.types.utils import CredentialItem + + +class TestAuthRecordPacking: + def test_pack_stringifies_all_fields(self): + packed = _pack_auth_record( + { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + ) + assert packed == { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": "1700000000", + } + + def test_pack_omits_none(self): + packed = _pack_auth_record({"access_token": "a", "refresh_token": None}) + assert packed == {"access_token": "a"} + + def test_unpack_coerces_expires_at_to_int(self): + record = _unpack_auth_record({"access_token": "a", "expires_at": "1700000000"}) + assert record == {"access_token": "a", "expires_at": 1700000000} + + def test_unpack_drops_invalid_expires_at(self): + record = _unpack_auth_record({"access_token": "a", "expires_at": "nope"}) + assert record == {"access_token": "a"} + + +class TestDBAuthenticator: + @pytest.fixture(autouse=True) + def _reset_credentials(self, monkeypatch): + original = list(litellm.credential_list) + monkeypatch.setattr(litellm, "credential_list", []) + yield + litellm.credential_list = original + + def test_read_returns_none_when_credential_missing(self): + auth = DBAuthenticator(credential_name="nope") + assert auth._read_auth_file() is None + + def test_read_returns_unpacked_record_from_cache(self): + litellm.credential_list = [ + CredentialItem( + credential_name="test", + credential_values={ + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": "1700000000", + }, + credential_info={"type": CREDENTIAL_TYPE}, + ) + ] + auth = DBAuthenticator(credential_name="test") + record = auth._read_auth_file() + assert record == { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + + def test_write_upserts_cache_and_schedules_db_persist(self): + auth = DBAuthenticator(credential_name="test") + with patch( + "litellm.llms.chatgpt.db_authenticator._schedule_db_persist" + ) as mock_persist: + auth._write_auth_file( + { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + ) + # Cache was updated + assert any( + c.credential_name == "test" and c.credential_values["access_token"] == "a" + for c in litellm.credential_list + ) + # DB persist was scheduled + mock_persist.assert_called_once() + item = mock_persist.call_args.args[0] + assert item.credential_name == "test" + assert item.credential_info == { + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "chatgpt", + } + + def test_write_then_read_roundtrip_via_cache(self): + auth = DBAuthenticator(credential_name="test") + with patch("litellm.llms.chatgpt.db_authenticator._schedule_db_persist"): + auth._write_auth_file( + { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + ) + assert auth._read_auth_file() == { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + + def test_ensure_token_dir_is_noop(self, tmp_path, monkeypatch): + # Parent uses os.makedirs; DBAuthenticator must not touch the filesystem. + calls = [] + monkeypatch.setattr("os.makedirs", lambda *a, **kw: calls.append((a, kw))) + monkeypatch.setattr("os.path.exists", lambda p: False) + auth = DBAuthenticator(credential_name="test") + auth._ensure_token_dir() + assert calls == [] + + +class TestPersistCredentialToDb: + @pytest.mark.asyncio + async def test_noop_when_prisma_missing(self, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", None) + item = CredentialItem( + credential_name="test", + credential_values={"access_token": "a"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + await persist_credential_to_db(item) # no raise + + @pytest.mark.asyncio + async def test_upserts_encrypted_values(self, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.upsert = MagicMock( + return_value=_AsyncNone() + ) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr( + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + lambda v, key=None: f"enc({v})", + ) + + item = CredentialItem( + credential_name="test", + credential_values={"access_token": "a", "refresh_token": "r"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + await persist_credential_to_db(item) + + fake_prisma.db.litellm_credentialstable.upsert.assert_called_once() + kwargs = fake_prisma.db.litellm_credentialstable.upsert.call_args.kwargs + assert kwargs["where"] == {"credential_name": "test"} + create = kwargs["data"]["create"] + assert create["credential_name"] == "test" + assert create["credential_values"] == { + "access_token": "enc(a)", + "refresh_token": "enc(r)", + } + assert create["credential_info"] == {"type": CREDENTIAL_TYPE} + update = kwargs["data"]["update"] + assert update["credential_values"] == create["credential_values"] + + +class _AsyncNone: + def __await__(self): + async def _coro(): + return None + + return _coro().__await__() + + +class TestPersistScheduling: + def test_schedule_starts_background_thread(self, monkeypatch): + from litellm.llms.chatgpt import db_authenticator as mod + + calls = [] + + def _fake_sync(item): + calls.append(item.credential_name) + + monkeypatch.setattr(mod, "_persist_item_sync", _fake_sync) + item = CredentialItem( + credential_name="c", + credential_values={"access_token": "a"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + mod._schedule_db_persist(item) + # Thread is daemon; give it a brief moment to run. + import time + + for _ in range(100): + if calls: + break + time.sleep(0.01) + assert calls == ["c"] + + def test_persist_item_sync_swallows_exceptions(self, monkeypatch): + """A DB failure on the worker thread must not crash the process.""" + from litellm.llms.chatgpt import db_authenticator as mod + + async def _boom(item): + raise RuntimeError("db offline") + + monkeypatch.setattr(mod, "persist_credential_to_db", _boom) + item = CredentialItem( + credential_name="c", + credential_values={"access_token": "a"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + # Should not raise; just logs. + mod._persist_item_sync(item) diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_pkce_login.py b/tests/test_litellm/llms/chatgpt/test_chatgpt_pkce_login.py new file mode 100644 index 00000000000..255e73602da --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/test_chatgpt_pkce_login.py @@ -0,0 +1,303 @@ +import base64 +import hashlib +import threading +import urllib.parse +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from litellm.llms.chatgpt.authenticator import Authenticator +from litellm.llms.chatgpt.common_utils import ( + CHATGPT_CLIENT_ID, + CHATGPT_OAUTH_TOKEN_URL, + GetAccessTokenError, +) +from litellm.llms.chatgpt.pkce import ( + OAUTH_AUTHORIZE_URL, + REDIRECT_PATH, + SCOPE, + _build_authorize_url, + _exchange_code_for_tokens, + _generate_code_challenge, + _generate_code_verifier, + _make_handler, +) + + +class TestPkceHelpers: + def test_verifier_length_and_charset(self): + verifier = _generate_code_verifier() + assert 43 <= len(verifier) <= 128 + # token_urlsafe uses RFC 7636 unreserved chars (A-Z, a-z, 0-9, -, _) + assert all(c.isalnum() or c in "-_" for c in verifier) + + def test_challenge_is_s256_of_verifier(self): + verifier = "fixed-verifier-for-testing-abc123-abc123-abc" + expected = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + assert _generate_code_challenge(verifier) == expected + + def test_authorize_url_contains_required_params(self): + url = _build_authorize_url( + redirect_uri="http://127.0.0.1:1455/auth/callback", + code_challenge="test-challenge", + state="test-state", + ) + assert url.startswith(OAUTH_AUTHORIZE_URL + "?") + parsed = urllib.parse.parse_qs(urllib.parse.urlparse(url).query) + assert parsed["response_type"] == ["code"] + assert parsed["client_id"] == [CHATGPT_CLIENT_ID] + assert parsed["redirect_uri"] == ["http://127.0.0.1:1455/auth/callback"] + assert parsed["scope"] == [SCOPE] + assert parsed["code_challenge"] == ["test-challenge"] + assert parsed["code_challenge_method"] == ["S256"] + assert parsed["state"] == ["test-state"] + assert parsed["id_token_add_organizations"] == ["true"] + assert parsed["codex_cli_simplified_flow"] == ["true"] + assert "originator" in parsed + + +class _FakeWFile: + def __init__(self): + self.written = b"" + + def write(self, data): + self.written += data + + +class _FakeHandler: + """Minimal stand-in that lets us invoke the generated do_GET without sockets.""" + + def __init__(self, handler_cls, path): + self.handler_cls = handler_cls + self.path = path + self.wfile = _FakeWFile() + self.response_status = None + self.headers_sent = {} + + def send_response(self, code): + self.response_status = code + + def send_header(self, k, v): + self.headers_sent[k] = v + + def end_headers(self): + pass + + +def _invoke_handler(handler_cls, path): + fake = _FakeHandler(handler_cls, path) + # Bind helper methods to the fake; they only read self.path / self.wfile / etc. + fake._respond_html = handler_cls._respond_html.__get__(fake, type(fake)) + fake._respond_error = handler_cls._respond_error.__get__(fake, type(fake)) + handler_cls.do_GET(fake) + return fake + + +class TestPkceCallbackHandler: + def test_success_captures_code_and_sets_event(self): + result = {} + event = threading.Event() + handler_cls = _make_handler(result, event, expected_state="abc") + fake = _invoke_handler(handler_cls, f"{REDIRECT_PATH}?code=xyz&state=abc") + assert fake.response_status == 200 + assert result == {"code": "xyz"} + assert event.is_set() + + def test_rejects_state_mismatch(self): + result = {} + event = threading.Event() + handler_cls = _make_handler(result, event, expected_state="abc") + fake = _invoke_handler(handler_cls, f"{REDIRECT_PATH}?code=xyz&state=wrong") + assert fake.response_status == 400 + assert result.get("error") == "state mismatch" + assert "code" not in result + assert event.is_set() + + def test_propagates_provider_error(self): + result = {} + event = threading.Event() + handler_cls = _make_handler(result, event, expected_state="abc") + fake = _invoke_handler( + handler_cls, + f"{REDIRECT_PATH}?error=access_denied&error_description=User+cancelled&state=abc", + ) + assert fake.response_status == 400 + assert result["error"] == "User cancelled" + assert event.is_set() + + def test_ignores_unrelated_paths(self): + result = {} + event = threading.Event() + handler_cls = _make_handler(result, event, expected_state="abc") + fake = _invoke_handler(handler_cls, "/favicon.ico") + assert fake.response_status == 404 + assert result == {} + assert not event.is_set() + + def test_escapes_error_description_to_prevent_xss(self): + """ + ``error_description`` comes from query params. The callback response + HTML must escape it so a malicious OAuth redirect cannot inject + script into the loopback server's response body. + """ + result = {} + event = threading.Event() + handler_cls = _make_handler(result, event, expected_state="abc") + # urlencoded "" in error_description + payload = "%3Cscript%3Ealert%281%29%3C%2Fscript%3E" + fake = _invoke_handler( + handler_cls, + f"{REDIRECT_PATH}?error=access_denied&error_description={payload}&state=abc", + ) + body = fake.wfile.written.decode("utf-8") + assert fake.response_status == 400 + # Raw tag must not appear in the response; escaped form must. + assert "" not in body + assert "<script>alert(1)</script>" in body + # The unescaped form is still what we store in `result["error"]` — the + # in-process dict is not rendered as HTML, only the response body is. + assert result["error"] == "" + + +class TestPkceTokenExchange: + def test_posts_form_data_and_returns_tokens(self): + fake_client = MagicMock() + fake_response = MagicMock() + fake_response.json.return_value = { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + } + fake_response.raise_for_status.return_value = None + fake_client.post.return_value = fake_response + + with patch( + "litellm.llms.chatgpt.pkce._get_httpx_client", + return_value=fake_client, + ): + tokens = _exchange_code_for_tokens( + code="the-code", + code_verifier="the-verifier", + redirect_uri="http://127.0.0.1:1455/auth/callback", + ) + + assert tokens == {"access_token": "a", "refresh_token": "r", "id_token": "i"} + args, kwargs = fake_client.post.call_args + assert args[0] == CHATGPT_OAUTH_TOKEN_URL + assert kwargs["data"] == { + "grant_type": "authorization_code", + "code": "the-code", + "redirect_uri": "http://127.0.0.1:1455/auth/callback", + "client_id": CHATGPT_CLIENT_ID, + "code_verifier": "the-verifier", + } + + def test_raises_on_missing_fields(self): + fake_client = MagicMock() + fake_response = MagicMock() + fake_response.json.return_value = {"access_token": "a"} # missing fields + fake_response.raise_for_status.return_value = None + fake_client.post.return_value = fake_response + + with patch( + "litellm.llms.chatgpt.pkce._get_httpx_client", + return_value=fake_client, + ): + with pytest.raises(GetAccessTokenError): + _exchange_code_for_tokens(code="c", code_verifier="v", redirect_uri="r") + + def test_raises_on_generic_exception(self): + """ + Non-HTTPStatusError exceptions (e.g. JSON decode, connection + reset) funnel through the generic Exception handler. + """ + fake_client = MagicMock() + fake_client.post.side_effect = RuntimeError("connection reset") + + with patch( + "litellm.llms.chatgpt.pkce._get_httpx_client", + return_value=fake_client, + ): + with pytest.raises(GetAccessTokenError) as excinfo: + _exchange_code_for_tokens(code="c", code_verifier="v", redirect_uri="r") + assert excinfo.value.status_code == 400 + assert "connection reset" in excinfo.value.message + + def test_raises_on_http_error(self): + fake_client = MagicMock() + request = httpx.Request("POST", CHATGPT_OAUTH_TOKEN_URL) + response = httpx.Response(400, request=request) + fake_client.post.side_effect = httpx.HTTPStatusError( + "bad", request=request, response=response + ) + + with patch( + "litellm.llms.chatgpt.pkce._get_httpx_client", + return_value=fake_client, + ): + with pytest.raises(GetAccessTokenError) as excinfo: + _exchange_code_for_tokens(code="c", code_verifier="v", redirect_uri="r") + assert excinfo.value.status_code == 400 + + +class TestCliExceptionHandling: + def test_broad_exception_produces_clean_message(self, capsys, monkeypatch): + """ + Non-ChatGPTAuthError exceptions (filesystem, network, etc.) should + surface a one-line error and exit 1 — never a raw traceback. + """ + from litellm.llms.chatgpt import cli + + monkeypatch.setattr("sys.argv", ["litellm-chatgpt-login", "--method", "device"]) + + def _boom(self): + raise RuntimeError("disk full") + + monkeypatch.setattr( + "litellm.llms.chatgpt.authenticator.Authenticator._login_device_code", + _boom, + ) + exit_code = cli.cli() + assert exit_code == 1 + captured = capsys.readouterr() + assert "disk full" in captured.err + assert "Traceback" not in captured.err + + def test_keyboard_interrupt_returns_130(self, capsys, monkeypatch): + from litellm.llms.chatgpt import cli + + monkeypatch.setattr("sys.argv", ["litellm-chatgpt-login", "--method", "device"]) + + def _interrupt(self): + raise KeyboardInterrupt() + + monkeypatch.setattr( + "litellm.llms.chatgpt.authenticator.Authenticator._login_device_code", + _interrupt, + ) + exit_code = cli.cli() + assert exit_code == 130 + assert "cancelled" in capsys.readouterr().err.lower() + + +class TestLoginPkcePortInUse: + @pytest.fixture + def authenticator(self): + with patch("os.path.exists", return_value=True): + return Authenticator() + + def test_raises_when_port_cannot_bind(self, authenticator): + with patch( + "litellm.llms.chatgpt.pkce.http.server.HTTPServer", + side_effect=OSError("address in use"), + ): + with pytest.raises(GetAccessTokenError) as excinfo: + authenticator.login_pkce( + open_browser=False, port=1455, timeout_seconds=1 + ) + assert "Failed to bind loopback server" in excinfo.value.message diff --git a/tests/test_litellm/llms/github_copilot/test_copilot_db_authenticator.py b/tests/test_litellm/llms/github_copilot/test_copilot_db_authenticator.py new file mode 100644 index 00000000000..95877f8d782 --- /dev/null +++ b/tests/test_litellm/llms/github_copilot/test_copilot_db_authenticator.py @@ -0,0 +1,281 @@ +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.llms.github_copilot.db_authenticator import ( + CREDENTIAL_TYPE, + OAUTH_CREDENTIAL_API_KEY_PREFIX, + DBAuthenticator, + persist_credential_to_db, + resolve_authenticator, +) +from litellm.llms.github_copilot.common_utils import GetAccessTokenError +from litellm.types.utils import CredentialItem + + +@pytest.fixture(autouse=True) +def _reset_credentials(monkeypatch): + original = list(litellm.credential_list) + monkeypatch.setattr(litellm, "credential_list", []) + DBAuthenticator._api_key_cache.clear() + yield + litellm.credential_list = original + DBAuthenticator._api_key_cache.clear() + + +class TestDBAuthenticatorAccessToken: + def test_raises_when_missing(self): + auth = DBAuthenticator(credential_name="nope") + with pytest.raises(GetAccessTokenError): + auth.get_access_token() + + def test_reads_from_cache(self): + litellm.credential_list = [ + CredentialItem( + credential_name="c", + credential_values={"access_token": "gho_abc"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + ] + auth = DBAuthenticator(credential_name="c") + assert auth.get_access_token() == "gho_abc" + + def test_store_access_token_upserts_and_invalidates_cache(self): + auth = DBAuthenticator(credential_name="c") + DBAuthenticator._api_key_cache["c"] = {"token": "stale", "expires_at": 9**12} + + with patch( + "litellm.llms.github_copilot.db_authenticator._schedule_db_persist" + ) as mock_persist: + auth.store_access_token("gho_new") + + # Cache contents updated + assert any( + c.credential_name == "c" + and c.credential_values["access_token"] == "gho_new" + for c in litellm.credential_list + ) + # Stale API key purged + assert "c" not in DBAuthenticator._api_key_cache + # DB persist scheduled + mock_persist.assert_called_once() + + +class TestDBAuthenticatorApiKey: + def test_uses_cached_api_key_when_not_expired(self): + import time + + DBAuthenticator._api_key_cache["c"] = { + "token": "cached-key", + "expires_at": int(time.time()) + 3600, + } + auth = DBAuthenticator(credential_name="c") + assert auth.get_api_key() == "cached-key" + + def test_refreshes_when_cache_stale(self): + import time + + DBAuthenticator._api_key_cache["c"] = { + "token": "stale", + "expires_at": int(time.time()) - 10, + } + litellm.credential_list = [ + CredentialItem( + credential_name="c", + credential_values={"access_token": "gho_abc"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + ] + auth = DBAuthenticator(credential_name="c") + with patch.object( + DBAuthenticator, + "_refresh_api_key", + return_value={ + "token": "fresh", + "expires_at": int(time.time()) + 3600, + "endpoints": {"api": "https://api.githubcopilot.com"}, + }, + ): + assert auth.get_api_key() == "fresh" + assert DBAuthenticator._api_key_cache["c"]["token"] == "fresh" + + def test_force_refresh_ignores_cache(self): + import time + + DBAuthenticator._api_key_cache["c"] = { + "token": "cached", + "expires_at": int(time.time()) + 3600, + } + litellm.credential_list = [ + CredentialItem( + credential_name="c", + credential_values={"access_token": "gho_abc"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + ] + auth = DBAuthenticator(credential_name="c") + with patch.object( + DBAuthenticator, + "_refresh_api_key", + return_value={"token": "forced", "expires_at": int(time.time()) + 7200}, + ) as mock_refresh: + info = auth.force_refresh_api_key() + mock_refresh.assert_called_once() + assert info["token"] == "forced" + assert DBAuthenticator._api_key_cache["c"]["token"] == "forced" + + def test_api_base_reads_from_cache(self): + import time + + DBAuthenticator._api_key_cache["c"] = { + "token": "tok", + "expires_at": int(time.time()) + 3600, + "endpoints": {"api": "https://api.example.copilot"}, + } + auth = DBAuthenticator(credential_name="c") + assert auth.get_api_base() == "https://api.example.copilot" + + def test_api_base_is_none_when_uncached(self): + auth = DBAuthenticator(credential_name="c") + assert auth.get_api_base() is None + + +class TestResolveAuthenticator: + def test_plain_api_key_returns_fallback(self): + fallback = MagicMock() + resolved = resolve_authenticator("sk-plain", None, fallback) + assert resolved is fallback + + def test_oauth_prefix_returns_db_authenticator(self): + fallback = MagicMock() + resolved = resolve_authenticator( + f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}my-creds", None, fallback + ) + assert isinstance(resolved, DBAuthenticator) + assert resolved.credential_name == "my-creds" + + def test_checks_litellm_params_api_key_when_top_level_plain(self): + """ + ``_get_openai_compatible_provider_info`` may rewrite ``api_key`` to the + resolved key before ``validate_environment`` runs. The raw marker + should still be picked up from ``litellm_params``. + """ + fallback = MagicMock() + resolved = resolve_authenticator( + "rewritten-copilot-api-key", + {"api_key": f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}my-creds"}, + fallback, + ) + assert isinstance(resolved, DBAuthenticator) + assert resolved.credential_name == "my-creds" + + def test_handles_pydantic_litellm_params(self): + fallback = MagicMock() + + class _FakeParams: + api_key = f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}from-pydantic" + + resolved = resolve_authenticator(None, _FakeParams(), fallback) + assert isinstance(resolved, DBAuthenticator) + assert resolved.credential_name == "from-pydantic" + + +class TestPersistCredentialToDb: + @pytest.mark.asyncio + async def test_noop_when_prisma_missing(self, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", None) + item = CredentialItem( + credential_name="c", + credential_values={"access_token": "a"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + await persist_credential_to_db(item) # no raise + + @pytest.mark.asyncio + async def test_upserts_encrypted_values(self, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + fake_prisma = MagicMock() + + async def _upsert(**kwargs): + return None + + fake_prisma.db.litellm_credentialstable.upsert = MagicMock( + side_effect=lambda **kwargs: _Awaitable(None) + ) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr( + "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", + lambda v, key=None: f"enc({v})", + ) + + item = CredentialItem( + credential_name="c", + credential_values={"access_token": "gho_abc"}, + credential_info={ + "type": CREDENTIAL_TYPE, + "custom_llm_provider": "github_copilot", + }, + ) + await persist_credential_to_db(item) + + fake_prisma.db.litellm_credentialstable.upsert.assert_called_once() + kwargs = fake_prisma.db.litellm_credentialstable.upsert.call_args.kwargs + assert kwargs["where"] == {"credential_name": "c"} + assert kwargs["data"]["create"]["credential_values"] == { + "access_token": "enc(gho_abc)" + } + assert kwargs["data"]["create"]["credential_info"]["type"] == CREDENTIAL_TYPE + + +class _Awaitable: + def __init__(self, value): + self._value = value + + def __await__(self): + async def _coro(): + return self._value + + return _coro().__await__() + + +class TestPersistScheduling: + def test_schedule_starts_background_thread(self, monkeypatch): + from litellm.llms.github_copilot import db_authenticator as mod + + calls = [] + + def _fake_sync(item): + calls.append(item.credential_name) + + monkeypatch.setattr(mod, "_persist_item_sync", _fake_sync) + item = CredentialItem( + credential_name="c", + credential_values={"access_token": "gho_abc"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + mod._schedule_db_persist(item) + import time + + for _ in range(100): + if calls: + break + time.sleep(0.01) + assert calls == ["c"] + + def test_persist_item_sync_swallows_exceptions(self, monkeypatch): + from litellm.llms.github_copilot import db_authenticator as mod + + async def _boom(item): + raise RuntimeError("db offline") + + monkeypatch.setattr(mod, "persist_credential_to_db", _boom) + item = CredentialItem( + credential_name="c", + credential_values={"access_token": "gho_abc"}, + credential_info={"type": CREDENTIAL_TYPE}, + ) + mod._persist_item_sync(item) # must not raise diff --git a/tests/test_litellm/llms/github_copilot/test_copilot_oauth_dispatch.py b/tests/test_litellm/llms/github_copilot/test_copilot_oauth_dispatch.py new file mode 100644 index 00000000000..aadb023df4a --- /dev/null +++ b/tests/test_litellm/llms/github_copilot/test_copilot_oauth_dispatch.py @@ -0,0 +1,84 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.llms.github_copilot.authenticator import Authenticator +from litellm.llms.github_copilot.chat.transformation import GithubCopilotConfig +from litellm.llms.github_copilot.db_authenticator import ( + OAUTH_CREDENTIAL_API_KEY_PREFIX, + DBAuthenticator, +) +from litellm.llms.github_copilot.responses.transformation import ( + GithubCopilotResponsesAPIConfig, +) +from litellm.types.router import GenericLiteLLMParams + + +@pytest.fixture(autouse=True) +def _reset_cache(): + DBAuthenticator._api_key_cache.clear() + yield + DBAuthenticator._api_key_cache.clear() + + +class TestChatTransformationDispatch: + def test_oauth_prefix_api_key_uses_db_authenticator(self): + config = GithubCopilotConfig() + fs_auth = MagicMock(spec=Authenticator) + fs_auth.get_api_base.side_effect = AssertionError( + "Filesystem authenticator must not be used for oauth: prefix" + ) + fs_auth.get_api_key.side_effect = AssertionError( + "Filesystem authenticator must not be used for oauth: prefix" + ) + config.authenticator = fs_auth + + with ( + patch.object( + DBAuthenticator, "get_api_base", return_value="https://x.example" + ), + patch.object(DBAuthenticator, "get_api_key", return_value="cop-key"), + ): + base, key, _ = config._get_openai_compatible_provider_info( + model="gpt-5", + api_base=None, + api_key=f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}my-creds", + custom_llm_provider="github_copilot", + ) + assert key == "cop-key" + assert base == "https://x.example" + + def test_plain_api_key_uses_filesystem_authenticator(self): + config = GithubCopilotConfig() + config.authenticator = MagicMock(spec=Authenticator) + config.authenticator.get_api_base.return_value = None + config.authenticator.get_api_key.return_value = "fs-key" + + _, key, _ = config._get_openai_compatible_provider_info( + model="gpt-5", + api_base=None, + api_key="sk-plain", + custom_llm_provider="github_copilot", + ) + assert key == "fs-key" + + +class TestResponsesTransformationDispatch: + def test_oauth_prefix_routes_to_db_authenticator(self): + config = GithubCopilotResponsesAPIConfig() + fs_auth = MagicMock(spec=Authenticator) + fs_auth.get_api_key.side_effect = AssertionError( + "Filesystem authenticator must not be used" + ) + config.authenticator = fs_auth + + with patch.object(DBAuthenticator, "get_api_key", return_value="cop-key"): + headers = config.validate_environment( + headers={}, + model="gpt-5.1-codex", + litellm_params=GenericLiteLLMParams( + api_key=f"{OAUTH_CREDENTIAL_API_KEY_PREFIX}my-creds" + ), + ) + # The Copilot responses config sets Authorization with the api_key. + assert any("cop-key" in v for v in headers.values()) diff --git a/tests/test_litellm/proxy/chatgpt_oauth_endpoints/__init__.py b/tests/test_litellm/proxy/chatgpt_oauth_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/chatgpt_oauth_endpoints/test_chatgpt_oauth_endpoints.py b/tests/test_litellm/proxy/chatgpt_oauth_endpoints/test_chatgpt_oauth_endpoints.py new file mode 100644 index 00000000000..761bfffab1d --- /dev/null +++ b/tests/test_litellm/proxy/chatgpt_oauth_endpoints/test_chatgpt_oauth_endpoints.py @@ -0,0 +1,470 @@ +import asyncio +import time +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.chatgpt_oauth_endpoints import endpoints as oauth_endpoints +from litellm.proxy.chatgpt_oauth_endpoints.endpoints import ( + RefreshRequest, + StartRequest, + _sessions, + _sessions_lock, + oauth_cancel, + oauth_refresh, + oauth_status, + start_oauth, +) + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="admin-1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + +def _non_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _view_only_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="view-1", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + + +@pytest.fixture(autouse=True) +def _clear_sessions(): + with _sessions_lock: + _sessions.clear() + yield + with _sessions_lock: + _sessions.clear() + + +class TestStartOAuth: + @pytest.mark.asyncio + async def test_admin_only_rejects_internal_user(self): + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_non_admin(), + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_admin_only_rejects_view_only_admin(self): + """PROXY_ADMIN_VIEW_ONLY must not be able to start OAuth flows.""" + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_view_only_admin(), + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_rejects_when_session_cap_reached(self): + from litellm.proxy.chatgpt_oauth_endpoints.endpoints import SESSIONS_MAX_SIZE + + with _sessions_lock: + for i in range(SESSIONS_MAX_SIZE): + _sessions[f"existing-{i}"] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + } + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert exc_info.value.status_code == 429 + + @pytest.mark.asyncio + async def test_device_code_failure_cleans_up_reserved_slot(self): + """ + When the device-code network call fails, the reserved session slot + must be released so the cap doesn't leak. + """ + from litellm.llms.chatgpt.common_utils import GetDeviceCodeError + + with patch.object( + oauth_endpoints.Authenticator, + "_request_device_code", + side_effect=GetDeviceCodeError(status_code=500, message="gh down"), + ): + with pytest.raises(HTTPException): + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + with _sessions_lock: + assert len(_sessions) == 0 + + @pytest.mark.asyncio + async def test_creates_session_and_spawns_worker(self): + fake_device_code = { + "device_auth_id": "d-1", + "user_code": "ABCD-1234", + "interval": "5", + } + + class _FakeThread: + def __init__(self, *args, **kwargs): + self.started = False + + def start(self): + self.started = True + + with ( + patch.object( + oauth_endpoints.Authenticator, + "_request_device_code", + return_value=fake_device_code, + ), + patch( + "litellm.proxy.chatgpt_oauth_endpoints.endpoints.threading.Thread", + _FakeThread, + ), + ): + response = await start_oauth( + StartRequest(credential_name="my-creds"), + user_api_key_dict=_admin(), + ) + + assert response.user_code == "ABCD-1234" + assert response.interval == 5 + assert response.verification_url.endswith("/codex/device") + + with _sessions_lock: + assert response.session_id in _sessions + entry = _sessions[response.session_id] + assert entry["status"] == "pending" + assert entry["credential_name"] == "my-creds" + + @pytest.mark.asyncio + async def test_returns_502_on_device_code_failure(self): + from litellm.llms.chatgpt.common_utils import GetDeviceCodeError + + with patch.object( + oauth_endpoints.Authenticator, + "_request_device_code", + side_effect=GetDeviceCodeError(status_code=500, message="upstream down"), + ): + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert exc_info.value.status_code == 502 + assert "upstream down" in exc_info.value.detail + + +class TestStatusEndpoint: + @pytest.mark.asyncio + async def test_returns_pending_session(self): + with _sessions_lock: + _sessions["s1"] = { + "status": "pending", + "credential_name": "my-creds", + "expires_at": time.time() + 600, + } + response = await oauth_status(session_id="s1", user_api_key_dict=_admin()) + assert response.status == "pending" + assert response.credential_name == "my-creds" + + @pytest.mark.asyncio + async def test_404_on_unknown(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_status(session_id="nope", user_api_key_dict=_admin()) + assert exc_info.value.status_code == 404 + + @pytest.mark.asyncio + async def test_admin_only(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_status(session_id="any", user_api_key_dict=_non_admin()) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_rejects_view_only_admin(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_status(session_id="any", user_api_key_dict=_view_only_admin()) + assert exc_info.value.status_code == 403 + + +class TestCancelEndpoint: + @pytest.mark.asyncio + async def test_cancel_flips_pending_to_cancelled(self): + with _sessions_lock: + _sessions["s1"] = { + "status": "pending", + "credential_name": "my-creds", + "expires_at": time.time() + 600, + "cancelled": False, + } + result = await oauth_cancel(session_id="s1", user_api_key_dict=_admin()) + assert result == {"success": True} + with _sessions_lock: + assert _sessions["s1"]["status"] == "cancelled" + assert _sessions["s1"]["cancelled"] is True + + @pytest.mark.asyncio + async def test_cancel_leaves_success_untouched(self): + with _sessions_lock: + _sessions["s1"] = { + "status": "success", + "credential_name": "my-creds", + "expires_at": time.time() + 600, + "cancelled": False, + } + await oauth_cancel(session_id="s1", user_api_key_dict=_admin()) + with _sessions_lock: + assert _sessions["s1"]["status"] == "success" + + +class TestRefreshEndpoint: + @pytest.mark.asyncio + async def test_admin_only(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), user_api_key_dict=_non_admin() + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_rejects_view_only_admin(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_view_only_admin(), + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_404_when_no_stored_refresh_token(self): + from litellm.llms.chatgpt.db_authenticator import DBAuthenticator + + with patch.object(DBAuthenticator, "_read_auth_file", return_value=None): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert exc_info.value.status_code == 404 + + @pytest.mark.asyncio + async def test_successful_refresh_returns_new_expires_at(self): + from litellm.llms.chatgpt.db_authenticator import DBAuthenticator + + stored = { + "access_token": "old", + "refresh_token": "r1", + "expires_at": 1699000000, + } + refreshed_stored = { + "access_token": "new", + "refresh_token": "r1", + "expires_at": 1700000000, + } + reads = [stored, refreshed_stored] + + with ( + patch.object( + DBAuthenticator, "_read_auth_file", side_effect=lambda: reads.pop(0) + ), + patch.object( + DBAuthenticator, + "_refresh_tokens", + return_value={"access_token": "new", "refresh_token": "r1"}, + ), + ): + response = await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert response.credential_name == "c" + assert response.expires_at == 1700000000 + + @pytest.mark.asyncio + async def test_502_on_refresh_failure(self): + from litellm.llms.chatgpt.common_utils import RefreshAccessTokenError + from litellm.llms.chatgpt.db_authenticator import DBAuthenticator + + with ( + patch.object( + DBAuthenticator, + "_read_auth_file", + return_value={"access_token": "a", "refresh_token": "r"}, + ), + patch.object( + DBAuthenticator, + "_refresh_tokens", + side_effect=RefreshAccessTokenError( + status_code=400, message="refresh failed" + ), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert exc_info.value.status_code == 502 + + +class TestBackgroundWorker: + def test_worker_marks_success_and_persists(self, monkeypatch): + from litellm.proxy.chatgpt_oauth_endpoints.endpoints import ( + _run_device_code_flow, + ) + + session_id = "s1" + with _sessions_lock: + _sessions[session_id] = { + "status": "pending", + "credential_name": "my-creds", + "expires_at": time.time() + 600, + "cancelled": False, + } + + auth = MagicMock() + auth._poll_for_authorization_code.return_value = { + "authorization_code": "ac", + "code_challenge": "cc", + "code_verifier": "cv", + } + auth._exchange_code_for_tokens.return_value = { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + } + auth._build_auth_record.return_value = { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + + persist_mock = MagicMock() + + async def _fake_persist(item): + persist_mock(item) + + monkeypatch.setattr( + "litellm.proxy.chatgpt_oauth_endpoints.endpoints.persist_credential_to_db", + _fake_persist, + ) + monkeypatch.setattr(litellm, "credential_list", []) + + _run_device_code_flow( + session_id=session_id, + credential_name="my-creds", + device_code={"interval": "5"}, + authenticator=auth, + ) + + with _sessions_lock: + assert _sessions[session_id]["status"] == "success" + persist_mock.assert_called_once() + persisted_item = persist_mock.call_args.args[0] + assert persisted_item.credential_name == "my-creds" + # Verify the item is in the in-memory cache too + assert any(c.credential_name == "my-creds" for c in litellm.credential_list) + + def test_worker_marks_error_on_db_persist_failure(self, monkeypatch): + """ + If tokens are obtained but the DB write fails, the session should + flip to ``error`` with an informative message (the in-memory cache + was already updated; next retry via UI will retry the DB write). + """ + from litellm.proxy.chatgpt_oauth_endpoints.endpoints import ( + _run_device_code_flow, + ) + + session_id = "s1" + with _sessions_lock: + _sessions[session_id] = { + "status": "pending", + "credential_name": "my-creds", + "expires_at": time.time() + 600, + "cancelled": False, + } + + auth = MagicMock() + auth._poll_for_authorization_code.return_value = { + "authorization_code": "ac", + "code_challenge": "cc", + "code_verifier": "cv", + } + auth._exchange_code_for_tokens.return_value = { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + } + auth._build_auth_record.return_value = { + "access_token": "a", + "refresh_token": "r", + "id_token": "i", + "account_id": "acct-1", + "expires_at": 1700000000, + } + + async def _boom(item): + raise RuntimeError("prisma disconnected") + + monkeypatch.setattr( + "litellm.proxy.chatgpt_oauth_endpoints.endpoints.persist_credential_to_db", + _boom, + ) + monkeypatch.setattr(litellm, "credential_list", []) + + _run_device_code_flow( + session_id=session_id, + credential_name="my-creds", + device_code={"interval": "5"}, + authenticator=auth, + ) + + with _sessions_lock: + assert _sessions[session_id]["status"] == "error" + assert "DB persist failed" in _sessions[session_id]["message"] + + def test_worker_marks_error_on_auth_failure(self, monkeypatch): + from litellm.llms.chatgpt.common_utils import GetAccessTokenError + from litellm.proxy.chatgpt_oauth_endpoints.endpoints import ( + _run_device_code_flow, + ) + + session_id = "s1" + with _sessions_lock: + _sessions[session_id] = { + "status": "pending", + "credential_name": "my-creds", + "expires_at": time.time() + 600, + "cancelled": False, + } + + auth = MagicMock() + auth._poll_for_authorization_code.side_effect = GetAccessTokenError( + status_code=408, message="timed out" + ) + + _run_device_code_flow( + session_id=session_id, + credential_name="my-creds", + device_code={"interval": "5"}, + authenticator=auth, + ) + + with _sessions_lock: + assert _sessions[session_id]["status"] == "error" + assert "timed out" in _sessions[session_id]["message"] diff --git a/tests/test_litellm/proxy/copilot_oauth_endpoints/__init__.py b/tests/test_litellm/proxy/copilot_oauth_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/copilot_oauth_endpoints/test_copilot_oauth_endpoints.py b/tests/test_litellm/proxy/copilot_oauth_endpoints/test_copilot_oauth_endpoints.py new file mode 100644 index 00000000000..265c14d3b6f --- /dev/null +++ b/tests/test_litellm/proxy/copilot_oauth_endpoints/test_copilot_oauth_endpoints.py @@ -0,0 +1,324 @@ +import time +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.llms.github_copilot.common_utils import ( + GetAccessTokenError, + GetDeviceCodeError, + RefreshAPIKeyError, +) +from litellm.llms.github_copilot.db_authenticator import DBAuthenticator +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.copilot_oauth_endpoints import endpoints as oauth_endpoints +from litellm.proxy.copilot_oauth_endpoints.endpoints import ( + RefreshRequest, + StartRequest, + _sessions, + _sessions_lock, + oauth_cancel, + oauth_refresh, + oauth_status, + start_oauth, +) + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _non_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="u-1", user_role=LitellmUserRoles.INTERNAL_USER) + + +def _view_only_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="view-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + + +@pytest.fixture(autouse=True) +def _clear_sessions(): + with _sessions_lock: + _sessions.clear() + DBAuthenticator._api_key_cache.clear() + yield + with _sessions_lock: + _sessions.clear() + DBAuthenticator._api_key_cache.clear() + + +class TestStartOAuth: + @pytest.mark.asyncio + async def test_admin_only_rejects_internal_user(self): + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), user_api_key_dict=_non_admin() + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_admin_only_rejects_view_only_admin(self): + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_view_only_admin(), + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_rejects_when_session_cap_reached(self): + from litellm.proxy.copilot_oauth_endpoints.endpoints import SESSIONS_MAX_SIZE + + with _sessions_lock: + for i in range(SESSIONS_MAX_SIZE): + _sessions[f"existing-{i}"] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + } + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), user_api_key_dict=_admin() + ) + assert exc_info.value.status_code == 429 + + @pytest.mark.asyncio + async def test_device_code_failure_cleans_up_reserved_slot(self): + with patch.object( + oauth_endpoints.Authenticator, + "_get_device_code", + side_effect=GetDeviceCodeError(status_code=500, message="gh down"), + ): + with pytest.raises(HTTPException): + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + with _sessions_lock: + assert len(_sessions) == 0 + + @pytest.mark.asyncio + async def test_creates_session_with_github_verification_url(self): + fake_device_code = { + "device_code": "dc-1", + "user_code": "ABCD-1234", + "verification_uri": "https://github.com/login/device", + "interval": 5, + } + + class _FakeThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + pass + + with ( + patch.object( + oauth_endpoints.Authenticator, + "_get_device_code", + return_value=fake_device_code, + ), + patch( + "litellm.proxy.copilot_oauth_endpoints.endpoints.threading.Thread", + _FakeThread, + ), + ): + response = await start_oauth( + StartRequest(credential_name="my-copilot"), + user_api_key_dict=_admin(), + ) + + assert response.user_code == "ABCD-1234" + assert response.verification_url == "https://github.com/login/device" + assert response.interval == 5 + + @pytest.mark.asyncio + async def test_returns_502_on_device_code_failure(self): + with patch.object( + oauth_endpoints.Authenticator, + "_get_device_code", + side_effect=GetDeviceCodeError(status_code=500, message="gh down"), + ): + with pytest.raises(HTTPException) as exc_info: + await start_oauth( + StartRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert exc_info.value.status_code == 502 + assert "gh down" in exc_info.value.detail + + +class TestStatusAndCancel: + @pytest.mark.asyncio + async def test_status_returns_pending(self): + with _sessions_lock: + _sessions["s1"] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + } + response = await oauth_status(session_id="s1", user_api_key_dict=_admin()) + assert response.status == "pending" + + @pytest.mark.asyncio + async def test_cancel_flips_pending(self): + with _sessions_lock: + _sessions["s1"] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + "cancelled": False, + } + result = await oauth_cancel(session_id="s1", user_api_key_dict=_admin()) + assert result == {"success": True} + with _sessions_lock: + assert _sessions["s1"]["status"] == "cancelled" + + +class TestRefreshEndpoint: + @pytest.mark.asyncio + async def test_admin_only(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), user_api_key_dict=_non_admin() + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_rejects_view_only_admin(self): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_view_only_admin(), + ) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_returns_expires_at_on_success(self): + with patch.object( + DBAuthenticator, + "force_refresh_api_key", + return_value={"token": "fresh", "expires_at": 1700000000}, + ): + response = await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert response.credential_name == "c" + assert response.api_key_expires_at == 1700000000 + + @pytest.mark.asyncio + async def test_502_on_refresh_failure(self): + with patch.object( + DBAuthenticator, + "force_refresh_api_key", + side_effect=RefreshAPIKeyError(status_code=401, message="bad token"), + ): + with pytest.raises(HTTPException) as exc_info: + await oauth_refresh( + RefreshRequest(credential_name="c"), + user_api_key_dict=_admin(), + ) + assert exc_info.value.status_code == 502 + + +class TestBackgroundWorker: + def test_worker_persists_access_token_on_success(self, monkeypatch): + from litellm.proxy.copilot_oauth_endpoints.endpoints import ( + _run_device_code_flow, + ) + + session_id = "s1" + with _sessions_lock: + _sessions[session_id] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + "cancelled": False, + } + + auth = MagicMock() + auth._poll_for_access_token.return_value = "gho_fresh_token" + + store_mock = MagicMock() + monkeypatch.setattr( + DBAuthenticator, "store_access_token", lambda self, tok: store_mock(tok) + ) + + _run_device_code_flow( + session_id=session_id, + credential_name="c", + device_code_info={"device_code": "dc", "user_code": "UC"}, + authenticator=auth, + ) + + store_mock.assert_called_once_with("gho_fresh_token") + with _sessions_lock: + assert _sessions[session_id]["status"] == "success" + + def test_worker_marks_error_on_db_persist_failure(self, monkeypatch): + from litellm.proxy.copilot_oauth_endpoints.endpoints import ( + _run_device_code_flow, + ) + + session_id = "s1" + with _sessions_lock: + _sessions[session_id] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + "cancelled": False, + } + + auth = MagicMock() + auth._poll_for_access_token.return_value = "gho_fresh" + + def _boom(self, tok): + raise RuntimeError("prisma disconnected") + + monkeypatch.setattr(DBAuthenticator, "store_access_token", _boom) + + _run_device_code_flow( + session_id=session_id, + credential_name="c", + device_code_info={"device_code": "dc", "user_code": "UC"}, + authenticator=auth, + ) + + with _sessions_lock: + assert _sessions[session_id]["status"] == "error" + assert "DB persist failed" in _sessions[session_id]["message"] + + def test_worker_marks_error_on_poll_failure(self): + from litellm.proxy.copilot_oauth_endpoints.endpoints import ( + _run_device_code_flow, + ) + + session_id = "s1" + with _sessions_lock: + _sessions[session_id] = { + "status": "pending", + "credential_name": "c", + "expires_at": time.time() + 600, + "cancelled": False, + } + + auth = MagicMock() + auth._poll_for_access_token.side_effect = GetAccessTokenError( + status_code=408, message="timed out" + ) + + _run_device_code_flow( + session_id=session_id, + credential_name="c", + device_code_info={"device_code": "dc", "user_code": "UC"}, + authenticator=auth, + ) + + with _sessions_lock: + assert _sessions[session_id]["status"] == "error" + assert "timed out" in _sessions[session_id]["message"] diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx index 694a98201c6..10989bd5de1 100644 --- a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx @@ -4,6 +4,8 @@ import type { UploadProps } from "antd/es/upload"; import React, { useState } from "react"; import ProviderSpecificFields from "../add_model/provider_specific_fields"; import { Providers, providerLogoMap } from "../provider_info_helpers"; +import ChatGPTLoginButton from "./ChatGPTLoginButton"; +import CopilotLoginButton from "./CopilotLoginButton"; const { Link } = Typography; interface AddCredentialsModalProps { @@ -16,6 +18,7 @@ interface AddCredentialsModalProps { const AddCredentialsModal: React.FC = ({ open, onCancel, onAddCredential, uploadProps }) => { const [form] = Form.useForm(); const [selectedProvider, setSelectedProvider] = useState(Providers.OpenAI); + const credentialName = Form.useWatch("credential_name", form); const handleSubmit = (values: any) => { const filteredValues = Object.entries(values).reduce((acc, [key, value]) => { @@ -89,7 +92,32 @@ const AddCredentialsModal: React.FC = ({ open, onCance - + {/* + `selectedProvider` holds the enum *key* string at runtime + (``AntdSelect.Option value={providerEnum}`` binds the key, and the + ``as Providers`` cast on setState is not enforced). Compare against + the key names, not ``Providers.ChatGPT`` / ``Providers.GITHUB_COPILOT`` + which resolve to enum *values* and would never match. + */} + {(selectedProvider as unknown as keyof typeof Providers) === "ChatGPT" ? ( + { + onCancel(); + form.resetFields(); + }} + /> + ) : (selectedProvider as unknown as keyof typeof Providers) === "GITHUB_COPILOT" ? ( + { + onCancel(); + form.resetFields(); + }} + /> + ) : ( + + )} {/* Modal Footer */}
@@ -107,7 +135,10 @@ const AddCredentialsModal: React.FC = ({ open, onCance > Cancel - + {(selectedProvider as unknown as keyof typeof Providers) !== "ChatGPT" && + (selectedProvider as unknown as keyof typeof Providers) !== "GITHUB_COPILOT" && ( + + )}
diff --git a/ui/litellm-dashboard/src/components/model_add/ChatGPTLoginButton.tsx b/ui/litellm-dashboard/src/components/model_add/ChatGPTLoginButton.tsx new file mode 100644 index 00000000000..e05862ead1b --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/ChatGPTLoginButton.tsx @@ -0,0 +1,25 @@ +import React from "react"; + +import { + chatgptOauthCancelCall, + chatgptOauthStartCall, + chatgptOauthStatusCall, +} from "@/components/networking"; +import OAuthDeviceLoginButton from "./OAuthDeviceLoginButton"; + +interface ChatGPTLoginButtonProps { + credentialName?: string; + onSuccess: () => void; +} + +const ChatGPTLoginButton: React.FC = (props) => ( + +); + +export default ChatGPTLoginButton; diff --git a/ui/litellm-dashboard/src/components/model_add/CopilotLoginButton.tsx b/ui/litellm-dashboard/src/components/model_add/CopilotLoginButton.tsx new file mode 100644 index 00000000000..f8fa44688b3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/CopilotLoginButton.tsx @@ -0,0 +1,25 @@ +import React from "react"; + +import { + copilotOauthCancelCall, + copilotOauthStartCall, + copilotOauthStatusCall, +} from "@/components/networking"; +import OAuthDeviceLoginButton from "./OAuthDeviceLoginButton"; + +interface CopilotLoginButtonProps { + credentialName?: string; + onSuccess: () => void; +} + +const CopilotLoginButton: React.FC = (props) => ( + +); + +export default CopilotLoginButton; diff --git a/ui/litellm-dashboard/src/components/model_add/OAuthDeviceLoginButton.tsx b/ui/litellm-dashboard/src/components/model_add/OAuthDeviceLoginButton.tsx new file mode 100644 index 00000000000..80c99597cfc --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/OAuthDeviceLoginButton.tsx @@ -0,0 +1,162 @@ +import { Alert, Button, Space, Spin, Typography } from "antd"; +import React, { useCallback, useEffect, useRef, useState } from "react"; + +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import type { + ChatGPTOAuthStartResponse, + ChatGPTOAuthStatusResponse, +} from "@/components/networking"; + +const { Text, Link, Paragraph } = Typography; + +const POLL_INTERVAL_MS = 3000; + +export interface OAuthDeviceLoginButtonProps { + providerLabel: string; // e.g. "ChatGPT" or "GitHub Copilot" + credentialName?: string; + onSuccess: () => void; + startCall: (accessToken: string, credentialName: string) => Promise; + statusCall: (accessToken: string, sessionId: string) => Promise; + cancelCall: (accessToken: string, sessionId: string) => Promise; +} + +type Phase = + | { kind: "idle" } + | { + kind: "active"; + sessionId: string; + userCode: string; + verificationUrl: string; + } + | { kind: "success" } + | { kind: "error"; message: string }; + +const OAuthDeviceLoginButton: React.FC = ({ + providerLabel, + credentialName, + onSuccess, + startCall, + statusCall, + cancelCall, +}) => { + const { accessToken } = useAuthorized(); + const [phase, setPhase] = useState({ kind: "idle" }); + const pollTimer = useRef | null>(null); + + const stopPolling = useCallback(() => { + if (pollTimer.current) { + clearInterval(pollTimer.current); + pollTimer.current = null; + } + }, []); + + useEffect(() => stopPolling, [stopPolling]); + + const startLogin = useCallback(async () => { + if (!accessToken || !credentialName) return; + setPhase({ kind: "idle" }); + try { + const start = await startCall(accessToken, credentialName); + setPhase({ + kind: "active", + sessionId: start.session_id, + userCode: start.user_code, + verificationUrl: start.verification_url, + }); + + stopPolling(); + pollTimer.current = setInterval(async () => { + try { + const status = await statusCall(accessToken, start.session_id); + if (status.status === "success") { + stopPolling(); + setPhase({ kind: "success" }); + onSuccess(); + } else if (status.status === "error" || status.status === "cancelled") { + stopPolling(); + setPhase({ + kind: "error", + message: status.message || status.status, + }); + } + } catch { + // transient error — keep polling + } + }, POLL_INTERVAL_MS); + } catch (err) { + setPhase({ + kind: "error", + message: err instanceof Error ? err.message : "Failed to start login", + }); + } + }, [accessToken, credentialName, onSuccess, startCall, statusCall, stopPolling]); + + const cancel = useCallback(async () => { + if (phase.kind !== "active" || !accessToken) return; + stopPolling(); + try { + await cancelCall(accessToken, phase.sessionId); + } catch { + // ignore — cancel is best-effort + } + setPhase({ kind: "idle" }); + }, [accessToken, phase, cancelCall, stopPolling]); + + if (phase.kind === "active") { + return ( + + + + Open the link below in a new tab and enter this code: + + + {phase.userCode} + + + {phase.verificationUrl} + + + + Waiting for browser confirmation… + + + } + /> + + + ); + } + + if (phase.kind === "success") { + return ( + + ); + } + + return ( + + {phase.kind === "error" && ( + + )} + + Enter a credential name above, then click Sign in to receive a short + code to enter in your browser. + + + + ); +}; + +export default OAuthDeviceLoginButton; diff --git a/ui/litellm-dashboard/src/components/model_add/credentials.tsx b/ui/litellm-dashboard/src/components/model_add/credentials.tsx index af3c757955e..61a0b08c5d3 100644 --- a/ui/litellm-dashboard/src/components/model_add/credentials.tsx +++ b/ui/litellm-dashboard/src/components/model_add/credentials.tsx @@ -1,10 +1,12 @@ import { + chatgptOauthRefreshCall, + copilotOauthRefreshCall, credentialCreateCall, credentialDeleteCall, CredentialItem, credentialUpdateCall, } from "@/components/networking"; // Assume this is your networking function -import { PencilAltIcon, TrashIcon } from "@heroicons/react/outline"; +import { PencilAltIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline"; import { Badge, Button, @@ -17,7 +19,7 @@ import { TableRow, Text, } from "@tremor/react"; -import { Form } from "antd"; +import { Button as AntdButton, Form, Tooltip } from "antd"; import { UploadProps } from "antd/es/upload"; import { useState } from "react"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; @@ -125,6 +127,29 @@ const CredentialsPanel: React.FC = ({ uploadProps }) => { } }; + const getOAuthRefreshCall = (credential: CredentialItem) => { + const type = (credential.credential_info as Record | undefined)?.type; + if (type === "chatgpt_oauth") return chatgptOauthRefreshCall; + if (type === "copilot_oauth") return copilotOauthRefreshCall; + return null; + }; + + const handleRefreshCredential = async (credential: CredentialItem) => { + const refreshCall = getOAuthRefreshCall(credential); + if (!accessToken || !refreshCall) return; + try { + await refreshCall(accessToken, credential.credential_name); + NotificationsManager.success( + `Refreshed tokens for ${credential.credential_name}`, + ); + await refetchCredentials(); + } catch (error) { + NotificationsManager.error( + error instanceof Error ? error.message : "Failed to refresh credential", + ); + } + }; + const openDeleteModal = (credential: CredentialItem) => { setCredentialToDelete(credential); setIsDeleteModalOpen(true); @@ -166,22 +191,37 @@ const CredentialsPanel: React.FC = ({ uploadProps }) => { {renderProviderBadge((credential.credential_info?.custom_llm_provider as string) || "-")} -