diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6eb13d2cba7..5341754132a 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -5,7 +5,10 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import ( + ANTHROPIC_WIF_KWARGS_KEYS, + AWS_CREDENTIAL_KWARGS_KEYS, +) from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch from litellm.types.utils import CallTypes, ModelInfo, Usage @@ -390,6 +393,9 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "max_retries", "_litellm_internal_model_credentials", *AWS_CREDENTIAL_KWARGS_KEYS, + # A federated deployment holds no api_key, so without these the fetch that reads a + # finished batch's output has nothing to authenticate with and its cost is never billed. + *sorted(ANTHROPIC_WIF_KWARGS_KEYS), ) for key in credential_keys: if key in litellm_params: diff --git a/litellm/llms/base_llm/auth/client_credentials.py b/litellm/llms/base_llm/auth/client_credentials.py index 7df03a334b2..d25dc86a535 100644 --- a/litellm/llms/base_llm/auth/client_credentials.py +++ b/litellm/llms/base_llm/auth/client_credentials.py @@ -24,6 +24,7 @@ from typing_extensions import assert_never from litellm.llms.base_llm.auth.identity_source import KeycloakSource, ref_for_error_message from litellm.llms.base_llm.auth.token_exchange import ( MAX_RESPONSE_BYTES, + endpoint_url_for_error_message, redact_oauth_error_body, validate_token_endpoint_url, ) @@ -146,7 +147,10 @@ def _resolve_client_secret(config: KeycloakSource, secret_reader: SecretReader) def _endpoint_error_message(config: KeycloakSource, response: httpx.Response, client_secret: str) -> str: endpoint_error: Final = redact_oauth_error_body(response.status_code, response.text, SecretStr(client_secret)) - return f"keycloak token endpoint {config.token_url} returned HTTP {endpoint_error.status_code}: {endpoint_error.redacted_body}" + return ( + f"keycloak token endpoint {endpoint_url_for_error_message(config.token_url)} " + f"returned HTTP {endpoint_error.status_code}: {endpoint_error.redacted_body}" + ) def _parse_success_body(response: httpx.Response) -> str: @@ -180,7 +184,10 @@ def fetch_keycloak_assertion( try: response: Final = poster.post(config.token_url, content=content, headers=headers, timeout=_TIMEOUT_SECONDS) except Exception as e: # noqa: BLE001 # injected posters may raise beyond httpx; every failure becomes a ValueError - raise ValueError(f"could not reach the keycloak token endpoint {config.token_url}: {type(e).__name__}") from e + raise ValueError( + f"could not reach the keycloak token endpoint {endpoint_url_for_error_message(config.token_url)}: " + f"{type(e).__name__}" + ) from e if not 200 <= response.status_code < 300: raise ValueError(_endpoint_error_message(config, response, client_secret)) return _parse_success_body(response) diff --git a/litellm/llms/base_llm/auth/token_exchange.py b/litellm/llms/base_llm/auth/token_exchange.py index 8518fe88905..73899d2b1d7 100644 --- a/litellm/llms/base_llm/auth/token_exchange.py +++ b/litellm/llms/base_llm/auth/token_exchange.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from math import inf from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias -from urllib.parse import urlencode, urlsplit +from urllib.parse import urlencode, urlsplit, urlunsplit import httpx from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError @@ -88,6 +88,19 @@ _RedactableBody: TypeAlias = Mapping[str, object] | list[object] | str | int | f _REDACTABLE_BODY_ADAPTER: Final = TypeAdapter[_RedactableBody](_RedactableBody) +def endpoint_url_for_error_message(url: str) -> str: + """``url`` reduced to scheme, host and path for operator-facing errors. + + A token endpoint is configuration, not a secret, and naming it is what makes these errors + actionable. But nothing stops an operator writing a credential into it, as a query parameter + or as userinfo, and these errors reach model callers, so neither part is echoed. + """ + parsed: Final = urlsplit(url) + host: Final = parsed.hostname or "" + authority: Final = f"{host}:{parsed.port}" if parsed.port is not None else host + return urlunsplit((parsed.scheme, authority, parsed.path, "", "")) + + def validate_token_endpoint_url(url: str) -> str | InsecureTokenUrl: parsed: Final = urlsplit(url) if parsed.scheme == "https": @@ -115,10 +128,18 @@ def _drop_reflected_assertion(rendered: str, assertion: SecretStr | None) -> str if assertion is None: return rendered secret: Final = assertion.get_secret_value() - probe: Final = secret[:_REFLECTION_PROBE_LENGTH] - if not probe or probe not in rendered: + if not secret: return rendered - return _REFLECTED_VALUE_MESSAGE + if len(secret) <= _REFLECTION_PROBE_LENGTH: + return _REFLECTED_VALUE_MESSAGE if secret in rendered else rendered + # Scan from the rendered side rather than probing the assertion's prefix: an endpoint that + # echoes the assertion from any offset (or only its tail) shares no prefix with it, but must + # still be caught. The rendered text is already capped, so this stays a handful of lookups. + windows: Final = ( + rendered[start : start + _REFLECTION_PROBE_LENGTH] + for start in range(len(rendered) - _REFLECTION_PROBE_LENGTH + 1) + ) + return _REFLECTED_VALUE_MESSAGE if any(window in secret for window in windows) else rendered def _redact_body_text(body_text: str) -> str: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index b64839423c2..ba5aba5d393 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -31,6 +31,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS +from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params from litellm.types.utils import CustomPricingLiteLLMParams, anthropic_wif_litellm_params @@ -190,20 +191,11 @@ def _allow_model_level_clientside_configurable_parameters( _ANTHROPIC_WIF_UNCONDITIONAL_BANNED: Final[tuple[str, ...]] = anthropic_wif_litellm_params # The Bedrock Claude Platform route reads a workspace from workspace_id or aws_workspace_id as # well, and neither is a federation parameter, so say so rather than leaving that caller stuck. -_BEDROCK_WORKSPACE_HINT: Final = " On the Bedrock Claude Platform route, pass workspace_id or aws_workspace_id instead." -def reject_server_owned_wif_params(body: Mapping[str, object]) -> None: - """Raise ``ValueError`` if a request-supplied mapping carries a server-owned workload-identity - federation field. These are never client-settable, on any surface, with or without a client-side - credential opt-in.""" - for param in _ANTHROPIC_WIF_UNCONDITIONAL_BANNED: - if param in body: - raise ValueError( - f"Rejected Request: {param} is a server-owned workload identity federation parameter " - "and cannot be set in a request body; configure it on the deployment instead." - + (_BEDROCK_WORKSPACE_HINT if param == "anthropic_workspace_id" else "") - ) +# Re-exported from litellm.types.router, where it lives so the router can call it on a +# post-authentication merge without core importing from the proxy package. +reject_server_owned_wif_params = _reject_server_owned_wif_params _NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body") diff --git a/litellm/proxy/common_utils/credential_hydration.py b/litellm/proxy/common_utils/credential_hydration.py index f9f62f01c69..446fb8544d7 100644 --- a/litellm/proxy/common_utils/credential_hydration.py +++ b/litellm/proxy/common_utils/credential_hydration.py @@ -5,6 +5,7 @@ Memory first (``litellm.credential_list``, already decrypted -- matching in-memory list has not yet picked up a credential another pod just wrote or updated. """ +from collections.abc import Mapping from types import MappingProxyType from typing import Final @@ -12,10 +13,43 @@ import litellm from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.utils import PrismaClient from litellm.repositories.credentials_repository import CredentialsRepository -from litellm.types.router import anthropic_wif_fields_named +from litellm.types.router import ( + GenericLiteLLMParams, + anthropic_wif_fields_named, + anthropic_wif_fields_present, +) from litellm.types.utils import CredentialItem +async def hydrate_named_credential_authoritative( + credential_name: str, + prisma_client: PrismaClient | None, +) -> CredentialItem | None: + """The stored credential, preferring the row over this pod's in-memory copy. + + ``hydrate_named_credential`` reads memory first, which is right when serving a request. A + management operation cannot: on a pod whose in-memory copy predates another pod's update, it + would export the superseded JWKS, or discover models against superseded values. Same reason + ``named_credential_wif_fields`` reads both. + """ + if prisma_client is None: + return await hydrate_named_credential(credential_name, prisma_client) + db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name) + if db_credential is None: + return await hydrate_named_credential(credential_name, prisma_client) + decrypted_values: Final = MappingProxyType( + { + key: decrypt_value_helper(value=value, key=key) or value + for key, value in db_credential.credential_values.items() + } + ) + return CredentialItem( + credential_name=db_credential.credential_name, + credential_values=decrypted_values, + credential_info=db_credential.credential_info, + ) + + async def hydrate_named_credential( credential_name: str, prisma_client: PrismaClient | None, @@ -63,3 +97,39 @@ async def named_credential_wif_fields( db_credential: Final = await CredentialsRepository(prisma_client).find_by_name(credential_name) stored: Final = () if db_credential is None else anthropic_wif_fields_named(db_credential.credential_values) return tuple(dict.fromkeys(in_memory + stored)) + + +async def effective_anthropic_wif_fields( + stored: Mapping[str, object] | None, + incoming: GenericLiteLLMParams | None, + prisma_client: PrismaClient | None, +) -> tuple[str, ...]: + """Federation field names the deployment would carry AFTER this write. + + Authorization has to read the resulting deployment, not the submitted payload. A patch that + names no federation field still lands on a deployment that has them, and a patch that only + attaches ``litellm_credential_name`` inherits whatever that credential holds. + + The two sides are matched differently on purpose. ``stored`` is matched by VALUE, because + ``GenericLiteLLMParams`` declares every ``anthropic_*`` field, so matching it by key would + report every deployment on the proxy as federated. ``incoming`` is matched by the keys the + write actually set, so an explicit null still counts as touching the field. + """ + from_stored: Final = () if stored is None else anthropic_wif_fields_present(stored) + from_incoming: Final = () if incoming is None else anthropic_wif_fields_named(incoming.model_fields_set) + credential_name: Final = _effective_credential_name(stored, incoming) + from_credential: Final = ( + () if credential_name is None else await named_credential_wif_fields(credential_name, prisma_client) + ) + return tuple(dict.fromkeys(from_stored + from_incoming + from_credential)) + + +def _effective_credential_name( + stored: Mapping[str, object] | None, + incoming: GenericLiteLLMParams | None, +) -> str | None: + if incoming is not None and "litellm_credential_name" in incoming.model_fields_set: + named: Final = incoming.litellm_credential_name + return named if isinstance(named, str) else None + from_stored: Final = None if stored is None else stored.get("litellm_credential_name") + return from_stored if isinstance(from_stored, str) else None diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 45cae85513c..357375ddea0 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKey from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.credential_hydration import ( hydrate_named_credential, + hydrate_named_credential_authoritative, named_credential_wif_fields, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper @@ -292,7 +293,7 @@ async def get_credential_internal_issuer_jwks( ) try: - credential: Final = await hydrate_named_credential(credential_name, prisma_client) + credential: Final = await hydrate_named_credential_authoritative(credential_name, prisma_client) if credential is None or credential.credential_info.get("custom_llm_provider") != "anthropic": raise HTTPException( status_code=404, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index ca989dfc2ff..7af079dfa8b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1963,6 +1963,9 @@ async def test_model_connection( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + # The Deployment above already carries the caller's merged params, so the effective + # state is model_params itself; there is no separate incoming patch here. + incoming_params=None, ) # Include health_check_params if provided litellm_params = _update_litellm_params_for_health_check( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index dfdabc326cd..52e4bb73d11 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -55,7 +55,11 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) -from litellm.proxy.common_utils.credential_hydration import hydrate_named_credential +from litellm.proxy.common_utils.credential_hydration import ( + effective_anthropic_wif_fields, + hydrate_named_credential, + hydrate_named_credential_authoritative, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin @@ -100,7 +104,6 @@ from litellm.types.router import ( Deployment, GenericLiteLLMParams, ModelInfo, - anthropic_wif_fields_present, updateDeployment, ) from litellm.types.utils import LlmProviders @@ -254,30 +257,6 @@ def _raise_on_strategy_router_write_violation( ) -def _reject_non_admin_wif_persistence( - litellm_params: GenericLiteLLMParams | None, - user_api_key_dict: UserAPIKeyAuth, -) -> None: - """Anthropic workload identity federation fields choose which server-side secret is read - and where it is sent. Only proxy admins may persist them on a deployment, mirroring the - ``blocked``-flag gate below: a team admin who otherwise manages a team-scoped deployment - must not be able to set these. - """ - if litellm_params is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: - return - wif_fields: Final = anthropic_wif_fields_present(litellm_params.model_dump(exclude_none=True)) - if not wif_fields: - return - raise ProxyException( - message=( - f"Only proxy admins can set {wif_fields[0]!r}, a server-owned workload identity federation parameter." - ), - type=ProxyErrorTypes.auth_error.value, - code=status.HTTP_403_FORBIDDEN, - param=wif_fields[0], - ) - - def _reject_non_admin_blocked_flag_on_create( blocked: bool | None, user_api_key_dict: UserAPIKeyAuth, @@ -688,6 +667,7 @@ async def patch_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=patch_data.litellm_params, ) # Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins @@ -701,8 +681,6 @@ async def patch_model( param="blocked", ) - _reject_non_admin_wif_persistence(patch_data.litellm_params, user_api_key_dict) - _raise_on_strategy_router_write_violation( incoming_params=patch_data.litellm_params, existing_params=db_model.litellm_params, @@ -1493,14 +1471,65 @@ class ModelManagementAuthChecks: ) return True + @staticmethod + async def _reject_non_admin_wif_write( + *, + model_params: Deployment, + incoming_params: GenericLiteLLMParams | None, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + ) -> None: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + stored: Final = model_params.litellm_params.model_dump(exclude_none=True) + wif_fields: Final = await effective_anthropic_wif_fields(stored, incoming_params, prisma_client) + if wif_fields: + # ProxyException rather than HTTPException so the offending field stays a structured + # `param`, which is the contract the narrower gate this replaced already published. + raise ProxyException( + message=( + f"Only proxy admins can modify a deployment configured for Anthropic workload identity " + f"federation ({wif_fields[0]!r})." + ), + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param=wif_fields[0], + ) + # A name the caller expects an admin to create later would resolve to nothing today and + # start federating the moment it exists, so a non-admin may only attach one that is already there. + if incoming_params is not None and "litellm_credential_name" in incoming_params.model_fields_set: + named: Final = incoming_params.litellm_credential_name + if isinstance(named, str) and await hydrate_named_credential(named, prisma_client) is None: + raise ProxyException( + message=f"No credential named {named!r} exists.", + type=ProxyErrorTypes.bad_request_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + @staticmethod async def can_user_make_model_call( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, premium_user: bool, + *, + incoming_params: GenericLiteLLMParams | None, allow_missing_team: bool = False, ) -> Literal[True]: + # Federation fields choose which server-side secret is read and where the org-scoped token + # it buys is sent, so only a proxy admin may touch a deployment that has them. Evaluated on + # the RESULTING deployment: a patch naming no federation field still lands on one that has + # them, and a patch attaching a credential by name inherits whatever that credential holds. + # `incoming_params` is keyword-only with no default so a new write path cannot typecheck + # without deciding what it writes. + await ModelManagementAuthChecks._reject_non_admin_wif_write( + model_params=model_params, + incoming_params=incoming_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( @@ -1595,6 +1624,7 @@ async def delete_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=None, allow_missing_team=True, ) @@ -1725,7 +1755,7 @@ async def _resolve_discovery_litellm_params( if data.litellm_credential_name is None: return MappingProxyType({k: v for k, v in (("api_key", data.api_key), ("api_base", data.api_base)) if v}) - credential: Final = await hydrate_named_credential(data.litellm_credential_name, prisma_client) + credential: Final = await hydrate_named_credential_authoritative(data.litellm_credential_name, prisma_client) if credential is None: raise HTTPException( status_code=404, @@ -1887,12 +1917,11 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=model_params.litellm_params, ) _reject_non_admin_blocked_flag_on_create(model_params.blocked, user_api_key_dict) - _reject_non_admin_wif_persistence(model_params.litellm_params, user_api_key_dict) - _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=None, @@ -2058,6 +2087,7 @@ async def update_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + incoming_params=model_params.litellm_params, ) _raise_on_strategy_router_write_violation( @@ -2072,8 +2102,6 @@ async def update_model( if model_params.litellm_params is None: raise Exception("litellm_params not provided") - _reject_non_admin_wif_persistence(model_params.litellm_params, user_api_key_dict) - _new_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) ### ENCRYPT PARAMS ### diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3c9a4097321..3e342e269d9 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -23,7 +23,7 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, ) -from litellm.types.router import LiteLLMParamsTypedDict +from litellm.types.router import LiteLLMParamsTypedDict, reject_server_owned_wif_params if TYPE_CHECKING: from litellm.router import Router as _Router @@ -347,6 +347,14 @@ async def run_async_fallback( ) attempted.record(original_model_group) + # A dict target is merged straight into kwargs below, and kwargs win over the deployment's own + # params, so a stored key/team/global fallback could otherwise set a federation field that the + # request itself is forbidden to carry. Checked here rather than at the merge: inside the loop + # the refusal would be caught as a per-target failure and quietly skipped to the next one. + for target in fallback_model_group: + if isinstance(target, dict): + reject_server_owned_wif_params(target) + for mg in fallback_model_group: if mg == original_model_group: continue diff --git a/litellm/types/router.py b/litellm/types/router.py index c812c89c25c..3996b50d687 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1069,3 +1069,23 @@ class AdaptiveRouterPreferences(BaseModel): quality_tier: int = Field(ge=1, le=3) strengths: list[RequestType] = Field(default_factory=list) + + +_BEDROCK_WORKSPACE_HINT: Final = " On the Bedrock Claude Platform route, pass workspace_id or aws_workspace_id instead." + + +def reject_server_owned_wif_params(body: Mapping[str, object]) -> None: + """Raise ``ValueError`` if a mapping that did not come from deployment config carries a + server-owned workload identity federation field. + + These are never client-settable on any surface, with or without a client-side credential + opt-in. This lives here rather than under ``litellm.proxy`` so the router can call it on a + post-authentication merge without core importing from the proxy package. + """ + for param in _anthropic_wif_litellm_params: + if param in body: + raise ValueError( + f"Rejected Request: {param} is a server-owned workload identity federation parameter " + "and cannot be set in a request body; configure it on the deployment instead." + + (_BEDROCK_WORKSPACE_HINT if param == "anthropic_workspace_id" else "") + ) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 41b4bb8cf76..5e4011f385a 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1048,6 +1048,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): result set - zero cost, zero usage, no models - instead of letting the file fetch raise "Output file id is None" on every aretrieve_batch logging poll. """ + # The output-file fetch must not even be attempted when there is no output file. async def _must_not_fetch(*args, **kwargs): pytest.fail("_fetch_batch_output_file_content should not be called") @@ -1166,7 +1167,10 @@ def test_anthropic_response_body_is_result_message(): def test_anthropic_usage_conversion_includes_cache_tokens(): - body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)} + body = { + "model": "claude-sonnet-4-5-20250929", + "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000), + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic") assert usage.prompt_tokens == 11000 assert usage.completion_tokens == 200 @@ -1181,7 +1185,9 @@ def test_bedrock_model_output_line_success_check(): "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True - assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + assert ( + bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + ) def test_bedrock_cost_uses_deployment_model_name(): @@ -1233,7 +1239,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): rows = [ { "custom_id": "req-1", - "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, + "response": { + "status_code": 200, + "body": { + "model": "gpt-5.2", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, } ] _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") @@ -1479,7 +1491,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> def test_bedrock_converse_shaped_batch_usage_is_parsed(): - body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}} + body = { + "model": "us.amazon.nova-lite-v1:0", + "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}, + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742) @@ -1523,3 +1538,37 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): assert usage.total_tokens == 0 assert "does not understand" in caplog.text assert "inputTextTokenCount" in caplog.text + + +class TestFileAccessCredentialsCarryFederation: + """A federated deployment holds no api_key, so the fetch that reads a finished batch's output + has to inherit the federation fields or it cannot authenticate and the batch is never billed.""" + + def test_federation_fields_survive_extraction(self): + from litellm.batches.batch_utils import _extract_file_access_credentials + + credentials = _extract_file_access_credentials( + { + "model": "anthropic/claude-sonnet-4-5", + "anthropic_federation_rule_id": "fdrl_x", + "anthropic_organization_id": "org-x", + "anthropic_identity_token_file": "/var/run/secrets/anthropic.com/token", + "something_unrelated": "dropped", + } + ) + + assert credentials["anthropic_federation_rule_id"] == "fdrl_x" + assert credentials["anthropic_organization_id"] == "org-x" + assert credentials["anthropic_identity_token_file"] == "/var/run/secrets/anthropic.com/token" + assert "something_unrelated" not in credentials + + def test_every_federation_field_is_carried(self): + """Derived from the kwargs set, so a new federation field is carried without an edit here.""" + from litellm.batches.batch_utils import _extract_file_access_credentials + from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS + + params = {name: f"value-{name}" for name in ANTHROPIC_WIF_KWARGS_KEYS} + + credentials = _extract_file_access_credentials(params) + + assert set(credentials) == set(ANTHROPIC_WIF_KWARGS_KEYS) diff --git a/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py b/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py index e006aa1a12a..76c309cae37 100644 --- a/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py +++ b/tests/test_litellm/llms/base_llm/auth/test_client_credentials.py @@ -429,7 +429,7 @@ class TestUnresolvedSecretRefIsNotEchoed: client_secret_ref=pasted_secret, ) - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match="could not be read") as excinfo: keycloak_assertion_source(config, secret_reader=lambda _ref: None)() assert pasted_secret not in str(excinfo.value) @@ -446,8 +446,39 @@ class TestUnresolvedSecretRefIsNotEchoed: signing_key_ref=pasted_pem, ) - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match="could not be read") as excinfo: internal_issuer_assertion_source(config, key_reader=lambda _ref: None)() assert pasted_pem not in str(excinfo.value) assert "withheld" in str(excinfo.value) + + +class TestTokenUrlIsNotEchoedWholesale: + """A token endpoint is configuration and naming it makes the error actionable, but nothing + stops an operator putting a credential in the URL, and these errors reach model callers.""" + + def test_query_string_is_dropped_from_a_status_error(self): + from litellm.llms.base_llm.auth.token_exchange import endpoint_url_for_error_message + + rendered = endpoint_url_for_error_message("https://idp.example/token?client_secret=supersecret") + + assert "supersecret" not in rendered + assert rendered == "https://idp.example/token" + + def test_userinfo_is_dropped_too(self): + from litellm.llms.base_llm.auth.token_exchange import endpoint_url_for_error_message + + rendered = endpoint_url_for_error_message("https://user:pw@idp.example:8443/token") + + assert "pw" not in rendered + assert rendered == "https://idp.example:8443/token" + + def test_transport_failure_message_carries_no_query_secret(self): + poster = RaisingPoster(httpx.ConnectTimeout("timed out")) + config = make_config(token_url="https://idp.example/token?client_secret=supersecret") + + with pytest.raises(ValueError, match="could not reach the keycloak token endpoint") as excinfo: + fetch_keycloak_assertion(config, poster=poster, secret_reader=DEFAULT_SECRET_READER) + + assert "supersecret" not in str(excinfo.value) + assert "idp.example/token" in str(excinfo.value) diff --git a/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py b/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py index 697997a5faa..344e87111a1 100644 --- a/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py +++ b/tests/test_litellm/llms/base_llm/auth/test_token_exchange.py @@ -591,6 +591,37 @@ class TestRedactionAndCaps: assert assertion.get_secret_value() not in result.redacted_body assert "REFLECTEDPAYLOAD" not in result.redacted_body + def test_assertion_reflected_from_an_offset_is_dropped(self): + """Regression: the probe only looked at the assertion's first 24 characters, so an + endpoint echoing it from any later offset shared no prefix and slipped through.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "A" * 40 + "PAYLOADMIDDLE" + "B" * 40 + ".signature") + tail = assertion.get_secret_value()[24:] + body = {"error": "invalid_grant", "error_description": tail} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert "PAYLOADMIDDLE" not in result.redacted_body + assert tail[:40] not in result.redacted_body + + def test_a_short_secret_is_still_matched_whole(self): + """A Keycloak client secret can be shorter than the probe length; the whole value is + compared in that case rather than a truncated prefix.""" + secret = SecretStr("short-secret") + body = {"error": "invalid_client", "error_description": "rejected short-secret"} + + result = redact_oauth_error_body(400, json.dumps(body), secret) + + assert "short-secret" not in result.redacted_body + + def test_an_unrelated_body_is_not_falsely_redacted(self): + """The scan must not fire on a body that merely shares short runs with the assertion.""" + assertion = SecretStr("eyJhbGciOiJSUzI1NiJ9." + "Z" * 60 + ".signature") + body = {"error": "invalid_grant", "error_description": "the federation rule was not found"} + + result = redact_oauth_error_body(400, json.dumps(body), assertion) + + assert "the federation rule was not found" in result.redacted_body + def test_json_array_body_constant_message(self): result = redact_oauth_error_body(400, json.dumps(["a", "b"])) assert result.redacted_body == "non-object error response omitted" diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py index 2c87a4295ea..2f9f3550b77 100644 --- a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -941,3 +941,59 @@ class TestNonAdminCannotTouchAStoredWifCredential: assert response.status_code == 200, response.text repository.create.assert_awaited_once() assert litellm.credential_list[0].credential_name == "ordinary-cred" + + +class TestManagementReadsTheStoredCredential: + """Serving a request reads memory first, which is right. A management operation cannot: on a + pod whose in-memory copy predates another pod's update it would act on superseded values.""" + + @pytest.mark.asyncio + async def test_authoritative_hydrate_prefers_the_row_over_a_stale_memory_copy(self): + import litellm + from litellm.proxy.common_utils.credential_hydration import ( + hydrate_named_credential, + hydrate_named_credential_authoritative, + ) + from litellm.types.utils import CredentialItem + + stale = CredentialItem( + credential_name="anthropic-wif", + credential_values={"anthropic_issuer_url": "https://old.example.com"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + row = MagicMock() + row.dict.return_value = { + "credential_name": "anthropic-wif", + "credential_values": {"anthropic_issuer_url": "https://new.example.com"}, + "credential_info": {"custom_llm_provider": "anthropic"}, + } + + prisma = MagicMock() + prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=row) + + with patch.object(litellm, "credential_list", [stale]): + served = await hydrate_named_credential("anthropic-wif", prisma) + managed = await hydrate_named_credential_authoritative("anthropic-wif", prisma) + + assert served is not None and served.credential_values["anthropic_issuer_url"] == "https://old.example.com" + assert managed is not None and managed.credential_values["anthropic_issuer_url"] == "https://new.example.com" + + @pytest.mark.asyncio + async def test_authoritative_hydrate_falls_back_to_memory_when_the_row_is_absent(self): + import litellm + from litellm.proxy.common_utils.credential_hydration import hydrate_named_credential_authoritative + from litellm.types.utils import CredentialItem + + only_in_memory = CredentialItem( + credential_name="config-yaml-credential", + credential_values={"anthropic_issuer_url": "https://configured.example.com"}, + credential_info={"custom_llm_provider": "anthropic"}, + ) + prisma = MagicMock() + prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=None) + + with patch.object(litellm, "credential_list", [only_in_memory]): + resolved = await hydrate_named_credential_authoritative("config-yaml-credential", prisma) + + assert resolved is not None + assert resolved.credential_values["anthropic_issuer_url"] == "https://configured.example.com" diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 4438b4eed69..991cce82578 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -231,6 +231,7 @@ class TestModelManagementAuthChecks: user_api_key_dict=self.admin_user, prisma_client=prisma_client, premium_user=True, + incoming_params=None, ) assert result is True @@ -252,6 +253,7 @@ class TestModelManagementAuthChecks: user_api_key_dict=self.normal_user, prisma_client=prisma_client, premium_user=True, + incoming_params=None, ) assert "403" in str(exc_info.value) @@ -4254,12 +4256,10 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "litellm.proxy.proxy_server.premium_user", True, ), - patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", - new=AsyncMock(return_value=None), - ), ): - with pytest.raises(Exception, match="Only proxy admins can set") as exc_info: + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for Anthropic" + ) as exc_info: await patch_model( model_id="m1", patch_data=updateDeployment( @@ -4312,10 +4312,6 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "litellm.proxy.proxy_server.premium_user", True, ), - patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", - new=AsyncMock(return_value=None), - ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value: value, @@ -4360,10 +4356,6 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "litellm.proxy.proxy_server.premium_user", True, ), - patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", - new=AsyncMock(return_value=None), - ), ): with pytest.raises(ProxyException) as exc_info: await add_new_model( @@ -4415,10 +4407,6 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["wif-gate-create-1"]}), ), - patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", - new=AsyncMock(return_value=None), - ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.proxy_server.proxy_config", MagicMock(add_deployment=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None))), @@ -4473,12 +4461,10 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "litellm.proxy.proxy_server.premium_user", True, ), - patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", - new=AsyncMock(return_value=None), - ), ): - with pytest.raises(Exception, match="Only proxy admins can set") as exc_info: + with pytest.raises( + Exception, match="Only proxy admins can modify a deployment configured for Anthropic" + ) as exc_info: await update_model( model_params=updateDeployment( litellm_params=updateLiteLLMParams( @@ -4531,10 +4517,6 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "litellm.proxy.proxy_server.premium_user", True, ), - patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", - new=AsyncMock(return_value=None), - ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value: value, @@ -4561,6 +4543,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel: assert "os.environ/ANTHROPIC_WIF_CLIENT_SECRET" in written_litellm_params +def _prisma_without_stored_credentials() -> MagicMock: + """A prisma double whose credential row is absent, so credential resolution falls back to + the in-memory list. Management reads consult the row first, which is the multi-pod case.""" + prisma = MagicMock() + prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=None) + return prisma + + class TestDiscoverProviderModels: """POST /provider/models/discover: proxy-admin-only, credential-name-only contract for server-owned auth (WIF), never a silent [] on failure.""" @@ -4665,7 +4655,7 @@ class TestDiscoverProviderModels: ], ) with patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.proxy_server.prisma_client", MagicMock() + "litellm.proxy.proxy_server.prisma_client", _prisma_without_stored_credentials() ): # test-quality-ok: the proxy wiring under test is what this patches with pytest.raises(HTTPException) as exc_info: await discover_provider_models( @@ -4727,7 +4717,7 @@ class TestDiscoverProviderModels: ) with ( patch( # test-quality-ok: the proxy wiring under test is what this patches - "litellm.proxy.proxy_server.prisma_client", MagicMock() + "litellm.proxy.proxy_server.prisma_client", _prisma_without_stored_credentials() ), patch( # test-quality-ok: the proxy wiring under test is what this patches "litellm.llms.anthropic.common_utils.AnthropicModelInfo.discover_models", @@ -4817,3 +4807,98 @@ class TestOneCredentialFeedsManyModelsNoWifCopy: assert "anthropic_federation_rule_id" not in written_litellm_params assert "anthropic_identity_token" not in written_litellm_params assert call.kwargs["data"]["blocked"] is False + + +class TestWifBoundaryReadsTheResultingDeployment: + """The proxy-admin rule has to be evaluated against the deployment the write PRODUCES. + Reading only the submitted payload let a team admin keep an existing federated deployment + and change it anyway, because the fields they sent named nothing federated.""" + + @staticmethod + def _existing_wif_row(): + row = MagicMock() + row.litellm_params = { + "model": "anthropic/claude-sonnet-4", + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + } + row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": row.litellm_params, + "model_info": {"id": "m1"}, + } + return row + + @pytest.mark.asyncio + async def test_non_admin_cannot_retarget_an_existing_wif_deployment_via_api_base(self): + """api_base is not a federation field, so the payload-only check saw nothing to refuse, + and the merged deployment then sent its assertion and minted token to the new host.""" + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=self._existing_wif_row()) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy wiring under test + patch( # test-quality-ok: proxy wiring under test + "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]}) + ), + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test + ): + with pytest.raises(Exception, match="Only proxy admins can modify a deployment configured for Anthropic"): + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(api_base="https://gateway.internal") + ), + user_api_key_dict=non_admin, + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_attach_a_federated_credential_by_name(self): + """litellm_credential_name names no federation field itself, but request-time hydration + imports whatever the credential holds, so the resulting deployment federates.""" + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + non_admin = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER) + plain_row = MagicMock() + plain_row.litellm_params = {"model": "anthropic/claude-sonnet-4"} + plain_row.model_dump.return_value = { + "model_name": "claude", + "litellm_params": plain_row.litellm_params, + "model_info": {"id": "m1"}, + } + # The credential is served from the row rather than this pod's memory, which is both the + # multi-pod case and the one the gate must not miss. + admin_credential_row = MagicMock() + admin_credential_row.dict.return_value = { + "credential_name": "admin-wif", + "credential_values": { + "anthropic_federation_rule_id": "fdrl_admin", + "anthropic_organization_id": "org-admin", + }, + "credential_info": {"custom_llm_provider": "anthropic"}, + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=plain_row) + mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=admin_credential_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy wiring under test + patch( # test-quality-ok: proxy wiring under test + "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]}) + ), + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test + ): + with pytest.raises(Exception, match="Only proxy admins can modify a deployment configured for Anthropic"): + await patch_model( + model_id="m1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name="admin-wif") + ), + user_api_key_dict=non_admin, + ) diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 24477248a8a..4c707755221 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -484,9 +484,7 @@ async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bar with pytest.raises(RuntimeError, match="fallback model also failed"): await run_async_fallback( litellm_router=router, - fallback_model_group=[ - {"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]} - ], + fallback_model_group=[{"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]}], original_model_group="primary-model", original_exception=RuntimeError("original failed"), max_fallbacks=3, @@ -843,3 +841,39 @@ class TestRunAsyncFallbackTriggersCooldown: ) mock_trigger.assert_not_called() + + +@pytest.mark.asyncio +async def test_a_stored_fallback_target_cannot_carry_a_federation_field(): + """A dict fallback target is merged into kwargs, and kwargs beat the deployment's own params, + so a stored key/team/global fallback could otherwise set the workspace a federation token is + minted for. The request itself is already forbidden to carry these, and a stored setting is + not a more trusted source than the request.""" + with pytest.raises(ValueError, match="server-owned workload identity federation parameter"): + await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=[{"model": "anthropic-backup", "anthropic_workspace_id": "wrkspc_other"}], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + +@pytest.mark.asyncio +async def test_the_refusal_is_not_swallowed_as_a_fallback_error(): + """Checked before the per-target loop on purpose: inside it, the refusal would be caught as + that target's failure and the run would quietly continue to the next one.""" + with pytest.raises(ValueError, match="anthropic_issuer_signing_key_ref"): + await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=[ + {"model": "anthropic-backup", "anthropic_issuer_signing_key_ref": "os.environ/ADMIN_KEY"}, + "a-perfectly-fine-model", + ], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + )