diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 4cfc6d4e081..abc2487a4a5 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -20,6 +20,7 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs from litellm.constants import ( BATCH_LINE_ITEMS_MAX_PART_BYTES, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, @@ -765,7 +766,7 @@ class CheckBatchCost: credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {} _file_content = await afile_content( file_id=cast(str, raw_output_file_id), - **cast(FileContentCallOptions, credentials), + **cast(FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(cast(str, raw_output_file_id))}), ) content_bytes: Final = await _file_content_bytes(_file_content) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 626a3b0bb83..bbb09e75b47 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -49,6 +49,7 @@ from litellm.constants import ( ) from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.internal_call_metadata import requests_background @@ -396,6 +397,7 @@ class _CursorPageArgs(TypedDict, total=False): class _RouterFileCallKwargs(TypedDict, total=False): client: ReadOnly[AsyncOpenAI | None] custom_llm_provider: ReadOnly[str | None] + stored_s3_bucket_name: ReadOnly[str] _MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str]) @@ -463,7 +465,7 @@ async def _delete_provider_file( await llm_router.afile_delete( model=model_id, file_id=provider_file_id, - **cast(_RouterFileCallKwargs, delete_data), + **cast(_RouterFileCallKwargs, {**delete_data, **stored_s3_bucket_kwargs(provider_file_id)}), ) @@ -1157,7 +1159,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve) retrieve_options: Final = cast( FileRetrieveCallOptions, - {**call_options, "max_retries": 0}, + {**call_options, **stored_s3_bucket_kwargs(provider_file_id), "max_retries": 0}, ) for attempt in range(len(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS) + 1): if attempt > 0: @@ -1795,8 +1797,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: + provider_file_id: Final = self.get_output_file_id_from_unified_file_id(potential_file_id) data["model"] = model_id - data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) + data["file_id"] = provider_file_id + data.update(stored_s3_bucket_kwargs(provider_file_id)) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -2894,7 +2898,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await llm_router.afile_content( model=model_id, file_id=provider_file_id, - **cast(_RouterFileCallKwargs, data), + **cast( + _RouterFileCallKwargs, + { + **data, + **( + stored_s3_bucket_kwargs(provider_file_id) + if isinstance(provider_file_id, str) + else {} + ), + }, + ), ), ) except Exception as e: diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index ca072368061..cdae3de47c8 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -11,6 +11,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils import batch_utils as core_batch_utils +from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo @@ -551,7 +552,7 @@ async def fetch_batch_managed_file_content( "file_id": provider_output_file_id, "custom_llm_provider": custom_llm_provider, **cast( # cast-ok: preserve dynamic provider credentials without validation - FileContentCallOptions, credentials + FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(provider_output_file_id)} ), } diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 477775bb423..d78fbd58a57 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -17,6 +17,28 @@ BEDROCK_MANAGED_S3_PREFIXES: Final = ( BEDROCK_MANAGED_S3_OUTPUT_PREFIX, ) MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://") +STORED_S3_BUCKET_PARAM: Final = "stored_s3_bucket_name" +REQUEST_S3_BUCKET_PARAMS: Final = frozenset({"s3_bucket_name", "s3_output_bucket_name"}) + + +def stored_s3_bucket_kwargs(stored_file_id: str) -> Mapping[str, str]: + """The bucket (and any prefix before LiteLLM's own path) of an ``s3://`` file id LiteLLM recorded itself, + as the param that lets Bedrock read it back. + + Only pass ids from LiteLLM's own records (managed file rows, provider batch responses), never one a caller sent. + """ + if not stored_file_id.startswith("s3://"): + return {} + bucket, _, object_key = stored_file_id.removeprefix("s3://").partition("/") + if not bucket: + return {} + segments: Final = object_key.split("/") + prefix_depth: Final = next( + (depth for depth in range(len(segments)) if "/".join(segments[depth:]).startswith(BEDROCK_MANAGED_S3_PREFIXES)), + 0, + ) + prefix: Final = "/".join(segments[:prefix_depth]) + return {STORED_S3_BUCKET_PARAM: f"{bucket}/{prefix}" if prefix else bucket} MAX_FILE_ID_DECODE_PASSES: Final = 8 diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 41c1baffab8..bbbbe9635c9 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -62,6 +62,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "bucket_name", "s3_bucket_name", "s3_output_bucket_name", + "stored_s3_bucket_name", "allow_legacy_cloud_file_ids", "s3_endpoint_url", "s3_region_name", diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 476d1395321..b19ad776051 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_OUTPUT_PREFIX, BEDROCK_MANAGED_S3_PREFIXES, BEDROCK_MANAGED_S3_UPLOAD_PREFIX, + STORED_S3_BUCKET_PARAM, build_managed_cloud_object_name, encode_s3_object_key_for_url, sanitize_cloud_object_component, @@ -313,11 +314,19 @@ def get_configured_s3_bucket_names(litellm_params: Mapping[str, object]) -> tupl Resolve the S3 buckets a Bedrock file id may live in. Bedrock batch outputs land in ``s3_output_bucket_name`` when it differs from - the input bucket, so retrieval validates against both. + the input bucket, so retrieval validates against both, plus the bucket LiteLLM + recorded for a managed file uploaded or written to a request-chosen bucket. """ input_bucket: Final = _input_bucket_name(litellm_params) output_bucket: Final = _output_bucket_name(litellm_params) - buckets: Final = tuple(dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket)) + stored_bucket: Final = litellm_params.get(STORED_S3_BUCKET_PARAM) + buckets: Final = tuple( + dict.fromkeys( + bucket + for bucket in (input_bucket, output_bucket, stored_bucket if isinstance(stored_bucket, str) else None) + if bucket + ) + ) if not buckets: raise ValueError(_S3_BUCKET_REQUIRED_ERROR) return buckets diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d3978a58c43..715b98a441a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -3,7 +3,7 @@ import os import re import sys from collections.abc import Collection, Iterator, Mapping, Sequence -from functools import lru_cache +from functools import lru_cache, partial from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType @@ -22,6 +22,7 @@ from litellm.constants import ( MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS, ) +from litellm.litellm_core_utils.cloud_storage_security import REQUEST_S3_BUCKET_PARAMS, STORED_S3_BUCKET_PARAM from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( @@ -444,6 +445,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( "s3_endpoint_url", "s3_bucket_name", "s3_output_bucket_name", + STORED_S3_BUCKET_PARAM, "allow_legacy_cloud_file_ids", "sagemaker_base_url", "deployment_url", @@ -494,6 +496,9 @@ _PROXY_ALLOWLIST_SETTING_BY_PARAM: Final[Mapping[str, _ClientSideAllowlistSettin } ) +TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY: Final = "allowed_s3_buckets" +_BUCKET_NAMES: Final = TypeAdapter(tuple[str, ...]) + class _ClientSideStorageAllowlist(BaseModel): model_config = ConfigDict(frozen=True) @@ -511,8 +516,15 @@ class _ClientSideStorageAllowlist(BaseModel): def _banned_param_error(param: str) -> ValueError: setting: Final = _PROXY_ALLOWLIST_SETTING_BY_PARAM.get(param) + team_hint: Final = ( + f" or in the team's `metadata.{TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY}`" + if param in REQUEST_S3_BUCKET_PARAMS + else "" + ) allowlist_hint: Final = ( - "" if setting is None else f" To allow specific values only, list them in `general_settings.{setting}`." + "" + if setting is None + else f" To allow specific values only, list them in `general_settings.{setting}`{team_hint}." ) return ValueError( f"Rejected Request: {param} is not allowed in request body. " @@ -532,11 +544,13 @@ def _check_banned_params( *, manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, + defer_s3_buckets: bool = False, ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. Shared between the root-level check and the nested-config check so a - new banned param only needs to be added in one place. + new banned param only needs to be added in one place. ``defer_s3_buckets`` lets + bucket params through for handlers that check them against the caller team's allowlist. """ wif_body: Final[Mapping[str, object]] = ( MappingProxyType({key: value for key, value in body.items() if key not in oauth_token_exchange_litellm_params}) @@ -557,6 +571,8 @@ def _check_banned_params( setting := _PROXY_ALLOWLIST_SETTING_BY_PARAM.get(param) ) is not None and _ClientSideStorageAllowlist.model_validate(general_settings).admits(setting, body[param]): continue + if defer_s3_buckets and param in REQUEST_S3_BUCKET_PARAMS: + continue if ( _allow_model_level_clientside_configurable_parameters( model=model, @@ -704,6 +720,7 @@ def is_request_body_safe( model: str, *, route: str | None = None, + defer_s3_buckets: bool = False, ) -> bool: """ Check if the request body is safe. @@ -733,33 +750,28 @@ def is_request_body_safe( if "model_list" in request_body: raise ValueError("Rejected Request: model_list is not allowed in the request body.") manages_deployments: Final = route_manages_deployments(route) - _check_banned_params( - request_body, - general_settings, - llm_router, - model, + check_banned_params: Final = partial( + _check_banned_params, + general_settings=general_settings, + llm_router=llm_router, + model=model, manages_deployments=manages_deployments, - mcp_server_fields_allowed=route_manages_mcp_servers(route), + defer_s3_buckets=defer_s3_buckets, ) + check_banned_params(request_body, mcp_server_fields_allowed=route_manages_mcp_servers(route)) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) if nested is not None: - _check_banned_params(nested, general_settings, llm_router, model, manages_deployments=manages_deployments) + check_banned_params(nested) for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) if metadata is not None: - _check_banned_params(metadata, general_settings, llm_router, model, manages_deployments=manages_deployments) + check_banned_params(metadata) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): - _check_banned_params( - extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["), - general_settings, - llm_router, - model, - manages_deployments=manages_deployments, - ) + check_banned_params(extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}[")) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): - _check_banned_params(target, general_settings, llm_router, model, manages_deployments=manages_deployments) + check_banned_params(target) target_model = target.get("model") if isinstance(target_model, str): _reject_url_valued_fallback_target(target_model) @@ -772,9 +784,7 @@ def is_request_body_safe( _reject_request_credential_params(litellm_params, manages_deployments=manages_deployments) litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata")) if litellm_params_metadata is not None: - _check_banned_params( - litellm_params_metadata, general_settings, llm_router, model, manages_deployments=manages_deployments - ) + check_banned_params(litellm_params_metadata) for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict reject_api_key_secret_reference(user_config_params) return True @@ -846,6 +856,48 @@ def _coerce_metadata_to_dict(value: object) -> dict[str, object] | None: return None +def _team_s3_buckets(team_metadata: Mapping[str, object] | None) -> tuple[str, ...]: + if team_metadata is None: + return () + try: + return _BUCKET_NAMES.validate_python(team_metadata.get(TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY) or ()) + except ValidationError: + return () + + +_TEAM_S3_BUCKET_ROUTE_TEMPLATES: Final = frozenset( + {"/v1/files", "/files", "/{provider}/v1/files", "/v1/batches", "/batches", "/{provider}/v1/batches"} +) + + +def checks_team_s3_buckets_in_handler(request: Request) -> bool: + """File upload and batch create check request buckets against the caller team's allowlist themselves.""" + return get_request_route_template(request) in _TEAM_S3_BUCKET_ROUTE_TEMPLATES and request.method == "POST" + + +def enforce_team_s3_bucket_allowlist( + request_body: Mapping[str, object], + general_settings: Mapping[str, object], + llm_router: Router | None, + team_metadata: Mapping[str, object] | None, +) -> None: + """Re-run the request body check that auth relaxed for bucket params, now admitting the caller team's buckets.""" + proxy_buckets: Final = _ClientSideStorageAllowlist.model_validate(general_settings).allowed_client_side_s3_buckets + model: Final = request_body.get("model") + try: + is_request_body_safe( + request_body=dict(request_body), + general_settings={ + **general_settings, + "allowed_client_side_s3_buckets": (*(proxy_buckets or ()), *_team_s3_buckets(team_metadata)), + }, + llm_router=llm_router, + model=model if isinstance(model, str) else "", + ) + except ValueError as rejection: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=f"Authentication Error, {rejection}") + + async def pre_db_read_auth_checks( request: Request, request_data: dict, @@ -875,6 +927,7 @@ async def pre_db_read_auth_checks( llm_router=llm_router, model=request_data.get("model", ""), # [TODO] use model passed in url as well (azure openai routes) route=route, + defer_s3_buckets=checks_team_s3_buckets_in_handler(request), ) # Check 3. Check if IP address is allowed @@ -1522,6 +1575,9 @@ class BatchLimitRequest(Protocol): def model_fields_set(self) -> Collection[str]: ... +_ADMIN_ONLY_METADATA_KEYS: Final = (*ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS, TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY) + + def enforce_batch_limits_are_admin_only( data: BatchLimitRequest, existing_metadata: Mapping[str, object] | None, @@ -1532,7 +1588,8 @@ def enforce_batch_limits_are_admin_only( Every key in ``ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS`` caps what the holder may do with batches, so a holder-writable copy would let a caller lift their - own quota. Gated on the resulting value rather than on presence, so a form + own quota, and the S3 bucket list grants access to buckets with the proxy's + own credentials. Gated on the resulting value rather than on presence, so a form resending the stored value stays a no-op. """ if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: @@ -1542,17 +1599,19 @@ def enforce_batch_limits_are_admin_only( (data.metadata or EMPTY_MAPPING) if "metadata" in data.model_fields_set else stored ) changed: Final = next( - (key for key in ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS if requested.get(key) != stored.get(key)), + (key for key in _ADMIN_ONLY_METADATA_KEYS if requested.get(key) != stored.get(key)), None, ) if changed is None: return + reason: Final = ( + "It lets the holder read and write those buckets with the proxy's AWS credentials." + if changed == TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY + else "It limits what the holder can do with batches, so the holder cannot raise it." + ) raise HTTPException( status_code=403, - detail={ - "error": f"Only proxy admins can set {changed} on a {entity}. " - "It limits what the holder can do with batches, so the holder cannot raise it." - }, + detail={"error": f"Only proxy admins can set {changed} on a {entity}. {reason}"}, ) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 512a9166e31..00fa7f3c33e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -12,6 +12,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -20,6 +21,7 @@ from litellm.litellm_core_utils.core_helpers import attribute_or_none, has_attri from litellm.litellm_core_utils.hidden_params import get_hidden_params, set_hidden_param from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config from litellm.proxy._types import * +from litellm.proxy.auth.auth_utils import enforce_team_s3_bucket_allowlist from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit from litellm.proxy.batches_endpoints.litellm_executed_batches import ( @@ -81,6 +83,7 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedObjectTable router: Final = APIRouter() +_REQUEST_BODY_ADAPTER: Final = TypeAdapter(Mapping[str, object]) class _UserBatchLister(Protocol): @@ -285,6 +288,7 @@ async def create_batch( """ from litellm.proxy.proxy_server import ( general_settings, + general_settings_view, llm_router, proxy_config, proxy_logging_obj, @@ -294,6 +298,12 @@ async def create_batch( data: dict = {} try: data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line + enforce_team_s3_bucket_allowlist( + request_body=_REQUEST_BODY_ADAPTER.validate_python(data), + general_settings=general_settings_view(), + llm_router=llm_router, + team_metadata=user_api_key_dict.team_metadata, + ) _raise_when_routing_headers_disagree(request, data.get("model")) verbose_proxy_logger.debug( "Request received by LiteLLM:\n%s", @@ -384,6 +394,7 @@ async def create_batch( prepare_data_with_credentials( data=_create_batch_data, credentials=credentials, + keep_request_buckets=True, ) # Create batch using model credentials @@ -488,6 +499,7 @@ async def create_batch( prepare_data_with_credentials( data=_create_batch_data, credentials=credentials, + keep_request_buckets=True, ) # Create batch using model credentials @@ -506,6 +518,7 @@ async def create_batch( llm_router=llm_router, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, + keep_request_buckets=True, ) _raise_not_found_when_openai_fallback_unservable( requested_provider=requested_provider, diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 0ece14eb90a..08a84ca91c3 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( UpdateKeyRequest, UpdateTeamRequest, ) +from litellm.proxy.auth.auth_utils import TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIM_MANAGED_TEAM_METADATA_KEY, SCIM_TEAM_DATA_METADATA_KEY, @@ -97,6 +98,7 @@ PROXY_CONTROLLED_TEAM_METADATA_KEYS: Final[frozenset[str]] = frozenset( "priority", "zguard_policy_id", "allowed_tools", + TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY, "allow_client_mock_response", "allow_client_message_redaction_opt_out", "allow_client_pricing_override", diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 9a044f1dee3..e5e68a930c9 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -24,6 +24,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.litellm_core_utils.batch_utils import batch_cost_is_final +from litellm.litellm_core_utils.cloud_storage_security import REQUEST_S3_BUCKET_PARAMS from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.base_llm.managed_resources.utils import is_base64_encoded_unified_id from litellm.proxy._types import ProxyException @@ -884,6 +885,8 @@ def apply_team_provider_credentials( llm_router: Optional["Router"], user_api_key_dict: "UserAPIKeyAuth", custom_llm_provider: str, + *, + keep_request_buckets: bool = False, ) -> None: """ Resolve credentials for a provider-only request (no model pinned) via @@ -901,7 +904,11 @@ def apply_team_provider_credentials( ) if credentials is None: return - prepare_data_with_credentials(data=data, credentials={**credentials, **explicit_request_credentials(data)}) + prepare_data_with_credentials( + data=data, + credentials={**credentials, **explicit_request_credentials(data)}, + keep_request_buckets=keep_request_buckets, + ) def add_deployment_model_info( @@ -925,10 +932,16 @@ def add_deployment_model_info( } +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) +_NO_REQUEST_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + def prepare_data_with_credentials( data: dict, credentials: dict, file_id: str | None = None, + *, + keep_request_buckets: bool = False, ) -> None: """ Update data dictionary with model credentials (in-place). @@ -937,8 +950,15 @@ def prepare_data_with_credentials( data: Data dictionary to update credentials: Credentials from router file_id: Optional original file_id to set (for decoded file IDs) + keep_request_buckets: Keep an S3 bucket the request already carries over the deployment's. Only + file upload and batch create set it, after checking the bucket against the caller's allowlist. + Reads keep the deployment's bucket so a request cannot widen which buckets a file id may point at. """ - data.update(credentials) + request_fields: Final = _REQUEST_FIELDS.validate_python(data) if keep_request_buckets else _NO_REQUEST_FIELDS + request_buckets: Final = MappingProxyType( + {key: request_fields[key] for key in REQUEST_S3_BUCKET_PARAMS if key in request_fields} + ) + data.update({**credentials, **request_buckets}) data.pop("custom_llm_provider", None) if file_id is not None: diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 5ce0ed8b493..ca1470bc94e 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -40,6 +40,7 @@ from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_ST from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope from litellm.llms.base_llm.managed_resources.isolation import build_list_page from litellm.proxy._types import * +from litellm.proxy.auth.auth_utils import enforce_team_s3_bucket_allowlist from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.batches_endpoints.litellm_executed_batches import ( litellm_executed_provider_of, @@ -181,6 +182,7 @@ async def _litellm_executed_batch_input_model( _MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None) _LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject]) +_REQUEST_BODY_ADAPTER: Final = TypeAdapter(Mapping[str, object]) class UploadedFileInfo(TypedDict): @@ -465,6 +467,7 @@ async def route_create_file( prepare_data_with_credentials( data=_create_file_request, credentials={**credentials, **deployment_drop_params(llm_router, model, _create_file_request)}, + keep_request_buckets=True, ) # Create the file with model credentials @@ -528,6 +531,7 @@ async def route_create_file( llm_router=llm_router, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, + keep_request_buckets=True, ) # get configs for custom_llm_provider llm_provider_config: Final = get_files_provider_config(custom_llm_provider=custom_llm_provider) @@ -630,6 +634,12 @@ async def create_file( # Extract file creation parameters using utility function request_body: Final = await read_request_body(request=request) or {} + enforce_team_s3_bucket_allowlist( + request_body=_REQUEST_BODY_ADAPTER.validate_python(request_body), + general_settings=general_settings_view(), + llm_router=llm_router, + team_metadata=user_api_key_dict.team_metadata, + ) file_params: Final = await extract_file_creation_params( request=request, request_body=request_body, @@ -707,6 +717,8 @@ async def create_file( # Add litellm_metadata to data if provided (from form field) if extracted_litellm_metadata is not None: data["litellm_metadata"] = extracted_litellm_metadata + if isinstance(s3_bucket_name := form_data_raw.get("s3_bucket_name"), str) and s3_bucket_name: + data["s3_bucket_name"] = s3_bucket_name if expires_after_anchor is not None or expires_after_seconds_str is not None: if expires_after_anchor is None or expires_after_seconds_str is None: diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 896fc9d975a..bef78554a36 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -1072,6 +1072,28 @@ def respx_interceptable_httpx_client(monkeypatch): litellm.in_memory_llm_clients_cache.flush_cache() +@pytest.mark.asyncio +@respx.mock +async def test_bedrock_output_written_to_a_request_bucket_is_read_from_that_bucket(respx_interceptable_httpx_client): + route: Final = respx.get( + "https://s3.us-west-2.amazonaws.com/team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out" + ).mock(return_value=httpx.Response(200, content=b'{"recordId": "1"}\n')) + + file_content: Final = await bu.fetch_batch_managed_file_content( + "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out", + custom_llm_provider="bedrock", # pyright: ignore[reportArgumentType] # the cost job passes bedrock at runtime + litellm_params={ + "s3_bucket_name": "model-bucket", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + }, + ) + + assert route.call_count == 1 + assert file_content == b'{"recordId": "1"}\n' + + @pytest.mark.asyncio @respx.mock async def test_output_file_content_vertex_managed_uri_accepted_by_real_validation(respx_interceptable_httpx_client): diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 8c8e970b80d..d8a647fda27 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -3520,6 +3520,87 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none(): assert managed_file_table.upsert_calls == [] +TEAM_BUCKET_FILE_ID: Final = "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out" + + +def _managed_files_holding_a_team_bucket_file() -> PROXY_LiteLLMManagedFiles: + internal_usage_cache: Final = MagicMock() + internal_usage_cache.async_get_cache = AsyncMock( + return_value={ + "unified_file_id": _managed_output_file_id(TEAM_BUCKET_FILE_ID), + "model_mappings": {"model-1": TEAM_BUCKET_FILE_ID}, + "flat_model_file_ids": [TEAM_BUCKET_FILE_ID], + "file_object": None, + "created_by": "user-a", + "updated_by": "user-a", + } + ) + internal_usage_cache.async_set_cache = AsyncMock() + prisma_client: Final = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + managed_files: Final = PROXY_LiteLLMManagedFiles( + internal_usage_cache=internal_usage_cache, prisma_client=prisma_client + ) + managed_files.delete_unified_file_id = AsyncMock(return_value=None) + return managed_files + + +@pytest.mark.asyncio +async def test_deleting_a_managed_file_in_a_team_bucket_names_that_bucket_to_the_provider(): + router: Final = MagicMock() + router.afile_delete = AsyncMock(return_value=MagicMock()) + + await _managed_files_holding_a_team_bucket_file().afile_delete( + file_id=_managed_output_file_id(TEAM_BUCKET_FILE_ID), litellm_parent_otel_span=None, llm_router=router + ) + + assert router.afile_delete.call_args.kwargs["file_id"] == TEAM_BUCKET_FILE_ID + assert router.afile_delete.call_args.kwargs["stored_s3_bucket_name"] == "team-a-bucket" + + +@pytest.mark.asyncio +async def test_reading_a_managed_file_in_a_team_bucket_names_that_bucket_to_the_provider(): + router: Final = MagicMock() + router.afile_content = AsyncMock(return_value=MagicMock()) + + await _managed_files_holding_a_team_bucket_file().afile_content( + file_id=_managed_output_file_id(TEAM_BUCKET_FILE_ID), litellm_parent_otel_span=None, llm_router=router + ) + + assert router.afile_content.call_args.kwargs["file_id"] == TEAM_BUCKET_FILE_ID + assert router.afile_content.call_args.kwargs["stored_s3_bucket_name"] == "team-a-bucket" + + +@pytest.mark.asyncio +async def test_refreshing_a_managed_file_in_a_team_bucket_names_that_bucket_to_the_provider(): + from litellm.types.llms.openai import OpenAIFileObject + + stored_file: Final = LiteLLM_ManagedFileTable( + unified_file_id="unified-team-file", + file_object=None, + model_mappings={"model-1": TEAM_BUCKET_FILE_ID}, + flat_model_file_ids=[TEAM_BUCKET_FILE_ID], + created_by="user-a", + ) + proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file) + router: Final = MagicMock() + router.registry.return_value.credentials = MagicMock( + return_value={"custom_llm_provider": "bedrock", "s3_bucket_name": "model-bucket"} + ) + provider_file: Final = OpenAIFileObject( + id=TEAM_BUCKET_FILE_ID, object="file", bytes=1, created_at=1, filename="input.jsonl.out", purpose="batch_output" + ) + + with patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=provider_file) as provider_retrieve: + refreshed: Final = await proxy_managed_files.afile_retrieve( + file_id="unified-team-file", litellm_parent_otel_span=None, llm_router=router + ) + + assert (refreshed.id, refreshed.filename) == ("unified-team-file", "input.jsonl.out") + assert provider_retrieve.call_args.kwargs["file_id"] == TEAM_BUCKET_FILE_ID + assert provider_retrieve.call_args.kwargs["stored_s3_bucket_name"] == "team-a-bucket" + + @pytest.mark.asyncio async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none(): """ @@ -5665,6 +5746,44 @@ async def test_provider_format_file_without_ownership_row_stays_accessible(): assert result["file_id"] == RAW_PROVIDER_FILE_ID +def _managed_output_file_id(provider_file_id: str) -> str: + unified: Final = ( + "litellm_proxy:application/json;unified_id,u-1;target_model_names,claude-batch;" + f"llm_output_file_id,{provider_file_id};llm_output_file_model_id,model-1" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_file_id,stored_bucket", + [("s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out", "team-a-bucket"), ("file-openai-1", None)], +) +async def test_owned_managed_file_content_is_read_from_the_bucket_litellm_recorded( + provider_file_id: str, stored_bucket: str | None +): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_first.return_value = _owned_record( + created_by="user_a", team_id="team_a" + ) + proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client) + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()), + cache=MagicMock(), + data={"file_id": _managed_output_file_id(provider_file_id)}, + call_type="afile_content", + ) + + assert (result["model"], result["file_id"], result.get("stored_s3_bucket_name")) == ( + "model-1", + provider_file_id, + stored_bucket, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID]) async def test_post_call_batch_create_stores_ownership_row(batch_id): diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index fdd6d5e61bc..f56910ec103 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2928,6 +2928,88 @@ class TestBedrockFileContentTransformation: assert response.content == b'{"recordId": "x"}' +class TestBedrockFileContentFromTheStoredBucket: + TEAM_OUTPUT_URI: Final = "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out" + TEAM_OUTPUT_URL: Final = "https://s3.us-west-2.amazonaws.com/team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out" + CREDENTIALS: Final = { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + } + + def test_file_in_another_bucket_is_refused_without_the_stored_bucket(self) -> None: + with pytest.raises(Exception, match="does not match the configured storage bucket"): + litellm.file_content( + file_id=self.TEAM_OUTPUT_URI, + custom_llm_provider="bedrock", + s3_bucket_name="model-bucket", + **self.CREDENTIALS, + ) + + def test_file_in_the_stored_bucket_is_read_from_that_bucket(self) -> None: + import httpx + import respx + + with respx.mock: + route: Final = respx.get(self.TEAM_OUTPUT_URL).mock(return_value=httpx.Response(200, content=b"{}")) + response: Final = litellm.file_content( + file_id=self.TEAM_OUTPUT_URI, + custom_llm_provider="bedrock", + s3_bucket_name="model-bucket", + stored_s3_bucket_name="team-a-bucket", + **self.CREDENTIALS, + ) + + assert route.call_count == 1 + assert response.content == b"{}" + + @pytest.mark.parametrize( + "stored_file_id", + [ + "s3://team-a-bucket/tenant/litellm-batch-outputs/job-1/input.jsonl.out", + "s3://team-a-bucket/tenant/litellm-bedrock-files-nova-1.jsonl", + "s3://team-a-bucket/tenant/deeper/litellm-bedrock-files/nova-1.jsonl", + ], + ) + def test_file_written_under_a_prefixed_request_bucket_is_read_back_from_it(self, stored_file_id: str) -> None: + import httpx + import respx + + from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs + + url: Final = f"https://s3.us-west-2.amazonaws.com/{stored_file_id.removeprefix('s3://')}" + with respx.mock: + route: Final = respx.get(url).mock(return_value=httpx.Response(200, content=b"{}")) + litellm.file_content( + file_id=stored_file_id, + custom_llm_provider="bedrock", + s3_bucket_name="model-bucket", + **stored_s3_bucket_kwargs(stored_file_id), + **self.CREDENTIALS, + ) + + assert route.call_count == 1 + + def test_file_under_the_model_bucket_prefix_keeps_its_url_when_a_stored_bucket_is_passed(self) -> None: + import httpx + import respx + + prefixed_url: Final = ( + "https://s3.us-west-2.amazonaws.com/model-bucket/tenant/litellm-batch-outputs/job-1/input.jsonl.out" + ) + with respx.mock: + route: Final = respx.get(prefixed_url).mock(return_value=httpx.Response(200, content=b"{}")) + litellm.file_content( + file_id="s3://model-bucket/tenant/litellm-batch-outputs/job-1/input.jsonl.out", + custom_llm_provider="bedrock", + s3_bucket_name="model-bucket/tenant", + stored_s3_bucket_name="model-bucket", + **self.CREDENTIALS, + ) + + assert route.call_count == 1 + + class TestBedrockFilesS3SignatureEncoding: """ S3 rebuilds the canonical request from the wire path with single percent-encoding, diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 8651d522ac8..82b80ad8da9 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -4,7 +4,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext import base64 import logging -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import SimpleNamespace from typing import Final, Optional from unittest.mock import MagicMock, patch @@ -12,7 +13,7 @@ import pytest from fastapi import HTTPException, Request import litellm -from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( _allow_model_level_clientside_configurable_parameters, _get_customer_id_from_standard_headers, @@ -33,7 +34,9 @@ from litellm.proxy.auth.auth_utils import ( get_project_model_rpm_limit, get_project_model_tpm_limit, get_request_route_template, + enforce_team_s3_bucket_allowlist, is_request_body_safe, + pre_db_read_auth_checks, ) from litellm.router import Router from litellm.types.utils import oauth_token_exchange_litellm_params, server_owned_wif_litellm_params @@ -5429,6 +5432,189 @@ def test_other_blocked_fields_do_not_point_at_the_storage_allowlist() -> None: assert "allowed_client_side" not in str(error.value) +_TEAM_A_METADATA: Final = {"allowed_s3_buckets": ["team-a-bucket"]} + + +def _enforce_for_team( + body: Mapping[str, object], team_metadata: Mapping[str, object] | None, **general_settings: object +) -> None: + enforce_team_s3_bucket_allowlist( + request_body=body, + general_settings=general_settings, + llm_router=None, + team_metadata=team_metadata, + ) + + +def _rejection_for_team(body: Mapping[str, object], team_metadata: Mapping[str, object] | None) -> HTTPException: + with pytest.raises(HTTPException) as rejection: + _enforce_for_team(body, team_metadata) + assert rejection.value.status_code == 401, rejection.value.detail + return rejection.value + + +@pytest.mark.parametrize("field", ["s3_bucket_name", "s3_output_bucket_name"]) +def test_team_allowlist_admits_its_own_bucket_and_rejects_another_teams(field: str) -> None: + _enforce_for_team({field: "team-a-bucket"}, _TEAM_A_METADATA) + rejection: Final = _rejection_for_team({field: "team-b-bucket"}, _TEAM_A_METADATA) + assert f"{field} is not allowed in request body" in str(rejection.detail) + assert "metadata.allowed_s3_buckets" in str(rejection.detail) + + +@pytest.mark.parametrize( + "team_metadata", + [None, {}, {"allowed_s3_buckets": None}, {"allowed_s3_buckets": "team-a-bucket"}, {"allowed_s3_buckets": [1]}], +) +def test_team_without_a_usable_allowlist_cannot_pick_a_bucket(team_metadata: Mapping[str, object] | None) -> None: + _rejection_for_team({"s3_bucket_name": "team-a-bucket"}, team_metadata) + + +def _team_bucket_admitted( + body: Mapping[str, object], team_metadata: Mapping[str, object] | None, **general_settings: object +) -> bool: + try: + _enforce_for_team(body, team_metadata, **general_settings) + except HTTPException: + return False + return True + + +def test_proxy_wide_allowlist_still_applies_alongside_the_team_allowlist() -> None: + assert _team_bucket_admitted( + {"s3_bucket_name": "team-a-bucket", "s3_output_bucket_name": "shared-bucket"}, + _TEAM_A_METADATA, + allowed_client_side_s3_buckets=["shared-bucket"], + ) + assert not _team_bucket_admitted({"s3_output_bucket_name": "shared-bucket"}, _TEAM_A_METADATA) + + +def test_team_allowlist_is_checked_for_buckets_nested_in_the_body() -> None: + _enforce_for_team({"extra_body": {"s3_bucket_name": "team-a-bucket"}}, _TEAM_A_METADATA) + _rejection_for_team({"litellm_metadata": {"s3_bucket_name": "team-b-bucket"}}, _TEAM_A_METADATA) + + +def test_team_allowlist_admits_no_other_blocked_field() -> None: + rejection: Final = _rejection_for_team({"s3_endpoint_url": "team-a-bucket"}, _TEAM_A_METADATA) + assert "s3_endpoint_url" in str(rejection.detail) + + +def _routed_request(method: str, route_template: str | None, path: str) -> Request: + scope: Final[dict[str, object]] = {"type": "http", "method": method, "path": path, "headers": [], "query_string": b""} + return Request(scope if route_template is None else {**scope, "route": SimpleNamespace(path=route_template)}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route_template,path", + [ + ("/v1/files", "/v1/files"), + ("/files", "/files"), + ("/{provider}/v1/files", "/bedrock/v1/files"), + ("/v1/batches", "/v1/batches"), + ("/batches", "/batches"), + ("/{provider}/v1/batches", "/bedrock/v1/batches"), + ], +) +async def test_auth_leaves_request_buckets_to_the_upload_and_batch_create_handlers( + route_template: str, path: str +) -> None: + assert await _auth_admits( + _routed_request("POST", route_template, path), + {"s3_bucket_name": "team-b-bucket", "s3_output_bucket_name": "team-b-bucket"}, + path, + ) + + +async def _auth_admits(request: Request, request_data: dict[str, object], route: str) -> bool: + try: + await pre_db_read_auth_checks(request=request, request_data=request_data, route=route) + except ValueError: + return False + return True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method,route_template,path", + [ + ("GET", "/v1/files", "/v1/files"), + ("POST", "/v1/files/{file_id}/content", "/v1/files/file-1/content"), + ("POST", "/v1/chat/completions", "/v1/chat/completions"), + ("POST", None, "/v1/batches"), + ], +) +async def test_auth_still_rejects_request_buckets_on_every_other_route( + method: str, route_template: str | None, path: str +) -> None: + with pytest.raises(ValueError, match="s3_bucket_name is not allowed"): + await pre_db_read_auth_checks( + request=_routed_request(method, route_template, path), + request_data={"s3_bucket_name": "team-b-bucket"}, + route=path, + ) + + +@pytest.mark.asyncio +async def test_a_caller_can_never_send_the_stored_bucket_param() -> None: + with pytest.raises(ValueError, match="stored_s3_bucket_name is not allowed"): + await pre_db_read_auth_checks( + request=_routed_request("POST", "/v1/batches", "/v1/batches"), + request_data={"stored_s3_bucket_name": "team-b-bucket"}, + route="/v1/batches", + ) + rejection: Final = _rejection_for_team( + {"stored_s3_bucket_name": "team-a-bucket"}, {"allowed_s3_buckets": ["team-a-bucket"]} + ) + assert "stored_s3_bucket_name" in str(rejection.detail) + with pytest.raises(ValueError, match="stored_s3_bucket_name"): + _batches_body_is_safe({"stored_s3_bucket_name": "team-a-bucket"}, _BUCKET_ALLOWLIST) + + +@pytest.mark.asyncio +async def test_auth_still_rejects_other_blocked_fields_on_the_upload_and_batch_create_routes() -> None: + with pytest.raises(ValueError, match="api_base"): + await pre_db_read_auth_checks( + request=_routed_request("POST", "/v1/batches", "/v1/batches"), + request_data={"s3_bucket_name": "team-b-bucket", "api_base": "https://attacker.example"}, + route="/v1/batches", + ) + + +@pytest.mark.parametrize( + "role,stored,sent,refused", + [ + (LitellmUserRoles.PROXY_ADMIN, None, ["any-bucket"], False), + (LitellmUserRoles.INTERNAL_USER, None, ["any-bucket"], True), + (LitellmUserRoles.INTERNAL_USER, ["team-a-bucket"], ["team-a-bucket", "other"], True), + (LitellmUserRoles.INTERNAL_USER, ["team-a-bucket"], None, True), + (LitellmUserRoles.INTERNAL_USER, ["team-a-bucket"], ["team-a-bucket"], False), + ], +) +def test_only_proxy_admins_change_a_teams_s3_bucket_list( + role: LitellmUserRoles, stored: list[str] | None, sent: list[str] | None, refused: bool +) -> None: + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only + + existing: Final = None if stored is None else {"allowed_s3_buckets": stored, "project": "x"} + request: Final = UpdateTeamRequest( + team_id="t1", metadata={"project": "x"} if sent is None else {"allowed_s3_buckets": sent, "project": "x"} + ) + caller: Final = UserAPIKeyAuth(user_role=role, user_id="u1") + + if not refused: + enforce_batch_limits_are_admin_only( + data=request, existing_metadata=existing, user_api_key_dict=caller, entity="team" + ) + return + with pytest.raises(HTTPException) as rejection: + enforce_batch_limits_are_admin_only( + data=request, existing_metadata=existing, user_api_key_dict=caller, entity="team" + ) + assert rejection.value.status_code == 403 + assert "Only proxy admins can set allowed_s3_buckets on a team." in str(rejection.value.detail) + + def _request_with_content_length(content_length: int) -> Request: return Request( { diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index b5da6f31eb4..cee8a129172 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -85,6 +85,12 @@ CREDS: Dict[str, Dict[str, str]] = { "api_base": "http://vllm.test/v1", "model": "hosted_vllm/qwen", }, + "bedrock-batch": { + "custom_llm_provider": "bedrock", + "model": "bedrock/us.amazon.nova-lite-v1:0", + "s3_bucket_name": "model-bucket", + "s3_output_bucket_name": "model-bucket", + }, } # A real model-encoded file id: decodes to "azure/gpt-4o", strips to "file-original123". @@ -426,6 +432,83 @@ async def test_create__model_from_body(harness): assert resp.id == encode_file_id_with_model("batch-provider-id", "vertex-model", id_type="batch") +_TEAM_A_KEY: Final = UserAPIKeyAuth( + api_key="sk-team-a", team_id="team-a", team_metadata={"allowed_s3_buckets": ["team-a-bucket"]} +) + + +def _batch_body_with_output_bucket(bucket: str) -> dict[str, object]: + return { + "input_file_id": "file-plain", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": "vertex-model", + "s3_output_bucket_name": bucket, + } + + +@pytest.mark.asyncio +async def test_create__output_bucket_on_the_team_allowlist_reaches_the_provider(harness): + set_body(harness, _batch_body_with_output_bucket("team-a-bucket")) + + await call_create(harness, user=_TEAM_A_KEY) + + assert harness.acreate_kwargs()["s3_output_bucket_name"] == "team-a-bucket" + + +@pytest.mark.asyncio +async def test_create__team_with_an_allowlist_and_no_request_bucket_leaves_the_bucket_to_the_model(harness): + set_body( + harness, + { + "input_file_id": "file-plain", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": "vertex-model", + }, + ) + + await call_create(harness, user=_TEAM_A_KEY) + + assert "s3_output_bucket_name" not in harness.acreate_kwargs() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_bucket,expected_bucket", [("team-a-bucket", "team-a-bucket"), (None, "model-bucket")] +) +@pytest.mark.parametrize( + "routing", + [ + {"input_file_id": "file-plain", "model": "bedrock-batch"}, + {"input_file_id": encode_file_id_with_model("file-plain", "bedrock-batch")}, + ], + ids=["model_in_body", "model_encoded_file_id"], +) +async def test_create__model_routed_batch_keeps_the_team_bucket_over_the_deployment_bucket( + harness, routing: dict[str, str], request_bucket: str | None, expected_bucket: str +): + body: Final = {"endpoint": "/v1/chat/completions", "completion_window": "24h", **routing} + set_body(harness, body if request_bucket is None else {**body, "s3_output_bucket_name": request_bucket}) + + await call_create(harness, user=_TEAM_A_KEY) + + assert harness.acreate_kwargs()["s3_output_bucket_name"] == expected_bucket + + +@pytest.mark.asyncio +async def test_create__output_bucket_off_the_team_allowlist_is_rejected_before_the_provider(harness): + set_body(harness, _batch_body_with_output_bucket("team-b-bucket")) + + with pytest.raises(ProxyException) as rejection: + await call_create(harness, user=_TEAM_A_KEY) + + assert str(rejection.value.code) == "401" + assert "s3_output_bucket_name is not allowed in request body" in rejection.value.message + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + @pytest.mark.asyncio async def test_create__model_from_header(harness): set_body( diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 09c76b2dd65..79524cfeda9 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -2297,6 +2297,10 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: return job async def _run(self, job: MagicMock) -> str: + output_file_id, _ = await self._run_capturing(job, self._RAW_OUTPUT_FILE_ID) + return output_file_id + + async def _run_capturing(self, job: MagicMock, raw_output_file_id: str) -> tuple[str, Mapping[str, object]]: from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -2342,7 +2346,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: object="batch", status="completed", ) - response.output_file_id = self._RAW_OUTPUT_FILE_ID + response.output_file_id = raw_output_file_id file_content = MagicMock() file_content.content = b'{"id":"req-1"}' @@ -2352,7 +2356,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: "litellm.files.main.afile_content", new_callable=AsyncMock, return_value=file_content, - ), + ) as provider_file_content, patch( "litellm.batches.batch_utils.get_file_content_as_dictionary", return_value=[{"id": "req-1"}], @@ -2376,7 +2380,17 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: prom_logger=None, ) - return response.output_file_id + return response.output_file_id, provider_file_content.call_args.kwargs + + @pytest.mark.asyncio + async def test_output_written_to_a_team_bucket_is_billed_from_that_bucket(self): + team_output: Final = "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out" + + _, download = await self._run_capturing( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)), team_output + ) + + assert (download["file_id"], download["stored_s3_bucket_name"]) == (team_output, "team-a-bucket") @pytest.mark.asyncio async def test_target_model_names_comes_from_input_file_not_provider_model(self): diff --git a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py index 4af57a83395..88b02ed1aa1 100644 --- a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -322,6 +322,7 @@ _PROXY_READ_TEAM_METADATA_KEYS: Final = ( "allowed_vector_store_indexes", "enforced_file_expires_after", "enforced_batch_output_expires_after", + "allowed_s3_buckets", "require_trace_id", "allowed_passthrough_routes", "denied_passthrough_routes", @@ -384,7 +385,9 @@ class TestTeamMetadataEdits: verdict: Final = _metadata_verdict(_STORED_METADATA, {**_STORED_METADATA, key: "set-by-team-admin"}) assert verdict == TeamAdminMetadataKeyProxyControlled(key=key) - @pytest.mark.parametrize("key", ["opted_out_global_guardrails", "callback_settings", "priority", "guardrails"]) + @pytest.mark.parametrize( + "key", ["opted_out_global_guardrails", "callback_settings", "priority", "guardrails", "allowed_s3_buckets"] + ) def test_changing_or_removing_a_stored_proxy_controlled_key_names_it(self, key: str): stored: Final = {**_STORED_METADATA, key: ["stored"]} changed: Final = _metadata_verdict(stored, {**stored, key: ["changed"]}) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 7e54f0893c0..83a121da21a 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -13062,6 +13062,67 @@ async def test_update_team_blocks_org_admin_require_trace_id_change(mock_db_clie assert "Only proxy admins can set `require_trace_id` on a team." in str(exc.value.message) +@pytest.mark.asyncio +async def test_new_team_blocks_org_admin_granting_s3_buckets(mock_db_client: MagicMock) -> None: + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + + with ( + _org_admins(("u-team-admin", "org-1")), + patch( + "litellm.proxy.management_endpoints.team_endpoints._should_auto_add_team_creator", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=MagicMock()), + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._check_org_team_limits", + new=AsyncMock(return_value=None), + ), + ): + exc: Final[pytest.ExceptionInfo[ProxyException]] + with pytest.raises(ProxyException) as exc: + await new_team( + data=NewTeamRequest( + team_alias="t", + organization_id="org-1", + metadata={"allowed_s3_buckets": ["another-orgs-bucket"]}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=_non_admin_auth(), + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set allowed_s3_buckets on a team." in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_team_blocks_org_admin_widening_s3_buckets(mock_db_client: MagicMock) -> None: + existing: Final[MagicMock] = MagicMock() + existing.metadata = {"allowed_s3_buckets": ["team-a-bucket"]} + existing.model_dump.return_value = { + "team_id": "t1", + "organization_id": "org-1", + "metadata": {"allowed_s3_buckets": ["team-a-bucket"]}, + } + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) + + with _org_admins(("u-team-admin", "org-1")): + exc: Final[pytest.ExceptionInfo[ProxyException]] + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest( + team_id="t1", metadata={"allowed_s3_buckets": ["team-a-bucket", "another-orgs-bucket"]} + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=_non_admin_auth(), + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set allowed_s3_buckets on a team." in str(exc.value.message) + + @pytest.mark.asyncio async def test_update_team_blocks_org_admin_null_metadata_clear_require_trace_id( mock_db_client: MagicMock, diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index bbe7f4383b6..2f469bb8684 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -21,6 +21,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_team_provider_deployment, is_litellm_executed_batch, map_raw_file_ids_to_unified, + prepare_data_with_credentials, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.utils import InternalUsageCache, handle_exception_on_proxy @@ -952,3 +953,19 @@ def test_a_provider_wildcard_grant_finds_only_deployments_named_under_it(deploye ) assert (None if credentials is None else credentials["api_key"]) == api_key + + +@pytest.mark.parametrize("field", ["s3_bucket_name", "s3_output_bucket_name"]) +def test_create_paths_keep_a_request_bucket_and_reads_keep_the_deployment_bucket(field: str) -> None: + deployment: Final = {"custom_llm_provider": "bedrock", "api_key": "sk-deployment", field: "model-bucket"} + create_with_bucket: Final[dict[str, object]] = {field: "team-a-bucket"} + create_without_bucket: Final[dict[str, object]] = {} + read_with_bucket: Final[dict[str, object]] = {field: "team-a-bucket"} + + prepare_data_with_credentials(data=create_with_bucket, credentials=dict(deployment), keep_request_buckets=True) + prepare_data_with_credentials(data=create_without_bucket, credentials=dict(deployment), keep_request_buckets=True) + prepare_data_with_credentials(data=read_with_bucket, credentials=dict(deployment)) + + assert create_with_bucket == {field: "team-a-bucket", "api_key": "sk-deployment"} + assert create_without_bucket == {field: "model-bucket", "api_key": "sk-deployment"} + assert read_with_bucket == {field: "model-bucket", "api_key": "sk-deployment"} diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index 27a65a91c41..1a57fc81a9f 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -951,6 +951,45 @@ def test_batch_upload_for_a_provider_model_still_goes_to_the_provider(batch_uplo assert provider_upload.call_args.kwargs["custom_llm_provider"] == "gemini" +def _upload_as_team_a(bucket: str | None): + import litellm.proxy.proxy_server as ps + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="sk-team-a", team_id="team-a", team_metadata={"allowed_s3_buckets": ["team-a-bucket"]} + ) + return _upload_batch_file( + {"x-litellm-model": "gemini-2.0-flash"}, {} if bucket is None else {"s3_bucket_name": bucket} + ) + + +def test_batch_upload_from_a_team_with_an_allowlist_and_no_bucket_leaves_the_bucket_to_the_model(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_as_team_a(None) + + assert response.status_code == 200, response.text + assert "s3_bucket_name" not in provider_upload.call_args.kwargs + + +def test_batch_upload_sends_a_bucket_on_the_team_allowlist_to_the_provider_upload(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_as_team_a("team-a-bucket") + + assert response.status_code == 200, response.text + assert provider_upload.call_args.kwargs["s3_bucket_name"] == "team-a-bucket" + + +def test_batch_upload_rejects_a_bucket_off_the_team_allowlist_before_the_provider_upload(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_as_team_a("team-b-bucket") + + assert response.status_code == 401, response.text + assert "s3_bucket_name is not allowed in request body" in response.text + provider_upload.assert_not_awaited() + + @pytest.mark.parametrize( "model, expected_drop_params", [("vertex-batch-drop", True), ("vertex-batch-keep", False)], diff --git a/tests/unit/test_router/core/test_endpoints_characterization_batch.py b/tests/unit/test_router/core/test_endpoints_characterization_batch.py index 6ba298accfe..626f25d00e6 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_batch.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_batch.py @@ -125,6 +125,32 @@ async def test_acreate_file_sends_target_model_names_to_a_litellm_proxy_deployme assert dict(host.calls[0].kwargs["extra_body"]) == {"keep": 1, "target_model_names": "remote-model"} # pyright: ignore[reportArgumentType] # a mapping +@pytest.mark.parametrize( + "request_bucket,expected_bucket", [({}, "model-bucket"), ({"s3_bucket_name": "team-a-bucket"}, "team-a-bucket")] +) +async def test_acreate_file_prefers_the_request_s3_bucket_over_the_deployment_bucket( + make_core: MakeCore, host: EndpointHost, request_bucket: dict[str, str], expected_bucket: str +): + core: Final = make_core(model_list=[deployment("g", "d1", s3_bucket_name="model-bucket")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch", **request_bucket) + assert host.calls[0].kwargs["s3_bucket_name"] == expected_bucket + + +@pytest.mark.parametrize( + "request_bucket,expected_bucket", + [({}, "model-bucket"), ({"s3_output_bucket_name": "team-a-bucket"}, "team-a-bucket")], +) +async def test_acreate_batch_prefers_the_request_output_bucket_over_the_deployment_bucket( + make_core: MakeCore, host: EndpointHost, request_bucket: dict[str, str], expected_bucket: str +): + host.answer("openai/d1", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1", s3_output_bucket_name="model-bucket")]) + await core.acreate_batch( + model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h", **request_bucket + ) + assert host.calls[0].kwargs["s3_output_bucket_name"] == expected_bucket + + async def test_acreate_file_sends_the_deployment_gcs_bucket_in_litellm_metadata( make_core: MakeCore, host: EndpointHost ):