From 0d7976116eff0a9090f846e0066067bd867e67fb Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 18:51:18 -0400
Subject: [PATCH] fix(batches): resolve team-scoped keys' org at creation for
 the snapshot

---
 .../proxy/hooks/managed_files.py              | 20 +++++++++++++-
 ..._batch_update_db_managed_output_file_id.py | 27 +++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py
index 4f0ca787280..2ed446ca44e 100644
--- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py
+++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py
@@ -277,6 +277,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         )
         verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
 
+    async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
+        """Organization to snapshot on the managed object row, like team_id. A key that
+        belongs to an org only through its team carries no org_id on the auth object, so
+        resolve the team's organization at creation time; costing then bills the org the
+        batch was submitted under even if the key or team moves before it completes."""
+        if user_api_key_dict.org_id:
+            return user_api_key_dict.org_id
+        if not user_api_key_dict.team_id:
+            return None
+        try:
+            team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
+                where={"team_id": user_api_key_dict.team_id}
+            )
+            return getattr(team_row, "organization_id", None) if team_row is not None else None
+        except Exception as e:
+            verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
+            return None
+
     async def store_unified_object_id(
         self,
         unified_object_id: str,
@@ -349,7 +367,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     "file_purpose": file_purpose,
                     "created_by": resolve_resource_owner_id(user_api_key_dict),
                     "team_id": user_api_key_dict.team_id,
-                    "org_id": user_api_key_dict.org_id,
+                    "org_id": await self._resolve_creator_org_id(user_api_key_dict),
                     "updated_by": user_api_key_dict.user_id,
                     "status": file_object.status,
                     **attribution_columns,
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 63884e0a779..b5865ab4a13 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
@@ -375,6 +375,7 @@ def _in_memory_managed_files():
     table.upsert = AsyncMock(side_effect=_upsert)
     prisma = MagicMock()
     prisma.db.litellm_managedobjecttable = table
+    prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
 
     cache = MagicMock()
     cache.async_set_cache = AsyncMock()
@@ -411,6 +412,32 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create():
     assert row["request_tags"].data == ["env:prod"]
 
 
+@pytest.mark.asyncio
+async def test_store_unified_object_id_resolves_org_through_the_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
+
+    instance, store = _in_memory_managed_files()
+    instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
+        return_value=SimpleNamespace(organization_id="org-via-team")
+    )
+    creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice")
+
+    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"
+
+
 @pytest.mark.asyncio
 async def test_store_unified_object_id_omits_key_and_tags_without_persist_attribution():
     """Regression (spend redirect): a caller that is not the batch create (a poll, or the