From 3aaa9bbd4519b136651900a12e48cb69772de18f Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 9 Jul 2026 08:10:08 +0000
Subject: [PATCH] test(managed-files): lock in store_unified_file_id
 idempotency on retrieve

---
 .../proxy/test_managed_files_hook.py          | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
index 3169b9b08e0..1526aad7a24 100644
--- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
+++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
@@ -385,3 +385,54 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri():
     message = str(exc_info.value)
     assert unified_file_id in message
     assert s3_uri not in message
+
+
+def _make_real_managed_files_instance():
+    """Create a _PROXY_LiteLLMManagedFiles with a real store_unified_file_id but
+    an AsyncMock prisma client, so the DB write path itself can be asserted."""
+    from litellm_enterprise.proxy.hooks.managed_files import (
+        _PROXY_LiteLLMManagedFiles,
+    )
+
+    mock_cache = MagicMock()
+    mock_cache.async_set_cache = AsyncMock()
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_managedfiletable.upsert = AsyncMock()
+    mock_prisma.db.litellm_managedfiletable.create = AsyncMock(
+        side_effect=AssertionError(
+            "store_unified_file_id must upsert, not create, on the retrieve path"
+        )
+    )
+
+    return (
+        _PROXY_LiteLLMManagedFiles(
+            internal_usage_cache=mock_cache,
+            prisma_client=mock_prisma,
+        ),
+        mock_prisma,
+    )
+
+
+@pytest.mark.asyncio
+async def test_store_unified_file_id_is_idempotent_via_upsert():
+    """Regression test for the managed-batch retrieve 500 (UniqueViolationError on
+    unified_file_id): re-registering an already-stored output file id must upsert on
+    unified_file_id, never do an unconditional create that raises on conflict."""
+    managed_files, mock_prisma = _make_real_managed_files_instance()
+    file_id = "litellm_proxy_unified_output_id_abc"
+
+    await managed_files.store_unified_file_id(
+        file_id=file_id,
+        file_object=_make_file_object(),
+        litellm_parent_otel_span=None,
+        model_mappings={"model-deploy-xyz": "file-output-abc"},
+        user_api_key_dict=_make_user_api_key_dict(),
+    )
+
+    mock_prisma.db.litellm_managedfiletable.create.assert_not_awaited()
+    mock_prisma.db.litellm_managedfiletable.upsert.assert_awaited_once()
+    assert (
+        mock_prisma.db.litellm_managedfiletable.upsert.await_args.kwargs["where"]
+        == {"unified_file_id": file_id}
+    )