From c99a1ab0d7978a85724fb81c94ab66e704ded309 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Tue, 11 Aug 2026 22:10:17 -0400 Subject: [PATCH 1/6] fix(bedrock): resolve the managed-batch output bucket on the model-routed and cost-poller paths get_configured_s3_bucket_name accepts the output bucket only from the immutable _litellm_internal_model_credentials snapshot or AWS_S3_BUCKET_NAME. That refusal to read litellm_params is deliberate: the bucket is what validate_managed_cloud_file_id checks a file id against, so trusting a request-supplied value would let a caller redirect reads to a bucket of their choosing Two live entry points reach the Bedrock file-content transformation without ever building that snapshot. The managed-files pre-call hook sets data["model"] for any id carrying llm_output_file_id, which is every batch output, so get_file_content always takes the model-routed branch; that branch called llm_router.afile_content directly, and managed_files_obj.afile_content, the only caller that built the snapshot, is therefore unreachable for batch output. CheckBatchCost spread the deployment credentials as plain kwargs, and get_litellm_params does not carry s3_bucket_name across (gcs_bucket_name is listed for exactly this reason, its S3 counterpart is not), so the poller lost the bucket the same way The result was that every completed Bedrock managed batch failed files.content with "S3 bucket_name is required" and never had its cost tracked, leaving the row to be re-polled every cycle. Both paths now resolve the deployment credentials and pass the same MappingProxyType snapshot the managed-files hook already builds --- .../proxy/common_utils/check_batch_cost.py | 2 + .../openai_files_endpoints/files_endpoints.py | 8 ++ .../proxy_unit_tests/test_check_batch_cost.py | 96 +++++++++++++++++++ .../test_files_endpoint.py | 90 +++++++++++++++++ 4 files changed, 196 insertions(+) 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 6fe37f0aacb..cfe60a79eed 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,6 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple from litellm._logging import verbose_proxy_logger @@ -537,6 +538,7 @@ class CheckBatchCost: credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} _file_content = await afile_content( file_id=raw_output_file_id, + _litellm_internal_model_credentials=MappingProxyType(dict(credentials)), **credentials, ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2432ea3729..1cbed2a68f9 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,6 +7,7 @@ import asyncio import traceback +from types import MappingProxyType from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -706,11 +707,18 @@ async def get_file_content( model: Final = cast(str | None, data.get("model")) if model: + deployment_credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model) + trusted_model_credentials: Final = ( + {"_litellm_internal_model_credentials": MappingProxyType(dict(deployment_credentials))} + if deployment_credentials is not None + else {} + ) response = await llm_router.afile_content( **{ "model": model, "file_id": file_id, **data, + **trusted_model_credentials, } ) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index fa274324fd6..20390159665 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -312,6 +312,102 @@ class TestCheckBatchCost: ), "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" + @pytest.mark.asyncio + async def test_output_fetch_passes_deployment_credentials_as_trusted_snapshot( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Bedrock resolves the output bucket ONLY from the immutable snapshot kwarg. + + Spreading the credentials as plain kwargs is not enough: get_litellm_params drops + s3_bucket_name, so without _litellm_internal_model_credentials the cost poller + cannot read the output file and every completed Bedrock batch stays unbilled. + """ + from types import MappingProxyType + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-bedrock-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "custom_llm_provider": "bedrock", + "s3_bucket_name": "configured-batch-bucket", + "aws_region_name": "us-east-1", + } + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "bedrock" + mock_deployment.litellm_params.model = "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"recordId":"req-1"}' + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"recordId": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("anthropic.claude-haiku-4-5-20251001-v1:0", "bedrock", None, None), + ), + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + + mock_afile_content.assert_awaited() + passed_kwargs = mock_afile_content.await_args[1] + snapshot = passed_kwargs.get("_litellm_internal_model_credentials") + assert snapshot is not None, "cost poller must pass the trusted credential snapshot" + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router 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 f27c8dfd2f4..0c26e5f7695 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 @@ -3149,6 +3149,96 @@ def test_require_managed_files_rejects_raw_provider_file_id( mock_call.assert_not_called() +def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkeypatch): + """A managed batch output id routes by model, and that branch must build the snapshot. + + The managed-files pre-call hook sets data["model"] for any id carrying + llm_output_file_id, so batch output retrieval always takes the model-routed branch + and never reaches managed_files_obj.afile_content. Bedrock resolves its output + bucket only from _litellm_internal_model_credentials, so without the snapshot every + Bedrock batch output retrieval fails with "S3 bucket_name is required". + """ + import base64 + from types import MappingProxyType + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.types.utils import SpecialEnums + + router = Router( + model_list=[ + { + "model_name": "anthropic.batch.claude-4.5-haiku", + "litellm_params": { + "model": "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "us-east-1", + "s3_bucket_name": "configured-batch-bucket", + }, + "model_info": {"id": "bedrock-batch-deployment-id"}, + } + ] + ) + + from unittest.mock import MagicMock + + managed_file_row = MagicMock() + managed_file_row.created_by = "test-user" + managed_file_row.team_id = None + managed_file_row.storage_backend = None + managed_file_row.storage_url = None + prisma_stub = MagicMock() + prisma_stub.db.litellm_managedfiletable.find_first = AsyncMock(return_value=managed_file_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_stub) + setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + + captured_kwargs: dict = {} + + async def _mock_router_afile_content(**kwargs): + captured_kwargs.update(kwargs) + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=b'{"recordId":"req-1"}', + headers={"content-type": "application/octet-stream"}, + ) + ) + + monkeypatch.setattr(router, "afile_content", _mock_router_afile_content) + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/jsonl", + "unified-output-id", + "anthropic.batch.claude-4.5-haiku", + "llm_output_file_id,s3://configured-batch-bucket/out/batch.jsonl", + "bedrock-batch-deployment-id", + ) + encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + try: + response = client.get( + f"/v1/files/{encoded_id}/content", + headers={"Authorization": "Bearer test-key", "custom-llm-provider": "bedrock"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + snapshot = captured_kwargs.get("_litellm_internal_model_credentials") + assert snapshot is not None, "model-routed branch must attach the trusted credential snapshot" + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + + def _unified_managed_file_id() -> str: import base64 From 460f0d29a95c25091fd375cd8f0f76297525ab78 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 12 Aug 2026 03:44:12 -0400 Subject: [PATCH 2/6] test(files): capture routed retrieval calls immutably The mock merged every call into one shared dict, so a second routed retrieval would overwrite the first and the assertions would still pass. Keep one frozen snapshot per call and assert exactly one call, which also makes an unintended second retrieval a failure rather than something the merge hides --- .../proxy/openai_files_endpoint/test_files_endpoint.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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 0c26e5f7695..e363a266688 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 @@ -3194,10 +3194,12 @@ def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkey monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) - captured_kwargs: dict = {} + # One frozen snapshot per call rather than one dict merged across calls, so a second + # invocation is visible instead of silently overwriting the first. + calls: list[MappingProxyType] = [] async def _mock_router_afile_content(**kwargs): - captured_kwargs.update(kwargs) + calls.append(MappingProxyType(dict(kwargs))) return HttpxBinaryResponseContent( response=httpx.Response( status_code=200, @@ -3231,7 +3233,8 @@ def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkey app.dependency_overrides.pop(ps.user_api_key_auth, None) assert response.status_code == 200, response.text - snapshot = captured_kwargs.get("_litellm_internal_model_credentials") + assert len(calls) == 1, f"expected exactly one routed retrieval, got {len(calls)}" + snapshot = calls[0].get("_litellm_internal_model_credentials") assert snapshot is not None, "model-routed branch must attach the trusted credential snapshot" assert isinstance( snapshot, MappingProxyType From 60fe4e464cc56847735d9c3d3889717f51bee371 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 00:58:31 -0400 Subject: [PATCH 3/6] fix(bedrock): resolve the managed-batch output bucket on the inline accounting path too A third path reads a completed batch's output file, and it could not resolve the bucket either. When cost is accounted from the retrieve itself rather than from the poller, the batch success handler calls _handle_completed_batch, which fetches the output file through _extract_file_access_credentials. That helper forwarded a whitelist covering Azure and Vertex, gcs_bucket_name included, but nothing for Bedrock, and retrieve_batch built its litellm_params through get_litellm_params, whose fixed signature drops the trusted credential snapshot. So the snapshot never reached the file read and it failed with "S3 bucket_name is required" for a bucket the deployment had configured, leaving the batch's cost unrecorded. Adding s3_bucket_name to that whitelist would not have worked. The Bedrock file config deliberately resolves the bucket only from the immutable server-side snapshot or the environment, never from a request param, because the bucket is what managed file ids are validated against. The snapshot is therefore what has to flow, exactly as it already does for the model-routed and cost-poller paths. retrieve_batch now re-adds the snapshot after get_litellm_params, the same way the file operations already do, the whitelist forwards it, and the proxy attaches it for router-routed managed batches from the deployment behind the unified id. Verified against a live proxy reading a real completed Bedrock batch: the cost row appears within seconds of the retrieve carrying the batch's real spend and usage, where before the read raised and no row was written. Resolving those credentials is best effort. A batch whose deployment no longer resolves, which happens when a model group is removed while batches are in flight, still serves its status instead of failing the request on the lookup. This matters for the OSS and polling-disabled configurations, where the retrieve path is the only thing that accounts for a batch at all. --- litellm/batches/batch_utils.py | 1 + litellm/batches/main.py | 2 + litellm/proxy/batches_endpoints/endpoints.py | 8 +++ .../openai_files_endpoints/common_utils.py | 25 +++++++ .../test_litellm/batches/test_batch_utils.py | 14 ++++ tests/test_litellm/batches/test_main.py | 31 +++++++++ .../proxy/batches_endpoints/test_endpoints.py | 6 +- .../test_files_common_utils.py | 68 +++++++++++++++++++ 8 files changed, 154 insertions(+), 1 deletion(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index e73b887ae0a..f7aa6c50de8 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -309,6 +309,7 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "bucket_name", "timeout", "max_retries", + "_litellm_internal_model_credentials", ] for key in credential_keys: if key in litellm_params: diff --git a/litellm/batches/main.py b/litellm/batches/main.py index ce52c12818e..bb04d495555 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -22,6 +22,7 @@ from openai.types.batch import BatchRequestCounts import litellm from litellm._logging import verbose_logger +from litellm.files.main import _add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.llms.azure.batches.handler import AzureBatchesAPI @@ -527,6 +528,7 @@ def retrieve_batch( custom_llm_provider=custom_llm_provider, **kwargs, ) + _add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) if litellm_logging_obj is not None: litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e442cefa360..1301b9327ec 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, apply_team_provider_credentials, decode_model_from_file_id, + add_internal_model_credentials_for_batch, encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, @@ -537,6 +538,13 @@ async def retrieve_batch( detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) + if unified_batch_id: + add_internal_model_credentials_for_batch( + data=data, + llm_router=llm_router, + model_id=get_model_id_from_unified_batch_id(unified_batch_id), + ) + response = await llm_router.aretrieve_batch(**data) response._hidden_params["unified_batch_id"] = unified_batch_id if unified_batch_id: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 56e986c89cf..32676bd1d9f 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -465,6 +465,31 @@ def apply_team_provider_credentials( prepare_data_with_credentials(data=data, credentials=credentials) +def add_internal_model_credentials_for_batch( + data: dict, + llm_router: "Router", + model_id: str | None, +) -> None: + """ + Attach the deployment's immutable server-side credential snapshot to a router-routed + batch call (in-place). + + Cost accounting for a completed batch reads the batch's output file, and the Bedrock + file config resolves its bucket only from this snapshot, never from a request param, + because the bucket is what managed file ids are validated against. Without it that + read fails and the batch's cost is never recorded. + """ + if model_id is None: + return + try: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + except Exception: # noqa: BLE001 # the snapshot only enables cost accounting; a batch whose deployment no longer resolves must still be retrievable + return + if credentials is None: + return + data["_litellm_internal_model_credentials"] = MappingProxyType(dict(credentials)) + + def prepare_data_with_credentials( data: dict, credentials: dict, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 523b512e4cf..cacae3624f3 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable. import json import os import sys +from types import MappingProxyType import httpx import pytest @@ -1229,3 +1230,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200) assert models == ["claude-sonnet-4-5"] + + +def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): + """Bedrock resolves a batch's output bucket only from the immutable server-side + snapshot, never from a request param, so cost accounting on the retrieve path cannot + read the output file unless this key is forwarded. Without it the accounting raises + "S3 bucket_name is required" for a bucket the deployment has configured, and the + batch's cost is never recorded.""" + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"}) + + credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot}) + + assert credentials["_litellm_internal_model_credentials"] is snapshot diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index 1f7a91a5511..17e9ee29d4d 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -28,6 +28,7 @@ import sys from contextlib import ExitStack from dataclasses import dataclass from typing import Any, Dict +from types import MappingProxyType from unittest.mock import MagicMock, patch import pytest @@ -742,3 +743,33 @@ def test_resolve_timeout__httpx_timeout_returns_float_read(): resolved = bm._resolve_timeout(_params(timeout=t), {}, "openai") assert isinstance(resolved, float) assert resolved == 99.0 + + +def test_retrieve__forwards_trusted_model_credentials_into_litellm_params(seams): + """The batch's cost is computed by reading its output file after the retrieve, and + Bedrock resolves that bucket only from this immutable snapshot. get_litellm_params has + a fixed signature that drops it, so without re-adding it here the snapshot never + reaches the logging object and cost accounting fails on a bucket that is configured.""" + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket"}) + logging_obj = MagicMock() + + bm.retrieve_batch( + batch_id="batch-1", + custom_llm_provider="openai", + litellm_logging_obj=logging_obj, + _litellm_internal_model_credentials=snapshot, + ) + + litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] + assert litellm_params["_litellm_internal_model_credentials"] is snapshot + + +def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): + """A retrieve with no snapshot must not invent an empty one, which would read as a + configured bucket of nothing.""" + logging_obj = MagicMock() + + bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="openai", litellm_logging_obj=logging_obj) + + litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] + assert "_litellm_internal_model_credentials" not in litellm_params diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a80c19f0708..aa5c63280b8 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1138,7 +1138,11 @@ async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness): # DISPATCH - router fired, direct litellm did not. assert retrieve_harness.router_aretrieve.call_count == 1 retrieve_harness.litellm_aretrieve.assert_not_called() - retrieve_harness.creds_resolver.assert_not_called() + + # Credentials are resolved for the deployment behind the unified id so the batch's + # output file can be read for cost accounting. This id resolves to nothing here, and + # the retrieve must still serve the batch rather than fail on the lookup. + retrieve_harness.creds_resolver.assert_called_once_with(model_id="gpt-4o-mini") # router receives the (still-encoded) batch id verbatim - this layer does # not decode it for the unified path. diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 4a021627c3e..a39f0c5f010 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -95,3 +95,71 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): "unified-out", "unified-err", ) + + +# =========================================================================== # +# add_internal_model_credentials_for_batch - the snapshot that lets a completed +# batch's output file be read, and therefore its cost be recorded +# =========================================================================== # + + +def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): + """Cost accounting for a completed batch reads its output file, and Bedrock resolves + that bucket only from this snapshot. It must be immutable so nothing downstream can + redirect the bucket that managed file ids are validated against.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials_for_batch, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"} + ) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-1") + + snapshot = data["_litellm_internal_model_credentials"] + assert snapshot["s3_bucket_name"] == "configured-bucket" + assert isinstance(snapshot, MappingProxyType) + with pytest.raises(TypeError): + snapshot["s3_bucket_name"] = "attacker-bucket" + router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="deployment-1") + + +@pytest.mark.parametrize( + "model_id, credentials", + [(None, {"s3_bucket_name": "b"}), ("deployment-1", None)], + ids=["no-model-id", "deployment-has-no-credentials"], +) +def test_add_internal_model_credentials_is_a_noop_without_a_resolvable_deployment(model_id, credentials): + """An unroutable batch must be left alone rather than given an empty snapshot, which + would look like a configured bucket of nothing.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials_for_batch, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock(return_value=credentials) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id=model_id) + + assert "_litellm_internal_model_credentials" not in data + + +def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): + """The snapshot only enables cost accounting, so a batch whose deployment no longer + resolves, which happens when a model group is removed while batches are in flight, + must still be retrievable rather than failing the request on the lookup.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials_for_batch, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock(side_effect=KeyError("deployment-gone")) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-gone") + + assert data == {"batch_id": "unified-batch-id"} From d7afc1797cf2c0e2c326bf4f2b379991b92a9498 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 01:32:59 -0400 Subject: [PATCH 4/6] refactor(batches): share the trusted-credentials helper across both call paths The helper that carries the credential snapshot into litellm_params lived private in files/main.py, and the batch retrieve needed it too. It now sits beside get_litellm_params, which is what it augments, so neither caller reaches into the other's private surface. Typed as Mapping/MutableMapping of object rather than Any, which the strict import rules ban. The file-content route builds the snapshot through the same helper as the batch route instead of assembling a conditional mapping inline, which drops two mutable constructions and leaves one way to attach it. Its name loses the batch suffix now that both routes use it. --- litellm/batches/main.py | 4 ++-- litellm/files/main.py | 16 ++++------------ .../litellm_core_utils/get_litellm_params.py | 18 ++++++++++++++++++ litellm/proxy/batches_endpoints/endpoints.py | 4 ++-- .../openai_files_endpoints/common_utils.py | 2 +- .../openai_files_endpoints/files_endpoints.py | 10 ++-------- .../test_files_common_utils.py | 14 +++++++------- 7 files changed, 36 insertions(+), 32 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index bb04d495555..20d38bbb77f 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -22,7 +22,7 @@ from openai.types.batch import BatchRequestCounts import litellm from litellm._logging import verbose_logger -from litellm.files.main import _add_trusted_model_credentials_to_litellm_params +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.llms.azure.batches.handler import AzureBatchesAPI @@ -528,7 +528,7 @@ def retrieve_batch( custom_llm_provider=custom_llm_provider, **kwargs, ) - _add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) + add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) if litellm_logging_obj is not None: litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/files/main.py b/litellm/files/main.py index 34421d13761..9a64c78552b 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -11,7 +11,6 @@ import time import uuid as uuid_module from collections.abc import Coroutine from functools import partial -from types import MappingProxyType from typing import Any, Final, Literal, cast import httpx @@ -34,6 +33,7 @@ import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse from litellm.files.types import FileContentProvider, FileContentStreamingResult +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.common_utils import get_azure_credentials @@ -85,14 +85,6 @@ bedrock_files_instance: Final = BedrockFilesHandler() ################################################# -def _add_trusted_model_credentials_to_litellm_params( - litellm_params_dict: dict[str, Any], kwargs: dict[str, Any] -) -> None: - trusted_model_credentials: Final = 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, @@ -372,7 +364,7 @@ def file_retrieve( ) if provider_config is not None: litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) @@ -494,7 +486,7 @@ def file_delete( pass optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) @@ -834,7 +826,7 @@ def file_content( try: optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index f251ab4d74a..3eb8c163d5c 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping, MutableMapping +from types import MappingProxyType from typing import Final from litellm.llms.openai.data_residency import infer_openai_data_residency @@ -184,3 +186,19 @@ def get_litellm_params( litellm_params[key] = kwargs[key] return litellm_params + + +def add_trusted_model_credentials_to_litellm_params( + litellm_params_dict: MutableMapping[str, object], kwargs: Mapping[str, object] +) -> None: + """ + Carry the immutable server-side credential snapshot into litellm_params. + + get_litellm_params has a fixed signature, so callers that need the snapshot to + survive into the logging object and the downstream file read have to re-add it. Only + a MappingProxyType is accepted, since providers resolve trusted configuration such + as a Bedrock file bucket from it and must not read a request-supplied mapping. + """ + trusted_model_credentials: Final = kwargs.get("_litellm_internal_model_credentials") + if isinstance(trusted_model_credentials, MappingProxyType): + litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1301b9327ec..9a6bb054d1f 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,7 +25,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, apply_team_provider_credentials, decode_model_from_file_id, - add_internal_model_credentials_for_batch, + add_internal_model_credentials, encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, @@ -539,7 +539,7 @@ async def retrieve_batch( ) if unified_batch_id: - add_internal_model_credentials_for_batch( + add_internal_model_credentials( data=data, llm_router=llm_router, model_id=get_model_id_from_unified_batch_id(unified_batch_id), diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 32676bd1d9f..f2e6fb633e1 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -465,7 +465,7 @@ def apply_team_provider_credentials( prepare_data_with_credentials(data=data, credentials=credentials) -def add_internal_model_credentials_for_batch( +def add_internal_model_credentials( data: dict, llm_router: "Router", model_id: str | None, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 1cbed2a68f9..361b5b920e2 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,7 +7,6 @@ import asyncio import traceback -from types import MappingProxyType from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -44,6 +43,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + add_internal_model_credentials, apply_team_provider_credentials, encode_file_id_with_model, extract_file_creation_params, @@ -707,18 +707,12 @@ async def get_file_content( model: Final = cast(str | None, data.get("model")) if model: - deployment_credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model) - trusted_model_credentials: Final = ( - {"_litellm_internal_model_credentials": MappingProxyType(dict(deployment_credentials))} - if deployment_credentials is not None - else {} - ) + add_internal_model_credentials(data=data, llm_router=llm_router, model_id=model) response = await llm_router.afile_content( **{ "model": model, "file_id": file_id, **data, - **trusted_model_credentials, } ) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index a39f0c5f010..ad7f5e4725a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -98,7 +98,7 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): # =========================================================================== # -# add_internal_model_credentials_for_batch - the snapshot that lets a completed +# add_internal_model_credentials - the snapshot that lets a completed # batch's output file be read, and therefore its cost be recorded # =========================================================================== # @@ -108,7 +108,7 @@ def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): that bucket only from this snapshot. It must be immutable so nothing downstream can redirect the bucket that managed file ids are validated against.""" from litellm.proxy.openai_files_endpoints.common_utils import ( - add_internal_model_credentials_for_batch, + add_internal_model_credentials, ) router = MagicMock() @@ -117,7 +117,7 @@ def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): ) data = {"batch_id": "unified-batch-id"} - add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-1") + add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-1") snapshot = data["_litellm_internal_model_credentials"] assert snapshot["s3_bucket_name"] == "configured-bucket" @@ -136,14 +136,14 @@ def test_add_internal_model_credentials_is_a_noop_without_a_resolvable_deploymen """An unroutable batch must be left alone rather than given an empty snapshot, which would look like a configured bucket of nothing.""" from litellm.proxy.openai_files_endpoints.common_utils import ( - add_internal_model_credentials_for_batch, + add_internal_model_credentials, ) router = MagicMock() router.get_deployment_credentials_with_provider = MagicMock(return_value=credentials) data = {"batch_id": "unified-batch-id"} - add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id=model_id) + add_internal_model_credentials(data=data, llm_router=router, model_id=model_id) assert "_litellm_internal_model_credentials" not in data @@ -153,13 +153,13 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): resolves, which happens when a model group is removed while batches are in flight, must still be retrievable rather than failing the request on the lookup.""" from litellm.proxy.openai_files_endpoints.common_utils import ( - add_internal_model_credentials_for_batch, + add_internal_model_credentials, ) router = MagicMock() router.get_deployment_credentials_with_provider = MagicMock(side_effect=KeyError("deployment-gone")) data = {"batch_id": "unified-batch-id"} - add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-gone") + add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-gone") assert data == {"batch_id": "unified-batch-id"} From 2a75381a9fc0f9aac2c03274e1fdfee68f70aa07 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 15 Aug 2026 11:44:06 -0700 Subject: [PATCH 5/6] style(batches): sort the common_utils import block --- litellm/proxy/batches_endpoints/endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 9a6bb054d1f..87f9927b191 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -23,9 +23,9 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + add_internal_model_credentials, apply_team_provider_credentials, decode_model_from_file_id, - add_internal_model_credentials, encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, From 1524880dceefe90d0ba92f03710a24ae609d931b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 15 Aug 2026 12:09:59 -0700 Subject: [PATCH 6/6] fix(batches): sign the retrieve-path output read with the deployment's AWS credentials --- litellm/batches/batch_utils.py | 6 +- .../test_litellm/batches/test_batch_utils.py | 56 +++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index f7aa6c50de8..ebef60b41c9 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -5,6 +5,7 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import _parse_prompt_tokens_details from litellm.types.llms.openai import Batch from litellm.types.utils import CallTypes, ModelInfo, Usage @@ -295,7 +296,7 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: if litellm_params: # List of credential keys that should be passed to file operations - credential_keys: Final = [ + credential_keys: Final = ( "api_key", "api_base", "api_version", @@ -310,7 +311,8 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "timeout", "max_retries", "_litellm_internal_model_credentials", - ] + *AWS_CREDENTIAL_KWARGS_KEYS, + ) for key in credential_keys: if key in litellm_params: credentials[key] = litellm_params[key] diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index cacae3624f3..d2074853f2b 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1243,3 +1243,59 @@ def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot}) assert credentials["_litellm_internal_model_credentials"] is snapshot + + +def test_extract_credentials_forwards_the_deployment_aws_credentials(): + """The retrieve path's logging object carries the deployment's AWS keys in its + litellm_params, and the S3 read of the output file signs with whatever afile_content + receives. Dropping them here sent the read to the ambient credential chain, so a + deployment whose only AWS credentials live in its litellm_params never recorded + batch cost on retrieve even once the bucket resolved.""" + params = { + "aws_access_key_id": "AKIA-deployment", + "aws_secret_access_key": "secret-deployment", + "aws_session_token": "token-deployment", + "aws_region_name": "us-west-2", + "aws_role_name": "arn:aws:iam::123456789012:role/batch-reader", + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + } + + credentials = bu._extract_file_access_credentials(params) + + assert credentials == {key: value for key, value in params.items() if key != "model"} + + +@pytest.mark.asyncio +async def test_output_file_content_bedrock_reads_with_deployment_aws_credentials(monkeypatch): + import litellm.files.main as files_main + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b""})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-west-2"}) + + await bu._fetch_batch_output_file_content( + _batch("s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out"), + custom_llm_provider="bedrock", + litellm_params={ + "aws_access_key_id": "AKIA-deployment", + "aws_secret_access_key": "secret-deployment", + "aws_session_token": "token-deployment", + "aws_region_name": "us-west-2", + "_litellm_internal_model_credentials": snapshot, + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + }, + ) + + assert captured["file_id"] == "s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out" + assert captured["custom_llm_provider"] == "bedrock" + assert captured["aws_access_key_id"] == "AKIA-deployment" + assert captured["aws_secret_access_key"] == "secret-deployment" + assert captured["aws_session_token"] == "token-deployment" + assert captured["aws_region_name"] == "us-west-2" + assert captured["_litellm_internal_model_credentials"] is snapshot + assert "model" not in captured