fix(batches): keep batch state in sync on a poll without claiming attribution (#34456)

A poll of a Vertex passthrough batch wrote nothing to the managed-object row,
so status and file_object stayed frozen at the create-time snapshot and
GET /v1/batches served a stale status and an empty output file id for the life
of the batch. Only the create may claim a batch, but every observation of one
may refresh its state.

store_unified_object_id takes create_if_missing, which the poll clears: it
refreshes status and file_object through update_many, and leaves a row that is
absent absent rather than creating one owned by the observer, since created_by
and team_id are written by whoever reaches the create branch. The update payload
is now shared with the upsert so it cannot drift into writing api_key,
request_tags, created_by or team_id.

The passthrough identity re-assertion that was previously part of this PR ships
separately in #36121, so this PR keeps only the batch attribution work.

The creating key owns user_api_key_alias only when it actually has one. Guarding
the overwrite on the presence of a key rather than on a resolved alias nulled the
field out for every key generated without key_alias, and for any key rotated or
deleted before its batch finished, losing the creating user's alias that the spend
row previously carried. The guard now matches the team-alias line below it.

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
 """
 
 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 @@ class CheckBatchCost:
         # 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 @@ 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_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 @@ class CheckBatchCost:
             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 @@ class CheckBatchCost:
                         "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 @@ class _ManagedObjectTableActions(Protocol):
         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 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         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 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             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 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     "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 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
     }
 )
 
+# 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 @@ class _ProxyDBLogger(CustomLogger):
             # 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 @@ class _ProxyDBLogger(CustomLogger):
             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 @@ class _ProxyDBLogger(CustomLogger):
         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 @@ class _ProxyDBLogger(CustomLogger):
         )
 
         # 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,

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -1,4 +1,6 @@
+import asyncio
 import re
+from collections.abc import Mapping
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
@@ -39,6 +41,32 @@ else:
 EndpointType = Any
 
 
+def _optional_str(value: object) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+def _optional_str_tuple(value: object) -> tuple[str, ...] | None:
+    if not isinstance(value, list):
+        return None
+    items: Final = cast(list[object], value)  # cast-ok: isinstance-narrowed; element type unknown
+    return tuple(tag for tag in items if isinstance(tag, str))
+
+
+def _request_tags(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None:
+    """Tags for the batch-cost spend row: the request's own tags when it sent any,
+    otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a
+    tagged key does not put its tags in the top-level metadata "tags" on the
+    passthrough path)
+    """
+    tags: Final = _optional_str_tuple(request_metadata.get("tags"))
+    if tags:
+        return tags
+    key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata")
+    if isinstance(key_auth_metadata, dict):
+        return _optional_str_tuple(key_auth_metadata.get("tags"))
+    return None
+
+
 class VertexPassthroughLoggingHandler:
     @staticmethod
     def vertex_passthrough_handler(
@@ -657,11 +685,13 @@ class VertexPassthroughLoggingHandler:
 
                 # Store the managed object for cost tracking
                 # This will be picked up by check_batch_cost polling mechanism
+                is_batch_create: Final = url_route.split("?")[0].rstrip("/").endswith("batchPredictionJobs")
                 VertexPassthroughLoggingHandler._store_batch_managed_object(
                     unified_object_id=unified_object_id,
                     batch_object=litellm_batch_response,
                     model_object_id=batch_id,
                     logging_obj=logging_obj,
+                    is_batch_create=is_batch_create,
                     **kwargs,
                 )
 
@@ -779,17 +809,45 @@ class VertexPassthroughLoggingHandler:
                 "kwargs": kwargs,
             }
 
+    @staticmethod
+    def _log_batch_registration_result(
+        finished: asyncio.Task, unified_object_id: str, model_object_id: str, is_batch_create: bool
+    ) -> None:
+        error: Final = finished.exception() if not finished.cancelled() else None
+        if finished.cancelled() or error is not None:
+            consequence: Final = (
+                "its cost will not be tracked" if is_batch_create else "its status and output file may be stale"
+            )
+            verbose_proxy_logger.error(
+                "Failed to store batch managed object with unified_object_id=%s, batch_id=%s; %s: %s",
+                unified_object_id,
+                model_object_id,
+                consequence,
+                error,
+            )
+            return
+        verbose_proxy_logger.info(
+            "Stored batch managed object with unified_object_id=%s, batch_id=%s",
+            unified_object_id,
+            model_object_id,
+        )
+
     @staticmethod
     def _store_batch_managed_object(
         unified_object_id: str,
         batch_object: LiteLLMBatch,
         model_object_id: str,
         logging_obj: LiteLLMLoggingObj,
+        is_batch_create: bool,
         **kwargs,
     ) -> None:
         """
         Store batch managed object for cost tracking.
         This will be picked up by the check_batch_cost polling mechanism.
+
+        A poll refreshes the batch status and file object but neither creates the row
+        nor writes attribution, so the creating key and its tags are persisted from
+        the create alone.
         """
         try:
             # Get the managed files hook from the logging object
@@ -805,7 +863,7 @@ class VertexPassthroughLoggingHandler:
 
                 user_api_key_dict: Final = UserAPIKeyAuth(
                     user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
-                    api_key="",
+                    api_key=_optional_str(_request_metadata.get("user_api_key")),
                     team_id=_request_metadata.get("user_api_key_team_id"),
                     team_alias=None,
                     user_role=LitellmUserRoles.CUSTOMER,  # Use proper enum value
@@ -827,9 +885,7 @@ class VertexPassthroughLoggingHandler:
                 )
 
                 # Store the unified object for batch cost tracking
-                import asyncio
-
-                asyncio.create_task(
+                task: Final = asyncio.create_task(
                     managed_files_hook.store_unified_object_id(
                         unified_object_id=unified_object_id,
                         file_object=batch_object,
@@ -837,13 +893,15 @@ class VertexPassthroughLoggingHandler:
                         model_object_id=model_object_id,
                         file_purpose="batch",
                         user_api_key_dict=user_api_key_dict,
+                        request_tags=_request_tags(_request_metadata),
+                        persist_attribution=is_batch_create,
+                        create_if_missing=is_batch_create,
                     )
                 )
-
-                verbose_proxy_logger.info(
-                    "Stored batch managed object with unified_object_id=%s, batch_id=%s",
-                    unified_object_id,
-                    model_object_id,
+                task.add_done_callback(
+                    lambda finished: VertexPassthroughLoggingHandler._log_batch_registration_result(
+                        finished, unified_object_id, model_object_id, is_batch_create
+                    )
                 )
             else:
                 verbose_proxy_logger.warning(

litellm/proxy/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?
 

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?
 

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -1641,3 +1641,153 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
 
         decoded = _is_base64_encoded_unified_file_id(output_file_id)
         assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP]
+class TestBatchCostAttribution:
+    """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so
+    the batch-cost log is attributed like a non-batch request."""
+
+    def _instance(self, key_row=None, team_row=None, user_row=None):
+        from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+        prisma = MagicMock()
+        prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row)
+        prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+        prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
+        return CheckBatchCost(
+            proxy_logging_obj=MagicMock(),
+            prisma_client=prisma,
+            llm_router=MagicMock(),
+        )
+
+    def _job(self, **overrides):
+        from types import SimpleNamespace
+
+        fields = {
+            "created_by": "alice",
+            "team_id": "team-alpha",
+            "api_key": "hash-alice",
+            "request_tags": ["env:prod"],
+        }
+        fields.update(overrides)
+        return SimpleNamespace(unified_object_id="uoi", **fields)
+
+    @pytest.mark.asyncio
+    async def test_metadata_carries_key_team_and_tags(self):
+        """The spend row names the creating key, its team, both aliases, and the tags."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias="prod-key"),
+            team_row=SimpleNamespace(team_alias="Team Alpha"),
+            user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key"] == "hash-alice"
+        assert metadata["user_api_key_user_id"] == "alice"
+        assert metadata["user_api_key_team_id"] == "team-alpha"
+        assert metadata["user_api_key_alias"] == "prod-key"
+        assert metadata["user_api_key_team_alias"] == "Team Alpha"
+        assert metadata["tags"] == ["env:prod"]
+
+    @pytest.mark.asyncio
+    async def test_metadata_tolerates_legacy_row_without_columns(self):
+        """Rows created before the columns existed carry only created_by/team_id and must
+        still produce an attributed row rather than raising."""
+        instance = self._instance()
+        job = self._job(api_key=None, request_tags=None)
+
+        metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
+
+        assert metadata["user_api_key"] is None
+        assert metadata["user_api_key_user_id"] == "alice"
+        assert metadata["user_api_key_team_id"] == "team-alpha"
+        assert "tags" not in metadata
+
+    @pytest.mark.asyncio
+    async def test_metadata_keeps_key_when_team_key_has_no_user(self):
+        """A team-scoped key carries no user id. The user lookup is skipped (prisma rejects
+        a None user_id) and the key hash still drives key-level attribution."""
+        from types import SimpleNamespace
+
+        instance = self._instance(key_row=SimpleNamespace(key_alias="svc-key"))
+        job = self._job(created_by=None)
+
+        metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
+
+        assert metadata["user_api_key"] == "hash-alice"
+        assert metadata["user_api_key_user_id"] is None
+        assert metadata["user_api_key_alias"] == "svc-key"
+        instance.prisma_client.db.litellm_usertable.find_unique.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_metadata_drops_non_string_tags(self):
+        """Non-string tags are dropped so a malformed stored value cannot slip past the
+        tag-budget checks that consume this metadata."""
+        instance = self._instance()
+        job = self._job(request_tags=["env:prod", 7, None, "team:ml"])
+
+        metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
+
+        assert metadata["tags"] == ["env:prod", "team:ml"]
+
+    @pytest.mark.asyncio
+    async def test_key_alias_lookup_failure_does_not_break_attribution(self):
+        """An alias lookup failure must not lose the spend row; the key hash and team still
+        attribute it."""
+        instance = self._instance()
+        instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+            side_effect=Exception("db down")
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key"] == "hash-alice"
+        assert metadata.get("user_api_key_alias") is None
+
+    @pytest.mark.asyncio
+    async def test_unnamed_key_keeps_the_creating_user_alias(self):
+        """Regression: a key generated without key_alias resolves to no alias, and the
+        overwrite must not null out the creating user's alias that _get_user_info supplied.
+        Most keys carry no alias, so this is the common batch, not an edge case."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias=None),
+            user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key_alias"] == "Alice Chen"
+        assert metadata["user_api_key"] == "hash-alice"
+
+    @pytest.mark.asyncio
+    async def test_rotated_key_keeps_the_creating_user_alias(self):
+        """Batches outlive keys. When the creating key has been rotated or deleted the
+        lookup returns no row, and the spend log keeps a resolvable name instead of null."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=None,
+            user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key_alias"] == "Alice Chen"
+
+    @pytest.mark.asyncio
+    async def test_named_key_still_owns_the_alias(self):
+        """The fallback must not weaken the intended precedence: a key that has its own
+        alias still overrides the creating user's."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias="prod-key"),
+            user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key_alias"] == "prod-key"

tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py

@@ -356,3 +356,142 @@ async def test_ensure_batch_response_returns_early_without_auth():
 
     assert response.output_file_id == "file-raw-output"
     mock_managed_files.get_unified_output_file_id.assert_not_called()
+
+
+def _in_memory_managed_files():
+    """Build a real _PROXY_LiteLLMManagedFiles whose prisma upsert hits an in-memory row."""
+    from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
+
+    store: dict = {}
+
+    async def _upsert(where, data):
+        key = where["unified_object_id"]
+        if key in store:
+            store[key].update(data["update"])
+        else:
+            store[key] = dict(data["create"])
+
+    table = MagicMock()
+    table.upsert = AsyncMock(side_effect=_upsert)
+    prisma = MagicMock()
+    prisma.db.litellm_managedobjecttable = table
+
+    cache = MagicMock()
+    cache.async_set_cache = AsyncMock()
+
+    return (
+        _PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma),
+        store,
+    )
+
+
+@pytest.mark.asyncio
+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")
+
+    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,
+        request_tags=["env:prod"],
+        persist_attribution=True,
+    )
+
+    row = store["unified-b"]
+    assert row["api_key"] == "hash-alice"
+    assert row["created_by"] == "alice"
+    assert row["team_id"] == "team-alpha"
+    assert row["request_tags"].data == ["env:prod"]
+
+
+@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
+    generic post-call hook on a retrieve) carries a real hashed key, but must never have it
+    recorded as the batch's paying key. created_by/team_id keep their existing behavior."""
+    instance, store = _in_memory_managed_files()
+    poller = UserAPIKeyAuth(user_id="bob", team_id="team-bravo", api_key="hash-bob")
+
+    await instance.store_unified_object_id(
+        unified_object_id="unified-b",
+        file_object=_build_batch_response(batch_id="b", status="in_progress"),
+        litellm_parent_otel_span=None,
+        model_object_id="b",
+        file_purpose="batch",
+        user_api_key_dict=poller,
+        request_tags=["env:dev"],
+    )
+
+    row = store["unified-b"]
+    assert "api_key" not in row
+    assert "request_tags" not in row
+    assert row["created_by"] == "bob"
+
+
+@pytest.mark.asyncio
+async def test_store_unified_object_id_attribution_columns_are_write_once():
+    """Identity is written only in the upsert create branch, so a later store for the same
+    batch (a status update, a poll) can neither reassign the paying key nor clear it."""
+    instance, store = _in_memory_managed_files()
+    creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice")
+    poller = UserAPIKeyAuth(user_id="bob", team_id="team-bravo", api_key="hash-bob")
+
+    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,
+        request_tags=["env:prod"],
+        persist_attribution=True,
+    )
+    await instance.store_unified_object_id(
+        unified_object_id="unified-b",
+        file_object=_build_batch_response(batch_id="b", status="completed"),
+        litellm_parent_otel_span=None,
+        model_object_id="b",
+        file_purpose="batch",
+        user_api_key_dict=poller,
+        request_tags=["poller-tag"],
+        persist_attribution=True,
+    )
+
+    row = store["unified-b"]
+    assert row["api_key"] == "hash-alice"
+    assert row["created_by"] == "alice"
+    assert row["status"] == "completed"
+
+    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"]
+
+
+@pytest.mark.asyncio
+async def test_store_unified_object_id_omits_unset_columns():
+    """A batch created with no tags (the common case) still registers: the optional columns
+    are omitted rather than passed as None, which prisma rejects for the Json column."""
+    instance, store = _in_memory_managed_files()
+    creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key=None)
+
+    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,
+        request_tags=None,
+        persist_attribution=True,
+    )
+
+    create_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"]["create"]
+    assert "api_key" not in create_data
+    assert "request_tags" not in create_data
+    assert "unified-b" in store

tests/test_litellm/enterprise/proxy/test_managed_files_hook.py

@@ -510,6 +510,117 @@ def _make_real_managed_files_instance():
     )
 
 
+def _make_object_store_instance():
+    """A real store_unified_object_id over an AsyncMock prisma client, so both the
+    upsert and the update-only write path 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_managedobjecttable.upsert = AsyncMock()
+    mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock()
+
+    return (
+        _PROXY_LiteLLMManagedFiles(
+            internal_usage_cache=mock_cache,
+            prisma_client=mock_prisma,
+        ),
+        mock_prisma,
+    )
+
+
+@pytest.mark.asyncio
+async def test_poll_refreshes_batch_state_without_claiming_the_row():
+    """Regression (stale batch state): a poll observes a batch it did not create, so it
+    must still refresh status and file_object -- otherwise GET /v1/batches serves the
+    create-time snapshot forever -- while writing none of the attribution columns and
+    never creating a row it would then own."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    poller = UserAPIKeyAuth(
+        api_key="sk-the-poller", user_id="bob", team_id="team-bravo", parent_otel_span=None
+    )
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="uoi-1",
+        file_object=_make_batch_response(status="completed"),
+        litellm_parent_otel_span=None,
+        model_object_id="batch-123",
+        file_purpose="batch",
+        user_api_key_dict=poller,
+        request_tags=("poller:tag",),
+        persist_attribution=False,
+        create_if_missing=False,
+    )
+
+    # the row is refreshed in place, and cannot be conjured by a poll
+    mock_prisma.db.litellm_managedobjecttable.upsert.assert_not_awaited()
+    update_many = mock_prisma.db.litellm_managedobjecttable.update_many
+    update_many.assert_awaited_once()
+    call = update_many.await_args
+    assert call.kwargs["where"] == {"unified_object_id": "uoi-1"}
+
+    written = call.kwargs["data"]
+    assert written["status"] == "completed"
+    assert json.loads(written["file_object"])["output_file_id"] == "file-output-abc"
+    # nothing the poller could be billed for
+    for owned in ("api_key", "request_tags", "created_by", "team_id"):
+        assert owned not in written
+
+
+@pytest.mark.asyncio
+async def test_create_still_upserts_and_claims_attribution():
+    """The create is the one caller that can speak for the batch, so it keeps the upsert
+    (creating the row when absent) and writes the attribution columns."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    creator = UserAPIKeyAuth(
+        api_key="sk-the-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None
+    )
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="uoi-2",
+        file_object=_make_batch_response(status="validating"),
+        litellm_parent_otel_span=None,
+        model_object_id="batch-456",
+        file_purpose="batch",
+        user_api_key_dict=creator,
+        request_tags=("env:prod",),
+        persist_attribution=True,
+    )
+
+    mock_prisma.db.litellm_managedobjecttable.update_many.assert_not_awaited()
+    upsert = mock_prisma.db.litellm_managedobjecttable.upsert
+    upsert.assert_awaited_once()
+    created = upsert.await_args.kwargs["data"]["create"]
+    # UserAPIKeyAuth hashes an sk- token on construction; the hash is what is billed
+    assert created["api_key"] == creator.api_key
+    assert created["api_key"] != "sk-the-creator"
+    assert created["created_by"] == "alice"
+    assert created["team_id"] == "team-alpha"
+
+
+@pytest.mark.asyncio
+async def test_default_callers_still_create_their_rows():
+    """create_if_missing defaults to True, so the fine-tune, Responses and Anthropic
+    callers, none of which pass it, keep upserting exactly as before."""
+    managed_files, mock_prisma = _make_object_store_instance()
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="uoi-3",
+        file_object=_make_batch_response(),
+        litellm_parent_otel_span=None,
+        model_object_id="ft-789",
+        file_purpose="fine-tune",
+        user_api_key_dict=_make_user_api_key_dict(),
+    )
+
+    mock_prisma.db.litellm_managedobjecttable.upsert.assert_awaited_once()
+    mock_prisma.db.litellm_managedobjecttable.update_many.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_store_unified_file_id_is_idempotent_via_upsert():
     """Regression test for the managed-batch retrieve 500 (UniqueViolationError on

tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py

@@ -17,6 +17,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import (
     _should_track_cost_callback,
     _update_database_and_spend_counters,
 )
+from litellm.types.utils import CallTypes
 
 
 @pytest.mark.asyncio
@@ -783,6 +784,166 @@ async def test_enrich_failure_metadata_skips_when_no_api_key():
         mock_get_key.assert_not_called()
 
 
+@pytest.mark.asyncio
+async def test_enrich_failure_metadata_keeps_captured_identity_when_not_resolving():
+    """
+    With resolve_missing_key_identity=False the key is not read, so a null user_id,
+    team_id and org_id captured earlier stay null instead of being refilled from the
+    key as it stands now. The team_alias lookup still runs off the captured team_id.
+    """
+    mock_key_obj = MagicMock()
+    mock_key_obj.key_alias = "alias-assigned-later"
+    mock_key_obj.user_id = "user-assigned-later"
+    mock_key_obj.team_id = "team-assigned-later"
+    mock_key_obj.org_id = "org-assigned-later"
+
+    mock_team_obj = MagicMock()
+    mock_team_obj.team_alias = "captured-team-alias"
+
+    with (
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+            new_callable=AsyncMock,
+            return_value=mock_key_obj,
+        ) as mock_get_key,
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+            new_callable=AsyncMock,
+            return_value=mock_team_obj,
+        ),
+    ):
+        metadata = {
+            "user_api_key": "hashed_key",
+            "user_api_key_alias": None,
+            "user_api_key_user_id": None,
+            "user_api_key_team_id": "captured-team-id",
+            "user_api_key_team_alias": None,
+            "user_api_key_org_id": None,
+        }
+        result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(
+            metadata, resolve_missing_key_identity=False
+        )
+
+        mock_get_key.assert_not_called()
+        assert result["user_api_key_user_id"] is None
+        assert result["user_api_key_team_id"] == "captured-team-id"
+        assert result["user_api_key_org_id"] is None
+        assert result["user_api_key_alias"] is None
+        assert result["user_api_key_team_alias"] == "captured-team-alias"
+
+
+@pytest.mark.asyncio
+async def test_enrich_failure_metadata_ignores_flag_when_alias_present():
+    """
+    A captured alias already closes the key lookup, so resolve_missing_key_identity
+    changes nothing for a key that has one; only the alias-less key depends on it.
+    """
+    mock_team_obj = MagicMock()
+    mock_team_obj.team_alias = "captured-team-alias"
+
+    with (
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+            new_callable=AsyncMock,
+        ) as mock_get_key,
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+            new_callable=AsyncMock,
+            return_value=mock_team_obj,
+        ),
+    ):
+        for resolve in (True, False):
+            metadata = {
+                "user_api_key": "hashed_key",
+                "user_api_key_alias": "captured-alias",
+                "user_api_key_user_id": None,
+                "user_api_key_team_id": "captured-team-id",
+                "user_api_key_team_alias": None,
+                "user_api_key_org_id": None,
+            }
+            result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(
+                metadata, resolve_missing_key_identity=resolve
+            )
+            mock_get_key.assert_not_called()
+            assert result["user_api_key_user_id"] is None
+            assert result["user_api_key_alias"] == "captured-alias"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "call_type, expect_key_read",
+    [
+        (CallTypes.aretrieve_batch.value, False),
+        (CallTypes.aretrieve_batch, False),
+        (CallTypes.acompletion.value, True),
+    ],
+)
+async def test_track_cost_callback_reads_key_only_for_in_request_logs(call_type, expect_key_read):
+    """
+    The batch cost row is logged long after the batch was created, so it keeps the
+    identity persisted at create time. Every other call type still backfills from
+    the key.
+    """
+    logger = _ProxyDBLogger()
+
+    mock_key_obj = MagicMock()
+    mock_key_obj.key_alias = "alias-assigned-later"
+    mock_key_obj.user_id = "user-assigned-later"
+    mock_key_obj.team_id = "team-assigned-later"
+    mock_key_obj.org_id = "org-assigned-later"
+
+    kwargs = {
+        "call_type": call_type,
+        "model": None,
+        "litellm_call_id": "test-call-id",
+        "stream": False,
+        "litellm_params": {
+            "metadata": {
+                "user_api_key": "hashed_key",
+                "user_api_key_alias": None,
+                "user_api_key_user_id": None,
+                "user_api_key_team_id": None,
+                "user_api_key_org_id": None,
+            }
+        },
+    }
+
+    with (
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
+            new_callable=AsyncMock,
+            return_value=mock_key_obj,
+        ) as mock_get_key,
+        patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
+            new_callable=AsyncMock,
+            return_value=MagicMock(team_alias=None),
+        ),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
+    ):
+        mock_proxy_logging.failed_tracking_alert = AsyncMock()
+        mock_proxy_logging.db_spend_update_writer = MagicMock()
+        mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
+
+        await logger._PROXY_track_cost_callback(
+            kwargs=kwargs,
+            completion_response=None,
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+        )
+
+    assert mock_get_key.called is expect_key_read
+
+    written = kwargs["litellm_params"]["metadata"]
+    if expect_key_read:
+        assert written["user_api_key_user_id"] == "user-assigned-later"
+        assert written["user_api_key_team_id"] == "team-assigned-later"
+    else:
+        assert written["user_api_key_user_id"] is None
+        assert written["user_api_key_team_id"] is None
+        assert written["user_api_key_org_id"] is None
+
+
 @pytest.mark.asyncio
 async def test_async_post_call_failure_hook_enriches_auth_error_metadata():
     """

tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py

@@ -258,6 +258,7 @@ class TestVertexAIBatchPassthroughHandler:
                     batch_object=batch_object,
                     model_object_id=model_object_id,
                     logging_obj=mock_logging_obj,
+                    is_batch_create=True,
                     user_api_key_dict={"user_id": "test-user"},
                 )
 
@@ -307,6 +308,7 @@ class TestVertexAIBatchPassthroughHandler:
                 batch_object={"id": "b1", "object": "batch", "status": "validating"},
                 model_object_id="b1",
                 logging_obj=mock_logging_obj,
+                is_batch_create=True,
                 **kwargs,
             )
 
@@ -315,6 +317,140 @@ class TestVertexAIBatchPassthroughHandler:
             assert call_kwargs["user_api_key_dict"].user_id == expected_user_id
             assert call_kwargs["user_api_key_dict"].team_id == expected_team_id
 
+    def _store_with_metadata(self, mock_logging_obj, mock_managed_files_hook, metadata):
+        with (
+            patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl,
+            patch(
+                "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger"
+            ),
+        ):
+            mock_pl.get_proxy_hook.return_value = mock_managed_files_hook
+            VertexPassthroughLoggingHandler._store_batch_managed_object(
+                unified_object_id="uoi",
+                batch_object={"id": "b1", "object": "batch", "status": "validating"},
+                model_object_id="b1",
+                logging_obj=mock_logging_obj,
+                is_batch_create=True,
+                litellm_params={"metadata": metadata},
+            )
+        mock_managed_files_hook.store_unified_object_id.assert_called_once()
+        return mock_managed_files_hook.store_unified_object_id.call_args[1]
+
+    def test_create_persists_key_hash_and_tags(
+        self, mock_logging_obj, mock_managed_files_hook
+    ):
+        """Regression (spend loss): the batch create must persist the creating key's hashed
+        token and its tags so CheckBatchCost can attribute the batch-cost spend row. Before
+        this fix the stored api_key was always "" and the row was dropped as unattributed."""
+        call_kwargs = self._store_with_metadata(
+            mock_logging_obj,
+            mock_managed_files_hook,
+            {
+                "user_api_key": "hashed-key-a",
+                "user_api_key_user_id": "alice",
+                "user_api_key_team_id": "team-alpha",
+                "user_api_key_auth_metadata": {"tags": ["env:prod", 7, "team:ml"]},
+            },
+        )
+
+        assert call_kwargs["user_api_key_dict"].api_key == "hashed-key-a"
+        # non-string tags are dropped so downstream tag budgets cannot be bypassed
+        assert call_kwargs["request_tags"] == ("env:prod", "team:ml")
+        assert call_kwargs["persist_attribution"] is True
+
+    @pytest.mark.parametrize(
+        "metadata, expected",
+        [
+            # a request that sent its own tags (x-litellm-tags header or body metadata)
+            ({"tags": ["req:a", "req:b"]}, ("req:a", "req:b")),
+            # request tags win over the key's own tags
+            (
+                {"tags": ["req:a"], "user_api_key_auth_metadata": {"tags": ["key:b"]}},
+                ("req:a",),
+            ),
+            # no request tags: fall back to the tags the key itself carries
+            ({"user_api_key_auth_metadata": {"tags": ["key:b"]}}, ("key:b",)),
+            # neither: no tags on the spend row
+            ({}, None),
+        ],
+    )
+    def test_request_tags_precedence(
+        self, mock_logging_obj, mock_managed_files_hook, metadata, expected
+    ):
+        """Request tags take precedence over the key's tags, and the key's tags are the
+        fallback because a tagged key does not put its tags in the top-level metadata."""
+        call_kwargs = self._store_with_metadata(
+            mock_logging_obj,
+            mock_managed_files_hook,
+            {"user_api_key": "hashed-key-a", **metadata},
+        )
+
+        assert call_kwargs["request_tags"] == expected
+
+    @pytest.mark.parametrize(
+        "url_route, expected",
+        [
+            ("/v1/projects/p/locations/us-central1/batchPredictionJobs", True),
+            ("/v1/projects/p/locations/us-central1/batchPredictionJobs/", True),
+            ("/v1/projects/p/locations/us-central1/batchPredictionJobs?alt=json", True),
+            ("/v1/projects/p/locations/us-central1/batchPredictionJobs/123456", False),
+            ("/v1/projects/p/locations/us-central1/batchPredictionJobs/123456?alt=json", False),
+        ],
+    )
+    def test_batch_is_registered_from_the_create_route_only(
+        self, mock_logging_obj, url_route, expected
+    ):
+        """Only a POST to the collection route is the create, and only the create claims
+        attribution. Every id-scoped route is a poll or retrieve, which still reports the
+        batch so its status and file object stay in sync, but carries is_batch_create=False
+        so it neither claims the batch nor creates a row it would then own."""
+        response = MagicMock()
+        response.status_code = 200
+        response.json.return_value = {
+            "name": "projects/p/locations/us-central1/batchPredictionJobs/123456",
+            "model": "publishers/google/models/gemini-2.5-flash",
+        }
+
+        with (
+            patch(
+                "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger"
+            ),
+            patch(
+                "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.VertexPassthroughLoggingHandler._store_batch_managed_object"
+            ) as mock_store,
+            patch(
+                "litellm.llms.vertex_ai.batches.transformation.VertexAIBatchTransformation"
+            ) as mock_transformation,
+            patch(
+                "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.VertexPassthroughLoggingHandler.get_actual_model_id_from_router",
+                return_value="gemini-2.5-flash",
+            ),
+        ):
+            mock_transformation.transform_vertex_ai_batch_response_to_openai_batch_response.return_value = {
+                "id": "123456",
+                "object": "batch",
+                "status": "validating",
+                "created_at": 1704067200,
+                "input_file_id": "gs://bucket/in.jsonl",
+                "completion_window": "24h",
+            }
+            mock_transformation._get_batch_id_from_vertex_ai_batch_response.return_value = "123456"
+
+            VertexPassthroughLoggingHandler.batch_prediction_jobs_handler(
+                httpx_response=response,
+                logging_obj=mock_logging_obj,
+                url_route=url_route,
+                result="",
+                start_time=datetime.now(),
+                end_time=datetime.now(),
+                cache_hit=False,
+            )
+
+        # every route reports the batch; only the create claims it
+        mock_store.assert_called_once()
+        assert mock_store.call_args[1]["unified_object_id"]
+        assert mock_store.call_args[1]["is_batch_create"] is expected
+
     def test_batch_cost_calculation_integration(self):
         """Single Vertex AI response → non-zero cost with correct token counts."""
         from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage