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 7acdd5dbdaf5..dc8f17fb665f 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,7 +3,7 @@ """ from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -43,11 +43,15 @@ def __init__( # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True - async def _get_user_info(self, batch_id, user_id) -> dict: + async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: """ Look up user email and key alias by user_id for enriching the S3 callback metadata. Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). + Returns an empty dict when user_id is None: batches created by a team or service + account key carry no user id, and find_unique(where={"user_id": None}) raises. """ + if not user_id: + return {} try: user_row = await self.prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_id} @@ -62,6 +66,66 @@ async def _get_user_info(self, batch_id, user_id) -> dict: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} + async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None: + """Resolve the creating virtual key's alias from its hashed token.""" + if not api_key: + return None + try: + key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) + return getattr(key_row, "key_alias", None) if key_row is not None else None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}") + return None + + async def _get_team_alias(self, team_id: str | None) -> str | None: + """Resolve a team's alias from its id.""" + if not team_id: + return None + try: + team_row = await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + return getattr(team_row, "team_alias", None) if team_row is not None else None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") + return None + + async def _build_creator_attribution_metadata( + self, job: "LiteLLM_ManagedObjectTable", batch_id: str + ) -> Dict[str, Any]: + """ + Rebuild the spend-tracking metadata for the key, team, and tags that created the + batch so the batch-cost spend log is attributed the same way a non-batch request + is. Rows created before api_key and request_tags were persisted carry only + created_by and team_id, and fall back to those. A named creating key owns + user_api_key_alias; when it has no alias, or the key has since been rotated or + deleted, the field keeps the creating user's alias that _get_user_info filled in, + because a resolvable name is more useful on the spend row than a null. + """ + api_key = getattr(job, "api_key", None) + team_id = getattr(job, "team_id", None) + request_tags = getattr(job, "request_tags", None) + + metadata: Dict[str, Any] = { + "user_api_key_user_id": job.created_by, + "user_api_key": api_key, + "user_api_key_team_id": team_id, + **(await self._get_user_info(batch_id, job.created_by)), + } + + key_alias = await self._get_key_alias(batch_id, api_key) + if key_alias is not None: + metadata["user_api_key_alias"] = key_alias + team_alias = await self._get_team_alias(team_id) + if team_alias is not None: + metadata["user_api_key_team_alias"] = team_alias + if isinstance(request_tags, list) and request_tags: + metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] + + return metadata + async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days @@ -485,9 +549,6 @@ async def _track_completed_batch_cost( function_id=str(uuid.uuid4()), ) - creator_user_id = job.created_by - user_info = await self._get_user_info(batch_id, job.created_by) - logging_obj.update_environment_variables( litellm_params={ # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks @@ -496,11 +557,7 @@ async def _track_completed_batch_cost( "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": { - "user_api_key_user_id": creator_user_id, - "user_api_key_team_id": getattr(job, "team_id", None), - **user_info, - }, + "metadata": await self._build_creator_attribution_metadata(job, batch_id), }, optional_params={}, ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index f0914240f797..a3994dccfd6b 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -163,6 +163,8 @@ async def upsert( self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] ) -> "PrismaManagedObjectRow": ... + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + class _CursorPageArgs(TypedDict, total=False): cursor: Mapping[str, str] @@ -263,7 +265,24 @@ async def store_unified_object_id( model_object_id: str, file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, + request_tags: Sequence[str] | None = None, + persist_attribution: bool = False, + create_if_missing: bool = True, ) -> None: + """Persist a managed object row, caching it and upserting it in the DB. + + persist_attribution is set only by the batch create, which is the one caller + that can speak for the creator; it gates the api_key and request_tags columns + that CheckBatchCost bills against, so a later poll or retrieve of the same + batch cannot record itself as the paying key. Like created_by and team_id, + both are written only in the upsert create branch, never on update. + + create_if_missing is cleared by callers that observe a batch they did not + create, such as a poll. They still refresh status and file_object, but a + row absent from the table is left absent rather than created with the + observer as its creator, because created_by and team_id are written from + whoever calls the create branch. + """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, @@ -277,6 +296,29 @@ async def store_unified_object_id( litellm_parent_otel_span=litellm_parent_otel_span, ) + from prisma import Json + + api_key = user_api_key_dict.api_key or None + attribution_columns = ( + { + **({"api_key": api_key} if api_key is not None else {}), + **({"request_tags": Json(list(request_tags))} if request_tags else {}), + } + if persist_attribution + else {} + ) + # FIX: Update status and file_object on every operation to keep state in sync + update_columns: Final = { + "file_object": file_object.model_dump_json(), + "status": file_object.status, + "updated_by": user_api_key_dict.user_id, + } + if not create_if_missing: + await _managed_object_table(self.prisma_client).update_many( + where={"unified_object_id": unified_object_id}, + data=update_columns, + ) + return await _managed_object_table(self.prisma_client).upsert( where={"unified_object_id": unified_object_id}, data={ @@ -289,12 +331,9 @@ async def store_unified_object_id( "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, + **attribution_columns, }, - "update": { - "file_object": file_object.model_dump_json(), - "status": file_object.status, - "updated_by": user_api_key_dict.user_id, - }, # FIX: Update status and file_object on every operation to keep state in sync + "update": update_columns, }, ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql new file mode 100644 index 000000000000..79bc6b24de89 --- /dev/null +++ b/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 '[]'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9c871b65f401..cabddf6f1a1a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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? diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3346f9d7e3b0..0e22b5324c16 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -48,6 +48,15 @@ } ) +# Both spellings, because call_type reaches the callback as str(...) of either the +# enum member or its value. +_CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( + ( + CallTypes.aretrieve_batch.value, + str(CallTypes.aretrieve_batch), + ) +) + class _ProxyDBLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -212,7 +221,10 @@ async def _PROXY_track_cost_callback( # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls). # Avoids a cache/DB lookup on every normal LLM request. if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"): - metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata) + metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( # rebind-ok: enriched metadata replaces the original + metadata=metadata, + resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES, + ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) @@ -337,7 +349,7 @@ async def _PROXY_track_cost_callback( spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) @staticmethod - async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: + async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict: """ Enriches failure spend log metadata by looking up the key object (and team object) from cache/DB when key fields are missing. @@ -349,6 +361,11 @@ async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: 2. Post-auth failures (provider errors, rate limits): key fields are populated but team_alias is missing because LiteLLM_VerificationTokenView SQL view doesn't include it. We look up the team object to fill in team_alias. + + Scenario 1 reads the key's identity as it stands right now, so it is only correct + for a log emitted within the request it describes. Callers that log after a delay, + against an identity captured earlier, pass resolve_missing_key_identity=False and + keep their own user_id, team_id and org_id. """ api_key_hash: Final = metadata.get("user_api_key") if not api_key_hash: @@ -361,7 +378,7 @@ async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: ) # Step 1: If key fields are missing, look up the full key object - if metadata.get("user_api_key_alias") is None: + if resolve_missing_key_identity and metadata.get("user_api_key_alias") is None: try: key_obj: Final = await get_key_object( hashed_token=api_key_hash, 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 2e3f7bb9aa62..9c5b7dc563ed 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 @@ -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 @@ 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 @@ def batch_prediction_jobs_handler( # 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 @@ def batch_prediction_jobs_handler( "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 @@ def _store_batch_managed_object( 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 @@ def _store_batch_managed_object( ) # 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 @@ def _store_batch_managed_object( 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( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9c871b65f401..cabddf6f1a1a 100644 --- a/litellm/proxy/schema.prisma +++ b/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? diff --git a/schema.prisma b/schema.prisma index 9c871b65f401..cabddf6f1a1a 100644 --- a/schema.prisma +++ b/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? diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 800c97d7ba12..6d7ada17ec52 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1641,3 +1641,153 @@ async def test_falls_back_to_deployment_model_group_without_managed_input_file(s 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" diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index 1b60c97b5101..ebd33aa2e53b 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/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 diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 63646ae53f84..34cd0cabc2cb 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/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 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 69f04ce2bbe5..2b162774aea6 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -17,6 +17,7 @@ _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(): """ 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 52da7a4a81d4..d53e6dedf0b9 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 @@ -258,6 +258,7 @@ def test_store_batch_managed_object( 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 @@ def test_store_batch_managed_object_propagates_user_identity_from_metadata( 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 @@ def test_store_batch_managed_object_propagates_user_identity_from_metadata( 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