From 90e941a3f8749eb037655b91fa0c689971ed47a0 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 15 Jul 2026 13:21:05 -0700 Subject: [PATCH] fix(proxy): attribute passthrough batch-cost spend to creating key/team/tags Vertex and Anthropic passthrough batches built a UserAPIKeyAuth with an empty api_key and forwarded no tags, so the managed-object row that CheckBatchCost later reads carried no identity. At cost time the poller had a blank key, so the batch-cost SpendLogs row was attributed to nobody (or dropped), leaving the creating key at $0 with blank api_key/team_id/request_tags. The request metadata already carries the authenticated hashed key, the user/team ids and the request tags; the old handlers threw them away. Add user_api_key and request_tags columns to LiteLLM_ManagedObjectTable (schema, migration, Pydantic model) and persist them in store_unified_object_id, keeping the identity in the DB row rather than the cached pydantic model. Both passthrough handlers now go through one shared helper, store_batch_managed_object, that reads the real identity out of the request metadata. At cost time CheckBatchCost._build_creator_attribution_metadata rebuilds the spend-tracking metadata from the row, resolving the key alias from the hashed token and the team alias from the team id, so the batch-cost row is attributed exactly like a non-batch request. _get_user_info short-circuits a null user_id so the poller no longer issues find_unique(where={"user_id": None}), and every lookup falls back gracefully for older rows that predate the two new columns. --- .../proxy/common_utils/check_batch_cost.py | 77 +++++++++-- .../proxy/hooks/managed_files.py | 6 + .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/managed_files.py | 2 + .../anthropic_passthrough_logging_handler.py | 69 ++-------- .../batch_managed_object_utils.py | 96 ++++++++++++++ .../vertex_passthrough_logging_handler.py | 69 ++-------- litellm/proxy/schema.prisma | 2 + schema.prisma | 2 + .../proxy_unit_tests/test_check_batch_cost.py | 123 ++++++++++++++++++ .../test_vertex_ai_batch_passthrough.py | 37 ++++++ 12 files changed, 362 insertions(+), 128 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260715020000_add_user_api_key_and_request_tags_to_managed_object_table/migration.sql create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_managed_object_utils.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index f209ab54f64..ae4ab5e9552 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -42,11 +42,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, object]: """ - 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). + Look up user email by user_id for enriching the S3 callback metadata. + Returns a dict with user_api_key_user_email (may be None). Returns an + empty dict when user_id is None (batches created by team/service-account + keys 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} @@ -55,12 +59,69 @@ class CheckBatchCost: return {} return { "user_api_key_user_email": getattr(user_row, "user_email", None), - "user_api_key_alias": getattr(user_row, "user_alias", None), } except Exception as e: 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, user_api_key: Optional[str]) -> Optional[str]: + """Resolve the creating virtual key's alias from its hashed token.""" + if not user_api_key: + return None + try: + key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": user_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 cost attribution: {e}") + return None + + async def _get_team_alias(self, team_id: Optional[str]) -> Optional[str]: + """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, object]: + """ + Rebuild the spend-tracking metadata for the key/team/tags that created the + batch so the batch-cost spend log is attributed the same way a non-batch + request is. Falls back gracefully to whatever identity the managed object + row carries (older rows created before user_api_key/request_tags were + persisted only have created_by/team_id). + """ + user_api_key = getattr(job, "user_api_key", None) + team_id = getattr(job, "team_id", None) + request_tags = getattr(job, "request_tags", None) + + metadata: dict[str, object] = { + "user_api_key_user_id": job.created_by, + "user_api_key": user_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(user_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"] = request_tags + + return metadata + async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days @@ -458,8 +519,7 @@ 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) + attribution_metadata = await self._build_creator_attribution_metadata(job, batch_id) logging_obj.update_environment_variables( litellm_params={ @@ -469,10 +529,7 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": { - "user_api_key_user_id": creator_user_id, - **user_info, - }, + "metadata": attribution_metadata, }, optional_params={}, ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 3f42867d90e..c30950e647f 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -165,10 +165,14 @@ 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: verbose_logger.info( f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" ) + from prisma import Json + + user_api_key = user_api_key_dict.api_key or None litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -191,6 +195,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, + "user_api_key": user_api_key, + "request_tags": Json(request_tags) if request_tags else None, "updated_by": user_api_key_dict.user_id, "status": file_object.status, }, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715020000_add_user_api_key_and_request_tags_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715020000_add_user_api_key_and_request_tags_to_managed_object_table/migration.sql new file mode 100644 index 00000000000..f87c759c712 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260715020000_add_user_api_key_and_request_tags_to_managed_object_table/migration.sql @@ -0,0 +1,5 @@ +-- Add user_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 "user_api_key" TEXT; +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a23cecc3911..96c5a880a24 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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? + user_api_key String? + request_tags Json? updated_at DateTime @updatedAt updated_by String? diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 99ba764dd98..b7e4658ed36 100644 --- a/litellm/models/managed_files.py +++ b/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 + user_api_key: Optional[str] = None + request_tags: Optional[list[str]] = None class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 50e90699194..6beb0b4539d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -19,6 +19,9 @@ from litellm.llms.anthropic.chat.handler import ( from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_managed_object_utils import ( + store_batch_managed_object, +) from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) @@ -966,65 +969,13 @@ class AnthropicPassthroughLoggingHandler: Store batch managed object for cost tracking. This will be picked up by the check_batch_cost polling mechanism. """ - try: - # Get the managed files hook from the logging object - # This is a bit of a hack, but we need access to the proxy logging system - from litellm.proxy.proxy_server import proxy_logging_obj - - managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None and hasattr(managed_files_hook, "store_unified_object_id"): - # Create a mock user API key dict for the managed object storage - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - - _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} - - user_api_key_dict = UserAPIKeyAuth( - user_id=_request_metadata.get("user_api_key_user_id", "default-user"), - api_key="", - team_id=_request_metadata.get("user_api_key_team_id"), - team_alias=None, - user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value - user_email=None, - max_budget=None, - spend=0.0, # Set to 0.0 instead of None - models=[], # Set to empty list instead of None - tpm_limit=None, - rpm_limit=None, - budget_duration=None, - budget_reset_at=None, - max_parallel_requests=None, - allowed_model_region=None, - metadata={}, # Set to empty dict instead of None - key_alias=None, - 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 - ) - - # Store the unified object for batch cost tracking - import asyncio - - asyncio.create_task( - managed_files_hook.store_unified_object_id( # type: ignore - unified_object_id=unified_object_id, - file_object=batch_object, - litellm_parent_otel_span=None, - model_object_id=model_object_id, - file_purpose="batch", - user_api_key_dict=user_api_key_dict, - ) - ) - - verbose_proxy_logger.info( - f"Stored Anthropic batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}" - ) - else: - verbose_proxy_logger.warning( - "Managed files hook not available, cannot store batch object for cost tracking" - ) - - except Exception as e: - verbose_proxy_logger.error(f"Error storing Anthropic batch managed object: {e}") + request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} + store_batch_managed_object( + unified_object_id=unified_object_id, + batch_object=batch_object, + model_object_id=model_object_id, + request_metadata=request_metadata, + ) @staticmethod def get_actual_model_id_from_router(model_name: str) -> str: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_managed_object_utils.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_managed_object_utils.py new file mode 100644 index 00000000000..78dd9d05ae6 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_managed_object_utils.py @@ -0,0 +1,96 @@ +""" +Shared helper for persisting a passthrough-created batch as a managed object. + +Vertex AI and Anthropic passthrough both create a provider batch job and then +need to register it for asynchronous cost tracking by CheckBatchCost. The +creating key's identity (hashed key, team, tags) is captured here so the +eventual batch-cost spend log can be attributed back to it, mirroring the +attribution a non-batch request receives. +""" + +import asyncio +from typing import TYPE_CHECKING, Optional, Protocol, cast + +from litellm._logging import verbose_proxy_logger +from litellm.types.utils import LiteLLMBatch + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + + +class _StoreUnifiedObjectHook(Protocol): + async def store_unified_object_id( + self, + unified_object_id: str, + file_object: LiteLLMBatch, + litellm_parent_otel_span: None, + model_object_id: str, + file_purpose: str, + user_api_key_dict: "UserAPIKeyAuth", + request_tags: Optional[list[str]], + ) -> None: ... + + +def _optional_str(value: object) -> Optional[str]: + return value if isinstance(value, str) else None + + +def _optional_str_list(value: object) -> Optional[list[str]]: + if isinstance(value, list): + items = cast(list[object], value) # cast-ok: isinstance-narrowed; element type unknown + return [str(tag) for tag in items] + return None + + +def store_batch_managed_object( + unified_object_id: str, + batch_object: LiteLLMBatch, + model_object_id: str, + request_metadata: dict[str, object], +) -> None: + """ + Register a provider batch job as a managed object so CheckBatchCost can + later compute its cost and write an attributed spend log. + + request_metadata is the request's ``litellm_params["metadata"]``, which + carries the authenticated key's ``user_api_key`` hash, user/team ids and + request tags. + """ + try: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import proxy_logging_obj + + managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") + if managed_files_hook is None or not hasattr(managed_files_hook, "store_unified_object_id"): + verbose_proxy_logger.warning( + "Managed files hook not available, cannot store batch object for cost tracking" + ) + return + + user_api_key_dict = UserAPIKeyAuth( + user_id=_optional_str(request_metadata.get("user_api_key_user_id")) or "default-user", + api_key=_optional_str(request_metadata.get("user_api_key")), + team_id=_optional_str(request_metadata.get("user_api_key_team_id")), + team_alias=_optional_str(request_metadata.get("user_api_key_team_alias")), + key_alias=_optional_str(request_metadata.get("user_api_key_alias")), + user_role=LitellmUserRoles.CUSTOMER, + ) + + hook = cast(_StoreUnifiedObjectHook, managed_files_hook) # cast-ok: presence checked via hasattr above + asyncio.create_task( + hook.store_unified_object_id( + unified_object_id=unified_object_id, + file_object=batch_object, + litellm_parent_otel_span=None, + model_object_id=model_object_id, + file_purpose="batch", + user_api_key_dict=user_api_key_dict, + request_tags=_optional_str_list(request_metadata.get("tags")), + ) + ) + + verbose_proxy_logger.info( + f"Stored batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}" + ) + except Exception as e: + verbose_proxy_logger.error(f"Error storing batch managed object: {e}") diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index c61d48eda8c..d6f932fbccd 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -16,6 +16,9 @@ from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( ) from litellm.llms.vertex_ai.videos.transformation import VertexAIVideoConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_managed_object_utils import ( + store_batch_managed_object, +) from litellm.types.utils import ( Choices, EmbeddingResponse, @@ -791,65 +794,13 @@ class VertexPassthroughLoggingHandler: Store batch managed object for cost tracking. This will be picked up by the check_batch_cost polling mechanism. """ - try: - # Get the managed files hook from the logging object - # This is a bit of a hack, but we need access to the proxy logging system - from litellm.proxy.proxy_server import proxy_logging_obj - - managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None and hasattr(managed_files_hook, "store_unified_object_id"): - # Create a mock user API key dict for the managed object storage - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - - _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} - - user_api_key_dict = UserAPIKeyAuth( - user_id=_request_metadata.get("user_api_key_user_id", "default-user"), - api_key="", - team_id=_request_metadata.get("user_api_key_team_id"), - team_alias=None, - user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value - user_email=None, - max_budget=None, - spend=0.0, # Set to 0.0 instead of None - models=[], # Set to empty list instead of None - tpm_limit=None, - rpm_limit=None, - budget_duration=None, - budget_reset_at=None, - max_parallel_requests=None, - allowed_model_region=None, - metadata={}, # Set to empty dict instead of None - key_alias=None, - 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 - ) - - # Store the unified object for batch cost tracking - import asyncio - - asyncio.create_task( - managed_files_hook.store_unified_object_id( # type: ignore - unified_object_id=unified_object_id, - file_object=batch_object, - litellm_parent_otel_span=None, - model_object_id=model_object_id, - file_purpose="batch", - user_api_key_dict=user_api_key_dict, - ) - ) - - verbose_proxy_logger.info( - f"Stored batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}" - ) - else: - verbose_proxy_logger.warning( - "Managed files hook not available, cannot store batch object for cost tracking" - ) - - except Exception as e: - verbose_proxy_logger.error(f"Error storing batch managed object: {e}") + request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} + store_batch_managed_object( + unified_object_id=unified_object_id, + batch_object=batch_object, + model_object_id=model_object_id, + request_metadata=request_metadata, + ) @staticmethod def get_actual_model_id_from_router(model_name: str) -> str: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a23cecc3911..96c5a880a24 100644 --- a/litellm/proxy/schema.prisma +++ b/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? + user_api_key String? + request_tags Json? updated_at DateTime @updatedAt updated_by String? diff --git a/schema.prisma b/schema.prisma index a23cecc3911..96c5a880a24 100644 --- a/schema.prisma +++ b/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? + user_api_key String? + request_tags Json? updated_at DateTime @updatedAt updated_by String? diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index b822799fb40..a56b6fe11ca 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1221,3 +1221,126 @@ class TestUnmanagedBatchCostFlagIsGeneralized: assert vertex_result == ("deploy-vertex", "8823717160934178816") assert bedrock_result == ("deploy-bedrock", TestUnmanagedBedrockRouting._ARN) + + +class TestBatchCostAttribution: + """A batch-cost spend log must be attributed to the creating key/team/tags the same + way a non-batch request is. Regression for batches created by keys with no user_id, + where the row was previously dropped and identity was blank.""" + + def _instance(self): + from litellm_enterprise.proxy.common_utils.check_batch_cost import ( + CheckBatchCost, + ) + + instance = CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=MagicMock(), + llm_router=MagicMock(), + ) + instance.prisma_client.db = MagicMock() + return instance + + @pytest.mark.asyncio + async def test_get_user_info_skips_query_when_user_id_none(self): + """find_unique(where={"user_id": None}) raises "A value is required but not set"; + the None user_id case (team/service-account keys) must short-circuit instead.""" + instance = self._instance() + instance.prisma_client.db.litellm_usertable.find_unique = AsyncMock() + + result = await instance._get_user_info("batch-1", None) + + assert result == {} + instance.prisma_client.db.litellm_usertable.find_unique.assert_not_called() + + @pytest.mark.asyncio + async def test_attribution_metadata_replays_key_team_and_tags(self): + from types import SimpleNamespace + + instance = self._instance() + instance.prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=SimpleNamespace(user_email="creator@example.test") + ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=SimpleNamespace(key_alias="prod-batch-key") + ) + instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=SimpleNamespace(team_alias="growth-team") + ) + job = SimpleNamespace( + created_by="user-1", + team_id="team-1", + user_api_key="hashed-key-abc", + request_tags=["env:prod", "team:growth"], + ) + + metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + + 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["user_api_key_alias"] == "prod-batch-key" + assert metadata["user_api_key_team_alias"] == "growth-team" + assert metadata["user_api_key_user_email"] == "creator@example.test" + assert metadata["tags"] == ["env:prod", "team:growth"] + + @pytest.mark.asyncio + async def test_attribution_metadata_keeps_key_when_no_user_id(self): + """A team/service-account key has no created_by. The persisted key hash and + team_id must still flow through so _should_track_cost_callback writes the row.""" + from types import SimpleNamespace + + from litellm.proxy.hooks.proxy_track_cost_callback import ( + _should_track_cost_callback, + ) + + instance = self._instance() + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=None + ) + job = SimpleNamespace( + created_by=None, + team_id="team-1", + user_api_key="hashed-key-abc", + request_tags=None, + ) + + metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + + assert metadata["user_api_key"] == "hashed-key-abc" + assert metadata["user_api_key_user_id"] is None + assert "tags" not in metadata + assert ( + _should_track_cost_callback( + user_api_key=metadata["user_api_key"], + user_id=metadata.get("user_api_key_user_id"), + team_id=metadata.get("user_api_key_team_id"), + end_user_id=None, + ) + is True + ) + + @pytest.mark.asyncio + async def test_attribution_metadata_tolerates_legacy_rows(self): + """Older managed-object rows predate user_api_key/request_tags. Attribution must + still fall back to created_by/team_id without raising.""" + from types import SimpleNamespace + + instance = self._instance() + instance.prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=None + ) + job = SimpleNamespace(created_by="user-1", team_id="team-1") + + metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + + assert metadata["user_api_key"] is None + assert metadata["user_api_key_user_id"] == "user-1" + assert metadata["user_api_key_team_id"] == "team-1" + assert "tags" not in metadata diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 5de682ec8a0..8ffb0b01a68 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -315,6 +315,43 @@ 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_persists_key_hash_and_tags( + self, mock_logging_obj, mock_managed_files_hook + ): + """The creating key hash and request tags must be persisted so CheckBatchCost + can later attribute the batch-cost spend log. Regression: previously api_key + was hardcoded to "" and tags were never forwarded. The metadata already carries + the hashed token, so it must survive the UserAPIKeyAuth round-trip unchanged.""" + hashed_key = "a" * 64 + 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, + "user_api_key_team_id": "team-456", + "user_api_key_alias": "prod-batch-key", + "tags": ["env:prod", "team:growth"], + } + }, + ) + + mock_managed_files_hook.store_unified_object_id.assert_called_once() + call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1] + assert call_kwargs["user_api_key_dict"].api_key == hashed_key + assert call_kwargs["user_api_key_dict"].key_alias == "prod-batch-key" + assert call_kwargs["request_tags"] == ["env:prod", "team:growth"] + 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