diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index a5fb1c9c6b9..c0c92373550 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -2,17 +2,23 @@ import json import os import time from collections.abc import Mapping +from dataclasses import dataclass from functools import lru_cache -from typing import TYPE_CHECKING, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol import httpx from pydantic import JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.constants import HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS +from litellm.exceptions import Timeout from litellm.litellm_core_utils.asyncify import can_block_current_thread from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout -from litellm.llms.custom_httpx.http_handler import get_httpx_client +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # its params arg is a bare dict in http_handler + get_httpx_client, +) from .common_utils import ( CHATGPT_API_BASE, @@ -22,6 +28,7 @@ from .common_utils import ( CHATGPT_DEVICE_TOKEN_URL, CHATGPT_DEVICE_VERIFY_URL, CHATGPT_OAUTH_TOKEN_URL, + ChatGPTAuthError, GetAccessTokenError, GetDeviceCodeError, JsonObject, @@ -46,6 +53,13 @@ def _optional_str(value: JsonValue | None) -> str | None: return value if isinstance(value, str) else None +def access_token_expires_at(token: str) -> int | None: + exp: Final = decode_jwt_claims(token).get("exp") + if isinstance(exp, (int, float)): + return int(exp) + return None + + def get_chatgpt_auth_file( litellm_params: "Mapping[str, object] | GenericLiteLLMParams | None", ) -> str | None: @@ -171,10 +185,7 @@ class Authenticator: return time.time() >= float(derived_expires_at) - TOKEN_EXPIRY_SKEW_SECONDS def _get_expires_at(self, token: str) -> int | None: - exp: Final = decode_jwt_claims(token).get("exp") - if isinstance(exp, (int, float)): - return int(exp) - return None + return access_token_expires_at(token) def _login_device_code(self) -> dict[str, str]: cooldown_remaining: Final = self._get_device_code_cooldown_remaining(self._read_auth_file()) @@ -223,7 +234,7 @@ class Authenticator: interval: Final = data.get("interval") if not device_auth_id or not user_code: raise GetDeviceCodeError( - message=f"Device code response missing fields: {data}", + message="Device code response is missing device_auth_id or user_code", status_code=400, ) return { @@ -315,7 +326,7 @@ class Authenticator: id_token: Final = _optional_str(data.get("id_token")) if not access_token or not refresh_token or not id_token: raise GetAccessTokenError( - message=f"Token exchange response missing fields: {data}", + message="Token exchange response is missing access_token, refresh_token, or id_token", status_code=400, ) return { @@ -354,7 +365,7 @@ class Authenticator: id_token: Final = _optional_str(data.get("id_token")) if not access_token or not id_token: raise RefreshAccessTokenError( - message=f"Refresh response missing fields: {data}", + message="Refresh response is missing access_token or id_token", status_code=400, ) @@ -412,3 +423,182 @@ class Authenticator: break time.sleep(sleep_for) return None + + +class ChatGPTOAuthClient(Protocol): + async def post( + self, + url: str, + data: dict[str, str] | str | bytes | None = None, # mutable-ok: mirrors the untyped client's dict parameters + json: dict[str, str] | None = None, # mutable-ok: mirrors the untyped client's dict parameters + params: dict[str, str] | None = None, # mutable-ok: mirrors the untyped client's dict parameters + headers: dict[str, str] | None = None, # mutable-ok: mirrors the untyped client's dict parameters + timeout: float | httpx.Timeout | None = None, + ) -> httpx.Response: ... + + +@dataclass(frozen=True, slots=True) +class ChatGPTDeviceLogin: + device_auth_id: str + user_code: str + interval: int + + +@dataclass(frozen=True, slots=True, repr=False) +class ChatGPTAuthorizationCode: + authorization_code: str + code_verifier: str + + +@dataclass(frozen=True, slots=True, repr=False) +class ChatGPTTokens: + access_token: str + refresh_token: str + id_token: str + + +def new_chatgpt_oauth_client() -> ChatGPTOAuthClient: + return get_async_httpx_client("chatgpt") + + +async def _apost_json_object( + client: ChatGPTOAuthClient, + url: str, + failure: str, + error_type: type[ChatGPTAuthError], + json: dict[str, str] | None = None, # mutable-ok: mirrors the untyped client's dict parameters + data: dict[str, str] | None = None, # mutable-ok: mirrors the untyped client's dict parameters +) -> JsonObject: + try: + response: Final = await client.post(url, json=json, data=data, timeout=_token_refresh_timeout()) + except httpx.HTTPStatusError as exc: + raise error_type( + message=f"{failure} with status {exc.response.status_code}", status_code=exc.response.status_code + ) from None + except (httpx.HTTPError, Timeout) as exc: + raise error_type(message=f"{failure} ({type(exc).__name__})", status_code=502) from None + try: + return _JSON_OBJECT_ADAPTER.validate_python(response.json()) + except ValueError: + raise error_type(message=f"{failure}: the response was not a JSON object", status_code=502) from None + + +async def astart_chatgpt_device_login(client: ChatGPTOAuthClient) -> ChatGPTDeviceLogin: + data: Final = await _apost_json_object( + client, + CHATGPT_DEVICE_CODE_URL, + "ChatGPT device code request failed", + GetDeviceCodeError, + json={"client_id": CHATGPT_CLIENT_ID}, + ) + device_auth_id: Final = _optional_str(data.get("device_auth_id")) + user_code: Final = _optional_str(data.get("user_code") or data.get("usercode")) + if not device_auth_id or not user_code: + raise GetDeviceCodeError(message="Device code response is missing device_auth_id or user_code", status_code=502) + return ChatGPTDeviceLogin( + device_auth_id=device_auth_id, user_code=user_code, interval=_poll_interval_seconds(data.get("interval")) + ) + + +def _poll_interval_seconds(raw: JsonValue | None) -> int: + try: + return max(int(str(raw)), DEVICE_CODE_POLL_SLEEP_SECONDS) + except ValueError: + return DEVICE_CODE_POLL_SLEEP_SECONDS + + +async def apoll_chatgpt_device_login( + client: ChatGPTOAuthClient, device_auth_id: str, user_code: str +) -> ChatGPTAuthorizationCode | None: + try: + response: Final = await client.post( + CHATGPT_DEVICE_TOKEN_URL, + json={"device_auth_id": device_auth_id, "user_code": user_code}, + timeout=_token_refresh_timeout(), + ) + except httpx.HTTPStatusError as exc: + if exc.response.status_code in (403, 404): + return None + raise GetAccessTokenError( + message=f"ChatGPT device login poll failed with status {exc.response.status_code}", + status_code=exc.response.status_code, + ) from None + except (httpx.HTTPError, Timeout) as exc: + raise GetAccessTokenError( + message=f"ChatGPT device login poll failed ({type(exc).__name__})", status_code=502 + ) from None + if response.status_code in (403, 404): + return None + try: + data: Final = _JSON_OBJECT_ADAPTER.validate_python(response.json()) + except ValueError: + return None + authorization_code: Final = _optional_str(data.get("authorization_code")) + code_verifier: Final = _optional_str(data.get("code_verifier")) + if not authorization_code or not code_verifier: + return None + return ChatGPTAuthorizationCode(authorization_code=authorization_code, code_verifier=code_verifier) + + +async def aexchange_chatgpt_authorization_code( + client: ChatGPTOAuthClient, code: ChatGPTAuthorizationCode +) -> ChatGPTTokens: + data: Final = await _apost_json_object( + client, + CHATGPT_OAUTH_TOKEN_URL, + "ChatGPT token exchange failed", + GetAccessTokenError, + data={ + "grant_type": "authorization_code", + "code": code.authorization_code, + "redirect_uri": f"{CHATGPT_AUTH_BASE}/deviceauth/callback", + "client_id": CHATGPT_CLIENT_ID, + "code_verifier": code.code_verifier, + }, + ) + access_token: Final = _optional_str(data.get("access_token")) + refresh_token: Final = _optional_str(data.get("refresh_token")) + id_token: Final = _optional_str(data.get("id_token")) + if not access_token or not refresh_token or not id_token: + raise GetAccessTokenError( + message="Token exchange response is missing access_token, refresh_token, or id_token", status_code=502 + ) + return ChatGPTTokens(access_token=access_token, refresh_token=refresh_token, id_token=id_token) + + +async def arefresh_chatgpt_tokens(client: ChatGPTOAuthClient, refresh_token: str) -> ChatGPTTokens: + data: Final = await _apost_json_object( + client, + CHATGPT_OAUTH_TOKEN_URL, + "ChatGPT token refresh failed", + RefreshAccessTokenError, + json={ + "client_id": CHATGPT_CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "scope": "openid profile email", + }, + ) + access_token: Final = _optional_str(data.get("access_token")) + id_token: Final = _optional_str(data.get("id_token")) + if not access_token or not id_token: + raise RefreshAccessTokenError(message="Refresh response is missing access_token or id_token", status_code=502) + return ChatGPTTokens( + access_token=access_token, + refresh_token=_optional_str(data.get("refresh_token")) or refresh_token, + id_token=id_token, + ) + + +def chatgpt_login_values(tokens: ChatGPTTokens) -> Mapping[str, str]: + expires_at: Final = access_token_expires_at(tokens.access_token) + account_id: Final = extract_chatgpt_account_id(tokens.id_token or tokens.access_token) + return MappingProxyType( + { + "chatgpt_access_token": tokens.access_token, + "chatgpt_refresh_token": tokens.refresh_token, + "chatgpt_id_token": tokens.id_token, + **({"chatgpt_account_id": account_id} if account_id else {}), + **({"chatgpt_expires_at": str(expires_at)} if expires_at is not None else {}), + } + ) diff --git a/litellm/models/credentials.py b/litellm/models/credentials.py index e6faf461104..99a77d2ae5f 100644 --- a/litellm/models/credentials.py +++ b/litellm/models/credentials.py @@ -125,3 +125,35 @@ class UserConnectionDeleteResponse(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) status: str + + +class ChatGPTLoginStartRequest(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + owner_team_id: str | None = None + owner_key_hash: str | None = None + display_name: str | None = None + + +class ChatGPTLoginFlowHandle(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + credential_name: str + device_auth_id: str = Field(repr=False) + user_code: str = Field(repr=False) + owner_team_id: str | None = None + owner_key_hash: str | None = None + display_name: str | None = None + interval: int + expires_at: float + + +ChatGPTLoginPollStatus: TypeAlias = Literal["pending", "connected"] + + +class ChatGPTLoginPollResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + status: ChatGPTLoginPollStatus + interval: int | None = None + account_id: str | None = None diff --git a/litellm/proxy/credential_endpoints/chatgpt_subscription_logins.py b/litellm/proxy/credential_endpoints/chatgpt_subscription_logins.py new file mode 100644 index 00000000000..158a03abb18 --- /dev/null +++ b/litellm/proxy/credential_endpoints/chatgpt_subscription_logins.py @@ -0,0 +1,328 @@ +import asyncio +import json +import math +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import partial +from types import MappingProxyType +from typing import Final, Protocol + +from pydantic import TypeAdapter + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.llms.chatgpt.authenticator import ( + TOKEN_EXPIRY_SKEW_SECONDS, + ChatGPTOAuthClient, + arefresh_chatgpt_tokens, + chatgpt_login_values, + new_chatgpt_oauth_client, +) +from litellm.llms.chatgpt.common_utils import ChatGPTAuthError +from litellm.llms.chatgpt.subscription_login import ( + CHATGPT_OWNER_KEY_HASH_KEY, + CHATGPT_OWNER_TEAM_ID_KEY, + CHATGPT_PROVIDER, + is_chatgpt_subscription_login, +) +from litellm.proxy.common_utils.credential_hydration import decrypted_credential, stored_credential_provider +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( + InProcessRefreshCoordinator, + OAuthToken, + RefreshCoordinator, +) +from litellm.types.utils import CredentialItem + +_COORDINATOR_NAMESPACE: Final = "chatgpt_subscription" +_IN_PROCESS_COORDINATOR: Final = InProcessRefreshCoordinator() +_NO_LOGINS: Final[Mapping[str, str]] = MappingProxyType({}) +_CREDENTIAL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +_REFRESH_TOKEN_KEY: Final = "chatgpt_refresh_token" + + +class CredentialRows(Protocol): + async def find_by_name(self, credential_name: str) -> CredentialItem | None: ... + + async def update_if_value_unchanged( + self, credential_name: str, value_key: str, stored_value: str, data: Mapping[str, object] + ) -> bool: ... + + +def _non_empty_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _credential_info(credential: CredentialItem | None) -> Mapping[str, object]: + return _CREDENTIAL_INFO_ADAPTER.validate_python( + (credential.credential_info if credential is not None else None) # pyright: ignore[reportUnknownMemberType] # untyped credential field + or {} + ) + + +def chatgpt_login_owner(credential: CredentialItem) -> tuple[str | None, str | None]: + info: Final = _credential_info(credential) + team_id: Final = _non_empty_str(info.get(CHATGPT_OWNER_TEAM_ID_KEY)) + key_hash: Final = _non_empty_str(info.get(CHATGPT_OWNER_KEY_HASH_KEY)) + if (team_id is None) == (key_hash is None): + return None, None + return team_id, key_hash + + +def _owned_by(credential: CredentialItem, team_id: str | None, key_hash: str | None) -> bool: + if not is_chatgpt_subscription_login(credential): + return False + owner_team_id, owner_key_hash = chatgpt_login_owner(credential) + return (owner_team_id is not None and owner_team_id == team_id) or ( + owner_key_hash is not None and owner_key_hash == key_hash + ) + + +def owned_chatgpt_subscription_logins(team_id: str | None, key_hash: str | None) -> tuple[CredentialItem, ...]: + return tuple(credential for credential in litellm.credential_list if _owned_by(credential, team_id, key_hash)) + + +def _expires_at(values: Mapping[str, object]) -> float | None: + raw: Final = values.get("chatgpt_expires_at") + if not isinstance(raw, (str, int, float)) or isinstance(raw, bool): + return None + try: + parsed: Final = float(raw) + except ValueError: + return None + return parsed if math.isfinite(parsed) else None + + +def _fresh_access_token(values: Mapping[str, object], now: float) -> str | None: + access_token: Final = _non_empty_str(values.get("chatgpt_access_token")) + expires_at: Final = _expires_at(values) + if access_token is None or expires_at is None or now >= expires_at - TOKEN_EXPIRY_SKEW_SECONDS: + return None + return access_token + + +def _in_memory_values(credential_name: str) -> Mapping[str, object] | None: + current: Final = CredentialAccessor.find_credential(credential_name) + return MappingProxyType(dict(current.credential_values)) if current is not None else None + + +@dataclass(frozen=True, slots=True) +class ChatGPTLoginRefresher: + rows: CredentialRows + oauth_client: ChatGPTOAuthClient + coordinator: RefreshCoordinator + shared_coordinator: RefreshCoordinator | None = None + clock: Callable[[], float] = time.time + encrypt: Callable[[str], str] = encrypt_value_helper + + async def access_token(self, owned: CredentialItem) -> str | None: + owner: Final = chatgpt_login_owner(owned) + refresh: Final = partial(self._refresh, owned.credential_name, owner) + reread: Final = partial(self._reread, owned.credential_name, owner) + flight: Final = json.dumps([owned.credential_name, *owner]) + shared: Final = self.shared_coordinator + token: Final = await self.coordinator.run( + _COORDINATOR_NAMESPACE, + flight, + refresh if shared is None else partial(shared.run, _COORDINATOR_NAMESPACE, flight, refresh, reread), + reread, + ) + return token.access_token if token is not None else None + + async def _stored(self, credential_name: str) -> CredentialItem | None: + try: + return await self.rows.find_by_name(credential_name) + except Exception: # noqa: BLE001 # a DB outage fails this login closed instead of breaking the request + verbose_proxy_logger.warning( + "Could not read ChatGPT subscription login %s from the database", credential_name + ) + return None + + def _remember( + self, stored: CredentialItem, values: Mapping[str, object], seen: Mapping[str, object] | None + ) -> None: + if seen is None or _in_memory_values(stored.credential_name) != seen: + return + CredentialAccessor.upsert_credentials( + [ + CredentialItem( + credential_name=stored.credential_name, + display_name=stored.display_name, + credential_values=dict(values), + credential_info=dict(_credential_info(stored)), + ) + ] + ) + + def _fresh_from( + self, stored: CredentialItem, seen: Mapping[str, object] | None + ) -> tuple[OAuthToken | None, Mapping[str, object]]: + values: Final = decrypted_credential(stored).credential_values + access_token: Final = _fresh_access_token(values, self.clock()) + if access_token is None: + return None, values + self._remember(stored, values, seen) + return OAuthToken(access_token=access_token, expires_at=_expires_at(values)), values + + async def _reread(self, credential_name: str, owner: tuple[str | None, str | None]) -> OAuthToken | None: + seen: Final = _in_memory_values(credential_name) + stored: Final = await self._stored(credential_name) + if stored is None or not _owned_by(stored, *owner): + return None + token, _ = self._fresh_from(stored, seen) + return token + + async def _refresh(self, credential_name: str, owner: tuple[str | None, str | None]) -> OAuthToken | None: + seen: Final = _in_memory_values(credential_name) + stored: Final = await self._stored(credential_name) + if stored is None or not _owned_by(stored, *owner): + return None + fresh, values = self._fresh_from(stored, seen) + if fresh is not None: + return fresh + refresh_token: Final = _non_empty_str(values.get(_REFRESH_TOKEN_KEY)) + stored_refresh_token: Final = _non_empty_str(stored.credential_values.get(_REFRESH_TOKEN_KEY)) + if refresh_token is None or stored_refresh_token is None: + return None + try: + tokens: Final = await arefresh_chatgpt_tokens(self.oauth_client, refresh_token) + except ChatGPTAuthError as exc: + verbose_proxy_logger.warning( + "ChatGPT subscription login %s could not be refreshed (status %s). Sign in again for it.", + credential_name, + exc.status_code, + ) + return None + login_values: Final = chatgpt_login_values(tokens) + if not await self._save(stored, stored_refresh_token, login_values): + return await self._reread(credential_name, owner) + refreshed_values: Final = {**values, **login_values} + self._remember(stored, refreshed_values, seen) + return OAuthToken(access_token=tokens.access_token, expires_at=_expires_at(refreshed_values)) + + async def _save(self, stored: CredentialItem, stored_refresh_token: str, login_values: Mapping[str, str]) -> bool: + encrypted: Final = {key: self.encrypt(value) for key, value in login_values.items()} + try: + return await self.rows.update_if_value_unchanged( + stored.credential_name, + _REFRESH_TOKEN_KEY, + stored_refresh_token, + {"credential_values": json.dumps({**stored.credential_values, **encrypted})}, + ) + except Exception: # noqa: BLE001 # the rotated tokens still serve this replica until they expire + verbose_proxy_logger.warning( + "ChatGPT subscription login %s was refreshed but could not be saved to the database. " + "Sign in again for it if requests start failing after the new token expires.", + stored.credential_name, + ) + return True + + +def runtime_chatgpt_login_refresher() -> ChatGPTLoginRefresher | None: + from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient + from litellm.proxy.mcp_server.outbound_credentials.runtime_refresh_coordinator import ( + runtime_refresh_coordinator, + ) + from litellm.proxy.proxy_server import prisma_client + from litellm.repositories.credentials_repository import CredentialsRepository + + if prisma_client is None: + return None + return ChatGPTLoginRefresher( + rows=CredentialsRepository(WriterPinnedClient(prisma_client.db)), + oauth_client=new_chatgpt_oauth_client(), + coordinator=_IN_PROCESS_COORDINATOR, + shared_coordinator=runtime_refresh_coordinator(), + ) + + +async def resolve_chatgpt_subscription_logins( + team_id: str | None, + key_hash: str | None, + refresher_factory: Callable[[], ChatGPTLoginRefresher | None] = runtime_chatgpt_login_refresher, + clock: Callable[[], float] = time.time, +) -> Mapping[str, str]: + owned: Final = owned_chatgpt_subscription_logins(team_id, key_hash) + if not owned: + return _NO_LOGINS + now: Final = clock() + fresh: Final = { + credential.credential_name: token + for credential in owned + if (token := _fresh_access_token(credential.credential_values, now)) is not None + } + stale: Final = tuple(credential for credential in owned if credential.credential_name not in fresh) + refresher: Final = refresher_factory() if stale else None + if refresher is None: + return MappingProxyType(fresh) + refreshed: Final = await asyncio.gather(*(refresher.access_token(credential) for credential in stale)) + return MappingProxyType( + { + **fresh, + **{ + credential.credential_name: token + for credential, token in zip(stale, refreshed, strict=True) + if token is not None + }, + } + ) + + +class CredentialWrites(Protocol): + async def find_by_name(self, credential_name: str) -> CredentialItem | None: ... + + async def create(self, data: Mapping[str, object]) -> object: ... + + async def update_by_name(self, credential_name: str, data: Mapping[str, object]) -> object: ... + + +_OWNER_KEYS: Final = frozenset({CHATGPT_OWNER_TEAM_ID_KEY, CHATGPT_OWNER_KEY_HASH_KEY}) + + +def blocks_chatgpt_login(existing: CredentialItem | None) -> bool: + if existing is None or not existing.credential_values or is_chatgpt_subscription_login(existing): + return False + return stored_credential_provider(_credential_info(existing).get("custom_llm_provider")) != CHATGPT_PROVIDER + + +async def save_chatgpt_login( + rows: CredentialWrites, + credential_name: str, + display_name: str | None, + login_values: Mapping[str, str], + owner: tuple[str | None, str | None], + actor: str, + encrypt: Callable[[str], str] = encrypt_value_helper, +) -> CredentialItem: + existing: Final = await rows.find_by_name(credential_name) + owner_team_id, owner_key_hash = owner + kept_info: Final = {key: value for key, value in _credential_info(existing).items() if key not in _OWNER_KEYS} + credential_info: Final = { + **kept_info, + "custom_llm_provider": CHATGPT_PROVIDER, + **({CHATGPT_OWNER_TEAM_ID_KEY: owner_team_id} if owner_team_id is not None else {}), + **({CHATGPT_OWNER_KEY_HASH_KEY: owner_key_hash} if owner_key_hash is not None else {}), + } + saved_display_name: Final = display_name or (existing.display_name if existing else None) + stored_fields: Final = { + "credential_values": json.dumps({key: encrypt(value) for key, value in login_values.items()}), + "credential_info": json.dumps(credential_info), + **({"display_name": saved_display_name} if saved_display_name is not None else {}), + "updated_by": actor, + } + if existing is None: + await rows.create({"credential_name": credential_name, **stored_fields, "created_by": actor}) + else: + await rows.update_by_name(credential_name, stored_fields) + saved: Final = CredentialItem( + credential_name=credential_name, + display_name=saved_display_name, + credential_values=dict(login_values), + credential_info=credential_info, + ) + CredentialAccessor.upsert_credentials([saved]) + return saved diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 57c7a753d71..1fc9383fd84 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -22,6 +22,7 @@ from litellm.constants import ( GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_HOST_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, + LITELLM_PROXY_ADMIN_NAME, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import get_masked_values, has_sensitive_keyword @@ -31,7 +32,20 @@ from litellm.llms.anthropic.wif import ( UnbuildableIdentitySource, anthropic_internal_issuer_jwks, ) +from litellm.llms.chatgpt.authenticator import ( + DEVICE_CODE_TIMEOUT_SECONDS, + ChatGPTOAuthClient, + aexchange_chatgpt_authorization_code, + apoll_chatgpt_device_login, + astart_chatgpt_device_login, + chatgpt_login_values, + new_chatgpt_oauth_client, +) +from litellm.llms.chatgpt.common_utils import CHATGPT_DEVICE_VERIFY_URL, ChatGPTAuthError from litellm.models.credentials import ( + ChatGPTLoginFlowHandle, + ChatGPTLoginPollResponse, + ChatGPTLoginStartRequest, CredentialView, UpdateCredentialItem, UserConnectionDeleteResponse, @@ -67,6 +81,7 @@ from litellm.repositories.credentials_repository import CredentialsRepository from litellm.types.router import server_owned_wif_fields_named, server_owned_wif_fields_written from litellm.types.utils import CreateCredentialItem, CredentialItem +from .chatgpt_subscription_logins import blocks_chatgpt_login, save_chatgpt_login from .user_provider_credentials import ( GithubCopilotUserConnectionPayload, decode_user_provider_credential, @@ -819,6 +834,175 @@ async def delete_user_connection( raise handle_exception_on_proxy(e) +def get_chatgpt_oauth_client() -> ChatGPTOAuthClient: + return new_chatgpt_oauth_client() + + +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException(status_code=403, detail="Only a proxy admin can sign in to ChatGPT for a team or key") + + +def _chatgpt_login_owner(body: ChatGPTLoginStartRequest) -> tuple[str | None, str | None]: + owner_team_id: Final = body.owner_team_id or None + owner_key_hash: Final = body.owner_key_hash or None + if (owner_team_id is None) == (owner_key_hash is None): + raise HTTPException(status_code=400, detail="Set exactly one of owner_team_id or owner_key_hash") + return owner_team_id, owner_key_hash + + +async def _require_chatgpt_login_owner_exists( + prisma_client: object, owner_team_id: str | None, owner_key_hash: str | None +) -> None: + from litellm.repositories.team_repository import TeamRepository + from litellm.repositories.verification_token_repository import VerificationTokenRepository + + if owner_team_id is not None and await TeamRepository(prisma_client).find_by_id(owner_team_id) is None: + raise HTTPException(status_code=400, detail=f"Team '{owner_team_id}' not found") + if ( + owner_key_hash is not None + and await VerificationTokenRepository(prisma_client).find_by_id(owner_key_hash) is None + ): + raise HTTPException(status_code=400, detail="Virtual key not found. Send the key's hashed token (Key ID)") + + +async def _reject_non_chatgpt_credential(credentials_repository: CredentialsRepository, credential_name: str) -> None: + _reject_config_defined_name(credential_name) + if blocks_chatgpt_login(await credentials_repository.find_by_name(credential_name)): + raise ProxyException( + message=f"Credential '{credential_name}' belongs to another provider. Pick another name.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_409_CONFLICT, + param="credential_name", + ) + + +def _decode_chatgpt_flow_handle(flow_handle: str) -> ChatGPTLoginFlowHandle | None: + decrypted: Final = decrypt_value_helper( + value=flow_handle, + key="chatgpt_login_flow_handle", + exception_type="debug", + return_original_value=False, + ) + if decrypted is None: + return None + try: + return ChatGPTLoginFlowHandle.model_validate_json(decrypted) + except ValidationError: + return None + + +def _chatgpt_upstream_error(error: ChatGPTAuthError) -> HTTPException: + return HTTPException(status_code=502, detail=error.message) + + +@router.post( + "/credentials/{credential_name:path}/chatgpt_login/start", + dependencies=[Depends(user_api_key_auth)], + tags=["credential management"], + response_model=UserConnectionStartResponse, +) +async def start_chatgpt_login( + request: Request, + fastapi_response: Response, + body: ChatGPTLoginStartRequest, + oauth_client: Annotated[ChatGPTOAuthClient, Depends(get_chatgpt_oauth_client)], + credential_name: str = Path(..., description="The credential name, percent-decoded"), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default +) -> UserConnectionStartResponse: + """Start a ChatGPT device-code sign-in whose login only the given team or virtual key can use. Proxy admin only.""" + from litellm.proxy.proxy_server import prisma_client + + try: + _require_proxy_admin(user_api_key_dict) + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + owner_team_id, owner_key_hash = _chatgpt_login_owner(body) + display_name: Final = _normalized_display_name(body.display_name) + await _require_chatgpt_login_owner_exists(prisma_client, owner_team_id, owner_key_hash) + await _reject_non_chatgpt_credential( + CredentialsRepository(WriterPinnedClient(prisma_client.db)), credential_name + ) + try: + login: Final = await astart_chatgpt_device_login(oauth_client) + except ChatGPTAuthError as e: + raise _chatgpt_upstream_error(e) from None + handle: Final = encrypt_value_helper( + ChatGPTLoginFlowHandle( + credential_name=credential_name, + device_auth_id=login.device_auth_id, + user_code=login.user_code, + owner_team_id=owner_team_id, + owner_key_hash=owner_key_hash, + display_name=display_name, + interval=login.interval, + expires_at=time.time() + DEVICE_CODE_TIMEOUT_SECONDS, + ).model_dump_json() + ) + return UserConnectionStartResponse( + user_code=login.user_code, + verification_uri=CHATGPT_DEVICE_VERIFY_URL, + expires_in=DEVICE_CODE_TIMEOUT_SECONDS, + interval=login.interval, + flow_handle=handle, + ) + except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy error contract + raise handle_exception_on_proxy(e) + + +@router.post( + "/credentials/{credential_name:path}/chatgpt_login/poll", + dependencies=[Depends(user_api_key_auth)], + tags=["credential management"], + response_model=ChatGPTLoginPollResponse, +) +async def poll_chatgpt_login( + request: Request, + fastapi_response: Response, + body: UserConnectionPollRequest, + oauth_client: Annotated[ChatGPTOAuthClient, Depends(get_chatgpt_oauth_client)], + credential_name: str = Path(..., description="The credential name, percent-decoded"), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI resolves the dependency from the default +) -> ChatGPTLoginPollResponse: + """Poll a ChatGPT sign-in once and save the login on the credential when it completes. Proxy admin only.""" + from litellm.proxy.proxy_server import prisma_client + + try: + _require_proxy_admin(user_api_key_dict) + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + handle: Final = _decode_chatgpt_flow_handle(body.flow_handle) + if handle is None or handle.credential_name != credential_name or handle.expires_at <= time.time(): + raise HTTPException(status_code=400, detail="invalid or expired flow_handle") + try: + code: Final = await apoll_chatgpt_device_login(oauth_client, handle.device_auth_id, handle.user_code) + if code is None: + return ChatGPTLoginPollResponse(status="pending", interval=handle.interval) + tokens: Final = await aexchange_chatgpt_authorization_code(oauth_client, code) + except ChatGPTAuthError as e: + raise _chatgpt_upstream_error(e) from None + credentials_repository: Final = CredentialsRepository(WriterPinnedClient(prisma_client.db)) + await _reject_non_chatgpt_credential(credentials_repository, credential_name) + login_values: Final = chatgpt_login_values(tokens) + await save_chatgpt_login( + credentials_repository, + credential_name=credential_name, + display_name=handle.display_name, + login_values=login_values, + owner=(handle.owner_team_id, handle.owner_key_hash), + actor=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + ) + return ChatGPTLoginPollResponse(status="connected", account_id=login_values.get("chatgpt_account_id")) + except Exception as e: # noqa: BLE001 # endpoint boundary: every failure becomes the proxy error contract + raise handle_exception_on_proxy(e) + + @with_service_target("user_provider_connections") async def _purge_user_connections_for_credential(credential_name: str) -> None: """Drop every user connection under a per-user credential and clear caches. diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 14d941ff62c..1938fc83ced 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3047,6 +3047,12 @@ async def add_litellm_data_to_request( ), ) + await _resolve_chatgpt_subscription_logins_for_request( + data=cast("dict[str, object]", data), # cast-ok: data is the untyped request body dict + team_id=user_api_key_dict.team_id, + key_hash=user_api_key_dict.token, + ) + if holds_api_key_secret_reference(data): # pyright: ignore[reportUnknownArgumentType] # data is an untyped dict raise HTTPException(status_code=400, detail=API_KEY_SECRET_REFERENCE_REJECTION) @@ -3474,6 +3480,27 @@ async def _resolve_user_provider_credentials_for_request( secret_fields["user_provider_credentials_user_id"] = user_id +async def _resolve_chatgpt_subscription_logins_for_request( + data: dict[str, object], # mutable-ok: writes resolved logins into the nested secret_fields dict + team_id: str | None, + key_hash: str | None, +) -> None: + from litellm.types.proxy.litellm_pre_call_utils import RedactedDict + + from .credential_endpoints.chatgpt_subscription_logins import resolve_chatgpt_subscription_logins + + secret_fields: Final = data.get("secret_fields") + if not isinstance(secret_fields, dict): + return + try: + logins: Final = await resolve_chatgpt_subscription_logins(team_id=team_id, key_hash=key_hash) + except Exception: # noqa: BLE001 # a login lookup failure fails only that login's deployments closed, never the request + verbose_proxy_logger.exception("Failed to load ChatGPT subscription logins for the request") + return + if logins: + secret_fields["chatgpt_subscription_logins"] = RedactedDict(dict(logins)) + + def _warn_stale_team_alias_once(warning_key: str, message: str, *args: str) -> None: if warning_key in _STALE_TEAM_ALIAS_WARNING_KEYS: return diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index cb1bc1d101b..29d69dc14d7 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -80,5 +80,19 @@ class CredentialsRepository: ) -> "prisma_models.LiteLLM_CredentialsTable | None": return await self.table.update(where={"credential_name": credential_name}, data=data) + async def update_if_value_unchanged( + self, credential_name: str, value_key: str, stored_value: str, data: Mapping[str, object] + ) -> bool: + from prisma import Json + + updated_rows: Final = await self.table.update_many( + where={ + "credential_name": credential_name, + "credential_values": {"path": [value_key], "equals": Json(stored_value)}, + }, + data=data, + ) + return updated_rows > 0 + async def delete_by_name(self, credential_name: str) -> "prisma_models.LiteLLM_CredentialsTable | None": return await self.table.delete(where={"credential_name": credential_name}) diff --git a/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py index 2e9e3462bee..4e4036d680a 100644 --- a/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py +++ b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py @@ -1,7 +1,10 @@ import base64 import json import time +import traceback +from collections.abc import Awaitable, Callable from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from pathlib import Path from typing import Final from unittest.mock import MagicMock, mock_open, patch @@ -13,12 +16,30 @@ import respx import litellm from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS from litellm.llms.chatgpt.authenticator import ( + DEVICE_CODE_POLL_SLEEP_SECONDS, TOKEN_REFRESH_TIMEOUT_SECONDS, Authenticator, + ChatGPTAuthorizationCode, + ChatGPTTokens, + aexchange_chatgpt_authorization_code, + apoll_chatgpt_device_login, + arefresh_chatgpt_tokens, + astart_chatgpt_device_login, + chatgpt_login_values, get_cached_authenticator, get_chatgpt_auth_file, ) -from litellm.llms.chatgpt.common_utils import GetAccessTokenError +from litellm.llms.chatgpt.common_utils import ( + CHATGPT_AUTH_BASE, + CHATGPT_CLIENT_ID, + CHATGPT_DEVICE_CODE_URL, + CHATGPT_DEVICE_TOKEN_URL, + CHATGPT_OAUTH_TOKEN_URL, + ChatGPTAuthError, + GetAccessTokenError, + GetDeviceCodeError, + RefreshAccessTokenError, +) from litellm.types.router import GenericLiteLLMParams @@ -309,3 +330,374 @@ def test_refresh_keeps_saved_workspace(tmp_path: Path, monkeypatch: pytest.Monke assert authenticator.get_access_token() == token assert authenticator.get_account_id(token) == "acct-selected-workspace" assert json.loads(auth_file.read_text())["refresh_token"] == "example-rotated-refresh" + + +SECRET_IN_RESPONSE: Final = "secret-from-the-oauth-response" + + +@dataclass(frozen=True) +class _Post: + url: str + data: object + json: object + + +class _FakeOAuthClient: + def __init__(self, status_code: int = 200, body: object = None, error: Exception | None = None) -> None: + self.status_code: Final = status_code + self.body: Final = body + self.error: Final = error + self.posts: Final[list[_Post]] = [] + + async def post( + self, + url: str, + data: dict[str, str] | str | bytes | None = None, + json: dict[str, str] | None = None, + params: dict[str, str] | None = None, + headers: dict[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> httpx.Response: + self.posts.append(_Post(url=url, data=data, json=json)) + if self.error is not None: + raise self.error + request: Final = httpx.Request("POST", url) + response: Final = ( + httpx.Response(self.status_code, content=self.body, request=request) + if isinstance(self.body, bytes) + else httpx.Response(self.status_code, json=self.body, request=request) + ) + response.raise_for_status() + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_code_field", ["user_code", "usercode"]) +async def test_astart_chatgpt_device_login_reads_the_user_code_under_either_field_name(user_code_field: str) -> None: + client: Final = _FakeOAuthClient( + body={"device_auth_id": "device-auth-example", user_code_field: "ABCD-1234", "interval": "9"} + ) + + login: Final = await astart_chatgpt_device_login(client) + + assert (login.device_auth_id, login.user_code, login.interval) == ("device-auth-example", "ABCD-1234", 9) + assert client.posts == [_Post(url=CHATGPT_DEVICE_CODE_URL, data=None, json={"client_id": CHATGPT_CLIENT_ID})] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("interval", "expected"), + [ + pytest.param(DEVICE_CODE_POLL_SLEEP_SECONDS + 7, DEVICE_CODE_POLL_SLEEP_SECONDS + 7, id="above_the_floor"), + pytest.param(DEVICE_CODE_POLL_SLEEP_SECONDS, DEVICE_CODE_POLL_SLEEP_SECONDS, id="at_the_floor"), + pytest.param(1, DEVICE_CODE_POLL_SLEEP_SECONDS, id="below_the_floor"), + pytest.param(0, DEVICE_CODE_POLL_SLEEP_SECONDS, id="zero"), + pytest.param("not-a-number", DEVICE_CODE_POLL_SLEEP_SECONDS, id="not_a_number"), + pytest.param(None, DEVICE_CODE_POLL_SLEEP_SECONDS, id="missing"), + ], +) +async def test_astart_chatgpt_device_login_never_polls_faster_than_the_floor(interval: object, expected: int) -> None: + client: Final = _FakeOAuthClient( + body={"device_auth_id": "device-auth-example", "user_code": "ABCD-1234", "interval": interval} + ) + + login: Final = await astart_chatgpt_device_login(client) + + assert login.interval == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + pytest.param({"user_code": "ABCD-1234", "note": SECRET_IN_RESPONSE}, id="no_device_auth_id"), + pytest.param({"device_auth_id": "device-auth-example", "note": SECRET_IN_RESPONSE}, id="no_user_code"), + pytest.param({"device_auth_id": "", "user_code": "ABCD-1234", "note": SECRET_IN_RESPONSE}, id="empty_id"), + ], +) +async def test_astart_chatgpt_device_login_rejects_an_incomplete_response(body: dict[str, str]) -> None: + with pytest.raises(GetDeviceCodeError) as raised: + await astart_chatgpt_device_login(_FakeOAuthClient(body=body)) + + assert raised.value.status_code == 502 + assert SECRET_IN_RESPONSE not in "".join(traceback.format_exception(raised.value)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [403, 404]) +async def test_apoll_chatgpt_device_login_is_pending_while_the_user_has_not_approved(status_code: int) -> None: + client: Final = _FakeOAuthClient(status_code=status_code, body={"error": "authorization_pending"}) + + assert await apoll_chatgpt_device_login(client, "device-auth-example", "ABCD-1234") is None + assert client.posts == [ + _Post( + url=CHATGPT_DEVICE_TOKEN_URL, + data=None, + json={"device_auth_id": "device-auth-example", "user_code": "ABCD-1234"}, + ) + ] + + +@pytest.mark.asyncio +async def test_apoll_chatgpt_device_login_returns_the_code_once_the_user_approves() -> None: + client: Final = _FakeOAuthClient( + body={"authorization_code": "code-example", "code_challenge": "challenge", "code_verifier": "verifier-example"} + ) + + code: Final = await apoll_chatgpt_device_login(client, "device-auth-example", "ABCD-1234") + + assert code is not None + assert (code.authorization_code, code.code_verifier) == ("code-example", "verifier-example") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + pytest.param({"authorization_code": "code-example"}, id="no_verifier"), + pytest.param({"code_verifier": "verifier-example"}, id="no_code"), + pytest.param(b"not json", id="not_json"), + pytest.param(["code-example"], id="not_an_object"), + ], +) +async def test_apoll_chatgpt_device_login_treats_an_unusable_success_as_still_pending(body: object) -> None: + assert await apoll_chatgpt_device_login(_FakeOAuthClient(body=body), "device-auth-example", "ABCD-1234") is None + + +@pytest.mark.asyncio +async def test_aexchange_chatgpt_authorization_code_redeems_the_code_with_its_verifier() -> None: + client: Final = _FakeOAuthClient( + body={"access_token": "access-new", "refresh_token": "refresh-new", "id_token": "id-new"} + ) + + tokens: Final = await aexchange_chatgpt_authorization_code( + client, ChatGPTAuthorizationCode(authorization_code="code-example", code_verifier="verifier-example") + ) + + assert (tokens.access_token, tokens.refresh_token, tokens.id_token) == ("access-new", "refresh-new", "id-new") + assert client.posts == [ + _Post( + url=CHATGPT_OAUTH_TOKEN_URL, + data={ + "grant_type": "authorization_code", + "code": "code-example", + "redirect_uri": f"{CHATGPT_AUTH_BASE}/deviceauth/callback", + "client_id": CHATGPT_CLIENT_ID, + "code_verifier": "verifier-example", + }, + json=None, + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", ["access_token", "refresh_token", "id_token"]) +async def test_aexchange_chatgpt_authorization_code_rejects_a_response_missing_a_token(missing: str) -> None: + complete: Final = {"access_token": "a", "refresh_token": "r", "id_token": "i", "note": SECRET_IN_RESPONSE} + body: Final = {key: value for key, value in complete.items() if key != missing} + with pytest.raises(GetAccessTokenError) as raised: + await aexchange_chatgpt_authorization_code( + _FakeOAuthClient(body=body), ChatGPTAuthorizationCode(authorization_code="c", code_verifier="v") + ) + + assert raised.value.status_code == 502 + assert SECRET_IN_RESPONSE not in "".join(traceback.format_exception(raised.value)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "expected_refresh_token"), + [ + pytest.param({"access_token": "access-new", "id_token": "id-new"}, "refresh-old", id="omitted"), + pytest.param( + {"access_token": "access-new", "id_token": "id-new", "refresh_token": ""}, "refresh-old", id="empty" + ), + pytest.param( + {"access_token": "access-new", "id_token": "id-new", "refresh_token": "refresh-rotated"}, + "refresh-rotated", + id="rotated", + ), + ], +) +async def test_arefresh_chatgpt_tokens_keeps_the_old_refresh_token_unless_a_new_one_arrives( + body: dict[str, str], expected_refresh_token: str +) -> None: + client: Final = _FakeOAuthClient(body=body) + + tokens: Final = await arefresh_chatgpt_tokens(client, "refresh-old") + + assert tokens.access_token == "access-new" + assert tokens.refresh_token == expected_refresh_token + assert tokens.id_token == "id-new" + assert client.posts == [ + _Post( + url=CHATGPT_OAUTH_TOKEN_URL, + data=None, + json={ + "client_id": CHATGPT_CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": "refresh-old", + "scope": "openid profile email", + }, + ) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", ["access_token", "id_token"]) +async def test_arefresh_chatgpt_tokens_rejects_a_response_missing_a_token(missing: str) -> None: + body: Final = { + key: value + for key, value in {"access_token": "a", "id_token": "i", "note": SECRET_IN_RESPONSE}.items() + if key != missing + } + with pytest.raises(RefreshAccessTokenError) as raised: + await arefresh_chatgpt_tokens(_FakeOAuthClient(body=body), "refresh-old") + + assert raised.value.status_code == 502 + assert SECRET_IN_RESPONSE not in "".join(traceback.format_exception(raised.value)) + + +_SENT_SECRET: Final = "secret-the-caller-sent" +_OAUTH_CALLS: Final[tuple[tuple[str, type[ChatGPTAuthError], Callable[[_FakeOAuthClient], Awaitable[object]]], ...]] = ( + ("start", GetDeviceCodeError, astart_chatgpt_device_login), + ("poll", GetAccessTokenError, lambda client: apoll_chatgpt_device_login(client, _SENT_SECRET, _SENT_SECRET)), + ( + "exchange", + GetAccessTokenError, + lambda client: aexchange_chatgpt_authorization_code( + client, ChatGPTAuthorizationCode(authorization_code=_SENT_SECRET, code_verifier=_SENT_SECRET) + ), + ), + ("refresh", RefreshAccessTokenError, lambda client: arefresh_chatgpt_tokens(client, _SENT_SECRET)), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("client", "expected_status"), + [ + pytest.param(_FakeOAuthClient(status_code=400, body={"error": SECRET_IN_RESPONSE}), 400, id="http_400"), + pytest.param(_FakeOAuthClient(status_code=401, body={"error": SECRET_IN_RESPONSE}), 401, id="http_401"), + pytest.param(_FakeOAuthClient(status_code=500, body=SECRET_IN_RESPONSE.encode()), 500, id="http_500"), + pytest.param( + _FakeOAuthClient(error=httpx.ConnectError(f"{SECRET_IN_RESPONSE} {_SENT_SECRET}")), 502, id="transport" + ), + pytest.param( + _FakeOAuthClient(error=litellm.Timeout(f"{SECRET_IN_RESPONSE} {_SENT_SECRET}", "chatgpt", "chatgpt")), + 502, + id="timeout", + ), + ], +) +@pytest.mark.parametrize(("call", "error_type", "invoke"), _OAUTH_CALLS, ids=[name for name, _, _ in _OAUTH_CALLS]) +async def test_chatgpt_oauth_failures_never_echo_the_response_body_or_a_token( + client: _FakeOAuthClient, + expected_status: int, + call: str, + error_type: type[ChatGPTAuthError], + invoke: Callable[[_FakeOAuthClient], Awaitable[object]], +) -> None: + with pytest.raises(error_type) as raised: + await invoke(client) + + rendered: Final = "".join(traceback.format_exception(raised.value)) + assert raised.value.status_code == expected_status + assert SECRET_IN_RESPONSE not in rendered + assert _SENT_SECRET not in rendered + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call", "error_type", "invoke"), + [entry for entry in _OAUTH_CALLS if entry[0] != "poll"], + ids=[name for name, _, _ in _OAUTH_CALLS if name != "poll"], +) +async def test_chatgpt_oauth_success_that_is_not_a_json_object_never_echoes_the_body( + call: str, error_type: type[ChatGPTAuthError], invoke: Callable[[_FakeOAuthClient], Awaitable[object]] +) -> None: + with pytest.raises(error_type) as raised: + await invoke(_FakeOAuthClient(body=[SECRET_IN_RESPONSE])) + + assert raised.value.status_code == 502 + assert SECRET_IN_RESPONSE not in "".join(traceback.format_exception(raised.value)) + + +@pytest.mark.parametrize( + ("url", "body", "start"), + [ + pytest.param( + CHATGPT_OAUTH_TOKEN_URL, + {"access_token": SECRET_IN_RESPONSE, "refresh_token": SECRET_IN_RESPONSE}, + lambda authenticator: authenticator._refresh_tokens("refresh-old"), + id="refresh", + ), + pytest.param( + CHATGPT_OAUTH_TOKEN_URL, + {"access_token": SECRET_IN_RESPONSE, "id_token": SECRET_IN_RESPONSE}, + lambda authenticator: authenticator._exchange_code_for_tokens( + {"authorization_code": "code-example", "code_verifier": "verifier-example"} + ), + id="exchange", + ), + pytest.param( + CHATGPT_DEVICE_CODE_URL, + {"user_code": "ABCD-1234", "note": SECRET_IN_RESPONSE}, + lambda authenticator: authenticator._request_device_code(), + id="device_code", + ), + ], +) +def test_sync_authenticator_missing_field_errors_never_echo_the_response_body( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + url: str, + body: dict[str, str], + start: Callable[[Authenticator], object], +) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + with respx.mock as transport: + transport.post(url).mock(return_value=httpx.Response(200, json=body)) + with pytest.raises(ChatGPTAuthError) as raised: + start(Authenticator()) + + assert SECRET_IN_RESPONSE not in "".join(traceback.format_exception(raised.value)) + + +def test_chatgpt_login_values_derive_the_account_id_and_expiry_from_the_tokens() -> None: + access_token: Final = _make_jwt({"exp": 4102444800}) + id_token: Final = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-team-a"}}) + + values: Final = chatgpt_login_values( + ChatGPTTokens(access_token=access_token, refresh_token="refresh-example", id_token=id_token) + ) + + assert dict(values) == { + "chatgpt_access_token": access_token, + "chatgpt_refresh_token": "refresh-example", + "chatgpt_id_token": id_token, + "chatgpt_account_id": "acct-team-a", + "chatgpt_expires_at": "4102444800", + } + + +def test_chatgpt_login_values_omit_what_the_tokens_do_not_carry() -> None: + values: Final = chatgpt_login_values( + ChatGPTTokens(access_token=_make_jwt({}), refresh_token="refresh-example", id_token=_make_jwt({})) + ) + + assert set(values) == {"chatgpt_access_token", "chatgpt_refresh_token", "chatgpt_id_token"} + + +@pytest.mark.parametrize( + "value", + [ + pytest.param( + ChatGPTTokens(access_token=_SENT_SECRET, refresh_token=_SENT_SECRET, id_token=_SENT_SECRET), id="tokens" + ), + pytest.param(ChatGPTAuthorizationCode(authorization_code=_SENT_SECRET, code_verifier=_SENT_SECRET), id="code"), + ], +) +def test_chatgpt_oauth_values_never_print_their_secrets(value: object) -> None: + assert _SENT_SECRET not in repr(value) + assert _SENT_SECRET not in str(value) diff --git a/tests/unit/proxy/credential_endpoints/test_chatgpt_subscription_logins.py b/tests/unit/proxy/credential_endpoints/test_chatgpt_subscription_logins.py new file mode 100644 index 00000000000..d26a7fa63dc --- /dev/null +++ b/tests/unit/proxy/credential_endpoints/test_chatgpt_subscription_logins.py @@ -0,0 +1,776 @@ +import asyncio +import base64 +import json +import logging +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +import httpx +import pytest + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.llms.chatgpt.authenticator import TOKEN_EXPIRY_SKEW_SECONDS +from litellm.llms.chatgpt.common_utils import CHATGPT_OAUTH_TOKEN_URL, OPENAI_AUTH_CLAIM_KEY +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper +from litellm.proxy.credential_endpoints.chatgpt_subscription_logins import ( + ChatGPTLoginRefresher, + blocks_chatgpt_login, + resolve_chatgpt_subscription_logins, + save_chatgpt_login, +) +from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import InProcessRefreshCoordinator, OAuthToken +from litellm.types.utils import CredentialItem + +NOW: Final = 1_900_000_000.0 +FRESH_EXPIRES_AT: Final = str(int(NOW) + 3600) +STALE_EXPIRES_AT: Final = str(int(NOW) + TOKEN_EXPIRY_SKEW_SECONDS) +REFRESHED_EXPIRES_AT: Final = int(NOW) + 7200 + + +def _now() -> float: + return NOW + + +def _jwt(claims: Mapping[str, object]) -> str: + def segment(value: Mapping[str, object]) -> str: + return base64.urlsafe_b64encode(json.dumps(value).encode()).rstrip(b"=").decode() + + return f"{segment({'alg': 'RS256', 'typ': 'JWT'})}.{segment(claims)}.c2lnbmF0dXJl" + + +def _access_token(tag: str, exp: int = REFRESHED_EXPIRES_AT, account_id: str = "acct-1") -> str: + return _jwt({"exp": exp, "jti": tag, OPENAI_AUTH_CLAIM_KEY: {"chatgpt_account_id": account_id}}) + + +def _id_token(account_id: str = "acct-1") -> str: + return _jwt({"email": "owner@example.com", OPENAI_AUTH_CLAIM_KEY: {"chatgpt_account_id": account_id}}) + + +def _login_values( + access_token: str, refresh_token: str, expires_at: str | None, account_id: str = "acct-1" +) -> dict[str, str]: + return { + "chatgpt_access_token": access_token, + "chatgpt_refresh_token": refresh_token, + "chatgpt_id_token": _id_token(account_id), + "chatgpt_account_id": account_id, + **({"chatgpt_expires_at": expires_at} if expires_at is not None else {}), + } + + +def _login( + name: str, + values: Mapping[str, str], + team_id: str | None = None, + key_hash: str | None = None, + **extra_info: str, +) -> CredentialItem: + return CredentialItem( + credential_name=name, + credential_values=dict(values), + credential_info={ + "custom_llm_provider": "chatgpt", + **({"chatgpt_owner_team_id": team_id} if team_id is not None else {}), + **({"chatgpt_owner_key_hash": key_hash} if key_hash is not None else {}), + **extra_info, + }, + ) + + +def _stored(credential: CredentialItem) -> CredentialItem: + return credential.model_copy( + update={ + "credential_values": { + key: encrypt_value_helper(str(value)) for key, value in credential.credential_values.items() + } + } + ) + + +def _decrypted(values: Mapping[str, object]) -> dict[str, str | None]: + return {key: decrypt_value_helper(str(value), key) for key, value in values.items()} + + +def _row_after_write(row: CredentialItem, data: Mapping[str, object]) -> CredentialItem: + return row.model_copy( + update={ + field: json.loads(value) if field in ("credential_values", "credential_info") else value + for field, value in data.items() + if field in ("credential_values", "credential_info", "display_name") and isinstance(value, str) + } + ) + + +class _CredentialRows: + def __init__(self, *rows: CredentialItem, writes_fail: bool = False, reads_fail: bool = False) -> None: + self.rows: dict[str, CredentialItem] = {row.credential_name: row for row in rows} + self.updates: list[tuple[str, Mapping[str, object]]] = [] + self.compared: list[tuple[str, str]] = [] + self.reads: list[str] = [] + self.creates: list[Mapping[str, object]] = [] + self.writes_fail = writes_fail + self.reads_fail = reads_fail + + def put(self, row: CredentialItem) -> None: + self.rows[row.credential_name] = row + + def remove(self, credential_name: str) -> None: + self.rows.pop(credential_name) + + async def find_by_name(self, credential_name: str) -> CredentialItem | None: + self.reads.append(credential_name) + if self.reads_fail: + raise ConnectionError("the database is unreachable") + return self.rows.get(credential_name) + + async def update_if_value_unchanged( + self, credential_name: str, value_key: str, stored_value: str, data: Mapping[str, object] + ) -> bool: + self.updates.append((credential_name, data)) + self.compared.append((value_key, stored_value)) + if self.writes_fail: + raise ConnectionError("the database is unreachable") + current: Final = self.rows.get(credential_name) + if current is None or current.credential_values.get(value_key) != stored_value: + return False + self.rows[credential_name] = _row_after_write(current, data) + return True + + async def update_by_name(self, credential_name: str, data: Mapping[str, object]) -> None: + self.updates.append((credential_name, data)) + if self.writes_fail: + raise ConnectionError("the database is unreachable") + self.rows[credential_name] = _row_after_write(self.rows[credential_name], data) + + async def create(self, data: Mapping[str, object]) -> None: + self.creates.append(data) + name: Final = str(data["credential_name"]) + self.rows[name] = _row_after_write( + CredentialItem(credential_name=name, credential_values={}, credential_info={}), data + ) + + +class _TokenEndpoint: + def __init__( + self, + status: int = 200, + body: Mapping[str, object] | None = None, + during_post: Callable[[], None] | None = None, + ) -> None: + self.status = status + self.body = body + self.during_post = during_post + self.posts: list[tuple[str, Mapping[str, str] | None]] = [] + + async def post( + self, + url: str, + data: Mapping[str, str] | str | bytes | None = None, + json: Mapping[str, str] | None = None, + params: Mapping[str, str] | None = None, + headers: Mapping[str, str] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> httpx.Response: + self.posts.append((url, json)) + await asyncio.sleep(0) + if self.during_post is not None: + self.during_post() + request: Final = httpx.Request("POST", url) + response: Final = httpx.Response(self.status, json=self.body, request=request) + if response.is_error: + raise httpx.HTTPStatusError( + f"{self.status} from {url}: {response.text}", request=request, response=response + ) + return response + + +def _refreshed( + tag: str, refresh_token: str | None = "rt-issued-by-refresh", during_post: Callable[[], None] | None = None +) -> tuple[_TokenEndpoint, str]: + access_token: Final = _access_token(tag) + return ( + _TokenEndpoint( + body={ + "access_token": access_token, + "id_token": _id_token(), + **({"refresh_token": refresh_token} if refresh_token is not None else {}), + }, + during_post=during_post, + ), + access_token, + ) + + +class _LosingCoordinator: + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[OAuthToken | None]], + reread: Callable[[], Awaitable[OAuthToken | None]], + ) -> OAuthToken | None: + return await reread() + + +class _SharedCoordinatorWithoutLock: + def __init__(self) -> None: + self.runs = 0 + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[OAuthToken | None]], + reread: Callable[[], Awaitable[OAuthToken | None]], + ) -> OAuthToken | None: + self.runs += 1 + return await refresh() + + +def _refresher( + rows: _CredentialRows, + oauth_client: _TokenEndpoint, + coordinator: InProcessRefreshCoordinator | _LosingCoordinator | None = None, + shared_coordinator: _SharedCoordinatorWithoutLock | None = None, +) -> ChatGPTLoginRefresher: + return ChatGPTLoginRefresher( + rows=rows, + oauth_client=oauth_client, + coordinator=coordinator or InProcessRefreshCoordinator(), + shared_coordinator=shared_coordinator, + clock=_now, + ) + + +async def _resolve( + team_id: str | None, key_hash: str | None, refresher: ChatGPTLoginRefresher | None = None +) -> Mapping[str, str]: + return await resolve_chatgpt_subscription_logins( + team_id, key_hash, refresher_factory=lambda: refresher, clock=_now + ) + + +def _served(credential_name: str) -> Mapping[str, object]: + credential: Final = CredentialAccessor.find_credential(credential_name) + assert credential is not None + return credential.credential_values + + +@pytest.fixture(autouse=True) +def _isolated(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "chatgpt-subscription-logins-test-salt") + monkeypatch.setattr(litellm, "credential_list", []) + + +def _serve(monkeypatch, *credentials: CredentialItem) -> None: + monkeypatch.setattr(litellm, "credential_list", list(credentials)) + + +@pytest.fixture +def proxy_log(caplog): + caplog.set_level(logging.DEBUG, logger=verbose_proxy_logger.name) + verbose_proxy_logger.addHandler(caplog.handler) + yield caplog + verbose_proxy_logger.removeHandler(caplog.handler) + + +def _logged_text(caplog: pytest.LogCaptureFixture) -> str: + return caplog.text + "\n".join(record.getMessage() for record in caplog.records) + + +async def test_a_fresh_login_is_served_from_memory_without_the_database_or_the_token_endpoint(monkeypatch): + access_token: Final = _access_token("fresh") + _serve(monkeypatch, _login("login-a", _login_values(access_token, "rt-a", FRESH_EXPIRES_AT), team_id="team-a")) + rows: Final = _CredentialRows() + oauth: Final = _TokenEndpoint(status=500) + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-a": access_token} + assert oauth.posts == [] + assert (rows.updates, rows.creates) == ([], []) + + +async def test_a_stale_login_another_replica_already_refreshed_is_served_from_the_row(monkeypatch): + memory_values: Final = _login_values(_access_token("old", exp=int(NOW)), "rt-old", STALE_EXPIRES_AT) + row_values: Final = _login_values(_access_token("refreshed-elsewhere"), "rt-current", FRESH_EXPIRES_AT) + _serve(monkeypatch, _login("login-a", memory_values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", row_values, team_id="team-a"))) + oauth: Final = _TokenEndpoint(status=500) + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-a": row_values["chatgpt_access_token"]} + assert oauth.posts == [] + assert rows.updates == [] + assert _served("login-a") == row_values + + +async def test_a_stale_login_redeems_the_rows_refresh_token_and_saves_the_result_encrypted(monkeypatch): + memory_values: Final = _login_values(_access_token("old"), "rt-memory-older", STALE_EXPIRES_AT) + row_values: Final = _login_values(_access_token("old-row"), "rt-row-current", STALE_EXPIRES_AT) + _serve(monkeypatch, _login("login-a", memory_values, team_id="team-a")) + stored_row: Final = _stored( + _login("login-a", {**row_values, "api_base": "https://chatgpt.com/backend-api/codex"}, team_id="team-a") + ) + rows: Final = _CredentialRows(stored_row) + oauth, new_access_token = _refreshed("new", refresh_token="rt-rotated") + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-a": new_access_token} + assert [(url, (body or {}).get("refresh_token")) for url, body in oauth.posts] == [ + (CHATGPT_OAUTH_TOKEN_URL, "rt-row-current") + ] + assert [name for name, _ in rows.updates] == ["login-a"] + written: Final = rows.updates[0][1]["credential_values"] + assert isinstance(written, str) + saved: Final = json.loads(written) + expected_login: Final = { + "chatgpt_access_token": new_access_token, + "chatgpt_refresh_token": "rt-rotated", + "chatgpt_id_token": _id_token(), + "chatgpt_account_id": "acct-1", + "chatgpt_expires_at": str(REFRESHED_EXPIRES_AT), + } + for key, plaintext in expected_login.items(): + assert saved[key] != plaintext + assert decrypt_value_helper(saved[key], key) == plaintext + for secret in (new_access_token, "rt-rotated", _id_token()): + assert secret not in written + assert saved["api_base"] == stored_row.credential_values["api_base"] + assert _served("login-a")["chatgpt_access_token"] == new_access_token + assert _served("login-a")["chatgpt_refresh_token"] == "rt-rotated" + + +async def test_a_refresh_response_without_a_refresh_token_keeps_the_old_one(monkeypatch): + values: Final = _login_values(_access_token("old"), "rt-still-valid", STALE_EXPIRES_AT) + _serve(monkeypatch, _login("login-a", values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", values, team_id="team-a"))) + oauth, new_access_token = _refreshed("new", refresh_token=None) + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-a": new_access_token} + assert _decrypted(rows.rows["login-a"].credential_values)["chatgpt_refresh_token"] == "rt-still-valid" + assert _served("login-a")["chatgpt_refresh_token"] == "rt-still-valid" + + +@pytest.mark.parametrize("expires_at", ["nan", "inf", "1e400", "soon"]) +async def test_a_login_whose_expiry_is_not_a_finite_time_is_refreshed(monkeypatch, expires_at): + values: Final = _login_values(_access_token("old"), "rt-a", expires_at) + _serve(monkeypatch, _login("login-a", values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", values, team_id="team-a"))) + oauth, new_access_token = _refreshed("new") + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-a": new_access_token} + assert len(oauth.posts) == 1 + + +async def test_concurrent_requests_for_one_stale_login_redeem_it_once(monkeypatch): + values: Final = _login_values(_access_token("old"), "rt-single-use", STALE_EXPIRES_AT) + _serve(monkeypatch, _login("login-a", values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", values, team_id="team-a"))) + oauth, new_access_token = _refreshed("new") + refresher: Final = _refresher(rows, oauth) + + resolved: Final = await asyncio.gather(*(_resolve("team-a", None, refresher) for _ in range(8))) + + assert resolved == [{"login-a": new_access_token}] * 8 + assert len(oauth.posts) == 1 + assert len(rows.updates) == 1 + + +async def test_a_refresh_lock_loser_serves_the_token_the_winner_saved(monkeypatch): + memory_values: Final = _login_values(_access_token("before-winner"), "rt-before-winner", STALE_EXPIRES_AT) + row_values: Final = _login_values(_access_token("winner"), "rt-winner", FRESH_EXPIRES_AT) + held: Final = _login("login-a", memory_values, team_id="team-a") + _serve(monkeypatch, held) + rows: Final = _CredentialRows(_stored(_login("login-a", row_values, team_id="team-a"))) + oauth: Final = _TokenEndpoint(status=500) + + token: Final = await _refresher(rows, oauth, _LosingCoordinator()).access_token(held) + + assert token == row_values["chatgpt_access_token"] + assert oauth.posts == [] + assert _served("login-a") == row_values + + +async def test_a_refresh_lock_loser_serves_nothing_while_the_row_is_still_stale(): + row_values: Final = _login_values(_access_token("stale"), "rt-stale", STALE_EXPIRES_AT) + held: Final = _login("login-a", row_values, team_id="team-a") + rows: Final = _CredentialRows(_stored(held)) + oauth: Final = _TokenEndpoint(status=500) + + token: Final = await _refresher(rows, oauth, _LosingCoordinator()).access_token(held) + + assert token is None + assert oauth.posts == [] + + +async def test_a_failed_redeem_serves_no_login_and_logs_no_secret(monkeypatch, proxy_log): + stale_access_token: Final = _access_token("stale-access") + other_access_token: Final = _access_token("other") + memory_values: Final = _login_values(stale_access_token, "rt-memory-copy", STALE_EXPIRES_AT) + row_values: Final = _login_values(stale_access_token, "rt-row-revoked", STALE_EXPIRES_AT) + _serve( + monkeypatch, + _login("login-a", memory_values, team_id="team-a"), + _login("login-b", _login_values(other_access_token, "rt-b", FRESH_EXPIRES_AT), team_id="team-a"), + ) + rows: Final = _CredentialRows(_stored(_login("login-a", row_values, team_id="team-a"))) + leaked_in_body: Final = _jwt({"sub": "token-echoed-in-error-body"}) + oauth: Final = _TokenEndpoint( + status=400, + body={"error": "invalid_grant", "error_description": f"refresh token already used: {leaked_in_body}"}, + ) + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-b": other_access_token} + assert len(oauth.posts) == 1 + assert rows.updates == [] + logged: Final = _logged_text(proxy_log) + assert "login-a" in logged + for secret in (stale_access_token, "rt-memory-copy", "rt-row-revoked", leaked_in_body, "already used"): + assert secret not in logged + assert secret not in repr(resolved) + + +async def test_a_failed_save_still_serves_the_refreshed_token(monkeypatch, proxy_log): + values: Final = _login_values(_access_token("old"), "rt-a", STALE_EXPIRES_AT) + _serve(monkeypatch, _login("login-a", values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", values, team_id="team-a")), writes_fail=True) + oauth, new_access_token = _refreshed("new", refresh_token="rt-rotated") + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-a": new_access_token} + assert len(oauth.posts) == 1 + assert len(rows.updates) == 1 + assert _served("login-a")["chatgpt_access_token"] == new_access_token + logged: Final = _logged_text(proxy_log) + assert "login-a" in logged + for secret in (new_access_token, "rt-rotated", "rt-a"): + assert secret not in logged + + +async def test_concurrent_calls_redeem_once_when_the_shared_lock_refreshes_without_locking(monkeypatch): + values: Final = _login_values(_access_token("old"), "rt-single-use", STALE_EXPIRES_AT) + held: Final = _login("login-a", values, team_id="team-a") + _serve(monkeypatch, held) + rows: Final = _CredentialRows(_stored(held)) + oauth, new_access_token = _refreshed("new") + shared: Final = _SharedCoordinatorWithoutLock() + refresher: Final = _refresher(rows, oauth, shared_coordinator=shared) + + tokens: Final = await asyncio.gather(*(refresher.access_token(held) for _ in range(5))) + + assert tokens == [new_access_token] * 5 + assert len(oauth.posts) == 1 + assert shared.runs == 1 + + +async def test_a_refresh_that_loses_its_write_to_a_new_sign_in_serves_the_new_sign_in(monkeypatch): + old_values: Final = _login_values(_access_token("old"), "rt-old", STALE_EXPIRES_AT) + signed_in_values: Final = _login_values(_access_token("signed-in-again"), "rt-signed-in-again", FRESH_EXPIRES_AT) + held: Final = _login("login-a", old_values, team_id="team-a") + _serve(monkeypatch, held) + stored_row: Final = _stored(held) + rows: Final = _CredentialRows(stored_row) + oauth, redeemed_access_token = _refreshed( + "redeemed", + refresh_token="rt-redeemed", + during_post=lambda: rows.put(_stored(_login("login-a", signed_in_values, team_id="team-a"))), + ) + + token: Final = await _refresher(rows, oauth).access_token(held) + + assert token == signed_in_values["chatgpt_access_token"] + assert rows.compared == [("chatgpt_refresh_token", stored_row.credential_values["chatgpt_refresh_token"])] + assert _decrypted(rows.rows["login-a"].credential_values) == signed_in_values + assert _served("login-a") == signed_in_values + assert redeemed_access_token != token + + +async def test_a_refresh_that_loses_its_write_to_another_owners_sign_in_serves_nothing(monkeypatch): + old_values: Final = _login_values(_access_token("old"), "rt-old", STALE_EXPIRES_AT) + team_b_values: Final = _login_values(_access_token("team-b-seat"), "rt-team-b", FRESH_EXPIRES_AT) + held: Final = _login("login-a", old_values, team_id="team-a") + _serve(monkeypatch, held) + rows: Final = _CredentialRows(_stored(held)) + oauth, _ = _refreshed( + "redeemed", + refresh_token="rt-redeemed", + during_post=lambda: rows.put(_stored(_login("login-a", team_b_values, team_id="team-b"))), + ) + + token: Final = await _refresher(rows, oauth).access_token(held) + + assert token is None + assert _decrypted(rows.rows["login-a"].credential_values) == team_b_values + assert _served("login-a") == old_values + + +@pytest.mark.parametrize("row_deleted", [False, True], ids=["memory-only", "memory-and-row"]) +async def test_a_login_deleted_while_its_refresh_is_in_flight_stays_deleted(monkeypatch, row_deleted): + values: Final = _login_values(_access_token("old"), "rt-a", STALE_EXPIRES_AT) + held: Final = _login("login-a", values, team_id="team-a") + _serve(monkeypatch, held) + rows: Final = _CredentialRows(_stored(held)) + + def delete() -> None: + _serve(monkeypatch) + if row_deleted: + rows.remove("login-a") + + oauth, _ = _refreshed("new", during_post=delete) + + await _refresher(rows, oauth).access_token(held) + + assert len(oauth.posts) == 1 + assert CredentialAccessor.find_credential("login-a") is None + + +async def test_a_login_deleted_before_its_refresh_starts_is_not_added_back(monkeypatch): + memory_values: Final = _login_values(_access_token("old"), "rt-old", STALE_EXPIRES_AT) + row_values: Final = _login_values(_access_token("row"), "rt-row", FRESH_EXPIRES_AT) + _serve(monkeypatch, _login("login-a", memory_values, team_id="team-a")) + refresher: Final = _refresher( + _CredentialRows(_stored(_login("login-a", row_values, team_id="team-a"))), _TokenEndpoint(status=500) + ) + + def delete_then_refresh() -> ChatGPTLoginRefresher: + _serve(monkeypatch) + return refresher + + await resolve_chatgpt_subscription_logins("team-a", None, refresher_factory=delete_then_refresh, clock=_now) + + assert CredentialAccessor.find_credential("login-a") is None + + +async def test_a_database_outage_serves_no_stale_login_and_keeps_the_others(monkeypatch, proxy_log): + other_access_token: Final = _access_token("other") + stale_values: Final = _login_values(_access_token("stale"), "rt-stale", STALE_EXPIRES_AT) + _serve( + monkeypatch, + _login("login-a", stale_values, team_id="team-a"), + _login("login-b", _login_values(other_access_token, "rt-b", FRESH_EXPIRES_AT), team_id="team-a"), + ) + rows: Final = _CredentialRows(_stored(_login("login-a", stale_values, team_id="team-a")), reads_fail=True) + oauth: Final = _TokenEndpoint(status=500) + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert resolved == {"login-b": other_access_token} + assert oauth.posts == [] + logged: Final = _logged_text(proxy_log) + assert "login-a" in logged + assert "rt-stale" not in logged + + +@pytest.mark.parametrize("owner_key", ["chatgpt_owner_team_id", "chatgpt_owner_key_hash"], ids=["team", "key"]) +@pytest.mark.parametrize("real_owner_first", [True, False], ids=["real-owner-first", "stale-owner-first"]) +async def test_callers_holding_different_owners_never_share_one_refresh(owner_key, real_owner_first): + row_values: Final = _login_values(_access_token("owner-b-seat"), "rt-owner-b", FRESH_EXPIRES_AT) + stale_values: Final = _login_values(_access_token("old"), "rt-old", STALE_EXPIRES_AT) + rows: Final = _CredentialRows(_stored(_login("login-a", row_values, **{owner_key: "owner-b"}))) + refresher: Final = _refresher(rows, _TokenEndpoint(status=500)) + held_by_owner_b: Final = _login("login-a", stale_values, **{owner_key: "owner-b"}) + held_by_owner_a: Final = _login("login-a", stale_values, **{owner_key: "owner-a"}) + callers: Final = (held_by_owner_b, held_by_owner_a) if real_owner_first else (held_by_owner_a, held_by_owner_b) + + tokens: Final = await asyncio.gather(*(refresher.access_token(caller) for caller in callers)) + + assert {caller.credential_info[owner_key]: token for caller, token in zip(callers, tokens)} == { + "owner-b": row_values["chatgpt_access_token"], + "owner-a": None, + } + + +def _ownership_fixture() -> tuple[CredentialItem, ...]: + def fresh(tag: str) -> dict[str, str]: + return _login_values(_access_token(tag), f"rt-{tag}", FRESH_EXPIRES_AT) + + return ( + _login("login-team-a", fresh("team-a"), team_id="team-a"), + _login("login-team-b", fresh("team-b"), team_id="team-b"), + _login("login-key-k", fresh("key-k"), key_hash="hash-k"), + _login("login-both-owners", fresh("both"), team_id="team-a", key_hash="hash-k"), + _login("login-no-owner", fresh("none")), + CredentialItem( + credential_name="openai-team-a", + credential_values={"api_key": "sk-not-a-subscription"}, + credential_info={"custom_llm_provider": "openai", "chatgpt_owner_team_id": "team-a"}, + ), + ) + + +@pytest.mark.parametrize( + ("team_id", "key_hash", "expected"), + [ + ("team-a", "hash-unbound", {"login-team-a"}), + (None, "hash-k", {"login-key-k"}), + ("team-b", "hash-k", {"login-team-b", "login-key-k"}), + ("team-a", "hash-k", {"login-team-a", "login-key-k"}), + (None, None, set()), + ("team-unbound", None, set()), + ], + ids=["team-a", "key-k", "team-b-and-key-k", "team-a-and-key-k", "no-owner", "unbound-team"], +) +async def test_a_caller_gets_only_the_logins_its_team_or_key_owns(monkeypatch, team_id, key_hash, expected): + credentials: Final = _ownership_fixture() + _serve(monkeypatch, *credentials) + tokens: Final = { + credential.credential_name: credential.credential_values.get("chatgpt_access_token") + for credential in credentials + } + + rows: Final = _CredentialRows() + oauth: Final = _TokenEndpoint(status=500) + + resolved: Final = await _resolve(team_id, key_hash, _refresher(rows, oauth)) + + assert resolved == {name: tokens[name] for name in expected} + assert (rows.reads, oauth.posts) == ([], []) + + +@pytest.mark.parametrize( + "expires_at", + [None, "not-a-number", ""], + ids=["missing", "unparsable", "blank"], +) +async def test_a_login_with_an_unknown_expiry_is_refreshed(monkeypatch, expires_at): + values: Final = _login_values(_access_token("unknown-expiry"), "rt-a", expires_at) + _serve(monkeypatch, _login("login-a", values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", values, team_id="team-a"))) + oauth, new_access_token = _refreshed("new") + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth)) + + assert len(oauth.posts) == 1 + assert resolved == {"login-a": new_access_token} + + +@pytest.mark.parametrize( + ("existing", "blocked"), + [ + (None, False), + ( + CredentialItem( + credential_name="taken", + credential_values={"api_key": "sk-openai"}, + credential_info={"custom_llm_provider": "openai"}, + ), + True, + ), + ( + CredentialItem(credential_name="taken", credential_values={"api_key": "sk-unknown"}, credential_info={}), + True, + ), + ( + CredentialItem( + credential_name="taken", credential_values={}, credential_info={"custom_llm_provider": "openai"} + ), + False, + ), + ( + CredentialItem( + credential_name="taken", + credential_values={"api_base": "https://chatgpt.com/backend-api/codex"}, + credential_info={"custom_llm_provider": "chatgpt"}, + ), + False, + ), + ( + CredentialItem( + credential_name="taken", + credential_values={"chatgpt_access_token": "at", "chatgpt_refresh_token": "rt"}, + credential_info={}, + ), + False, + ), + ], + ids=["free", "other-provider", "no-provider", "no-values", "chatgpt-provider", "existing-login"], +) +def test_a_sign_in_is_blocked_only_by_another_providers_credential(existing, blocked): + assert blocks_chatgpt_login(existing) is blocked + + +async def test_saving_a_new_login_creates_an_encrypted_credential_its_owner_can_use(): + rows: Final = _CredentialRows() + access_token: Final = _access_token("signed-in") + values: Final = _login_values(access_token, "rt-signed-in", FRESH_EXPIRES_AT) + + await save_chatgpt_login(rows, "login-new", "Team A ChatGPT", values, ("team-a", None), "admin-user") + + assert rows.updates == [] + assert len(rows.creates) == 1 + created: Final = rows.creates[0] + written_values: Final = created["credential_values"] + assert isinstance(written_values, str) + saved: Final = json.loads(written_values) + assert saved.keys() == values.keys() + for key, plaintext in values.items(): + assert saved[key] != plaintext + assert decrypt_value_helper(saved[key], key) == plaintext + for secret in (access_token, "rt-signed-in"): + assert secret not in written_values + assert json.loads(str(created["credential_info"])) == { + "custom_llm_provider": "chatgpt", + "chatgpt_owner_team_id": "team-a", + } + assert (created["credential_name"], created["display_name"], created["created_by"]) == ( + "login-new", + "Team A ChatGPT", + "admin-user", + ) + assert await _resolve("team-a", None) == {"login-new": access_token} + assert await _resolve("team-b", "hash-k") == {} + + +async def test_saving_again_on_the_same_name_replaces_the_owner_and_keeps_the_rest(monkeypatch): + old_values: Final = _login_values(_access_token("old"), "rt-old", FRESH_EXPIRES_AT) + existing: Final = _login("login-a", old_values, team_id="team-a", note="billing seat 3").model_copy( + update={"display_name": "Old label"} + ) + _serve(monkeypatch, existing) + rows: Final = _CredentialRows(_stored(existing)) + new_access_token: Final = _access_token("rotated") + new_values: Final = _login_values(new_access_token, "rt-rotated", FRESH_EXPIRES_AT) + + await save_chatgpt_login(rows, "login-a", None, new_values, (None, "hash-k"), "admin-user") + + assert rows.creates == [] + assert [name for name, _ in rows.updates] == ["login-a"] + written: Final = rows.updates[0][1] + assert json.loads(str(written["credential_info"])) == { + "note": "billing seat 3", + "custom_llm_provider": "chatgpt", + "chatgpt_owner_key_hash": "hash-k", + } + assert written["display_name"] == "Old label" + assert _decrypted(json.loads(str(written["credential_values"]))) == new_values + assert await _resolve("team-a", None) == {} + assert await _resolve(None, "hash-k") == {"login-a": new_access_token} + + +@pytest.mark.parametrize("coordinator", [InProcessRefreshCoordinator, _LosingCoordinator], ids=["refresh", "reread"]) +async def test_a_login_rebound_to_another_owner_is_not_served_to_the_old_owner(monkeypatch, coordinator): + memory_values: Final = _login_values(_access_token("team-a-seat"), "rt-team-a", STALE_EXPIRES_AT) + row_values: Final = _login_values( + _access_token("team-b-seat", account_id="acct-team-b"), "rt-team-b", FRESH_EXPIRES_AT, "acct-team-b" + ) + _serve(monkeypatch, _login("login-a", memory_values, team_id="team-a")) + rows: Final = _CredentialRows(_stored(_login("login-a", row_values, team_id="team-b"))) + oauth: Final = _TokenEndpoint(status=500) + + resolved: Final = await _resolve("team-a", None, _refresher(rows, oauth, coordinator())) + + assert resolved == {} + assert oauth.posts == [] + assert _served("login-a") == memory_values diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 9dd36a535e4..c7809ddf9b9 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -1,5 +1,6 @@ """Tests for the credential management endpoints.""" +import base64 import json from contextlib import contextmanager from functools import reduce @@ -18,12 +19,27 @@ from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import get_masked_values +from litellm.llms.chatgpt.common_utils import ( + CHATGPT_DEVICE_CODE_URL, + CHATGPT_DEVICE_TOKEN_URL, + CHATGPT_DEVICE_VERIFY_URL, + CHATGPT_OAUTH_TOKEN_URL, + OPENAI_AUTH_CLAIM_KEY, +) from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper -from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + V2_GCM_PREFIX, + decrypt_value_helper, + encrypt_value_helper, +) +from litellm.proxy.credential_endpoints.endpoints import ( + _decode_flow_handle, + get_chatgpt_oauth_client, + get_llm_router, +) from litellm.proxy.proxy_server import app from litellm.types.utils import CredentialItem @@ -3275,3 +3291,315 @@ class TestCredentialDisplayName: assert (by_name["display_name"], by_name["source"]) == ("Prod OpenAI", "db") assert config_by_name["source"] == "config" assert listed["openai-prod"]["credential_values"]["api_key"] != "sk-old" + + +_CHATGPT_DEVICE_AUTH_ID: Final = "deviceauth_7Qm2x9Lp4Rt8Vw3Kz" +_CHATGPT_USER_CODE: Final = "WXYZ-1234" +_CHATGPT_TOKEN_EXPIRES_AT: Final = 4102444800 + + +def _chatgpt_jwt(claims: dict) -> str: + def segment(value: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(value).encode()).rstrip(b"=").decode() + + return f"{segment({'alg': 'RS256', 'typ': 'JWT'})}.{segment(claims)}.c2lnbmF0dXJl" + + +def _chatgpt_tokens(tag: str, account_id: str) -> dict[str, str]: + auth_claims: Final = {OPENAI_AUTH_CLAIM_KEY: {"chatgpt_account_id": account_id}} + return { + "access_token": _chatgpt_jwt({"exp": _CHATGPT_TOKEN_EXPIRES_AT, "jti": tag, **auth_claims}), + "refresh_token": f"rt_{tag}_9f8e7d6c5b4a", + "id_token": _chatgpt_jwt({"email": f"{tag}@example.com", **auth_claims}), + } + + +class _ChatGPTAuthServer: + """Stands in for the async HTTP client the sign-in routes reach auth.openai.com through. It answers + each URL from ``answers`` and raises ``httpx.HTTPStatusError`` on a non-2xx answer, as litellm's + ``AsyncHTTPHandler`` does, with the response body in the error text.""" + + def __init__(self) -> None: + self.answers: dict[str, tuple[int, dict]] = { + CHATGPT_DEVICE_CODE_URL: ( + 200, + {"device_auth_id": _CHATGPT_DEVICE_AUTH_ID, "user_code": _CHATGPT_USER_CODE, "interval": "7"}, + ), + CHATGPT_DEVICE_TOKEN_URL: (403, {"error": "deviceauth_authorization_pending"}), + } + self.requests: list[tuple[str, object]] = [] + + def approve(self, tag: str, account_id: str) -> dict[str, str]: + tokens: Final = _chatgpt_tokens(tag, account_id) + self.answers[CHATGPT_DEVICE_TOKEN_URL] = ( + 200, + {"authorization_code": f"code-{tag}", "code_verifier": f"verifier-{tag}"}, + ) + self.answers[CHATGPT_OAUTH_TOKEN_URL] = (200, tokens) + return tokens + + async def post(self, url, data=None, json=None, params=None, headers=None, timeout=None) -> httpx.Response: + self.requests.append((url, json if json is not None else data)) + status, body = self.answers[url] + request: Final = httpx.Request("POST", url) + response: Final = httpx.Response(status, json=body, request=request) + if response.is_error: + raise httpx.HTTPStatusError(f"{status} from {url}: {response.text}", request=request, response=response) + return response + + +class _RowsById: + def __init__(self, id_field: str, *ids: str) -> None: + self._id_field = id_field + self._ids = frozenset(ids) + + async def find_unique(self, where: dict) -> dict | None: + row_id: Final = where[self._id_field] + return {self._id_field: row_id} if row_id in self._ids else None + + +class _CredentialsTable: + """The Prisma credentials table: rows as Prisma hands them back, with the JSON columns parsed.""" + + def __init__(self) -> None: + self.rows: dict[str, dict] = {} + self.writes: list[str] = [] + + @staticmethod + def _as_read(data: dict) -> dict: + return { + key: json.loads(value) if key in ("credential_values", "credential_info") else value + for key, value in data.items() + } + + async def find_unique(self, where: dict) -> dict | None: + return self.rows.get(where["credential_name"]) + + async def create(self, data: dict) -> dict: + self.writes.append("create") + self.rows[data["credential_name"]] = self._as_read(data) + return self.rows[data["credential_name"]] + + async def update(self, where: dict, data: dict) -> dict: + self.writes.append("update") + name: Final = where["credential_name"] + self.rows[name] = {**self.rows[name], **self._as_read(data)} + return self.rows[name] + + +@pytest.fixture +def chatgpt_sign_in(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "chatgpt-sign-in-endpoints-test-salt") + monkeypatch.setattr(litellm, "credential_list", []) + credentials: Final = _CredentialsTable() + auth_server: Final = _ChatGPTAuthServer() + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + SimpleNamespace( + db=SimpleNamespace( + litellm_credentialstable=credentials, + litellm_teamtable=_RowsById("team_id", "team-a", "team-b"), + litellm_verificationtoken=_RowsById("token", "hash-k"), + ) + ), + ) + app.dependency_overrides[get_chatgpt_oauth_client] = lambda: auth_server + yield SimpleNamespace(credentials=credentials, auth_server=auth_server) + app.dependency_overrides.pop(get_chatgpt_oauth_client, None) + + +def _start_chatgpt_login(name: str, body: dict, auth=_as_admin): + return _call_as("POST", f"/credentials/{name}/chatgpt_login/start", body, auth) + + +def _poll_chatgpt_login(name: str, flow_handle: str, auth=_as_admin): + return _call_as("POST", f"/credentials/{name}/chatgpt_login/poll", {"flow_handle": flow_handle}, auth) + + +def _started_chatgpt_login(name: str, owner: dict) -> str: + response: Final = _start_chatgpt_login(name, owner) + assert response.status_code == 200, response.text + return response.json()["flow_handle"] + + +def _stored_login(row: dict) -> dict[str, str | None]: + return {key: decrypt_value_helper(value, key) for key, value in row["credential_values"].items()} + + +def _expected_login(tokens: dict[str, str], account_id: str) -> dict[str, str]: + return { + "chatgpt_access_token": tokens["access_token"], + "chatgpt_refresh_token": tokens["refresh_token"], + "chatgpt_id_token": tokens["id_token"], + "chatgpt_account_id": account_id, + "chatgpt_expires_at": str(_CHATGPT_TOKEN_EXPIRES_AT), + } + + +def test_chatgpt_sign_in_start_is_proxy_admin_only(chatgpt_sign_in): + response = _start_chatgpt_login("login-a", {"owner_team_id": "team-a"}, auth=_as_non_admin) + + assert response.status_code == 403, response.text + assert chatgpt_sign_in.auth_server.requests == [] + + +def test_chatgpt_sign_in_poll_is_proxy_admin_only(chatgpt_sign_in): + flow_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a"}) + chatgpt_sign_in.auth_server.approve("first", "acct-a") + + response = _poll_chatgpt_login("login-a", flow_handle, auth=_as_non_admin) + + assert response.status_code == 403, response.text + assert [url for url, _ in chatgpt_sign_in.auth_server.requests] == [CHATGPT_DEVICE_CODE_URL] + assert chatgpt_sign_in.credentials.rows == {} + + +@pytest.mark.parametrize( + "owner", + [{"owner_team_id": "team-a", "owner_key_hash": "hash-k"}, {}, {"owner_team_id": "", "owner_key_hash": ""}], + ids=["both", "neither", "blank"], +) +def test_chatgpt_sign_in_needs_exactly_one_owner(chatgpt_sign_in, owner): + response = _start_chatgpt_login("login-a", owner) + + assert response.status_code == 400, response.text + assert chatgpt_sign_in.auth_server.requests == [] + + +@pytest.mark.parametrize( + "owner", + [{"owner_team_id": "team-deleted"}, {"owner_key_hash": "hash-of-a-deleted-key"}], + ids=["team", "key"], +) +def test_chatgpt_sign_in_needs_an_owner_that_exists(chatgpt_sign_in, owner): + response = _start_chatgpt_login("login-a", owner) + + assert response.status_code == 400, response.text + assert chatgpt_sign_in.auth_server.requests == [] + + +@pytest.mark.parametrize("owner", [{"owner_team_id": "team-a"}, {"owner_key_hash": "hash-k"}], ids=["team", "key"]) +def test_chatgpt_sign_in_start_returns_the_device_code_without_the_device_auth_id(chatgpt_sign_in, owner): + response = _start_chatgpt_login("login-a", owner) + + assert response.status_code == 200, response.text + body = response.json() + assert (body["user_code"], body["verification_uri"], body["interval"]) == ( + _CHATGPT_USER_CODE, + CHATGPT_DEVICE_VERIFY_URL, + 7, + ) + flow_handle = body["flow_handle"] + assert isinstance(flow_handle, str) and flow_handle + assert _CHATGPT_DEVICE_AUTH_ID not in response.text + assert _CHATGPT_DEVICE_AUTH_ID not in flow_handle + encoded = flow_handle.removeprefix(V2_GCM_PREFIX) + assert _CHATGPT_DEVICE_AUTH_ID.encode() not in base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)) + + +def test_chatgpt_sign_in_poll_is_pending_until_the_user_approves(chatgpt_sign_in): + flow_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a"}) + + response = _poll_chatgpt_login("login-a", flow_handle) + + assert response.status_code == 200, response.text + assert response.json()["status"] == "pending" + assert chatgpt_sign_in.credentials.rows == {} + + +def test_chatgpt_sign_in_saves_an_encrypted_login_and_answers_without_a_token(chatgpt_sign_in): + flow_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a", "display_name": "Team A seat"}) + tokens = chatgpt_sign_in.auth_server.approve("first", "acct-team-a") + + response = _poll_chatgpt_login("login-a", flow_handle) + + assert response.status_code == 200, response.text + assert response.json() == {"status": "connected", "interval": None, "account_id": "acct-team-a"} + for secret in tokens.values(): + assert secret not in response.text + assert chatgpt_sign_in.auth_server.requests[1] == ( + CHATGPT_DEVICE_TOKEN_URL, + {"device_auth_id": _CHATGPT_DEVICE_AUTH_ID, "user_code": _CHATGPT_USER_CODE}, + ) + row = chatgpt_sign_in.credentials.rows["login-a"] + assert _stored_login(row) == _expected_login(tokens, "acct-team-a") + for secret in tokens.values(): + assert secret not in json.dumps(row) + assert row["credential_info"] == {"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-a"} + assert row["display_name"] == "Team A seat" + assert CredentialAccessor.get_credential_values("login-a") == _expected_login(tokens, "acct-team-a") + + +def test_signing_in_again_on_the_same_name_updates_that_credential(chatgpt_sign_in): + first_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a"}) + chatgpt_sign_in.auth_server.approve("first", "acct-first") + assert _poll_chatgpt_login("login-a", first_handle).json()["status"] == "connected" + second_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-b"}) + second_tokens = chatgpt_sign_in.auth_server.approve("second", "acct-second") + + response = _poll_chatgpt_login("login-a", second_handle) + + assert response.json() == {"status": "connected", "interval": None, "account_id": "acct-second"} + assert chatgpt_sign_in.credentials.writes == ["create", "update"] + assert list(chatgpt_sign_in.credentials.rows) == ["login-a"] + row = chatgpt_sign_in.credentials.rows["login-a"] + assert _stored_login(row) == _expected_login(second_tokens, "acct-second") + assert row["credential_info"] == {"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-b"} + + +def test_a_flow_handle_minted_for_another_name_is_rejected(chatgpt_sign_in): + flow_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a"}) + chatgpt_sign_in.auth_server.approve("first", "acct-a") + + response = _poll_chatgpt_login("login-b", flow_handle) + + assert response.status_code == 400, response.text + assert [url for url, _ in chatgpt_sign_in.auth_server.requests] == [CHATGPT_DEVICE_CODE_URL] + assert chatgpt_sign_in.credentials.rows == {} + + +def _another_providers_row(name: str) -> dict: + return { + "credential_name": name, + "credential_values": {"api_key": encrypt_value_helper("sk-openai-prod")}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def test_chatgpt_sign_in_start_refuses_a_name_held_by_another_provider(chatgpt_sign_in): + existing = _another_providers_row("openai-prod") + chatgpt_sign_in.credentials.rows["openai-prod"] = existing + + response = _start_chatgpt_login("openai-prod", {"owner_team_id": "team-a"}) + + assert response.status_code == 409, response.text + assert chatgpt_sign_in.auth_server.requests == [] + assert chatgpt_sign_in.credentials.rows == {"openai-prod": existing} + + +def test_chatgpt_sign_in_poll_refuses_a_name_another_provider_took_meanwhile(chatgpt_sign_in): + flow_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a"}) + existing = _another_providers_row("login-a") + chatgpt_sign_in.credentials.rows["login-a"] = existing + chatgpt_sign_in.auth_server.approve("first", "acct-a") + + response = _poll_chatgpt_login("login-a", flow_handle) + + assert response.status_code == 409, response.text + assert chatgpt_sign_in.credentials.writes == [] + assert chatgpt_sign_in.credentials.rows == {"login-a": existing} + + +@pytest.mark.parametrize("failing_url", [CHATGPT_DEVICE_TOKEN_URL, CHATGPT_OAUTH_TOKEN_URL], ids=["poll", "exchange"]) +def test_an_upstream_failure_during_poll_is_a_502_without_the_upstream_body(chatgpt_sign_in, failing_url): + flow_handle = _started_chatgpt_login("login-a", {"owner_team_id": "team-a"}) + chatgpt_sign_in.auth_server.approve("first", "acct-a") + leaked = _chatgpt_jwt({"sub": "token-echoed-in-error-body"}) + chatgpt_sign_in.auth_server.answers[failing_url] = (500, {"error": "upstream-body-7f3k", "detail": leaked}) + + response = _poll_chatgpt_login("login-a", flow_handle) + + assert response.status_code == 502, response.text + assert "upstream-body-7f3k" not in response.text + assert leaked not in response.text + assert chatgpt_sign_in.credentials.rows == {} diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 81a196af676..7c04b9b25ec 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -46,6 +46,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_enforced_params, _match_and_track_policies, _promoted_trace_control_fields, + _resolve_chatgpt_subscription_logins_for_request, _resolve_credential_from_model_config, _resolve_provider_from_deployment, _update_model_if_key_alias_exists, @@ -61,6 +62,7 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES +from litellm.types.proxy.litellm_pre_call_utils import RedactedDict from litellm.types.utils import CredentialItem @@ -10706,6 +10708,107 @@ class TestResolveUserProviderCredentials: ) + +class TestResolveChatGPTSubscriptionLogins: + """``_resolve_chatgpt_subscription_logins_for_request`` hands a request only the ChatGPT subscription + logins its team or its key owns, inside secret_fields where no log or body snapshot shows them.""" + + _LOGIN_NAMES: Final = ("login-team-a", "login-team-b", "login-key-k") + + @staticmethod + def _access_token(name: str) -> str: + return f"eyJhbGciOiJSUzI1NiJ9.{name}-subscription-access.c2lnbmF0dXJl" + + def _serve_logins(self, monkeypatch) -> None: + def login(name: str, **owner: str) -> CredentialItem: + return CredentialItem( + credential_name=name, + credential_values={ + "chatgpt_access_token": self._access_token(name), + "chatgpt_refresh_token": f"rt-{name}-refresh", + "chatgpt_account_id": f"acct-{name}", + "chatgpt_expires_at": "4102444800", + }, + credential_info={"custom_llm_provider": "chatgpt", **owner}, + ) + + monkeypatch.setattr( + litellm, + "credential_list", + [ + login("login-team-a", chatgpt_owner_team_id="team-a"), + login("login-team-b", chatgpt_owner_team_id="team-b"), + login("login-key-k", chatgpt_owner_key_hash="hash-k"), + ], + ) + + def _assert_no_token_in(self, text: str) -> None: + for name in self._LOGIN_NAMES: + assert self._access_token(name) not in text + assert f"rt-{name}-refresh" not in text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("team_id", "key_hash", "owned"), + [ + ("team-a", "hash-unbound", ("login-team-a",)), + (None, "hash-k", ("login-key-k",)), + ("team-b", "hash-k", ("login-team-b", "login-key-k")), + ], + ids=["by-team", "by-key", "by-team-and-key"], + ) + async def test_secret_fields_carry_only_the_callers_logins(self, monkeypatch, team_id, key_hash, owned): + self._serve_logins(monkeypatch) + data = {"model": "chatgpt-seat", "secret_fields": {"raw_headers": {}}} + + await _resolve_chatgpt_subscription_logins_for_request(data=data, team_id=team_id, key_hash=key_hash) + + logins = data["secret_fields"]["chatgpt_subscription_logins"] + assert logins == {name: self._access_token(name) for name in owned} + assert isinstance(logins, RedactedDict) + for shown in (str(logins), repr(logins), str(data), repr(data)): + self._assert_no_token_in(shown) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("team_id", "key_hash"), [("team-unbound", "hash-unbound"), (None, None)], ids=["unbound", "no-owner"] + ) + async def test_a_caller_who_owns_no_login_gets_no_logins_entry(self, monkeypatch, team_id, key_hash): + self._serve_logins(monkeypatch) + data = {"model": "chatgpt-seat", "secret_fields": {"raw_headers": {}}} + + await _resolve_chatgpt_subscription_logins_for_request(data=data, team_id=team_id, key_hash=key_hash) + + assert data["secret_fields"] == {"raw_headers": {}} + + @pytest.mark.asyncio + async def test_a_request_gets_its_keys_and_teams_logins_and_nothing_else_shows_them(self, monkeypatch): + self._serve_logins(monkeypatch) + request = MagicMock(spec=Request) + request.url = MagicMock() + request.url.__str__.return_value = "http://localhost/v1/responses" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json", "Authorization": "Bearer sk-caller"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + + updated = await add_litellm_data_to_request( + data={"model": "chatgpt-seat", "input": "hello"}, + request=request, + user_api_key_dict=UserAPIKeyAuth(token="hash-k", team_id="team-b", user_id="user-1", metadata={}), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["secret_fields"]["chatgpt_subscription_logins"] == { + "login-team-b": self._access_token("login-team-b"), + "login-key-k": self._access_token("login-key-k"), + } + self._assert_no_token_in(json.dumps({k: v for k, v in updated.items() if k != "secret_fields"}, default=str)) + + def _per_user_env(monkeypatch, credential_names=("copilot-cred",)): from litellm.caching.dual_cache import DualCache