mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(chatgpt): bind a ChatGPT subscription login to one team or key
A chatgpt/ deployment whose litellm_credential_name names a subscription login now runs on the access token the caller's own team or key owns, and answers 401 for any other caller instead of falling back to the host's auth file. A request body cannot override that credential name, the login values never merge into call kwargs, and their keys are server owned so a body or non-admin credential write cannot set them. acompletion also passes the call's api_key into its first provider lookup, so a bound login no longer triggers the host device-code login
This commit is contained in:
mateo-berri
parent
b25ea48c4b
commit
2132c83dd1
14 changed files with 1116 additions and 6 deletions
155
litellm/llms/chatgpt/subscription_login.py
Normal file
155
litellm/llms/chatgpt/subscription_login.py
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.exceptions import CallerCredentialAuthenticationError
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.llms.chatgpt.common_utils import CHATGPT_ACCOUNT_ID_HEADER, CHATGPT_API_BASE, chatgpt_client_api_base
|
||||
from litellm.types.utils import CredentialItem, chatgpt_subscription_login_litellm_params
|
||||
|
||||
CHATGPT_SUBSCRIPTION_LOGINS_SECRET_FIELD: Final = "chatgpt_subscription_logins"
|
||||
CHATGPT_OWNER_TEAM_ID_KEY: Final = "chatgpt_owner_team_id"
|
||||
CHATGPT_OWNER_KEY_HASH_KEY: Final = "chatgpt_owner_key_hash"
|
||||
CHATGPT_PROVIDER: Final = "chatgpt"
|
||||
_LOGIN_HEADER_NAMES: Final = frozenset({"authorization", CHATGPT_ACCOUNT_ID_HEADER.lower()})
|
||||
_STR_KEYED: Final = TypeAdapter(Mapping[str, object])
|
||||
_ENTRIES: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def is_chatgpt_subscription_login(credential: CredentialItem) -> bool:
|
||||
return any(key in credential.credential_values for key in chatgpt_subscription_login_litellm_params)
|
||||
|
||||
|
||||
def names_chatgpt_subscription_login(credential_name: object) -> bool:
|
||||
if not isinstance(credential_name, str) or not credential_name:
|
||||
return False
|
||||
credential: Final = CredentialAccessor.find_credential(credential_name)
|
||||
return credential is not None and is_chatgpt_subscription_login(credential)
|
||||
|
||||
|
||||
def locked_chatgpt_credential_name(deployment_params: object) -> str | None:
|
||||
params: Final = _str_keyed(deployment_params)
|
||||
credential_name: Final = params.get("litellm_credential_name") if params is not None else None
|
||||
if params is None or not isinstance(credential_name, str) or not credential_name:
|
||||
return None
|
||||
if names_chatgpt_subscription_login(credential_name) or _is_chatgpt_call(params, params.get("model")):
|
||||
return credential_name
|
||||
return None
|
||||
|
||||
|
||||
def _str_keyed(value: object) -> Mapping[str, object] | None:
|
||||
if not isinstance(value, Mapping):
|
||||
return None
|
||||
try:
|
||||
return _STR_KEYED.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _is_chatgpt_call(kwargs: Mapping[str, object], model: object) -> bool:
|
||||
return kwargs.get("custom_llm_provider") == CHATGPT_PROVIDER or any(
|
||||
isinstance(candidate, str) and candidate.startswith(f"{CHATGPT_PROVIDER}/")
|
||||
for candidate in (model, kwargs.get("model"))
|
||||
)
|
||||
|
||||
|
||||
def _login_unavailable(credential_name: str) -> CallerCredentialAuthenticationError:
|
||||
return CallerCredentialAuthenticationError(
|
||||
message=(
|
||||
f"ChatGPT subscription login '{credential_name}' is not available to this key. "
|
||||
"An admin can sign in to ChatGPT for this team or key under LLM Credentials in the LiteLLM UI"
|
||||
),
|
||||
llm_provider=CHATGPT_PROVIDER,
|
||||
model="",
|
||||
)
|
||||
|
||||
|
||||
def _owned_access_token(kwargs: Mapping[str, object], credential_name: str) -> str | None:
|
||||
secret_fields: Final = _str_keyed(kwargs.get("secret_fields"))
|
||||
logins: Final = _str_keyed(
|
||||
secret_fields.get(CHATGPT_SUBSCRIPTION_LOGINS_SECRET_FIELD) if secret_fields is not None else None
|
||||
)
|
||||
token: Final = logins.get(credential_name) if logins is not None else None
|
||||
return token if isinstance(token, str) and token else None
|
||||
|
||||
|
||||
def _without_login_headers(headers: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {key: value for key, value in headers.items() if key.lower() not in _LOGIN_HEADER_NAMES}
|
||||
|
||||
|
||||
def _scoped_header_without_login(entry: object) -> object:
|
||||
scoped: Final = _str_keyed(entry)
|
||||
extra_headers: Final = _str_keyed(scoped.get("extra_headers")) if scoped is not None else None
|
||||
if scoped is None or extra_headers is None:
|
||||
return entry
|
||||
return {**scoped, "extra_headers": _without_login_headers(extra_headers)}
|
||||
|
||||
|
||||
def _provider_specific_header_without_login(value: object) -> object:
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [_scoped_header_without_login(entry) for entry in _ENTRIES.validate_python(value)]
|
||||
return _scoped_header_without_login(value)
|
||||
|
||||
|
||||
def _nested_extra_headers_without_login(params: object) -> object:
|
||||
mapping: Final = _str_keyed(params)
|
||||
extra_headers: Final = _str_keyed(mapping.get("extra_headers")) if mapping is not None else None
|
||||
if mapping is None or extra_headers is None:
|
||||
return params
|
||||
return {**mapping, "extra_headers": _without_login_headers(extra_headers)}
|
||||
|
||||
|
||||
def _header_mapping_without_login(value: object) -> object:
|
||||
headers: Final = _str_keyed(value)
|
||||
return value if headers is None else _without_login_headers(headers)
|
||||
|
||||
|
||||
def _subscription_api_base(kwargs: Mapping[str, object]) -> str:
|
||||
validated: Final = tuple(
|
||||
chatgpt_client_api_base(value)
|
||||
for value in (kwargs.get("base_url"), kwargs.get("api_base"))
|
||||
if isinstance(value, str) and value
|
||||
)
|
||||
return validated[0] if validated else CHATGPT_API_BASE
|
||||
|
||||
|
||||
_HEADER_CARRIERS: Final[tuple[tuple[str, Callable[[object], object]], ...]] = (
|
||||
("provider_specific_header", _provider_specific_header_without_login),
|
||||
("extra_headers", _header_mapping_without_login),
|
||||
("headers", _header_mapping_without_login),
|
||||
("optional_params", _nested_extra_headers_without_login),
|
||||
("litellm_params", _nested_extra_headers_without_login),
|
||||
)
|
||||
|
||||
|
||||
def _request_updates(kwargs: Mapping[str, object], access_token: str) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
**{key: rebuild(kwargs[key]) for key, rebuild in _HEADER_CARRIERS if key in kwargs},
|
||||
"api_key": access_token,
|
||||
"api_base": _subscription_api_base(kwargs),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def attach_chatgpt_subscription_login(
|
||||
kwargs: dict[str, object], # mutable-ok: kwargs is the request mutation channel
|
||||
model: object,
|
||||
) -> None:
|
||||
credential_name: Final = kwargs.get("litellm_credential_name")
|
||||
if not isinstance(credential_name, str) or not credential_name:
|
||||
return
|
||||
credential: Final = CredentialAccessor.find_credential(credential_name)
|
||||
if credential is None and not _is_chatgpt_call(kwargs, model):
|
||||
return
|
||||
if credential is not None and not is_chatgpt_subscription_login(credential):
|
||||
return
|
||||
access_token: Final = _owned_access_token(kwargs, credential_name)
|
||||
if credential is None or access_token is None or not _is_chatgpt_call(kwargs, model):
|
||||
raise _login_unavailable(credential_name)
|
||||
updates: Final = _request_updates(kwargs, access_token)
|
||||
kwargs.pop("chatgpt_auth_file", None)
|
||||
kwargs.pop("base_url", None)
|
||||
kwargs.update(updates)
|
||||
|
|
@ -698,6 +698,7 @@ async def acompletion(
|
|||
"str | None", kwargs.get("api_base")
|
||||
)
|
||||
or base_url,
|
||||
api_key=api_key,
|
||||
litellm_params=(GenericLiteLLMParams.model_validate(_supplemental) if _supplemental else None),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -469,6 +469,7 @@ from litellm.types.utils import (
|
|||
StandardLoggingRoutingDecision,
|
||||
Usage,
|
||||
all_litellm_params,
|
||||
chatgpt_subscription_login_litellm_params,
|
||||
is_server_derived_pricing_key,
|
||||
shared_backend_model_info,
|
||||
)
|
||||
|
|
@ -5075,6 +5076,22 @@ class Router:
|
|||
model=deployment_model_name,
|
||||
llm_provider="",
|
||||
)
|
||||
from litellm.llms.chatgpt.subscription_login import locked_chatgpt_credential_name
|
||||
|
||||
locked_chatgpt_name: Final = locked_chatgpt_credential_name(deployment_context)
|
||||
if (
|
||||
locked_chatgpt_name is not None
|
||||
and "litellm_credential_name" in kwargs
|
||||
and caller_credential_name != locked_chatgpt_name
|
||||
):
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
"litellm_credential_name cannot be overridden on a deployment that uses "
|
||||
"a ChatGPT subscription login"
|
||||
),
|
||||
model=declared_model if isinstance(declared_model, str) else "",
|
||||
llm_provider="",
|
||||
)
|
||||
|
||||
kwargs["model_info"] = model_info
|
||||
|
||||
|
|
@ -12377,7 +12394,11 @@ class Router:
|
|||
verbose_router_logger.warning(
|
||||
"Credential '%s' not found in credential_list", deployment.litellm_params.litellm_credential_name
|
||||
)
|
||||
credentials.update(credential_values)
|
||||
credentials.update(
|
||||
(key, value)
|
||||
for key, value in credential_values.items()
|
||||
if key not in chatgpt_subscription_login_litellm_params
|
||||
)
|
||||
# Remove the credential name since we've resolved it
|
||||
credentials.pop("litellm_credential_name", None)
|
||||
|
||||
|
|
|
|||
|
|
@ -115,6 +115,7 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di
|
|||
# admin's value in ``litellm_params`` and have it forwarded to the
|
||||
# redirected upstream.
|
||||
if "api_base" in request_kwargs or "base_url" in request_kwargs:
|
||||
from litellm.llms.chatgpt.subscription_login import names_chatgpt_subscription_login
|
||||
from litellm.llms.github_copilot.per_user_auth import github_copilot_auth_mode
|
||||
|
||||
if sends_api_key_secret_reference(litellm_params): # pyright: ignore[reportUnknownArgumentType] # litellm_params is the untyped request dict
|
||||
|
|
@ -141,6 +142,8 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di
|
|||
per_user_credential_locked: Final = github_copilot_auth_mode(
|
||||
credential_name_value, # pyright: ignore[reportUnknownArgumentType] # value comes from the untyped request dict
|
||||
auth_type, # pyright: ignore[reportUnknownArgumentType] # value comes from the untyped request dict
|
||||
) or names_chatgpt_subscription_login(
|
||||
credential_name_value # pyright: ignore[reportUnknownArgumentType] # value comes from the untyped request dict
|
||||
)
|
||||
for field in _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE:
|
||||
if per_user_credential_locked and field in ("litellm_credential_name", "github_copilot_auth_type"):
|
||||
|
|
|
|||
|
|
@ -26,3 +26,4 @@ class SecretFields(TypedDict):
|
|||
raw_headers: dict
|
||||
user_provider_credentials: NotRequired[ReadOnly[Mapping[str, Mapping[str, str]]]]
|
||||
user_provider_credentials_user_id: NotRequired[ReadOnly[str]]
|
||||
chatgpt_subscription_logins: NotRequired[ReadOnly[Mapping[str, str]]]
|
||||
|
|
|
|||
|
|
@ -4267,6 +4267,15 @@ oauth_token_exchange_litellm_params: Final = (
|
|||
"token_exchange_scope",
|
||||
)
|
||||
github_copilot_oauth_litellm_params: Final = ("github_copilot_auth_type", "github_copilot_host")
|
||||
chatgpt_subscription_login_secret_params: Final = (
|
||||
"chatgpt_access_token",
|
||||
"chatgpt_refresh_token",
|
||||
"chatgpt_id_token",
|
||||
)
|
||||
chatgpt_subscription_login_litellm_params: Final = chatgpt_subscription_login_secret_params + (
|
||||
"chatgpt_account_id",
|
||||
"chatgpt_expires_at",
|
||||
)
|
||||
OIDC_SECRET_REFERENCE_PREFIX: Final = "oidc/"
|
||||
API_KEY_SECRET_REFERENCE_REJECTION: Final = (
|
||||
"Rejected Request: api_key holds an oidc/ secret reference, which only a deployment or credential a proxy "
|
||||
|
|
@ -4376,8 +4385,11 @@ server_owned_wif_litellm_params: Final = (
|
|||
+ oauth_token_exchange_litellm_params
|
||||
+ github_copilot_oauth_litellm_params
|
||||
+ ("chatgpt_auth_file",)
|
||||
+ chatgpt_subscription_login_litellm_params
|
||||
)
|
||||
secret_bearing_wif_litellm_params: Final = tuple(
|
||||
sorted(WIF_SECRET_BEARING_KEYS.union(chatgpt_subscription_login_secret_params))
|
||||
)
|
||||
secret_bearing_wif_litellm_params: Final = tuple(sorted(WIF_SECRET_BEARING_KEYS))
|
||||
|
||||
all_litellm_params = [ # rebind-ok: two star imports in litellm/__init__.py re-bind it
|
||||
*OWNED_KWARG_NAMES,
|
||||
|
|
|
|||
|
|
@ -267,6 +267,7 @@ from litellm.types.utils import (
|
|||
StreamingChoices,
|
||||
TranscriptionResponse,
|
||||
Usage,
|
||||
chatgpt_subscription_login_litellm_params,
|
||||
credential_api_key_secret_reference,
|
||||
is_litellm_owned_kwarg,
|
||||
)
|
||||
|
|
@ -779,7 +780,7 @@ def _named_credential_values(kwargs: Mapping[str, object]) -> Mapping[str, objec
|
|||
{
|
||||
key: value
|
||||
for key, value in credential_api_key_secret_reference(credential.credential_values).items()
|
||||
if key not in kwargs
|
||||
if key not in kwargs and key not in chatgpt_subscription_login_litellm_params
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -1940,6 +1941,9 @@ def client(original_function):
|
|||
from litellm.llms.github_copilot.per_user_auth import attach_github_copilot_user_session
|
||||
|
||||
attach_github_copilot_user_session(kwargs) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict
|
||||
from litellm.llms.chatgpt.subscription_login import attach_chatgpt_subscription_login
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, model) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict
|
||||
kwargs["litellm_logging_obj"] = logging_obj
|
||||
LLMCachingHandler: Final = _get_cached_llm_caching_handler()
|
||||
_llm_caching_handler: Final[LLMCachingHandler] = LLMCachingHandler(
|
||||
|
|
@ -2234,6 +2238,9 @@ def client(original_function):
|
|||
from litellm.llms.github_copilot.per_user_auth import aattach_github_copilot_user_session
|
||||
|
||||
await aattach_github_copilot_user_session(kwargs) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict
|
||||
from litellm.llms.chatgpt.subscription_login import attach_chatgpt_subscription_login
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, model) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict
|
||||
logging_obj.llm_caching_handler = _llm_caching_handler
|
||||
# [OPTIONAL] CHECK BUDGET
|
||||
if litellm.max_budget:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import pytest
|
|||
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig
|
||||
from litellm.llms.chatgpt.common_utils import GetAccessTokenError
|
||||
from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE, GetAccessTokenError
|
||||
|
||||
GATEWAY_BASE = "https://chatgpt.example/backend-api/codex"
|
||||
|
||||
|
|
@ -72,7 +72,9 @@ import litellm
|
|||
from litellm.integrations.dotprompt.dotprompt_manager import DotpromptManager
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
||||
def _write_auth_record(path: Path, token: str, account_id: str) -> None:
|
||||
|
|
@ -636,3 +638,54 @@ async def test_requests_preserve_gateway_token_and_selected_workspace(
|
|||
else await litellm.aresponses(**params, input=[{"role": "user", "content": "Hello"}])
|
||||
)
|
||||
assert response.output[0].content[0].text == expected_account
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync", [False, True])
|
||||
@pytest.mark.parametrize("through_router", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_on_a_subscription_deployment_sends_the_owned_login_instead_of_the_host_auth_file(
|
||||
auth_files: tuple[Path, Path], through_router: bool, sync: bool, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
subscription_token: Final = _client_token("acct-team-a")
|
||||
login: Final = CredentialItem(
|
||||
credential_name="team-a-chatgpt",
|
||||
credential_values={
|
||||
"chatgpt_access_token": subscription_token,
|
||||
"chatgpt_refresh_token": "refresh-team-a-example",
|
||||
"chatgpt_id_token": "id-team-a-example",
|
||||
"chatgpt_account_id": "acct-team-a",
|
||||
"chatgpt_expires_at": "4102444800",
|
||||
},
|
||||
credential_info={"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-a"},
|
||||
)
|
||||
monkeypatch.setattr(litellm, "credential_list", [login])
|
||||
deployment_params: Final = {"model": "chatgpt/gpt-6.1-sol", "litellm_credential_name": login.credential_name}
|
||||
router: Final = (
|
||||
litellm.Router(model_list=[{"model_name": "team-a-gpt", "litellm_params": deployment_params}])
|
||||
if through_router
|
||||
else None
|
||||
)
|
||||
params: Final = {
|
||||
**({"model": "team-a-gpt"} if router else deployment_params),
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"secret_fields": {
|
||||
"raw_headers": {},
|
||||
"chatgpt_subscription_logins": RedactedDict({login.credential_name: subscription_token}),
|
||||
},
|
||||
}
|
||||
complete: Final = (router.completion if sync else router.acompletion) if router else (
|
||||
litellm.completion if sync else litellm.acompletion
|
||||
)
|
||||
|
||||
with respx.mock as transport:
|
||||
route: Final = transport.post(f"{CHATGPT_API_BASE}/responses").mock(side_effect=_subscription_response)
|
||||
result: Final = await asyncio.to_thread(complete, **params) if sync else await complete(**params)
|
||||
|
||||
assert result.choices[0].message.content == "acct-team-a"
|
||||
assert route.call_count == 1
|
||||
sent: Final = route.calls.last.request
|
||||
assert sent.headers.get_list("authorization") == [f"Bearer {subscription_token}"]
|
||||
body: Final = sent.content.decode()
|
||||
assert [value for value in login.credential_values.values() if value in body] == []
|
||||
assert [key for key in login.credential_values if key in body] == []
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Tests for ChatGPT subscription Responses API transformation
|
|||
Source: litellm/llms/chatgpt/responses/transformation.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import Generator
|
||||
|
|
@ -13,15 +14,18 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.chatgpt.authenticator import Authenticator
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE
|
||||
from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.main import responses_api_bridge_check
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.types.utils import CredentialItem, LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
|
||||
|
|
@ -667,3 +671,127 @@ def test_prompt_cache_key_keeps_one_session_header_across_requests(
|
|||
)
|
||||
|
||||
assert [call.args[0].headers["session_id"] for call in upstream.call_args_list] == ["conv-123", "conv-123"]
|
||||
|
||||
|
||||
def _team_subscription_login(access_token: str) -> CredentialItem:
|
||||
return CredentialItem(
|
||||
credential_name="team-a-chatgpt",
|
||||
credential_values={
|
||||
"chatgpt_access_token": access_token,
|
||||
"chatgpt_refresh_token": "refresh-team-a-example",
|
||||
"chatgpt_id_token": "id-team-a-example",
|
||||
"chatgpt_account_id": "acct-team-a",
|
||||
"chatgpt_expires_at": "4102444800",
|
||||
},
|
||||
credential_info={"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-a"},
|
||||
)
|
||||
|
||||
|
||||
def _completed_responses_reply(_request: httpx.Request) -> httpx.Response:
|
||||
completed: Final = {
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-5.5",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "OK", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
}
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
content=f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_subscription_login_replaces_the_host_auth_file_login_on_the_wire(
|
||||
sync: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
|
||||
host_token: Final = _fake_chatgpt_oauth_token("acct-host")
|
||||
(tmp_path / "auth.json").write_text(
|
||||
json.dumps({"access_token": host_token, "account_id": "acct-host", "expires_at": 4102444800})
|
||||
)
|
||||
subscription_token: Final = _fake_chatgpt_oauth_token("acct-team-a")
|
||||
login: Final = _team_subscription_login(subscription_token)
|
||||
monkeypatch.setattr(litellm, "credential_list", [login])
|
||||
params: Final = {
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"input": "hi",
|
||||
"litellm_credential_name": login.credential_name,
|
||||
"secret_fields": {
|
||||
"raw_headers": {},
|
||||
"chatgpt_subscription_logins": RedactedDict({login.credential_name: subscription_token}),
|
||||
},
|
||||
}
|
||||
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
route: Final = upstream.post(f"{CHATGPT_API_BASE}/responses").mock(side_effect=_completed_responses_reply)
|
||||
response: Final = (
|
||||
await asyncio.to_thread(litellm.responses, **params) if sync else await litellm.aresponses(**params)
|
||||
)
|
||||
|
||||
assert response.output_text == "OK"
|
||||
assert route.call_count == 1
|
||||
sent: Final = route.calls.last.request
|
||||
assert sent.headers.get_list("authorization") == [f"Bearer {subscription_token}"]
|
||||
assert sent.headers.get_list("chatgpt-account-id") == ["acct-team-a"]
|
||||
body: Final = sent.content.decode()
|
||||
assert [value for value in login.credential_values.values() if str(value) in body] == []
|
||||
assert [key for key in login.credential_values if key in body] == []
|
||||
assert host_token not in body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync", [False, True])
|
||||
@pytest.mark.parametrize("owned_logins", [None, {"team-b-chatgpt": "team-b-example-token"}])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_subscription_login_the_caller_does_not_own_never_falls_back_to_the_host_auth_file(
|
||||
sync: bool, owned_logins: dict[str, str] | None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
|
||||
(tmp_path / "auth.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"access_token": _fake_chatgpt_oauth_token("acct-host"),
|
||||
"account_id": "acct-host",
|
||||
"expires_at": 4102444800,
|
||||
}
|
||||
)
|
||||
)
|
||||
login: Final = _team_subscription_login(_fake_chatgpt_oauth_token("acct-team-a"))
|
||||
monkeypatch.setattr(litellm, "credential_list", [login])
|
||||
params: Final = {
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"input": "hi",
|
||||
"litellm_credential_name": login.credential_name,
|
||||
**(
|
||||
{}
|
||||
if owned_logins is None
|
||||
else {"secret_fields": {"raw_headers": {}, "chatgpt_subscription_logins": RedactedDict(owned_logins)}}
|
||||
),
|
||||
}
|
||||
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
route: Final = upstream.post(f"{CHATGPT_API_BASE}/responses").mock(side_effect=_completed_responses_reply)
|
||||
with pytest.raises(litellm.CallerCredentialAuthenticationError) as raised:
|
||||
await (asyncio.to_thread(litellm.responses, **params) if sync else litellm.aresponses(**params))
|
||||
|
||||
assert raised.value.status_code == 401
|
||||
assert route.call_count == 0
|
||||
|
|
|
|||
309
tests/unit/llms/chatgpt/test_subscription_login.py
Normal file
309
tests/unit/llms/chatgpt/test_subscription_login.py
Normal file
|
|
@ -0,0 +1,309 @@
|
|||
import base64
|
||||
import copy
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import CallerCredentialAuthenticationError
|
||||
from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE, ChatGPTAuthError
|
||||
from litellm.llms.chatgpt.subscription_login import (
|
||||
attach_chatgpt_subscription_login,
|
||||
names_chatgpt_subscription_login,
|
||||
)
|
||||
from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
||||
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'})}.{segment(claims)}.fake-signature"
|
||||
|
||||
|
||||
TEAM_A_LOGIN: Final = "team-a-chatgpt"
|
||||
TEAM_B_LOGIN: Final = "team-b-chatgpt"
|
||||
TEAM_A_TOKEN: Final = _jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-team-a"}, "exp": 4102444800})
|
||||
TEAM_B_TOKEN: Final = _jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-team-b"}, "exp": 4102444800})
|
||||
CLIENT_TOKEN: Final = _jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-client"}})
|
||||
|
||||
|
||||
def _subscription(name: str, access_token: str, owner_team_id: str) -> CredentialItem:
|
||||
return CredentialItem(
|
||||
credential_name=name,
|
||||
credential_values={
|
||||
"chatgpt_access_token": access_token,
|
||||
"chatgpt_refresh_token": f"refresh-{name}",
|
||||
"chatgpt_id_token": access_token,
|
||||
"chatgpt_account_id": f"acct-{owner_team_id}",
|
||||
"chatgpt_expires_at": "4102444800",
|
||||
},
|
||||
credential_info={"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": owner_team_id},
|
||||
)
|
||||
|
||||
|
||||
ORDINARY: Final = CredentialItem(
|
||||
credential_name="plain-openai", credential_values={"api_key": "sk-plain-example"}, credential_info={}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def credentials(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
_subscription(TEAM_A_LOGIN, TEAM_A_TOKEN, "team-a"),
|
||||
_subscription(TEAM_B_LOGIN, TEAM_B_TOKEN, "team-b"),
|
||||
ORDINARY,
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _owned(*logins: tuple[str, str]) -> dict[str, object]:
|
||||
return {"chatgpt_subscription_logins": RedactedDict(dict(logins)), "raw_headers": {}}
|
||||
|
||||
|
||||
def _request(**overrides: object) -> dict[str, object]:
|
||||
return {
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"litellm_credential_name": TEAM_A_LOGIN,
|
||||
"chatgpt_auth_file": "/host/chatgpt/auth.json",
|
||||
"secret_fields": _owned((TEAM_A_LOGIN, TEAM_A_TOKEN)),
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "overrides"),
|
||||
[
|
||||
("chatgpt/gpt-5.5", {}),
|
||||
(None, {}),
|
||||
("gpt-5.5", {"model": "gpt-5.5", "custom_llm_provider": "chatgpt"}),
|
||||
],
|
||||
ids=["model_argument", "model_kwarg", "custom_llm_provider"],
|
||||
)
|
||||
def test_attach_sends_the_owned_login_to_chatgpt_instead_of_the_host_auth_file(
|
||||
model: str | None, overrides: dict[str, object]
|
||||
) -> None:
|
||||
kwargs: Final = _request(**overrides)
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, model)
|
||||
|
||||
assert kwargs["api_key"] == TEAM_A_TOKEN
|
||||
assert kwargs["api_base"] == CHATGPT_API_BASE
|
||||
assert "chatgpt_auth_file" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"secret_fields",
|
||||
[
|
||||
_owned((TEAM_B_LOGIN, TEAM_B_TOKEN)),
|
||||
_owned(),
|
||||
_owned((TEAM_A_LOGIN, "")),
|
||||
{"raw_headers": {}},
|
||||
None,
|
||||
],
|
||||
ids=["another_teams_login", "no_logins", "empty_token", "no_login_field", "no_secret_fields"],
|
||||
)
|
||||
def test_a_login_the_caller_does_not_own_fails_closed_with_a_401(secret_fields: dict[str, object] | None) -> None:
|
||||
kwargs: Final = _request(secret_fields=secret_fields) if secret_fields is not None else _request()
|
||||
if secret_fields is None:
|
||||
kwargs.pop("secret_fields")
|
||||
before: Final = copy.deepcopy(kwargs)
|
||||
|
||||
with pytest.raises(CallerCredentialAuthenticationError) as raised:
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
assert raised.value.status_code == 401
|
||||
assert TEAM_A_LOGIN in str(raised.value)
|
||||
assert TEAM_A_TOKEN not in str(raised.value)
|
||||
assert TEAM_B_TOKEN not in str(raised.value)
|
||||
assert kwargs == before
|
||||
|
||||
|
||||
def test_a_chatgpt_deployment_naming_a_missing_credential_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [])
|
||||
kwargs: Final = _request()
|
||||
|
||||
with pytest.raises(CallerCredentialAuthenticationError) as raised:
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
assert raised.value.status_code == 401
|
||||
assert TEAM_A_TOKEN not in str(raised.value)
|
||||
assert "api_key" not in kwargs
|
||||
assert kwargs["chatgpt_auth_file"] == "/host/chatgpt/auth.json"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "model"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"model": "openai/gpt-5.4-mini",
|
||||
"litellm_credential_name": "deleted-credential",
|
||||
"secret_fields": _owned(),
|
||||
},
|
||||
"openai/gpt-5.4-mini",
|
||||
),
|
||||
(
|
||||
{"model": "chatgpt/gpt-5.5", "litellm_credential_name": "plain-openai", "chatgpt_auth_file": "/a.json"},
|
||||
"chatgpt/gpt-5.5",
|
||||
),
|
||||
(
|
||||
{
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"litellm_credential_name": "plain-openai",
|
||||
"secret_fields": _owned(("plain-openai", TEAM_A_TOKEN)),
|
||||
},
|
||||
"chatgpt/gpt-5.5",
|
||||
),
|
||||
({"model": "chatgpt/gpt-5.5", "secret_fields": _owned((TEAM_A_LOGIN, TEAM_A_TOKEN))}, "chatgpt/gpt-5.5"),
|
||||
({"model": "chatgpt/gpt-5.5", "litellm_credential_name": ""}, "chatgpt/gpt-5.5"),
|
||||
],
|
||||
ids=[
|
||||
"missing_credential_on_another_provider",
|
||||
"ordinary_credential",
|
||||
"ordinary_credential_with_a_same_named_login",
|
||||
"no_credential_name",
|
||||
"empty_credential_name",
|
||||
],
|
||||
)
|
||||
def test_requests_that_name_no_subscription_login_are_left_untouched(
|
||||
kwargs: dict[str, object], model: str
|
||||
) -> None:
|
||||
before: Final = copy.deepcopy(kwargs)
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, model)
|
||||
|
||||
assert kwargs == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "overrides"),
|
||||
[
|
||||
("openai/gpt-5.4-mini", {"model": "openai/gpt-5.4-mini"}),
|
||||
("gpt-5.4-mini", {"model": "gpt-5.4-mini", "custom_llm_provider": "openai"}),
|
||||
],
|
||||
)
|
||||
def test_a_subscription_login_is_never_sent_to_another_provider(model: str, overrides: dict[str, object]) -> None:
|
||||
kwargs: Final = _request(**overrides)
|
||||
|
||||
with pytest.raises(CallerCredentialAuthenticationError) as raised:
|
||||
attach_chatgpt_subscription_login(kwargs, model)
|
||||
|
||||
assert raised.value.status_code == 401
|
||||
assert TEAM_A_TOKEN not in str(raised.value)
|
||||
assert "api_key" not in kwargs
|
||||
|
||||
|
||||
def _login_headers() -> dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {CLIENT_TOKEN}",
|
||||
"chatgpt-account-id": "acct-client",
|
||||
"originator": "codex_cli_rs",
|
||||
}
|
||||
|
||||
|
||||
def _header_names(headers: object) -> set[str]:
|
||||
assert isinstance(headers, Mapping)
|
||||
return {str(name).lower() for name in headers}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scoped_as_list", [False, True])
|
||||
def test_attach_strips_a_forwarded_chatgpt_login_from_every_header_carrier(scoped_as_list: bool) -> None:
|
||||
scoped: Final = {"custom_llm_provider": "chatgpt", "extra_headers": _login_headers()}
|
||||
kwargs: Final = _request(
|
||||
provider_specific_header=[scoped, {"custom_llm_provider": "openai", "extra_headers": _login_headers()}]
|
||||
if scoped_as_list
|
||||
else scoped,
|
||||
extra_headers={"AUTHORIZATION": f"Bearer {CLIENT_TOKEN}", "ChatGPT-Account-Id": "acct-client", "x-trace": "1"},
|
||||
headers={"authorization": f"Bearer {CLIENT_TOKEN}", "x-trace": "1"},
|
||||
optional_params={"extra_headers": _login_headers(), "temperature": 0.2},
|
||||
litellm_params={"extra_headers": _login_headers(), "metadata": {"tag": "kept"}},
|
||||
)
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
entries: Final = kwargs["provider_specific_header"] if scoped_as_list else [kwargs["provider_specific_header"]]
|
||||
assert isinstance(entries, list)
|
||||
expected_providers: Final = ["chatgpt", "openai"] if scoped_as_list else ["chatgpt"]
|
||||
assert [entry["custom_llm_provider"] for entry in entries] == expected_providers
|
||||
assert all(_header_names(entry["extra_headers"]) == {"originator"} for entry in entries)
|
||||
assert _header_names(kwargs["extra_headers"]) == {"x-trace"}
|
||||
assert _header_names(kwargs["headers"]) == {"x-trace"}
|
||||
assert kwargs["optional_params"] == {"extra_headers": {"originator": "codex_cli_rs"}, "temperature": 0.2}
|
||||
assert kwargs["litellm_params"] == {"extra_headers": {"originator": "codex_cli_rs"}, "metadata": {"tag": "kept"}}
|
||||
assert kwargs["api_key"] == TEAM_A_TOKEN
|
||||
assert CLIENT_TOKEN not in json.dumps({key: value for key, value in kwargs.items() if key != "secret_fields"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides",
|
||||
[
|
||||
{"api_base": "https://attacker.example/backend-api/codex"},
|
||||
{"base_url": "https://attacker.example/backend-api/codex"},
|
||||
{"api_base": "http://chatgpt.com/backend-api/codex"},
|
||||
{"api_base": "https://chatgpt.com.attacker.example/backend-api/codex"},
|
||||
{"base_url": CHATGPT_API_BASE, "api_base": "https://attacker.example/backend-api/codex"},
|
||||
{"base_url": "https://attacker.example/backend-api/codex", "api_base": CHATGPT_API_BASE},
|
||||
],
|
||||
)
|
||||
def test_a_client_base_off_chatgpt_is_refused_before_the_login_is_attached(overrides: dict[str, object]) -> None:
|
||||
kwargs: Final = _request(**overrides)
|
||||
before: Final = copy.deepcopy(kwargs)
|
||||
|
||||
with pytest.raises(ChatGPTAuthError) as raised:
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert TEAM_A_TOKEN not in str(raised.value)
|
||||
assert kwargs == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["api_base", "base_url"])
|
||||
def test_a_client_base_on_chatgpt_is_kept_as_api_base_and_base_url_is_removed(field: str) -> None:
|
||||
client_base: Final = "https://chatgpt.com:443/backend-api/codex"
|
||||
kwargs: Final = _request(**{field: client_base})
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
assert kwargs["api_base"] == client_base
|
||||
assert "base_url" not in kwargs
|
||||
assert kwargs["api_key"] == TEAM_A_TOKEN
|
||||
|
||||
|
||||
def test_a_caller_api_base_override_on_a_subscription_deployment_never_sends_the_login_off_chatgpt() -> None:
|
||||
redirected: Final = get_dynamic_litellm_params(
|
||||
litellm_params={"model": "chatgpt/gpt-5.5", "litellm_credential_name": TEAM_A_LOGIN},
|
||||
request_kwargs={"api_base": "https://attacker.example/backend-api/codex"},
|
||||
)
|
||||
kwargs: Final = {**redirected, "secret_fields": _owned((TEAM_A_LOGIN, TEAM_A_TOKEN))}
|
||||
|
||||
with pytest.raises(ChatGPTAuthError):
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
assert "api_key" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("credential_name", "expected"),
|
||||
[(TEAM_A_LOGIN, True), ("plain-openai", False), ("deleted-credential", False), ("", False), (None, False)],
|
||||
)
|
||||
def test_names_chatgpt_subscription_login(credential_name: object, expected: bool) -> None:
|
||||
assert names_chatgpt_subscription_login(credential_name) is expected
|
||||
|
||||
|
||||
def test_attach_is_repeatable_on_the_same_request() -> None:
|
||||
kwargs: Final = _request()
|
||||
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
first: Final = dict(kwargs)
|
||||
attach_chatgpt_subscription_login(kwargs, "chatgpt/gpt-5.5")
|
||||
|
||||
assert kwargs == first
|
||||
assert kwargs["api_key"] == TEAM_A_TOKEN
|
||||
|
|
@ -237,6 +237,110 @@ def test_a_request_body_cannot_pick_a_federated_identity_by_credential_name(monk
|
|||
)
|
||||
|
||||
|
||||
_CHATGPT_LOGIN_FIELDS: Final = (
|
||||
"chatgpt_access_token",
|
||||
"chatgpt_refresh_token",
|
||||
"chatgpt_id_token",
|
||||
"chatgpt_account_id",
|
||||
"chatgpt_expires_at",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", _CHATGPT_LOGIN_FIELDS)
|
||||
@pytest.mark.parametrize("nested", [False, True], ids=["top_level", "nested_litellm_params"])
|
||||
@pytest.mark.parametrize("allow_client_side_credentials", [False, True])
|
||||
def test_a_request_body_cannot_carry_a_chatgpt_subscription_login_value(
|
||||
field: str, nested: bool, allow_client_side_credentials: bool
|
||||
) -> None:
|
||||
login: Final = {field: "attacker-chosen"}
|
||||
with pytest.raises(
|
||||
ValueError, match="server-owned workload identity federation or OAuth token exchange parameter"
|
||||
) as error:
|
||||
is_request_body_safe(
|
||||
request_body={"model": "chatgpt/gpt-5.5", **({"litellm_params": login} if nested else login)},
|
||||
general_settings={"allow_client_side_credentials": allow_client_side_credentials},
|
||||
llm_router=None,
|
||||
model="chatgpt/gpt-5.5",
|
||||
)
|
||||
|
||||
assert field in str(error.value)
|
||||
assert "attacker-chosen" not in str(error.value)
|
||||
|
||||
|
||||
def test_a_deployment_opt_in_cannot_let_a_request_body_carry_a_chatgpt_refresh_token() -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "chatgpt/gpt-5.5",
|
||||
"litellm_params": {
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"configurable_clientside_auth_params": ["chatgpt_refresh_token", "api_key"],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="chatgpt_refresh_token"):
|
||||
is_request_body_safe(
|
||||
request_body={"model": "chatgpt/gpt-5.5", "chatgpt_refresh_token": "attacker-chosen"},
|
||||
general_settings={},
|
||||
llm_router=router,
|
||||
model="chatgpt/gpt-5.5",
|
||||
)
|
||||
|
||||
|
||||
_NAMES_TEAM_A_LOGIN: Final = {"litellm_credential_name": "team-a-chatgpt"}
|
||||
_CLIENT_ROUTER_MODEL_NAMING_TEAM_A_LOGIN: Final = {
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"user_config": _client_router_config("model_list", {"model": "chatgpt/gpt-5.5", **_NAMES_TEAM_A_LOGIN}),
|
||||
}
|
||||
_CLIENT_ROUTER_DEFAULTS_NAMING_TEAM_A_LOGIN: Final = {
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"user_config": _client_router_config("default_litellm_params", dict(_NAMES_TEAM_A_LOGIN)),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "allow_client_side_credentials"),
|
||||
[
|
||||
pytest.param({"model": "chatgpt/gpt-5.5", **_NAMES_TEAM_A_LOGIN}, False, id="top_level"),
|
||||
pytest.param({"model": "chatgpt/gpt-5.5", **_NAMES_TEAM_A_LOGIN}, True, id="top_level_opted_in"),
|
||||
pytest.param({"model": "chatgpt/gpt-5.5", "litellm_params": _NAMES_TEAM_A_LOGIN}, False, id="nested"),
|
||||
pytest.param(
|
||||
{"model": "chatgpt/gpt-5.5", "litellm_params": _NAMES_TEAM_A_LOGIN}, True, id="nested_opted_in"
|
||||
),
|
||||
pytest.param(_CLIENT_ROUTER_MODEL_NAMING_TEAM_A_LOGIN, True, id="client_router_model_list"),
|
||||
pytest.param(_CLIENT_ROUTER_DEFAULTS_NAMING_TEAM_A_LOGIN, True, id="client_router_defaults"),
|
||||
],
|
||||
)
|
||||
def test_a_request_body_cannot_pick_a_chatgpt_subscription_login_by_credential_name(
|
||||
monkeypatch: pytest.MonkeyPatch, body: dict[str, object], allow_client_side_credentials: bool
|
||||
) -> None:
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="team-a-chatgpt",
|
||||
credential_values={"chatgpt_access_token": "access-team-a-example", "chatgpt_account_id": "acct-a"},
|
||||
credential_info={"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-a"},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="names a credential configured for") as error:
|
||||
is_request_body_safe(
|
||||
request_body=body,
|
||||
general_settings={"allow_client_side_credentials": allow_client_side_credentials},
|
||||
llm_router=None,
|
||||
model="chatgpt/gpt-5.5",
|
||||
)
|
||||
|
||||
assert "access-team-a-example" not in str(error.value)
|
||||
|
||||
|
||||
def test_a_request_body_may_still_name_a_credential_that_does_not_federate(monkeypatch):
|
||||
"""Only federation makes a credential a deployment decision. An ordinary named credential stays
|
||||
usable from a request body, so the ban must read what the credential holds, not its presence."""
|
||||
|
|
@ -3193,6 +3297,33 @@ class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
|||
)
|
||||
assert "github_copilot_auth_type" not in out
|
||||
|
||||
@pytest.mark.parametrize("override", ["api_base", "base_url"])
|
||||
@pytest.mark.parametrize(
|
||||
"credential_values",
|
||||
[
|
||||
pytest.param({"chatgpt_access_token": "access-team-a-example"}, id="chatgpt_subscription_login"),
|
||||
pytest.param({"chatgpt_account_id": "acct-team-a"}, id="chatgpt_login_without_token_yet"),
|
||||
],
|
||||
)
|
||||
def test_a_chatgpt_subscription_credential_name_survives_a_base_override(
|
||||
self, override: str, credential_values: dict[str, str]
|
||||
):
|
||||
from litellm.router_utils.clientside_credential_handler import get_dynamic_litellm_params
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
with patch.object(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[CredentialItem(credential_name="team-a-login", credential_values=credential_values, credential_info={})],
|
||||
):
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={"model": "chatgpt/gpt-5.5", "litellm_credential_name": "team-a-login"},
|
||||
request_kwargs={override: "https://attacker.example"},
|
||||
)
|
||||
|
||||
assert out["litellm_credential_name"] == "team-a-login"
|
||||
assert out[override] == "https://attacker.example"
|
||||
|
||||
def test_field_echo_does_not_preserve_admin_value(self):
|
||||
# Regression: a caller that echoes an admin-config field name with
|
||||
# an *empty* value (or any value) must not be able to keep the
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import functools
|
||||
import gc
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
import weakref
|
||||
from collections.abc import Iterator, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from types import MethodType
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -3354,3 +3357,169 @@ def test_upsert_deployment_clears_stale_budget_config(monkeypatch):
|
|||
|
||||
router.upsert_deployment(deployment=unbudgeted)
|
||||
assert budget_limiter._get_budget_config_for_deployment(model_id) is None
|
||||
|
||||
|
||||
def test_get_deployment_credentials_with_provider_never_returns_chatgpt_subscription_login_values(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
login_values: Final = {
|
||||
"chatgpt_access_token": "access-team-a-example",
|
||||
"chatgpt_refresh_token": "refresh-team-a-example",
|
||||
"chatgpt_id_token": "id-team-a-example",
|
||||
"chatgpt_account_id": "acct-team-a",
|
||||
"chatgpt_expires_at": "4102444800",
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(
|
||||
credential_name="team-a-chatgpt",
|
||||
credential_values={**login_values, "api_version": "2026-01-01"},
|
||||
credential_info={"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-a"},
|
||||
)
|
||||
],
|
||||
)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "team-a-gpt",
|
||||
"litellm_params": {"model": "chatgpt/gpt-5.5", "litellm_credential_name": "team-a-chatgpt"},
|
||||
"model_info": {"id": "team-a-subscription"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
credentials: Final = router.get_deployment_credentials_with_provider(model_id="team-a-subscription")
|
||||
|
||||
assert credentials == {"model": "chatgpt/gpt-5.5", "custom_llm_provider": "chatgpt", "api_version": "2026-01-01"}
|
||||
|
||||
|
||||
def _fake_chatgpt_jwt(account_id: str) -> str:
|
||||
claims: Final = {"exp": 4102444800, "https://api.openai.com/auth": {"chatgpt_account_id": account_id}}
|
||||
payload: Final = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=")
|
||||
return f"e30.{payload}.fake-signature"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def chatgpt_logins(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False)
|
||||
(tmp_path / "auth.json").write_text(
|
||||
json.dumps(
|
||||
{"access_token": _fake_chatgpt_jwt("acct-host"), "account_id": "acct-host", "expires_at": 4102444800}
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
*(
|
||||
CredentialItem(
|
||||
credential_name=name,
|
||||
credential_values={"chatgpt_access_token": _fake_chatgpt_jwt(f"acct-{name}")},
|
||||
credential_info={"custom_llm_provider": "chatgpt"},
|
||||
)
|
||||
for name in ("team-a-chatgpt", "team-b-chatgpt")
|
||||
),
|
||||
CredentialItem(credential_name="plain-key", credential_values={"api_key": "sk-plain"}, credential_info={}),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
_LOCKED_CHATGPT_DEPLOYMENTS: Final = {
|
||||
"subscription": {"model": "chatgpt/gpt-5.5", "litellm_credential_name": "team-a-chatgpt"},
|
||||
"subscription_by_provider": {
|
||||
"model": "gpt-5.5",
|
||||
"custom_llm_provider": "chatgpt",
|
||||
"litellm_credential_name": "team-a-chatgpt",
|
||||
},
|
||||
"chatgpt_naming_a_credential": {"model": "chatgpt/gpt-5.5", "litellm_credential_name": "plain-key"},
|
||||
}
|
||||
_CREDENTIAL_NAME_OVERRIDES: Final = ("", None, "team-b-chatgpt", "plain-key", "unknown-credential")
|
||||
|
||||
|
||||
def _owned_logins() -> dict[str, object]:
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
|
||||
return {"chatgpt_subscription_logins": RedactedDict({"team-a-chatgpt": _fake_chatgpt_jwt("acct-team-a-chatgpt")})}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("chatgpt_logins")
|
||||
@pytest.mark.parametrize(
|
||||
("deployment", "override"),
|
||||
[
|
||||
pytest.param(deployment, override, id=f"{name}-{override!r}")
|
||||
for (name, deployment), override in itertools.product(
|
||||
_LOCKED_CHATGPT_DEPLOYMENTS.items(), _CREDENTIAL_NAME_OVERRIDES
|
||||
)
|
||||
if override != deployment["litellm_credential_name"]
|
||||
],
|
||||
)
|
||||
async def test_a_request_cannot_swap_the_credential_a_chatgpt_deployment_names(
|
||||
deployment: dict[str, str], override: str | None
|
||||
) -> None:
|
||||
router: Final = Router(model_list=[{"model_name": "team-a-gpt", "litellm_params": deployment}])
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match="cannot be overridden on a deployment that uses a ChatGPT"):
|
||||
await router.acompletion(
|
||||
model="team-a-gpt",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="ran",
|
||||
secret_fields=_owned_logins(),
|
||||
litellm_credential_name=override,
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("chatgpt_logins")
|
||||
@pytest.mark.parametrize("deployment", _LOCKED_CHATGPT_DEPLOYMENTS.values(), ids=_LOCKED_CHATGPT_DEPLOYMENTS.keys())
|
||||
@pytest.mark.parametrize("repeats_its_own_name", [False, True])
|
||||
async def test_a_chatgpt_deployment_runs_when_the_request_leaves_its_credential_alone(
|
||||
deployment: dict[str, str], repeats_its_own_name: bool
|
||||
) -> None:
|
||||
router: Final = Router(model_list=[{"model_name": "team-a-gpt", "litellm_params": deployment}])
|
||||
own_name: Final = {"litellm_credential_name": deployment["litellm_credential_name"]} if repeats_its_own_name else {}
|
||||
|
||||
response: Final = await router.acompletion(
|
||||
model="team-a-gpt",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="ran",
|
||||
secret_fields=_owned_logins(),
|
||||
num_retries=0,
|
||||
**own_name,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "ran"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("chatgpt_logins")
|
||||
@pytest.mark.parametrize("override", ["", None, "unknown-credential"])
|
||||
async def test_a_deployment_that_is_not_chatgpt_keeps_its_existing_credential_override_behavior(
|
||||
override: str | None,
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "plain-gpt",
|
||||
"litellm_params": {"model": "openai/gpt-5.4-mini", "litellm_credential_name": "plain-key"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
response: Final = await router.acompletion(
|
||||
model="plain-gpt",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="ran",
|
||||
litellm_credential_name=override,
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "ran"
|
||||
|
|
|
|||
|
|
@ -7183,6 +7183,124 @@ def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_with
|
|||
assert _credential_warnings(caplog) == []
|
||||
|
||||
|
||||
_CHATGPT_SUBSCRIPTION_LOGIN_VALUES: Final[Mapping[str, str]] = {
|
||||
"chatgpt_access_token": "access-team-a-example",
|
||||
"chatgpt_refresh_token": "refresh-team-a-example",
|
||||
"chatgpt_id_token": "id-team-a-example",
|
||||
"chatgpt_account_id": "acct-team-a",
|
||||
"chatgpt_expires_at": "4102444800",
|
||||
}
|
||||
|
||||
|
||||
def _chatgpt_subscription_credential(**extra_values: str) -> CredentialItem:
|
||||
return CredentialItem(
|
||||
credential_name="team-a-chatgpt",
|
||||
credential_values={**_CHATGPT_SUBSCRIPTION_LOGIN_VALUES, **extra_values},
|
||||
credential_info={"custom_llm_provider": "chatgpt", "chatgpt_owner_team_id": "team-a"},
|
||||
)
|
||||
|
||||
|
||||
def test_load_credentials_from_list_never_merges_chatgpt_subscription_login_values_into_kwargs(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.utils import load_credentials_from_list, with_named_credential
|
||||
|
||||
monkeypatch.setattr(litellm, "credential_list", [_chatgpt_subscription_credential(api_version="2026-01-01")])
|
||||
request_kwargs: Final[dict[str, object]] = {
|
||||
"litellm_credential_name": "team-a-chatgpt",
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
}
|
||||
load_credentials_from_list(request_kwargs)
|
||||
|
||||
assert request_kwargs == {
|
||||
"litellm_credential_name": "team-a-chatgpt",
|
||||
"model": "chatgpt/gpt-5.5",
|
||||
"api_version": "2026-01-01",
|
||||
}
|
||||
assert with_named_credential({"litellm_credential_name": "team-a-chatgpt"}) == {
|
||||
"litellm_credential_name": "team-a-chatgpt",
|
||||
"api_version": "2026-01-01",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"secret_fields",
|
||||
[
|
||||
pytest.param(None, id="no_secret_fields"),
|
||||
pytest.param({"chatgpt_subscription_logins": {}}, id="no_owned_logins"),
|
||||
pytest.param({"chatgpt_subscription_logins": {"team-b-chatgpt": "access-team-b-example"}}, id="another_login"),
|
||||
],
|
||||
)
|
||||
async def test_acompletion_on_a_chatgpt_subscription_the_caller_does_not_own_is_a_401(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, secret_fields: dict[str, object] | None
|
||||
) -> None:
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(litellm, "credential_list", [_chatgpt_subscription_credential()])
|
||||
with pytest.raises(litellm.CallerCredentialAuthenticationError) as raised:
|
||||
await litellm.acompletion(
|
||||
model="chatgpt/gpt-5.5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
litellm_credential_name="team-a-chatgpt",
|
||||
mock_response="owned",
|
||||
**({} if secret_fields is None else {"secret_fields": secret_fields}),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 401
|
||||
assert "team-a-chatgpt" in str(raised.value)
|
||||
assert not any(token in str(raised.value) for token in _CHATGPT_SUBSCRIPTION_LOGIN_VALUES.values())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("through_router", [False, True], ids=["direct", "through_router"])
|
||||
async def test_acompletion_on_a_chatgpt_subscription_the_caller_owns_returns_the_response_with_no_host_login(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, through_router: bool
|
||||
) -> None:
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(litellm, "credential_list", [_chatgpt_subscription_credential()])
|
||||
request: Final = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"mock_response": "owned",
|
||||
"secret_fields": {"chatgpt_subscription_logins": RedactedDict({"team-a-chatgpt": "access-team-a-example"})},
|
||||
}
|
||||
deployment: Final = {"model": "chatgpt/gpt-5.5", "litellm_credential_name": "team-a-chatgpt"}
|
||||
router: Final = litellm.Router(model_list=[{"model_name": "team-a-gpt", "litellm_params": deployment}])
|
||||
response: Final = await (
|
||||
router.acompletion(model="team-a-gpt", **request)
|
||||
if through_router
|
||||
else litellm.acompletion(**deployment, **request)
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "owned"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("through_router", [False, True], ids=["direct", "through_router"])
|
||||
async def test_aresponses_on_a_chatgpt_subscription_the_caller_owns_returns_the_response_with_no_host_login(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, through_router: bool
|
||||
) -> None:
|
||||
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict
|
||||
|
||||
monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(litellm, "credential_list", [_chatgpt_subscription_credential()])
|
||||
request: Final = {
|
||||
"input": "hi",
|
||||
"mock_response": "owned",
|
||||
"secret_fields": {"chatgpt_subscription_logins": RedactedDict({"team-a-chatgpt": "access-team-a-example"})},
|
||||
}
|
||||
deployment: Final = {"model": "chatgpt/gpt-5.5", "litellm_credential_name": "team-a-chatgpt"}
|
||||
router: Final = litellm.Router(model_list=[{"model_name": "team-a-gpt", "litellm_params": deployment}])
|
||||
response: Final = await (
|
||||
router.aresponses(model="team-a-gpt", **request)
|
||||
if through_router
|
||||
else litellm.aresponses(**deployment, **request)
|
||||
)
|
||||
|
||||
assert response.output[0].content[0].text == "owned"
|
||||
|
||||
|
||||
_MISTRAL_CHAT_URL: Final = "https://api.mistral.ai/v1/chat/completions"
|
||||
_MISTRAL_REPLY: Final = {
|
||||
"id": "chatcmpl-oidc",
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.types.utils import (
|
|||
CustomPricingLiteLLMParams,
|
||||
MirroredPricingParams,
|
||||
anthropic_wif_litellm_params,
|
||||
chatgpt_subscription_login_litellm_params,
|
||||
github_copilot_oauth_litellm_params,
|
||||
oauth_token_exchange_litellm_params,
|
||||
openai_wif_litellm_params,
|
||||
|
|
@ -368,6 +369,7 @@ def test_server_owned_registry_includes_anthropic_openai_and_oauth_token_exchang
|
|||
+ oauth_token_exchange_litellm_params
|
||||
+ github_copilot_oauth_litellm_params
|
||||
+ ("chatgpt_auth_file",)
|
||||
+ chatgpt_subscription_login_litellm_params
|
||||
)
|
||||
assert set(openai_wif_litellm_params) == {
|
||||
"openai_identity_provider_id",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue