fix(proxy): attribute vertex batch passthrough cost to creating key/team/tags

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -61,6 +61,45 @@ class CheckBatchCost:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
             return {}
 
+    async def _build_spend_attribution_metadata(
+        self,
+        job: "LiteLLM_ManagedObjectTable",
+        creator_user_id: Optional[str],
+        user_info: dict,
+    ) -> dict:
+        """
+        Build spend-log metadata that attributes the batch cost to the virtual key, team,
+        and tags captured at batch-create time. Key and team aliases are resolved via the
+        shared cost-callback enrichment so the row matches a non-batch request's attribution.
+
+        Rows created before api_key was persisted (legacy) fall back to the previous
+        created_by/user_info-only behavior.
+        """
+        from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
+
+        api_key = job.api_key if isinstance(job.api_key, str) and job.api_key else None
+        team_id = job.team_id if isinstance(job.team_id, str) and job.team_id else None
+        request_tags = job.request_tags if isinstance(job.request_tags, list) else None
+
+        if not api_key:
+            return {
+                "user_api_key_user_id": creator_user_id,
+                **user_info,
+            }
+
+        metadata = {
+            "user_api_key_user_id": creator_user_id,
+            "user_api_key": api_key,
+            **({"user_api_key_team_id": team_id} if team_id else {}),
+            **({"tags": request_tags} if request_tags else {}),
+        }
+        metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata)
+
+        user_email = user_info.get("user_api_key_user_email")
+        if user_email is not None and metadata.get("user_api_key_user_email") is None:
+            metadata = {**metadata, "user_api_key_user_email": user_email}
+        return metadata
+
     async def _cleanup_stale_managed_objects(self) -> None:
         """
         Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@@ -460,6 +499,11 @@ class CheckBatchCost:
 
         creator_user_id = job.created_by
         user_info = await self._get_user_info(batch_id, job.created_by)
+        spend_metadata = await self._build_spend_attribution_metadata(
+            job=job,
+            creator_user_id=creator_user_id,
+            user_info=user_info,
+        )
 
         logging_obj.update_environment_variables(
             litellm_params={
@@ -469,10 +513,7 @@ class CheckBatchCost:
                         "user-agent": CHECK_BATCH_COST_USER_AGENT,
                     }
                 },
-                "metadata": {
-                    "user_api_key_user_id": creator_user_id,
-                    **user_info,
-                },
+                "metadata": spend_metadata,
             },
             optional_params={},
         )

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -165,7 +165,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         model_object_id: str,
         file_purpose: Literal["batch", "fine-tune", "response"],
         user_api_key_dict: UserAPIKeyAuth,
+        request_tags: Optional[List[str]] = None,
     ) -> None:
+        from prisma import Json
+
         verbose_logger.info(
             f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache"
         )
@@ -191,6 +194,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     "file_purpose": file_purpose,
                     "created_by": user_api_key_dict.user_id,
                     "team_id": user_api_key_dict.team_id,
+                    "api_key": user_api_key_dict.api_key or None,
+                    "request_tags": Json(request_tags) if request_tags else None,
                     "updated_by": user_api_key_dict.user_id,
                     "status": file_object.status,
                 },

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

@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "api_key" TEXT;
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -942,6 +942,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?
   updated_at DateTime @updatedAt
   updated_by String?
 

litellm/models/managed_files.py

@@ -32,6 +32,8 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
     file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse]
     created_by: Optional[str] = None
     team_id: Optional[str] = None
+    api_key: Optional[str] = None
+    request_tags: Optional[list[str]] = None
 
 
 class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase):

litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py

@@ -978,11 +978,12 @@ class AnthropicPassthroughLoggingHandler:
 
                 _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {}
 
+                _request_tags = _request_metadata.get("tags")
                 user_api_key_dict = UserAPIKeyAuth(
                     user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
-                    api_key="",
+                    api_key=_request_metadata.get("user_api_key") or "",
                     team_id=_request_metadata.get("user_api_key_team_id"),
-                    team_alias=None,
+                    team_alias=_request_metadata.get("user_api_key_team_alias"),
                     user_role=LitellmUserRoles.CUSTOMER,  # Use proper enum value
                     user_email=None,
                     max_budget=None,
@@ -995,7 +996,7 @@ class AnthropicPassthroughLoggingHandler:
                     max_parallel_requests=None,
                     allowed_model_region=None,
                     metadata={},  # Set to empty dict instead of None
-                    key_alias=None,
+                    key_alias=_request_metadata.get("user_api_key_alias"),
                     permissions={},  # Set to empty dict instead of None
                     model_max_budget={},  # Set to empty dict instead of None
                     model_spend={},  # Set to empty dict instead of None
@@ -1012,6 +1013,7 @@ class AnthropicPassthroughLoggingHandler:
                         model_object_id=model_object_id,
                         file_purpose="batch",
                         user_api_key_dict=user_api_key_dict,
+                        request_tags=_request_tags if isinstance(_request_tags, list) else None,
                     )
                 )
 

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -803,11 +803,12 @@ class VertexPassthroughLoggingHandler:
 
                 _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {}
 
+                _request_tags = _request_metadata.get("tags")
                 user_api_key_dict = UserAPIKeyAuth(
                     user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
-                    api_key="",
+                    api_key=_request_metadata.get("user_api_key") or "",
                     team_id=_request_metadata.get("user_api_key_team_id"),
-                    team_alias=None,
+                    team_alias=_request_metadata.get("user_api_key_team_alias"),
                     user_role=LitellmUserRoles.CUSTOMER,  # Use proper enum value
                     user_email=None,
                     max_budget=None,
@@ -820,7 +821,7 @@ class VertexPassthroughLoggingHandler:
                     max_parallel_requests=None,
                     allowed_model_region=None,
                     metadata={},  # Set to empty dict instead of None
-                    key_alias=None,
+                    key_alias=_request_metadata.get("user_api_key_alias"),
                     permissions={},  # Set to empty dict instead of None
                     model_max_budget={},  # Set to empty dict instead of None
                     model_spend={},  # Set to empty dict instead of None
@@ -837,6 +838,7 @@ class VertexPassthroughLoggingHandler:
                         model_object_id=model_object_id,
                         file_purpose="batch",
                         user_api_key_dict=user_api_key_dict,
+                        request_tags=_request_tags if isinstance(_request_tags, list) else None,
                     )
                 )
 

litellm/proxy/schema.prisma

@@ -942,6 +942,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?
   updated_at DateTime @updatedAt
   updated_by String?
 

schema.prisma

@@ -942,6 +942,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?
   updated_at DateTime @updatedAt
   updated_by String?
 

tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py

@@ -2367,3 +2367,104 @@ async def test_same_user_different_keys_can_access_batch():
     assert "batch_id" in result2
     # Both keys should get the same result
     assert result1["batch_id"] == result2["batch_id"]
+
+
+@pytest.mark.asyncio
+async def test_store_unified_object_id_persists_request_identity():
+    """
+    store_unified_object_id must persist the creating key's hashed token, team, and request
+    tags onto the managed object so CheckBatchCost can later attribute the batch-cost SpendLogs
+    row. Regression for batches whose cost row was dropped because identity was never carried.
+    """
+    from prisma import Json
+
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.utils import LiteLLMBatch
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.upsert = AsyncMock(
+        return_value=MagicMock()
+    )
+    internal_usage_cache = MagicMock()
+    internal_usage_cache.async_set_cache = AsyncMock()
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        internal_usage_cache=internal_usage_cache,
+        prisma_client=prisma_client,
+    )
+
+    batch_object = LiteLLMBatch(
+        id="batch-123",
+        completion_window="24h",
+        created_at=1,
+        endpoint="/v1/chat/completions",
+        input_file_id="gs://bucket/in.jsonl",
+        object="batch",
+        status="validating",
+    )
+
+    await proxy_managed_files.store_unified_object_id(
+        unified_object_id="unified-123",
+        file_object=batch_object,
+        litellm_parent_otel_span=None,
+        model_object_id="batch-123",
+        file_purpose="batch",
+        user_api_key_dict=UserAPIKeyAuth(
+            user_id="user-1", team_id="team-1", api_key="hashed-key-abc"
+        ),
+        request_tags=["tag-a", "tag-b"],
+    )
+
+    create_data = prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs[
+        "data"
+    ]["create"]
+    assert create_data["api_key"] == "hashed-key-abc"
+    assert create_data["created_by"] == "user-1"
+    assert create_data["team_id"] == "team-1"
+    assert create_data["request_tags"] == Json(["tag-a", "tag-b"])
+
+
+@pytest.mark.asyncio
+async def test_store_unified_object_id_no_tags_stores_null_request_tags():
+    """When no request tags are present the column is left NULL rather than an empty-list Json."""
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.utils import LiteLLMBatch
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.upsert = AsyncMock(
+        return_value=MagicMock()
+    )
+    internal_usage_cache = MagicMock()
+    internal_usage_cache.async_set_cache = AsyncMock()
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        internal_usage_cache=internal_usage_cache,
+        prisma_client=prisma_client,
+    )
+
+    batch_object = LiteLLMBatch(
+        id="batch-123",
+        completion_window="24h",
+        created_at=1,
+        endpoint="/v1/chat/completions",
+        input_file_id="gs://bucket/in.jsonl",
+        object="batch",
+        status="validating",
+    )
+
+    await proxy_managed_files.store_unified_object_id(
+        unified_object_id="unified-123",
+        file_object=batch_object,
+        litellm_parent_otel_span=None,
+        model_object_id="batch-123",
+        file_purpose="batch",
+        user_api_key_dict=UserAPIKeyAuth(user_id="user-1", api_key="hashed-key-abc"),
+    )
+
+    create_data = prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs[
+        "data"
+    ]["create"]
+    assert create_data["request_tags"] is None
+    assert create_data["api_key"] == "hashed-key-abc"
+
+

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -1221,3 +1221,87 @@ class TestUnmanagedBatchCostFlagIsGeneralized:
 
         assert vertex_result == ("deploy-vertex", "8823717160934178816")
         assert bedrock_result == ("deploy-bedrock", TestUnmanagedBedrockRouting._ARN)
+
+
+class _Job:
+    """Minimal stand-in for a LiteLLM_ManagedObjectTable row with just the identity
+    fields _build_spend_attribution_metadata reads."""
+
+    def __init__(self, created_by=None, api_key=None, team_id=None, request_tags=None):
+        self.created_by = created_by
+        self.api_key = api_key
+        self.team_id = team_id
+        self.request_tags = request_tags
+
+
+class TestSpendAttributionMetadata:
+    """_build_spend_attribution_metadata must carry the create-time identity into the
+    spend-log metadata so the batch-cost SpendLogs row is written and attributed to the
+    creating key/team/tags (issue #33316)."""
+
+    def _instance(self):
+        from litellm_enterprise.proxy.common_utils.check_batch_cost import (
+            CheckBatchCost,
+        )
+
+        return CheckBatchCost(
+            proxy_logging_obj=MagicMock(),
+            prisma_client=MagicMock(),
+            llm_router=MagicMock(),
+        )
+
+    @pytest.mark.asyncio
+    async def test_carries_key_team_and_tags_and_resolves_aliases(self):
+        instance = self._instance()
+        job = _Job(
+            created_by="user-1",
+            api_key="hashed-key-abc",
+            team_id="team-1",
+            request_tags=["tag-a", "tag-b"],
+        )
+
+        async def _fake_enrich(metadata):
+            return {
+                **metadata,
+                "user_api_key_alias": "my-key",
+                "user_api_key_team_alias": "my-team",
+            }
+
+        with patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback._ProxyDBLogger._enrich_failure_metadata_with_key_info",
+            side_effect=_fake_enrich,
+        ):
+            metadata = await instance._build_spend_attribution_metadata(
+                job=job,
+                creator_user_id="user-1",
+                user_info={"user_api_key_user_email": "u@example.com"},
+            )
+
+        assert metadata["user_api_key"] == "hashed-key-abc"
+        assert metadata["user_api_key_user_id"] == "user-1"
+        assert metadata["user_api_key_team_id"] == "team-1"
+        assert metadata["tags"] == ["tag-a", "tag-b"]
+        assert metadata["user_api_key_alias"] == "my-key"
+        assert metadata["user_api_key_team_alias"] == "my-team"
+        assert metadata["user_api_key_user_email"] == "u@example.com"
+
+    @pytest.mark.asyncio
+    async def test_legacy_row_without_api_key_falls_back(self):
+        """Rows created before api_key was persisted keep the previous created_by/user_info
+        behavior and never invoke key enrichment."""
+        instance = self._instance()
+        job = _Job(created_by="user-1", api_key=None, team_id=None, request_tags=None)
+
+        with patch(
+            "litellm.proxy.hooks.proxy_track_cost_callback._ProxyDBLogger._enrich_failure_metadata_with_key_info",
+        ) as mock_enrich:
+            metadata = await instance._build_spend_attribution_metadata(
+                job=job,
+                creator_user_id="user-1",
+                user_info={"user_api_key_user_email": "u@example.com"},
+            )
+
+        mock_enrich.assert_not_called()
+        assert "user_api_key" not in metadata
+        assert metadata["user_api_key_user_id"] == "user-1"
+        assert metadata["user_api_key_user_email"] == "u@example.com"

tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py

@@ -315,6 +315,46 @@ 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 test_store_batch_managed_object_propagates_key_aliases_and_tags(
+        self, mock_logging_obj, mock_managed_files_hook
+    ):
+        """The fabricated UserAPIKeyAuth must carry the hashed key and aliases, and the request
+        tags must be forwarded, so CheckBatchCost can later write an attributed batch-cost
+        SpendLogs row. Regression for issue #33316 where api_key was hardcoded to "" and tags
+        were dropped, causing the cost row to be silently discarded."""
+        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,
+                litellm_params={
+                    "metadata": {
+                        "user_api_key": "hashed-key-abc",
+                        "user_api_key_user_id": "real-user-123",
+                        "user_api_key_team_id": "team-456",
+                        "user_api_key_alias": "my-key",
+                        "user_api_key_team_alias": "my-team",
+                        "tags": ["tag-a", "tag-b"],
+                    }
+                },
+            )
+
+            mock_managed_files_hook.store_unified_object_id.assert_called_once()
+            call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1]
+            user_api_key_dict = call_kwargs["user_api_key_dict"]
+            assert user_api_key_dict.api_key == "hashed-key-abc"
+            assert user_api_key_dict.key_alias == "my-key"
+            assert user_api_key_dict.team_alias == "my-team"
+            assert call_kwargs["request_tags"] == ["tag-a", "tag-b"]
+
     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