fix(batches): attribute Vertex passthrough batch cost to key/team/tags

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -3,7 +3,7 @@
 """
 
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
@@ -43,11 +43,15 @@ def __init__(
         # the guaranteed-failing primary query on every subsequent cycle.
         self._has_batch_processed_column: bool = True
 
-    async def _get_user_info(self, batch_id, user_id) -> dict:
+    async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
         """
         Look up user email and key alias by user_id for enriching the S3 callback metadata.
         Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
+        Returns an empty dict when user_id is None: batches created by a team or service
+        account key carry no user id, and find_unique(where={"user_id": None}) raises.
         """
+        if not user_id:
+            return {}
         try:
             user_row = await self.prisma_client.db.litellm_usertable.find_unique(
                 where={"user_id": user_id}
@@ -62,6 +66,66 @@ async def _get_user_info(self, batch_id, user_id) -> dict:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
             return {}
 
+    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
+        try:
+            key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique(
+                where={"token": api_key}
+            )
+            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
+
+    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
+        try:
+            team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
+                where={"team_id": team_id}
+            )
+            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
+
+    async def _build_creator_attribution_metadata(
+        self, job: "LiteLLM_ManagedObjectTable", batch_id: str
+    ) -> Dict[str, Any]:
+        """
+        Rebuild the spend-tracking metadata for the key, team, and tags that created the
+        batch so the batch-cost spend log is attributed the same way a non-batch request
+        is. Rows created before api_key and request_tags were persisted carry only
+        created_by and team_id, and fall back to those. A named creating key owns
+        user_api_key_alias; when it has no alias, or the key has since been rotated or
+        deleted, the field keeps the creating user's alias that _get_user_info filled in,
+        because a resolvable name is more useful on the spend row than a null.
+        """
+        api_key = getattr(job, "api_key", None)
+        team_id = getattr(job, "team_id", None)
+        request_tags = getattr(job, "request_tags", None)
+
+        metadata: Dict[str, Any] = {
+            "user_api_key_user_id": job.created_by,
+            "user_api_key": api_key,
+            "user_api_key_team_id": team_id,
+            **(await self._get_user_info(batch_id, job.created_by)),
+        }
+
+        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 = await self._get_team_alias(team_id)
+        if team_alias is not None:
+            metadata["user_api_key_team_alias"] = team_alias
+        if isinstance(request_tags, list) and request_tags:
+            metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
+
+        return metadata
+
     async def _cleanup_stale_managed_objects(self) -> None:
         """
         Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@@ -485,9 +549,6 @@ async def _track_completed_batch_cost(
             function_id=str(uuid.uuid4()),
         )
 
-        creator_user_id = job.created_by
-        user_info = await self._get_user_info(batch_id, job.created_by)
-
         logging_obj.update_environment_variables(
             litellm_params={
                 # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@@ -496,11 +557,7 @@ async def _track_completed_batch_cost(
                         "user-agent": CHECK_BATCH_COST_USER_AGENT,
                     }
                 },
-                "metadata": {
-                    "user_api_key_user_id": creator_user_id,
-                    "user_api_key_team_id": getattr(job, "team_id", None),
-                    **user_info,
-                },
+                "metadata": await self._build_creator_attribution_metadata(job, batch_id),
             },
             optional_params={},
         )

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -163,6 +163,8 @@ async def upsert(
         self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]
     ) -> "PrismaManagedObjectRow": ...
 
+    async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
+
 
 class _CursorPageArgs(TypedDict, total=False):
     cursor: Mapping[str, str]
@@ -263,7 +265,24 @@ async def store_unified_object_id(
         model_object_id: str,
         file_purpose: Literal["batch", "fine-tune", "response"],
         user_api_key_dict: UserAPIKeyAuth,
+        request_tags: Sequence[str] | None = None,
+        persist_attribution: bool = False,
+        create_if_missing: bool = True,
     ) -> None:
+        """Persist a managed object row, caching it and upserting it in the DB.
+
+        persist_attribution is set only by the batch create, which is the one caller
+        that can speak for the creator; it gates the api_key and request_tags columns
+        that CheckBatchCost bills against, so a later poll or retrieve of the same
+        batch cannot record itself as the paying key. Like created_by and team_id,
+        both are written only in the upsert create branch, never on update.
+
+        create_if_missing is cleared by callers that observe a batch they did not
+        create, such as a poll. They still refresh status and file_object, but a
+        row absent from the table is left absent rather than created with the
+        observer as its creator, because created_by and team_id are written from
+        whoever calls the create branch.
+        """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
@@ -277,6 +296,29 @@ async def store_unified_object_id(
             litellm_parent_otel_span=litellm_parent_otel_span,
         )
 
+        from prisma import Json
+
+        api_key = user_api_key_dict.api_key or None
+        attribution_columns = (
+            {
+                **({"api_key": api_key} if api_key is not None else {}),
+                **({"request_tags": Json(list(request_tags))} if request_tags else {}),
+            }
+            if persist_attribution
+            else {}
+        )
+        # FIX: Update status and file_object on every operation to keep state in sync
+        update_columns: Final = {
+            "file_object": file_object.model_dump_json(),
+            "status": file_object.status,
+            "updated_by": user_api_key_dict.user_id,
+        }
+        if not create_if_missing:
+            await _managed_object_table(self.prisma_client).update_many(
+                where={"unified_object_id": unified_object_id},
+                data=update_columns,
+            )
+            return
         await _managed_object_table(self.prisma_client).upsert(
             where={"unified_object_id": unified_object_id},
             data={
@@ -289,12 +331,9 @@ async def store_unified_object_id(
                     "team_id": user_api_key_dict.team_id,
                     "updated_by": user_api_key_dict.user_id,
                     "status": file_object.status,
+                    **attribution_columns,
                 },
-                "update": {
-                    "file_object": file_object.model_dump_json(),
-                    "status": file_object.status,
-                    "updated_by": user_api_key_dict.user_id,
-                },  # FIX: Update status and file_object on every operation to keep state in sync
+                "update": update_columns,
             },
         )
 

litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql

@@ -0,0 +1,5 @@
+-- Add api_key and request_tags columns to LiteLLM_ManagedObjectTable
+-- Captured at batch-create time so CheckBatchCost can attribute batch-cost spend
+-- back to the creating virtual key (and its tags) even when created_by is null.
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "api_key" TEXT;
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB DEFAULT '[]';

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -985,6 +985,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   created_at DateTime @default(now())
   created_by String?
   team_id String?
+  api_key String?
+  request_tags Json? @default("[]")
   updated_at DateTime @updatedAt
   updated_by String?
 

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -48,6 +48,15 @@
     }
 )
 
+# Both spellings, because call_type reaches the callback as str(...) of either the
+# enum member or its value.
+_CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
+    (
+        CallTypes.aretrieve_batch.value,
+        str(CallTypes.aretrieve_batch),
+    )
+)
+
 
 class _ProxyDBLogger(CustomLogger):
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
@@ -212,7 +221,10 @@ async def _PROXY_track_cost_callback(
             # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls).
             # Avoids a cache/DB lookup on every normal LLM request.
             if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"):
-                metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata)
+                metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(  # rebind-ok: enriched metadata replaces the original
+                    metadata=metadata,
+                    resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES,
+                )
                 _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
             budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
             user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
@@ -337,7 +349,7 @@ async def _PROXY_track_cost_callback(
             spend_log_error("Error in tracking cost callback - %s", str(e), exc=e)
 
     @staticmethod
-    async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict:
+    async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict:
         """
         Enriches failure spend log metadata by looking up the key object (and team object)
         from cache/DB when key fields are missing.
@@ -349,6 +361,11 @@ async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict:
         2. Post-auth failures (provider errors, rate limits): key fields are populated
            but team_alias is missing because LiteLLM_VerificationTokenView SQL view
            doesn't include it. We look up the team object to fill in team_alias.
+
+        Scenario 1 reads the key's identity as it stands right now, so it is only correct
+        for a log emitted within the request it describes. Callers that log after a delay,
+        against an identity captured earlier, pass resolve_missing_key_identity=False and
+        keep their own user_id, team_id and org_id.
         """
         api_key_hash: Final = metadata.get("user_api_key")
         if not api_key_hash:
@@ -361,7 +378,7 @@ async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict:
         )
 
         # Step 1: If key fields are missing, look up the full key object
-        if metadata.get("user_api_key_alias") is None:
+        if resolve_missing_key_identity and metadata.get("user_api_key_alias") is None:
             try:
                 key_obj: Final = await get_key_object(
                     hashed_token=api_key_hash,