From 538f054a907bf249df27730284e328d95b6bbc19 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 1 Aug 2026 22:36:09 +0000 Subject: [PATCH] fix(proxy): scope batch file-id redaction to managed deployments and harden edge cases Guard the fail-closed redaction so it only runs when the managed-files hook is active, so proxies without managed files keep returning the raw provider file ids they have always used. Cover input_file_id alongside output/error, recover from a store that raced a concurrent terminal retrieval by re-resolving the raw id to a sibling's unified id, and forward unified_batch_id from update_batch_in_database so model_id still derives when _hidden_params is empty Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai_files_endpoints/common_utils.py | 23 ++-- ..._batch_update_db_managed_output_file_id.py | 104 +++++++++++++++++- 2 files changed, 116 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 6ca70ff5249..7ec5b6d1629 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -944,14 +944,20 @@ def _resolve_batch_file_owner_auth(db_batch_object, requester_api_key_dict): return requester_api_key_dict -def _redact_unconverted_batch_output_file_ids(response, verbose_proxy_logger) -> None: - """Fail closed: never return a raw provider output/error file ID. +def _redact_unconverted_batch_output_file_ids(response, verbose_proxy_logger, managed_files_obj) -> None: + """Fail closed: never return a raw provider input/output/error file ID. - Any ID that could not be rewritten to a managed unified ID (hook missing, - model metadata unresolved, or registration raised) is dropped so it can't be - replayed against /v1/files to bypass the managed-file ownership check. + Only meaningful when the managed-files hook is active. Without it there is no + managed-file ownership check to bypass, and the raw ID is exactly what + non-managed deployments have always returned and use to fetch results, so + redacting there would be a breaking change rather than a security fix. When + the hook is active, any ID that could not be rewritten to a managed unified + ID (model metadata unresolved or registration raised) is dropped so it can't + be replayed against /v1/files to bypass the ownership check. """ - for file_attr in ("output_file_id", "error_file_id"): + if managed_files_obj is None: + return + for file_attr in ("input_file_id", "output_file_id", "error_file_id"): file_id = getattr(response, file_attr, None) if file_id and not _is_base64_encoded_unified_file_id(file_id): setattr(response, file_attr, None) @@ -989,7 +995,9 @@ async def ensure_batch_response_managed_file_ids( verbose_proxy_logger=verbose_proxy_logger, ) - _redact_unconverted_batch_output_file_ids(response, verbose_proxy_logger) + await resolve_output_file_ids_to_unified(response, prisma_client) + + _redact_unconverted_batch_output_file_ids(response, verbose_proxy_logger, managed_files_obj) async def _register_raw_batch_output_file_ids( @@ -1159,6 +1167,7 @@ async def update_batch_in_database( verbose_proxy_logger=verbose_proxy_logger, user_api_key_dict=user_api_key_dict, db_batch_object=db_batch_object, + unified_batch_id=unified_batch_id, ) # Only update if status has changed (when db_batch_object is provided) diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index 29d2bd73a48..472f0b50b5f 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -242,11 +242,14 @@ async def test_ensure_batch_response_derives_model_id_from_unified_batch_id(): @pytest.mark.asyncio -async def test_ensure_batch_response_redacts_raw_id_without_managed_files_obj(): - """Fail closed: without the managed-files hook, conversion is impossible, so the - raw provider id must be redacted rather than returned.""" +async def test_ensure_batch_response_preserves_raw_ids_without_managed_files_obj(): + """Without the managed-files hook there is no ownership ACL to bypass, and raw + provider ids are exactly what non-managed deployments have always returned and + use to fetch results, so they must be preserved rather than redacted; redacting + here would be a backwards-incompatible break, not a security fix.""" response = _build_batch_response( output_file_id="file-raw-output", + error_file_id="file-raw-error", hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, ) @@ -258,7 +261,9 @@ async def test_ensure_batch_response_redacts_raw_id_without_managed_files_obj(): user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), ) - assert response.output_file_id is None + assert response.output_file_id == "file-raw-output" + assert response.error_file_id == "file-raw-error" + assert response.input_file_id == "file-input123" @pytest.mark.asyncio @@ -356,3 +361,94 @@ async def test_ensure_batch_response_mirrors_none_owner_instead_of_synthetic_def assert forwarded_auth.user_id is None assert forwarded_auth.team_id == "victim-team" assert response.output_file_id == UNIFIED_FILE_ID + + +@pytest.mark.asyncio +async def test_ensure_batch_response_redacts_unconverted_raw_input_file_id(): + """Regression for GH #33989 review: a raw provider input_file_id that no managed + row resolves is replayable against /v1/files just like a raw output id, so with + the managed-files hook active it must be redacted, not returned.""" + response = _build_batch_response( + output_file_id=None, + error_file_id=None, + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + assert response.input_file_id == "file-input123" + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=_build_managed_files_mock(), + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert response.input_file_id is None + + +@pytest.mark.asyncio +async def test_ensure_batch_response_recovers_unified_id_when_store_races(): + """Regression for GH #33989 review: if store_unified_file_id raises because a + concurrent terminal retrieval already registered the row, the raw id must be + recovered to the sibling's unified id via re-lookup rather than redacted to + None, which would return a null file id for a batch that was in fact registered.""" + raw_output_file_id = "file-raw-output" + sibling_unified_id = _make_unified_file_id(raw_output_file_id="file-sibling-raw") + response = _build_batch_response( + output_file_id=raw_output_file_id, + error_file_id=None, + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + response.input_file_id = UNIFIED_FILE_ID + + mock_managed_files = _build_managed_files_mock() + mock_managed_files.store_unified_file_id = AsyncMock(side_effect=RuntimeError("duplicate row")) + + mock_prisma = _build_prisma_mock() + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=[None, SimpleNamespace(unified_file_id=sibling_unified_id)] + ) + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert response.output_file_id == sibling_unified_id + + +@pytest.mark.asyncio +async def test_update_batch_in_database_forwards_unified_batch_id_for_model_id(): + """Regression for GH #33989 review: update_batch_in_database must forward + unified_batch_id so a response with empty _hidden_params can still derive + model_id and register its raw output id instead of redacting it to None.""" + raw_output_file_id = "file-raw-output" + batch_id = "batch_managed_ids_test" + unified_batch_id = "litellm_proxy;model_id:deployment-42;llm_batch_id:batch_abc" + + response = _build_batch_response( + batch_id=batch_id, + output_file_id=raw_output_file_id, + hidden_params={}, + ) + + mock_managed_files = _build_managed_files_mock() + mock_prisma = _build_prisma_mock() + + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert mock_managed_files.get_unified_output_file_id.call_args.kwargs["model_id"] == "deployment-42" + stored = json.loads(mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs["data"]["file_object"]) + assert stored["output_file_id"] == UNIFIED_FILE_ID + assert stored["output_file_id"] != raw_output_file_id