From df029adeb4cab2029a2fd782e528beab3ca0092b Mon Sep 17 00:00:00 2001 From: AD Mohanraj Date: Thu, 18 Jun 2026 05:34:25 -0700 Subject: [PATCH] feat(bedrock): support file content retrieval for batch output files (#30595) Implements transform_file_content_request and transform_file_content_response in BedrockFilesConfig so GET /v1/files/{id}/content works for Bedrock batch files. The request transform resolves the file id (direct s3:// URI or base64 unified id) to its S3 object, validates bucket and key prefix against the server-configured bucket, and SigV4-signs an S3 GetObject using the same credential and region resolution as the existing upload path. The credential and region params are validated into a typed model at the boundary, so the only untyped values left are the botocore signing primitives. Also fixes the proxy managed-files path: CredentialLiteLLMParams now carries s3_bucket_name (previously dropped when building deployment credentials) and the managed-files hook passes the deployment credential snapshot when routing afile_content, so unified-id content retrieval works with per-model bucket config instead of only the AWS_S3_BUCKET_NAME env var. Preserves managed-file access control: the proxy file-content endpoint now rejects raw cloud-storage ids (s3://, gs://), which would otherwise skip the owner/team check that only runs for unified ids and let a caller read another tenant's batch output by its object key. Managed outputs are reachable only through their unified file id. The afile_content "not found" error now reports the caller's unified id rather than the resolved internal S3 URI. Fixes #16186, #15563 --- .../proxy/hooks/managed_files.py | 23 +- .../cloud_storage_security.py | 15 + litellm/llms/bedrock/files/handler.py | 68 +--- litellm/llms/bedrock/files/transformation.py | 224 ++++++++++++- .../openai_files_endpoints/files_endpoints.py | 12 + litellm/types/router.py | 1 + .../proxy/test_managed_files_hook.py | 130 ++++++++ .../test_cloud_storage_security.py | 15 + .../test_bedrock_files_transformation.py | 314 ++++++++++++++++++ .../test_files_endpoint.py | 19 ++ 10 files changed, 741 insertions(+), 80 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 6830147116d..8486e37384e 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1,9 +1,9 @@ # What is this? ## This hook is used to check for LiteLLM managed files in the request body, and replace them with model-specific file id -import asyncio import base64 import json +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast from fastapi import HTTPException @@ -1472,8 +1472,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. " error_message += ( - f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. " - f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)." + "To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. " + "Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)." ) # Record blocked deletion metric @@ -1550,9 +1550,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if specific_model_file_id_mapping: exception_dict = {} - for model_id, file_id in specific_model_file_id_mapping.items(): + for model_id, provider_file_id in specific_model_file_id_mapping.items(): try: - return await llm_router.afile_content(model=model_id, file_id=file_id, **data) # type: ignore + # Cloud-storage providers (e.g. Bedrock S3) validate file ids + # against the deployment's configured bucket, which they only + # trust from this immutable server-side snapshot, never from + # request params. + credentials = llm_router.get_deployment_credentials_with_provider( + model_id=model_id + ) + if credentials is not None: + data["_litellm_internal_model_credentials"] = cast( + Dict, MappingProxyType(dict(credentials)) + ) + else: + data.pop("_litellm_internal_model_credentials", None) + return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore except Exception as e: exception_dict[model_id] = str(e) raise Exception( diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index daa3dc60320..a75d1178d5a 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -15,8 +15,23 @@ BEDROCK_MANAGED_S3_PREFIXES = ( BEDROCK_MANAGED_S3_UPLOAD_PREFIX, BEDROCK_MANAGED_S3_OUTPUT_PREFIX, ) +MANAGED_CLOUD_STORAGE_SCHEMES = ("s3://", "gs://") _MAPPING_PROXY_TYPE: type = type(MappingProxyType({})) + +def is_managed_cloud_storage_uri(file_id: str) -> bool: + """ + True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``). + + These are internal provider artifacts. On the multi-tenant proxy they must be + retrieved through their managed unified file id so owner/team access is enforced; + a raw URI supplied by a caller bypasses that check. + """ + return isinstance(file_id, str) and file_id.startswith( + MANAGED_CLOUD_STORAGE_SCHEMES + ) + + _SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+") diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index ecf157e12ee..b6aae2159c1 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -1,8 +1,6 @@ import asyncio -import base64 -import os -from types import MappingProxyType -from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast +from collections.abc import Mapping +from typing import Any, Coroutine, Optional, Tuple, Union import httpx @@ -17,7 +15,6 @@ from litellm.types.llms.openai import ( FileContentRequest, HttpxBinaryResponseContent, ) -from litellm.types.utils import SpecialEnums from ..base_aws_llm import BaseAWSLLM @@ -37,40 +34,9 @@ class BedrockFilesHandler(BaseAWSLLM): ) def _extract_s3_uri_from_file_id(self, file_id: str) -> str: - """ - Extract S3 URI from encoded file ID. + from .transformation import extract_s3_uri_from_file_id - The file ID can be in two formats: - 1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path - 2. Direct S3 URI: s3://bucket/litellm-managed-prefix/path - - Args: - file_id: Encoded file ID or direct S3 URI - - Returns: - S3 URI (e.g., "s3://bucket-name/path/to/file") - """ - # First, try to decode if it's a base64-encoded unified file ID - try: - # Add padding if needed - padded = file_id + "=" * (-len(file_id) % 4) - decoded = base64.urlsafe_b64decode(padded).decode() - - # Check if it's a unified file ID format - if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): - # Extract llm_output_file_id from the decoded string - if "llm_output_file_id," in decoded: - s3_uri = decoded.split("llm_output_file_id,")[1].split(";")[0] - return s3_uri - except Exception: - pass - - # If not base64 encoded or doesn't contain llm_output_file_id, accept only - # explicit S3 URIs. Bucket and key validation happens before any S3 call. - if file_id.startswith("s3://"): - return file_id - - raise ValueError("file_id must be a managed LiteLLM S3 file id") + return extract_s3_uri_from_file_id(file_id) def _parse_s3_uri( self, @@ -95,26 +61,12 @@ class BedrockFilesHandler(BaseAWSLLM): allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, ) - def _get_configured_s3_bucket_name(self, litellm_params: dict) -> str: - trusted_model_credentials = litellm_params.get( - "_litellm_internal_model_credentials" - ) - bucket_name = None - if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - trusted_model_credentials_mapping = cast( - Mapping[str, Any], trusted_model_credentials - ) - candidate_bucket_name = trusted_model_credentials_mapping.get( - "s3_bucket_name" - ) - if isinstance(candidate_bucket_name, str): - bucket_name = candidate_bucket_name - bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME") - if not bucket_name: - raise ValueError( - "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." - ) - return bucket_name + def _get_configured_s3_bucket_name( + self, litellm_params: Mapping[str, object] + ) -> str: + from .transformation import get_configured_s3_bucket_name + + return get_configured_s3_bucket_name(litellm_params) async def afile_content( self, diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index cec2e934af8..6cfaa88275d 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1,23 +1,37 @@ +import base64 import json import os import time -from typing import Any, Dict, List, Optional, Tuple, Union +from collections.abc import Mapping, MutableMapping +from types import MappingProxyType +from typing import ( + Any, + Dict, + List, + Optional, + Tuple, + Union, +) from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, ConfigDict from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.files.utils import FilesAPIUtils from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, + BEDROCK_MANAGED_S3_PREFIXES, BEDROCK_MANAGED_S3_UPLOAD_PREFIX, build_managed_cloud_object_name, encode_s3_object_key_for_url, sanitize_cloud_object_component, + should_allow_legacy_cloud_file_ids, split_configured_cloud_bucket_name, + validate_managed_cloud_file_id, ) from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -28,18 +42,98 @@ from litellm.llms.base_llm.files.transformation import ( from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContentRequest, FileTypes, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, ) -from litellm.types.utils import ExtractedFileData, LlmProviders +from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError +# litellm_params key used to hand the SigV4-signed GET headers from +# `transform_file_content_request` to `validate_environment` (the only hook +# the shared file-content HTTP handler exposes for setting request headers). +# Same pattern as the `upload_url` handoff in `transform_create_file_request`. +S3_SIGNED_GET_HEADERS_PARAM = "_s3_signed_get_headers" + + +class _BedrockS3RequestParams(BaseModel): + """Typed view of the credential/region params the S3 GetObject path reads.""" + + model_config = ConfigDict(extra="ignore") + + aws_access_key_id: str | None = None + aws_secret_access_key: str | None = None + aws_session_token: str | None = None + aws_region_name: str | None = None + aws_session_name: str | None = None + aws_profile_name: str | None = None + aws_role_name: str | None = None + aws_web_identity_token: str | None = None + aws_sts_endpoint: str | None = None + s3_region_name: str | None = None + s3_endpoint_url: str | None = None + + +class _TrustedS3ModelCredentials(BaseModel): + """The S3 bucket the server trusts file ids against, from the deployment snapshot.""" + + model_config = ConfigDict(extra="ignore") + + s3_bucket_name: str | None = None + + +def extract_s3_uri_from_file_id(file_id: str) -> str: + """ + Resolve a Bedrock file id to its S3 URI. + + Accepts either a base64-encoded LiteLLM unified file id (whose decoded + form carries `llm_output_file_id,s3://...`) or a direct `s3://` URI. + """ + try: + padded = file_id + "=" * (-len(file_id) % 4) + decoded = base64.urlsafe_b64decode(padded).decode() + + if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): + if "llm_output_file_id," in decoded: + return decoded.split("llm_output_file_id,")[1].split(";")[0] + except Exception: + pass + + if file_id.startswith("s3://"): + return file_id + + raise ValueError("file_id must be a managed LiteLLM S3 file id") + + +def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: + """ + Resolve the server-configured S3 bucket for Bedrock file operations. + + Only trusts the immutable server-side credential snapshot or the + environment; never a request-supplied param, since the bucket is what + `validate_managed_cloud_file_id` checks file ids against. + """ + trusted_model_credentials = litellm_params.get( + "_litellm_internal_model_credentials" + ) + bucket_name: str | None = None + if isinstance(trusted_model_credentials, MappingProxyType): + snapshot: dict[str, object] = {} + snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot + bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_bucket_name + bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME") + if not bucket_name: + raise ValueError( + "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." + ) + return bucket_name + class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ @@ -63,16 +157,21 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def validate_environment( self, - headers: dict, + headers: MutableMapping[str, object], model: str, messages: List[AllMessageValues], optional_params: dict, - litellm_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, + litellm_params: MutableMapping[str, object], + api_key: str | None = None, + api_base: str | None = None, ) -> dict: - # No additional headers needed for S3 uploads - AWS credentials handled by BaseAWSLLM - return headers + result: dict[str, object] = {} + result.update(headers) + signed_headers = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None) + if isinstance(signed_headers, Mapping): + result.update(signed_headers) # any-ok: untyped handoff headers + # otherwise no extra headers - AWS credentials are handled by BaseAWSLLM + return result def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: """ @@ -927,23 +1026,114 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def transform_file_content_request( self, - file_content_request, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError( - "BedrockFilesConfig does not support file content retrieval" + file_content_request: FileContentRequest, + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + """ + Build a SigV4-signed S3 GetObject request for a Bedrock batch file. + + Bedrock batch file ids are `s3://bucket/key` URIs (or unified ids + that decode to one); the bucket and key are validated against the + server-configured bucket before any request is signed. + """ + file_id = file_content_request.get("file_id") + if not file_id: + raise ValueError("file_id is required for Bedrock file content retrieval") + + s3_uri = extract_s3_uri_from_file_id(file_id) + bucket_name, object_key = validate_managed_cloud_file_id( + file_id=s3_uri, + scheme="s3://", + configured_bucket_name=get_configured_s3_bucket_name(litellm_params), + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( + litellm_params + ), ) + # The shared file-content handler passes optional_params={}, so AWS + # credentials/region arrive via litellm_params here (unlike the upload + # path). s3_region_name wins over aws_region_name, same priority as + # get_complete_file_url above. + merged_params: dict[str, object] = {} + merged_params.update(litellm_params) + merged_params.update(optional_params) + request_params = _BedrockS3RequestParams.model_validate(merged_params) + + region_preference = ( + request_params.s3_region_name or request_params.aws_region_name + ) + region_params: dict[str, str | None] = {"aws_region_name": region_preference} + aws_region_name = self._get_aws_region_name( + optional_params=region_params, model="" + ) + + s3_endpoint_url = ( + request_params.s3_endpoint_url + or f"https://s3.{aws_region_name}.amazonaws.com" + ).rstrip("/") + url = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" + + litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( + api_base=url, + aws_region_name=aws_region_name, + request_params=request_params, + ) + return url, {} + + def _sign_s3_get_request( + self, + api_base: str, + aws_region_name: str, + request_params: _BedrockS3RequestParams, + ) -> dict[str, str]: + """ + SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT). + """ + try: + import hashlib + + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + credentials = self.get_credentials( # any-ok: boto3 Credentials is untyped + aws_access_key_id=request_params.aws_access_key_id, + aws_secret_access_key=request_params.aws_secret_access_key, + aws_session_token=request_params.aws_session_token, + aws_region_name=aws_region_name, + aws_session_name=request_params.aws_session_name, + aws_profile_name=request_params.aws_profile_name, + aws_role_name=request_params.aws_role_name, + aws_web_identity_token=request_params.aws_web_identity_token, + aws_sts_endpoint=request_params.aws_sts_endpoint, + ) + + empty_body_hash = hashlib.sha256(b"").hexdigest() + aws_request = AWSRequest( # any-ok: botocore AWSRequest is untyped + method="GET", + url=api_base, + headers={"x-amz-content-sha256": empty_body_hash}, + ) + auth = SigV4Auth(credentials, "s3", aws_region_name) # any-ok: botocore untyped + auth.add_auth(aws_request) # any-ok: botocore request mutation is untyped + return dict(aws_request.headers) # any-ok: botocore headers are untyped + def transform_file_content_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError( - "BedrockFilesConfig does not support file content retrieval" - ) + if raw_response.status_code >= 400: + raise BedrockError( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + ) + return HttpxBinaryResponseContent(response=raw_response) class BedrockJsonlFilesTransformation: diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d7dab350154..944423632ef 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -38,6 +38,9 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.litellm_core_utils.cloud_storage_security import ( + is_managed_cloud_storage_uri, +) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, encode_file_id_with_model, @@ -726,6 +729,15 @@ async def get_file_content( } ) else: + # A raw cloud-storage URI (s3://, gs://) supplied here would skip the + # managed-file owner/team check that only runs for unified ids, letting + # a caller read another tenant's object by its key. Such objects are only + # reachable through their managed unified id. + if is_managed_cloud_storage_uri(file_id): + raise HTTPException( + status_code=400, + detail="Raw cloud storage file ids cannot be retrieved directly. Use the LiteLLM managed file id returned when the file was created.", + ) # Check for model-based credential routing ( should_route, diff --git a/litellm/types/router.py b/litellm/types/router.py index 1611f1e5538..607bfd584fd 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -186,6 +186,7 @@ class CredentialLiteLLMParams(BaseModel): aws_region_name: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None + s3_bucket_name: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 1336490a344..3169b9b08e0 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -255,3 +255,133 @@ async def test_should_skip_non_file_unified_id_on_output_file_id(): assert batch_response.output_file_id == batch_unified mock_afile_retrieve.assert_not_called() managed_files.store_unified_file_id.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_afile_content_passes_trusted_model_credentials_to_router(): + """ + afile_content must hand the deployment's credential snapshot to the router + call as an immutable server-side mapping. Cloud-storage providers (Bedrock + S3) validate file ids against the bucket in that snapshot, so without it + unified-id content retrieval only works when AWS_S3_BUCKET_NAME is set. + """ + from types import MappingProxyType + + managed_files = _make_managed_files_instance() + unified_file_id = "unified-file-id" + s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + managed_files.get_model_file_id_mapping = AsyncMock( + return_value={unified_file_id: {"model-123": s3_uri}} + ) + + mock_router = MagicMock() + mock_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "custom_llm_provider": "bedrock", + "s3_bucket_name": "my-bucket", + "aws_region_name": "us-west-2", + } + ) + mock_router.afile_content = AsyncMock(return_value=MagicMock()) + + await managed_files.afile_content( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=mock_router, + ) + + call_kwargs = mock_router.afile_content.call_args.kwargs + assert call_kwargs["model"] == "model-123" + assert call_kwargs["file_id"] == s3_uri + trusted_credentials = call_kwargs["_litellm_internal_model_credentials"] + assert isinstance(trusted_credentials, MappingProxyType) + assert trusted_credentials["s3_bucket_name"] == "my-bucket" + + +@pytest.mark.asyncio +async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch): + """ + Proxy repro for Bedrock batch output retrieval: a unified file id that + resolves to an s3:// output object must be fetched via a SigV4-signed S3 + GET using the deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env). + + Regression test for "BedrockFilesConfig does not support file content + retrieval" raised on this path. + """ + import httpx + import respx + + import litellm + from litellm import Router + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + router = Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + "model_info": {"id": "model-123"}, + } + ] + ) + + managed_files = _make_managed_files_instance() + unified_file_id = "unified-file-id" + s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + managed_files.get_model_file_id_mapping = AsyncMock( + return_value={unified_file_id: {"model-123": s3_uri}} + ) + + expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + with respx.mock: + route = respx.get(expected_url).mock( + return_value=httpx.Response(200, content=b'{"recordId": "x"}') + ) + + response = await managed_files.afile_content( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + ) + + assert route.called + assert ( + route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + ) + assert response.content == b'{"recordId": "x"}' + + +@pytest.mark.asyncio +async def test_afile_content_error_reports_unified_id_not_provider_uri(): + """When every model attempt fails, the error must name the caller's unified + file id, never the resolved internal s3:// URI (no internal-path leak).""" + managed_files = _make_managed_files_instance() + unified_file_id = "litellm_proxy_unified_id_abc" + s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + managed_files.get_model_file_id_mapping = AsyncMock( + return_value={unified_file_id: {"model-123": s3_uri}} + ) + + mock_router = MagicMock() + mock_router.get_deployment_credentials_with_provider = MagicMock(return_value=None) + mock_router.afile_content = AsyncMock(side_effect=Exception("deployment failed")) + + with pytest.raises(Exception) as exc_info: + await managed_files.afile_content( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=mock_router, + ) + + message = str(exc_info.value) + assert unified_file_id in message + assert s3_uri not in message diff --git a/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py b/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py new file mode 100644 index 00000000000..c3a2511a263 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py @@ -0,0 +1,15 @@ +from litellm.litellm_core_utils.cloud_storage_security import ( + is_managed_cloud_storage_uri, +) + + +def test_is_managed_cloud_storage_uri_detects_raw_object_uris(): + assert is_managed_cloud_storage_uri("s3://bucket/litellm-batch-outputs/x.jsonl.out") + assert is_managed_cloud_storage_uri("gs://bucket/litellm-vertex-files/x") + + +def test_is_managed_cloud_storage_uri_ignores_provider_and_unified_ids(): + # Plain provider ids and base64 unified ids carry no storage scheme. + assert not is_managed_cloud_storage_uri("file-abc123") + assert not is_managed_cloud_storage_uri("bGl0ZWxsbV9wcm94eQ==") + assert not is_managed_cloud_storage_uri("") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 4731be13e78..c548fe53e15 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -4,8 +4,11 @@ Test bedrock files transformation functionality import json import os +from unittest.mock import MagicMock from urllib.parse import unquote, urlparse +import pytest + from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation @@ -1173,3 +1176,314 @@ class TestBedrockFilesEmbeddingTransformation: assert not BedrockFilesConfig._is_embedding_record( {"url": "/v1/responses", "body": {"input": "x"}} ) + + +class TestBedrockFileContentTransformation: + """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" + + S3_URI = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + EXPECTED_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + + def _litellm_params(self) -> dict: + return { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + } + + def test_transform_file_content_request_signs_s3_get(self, monkeypatch): + """The request transform must produce the S3 object URL plus SigV4 GET headers.""" + import hashlib + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_GET_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + + url, params = BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": self.S3_URI}, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.EXPECTED_URL + assert params == {} + + signed_headers = litellm_params[S3_SIGNED_GET_HEADERS_PARAM] + assert ( + signed_headers["x-amz-content-sha256"] == hashlib.sha256(b"").hexdigest() + ), "GET has no payload, so the content hash must be the empty-body hash" + authorization = signed_headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIAEXAMPLE/") + assert "/us-west-2/s3/aws4_request" in authorization + assert "x-amz-content-sha256" in authorization + assert "X-Amz-Date" in signed_headers + + def test_transform_file_content_request_decodes_unified_file_id(self, monkeypatch): + """Base64 unified ids carrying llm_output_file_id must resolve to their S3 object.""" + import base64 + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.utils import SpecialEnums + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "unified-id", "", self.S3_URI, "model-id" + ) + encoded_file_id = ( + base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") + ) + + url, _ = BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": encoded_file_id}, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + assert url == self.EXPECTED_URL + + def test_transform_file_content_request_rejects_foreign_bucket(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with pytest.raises(ValueError, match="configured storage bucket"): + BedrockFilesConfig().transform_file_content_request( + file_content_request={ + "file_id": "s3://other-bucket/litellm-batch-outputs/job/x.jsonl.out" + }, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + def test_transform_file_content_request_rejects_unmanaged_key(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with pytest.raises(ValueError, match="LiteLLM-managed"): + BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": "s3://my-bucket/private/x.jsonl"}, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + def test_extract_s3_uri_rejects_non_managed_file_id(self): + """A file id that is neither an s3:// URI nor a unified id must be rejected.""" + from litellm.llms.bedrock.files.transformation import ( + extract_s3_uri_from_file_id, + ) + + with pytest.raises(ValueError, match="managed LiteLLM S3 file id"): + extract_s3_uri_from_file_id("file-1234567890") + + def test_transform_file_content_request_requires_configured_bucket( + self, monkeypatch + ): + """Without a server-configured bucket (env or snapshot), the request must fail + before any S3 call rather than guessing a bucket from the file id.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + + with pytest.raises(ValueError, match="S3 bucket_name is required"): + BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": self.S3_URI}, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + def test_transform_file_content_request_requires_file_id(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with pytest.raises(ValueError, match="file_id is required"): + BedrockFilesConfig().transform_file_content_request( + file_content_request={}, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + def test_sign_request_without_botocore_raises_helpful_error(self, monkeypatch): + """A missing botocore must surface an actionable 'install boto3' error + rather than a raw import failure.""" + import sys + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setitem(sys.modules, "botocore.auth", None) + + with pytest.raises(ImportError, match="boto3"): + BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": self.S3_URI}, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + def test_bucket_resolved_from_trusted_model_credentials(self, monkeypatch): + """Per-model s3_bucket_name must be honored via the server-side credential snapshot.""" + from types import MappingProxyType + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + litellm_params = self._litellm_params() + litellm_params["_litellm_internal_model_credentials"] = MappingProxyType( + {"s3_bucket_name": "my-bucket"} + ) + + url, _ = BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": self.S3_URI}, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.EXPECTED_URL + + def test_s3_region_name_wins_for_content_signing(self, monkeypatch): + """s3_region_name must override aws_region_name for both the URL and the signature.""" + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_GET_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = self._litellm_params() + litellm_params["s3_region_name"] = "eu-west-1" + + url, _ = BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": self.S3_URI}, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url.startswith("https://s3.eu-west-1.amazonaws.com/") + authorization = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]["Authorization"] + assert "/eu-west-1/s3/aws4_request" in authorization + + def test_validate_environment_merges_and_pops_signed_get_headers(self): + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_GET_HEADERS_PARAM, + BedrockFilesConfig, + ) + + litellm_params = { + S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"} + } + + headers = BedrockFilesConfig().validate_environment( + headers={"x-custom": "kept"}, + model="", + messages=[], + optional_params={}, + litellm_params=litellm_params, + ) + + assert headers == { + "x-custom": "kept", + "Authorization": "AWS4-HMAC-SHA256 test", + } + assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params + + def test_transform_file_content_response_wraps_binary_content(self): + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.openai import HttpxBinaryResponseContent + + raw_response = httpx.Response( + status_code=200, + content=b'{"recordId": "CALL0000001"}', + request=httpx.Request("GET", self.EXPECTED_URL), + ) + + result = BedrockFilesConfig().transform_file_content_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == b'{"recordId": "CALL0000001"}' + + def test_transform_file_content_response_raises_on_s3_error(self): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + raw_response = httpx.Response( + status_code=403, + content=b"AccessDenied", + request=httpx.Request("GET", self.EXPECTED_URL), + ) + + with pytest.raises(BedrockError, match="AccessDenied"): + BedrockFilesConfig().transform_file_content_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + def test_file_content_end_to_end_sends_signed_get(self, monkeypatch): + """litellm.file_content must issue a SigV4-signed GET and return the S3 object bytes.""" + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock: + route = respx.get(self.EXPECTED_URL).mock( + return_value=httpx.Response(200, content=b'{"recordId": "x"}') + ) + + response = litellm.file_content( + file_id=self.S3_URI, + custom_llm_provider="bedrock", + **self._litellm_params(), + ) + + assert route.called + request = route.calls[0].request + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "x-amz-content-sha256" in request.headers + assert response.content == b'{"recordId": "x"}' + + @pytest.mark.asyncio + async def test_afile_content_end_to_end_sends_signed_get(self, monkeypatch): + """Async variant: litellm.afile_content over the same signed GET path.""" + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + # respx can only intercept httpx transports + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock: + route = respx.get(self.EXPECTED_URL).mock( + return_value=httpx.Response(200, content=b'{"recordId": "x"}') + ) + + response = await litellm.afile_content( + file_id=self.S3_URI, + custom_llm_provider="bedrock", + **self._litellm_params(), + ) + + assert route.called + assert ( + route.calls[0] + .request.headers["Authorization"] + .startswith("AWS4-HMAC-SHA256") + ) + assert response.content == b'{"recordId": "x"}' diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index cdb09215aa0..f42639cee8a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -228,6 +228,25 @@ def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router) assert "Invalid purpose: my-bad-purpose" in response.json()["error"]["message"] +def test_get_file_content_rejects_raw_cloud_storage_uri(llm_router: Router): + """A raw s3:// file id must be rejected on the proxy content endpoint. + + Such an id is not a managed unified id, so it would otherwise skip the + owner/team access check and let a caller read another tenant's batch output + object by its key. Callers must use the managed unified file id. + """ + from urllib.parse import quote + + s3_file_id = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + response = client.get( + f"/v1/files/{quote(s3_file_id, safe='')}/content?provider=bedrock", + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 400 + assert "managed file id" in response.json()["error"]["message"].lower() + + def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: Router): """ Asserts 'create_file' is called with the correct arguments