fix(managed_batches): convert raw output_file_id to managed ID in CheckBatchCost poller

[Sameerlite]

Problem

LiteLLM_ManagedObjectTable.file_object can be written with a raw provider output_file_id (e.g. file-b83bb643-...) instead of a LiteLLM managed base64 ID. When a client later calls GET /files/{output_file_id}/content using this raw ID, the proxy cannot recover the Azure/model routing context and falls back to default OpenAI credentials, producing:

The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable

Root Cause

CheckBatchCost (the background enterprise poller) writes the completed batch's file_object directly to the DB without going through async_post_call_success_hook, which is the HTTP path responsible for converting raw provider file IDs to managed base64 IDs.

Fix

In check_batch_cost.py, after resolving the deployment and provider, retrieve the managed_files hook and convert any raw output_file_id / error_file_id on the response object to managed base64 IDs (and register the mapping in litellm_managedfiletable) before the model_dump_json() DB write.

Test

Added test_raw_output_file_id_converted_to_managed_id to the existing test_check_batch_cost.py suite. It verifies that get_unified_output_file_id and store_unified_file_id are called with the correct arguments and that response.output_file_id is updated to the managed ID before the DB write.

---

Note

Medium Risk
Touches batch polling and managed-file ID translation used for later /files/{id}/content access, so a mistake could break batch retrieval/file downloads. Changes are localized and covered by new unit/regression tests.

Overview
Fixes the CheckBatchCost background poller to translate raw provider output_file_id/error_file_id values into managed base64 IDs (via the managed_files hook) and store the mapping before writing the batch file_object back to LiteLLM_ManagedObjectTable.

Adds a focused unit test asserting the ID conversion + store_unified_file_id registration, and updates the managed-files access regression test setup to include team_id and the new hook calls.

^{Reviewed by Cursor Bugbot for commit e8f2324. Bugbot is set up for automated code reviews on this repo. Configure here.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

Fixes the CheckBatchCost background poller so that raw provider output_file_id / error_file_id values are converted to managed base64 IDs (and registered in litellm_managedfiletable) before the batch file_object is written back to LiteLLM_ManagedObjectTable, preventing later GET /files/{id}/content calls from falling through to default OpenAI credentials.

• check_batch_cost.py: After resolving deployment_info and model_name, retrieves the managed_files hook and, for each of output_file_id / error_file_id, calls get_unified_output_file_id then store_unified_file_id (in that order) before setattr-ing the managed ID onto the response; failures are caught and warned without blocking the DB write.
• test_check_batch_cost.py: New test_raw_output_file_id_converted_to_managed_id covers both file-ID branches with a non-None error_file_id, asserts correct call counts and model_mappings values, and verifies both attributes are updated before the DB flush.
• test_managed_files_access_check.py: Minimal stub additions (team_id, AsyncMock for store_unified_file_id, return_value for get_unified_output_file_id) prevent the existing regression test from failing on the new code path while leaving its original assertions intact.

Confidence Score: 5/5

Safe to merge — the change is localized to the batch cost poller, adds an opt-out conversion path that silently warns on failure, and does not alter existing behavior when the managed_files hook is absent.

The store_unified_file_id call correctly precedes the setattr, so a transient DB error leaves the response with its original raw file ID rather than a dangling managed ID. The idempotency guard prevents double-conversion. The new test covers both output_file_id and error_file_id branches, and the existing regression test's core assertions are unchanged.

No files require special attention.

Important Files Changed

Filename	Overview
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py	Adds post-completion conversion of raw provider file IDs to managed base64 IDs in the batch cost poller; ordering is correct (store before setattr), and the guard against double-conversion is in place.
tests/proxy_unit_tests/test_check_batch_cost.py	Adds a focused regression test covering both output_file_id and error_file_id conversion; side_effect ordering for _is_base64_encoded_unified_file_id is consistent with the four production call sites; existing fixture now stubs get_proxy_hook to return None so pre-existing tests are unaffected.
tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py	Minimal additions to existing test: team_id stub, AsyncMock for store_unified_file_id, and a return_value for get_unified_output_file_id — all required to prevent the new code path from raising when the hook is present; no assertions weakened.

_{Reviews (3): Last reviewed commit: "fix managed file test" | Re-trigger Greptile}

[cursor]

Cursor Bugbot has reviewed your changes using high mode and found 2 potential issues.

Bugbot Autofix prepared fixes for both issues found in the latest run.

• ✅ Fixed: setattr before store creates orphaned managed file ID
  • Moved the setattr(response, _file_attr, _unified_file_id) call to after the await managed_files_hook.store_unified_file_id(...) call so the response only retains the managed ID when the DB mapping has been successfully persisted.
• ✅ Fixed: Missing team_id breaks team-level file access control
  • Populated team_id on the minimal UserAPIKeyAuth from job.team_id so the managed file record is created with the correct team ownership for team-level access checks.

Preview (b99b935b2f)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -300,6 +300,42 @@
                     custom_llm_provider=custom_llm_provider,
                 )
 
+                # CheckBatchCost bypasses async_post_call_success_hook, so convert raw
+                # output/error file IDs to managed base64 IDs before the DB write here.
+                managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files")
+                if managed_files_hook is not None:
+                    from litellm.proxy._types import UserAPIKeyAuth
+                    _minimal_auth = UserAPIKeyAuth(
+                        user_id=job.created_by or "default-user-id",
+                        team_id=getattr(job, "team_id", None),
+                    )
+                    for _file_attr in ["output_file_id", "error_file_id"]:
+                        _raw_file_id = getattr(response, _file_attr, None)
+                        if _raw_file_id and not _is_base64_encoded_unified_file_id(_raw_file_id):
+                            try:
+                                _unified_file_id = managed_files_hook.get_unified_output_file_id(
+                                    output_file_id=_raw_file_id,
+                                    model_id=model_id,
+                                    model_name=str(model_name) if model_name else deployment_info.model_name or None,
+                                )
+                                await managed_files_hook.store_unified_file_id(
+                                    file_id=_unified_file_id,
+                                    file_object=None,
+                                    litellm_parent_otel_span=None,
+                                    model_mappings={model_id: _raw_file_id},
+                                    user_api_key_dict=_minimal_auth,
+                                )
+                                setattr(response, _file_attr, _unified_file_id)
+                                verbose_proxy_logger.info(
+                                    f"CheckBatchCost: converted {_file_attr} "
+                                    f"{_raw_file_id!r} -> managed ID for batch {batch_id}"
+                                )
+                            except Exception as _e:
+                                verbose_proxy_logger.warning(
+                                    f"CheckBatchCost: failed to create managed file ID for "
+                                    f"{_file_attr}={_raw_file_id!r}: {_e}"
+                                )
+
                 # Pass deployment model_info so custom batch pricing
                 # (input_cost_per_token_batches etc.) is used for cost calc
                 deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {}

diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -22,7 +22,9 @@
 
     @pytest.fixture
     def mock_proxy_logging_obj(self):
-        return MagicMock()
+        mock = MagicMock()
+        mock.get_proxy_hook.return_value = None
+        return mock
 
     @pytest.fixture
     def mock_llm_router(self):
@@ -372,3 +374,122 @@
             update_data["batch_processed"] is True
         ), "update() must include batch_processed=True when column is present"
         assert update_data["status"] == "complete"
+
+    @pytest.mark.asyncio
+    async def test_raw_output_file_id_converted_to_managed_id(
+        self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
+    ):
+        """CheckBatchCost must convert a raw provider output_file_id to a managed base64 ID.
+
+        Without this, GET /batches/{id} returns a raw file ID that cannot be routed
+        through the proxy, causing API_KEY errors when clients call GET /files/{id}/content.
+        """
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        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-raw-file-1"
+        mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
+        mock_job.created_by = "user-1"
+
+        check_batch_cost_instance._has_batch_processed_column = True
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[mock_job]
+        )
+
+        raw_output_file_id = "file-batch-output-abc123"
+        fake_managed_id = "bGl0ZWxsbV9wcm94eTo6bWFuYWdlZA=="
+
+        mock_response = MagicMock()
+        mock_response.status = "completed"
+        mock_response.output_file_id = raw_output_file_id
+        mock_response.error_file_id = None
+        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={"api_key": "sk-test"}
+        )
+
+        mock_deployment = MagicMock()
+        mock_deployment.litellm_params.custom_llm_provider = "azure"
+        mock_deployment.litellm_params.model = "azure/gpt-5-mini"
+        mock_deployment.model_name = "gpt-5-batch"
+        mock_deployment.model_info.model_dump.return_value = {}
+        mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
+
+        mock_hook = MagicMock()
+        mock_hook.get_unified_output_file_id.return_value = fake_managed_id
+        mock_hook.store_unified_file_id = AsyncMock()
+        check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = (
+            mock_hook
+        )
+
+        mock_file_content = MagicMock()
+        mock_file_content.content = b'{"id":"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",
+                # call 1: job unified_object_id decode, call 2: existing raw check,
+                # call 3: fix's guard for output_file_id
+                side_effect=[decoded_id, None, 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,
+            ),
+            patch(
+                "litellm.batches.batch_utils._get_file_content_as_dictionary",
+                return_value=[{"id": "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},
+                    ["gpt-4"],
+                ),
+            ),
+            patch(
+                "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
+                return_value=("gpt-5-mini", "azure", 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_hook.get_unified_output_file_id.assert_called_once_with(
+            output_file_id=raw_output_file_id,
+            model_id="model-123",
+            model_name="gpt-5-mini",
+        )
+        mock_hook.store_unified_file_id.assert_awaited_once()
+        store_kwargs = mock_hook.store_unified_file_id.call_args[1]
+        assert store_kwargs["file_id"] == fake_managed_id
+        assert store_kwargs["model_mappings"] == {"model-123": raw_output_file_id}
+        assert mock_response.output_file_id == fake_managed_id

_{You can send follow-ups to the cloud agent here.}

^{Reviewed by Cursor Bugbot for commit 13d34a1. Configure here.}

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ Sameerlite
❌ cursoragent
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Sameerlite]

@greptile re review

[mateo-berri]

@greptileai

[mateo-berri]

LGTM; thanks!