fix(proxy): guard managed file id parsing for non-output unified formats

Only treat decoded unified ids as already-wrapped output files when they contain llm_output_file_id. Skip other litellm_proxy id shapes instead of IndexError on split.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute 2026-06-09 13:10:45 +05:30
parent 6748607d91
commit a7f0a5baa5
No known key found for this signature in database
2 changed files with 50 additions and 3 deletions

View file

@ -504,7 +504,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if retrieve_file_id
else False
)
if potential_file_id:
if potential_file_id and "llm_output_file_id," in potential_file_id:
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
if model_id:
data["model"] = model_id
@ -1058,7 +1058,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return file_id.split("llm_output_file_model_id,")[1].split(";")[0]
def get_output_file_id_from_unified_file_id(self, file_id: str) -> str:
return file_id.split("llm_output_file_id,")[1].split(";")[0]
marker = "llm_output_file_id,"
if marker not in file_id:
raise ValueError(
f"Unified id does not contain {marker!r}: {file_id[:80]!r}"
)
return file_id.split(marker, 1)[1].split(";")[0]
async def async_post_call_success_hook(
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
@ -1102,13 +1107,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
decoded_unified_file_id = _is_base64_encoded_unified_file_id(
file_id_value
)
if decoded_unified_file_id:
if (
decoded_unified_file_id
and "llm_output_file_id," in decoded_unified_file_id
):
provider_file_id = (
self.get_output_file_id_from_unified_file_id(
decoded_unified_file_id
)
)
unified_file_id = file_id_value
elif decoded_unified_file_id:
verbose_logger.warning(
f"Skipping {file_attr}={file_id_value!r}: "
"unified id is not a managed file output id"
)
continue
else:
provider_file_id = file_id_value
unified_file_id = self.get_unified_output_file_id(

View file

@ -222,3 +222,36 @@ async def test_should_not_double_wrap_already_unified_output_file_id():
already_unified + "=" * (-len(already_unified) % 4)
).decode()
assert decoded.count(f"llm_output_file_id,{provider_file_id}") == 1
@pytest.mark.asyncio
async def test_should_skip_non_file_unified_id_on_output_file_id():
"""Batch-style unified ids lack llm_output_file_id; must not IndexError or re-wrap."""
import base64
managed_files = _make_managed_files_instance()
batch_unified = (
base64.urlsafe_b64encode(
b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc"
)
.decode()
.rstrip("=")
)
batch_response = _make_batch_response(
model_id="openai/openai/gpt-5.5-batch",
model_name="openai/openai/gpt-5.5-batch",
output_file_id=batch_unified,
)
user_api_key_dict = _make_user_api_key_dict()
with patch("litellm.afile_retrieve", AsyncMock()) as mock_afile_retrieve:
await managed_files.async_post_call_success_hook(
data={},
user_api_key_dict=user_api_key_dict,
response=batch_response,
)
assert batch_response.output_file_id == batch_unified
mock_afile_retrieve.assert_not_called()
managed_files.store_unified_file_id.assert_not_awaited()