feat(batches): snapshot the creating key's org on the managed object row

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -112,39 +112,66 @@ class CheckBatchCost:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
             return {}
 
-    async def _get_key_attribution(self, batch_id: str, api_key: str | None) -> tuple[str | None, str | None]:
-        """Resolve the creating virtual key's (alias, org_id) from its hashed token."""
+    async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None:
+        """Resolve the creating virtual key's alias from its hashed token."""
         if not api_key:
-            return None, None
+            return None
         try:
             key_row: prisma_models.LiteLLM_VerificationToken | None = (
                 await self.prisma_client.db.litellm_verificationtoken.find_unique(
                     where={"token": api_key}
                 )
             )
-            if key_row is None:
-                return None, None
-            return getattr(key_row, "key_alias", None), getattr(key_row, "org_id", None)
+            return getattr(key_row, "key_alias", None) if key_row is not None else None
         except Exception as e:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
-            return None, None
+            return None
 
-    async def _get_team_attribution(self, team_id: str | None) -> tuple[str | None, str | None]:
-        """Resolve a team's (alias, organization_id) from its id."""
+    async def _get_team_alias(self, team_id: str | None) -> str | None:
+        """Resolve a team's alias from its id."""
         if not team_id:
-            return None, None
+            return None
         try:
             team_row: prisma_models.LiteLLM_TeamTable | None = (
                 await self.prisma_client.db.litellm_teamtable.find_unique(
                     where={"team_id": team_id}
                 )
             )
-            if team_row is None:
-                return None, None
-            return getattr(team_row, "team_alias", None), getattr(team_row, "organization_id", None)
+            return getattr(team_row, "team_alias", None) if team_row is not None else None
         except Exception as e:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
-            return None, None
+            return None
+
+    async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None:
+        """Organization to bill the batch against, snapshotted on the row at creation
+        like team_id. Rows created before the org_id column existed carry None, so they
+        fall back to the creating key's org (or its team's) as resolved today."""
+        org_id = getattr(job, "org_id", None)
+        if org_id:
+            return org_id
+        api_key = getattr(job, "api_key", None)
+        team_id = getattr(job, "team_id", None)
+        try:
+            if api_key:
+                key_row: prisma_models.LiteLLM_VerificationToken | None = (
+                    await self.prisma_client.db.litellm_verificationtoken.find_unique(
+                        where={"token": api_key}
+                    )
+                )
+                key_org_id = getattr(key_row, "org_id", None) if key_row is not None else None
+                if key_org_id:
+                    return key_org_id
+            if team_id:
+                team_row: prisma_models.LiteLLM_TeamTable | None = (
+                    await self.prisma_client.db.litellm_teamtable.find_unique(
+                        where={"team_id": team_id}
+                    )
+                )
+                return getattr(team_row, "organization_id", None) if team_row is not None else None
+            return None
+        except Exception as e:
+            verbose_proxy_logger.error(f"CheckBatchCost: could not resolve org for batch {batch_id}: {e}")
+            return None
 
     async def _build_creator_attribution_metadata(
         self, job: "LiteLLM_ManagedObjectTable", batch_id: str
@@ -173,13 +200,13 @@ class CheckBatchCost:
             **(await self._get_user_info(batch_id, job.created_by)),
         }
 
-        key_alias, key_org_id = await self._get_key_attribution(batch_id, api_key)
+        key_alias = await self._get_key_alias(batch_id, api_key)
         if key_alias is not None:
             metadata["user_api_key_alias"] = key_alias
-        team_alias, team_org_id = await self._get_team_attribution(team_id)
+        team_alias = await self._get_team_alias(team_id)
         if team_alias is not None:
             metadata["user_api_key_team_alias"] = team_alias
-        org_id: Final = key_org_id or team_org_id
+        org_id: Final = await self._get_org_id(job, batch_id)
         if org_id is not None:
             metadata["user_api_key_org_id"] = org_id
         if isinstance(request_tags, list) and request_tags:

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -349,6 +349,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,
                     "updated_by": user_api_key_dict.user_id,
                     "status": file_object.status,
                     **attribution_columns,

litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql

@@ -0,0 +1,4 @@
+-- Add org_id column to LiteLLM_ManagedObjectTable
+-- Snapshots the creating key's organization at submission time, like team_id,
+-- so CheckBatchCost can bill organization spend hours later without re-resolving
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   created_at DateTime @default(now())
   created_by String?
   team_id String?
+  org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
   api_key String?
   request_tags Json? @default("[]")
   updated_at DateTime @updatedAt

litellm/models/managed_files.py

@@ -32,6 +32,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
     file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse
     created_by: str | None = None
     team_id: str | None = None
+    org_id: str | None = None
 
 
 class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase):

schema.prisma

@@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   created_at DateTime @default(now())
   created_by String?
   team_id String?
+  org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
   api_key String?
   request_tags Json? @default("[]")
   updated_at DateTime @updatedAt

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -2553,10 +2553,27 @@ class TestBatchCostAttribution:
 
         assert metadata["user_api_key_alias"] == "prod-key"
 
+    @pytest.mark.asyncio
+    async def test_org_id_snapshotted_on_the_row_wins(self):
+        """The org_id column captures the creating key's organization at submission time,
+        like team_id, so a key later moved to another org still bills the original one."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias="prod-key", org_id="org-moved-to"),
+            team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(
+            self._job(org_id="org-at-creation"), "batch-1"
+        )
+
+        assert metadata["user_api_key_org_id"] == "org-at-creation"
+
     @pytest.mark.asyncio
     async def test_org_id_comes_from_the_creating_key(self):
-        """The spend update writer increments organization spend from user_api_key_org_id,
-        so an org-scoped key's batch cost must carry the key's org id."""
+        """The spend update writer increments organization spend from user_api_key_org_id.
+        A legacy row without the org_id column falls back to the creating key's org."""
         from types import SimpleNamespace
 
         instance = self._instance(

tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py

@@ -390,7 +390,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create():
     """Regression (spend loss): the batch create persists the creating key hash and tags so
     CheckBatchCost can write an attributed spend row instead of a blank one the DB drops."""
     instance, store = _in_memory_managed_files()
-    creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice")
+    creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice", org_id="org-acme")
 
     await instance.store_unified_object_id(
         unified_object_id="unified-b",
@@ -407,6 +407,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create():
     assert row["api_key"] == "hash-alice"
     assert row["created_by"] == "alice"
     assert row["team_id"] == "team-alpha"
+    assert row["org_id"] == "org-acme"
     assert row["request_tags"].data == ["env:prod"]
 
 
@@ -471,6 +472,7 @@ async def test_store_unified_object_id_attribution_columns_are_write_once():
     upsert_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"]
     assert "api_key" not in upsert_data["update"]
     assert "request_tags" not in upsert_data["update"]
+    assert "org_id" not in upsert_data["update"]
 
 
 @pytest.mark.asyncio