fix(managed_files): resolve the creator org through the cached team lookup

Batch creation snapshotted the team's organization with a direct
litellm_teamtable query on every create. Go through get_team_object
instead, which serves the team auth already cached and only falls back
to the database when the team was never cached.

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -288,11 +288,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             return user_api_key_dict.org_id
         if not user_api_key_dict.team_id:
             return None
+        from litellm.proxy.auth.auth_checks import get_team_object
+        from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
+
         try:
-            team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
-                where={"team_id": user_api_key_dict.team_id}
+            team: Final = await get_team_object(
+                team_id=user_api_key_dict.team_id,
+                prisma_client=self.prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=user_api_key_dict.parent_otel_span,
+                proxy_logging_obj=proxy_logging_obj,
             )
-            return getattr(team_row, "organization_id", None) if team_row is not None else None
+            return team.organization_id
         except Exception as e:
             verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
             return None

tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py

@@ -413,29 +413,63 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create():
 
 
 @pytest.mark.asyncio
-async def test_store_unified_object_id_resolves_org_through_the_team():
+async def test_store_unified_object_id_resolves_org_through_the_cached_team():
     """Most keys belong to an org only through their team, so the auth object carries no
-    org_id. The create resolves the team's organization so org spend is snapshotted at
-    submission time instead of never being billed."""
-    from types import SimpleNamespace
+    org_id. The create reads the team that auth already cached, so org spend is snapshotted
+    at submission time without a database query in the request path."""
+    from litellm.models.team import LiteLLM_TeamTableCachedObj
+    from litellm.proxy.proxy_server import user_api_key_cache
+
+    instance, store = _in_memory_managed_files()
+    creator = UserAPIKeyAuth(user_id="alice", team_id="team-cached", api_key="hash-alice")
+    await user_api_key_cache.async_set_cache(
+        key="team_id:team-cached",
+        value=LiteLLM_TeamTableCachedObj(team_id="team-cached", organization_id="org-via-team"),
+        model_type=LiteLLM_TeamTableCachedObj,
+    )
+    try:
+        await instance.store_unified_object_id(
+            unified_object_id="unified-b",
+            file_object=_build_batch_response(batch_id="b", status="validating"),
+            litellm_parent_otel_span=None,
+            model_object_id="b",
+            file_purpose="batch",
+            user_api_key_dict=creator,
+            persist_attribution=True,
+        )
+    finally:
+        user_api_key_cache.delete_cache(key="team_id:team-cached")
+
+    assert store["unified-b"]["org_id"] == "org-via-team"
+    instance.prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_store_unified_object_id_resolves_org_from_the_db_when_the_team_is_not_cached():
+    """A team no request has run under yet is absent from the auth cache; its organization
+    still comes back from the table so the org is billed rather than dropped."""
+    from litellm.models.team import LiteLLM_TeamTable
+    from litellm.proxy.proxy_server import user_api_key_cache
 
     instance, store = _in_memory_managed_files()
     instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
-        return_value=SimpleNamespace(organization_id="org-via-team")
+        return_value=LiteLLM_TeamTable(team_id="team-uncached", organization_id="org-via-db")
     )
-    creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice")
+    creator = UserAPIKeyAuth(user_id="alice", team_id="team-uncached", api_key="hash-alice")
+    try:
+        await instance.store_unified_object_id(
+            unified_object_id="unified-b",
+            file_object=_build_batch_response(batch_id="b", status="validating"),
+            litellm_parent_otel_span=None,
+            model_object_id="b",
+            file_purpose="batch",
+            user_api_key_dict=creator,
+            persist_attribution=True,
+        )
+    finally:
+        user_api_key_cache.delete_cache(key="team_id:team-uncached")
 
-    await instance.store_unified_object_id(
-        unified_object_id="unified-b",
-        file_object=_build_batch_response(batch_id="b", status="validating"),
-        litellm_parent_otel_span=None,
-        model_object_id="b",
-        file_purpose="batch",
-        user_api_key_dict=creator,
-        persist_attribution=True,
-    )
-
-    assert store["unified-b"]["org_id"] == "org-via-team"
+    assert store["unified-b"]["org_id"] == "org-via-db"
 
 
 @pytest.mark.asyncio