mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(proxy): sign in, refresh, and store ChatGPT subscription logins as credentials
Adds admin-only POST /credentials/{name}/chatgpt_login/start and /poll endpoints that run the ChatGPT device-code sign-in and save the login, encrypted, on a credential owned by exactly one team or virtual key. Each request resolves only its own team's or key's logins, and refreshes a stale one once per process (inside the Redis refresh lock when Redis is set). The refresh re-reads the row from the primary database, checks the row still belongs to the caller's owner, and writes the rotated tokens back only if the row still holds the refresh token it redeemed. OAuth errors carry the status code, never the response body
This commit is contained in:
mateo-berri
parent
2132c83dd1
commit
9a1e20f674
10 changed files with 2386 additions and 12 deletions
|
|
@ -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 {}),
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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 == {}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue