From e7c23fc0413cda59fec82f2c973323d0ba1f5a17 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 18:14:36 -0700 Subject: [PATCH] harden cloud file compatibility path --- litellm/files/main.py | 27 +++++++++++++++---- .../cloud_storage_security.py | 21 ++++++++++----- litellm/llms/bedrock/files/handler.py | 11 ++++++-- litellm/llms/vertex_ai/files/handler.py | 2 -- .../openai_files_endpoints/files_endpoints.py | 2 ++ .../files/test_bedrock_files_handler.py | 20 ++++++++++++++ .../files/test_vertex_ai_files_handler.py | 17 ------------ .../test_vertex_ai_files_transformation.py | 26 ++++++++++++++++++ 8 files changed, 94 insertions(+), 32 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index a516ab372d7..65802be683f 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -86,6 +86,16 @@ bedrock_files_instance = BedrockFilesHandler() ################################################# +def _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict: Dict[str, Any], kwargs: Dict[str, Any] +) -> None: + trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials") + if isinstance(trusted_model_credentials, type(MappingProxyType({}))): + litellm_params_dict["_litellm_internal_model_credentials"] = ( + trusted_model_credentials + ) + + @client async def acreate_file( file: FileTypes, @@ -374,6 +384,10 @@ def file_retrieve( ) if provider_config is not None: litellm_params_dict = get_litellm_params(**kwargs) + _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base @@ -498,6 +512,10 @@ def file_delete( pass optional_params = GenericLiteLLMParams(**kwargs) litellm_params_dict = get_litellm_params(**kwargs) + _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -847,11 +865,10 @@ def file_content( try: optional_params = GenericLiteLLMParams(**kwargs) litellm_params_dict = get_litellm_params(**kwargs) - trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials") - if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - litellm_params_dict["_litellm_internal_model_credentials"] = ( - trusted_model_credentials - ) + _add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 client = kwargs.get("client") diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 4766331fae1..cdb36933c59 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,7 +1,7 @@ -import os import posixpath import re -from typing import Optional, Sequence, Tuple +from types import MappingProxyType +from typing import Any, Mapping, Optional, Sequence, Tuple, cast from urllib.parse import quote, unquote from litellm._uuid import uuid @@ -15,6 +15,7 @@ BEDROCK_MANAGED_S3_PREFIXES = ( BEDROCK_MANAGED_S3_UPLOAD_PREFIX, BEDROCK_MANAGED_S3_OUTPUT_PREFIX, ) +_MAPPING_PROXY_TYPE: type = type(MappingProxyType({})) _SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+") @@ -110,12 +111,20 @@ def encode_s3_object_key_for_url(object_key: str) -> str: return quote(unquote(object_key), safe="/") -def should_allow_legacy_cloud_file_ids(litellm_params: Optional[dict] = None) -> bool: +def should_allow_legacy_cloud_file_ids( + litellm_params: Optional[Mapping[str, Any]] = None, +) -> bool: value = None - if isinstance(litellm_params, dict): + if isinstance(litellm_params, Mapping): value = litellm_params.get("allow_legacy_cloud_file_ids") - if value is None: - value = os.getenv("LITELLM_ALLOW_LEGACY_CLOUD_FILE_IDS") + if value is None: + trusted_model_credentials = litellm_params.get( + "_litellm_internal_model_credentials" + ) + if isinstance(trusted_model_credentials, _MAPPING_PROXY_TYPE): + value = cast(Mapping[str, Any], trusted_model_credentials).get( + "allow_legacy_cloud_file_ids" + ) if isinstance(value, bool): return value diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index faa34ac6ca7..ecf157e12ee 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -2,7 +2,7 @@ import asyncio import base64 import os from types import MappingProxyType -from typing import Any, Coroutine, Optional, Tuple, Union +from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast import httpx @@ -101,7 +101,14 @@ class BedrockFilesHandler(BaseAWSLLM): ) bucket_name = None if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - bucket_name = trusted_model_credentials.get("s3_bucket_name") + 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( diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index 402a21f56c4..48683ba41ae 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -11,7 +11,6 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import ( ) from litellm.litellm_core_utils.cloud_storage_security import ( VERTEX_AI_MANAGED_GCS_PREFIX, - should_allow_legacy_cloud_file_ids, validate_managed_cloud_file_id, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client @@ -135,7 +134,6 @@ class VertexAIFilesHandler(GCSBucketBase): scheme="gs://", configured_bucket_name=configured_bucket_name, allowed_object_prefixes=(VERTEX_AI_MANAGED_GCS_PREFIX,), - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(), ) async def afile_content( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c69898621e4..378cbbda89c 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -950,6 +950,7 @@ async def get_file( data=data, credentials=credentials, # type: ignore file_id=original_file_id, + include_internal_credentials=True, ) response = await litellm.afile_retrieve(**data) # type: ignore @@ -1150,6 +1151,7 @@ async def delete_file( data=data, credentials=credentials, # type: ignore file_id=original_file_id, + include_internal_credentials=True, ) response = await litellm.afile_delete( diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py index e73be764c66..7f91b49a6f5 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py @@ -184,3 +184,23 @@ def test_should_forward_trusted_model_credentials_to_bedrock_provider_config(): litellm_params = mock_retrieve_file_content.call_args.kwargs["litellm_params"] assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials assert "s3_bucket_name" not in litellm_params + + +def test_should_forward_trusted_model_credentials_to_retrieve_provider_config(): + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + mock_response = MagicMock() + + with patch.object( + files_main.base_llm_http_handler, + "retrieve_file", + return_value=mock_response, + ) as mock_retrieve_file: + response = files_main.file_retrieve( + file_id="gs://safe-bucket/private/file.jsonl", + custom_llm_provider="vertex_ai", + _litellm_internal_model_credentials=trusted_credentials, + ) + + assert response is mock_response + litellm_params = mock_retrieve_file.call_args.kwargs["litellm_params"] + assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 6558cc2969f..2b71e6c18dc 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -3,7 +3,6 @@ Test Vertex AI files handler functionality """ import asyncio -import os import pytest from unittest.mock import AsyncMock, patch @@ -67,22 +66,6 @@ class TestVertexAIFilesHandler: configured_bucket_name="test-bucket", ) - def test_extract_bucket_and_object_from_file_id_allows_legacy_path_with_env_flag( - self, - ): - file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" - - with patch.dict(os.environ, {"LITELLM_ALLOW_LEGACY_CLOUD_FILE_IDS": "true"}): - bucket_name, object_path = ( - self.handler._extract_bucket_and_object_from_file_id( - file_id=file_id, - configured_bucket_name="test-bucket", - ) - ) - - assert bucket_name == "test-bucket" - assert object_path == "test-file.txt" - def test_extract_bucket_and_object_from_file_id_rejects_no_gs_prefix(self): """Test extraction when gs:// prefix is missing""" file_id = "test-bucket%2Flitellm-vertex-files%2Ftest-file.txt" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 5c018513b0d..91d5e87371f 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -3,6 +3,7 @@ Tests for VertexAIFilesConfig transformation methods (Issues 5-7). """ import urllib.parse +from types import MappingProxyType from urllib.parse import parse_qs, urlparse import httpx @@ -85,6 +86,31 @@ class TestParseGcsUri: assert bucket == "my-bucket" assert encoded == urllib.parse.quote("private/object.txt", safe="") + def test_should_allow_legacy_object_path_with_trusted_server_flag(self, config): + trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True}) + bucket, encoded = config._parse_gcs_uri( + "gs://my-bucket/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket", + "_litellm_internal_model_credentials": trusted_credentials, + }, + ) + + assert bucket == "my-bucket" + assert encoded == urllib.parse.quote("private/object.txt", safe="") + + def test_should_reject_user_supplied_legacy_flag_snapshot(self, config): + with pytest.raises(ValueError, match="LiteLLM-managed"): + config._parse_gcs_uri( + "gs://my-bucket/private/object.txt", + litellm_params={ + "bucket_name": "my-bucket", + "_litellm_internal_model_credentials": { + "allow_legacy_cloud_file_ids": True + }, + }, + ) + def test_should_keep_configured_prefix_for_legacy_object_path(self, config): bucket, encoded = config._parse_gcs_uri( "gs://my-bucket/team-a/private/object.txt",