diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 31af5a144eb..99b83c9a146 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,6 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from itertools import accumulate, chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, Union, cast @@ -69,7 +70,7 @@ if TYPE_CHECKING: ChatCompletionToolReferenceObject, OpenAIMessageContentListBlock, ) - from litellm.types.utils import Choices + from litellm.types.utils import Choices, Message _CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) @@ -194,6 +195,67 @@ def _as_chat_reasoning_items( return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) +def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation: + shifted: Final = { + key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value + for key, value in annotation.items() + } + return cast(ChatCompletionAnnotation, shifted) + + +def _message_annotations(message: "Message") -> Sequence[ChatCompletionAnnotation]: + annotations: Sequence[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) + return annotations or () + + +def _message_reasoning_content(message: "Message") -> str | None: + reasoning_content: str | None = getattr(message, "reasoning_content", None) + return reasoning_content + + +def _message_reasoning_items(message: "Message") -> Sequence[ChatCompletionReasoningItem]: + reasoning_items: Sequence[ChatCompletionReasoningItem] | None = getattr(message, "reasoning_items", None) + return reasoning_items or () + + +def _shifted_annotations(messages: Sequence["Message"], offsets: Sequence[int]) -> Iterator[ChatCompletionAnnotation]: + for message, offset in zip(messages, offsets): + for annotation in _message_annotations(message): + yield _shift_annotation(annotation, offset) + + +def _fold_choices_into_one( + choices: Sequence["Choices"], + tool_calls: Sequence[Mapping[str, object]], + trailing_reasoning_content: str | None, + trailing_reasoning_item: _BuiltReasoningItem | None, +) -> "Choices": + """Clients that only read choices[0] must still see every text and tool call.""" + from litellm.types.utils import Choices, Message + + messages: Final = tuple(choice.message for choice in choices) + texts: Final = tuple(str(message.content or "") for message in messages) + offsets: Final = (0, *accumulate(len(text) for text in texts)) + annotations: Final = tuple(_shifted_annotations(messages, offsets)) + reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) + reasoning_items: Final = ( + *chain.from_iterable(_message_reasoning_items(m) for m in messages), + *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + ) + return Choices( + message=Message( + role=messages[0].role if messages else "assistant", + content="".join(texts) if messages else None, + annotations=list(annotations) or None, + tool_calls=list(tool_calls) or None, + reasoning_content=" ".join(value for value in reasoning_contents if value) or None, + reasoning_items=_as_chat_reasoning_items(reasoning_items), + ), + finish_reason="tool_calls" if tool_calls else "stop", + index=0, + ) + + _ToolChoiceT = TypeVar("_ToolChoiceT") @@ -784,22 +846,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: pass # don't fail request if item in list is not supported - # If we accumulated tool calls, create a single choice with all of them - if accumulated_tool_calls: - msg = Message( - content=None, + if not choices and not accumulated_tool_calls: + return choices + return [ + _fold_choices_into_one( + choices=choices, tool_calls=accumulated_tool_calls, - reasoning_content=reasoning_content, - reasoning_items=cast( - list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), - ), + trailing_reasoning_content=reasoning_content, + trailing_reasoning_item=pending_reasoning_item, ) - choices.append(Choices(message=msg, finish_reason="tool_calls", index=index)) - reasoning_content = None - pending_reasoning_item = None - - return choices + ] @staticmethod def _build_empty_incomplete_choice( diff --git a/litellm/constants.py b/litellm/constants.py index 4420fd73e77..4f0c34a7bb8 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1899,6 +1899,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 +SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 0d6cbc2232e..1db144b5fdc 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -695,6 +695,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): tools: list | None = None, cache_control: object = None, request_kwargs: object = None, + on_messages_route: bool = False, ) -> bool: """Return True if the request already carries any client-supplied cache_control. @@ -704,10 +705,14 @@ class AnthropicCacheControlHook(CustomPromptManagement): envelope. Configured injection points are an explicit instruction and are applied alongside the client's marks, bounded by the provider cap. """ - return ( - AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) - + AnthropicCacheControlHook.count_external_cache_breakpoints(tools, cache_control, request_kwargs) - ) > 0 + external_breakpoints: Final = ( + AnthropicCacheControlHook.count_external_cache_breakpoints_on_messages_route( + tools, cache_control, request_kwargs + ) + if on_messages_route + else AnthropicCacheControlHook.count_external_cache_breakpoints(tools, cache_control, request_kwargs) + ) + return AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) + external_breakpoints > 0 @staticmethod def get_default_injection_points( @@ -719,6 +724,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): enable_prompt_caching: bool | None = None, cache_control: object = None, request_kwargs: object = None, + on_messages_route: bool = False, ) -> list[CacheControlInjectionPoint]: """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. @@ -739,7 +745,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not supports_anthropic_cache_control(model, custom_llm_provider): return [] - if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control, request_kwargs): + if AnthropicCacheControlHook._request_has_cache_control( + messages, system, tools, cache_control, request_kwargs, on_messages_route + ): return [] if is_claude_code_one_shot_subagent_request( @@ -968,6 +976,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): enable_prompt_caching=enable_prompt_caching, cache_control=cache_control, request_kwargs=kwargs, + on_messages_route=True, ) if model is not None else () diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 821abe56580..3d6127aa63a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -672,6 +672,9 @@ class Logging(LiteLLMLoggingBaseClass): self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: datetime.datetime | None = None + # The model the proxy shows the client on streamed chunks. The logged streamed response carries it + # once that response is priced, the same way a non-streamed response is logged + self.client_facing_stream_model: str | None = None self.zero_cost_warned: bool = False self._llm_caching_handler: LLMCachingHandler | None = None @@ -710,6 +713,7 @@ class Logging(LiteLLMLoggingBaseClass): # enqueue closure here instead of firing it immediately. self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None + self._async_success_scheduled: bool = False self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self._async_success_enqueued: bool = False self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None @@ -2280,6 +2284,15 @@ class Logging(LiteLLMLoggingBaseClass): def mark_async_success_log_enqueued(self) -> None: self._async_success_enqueued = True + def claim_async_success_log(self) -> bool: + """One async success log per request. The innermost @client wrapper always exits first, + since the outer one is awaiting it, so it claims the log here and the outer wrapper gets + False and schedules nothing.""" + if self._async_success_scheduled: + return False + self._async_success_scheduled = True + return True + def has_run_logging( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], @@ -2508,6 +2521,15 @@ class Logging(LiteLLMLoggingBaseClass): setattr(result, "usage", transformed_usage) return result + def _with_client_facing_stream_model( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + model: Final = self.client_facing_stream_model + if model is None or getattr(response, "model", None) in (None, model): + return response + return response.model_copy(update={"model": model}) + def _success_handler_helper_fn( self, result=None, @@ -2824,9 +2846,11 @@ class Logging(LiteLLMLoggingBaseClass): result=complete_streaming_response ) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) + logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + self.model_call_details["complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + logged_streaming_response, start_time, end_time ) standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( "standard_logging_object" @@ -3370,10 +3394,13 @@ class Logging(LiteLLMLoggingBaseClass): await self._prepare_baseline_cache_estimate(complete_streaming_response) + logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + self.model_call_details["async_complete_streaming_response"] = logged_streaming_response + ## STANDARDIZED LOGGING PAYLOAD try: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + logged_streaming_response, start_time, end_time ) except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release) verbose_logger.exception( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 0bd46382fef..c28b9a43301 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -44,22 +44,31 @@ def _is_message_stop_chunk(chunk: object) -> bool: def is_anthropic_ping_chunk(chunk: object) -> bool: """ - Whether a chunk is a pure ``ping`` keepalive frame. It carries no content - and can recur indefinitely on a slow-starting or idle connection, so a - mid-stream fallback wrapper drops it outright while still deciding - whether to commit to the primary stream, rather than buffering it. + Whether a chunk is made only of whole ``ping`` keepalive frames. A ping + carries no content or lifecycle, so a mid-stream fallback wrapper can + forward it live while still deciding whether to commit to the primary + stream, without risking two overlapping message lifecycles on the wire. A physical transport chunk that coalesces a ping with any other SSE event (``message_start``, ``content_block_delta``, ``event: error``, ...) - is NOT a pure ping - dropping it whole would discard those events - so - only a chunk whose every ``event:`` line is ``event: ping`` qualifies. + is NOT a pure ping, and neither is a fragment of a ping frame split + across two reads, or a chunk that opens with the tail of an earlier + frame: forwarding either live would interleave it with frames still + held back for a fallback. Only a chunk that begins with ``event: ping``, + ends on a frame boundary, and whose every ``event:`` line is + ``event: ping`` qualifies. """ if isinstance(chunk, dict): return chunk.get("type") == "ping" - if isinstance(chunk, (bytes, bytearray)): - event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:")) - return bool(event_lines) and all(line == b"event: ping" for line in event_lines) - return False + if not isinstance(chunk, (bytes, bytearray)): + return False + event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:")) + return ( + bool(event_lines) + and all(line == b"event: ping" for line in event_lines) + and chunk.startswith(b"event: ping") + and chunk.endswith((b"\n\n", b"\r\n\r\n")) + ) def is_anthropic_content_delta_chunk(chunk: object) -> bool: diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index cecf3e50f5c..c2a0a41c3e2 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -16,6 +16,7 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import ( recover_cli_session_key_metadata, recover_double_hashed_key_metadata, recover_key_metadata_from_spend_logs, + recover_key_owner_from_daily_spend, ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient @@ -468,6 +469,17 @@ def _parse_spend_date(raw: str | None) -> datetime | None: _EMPTY_KEY_METADATA: Final[Mapping[str, _KeyMetadataDict]] = MappingProxyType({}) +def _metadata_with_recovered_owner( + metadata: Mapping[str, _KeyMetadataDict], + key: str, + owner: str, +) -> _KeyMetadataDict: + current: Final = metadata.get(key) + if current is None: + return {"user_id": owner} + return {**current, "user_id": owner} + + async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], @@ -530,7 +542,19 @@ async def get_api_key_metadata( else _EMPTY_KEY_METADATA ) combined: Final = MappingProxyType({**after_token_recovery, **from_spend_logs}) - return await attach_user_details(prisma_client, combined) + ownerless: Final = frozenset( + key + for key in api_keys + if not combined.get(key, {}).get("user_id") and not combined.get(key, {}).get("key_exists") + ) + owners: Final = await recover_key_owner_from_daily_spend(prisma_client, ownerless) + metadata_with_owners: Final[Mapping[str, _KeyMetadataDict]] = MappingProxyType( + { + **combined, + **{key: _metadata_with_recovered_owner(combined, key, owner) for key, owner in owners.items()}, + } + ) + return await attach_user_details(prisma_client, metadata_with_owners) def _adjust_dates_for_timezone( @@ -944,7 +968,7 @@ async def _aggregate_spend_records( record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY } - api_key_metadata: dict[str, _KeyMetadataDict] = {} + api_key_metadata: Mapping[str, _KeyMetadataDict] = MappingProxyType({}) if api_keys: api_key_metadata = await get_api_key_metadata( prisma_client, api_keys, _spend_logs_window(frozenset(record.date for record in records)) @@ -1144,7 +1168,7 @@ async def _aggregate_grouping_sets_records( """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread.""" api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY} - api_key_metadata: dict[str, _KeyMetadataDict] = {} + api_key_metadata: Mapping[str, _KeyMetadataDict] = MappingProxyType({}) if api_keys: api_key_metadata = await get_api_key_metadata( prisma_client, api_keys, _spend_logs_window(frozenset(r.date for r in records)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 79b4b8aaa29..9aa606599c1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9222,12 +9222,17 @@ def _restamp_streaming_chunk_model( ) model_mismatch_logged = True + # The streaming wrapper keeps these same chunk objects to assemble the response it + # prices, so stamp a copy for the client and leave the provider's model for pricing. + # The logging object stamps the same model on the assembled response after pricing it. + logging_obj: Final = request_data.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.client_facing_stream_model = target_model if isinstance(chunk, dict): - chunk["model"] = target_model - return chunk, model_mismatch_logged + return {**chunk, "model": target_model}, model_mismatch_logged try: - chunk.model = target_model + return chunk.model_copy(update={"model": target_model}), model_mismatch_logged except Exception as e: verbose_proxy_logger.error( "litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s", diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 08e8fff8f1d..9cc02a74c99 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -4,7 +4,7 @@ from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timedelta from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, Literal, TypeVar from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -17,6 +17,7 @@ from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, ) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient @@ -38,29 +39,69 @@ WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) ORDER BY token, deleted_at DESC """ -_SPEND_LOG_ALIAS_SQL: Final = """ -SELECT api_key AS digest, - MIN(key_alias) AS first_alias, - MAX(key_alias) AS last_alias, - MIN(team_id) AS first_team, - MAX(team_id) AS last_team, - MIN(user_id) AS first_owner, - MAX(user_id) AS last_owner -FROM ( - SELECT api_key, - NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, - COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, - COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id - FROM "LiteLLM_SpendLogs" - WHERE api_key = ANY($1::text[]) - AND "startTime" >= $2::timestamp - AND "startTime" < $3::timestamp -) named -WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL + +def _named_spend_log_edge_row_sql( + direction: Literal["ASC", "DESC"], since: Literal["$2::timestamp", "oldest_probe.stopped_at"] +) -> str: + return f""" + SELECT "startTime", key_alias, team_id, user_id + FROM ( + SELECT "startTime", + NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, + COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, + COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id + FROM ( + SELECT "startTime", metadata, team_id, "user" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= {since} + AND "startTime" < $3::timestamp + ORDER BY "startTime" {direction} + LIMIT {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE} + ) edge + ) named + WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL + ORDER BY "startTime" {direction} + LIMIT 1 + """ + + +_OLDEST_PROBE_STOPPED_AT_SQL: Final = f""" + SELECT COALESCE(first_row."startTime", ( + SELECT "startTime" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + ORDER BY "startTime" ASC + OFFSET {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE - 1} + LIMIT 1 + )) AS stopped_at +""" + +_SPEND_LOG_ALIAS_SQL: Final = f""" +SELECT keys.digest, + first_row.key_alias AS first_alias, + last_row.key_alias AS last_alias, + first_row.team_id AS first_team, + last_row.team_id AS last_team, + first_row.user_id AS first_owner, + last_row.user_id AS last_owner +FROM unnest($1::text[]) AS keys(digest) +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("ASC", "$2::timestamp")}) first_row ON true +LEFT JOIN LATERAL ({_OLDEST_PROBE_STOPPED_AT_SQL}) oldest_probe ON true +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("DESC", "oldest_probe.stopped_at")}) last_row ON true +""" + +_DAILY_USER_SPEND_OWNER_SQL: Final = """ +SELECT api_key, MIN(user_id) AS first_owner, MAX(user_id) AS last_owner +FROM "LiteLLM_DailyUserSpend" +WHERE api_key = ANY($1::text[]) AND user_id IS NOT NULL AND user_id <> '' GROUP BY api_key """ _SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" +_SPEND_LOG_NO_BITMAP_SCAN_SQL: Final = "SET LOCAL enable_bitmapscan = off" _SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS) _HASHED_JWT_PREFIX: Final = "hashed-jwt-" @@ -83,7 +124,9 @@ class _TokenDigestRow(BaseModel): def _unanimous(first: str | None, last: str | None) -> str | None: - return first if first == last else None + if first is None: + return last + return first if last is None or first == last else None class _SpendLogDigestRow(BaseModel): @@ -103,8 +146,15 @@ class _SpendLogDigestRow(BaseModel): ) +class _DailyUserSpendOwnerRow(BaseModel): + api_key: str + first_owner: str | None = None + last_owner: str | None = None + + _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) _SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...]) +_DAILY_USER_SPEND_OWNER_ROWS: Final = TypeAdapter(tuple[_DailyUserSpendOwnerRow, ...]) _CACHED_KEY_METADATA: Final = TypeAdapter(KeyMetadataDict) _SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( max_size_in_memory=SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, @@ -112,6 +162,7 @@ _SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( ) _SPEND_LOG_QUERY_LOCK: Final = asyncio.Lock() _EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) +_EMPTY_KEY_OWNERS: Final[Mapping[str, str]] = MappingProxyType({}) async def _db_or_empty( @@ -128,6 +179,19 @@ async def _db_or_empty( return None +async def _rows_within_the_statement_timeout( + prisma_client: PrismaClient, + sql: str, + *params: object, + planner_settings: tuple[str, ...] = (), +) -> Sequence[Mapping[str, object]]: + async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction: + await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) + for setting in planner_settings: + await transaction.execute_raw(setting) + return await transaction.query_raw(sql, *params) + + async def _reverse_hash_key_metadata( prisma_client: PrismaClient, sql: str, @@ -151,6 +215,29 @@ async def _reverse_hash_key_metadata( ) +async def recover_key_owner_from_daily_spend( + prisma_client: PrismaClient, + keys: AbstractSet[str], +) -> Mapping[str, str]: + if not keys: + return _EMPTY_KEY_OWNERS + rows: Final = await _db_or_empty( + lambda: _rows_within_the_statement_timeout(prisma_client, _DAILY_USER_SPEND_OWNER_SQL, sorted(keys)), + "Failed daily-spend key owner recovery for %d keys: %s", + len(keys), + ) + if rows is None: + return _EMPTY_KEY_OWNERS + return MappingProxyType( + { + row.api_key: owner + for row in _DAILY_USER_SPEND_OWNER_ROWS.validate_python(rows) + for owner in (_unanimous(row.first_owner, row.last_owner),) + if row.api_key in keys and owner is not None + } + ) + + @dataclass(frozen=True, slots=True) class _UserDetails: email: str | None @@ -304,24 +391,21 @@ def _cached_spend_log_metadata( ) -async def _spend_log_rows_within_the_statement_timeout( - prisma_client: PrismaClient, - digests: AbstractSet[str], - window: tuple[datetime, datetime], -) -> Sequence[Mapping[str, object]]: - start, end = window - async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction: - await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) - return await transaction.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end) - - async def _query_spend_log_metadata( prisma_client: PrismaClient, digests: AbstractSet[str], window: tuple[datetime, datetime], ) -> Mapping[str, KeyMetadataDict] | None: + start, end = window rows: Final = await _db_or_empty( - lambda: _spend_log_rows_within_the_statement_timeout(prisma_client, digests, window), + lambda: _rows_within_the_statement_timeout( + prisma_client, + _SPEND_LOG_ALIAS_SQL, + sorted(digests), + start, + end, + planner_settings=(_SPEND_LOG_NO_BITMAP_SCAN_SQL,), + ), "Failed spend-log alias recovery for %d missing keys: %s", len(digests), ) diff --git a/litellm/router.py b/litellm/router.py index d863d3fdb21..ec6d608a0d8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -526,31 +526,20 @@ def _with_router_resolved_session_model(session: object, model_name: str) -> Map # Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks -# until real content commits the primary stream; a hostile or slow-starting -# upstream that never emits content or an error could otherwise grow that -# buffer without bound, so hitting this cap forces an early commit instead. +# until real content commits the primary stream, and only while a fallback +# can still take over; a hostile or slow-starting upstream that never emits +# content or an error could otherwise grow that buffer without bound, so +# hitting this cap forces an early commit instead. MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 -def _anthropic_stream_should_drop_pre_content_ping(chunk: object, has_generated_content: bool) -> bool: - """A `ping` keepalive seen before any real content is dropped outright - it recurs indefinitely on a - slow-starting connection and carries nothing worth buffering toward a possible fallback.""" +def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool) -> bool: + """A `ping` keepalive reaches the client live whenever the stream has not committed: it carries no + lifecycle, so it cannot create overlapping lifecycles on the wire, and it keeps the connection alive + while lifecycle frames sit buffered for a possible fallback during a long thinking pass.""" from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import is_anthropic_ping_chunk - if has_generated_content: - return False - return is_anthropic_ping_chunk(chunk) - - -def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: - """A `ping` that no lifecycle frame precedes reaches the client live: a fallback's own message_start can still - follow it without overlapping lifecycles, and AgenticAnthropicStreamingIterator's hold-back keepalive is exactly - such a ping.""" - from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import is_anthropic_ping_chunk - - if has_generated_content or buffered_chunk_count: - return False - return is_anthropic_ping_chunk(chunk) + return not has_generated_content and is_anthropic_ping_chunk(chunk) def _is_retriable_anthropic_status(status_code: int) -> bool: @@ -5451,14 +5440,19 @@ class Router: Lifecycle/bookkeeping frames (message_start, content_block_start, ping, ...) do not by themselves disqualify a fallback attempt - - Anthropic routinely sends message_start before an overload error - - but they are BUFFERED rather than forwarded immediately, since - forwarding one and then appending a fallback attempt's own - message_start would produce two overlapping message lifecycles on - one SSE stream. Buffered frames are flushed, in order, the moment - real content arrives (the primary attempt has committed by then - anyway) or once the stream ends without ever producing content or - an error. + Anthropic routinely sends message_start before an overload error. + When a fallback can still take over they are BUFFERED rather than + forwarded immediately, since forwarding one and then appending a + fallback attempt's own message_start would produce two overlapping + message lifecycles on one SSE stream; a `ping` carries no lifecycle, + so it is forwarded live even while lifecycle frames sit buffered, + keeping the connection alive during a long thinking pass. Buffered + frames are flushed, in order, the moment real content arrives (the + primary attempt has committed by then anyway) or once the stream + ends without ever producing content or an error. When no fallback + can take over the request is already committed, so every frame, + including pings and provider error frames, is forwarded live and + verbatim instead. """ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, @@ -5475,34 +5469,33 @@ class Router: from litellm.exceptions import MidStreamFallbackError # Lifecycle/bookkeeping frames (message_start, content_block_start, - # ping, ...) are held back rather than forwarded immediately: - # Anthropic routinely sends message_start before an overload - # error, and once a byte reaches the client a fallback attempt - # can only append its OWN message_start, producing two - # overlapping message lifecycles on one SSE stream. Buffered - # frames are flushed the moment real content (content_block_delta) + # ...) are held back rather than forwarded immediately, but only + # while a fallback can still take over: Anthropic routinely sends + # message_start before an overload error, and once a byte reaches + # the client a fallback attempt can only append its OWN + # message_start, producing two overlapping message lifecycles on + # one SSE stream. A `ping` keepalive carries no lifecycle, so it + # is forwarded live even behind buffered frames, keeping the + # connection alive through a long thinking pass. Buffered frames + # are flushed the moment real content (content_block_delta) # arrives - at that point the primary attempt has committed and a # clean retry is no longer possible anyway - or once the primary - # stream ends without ever producing content. A `ping` keepalive - # that nothing precedes is forwarded live (it is how a hold-back - # turn keeps its connection alive); one behind buffered frames is - # dropped outright rather than buffered, since it can recur - # indefinitely on a slow-starting connection and carries nothing - # worth preserving; hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - # forces the same early commit as real content arriving, so a - # hostile or pathological upstream can't grow the buffer forever. - has_generated_content = False # rebind-ok: set once real content is seen, or the buffer cap is hit - buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + # stream ends without ever producing content. Hitting + # MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS forces the same early + # commit as real content arriving, so a hostile or pathological + # upstream can't grow the buffer forever. With no fallback able + # to take over there is nothing to buffer for, so every frame, + # including pings and provider error frames, is forwarded live. model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group + has_generated_content = not self._anthropic_messages_stream_can_fall_back( # rebind-ok: set once real content is seen, the buffer cap is hit, or no fallback can take over + model, initial_kwargs + ) + buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline try: async for chunk in source_iterator: - if _anthropic_stream_forwards_ping_live( - chunk, has_generated_content, len(buffered_lifecycle_chunks) - ): + if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): yield chunk continue - if _anthropic_stream_should_drop_pre_content_ping(chunk, has_generated_content): - continue if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): has_generated_content = True # A transport can split one SSE data line across byte chunks, so pre-content @@ -8466,6 +8459,56 @@ class Router: ) return has_unattempted_fallback_target(resolved, kwargs) + def _anthropic_messages_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: + """ + The distinct deployment order levels the fallback dispatcher would see for this request, + computed the same way: the tier a pre-routing hook selected wins over the requested group. + """ + request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") + order_model_group: Final = get_pre_routing_selection(kwargs) or model_group + all_deployments: Final = self.get_model_list(model_name=order_model_group, team_id=request_team_id) or () + return tuple( + sorted( + { + litellm.utils._get_deployment_order(d) + for d in all_deployments + if litellm.utils._get_deployment_order(d) is not None + } + ) + ) + + def _anthropic_messages_stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + """ + Whether async_function_with_fallbacks_common_utils could still route a + MidStreamFallbackError somewhere for this request (order levels, weighted + failover, content-policy or generic fallbacks), which is the only case where + holding lifecycle frames back from the client buys a clean retry. Errs toward + True whenever a dispatcher path might reach a fallback. + """ + if fallbacks_disabled_for_request(kwargs): + return False + if self.enable_weighted_failover: + return True + order_levels: Final = self._anthropic_messages_order_levels(model_group, kwargs) + if len(order_levels) > 1: + current_target: Final = kwargs.get("_target_order") + skip_up_to: Final = current_target if current_target is not None else order_levels[0] + if any(o > skip_up_to for o in order_levels): + return True + content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + if content_policy_fallbacks is not None and self._has_content_policy_fallback(model_group, kwargs): + return True + fallbacks: Final = kwargs.get("fallbacks", self.fallbacks) + if not fallbacks: + return False + if _check_non_standard_fallback_format(fallbacks=fallbacks): + return True + resolved, _ = get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, + lookup_groups=fallback_lookup_groups(kwargs, model_group), + ) + return has_unattempted_fallback_target(resolved, kwargs) + def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: """ Determines if a content policy error should be raised. diff --git a/litellm/utils.py b/litellm/utils.py index 59a93965993..d6fc6399fad 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1358,11 +1358,14 @@ def _schedule_async_success_logging( Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses bridge) each exit through here with the same logging object and their own shape of the same - response. The immediate path already logs one request once, since the first task marks - ``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same - first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log - reads usage from, and a later wrapper never swaps in its client-shaped translation. + response. The innermost wrapper always exits first, since the outer one is awaiting it, so it + claims the async success log here, synchronously, and the outer wrapper returns without + enqueuing: one request queues one handler, whatever the handler tasks do later. The deferred + slot keeps the same innermost-wins rule: the provider-shaped result is the one the spend log + reads usage from, and the outer wrapper never swaps in its client-shaped translation. """ + if not logging_obj.claim_async_success_log(): + return def _enqueue_async_logging() -> None: logging_obj.mark_async_success_log_enqueued() diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index e07cbe6b2a3..6590a9e414f 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -3,7 +3,7 @@ from __future__ import annotations import os import time import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass from hashlib import sha256 @@ -22,6 +22,11 @@ def object_value(value: JsonValue) -> dict[str, JsonValue]: return JSON_OBJECT.validate_python(value) +def list_value(value: JsonValue) -> Sequence[JsonValue]: + assert isinstance(value, list), f"Expected a list, received {type(value).__name__}" + return value + + def string_value(value: JsonValue) -> str: assert isinstance(value, str), f"Expected a string, received {type(value).__name__}" return value diff --git a/tests/integration/_support/daily_activity.py b/tests/integration/_support/daily_activity.py new file mode 100644 index 00000000000..346fc156ea9 --- /dev/null +++ b/tests/integration/_support/daily_activity.py @@ -0,0 +1,302 @@ +import os +import uuid +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta +from hashlib import sha256 +from itertools import chain +from typing import Final + +import httpx +import psycopg +import pytest +from integration._support.client import Gateway, Scenario, object_value +from psycopg import sql +from psycopg.types.json import Jsonb +from pydantic import JsonValue + +USER_SPEND: Final = "LiteLLM_DailyUserSpend" +TEAM_SPEND: Final = "LiteLLM_DailyTeamSpend" +TAG_SPEND: Final = "LiteLLM_DailyTagSpend" +ORGANIZATION_SPEND: Final = "LiteLLM_DailyOrganizationSpend" +END_USER_SPEND: Final = "LiteLLM_DailyEndUserSpend" +AGENT_SPEND: Final = "LiteLLM_DailyAgentSpend" +DAY: Final = "2026-02-03" +AGGREGATED_USER_ACTIVITY: Final = "/user/daily/activity/aggregated" + +INSERT_DAILY_ROW: Final = sql.SQL( + "INSERT INTO {table} (id, {entity}, date, api_key, model, model_group, custom_llm_provider, prompt_tokens," + " completion_tokens, spend, api_requests, successful_requests, failed_requests, updated_at)" + " VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, now())" +) +DELETE_DAILY_ROWS: Final = sql.SQL("DELETE FROM {table} WHERE api_key = ANY(%s)") +INSERT_SPEND_LOG: Final = ( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, "startTime", "endTime", metadata)' + " VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s)" +) +DELETE_SPEND_LOG: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +INSERT_SPEND_LOG_ROW: Final = ( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, "startTime", "endTime", metadata, team_id, "user")' + " VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s, %s, %s)" +) +DELETE_SPEND_LOG_ROWS: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)' +DELETE_KEY_ROW: Final = 'DELETE FROM "LiteLLM_VerificationToken" WHERE token = %s' +DELETE_ARCHIVED_KEY_ROW: Final = 'DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s' +LOCK_TABLE: Final = sql.SQL("LOCK TABLE {table} IN ACCESS EXCLUSIVE MODE") +SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs" +FIRST_SPEND_LOG_AT: Final = datetime(2026, 2, 3, 12, 0, 0) + + +@dataclass(frozen=True, slots=True) +class Route: + path: str + table: str + entity_column: str + entity_filter: str | None + + +ROUTES: Final = ( + Route("/user/daily/activity", USER_SPEND, "user_id", None), + Route(AGGREGATED_USER_ACTIVITY, USER_SPEND, "user_id", None), + Route("/team/daily/activity", TEAM_SPEND, "team_id", "team_ids"), + Route("/team/daily/activity/aggregated", TEAM_SPEND, "team_id", "team_ids"), + Route("/tag/daily/activity", TAG_SPEND, "tag", "tags"), + Route("/organization/daily/activity", ORGANIZATION_SPEND, "organization_id", "organization_ids"), + Route("/customer/daily/activity", END_USER_SPEND, "end_user_id", "end_user_ids"), + Route("/end_user/daily/activity", END_USER_SPEND, "end_user_id", "end_user_ids"), + Route("/agent/daily/activity", AGENT_SPEND, "agent_id", "agent_ids"), +) + + +def user_with_an_email(scenario: Scenario) -> tuple[str, str]: + email: Final = f"integration-{uuid.uuid4().hex}@example.com" + return scenario.user(user_email=email), email + + +def key_no_key_table_holds() -> str: + return f"integration-ownerless-{uuid.uuid4().hex}" + + +def digest_no_key_table_holds() -> str: + return sha256(uuid.uuid4().bytes).hexdigest() + + +def activity_of_key( + gateway: Gateway, path: str, api_key: str, *, reader: str | None = None, **filters: str +) -> httpx.Response: + return gateway.request( + "GET", path, params={"start_date": DAY, "end_date": DAY, "api_key": api_key, **filters}, key=reader + ) + + +@dataclass(frozen=True, slots=True) +class DailyRow: + table: str + entity_column: str + entity: str | None + api_key: str + date: str + model: str + provider: str + prompt_tokens: int + completion_tokens: int + spend: float + successful_requests: int + failed_requests: int + + +def _insert(connection: psycopg.Connection[tuple[object, ...]], row: DailyRow) -> None: + connection.execute( + INSERT_DAILY_ROW.format(table=sql.Identifier(row.table), entity=sql.Identifier(row.entity_column)), + ( + row.entity, + row.date, + row.api_key, + row.model, + row.model, + row.provider, + row.prompt_tokens, + row.completion_tokens, + row.spend, + row.successful_requests + row.failed_requests, + row.successful_requests, + row.failed_requests, + ), + ) + + +def insert_daily_rows(rows: Sequence[DailyRow], *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + for row in rows: + _insert(connection, row) + + +def delete_daily_rows(rows: Sequence[DailyRow], *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + for table in sorted({row.table for row in rows}): + connection.execute( + DELETE_DAILY_ROWS.format(table=sql.Identifier(table)), + (sorted({row.api_key for row in rows if row.table == table}),), + ) + + +@contextmanager +def daily_rows(rows: Sequence[DailyRow], *, database_url: str | None = None) -> Iterator[None]: + insert_daily_rows(rows, database_url=database_url) + try: + yield + finally: + delete_daily_rows(rows, database_url=database_url) + + +@contextmanager +def spend_log_naming_only_an_alias(request_id: str, api_key: str, started: str, alias: str) -> Iterator[None]: + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + INSERT_SPEND_LOG, (request_id, api_key, started, started, Jsonb({"user_api_key_alias": alias})) + ) + try: + yield + finally: + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_SPEND_LOG, (request_id,)) + + +@dataclass(frozen=True, slots=True) +class SpendLogRow: + started: str + metadata: JsonValue = None + team_id: str | None = None + user: str | None = None + + +def started_at(index: int) -> str: + return (FIRST_SPEND_LOG_AT + timedelta(seconds=index)).strftime("%Y-%m-%d %H:%M:%S") + + +def nameless_rows(count: int, first_index: int = 0) -> tuple[SpendLogRow, ...]: + return tuple(SpendLogRow(started_at(first_index + offset), {}) for offset in range(count)) + + +def named_row(index: int, alias: str) -> SpendLogRow: + return SpendLogRow(started_at(index), {"user_api_key_alias": alias}) + + +@contextmanager +def spend_logs_of_key( + api_key: str, rows: Sequence[SpendLogRow], *, database_url: str | None = None +) -> Iterator[tuple[str, ...]]: + request_ids: Final = tuple(f"integration-{uuid.uuid4().hex}" for _ in rows) + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.cursor().executemany( + INSERT_SPEND_LOG_ROW, + tuple( + (request_id, api_key, row.started, row.started, Jsonb(row.metadata), row.team_id, row.user) + for request_id, row in zip(request_ids, rows, strict=True) + ), + ) + try: + yield request_ids + finally: + delete_spend_logs(request_ids, database_url=database_url) + + +def delete_spend_logs(request_ids: Sequence[str], *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_SPEND_LOG_ROWS, (list(request_ids),)) + + +def purge_key_from_the_key_tables(digest: str, *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_KEY_ROW, (digest,)) + connection.execute(DELETE_ARCHIVED_KEY_ROW, (digest,)) + + +@contextmanager +def locked_table(table: str, *, database_url: str | None = None) -> Iterator[None]: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(LOCK_TABLE.format(table=sql.Identifier(table))) + try: + yield + finally: + connection.rollback() + + +def records_of_key(node: JsonValue, api_key: str) -> tuple[JsonValue, ...]: + if isinstance(node, list): + return tuple(chain.from_iterable(records_of_key(item, api_key) for item in node)) + if not isinstance(node, dict): + return () + nested: Final = tuple(chain.from_iterable(records_of_key(value, api_key) for value in node.values())) + return (node[api_key], *nested) if api_key in node else nested + + +def seeded_row(table: str, entity_column: str, entity: str | None, api_key: str, date: str) -> DailyRow: + return DailyRow(table, entity_column, entity, api_key, date, "gpt-4o-mini", "openai", 10, 5, 0.25, 1, 0) + + +def user_row(user: str | None, api_key: str, date: str) -> DailyRow: + return seeded_row(USER_SPEND, "user_id", user, api_key, date) + + +def seeded_metrics(rows: int) -> dict[str, float]: + return { + "spend": 0.25 * rows, + "prompt_tokens": 10 * rows, + "completion_tokens": 5 * rows, + "total_tokens": 15 * rows, + "api_requests": rows, + "successful_requests": rows, + } + + +def key_metadata( + *, + alias: str | None = None, + team: str | None = None, + user: str | None = None, + email: str | None = None, + exists: bool = False, +) -> dict[str, JsonValue]: + return {"key_alias": alias, "team_id": team, "user_id": user, "user_email": email, "key_exists": exists} + + +def counted(metrics: JsonValue) -> dict[str, JsonValue]: + return {name: value for name, value in object_value(metrics).items() if value} + + +def assert_key_reported( + response: httpx.Response, + api_key: str, + date: str, + metadata: Mapping[str, JsonValue], + metrics: Mapping[str, float], +) -> None: + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + records: Final = tuple(object_value(record) for record in records_of_key(body, api_key)) + assert records, response.text + assert all(record["metadata"] == metadata for record in records), response.text + assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text + days: Final = body["results"] + assert isinstance(days, list) and len(days) == 1, response.text + day: Final = object_value(days[0]) + assert day["date"] == date, response.text + assert counted(day["metrics"]) == pytest.approx(metrics), response.text + assert object_value(body["metadata"])["total_spend"] == pytest.approx(metrics["spend"]), response.text + + +def assert_key_owner_and_totals( + response: httpx.Response, + api_key: str, + metadata: Mapping[str, JsonValue], + totals: Mapping[str, float], +) -> None: + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + records: Final = tuple(object_value(record) for record in records_of_key(body, api_key)) + assert records, response.text + assert all(record["metadata"] == metadata for record in records), response.text + reported: Final = object_value(body["metadata"]) + assert {name: reported[name] for name in totals} == pytest.approx(totals), response.text diff --git a/tests/integration/providers/test_anthropic_messages_live_lifecycle_wire.py b/tests/integration/providers/test_anthropic_messages_live_lifecycle_wire.py new file mode 100644 index 00000000000..cb7043c0362 --- /dev/null +++ b/tests/integration/providers/test_anthropic_messages_live_lifecycle_wire.py @@ -0,0 +1,82 @@ +import json +import threading +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "claude-sonnet-4-5-20250929" +_API_KEY: Final = "synthetic-anthropic-key" + + +def _sse(event: str, payload: dict[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode() + + +def test_messages_stream_message_start_reaches_client_before_content_without_fallback( + gateway: Gateway, +) -> None: + """With no fallback able to take over, the proxy must not hold lifecycle + frames back for a retry that cannot happen: message_start reaches the + client while the upstream is still thinking.""" + gate: Final = threading.Event() + head: Final = _sse("message_start", {"type": "message_start", "message": {"id": "msg_live_1"}}) + _sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ) + tail: Final = ( + _sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + ) + + _sse("content_block_stop", {"type": "content_block_stop", "index": 0}) + + _sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 3}}, + ) + + _sse("message_stop", {"type": "message_stop"}) + ) + prompt: Final = "live-lifecycle-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages" + assert request.headers["x-api-key"] == _API_KEY + body: Final = json.loads(request.body) + assert body["model"] == _MODEL + assert body["stream"] is True + assert body["messages"] == [{"role": "user", "content": prompt}] + return Reply(content_type="text/event-stream", chunks=(head, tail), gate_after_first=gate) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + with gateway.client.stream( + "POST", + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "stream": True, + "messages": [{"role": "user", "content": prompt}], + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read().decode() + lines = response.iter_lines() + first_event: Final = next( + json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ") + ) + assert first_event["type"] == "message_start" + gate.set() + events: Final = (first_event,) + tuple( + json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ") + ) + assert tuple(event["type"] for event in events) == ( + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ), f"observed events: {events!r}" + assert [request.target for request in wire.drain()] == ["/v1/messages"] diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index e700d17ea88..4357ae9a6d6 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -3,8 +3,9 @@ import uuid from typing import Final import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, list_value, object_value from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue @pytest.mark.covers("other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length") @@ -188,3 +189,256 @@ def test_messages_over_responses_deployment_with_max_tokens_one_reaches_openai_a assert len(wire.drain()) == 1 assert body["content"] == [{"type": "text", "text": "ok"}], response.text assert body["usage"]["input_tokens"] == 9 and body["usage"]["output_tokens"] == 1, response.text + + +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + + +def _message_item(item_id: str, text: str) -> dict[str, JsonValue]: + return { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + +_WEATHER_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", +} + +_TWO_MESSAGES_THEN_CALL: Final = ( + _message_item("msg_first", "First message."), + _message_item("msg_second", "Second message."), + _WEATHER_CALL, +) + + +def test_chat_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_two_messages", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "First message.Second message.", + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def _responses_frame(event: str, payload: dict[str, JsonValue]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _message_item_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + item_id: Final = item["id"] + text: Final = object_value(list_value(item["content"])[0])["text"] + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "content": []}, + }, + ), + _responses_frame( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "delta": text, + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _function_call_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "arguments": ""}, + }, + ), + _responses_frame( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "item_id": item["id"], + "output_index": output_index, + "delta": item["arguments"], + }, + ), + _responses_frame( + "response.function_call_arguments.done", + { + "type": "response.function_call_arguments.done", + "item_id": item["id"], + "output_index": output_index, + "arguments": item["arguments"], + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _sse_payloads(lines: list[str]) -> list[dict[str, JsonValue]]: + return [json.loads(line[6:]) for line in lines if line.startswith("data: ") and line != "data: [DONE]"] + + +def test_chat_stream_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-stream-" + uuid.uuid4().hex + response_payload: Final[dict[str, JsonValue]] = { + "id": "resp_two_messages_stream", + "object": "response", + "created_at": 1789788253, + "model": "gpt-6-sol", + "status": "in_progress", + "output": [], + } + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return Reply( + content_type="text/event-stream", + chunks=( + _responses_frame("response.created", {"type": "response.created", "response": response_payload}), + *_message_item_frames(0, _TWO_MESSAGES_THEN_CALL[0]), + *_message_item_frames(1, _TWO_MESSAGES_THEN_CALL[1]), + *_function_call_frames(2, _TWO_MESSAGES_THEN_CALL[2]), + _responses_frame( + "response.completed", + { + "type": "response.completed", + "response": { + **response_payload, + "status": "completed", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + chunks: Final = _sse_payloads(list(response.iter_lines())) + choices: Final = [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])] + assert [choice["index"] for choice in choices] == [0] * len(choices), chunks + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert ( + "".join(str(delta["content"]) for delta in deltas if delta.get("content")) + == "First message.Second message." + ) + tool_calls: Final = [ + object_value(call) + for delta in deltas + if delta.get("tool_calls") + for call in list_value(delta["tool_calls"]) + ] + assert [ + object_value(call["function"])["name"] for call in tool_calls if object_value(call["function"]).get("name") + ] == ["get_weather"], chunks + assert ( + "".join(str(object_value(call["function"]).get("arguments") or "") for call in tool_calls) + == '{"city":"Paris"}' + ), chunks + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], chunks + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/integration/spend/test_daily_activity_key_alias_probes.py b/tests/integration/spend/test_daily_activity_key_alias_probes.py new file mode 100644 index 00000000000..8d9b8616435 --- /dev/null +++ b/tests/integration/spend/test_daily_activity_key_alias_probes.py @@ -0,0 +1,490 @@ +import time +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.daily_activity import ( + AGGREGATED_USER_ACTIVITY, + DAY, + ROUTES, + SPEND_LOGS_TABLE, + USER_SPEND, + Route, + SpendLogRow, + activity_of_key, + assert_key_reported, + daily_rows, + digest_no_key_table_holds, + key_metadata, + locked_table, + named_row, + nameless_rows, + records_of_key, + seeded_metrics, + seeded_row, + spend_logs_of_key, + started_at, + user_row, + user_with_an_email, +) +from integration._support.database import read_rows, scratch_database +from integration._support.process import OwnedProxy, owned_proxy_process +from pydantic import JsonValue + +DAY_OUTSIDE_THE_WINDOW: Final = "2026-02-10" +GIVES_UP_WITHIN_SECONDS: Final = 10 +CONCURRENT_READS: Final = 20 +CACHED_MISS_CLEARS_WITHIN_SECONDS: Final = 45 +ALIAS_OF_ONE_SPEND_LOG: Final = ( + "SELECT metadata->>'user_api_key_alias' AS alias FROM \"LiteLLM_SpendLogs\" WHERE request_id = %s" +) + + +def _alias() -> str: + return f"integration-alias-{uuid.uuid4().hex}" + + +def _named_between_fifty_and_fifty(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(50), named_row(50, alias), *nameless_rows(50, 51)) + + +def _oldest_named(alias: str) -> tuple[SpendLogRow, ...]: + return (named_row(0, alias), *nameless_rows(150, 1)) + + +def _newest_named(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(150), named_row(150, alias)) + + +def _both_edges_named(alias: str) -> tuple[SpendLogRow, ...]: + return (named_row(0, alias), *nameless_rows(150, 1), named_row(151, alias)) + + +def _named_after_one_hundred(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(100), named_row(100, alias), *nameless_rows(99, 101)) + + +def _named_after_ninety_nine(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(99), named_row(99, alias), *nameless_rows(100, 100)) + + +def _named_only_in_the_middle(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(100), named_row(100, alias), *nameless_rows(100, 101)) + + +def _renamed_and_renamed_back(alias: str, other: str) -> tuple[SpendLogRow, ...]: + return ( + named_row(0, alias), + *nameless_rows(100, 1), + named_row(101, other), + *nameless_rows(100, 102), + named_row(202, alias), + ) + + +def _team_in_the_column(team: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {}, team_id=team) + + +def _team_in_the_metadata(team: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {"user_api_key_team_id": team}) + + +def _user_in_the_column(user: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {}, user=user) + + +def _user_in_the_metadata(user: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {"user_api_key_user_id": user}) + + +def _activity_on_route(gateway: Gateway, route: Route, api_key: str, entity: str) -> httpx.Response: + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + return activity_of_key(gateway, route.path, api_key, **filters) + + +def _reported_aliases(response: httpx.Response, api_key: str) -> tuple[JsonValue, ...]: + if response.status_code != 200: + return () + return tuple( + object_value(object_value(record)["metadata"])["key_alias"] + for record in records_of_key(object_value(response.json()), api_key) + ) + + +def _names_the_key(api_key: str, alias: str) -> Callable[[httpx.Response], bool]: + def names(response: httpx.Response) -> bool: + reported: Final = _reported_aliases(response, api_key) + return bool(reported) and frozenset(reported) == frozenset((alias,)) + + return names + + +@contextmanager +def _proxy_on(gateway: Gateway, directory: Path, database_url: str, *, workers: int = 1) -> Iterator[OwnedProxy]: + with owned_proxy_process( + gateway, + directory, + {"DATABASE_URL": database_url}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=workers, + ) as owned: + yield owned + + +def _owner_on(candidate: Gateway) -> tuple[str, str]: + owner: Final = f"integration-{uuid.uuid4().hex}" + email: Final = f"{owner}@example.com" + candidate.post("/user/new", {"user_id": owner, "user_email": email, "auto_create_key": False}) + return owner, email + + +@pytest.mark.parametrize("route", ROUTES, ids=lambda route: route.path.strip("/").replace("/", "_")) +def test_alias_named_only_by_a_spend_log_is_reported_on_every_daily_activity_route( + gateway: Gateway, route: Route +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + entity_rows: Final = ( + () if route.table == USER_SPEND else (seeded_row(route.table, route.entity_column, entity, api_key, DAY),) + ) + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + with ( + daily_rows((user_row(owner, api_key, DAY), *entity_rows)), + spend_logs_of_key(api_key, (named_row(0, alias),)), + ): + assert_key_reported( + activity_of_key(gateway, route.path, api_key, **filters), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "layout", + ( + pytest.param(_named_between_fifty_and_fifty, id="named_between_50_and_50_nameless"), + pytest.param(_oldest_named, id="oldest_named_150_nameless_newer"), + pytest.param(_newest_named, id="newest_named_150_nameless_older"), + pytest.param(_both_edges_named, id="both_edges_named_150_nameless_between"), + pytest.param(_named_after_one_hundred, id="100_nameless_named_99_nameless"), + pytest.param(_named_after_ninety_nine, id="99_nameless_named_100_nameless"), + ), +) +def test_alias_on_an_edge_of_the_window_is_reported_whatever_surrounds_it( + gateway: Gateway, layout: Callable[[str], tuple[SpendLogRow, ...]] +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, layout(alias)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_alias_named_only_in_the_middle_of_two_hundred_nameless_rows_is_not_picked_up(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with ( + daily_rows((user_row(owner, api_key, DAY),)), + spend_logs_of_key(api_key, _named_only_in_the_middle(_alias())), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +def test_key_renamed_and_renamed_back_is_reported_with_the_alias_on_both_edges(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + rows: Final = _renamed_and_renamed_back(alias, _alias()) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "spend_log_of_team", + ( + pytest.param(_team_in_the_column, id="team_id_column"), + pytest.param(_team_in_the_metadata, id="team_id_in_metadata"), + ), +) +def test_team_named_only_by_a_spend_log_is_reported_next_to_the_daily_owner( + gateway: Gateway, spend_log_of_team: Callable[[str], SpendLogRow] +) -> None: + api_key: Final = digest_no_key_table_holds() + team: Final = f"integration-team-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (spend_log_of_team(team),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(team=team, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "spend_log_of_user", + ( + pytest.param(_user_in_the_column, id="user_column"), + pytest.param(_user_in_the_metadata, id="user_id_in_metadata"), + ), +) +def test_user_named_by_a_spend_log_beats_the_owner_the_daily_rows_name( + gateway: Gateway, spend_log_of_user: Callable[[str], SpendLogRow] +) -> None: + api_key: Final = digest_no_key_table_holds() + with gateway.scenario() as scenario: + daily_owner, _ = user_with_an_email(scenario) + log_user, log_email = user_with_an_email(scenario) + with ( + daily_rows((user_row(daily_owner, api_key, DAY),)), + spend_logs_of_key(api_key, (spend_log_of_user(log_user),)), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=log_user, email=log_email), + seeded_metrics(1), + ) + + +def test_hashed_jwt_digest_is_named_by_its_spend_log(gateway: Gateway) -> None: + api_key: Final = f"hashed-jwt-{sha256(uuid.uuid4().bytes).hexdigest()}" + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (named_row(0, alias),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + ("started", "inside_the_window"), + ( + pytest.param("2026-02-01 23:59:59", False, id="second_before_the_window"), + pytest.param("2026-02-02 00:00:00", True, id="first_second_of_the_window"), + pytest.param("2026-02-04 23:59:59", True, id="last_second_of_the_window"), + pytest.param("2026-02-05 00:00:00", False, id="first_second_after_the_window"), + ), +) +def test_spend_log_names_the_key_only_from_one_day_before_to_two_days_after_the_read( + gateway: Gateway, started: str, inside_the_window: bool +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + row: Final = SpendLogRow(started, {"user_api_key_alias": alias}) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias if inside_the_window else None, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_two_aliases_on_the_two_edges_leave_the_key_unnamed(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + rows: Final = (named_row(0, _alias()), *nameless_rows(150, 1), named_row(151, _alias())) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "unnamed_rows", + ( + pytest.param((SpendLogRow(started_at(0), {"user_api_key_alias": ""}),), id="empty_string_alias"), + pytest.param( + (SpendLogRow(started_at(0), ["x"]), SpendLogRow(started_at(1), "x")), id="array_then_string_metadata" + ), + ), +) +def test_rows_without_a_usable_alias_do_not_hide_the_named_row_after_them( + gateway: Gateway, unnamed_rows: tuple[SpendLogRow, ...] +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + rows: Final = (*unnamed_rows, named_row(len(unnamed_rows), alias)) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "stored_alias", + ( + pytest.param(123, id="json_int"), + pytest.param(["a"], id="json_list"), + pytest.param("a" * 5000, id="five_kb_string"), + ), +) +def test_alias_of_an_unexpected_shape_is_reported_as_postgres_renders_it( + gateway: Gateway, stored_alias: JsonValue +) -> None: + api_key: Final = digest_no_key_table_holds() + row: Final = SpendLogRow(started_at(0), {"user_api_key_alias": stored_alias}) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)) as request_ids: + rendered: Final = read_rows(ALIAS_OF_ONE_SPEND_LOG, (request_ids[0],))[0]["alias"] + assert isinstance(rendered, str) and rendered, rendered + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=rendered, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.timeout(300) +def test_alias_found_once_is_served_from_the_cache_for_the_same_window_only(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + rows: Final = (user_row(owner, api_key, DAY), user_row(owner, api_key, DAY_OUTSIDE_THE_WINDOW)) + with daily_rows(rows, database_url=database_url): + with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url): + first: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + cached: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + other_window: Final = owned.gateway.request( + "GET", + AGGREGATED_USER_ACTIVITY, + params={"start_date": DAY_OUTSIDE_THE_WINDOW, "end_date": DAY_OUTSIDE_THE_WINDOW, "api_key": api_key}, + ) + named: Final = key_metadata(alias=alias, user=owner, email=email) + assert_key_reported(first, api_key, DAY, named, seeded_metrics(1)) + assert_key_reported(cached, api_key, DAY, named, seeded_metrics(1)) + assert_key_reported( + other_window, api_key, DAY_OUTSIDE_THE_WINDOW, key_metadata(user=owner, email=email), seeded_metrics(1) + ) + + +@pytest.mark.timeout(300) +def test_alias_logged_after_a_cached_miss_shows_once_the_miss_expires(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + with daily_rows((user_row(owner, api_key, DAY),), database_url=database_url): + missed: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url): + named: Final = eventually( + lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), + _names_the_key(api_key, alias), + seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, + ) + assert_key_reported(missed, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_key_reported(named, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) + + +@pytest.mark.timeout(300) +def test_alias_lookup_gives_up_while_spend_logs_are_locked_and_answers_once_they_are_not( + gateway: Gateway, tmp_path: Path +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url, workers=2) as owned: + owner, email = _owner_on(owned.gateway) + with ( + daily_rows((user_row(owner, api_key, DAY),), database_url=database_url), + spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url), + ): + with locked_table(SPEND_LOGS_TABLE, database_url=database_url): + started: Final = time.monotonic() + locked: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + waited: Final = time.monotonic() - started + unlocked: Final = eventually( + lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), + _names_the_key(api_key, alias), + seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, + ) + assert waited < GIVES_UP_WITHIN_SECONDS, waited + assert_key_reported(locked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_key_reported(unlocked, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) + + +def test_concurrent_reads_over_every_route_all_name_a_fresh_key(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + entity_columns: Final = {route.table: route.entity_column for route in ROUTES if route.table != USER_SPEND} + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + rows: Final = ( + user_row(owner, api_key, DAY), + *(seeded_row(table, column, entity, api_key, DAY) for table, column in entity_columns.items()), + ) + with ( + daily_rows(rows), + spend_logs_of_key(api_key, (named_row(0, alias),)), + ThreadPoolExecutor(CONCURRENT_READS) as pool, + ): + reads: Final = tuple( + pool.submit(_activity_on_route, gateway, ROUTES[index % len(ROUTES)], api_key, entity) + for index in range(CONCURRENT_READS) + ) + responses: Final = tuple(read.result() for read in reads) + for response in responses: + assert_key_reported( + response, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1) + ) diff --git a/tests/integration/spend/test_daily_activity_key_owner.py b/tests/integration/spend/test_daily_activity_key_owner.py new file mode 100644 index 00000000000..cec19ce5ea0 --- /dev/null +++ b/tests/integration/spend/test_daily_activity_key_owner.py @@ -0,0 +1,196 @@ +import uuid +from hashlib import sha256 +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario, string_value +from integration._support.daily_activity import ( + AGGREGATED_USER_ACTIVITY, + DAY, + ROUTES, + USER_SPEND, + Route, + activity_of_key, + assert_key_reported, + daily_rows, + key_metadata, + key_no_key_table_holds, + seeded_metrics, + seeded_row, + spend_log_naming_only_an_alias, + user_row, + user_with_an_email, +) + + +@pytest.mark.parametrize("route", ROUTES, ids=lambda route: route.path.strip("/").replace("/", "_")) +def test_key_missing_from_the_key_tables_is_reported_with_the_one_user_its_daily_spend_names( + gateway: Gateway, route: Route +) -> None: + api_key: Final = key_no_key_table_holds() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + entity_rows: Final = ( + () if route.table == USER_SPEND else (seeded_row(route.table, route.entity_column, entity, api_key, DAY),) + ) + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + with daily_rows((user_row(owner, api_key, DAY), *entity_rows)): + assert_key_reported( + activity_of_key(gateway, route.path, api_key, **filters), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +def test_key_whose_daily_spend_names_two_users_is_reported_with_no_owner(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + first, _ = user_with_an_email(scenario) + second, _ = user_with_an_email(scenario) + with daily_rows((user_row(first, api_key, DAY), user_row(second, api_key, DAY))): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(), + seeded_metrics(2), + ) + + +@pytest.mark.parametrize("unnamed", ["", None], ids=["blank_user", "null_user"]) +def test_daily_spend_rows_naming_no_user_do_not_hide_the_one_user_the_others_name( + gateway: Gateway, unnamed: str | None +) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY), user_row(unnamed, api_key, DAY))): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(2), + ) + + +def test_key_whose_daily_spend_names_no_user_at_all_is_reported_with_no_owner(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with daily_rows((user_row("", api_key, DAY), user_row(None, api_key, DAY))): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(), + seeded_metrics(2), + ) + + +def test_owner_the_user_table_does_not_hold_is_reported_by_id_with_no_email(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + owner: Final = f"integration-departed-{uuid.uuid4().hex}" + with daily_rows((user_row(owner, api_key, DAY),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner), + seeded_metrics(1), + ) + + +def _stored_form(token: str) -> str: + return sha256(token.encode()).hexdigest() + + +def _deleted_key(gateway: Gateway, scenario: Scenario, alias: str, **fields: str) -> str: + token: Final = string_value(gateway.post("/key/generate", {"key_alias": alias, **fields})["key"]) + scenario.delete_key(token) + return _stored_form(token) + + +def test_live_key_keeps_its_own_user_when_its_daily_spend_names_another(gateway: Gateway) -> None: + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + other, _ = user_with_an_email(scenario) + api_key: Final = _stored_form(scenario.key(user_id=owner, key_alias=alias)) + with daily_rows((user_row(other, api_key, DAY),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email, exists=True), + seeded_metrics(1), + ) + + +def test_live_key_with_no_user_is_not_given_the_user_its_daily_spend_names(gateway: Gateway) -> None: + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + spender, _ = user_with_an_email(scenario) + api_key: Final = _stored_form(scenario.key(key_alias=alias)) + with daily_rows((user_row(spender, api_key, DAY),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, exists=True), + seeded_metrics(1), + ) + + +def test_deleted_key_keeps_its_own_user_when_its_daily_spend_names_another(gateway: Gateway) -> None: + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + other, _ = user_with_an_email(scenario) + api_key: Final = _deleted_key(gateway, scenario, alias, user_id=owner) + with daily_rows((user_row(other, api_key, DAY),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_deleted_key_with_no_user_keeps_its_alias_and_gains_the_one_user_its_daily_spend_names( + gateway: Gateway, +) -> None: + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + api_key: Final = _deleted_key(gateway, scenario, alias) + with daily_rows((user_row(owner, api_key, DAY),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_key_named_only_by_a_spend_log_alias_keeps_that_alias_and_gains_the_one_user_its_daily_spend_names( + gateway: Gateway, +) -> None: + api_key: Final = sha256(uuid.uuid4().bytes).hexdigest() + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with ( + spend_log_naming_only_an_alias(f"integration-{uuid.uuid4().hex}", api_key, f"{DAY} 12:00:00", alias), + daily_rows((user_row(owner, api_key, DAY),)), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) diff --git a/tests/integration/spend/test_daily_activity_key_owner_faults.py b/tests/integration/spend/test_daily_activity_key_owner_faults.py new file mode 100644 index 00000000000..998cd2396ae --- /dev/null +++ b/tests/integration/spend/test_daily_activity_key_owner_faults.py @@ -0,0 +1,264 @@ +import os +import signal +import time +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from itertools import chain +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.daily_activity import ( + AGGREGATED_USER_ACTIVITY, + DAY, + TEAM_SPEND, + USER_SPEND, + activity_of_key, + assert_key_reported, + daily_rows, + insert_daily_rows, + key_metadata, + key_no_key_table_holds, + locked_table, + records_of_key, + seeded_metrics, + seeded_row, + user_row, + user_with_an_email, +) +from integration._support.database import scratch_database +from integration._support.process import OwnedProxy, group_members, owned_proxy_process + +USER_ACTIVITY: Final = "/user/daily/activity" +TEAM_ACTIVITY: Final = "/team/daily/activity" +AGGREGATED_TEAM_ACTIVITY: Final = "/team/daily/activity/aggregated" +KEYS_OF_ONE_TEAM: Final = 300 +GIVES_UP_WITHIN_SECONDS: Final = 10 +READS_AFTER_THE_WORKER_IS_REPLACED: Final = 6 + + +@contextmanager +def _proxy_on(gateway: Gateway, directory: Path, database_url: str, *, workers: int = 1) -> Iterator[OwnedProxy]: + with owned_proxy_process( + gateway, + directory, + {"DATABASE_URL": database_url}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=workers, + ) as owned: + yield owned + + +def _owner_on(candidate: Gateway) -> tuple[str, str]: + owner: Final = f"integration-{uuid.uuid4().hex}" + email: Final = f"{owner}@example.com" + candidate.post("/user/new", {"user_id": owner, "user_email": email, "auto_create_key": False}) + return owner, email + + +def _read_on_a_new_connection(candidate: Gateway, api_key: str) -> httpx.Response: + return candidate.request( + "GET", + AGGREGATED_USER_ACTIVITY, + params={"start_date": DAY, "end_date": DAY, "api_key": api_key}, + headers={"Connection": "close"}, + ) + + +def _running_children(owned: OwnedProxy) -> tuple[int, ...]: + return tuple( + member.pid + for member in group_members(owned.process.pid) + if member.pid != owned.process.pid and member.is_running() and member.status() != psutil.STATUS_ZOMBIE + ) + + +def test_user_reading_a_key_shared_with_another_user_is_shown_no_owner_and_nothing_of_the_other_user( + gateway: Gateway, +) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + reader, _ = user_with_an_email(scenario) + other, other_email = user_with_an_email(scenario) + reader_key: Final = scenario.key(user_id=reader) + with daily_rows((user_row(reader, api_key, DAY), user_row(other, api_key, DAY))): + response: Final = activity_of_key(gateway, USER_ACTIVITY, api_key, reader=reader_key) + assert_key_reported(response, api_key, DAY, key_metadata(), seeded_metrics(1)) + assert other not in response.text + assert other_email not in response.text + + +def test_user_reading_a_key_only_they_spent_with_is_shown_themselves_as_its_owner(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + reader, email = user_with_an_email(scenario) + reader_key: Final = scenario.key(user_id=reader) + with daily_rows((user_row(reader, api_key, DAY),)): + response: Final = activity_of_key(gateway, USER_ACTIVITY, api_key, reader=reader_key) + assert_key_reported(response, api_key, DAY, key_metadata(user=reader, email=email), seeded_metrics(1)) + + +def test_user_reading_a_key_only_another_user_spent_with_is_shown_nothing_of_it(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + reader, _ = user_with_an_email(scenario) + other, other_email = user_with_an_email(scenario) + reader_key: Final = scenario.key(user_id=reader) + with daily_rows((user_row(other, api_key, DAY),)): + response: Final = activity_of_key(gateway, USER_ACTIVITY, api_key, reader=reader_key) + assert response.status_code == 200, response.text + assert object_value(response.json())["results"] == [], response.text + assert other not in response.text + assert other_email not in response.text + + +def test_invalid_key_is_refused_without_naming_the_owner(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)): + response: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key, reader="sk-not-a-key") + assert response.status_code == 401, response.text + assert owner not in response.text + assert email not in response.text + + +def test_five_kilobyte_key_is_reported_with_the_one_user_its_daily_spend_names(gateway: Gateway) -> None: + api_key: Final = f"integration-5kb-{uuid.uuid4().hex}-{'k' * 5000}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +def test_key_with_no_daily_spend_is_reported_as_no_activity(gateway: Gateway) -> None: + response: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, key_no_key_table_holds()) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + assert body["results"] == [], response.text + totals: Final = object_value(body["metadata"]) + assert [totals["total_spend"], totals["total_api_requests"]] == [0.0, 0], response.text + + +def test_every_key_of_a_team_is_reported_with_its_own_user(gateway: Gateway) -> None: + team: Final = f"integration-entity-{uuid.uuid4().hex}" + owners: Final = {key_no_key_table_holds(): f"integration-owner-{uuid.uuid4().hex}" for _ in range(KEYS_OF_ONE_TEAM)} + rows: Final = tuple( + chain.from_iterable( + (user_row(owner, api_key, DAY), seeded_row(TEAM_SPEND, "team_id", team, api_key, DAY)) + for api_key, owner in owners.items() + ) + ) + with daily_rows(rows): + response: Final = gateway.request( + "GET", AGGREGATED_TEAM_ACTIVITY, params={"start_date": DAY, "end_date": DAY, "team_ids": team} + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + days: Final = body["results"] + assert isinstance(days, list) and len(days) == 1, response.text + reported: Final = object_value(object_value(object_value(days[0])["breakdown"])["api_keys"]) + assert {api_key: object_value(record)["metadata"] for api_key, record in reported.items()} == { + api_key: key_metadata(user=owner) for api_key, owner in owners.items() + }, response.text + totals: Final = object_value(body["metadata"]) + assert totals["total_api_requests"] == KEYS_OF_ONE_TEAM, response.text + assert totals["total_spend"] == pytest.approx(0.25 * KEYS_OF_ONE_TEAM), response.text + + +def test_reading_the_same_activity_twice_gives_the_same_answer(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + owner, _ = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)): + first: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key) + second: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key) + assert [first.status_code, second.status_code] == [200, 200], [first.text, second.text] + assert records_of_key(first.json(), api_key), first.text + assert first.json() == second.json(), [first.text, second.text] + + +def test_key_stops_being_reported_with_an_owner_once_a_second_user_spends_with_it(gateway: Gateway) -> None: + api_key: Final = key_no_key_table_holds() + with gateway.scenario() as scenario: + first, email = user_with_an_email(scenario) + second, _ = user_with_an_email(scenario) + with daily_rows((user_row(first, api_key, DAY),)): + alone: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key) + with daily_rows((user_row(second, api_key, DAY),)): + shared: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key) + assert_key_reported(alone, api_key, DAY, key_metadata(user=first, email=email), seeded_metrics(1)) + assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2)) + + +@pytest.mark.timeout(300) +def test_owner_lookup_gives_up_while_daily_user_spend_is_locked_and_answers_once_it_is_not( + gateway: Gateway, tmp_path: Path +) -> None: + api_key: Final = key_no_key_table_holds() + team: Final = f"integration-entity-{uuid.uuid4().hex}" + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + rows: Final = (user_row(owner, api_key, DAY), seeded_row(TEAM_SPEND, "team_id", team, api_key, DAY)) + with daily_rows(rows, database_url=database_url): + with locked_table(USER_SPEND, database_url=database_url): + started: Final = time.monotonic() + locked: Final = activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team) + waited: Final = time.monotonic() - started + unlocked: Final = activity_of_key(owned.gateway, TEAM_ACTIVITY, api_key, team_ids=team) + assert waited < GIVES_UP_WITHIN_SECONDS, waited + assert_key_reported(locked, api_key, DAY, key_metadata(), seeded_metrics(1)) + assert_key_reported(unlocked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + + +@pytest.mark.timeout(300) +def test_owner_is_reported_while_a_worker_is_killed_and_after_it_is_replaced(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = key_no_key_table_holds() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url, workers=2) as owned: + owner, email = _owner_on(owned.gateway) + with daily_rows((user_row(owner, api_key, DAY),), database_url=database_url): + before: Final = _read_on_a_new_connection(owned.gateway, api_key) + members: Final = tuple( + member for member in group_members(owned.process.pid) if member.pid != owned.process.pid + ) + children: Final = tuple(member.pid for member in members) + workers: Final = tuple( + member.pid for member in members if any("spawn_main" in part for part in member.cmdline()) + ) + assert len(workers) >= 2, workers + os.kill(workers[0], signal.SIGKILL) + during: Final = _read_on_a_new_connection(owned.gateway, api_key) + eventually( + lambda: _running_children(owned), + lambda pids: len(pids) >= len(children) and any(pid not in children for pid in pids), + seconds=30, + ) + after: Final = tuple( + _read_on_a_new_connection(owned.gateway, api_key) for _ in range(READS_AFTER_THE_WORKER_IS_REPLACED) + ) + for response in (before, during, *after): + assert_key_reported(response, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + + +@pytest.mark.timeout(300) +def test_owner_is_reported_again_after_the_proxy_restarts(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = key_no_key_table_holds() + with scratch_database() as database_url: + with _proxy_on(gateway, tmp_path, database_url) as first: + owner, email = _owner_on(first.gateway) + insert_daily_rows((user_row(owner, api_key, DAY),), database_url=database_url) + before: Final = activity_of_key(first.gateway, AGGREGATED_USER_ACTIVITY, api_key) + with _proxy_on(gateway, tmp_path, database_url) as second: + after: Final = activity_of_key(second.gateway, AGGREGATED_USER_ACTIVITY, api_key) + for response in (before, after): + assert_key_reported(response, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) diff --git a/tests/integration/spend/test_daily_activity_key_owner_traffic.py b/tests/integration/spend/test_daily_activity_key_owner_traffic.py new file mode 100644 index 00000000000..b8f113aca49 --- /dev/null +++ b/tests/integration/spend/test_daily_activity_key_owner_traffic.py @@ -0,0 +1,443 @@ +import json +import os +import threading +import uuid +from collections.abc import Iterable +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta +from hashlib import sha256 +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, string_value +from integration._support.daily_activity import ( + AGGREGATED_USER_ACTIVITY, + DAY, + ROUTES, + USER_SPEND, + Route, + activity_of_key, + assert_key_owner_and_totals, + assert_key_reported, + daily_rows, + key_metadata, + key_no_key_table_holds, + purge_key_from_the_key_tables, + seeded_metrics, + seeded_row, + user_row, + user_with_an_email, +) +from integration._support.database import read_rows, scratch_database +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +from litellm.proxy._types import LiteLLM_UserTable +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + +REQUESTS_OF_KEY: Final = ( + 'SELECT COALESCE(SUM(api_requests), 0)::int AS requests FROM "LiteLLM_DailyUserSpend" ' + "WHERE api_key=%s AND user_id=%s" +) +NAMED_SPEND_LOGS_OF_KEY: Final = ( + 'SELECT COUNT(*)::int AS named FROM "LiteLLM_SpendLogs" ' + "WHERE api_key=%s AND NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL" +) +UNIFIED_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses") +REQUESTS_OF_A_BURST: Final = 21 +READS_DURING_A_BURST: Final = 30 +TOKEN_LIMIT_DISCOVERY: Final = ("GET", "/v1/models") +TOOL_CALL: Final = "call_integration_usage" +ANSWER: Final = "One request cost $0.25" +SUMMARY_OF_ONE_SEEDED_ROW: Final = "\n".join( + ( + "Total Spend: $0.2500", + "Total Requests: 1", + "Successful: 1 | Failed: 0", + "Total Tokens: 15", + "", + "Top Models by Spend:", + " - gpt-4o-mini: $0.2500 (1 reqs, 15 tokens)", + "", + "Top Providers by Spend:", + " - openai: $0.2500 (1 reqs)", + ) +) + + +def _chat_completion() -> dict[str, JsonValue]: + return { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _response() -> dict[str, JsonValue]: + return { + "id": f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{uuid.uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + + +def _provider(request: Request) -> Reply: + body: Final = _response() if request.target.endswith("/responses") else _chat_completion() + return Reply(body=json.dumps(body).encode()) + + +def _usage_tool_call() -> dict[str, JsonValue]: + call: Final[dict[str, JsonValue]] = { + "id": TOOL_CALL, + "type": "function", + "function": { + "name": "get_usage_data", + "arguments": json.dumps({"start_date": DAY, "end_date": DAY}), + }, + } + return { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": None, "tool_calls": [call]}, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _streamed_chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": "chatcmpl-integration-usage", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(chunk)}\n\n".encode() + + +def _usage_analyst(request: Request) -> Reply: + if json.loads(request.body).get("stream"): + return Reply( + chunks=( + _streamed_chunk({"role": "assistant", "content": ANSWER}, None), + _streamed_chunk({}, "stop"), + b"data: [DONE]\n\n", + ), + content_type="text/event-stream", + ) + return Reply(body=json.dumps(_usage_tool_call()).encode()) + + +def _sent_for_callers(requests: Iterable[Request]) -> tuple[Request, ...]: + return tuple(request for request in requests if (request.method, request.target) != TOKEN_LIMIT_DISCOVERY) + + +def _priced_model(scenario: Scenario, provider_url: str) -> str: + return scenario.model( + api_base=f"{provider_url}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002, num_retries=0 + ) + + +def _request_body(endpoint: str, model: str, prompt: str) -> dict[str, JsonValue]: + if endpoint == "/v1/chat/completions": + return {"model": model, "messages": [{"role": "user", "content": prompt}]} + if endpoint == "/v1/messages": + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + return {"model": model, "input": prompt} + + +def _activity_on_route(gateway: Gateway, route: Route, api_key: str, entity: str) -> httpx.Response: + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + return activity_of_key(gateway, route.path, api_key, **filters) + + +def _prompt() -> str: + return f"daily activity owner {uuid.uuid4().hex}" + + +def _totals_of_requests(requests: int) -> dict[str, float]: + return { + "total_spend": 0.02 * requests, + "total_prompt_tokens": 10 * requests, + "total_completion_tokens": 5 * requests, + "total_tokens": 15 * requests, + "total_api_requests": requests, + "total_successful_requests": requests, + "total_failed_requests": 0, + } + + +def _activity_around_today(gateway: Gateway, api_key: str) -> httpx.Response: + today: Final = datetime.now(UTC).date() + return gateway.request( + "GET", + AGGREGATED_USER_ACTIVITY, + params={ + "start_date": str(today - timedelta(days=1)), + "end_date": str(today + timedelta(days=1)), + "timezone": "0", + "api_key": api_key, + }, + ) + + +def _wait_for_requests(api_key: str, user: str, requests: int) -> None: + eventually( + lambda: read_rows(REQUESTS_OF_KEY, (api_key, user)), + lambda rows: rows[0]["requests"] == requests, + seconds=70, + ) + + +def _wait_for_named_spend_logs(api_key: str, requests: int) -> None: + eventually( + lambda: read_rows(NAMED_SPEND_LOGS_OF_KEY, (api_key,)), + lambda rows: rows[0]["named"] == requests, + seconds=70, + ) + + +def _cli_session_token(user: str, team: str) -> str: + cli_user: Final = LiteLLM_UserTable(user_id=user, user_role="internal_user", teams=[team], models=[]) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=cli_user, team_id=team, team_alias="cli-team") + + +def test_key_used_on_every_unified_endpoint_is_reported_with_its_own_alias_and_user(gateway: Gateway) -> None: + chat_prompt, messages_prompt, responses_prompt = _prompt(), _prompt(), _prompt() + with wire_server(_provider) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire.url) + owner, email = user_with_an_email(scenario) + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + key: Final = scenario.key(user_id=owner, key_alias=alias, models=[model]) + stored: Final = sha256(key.encode()).hexdigest() + prompts: Final = (chat_prompt, messages_prompt, responses_prompt) + answers: Final = tuple( + gateway.request("POST", endpoint, _request_body(endpoint, model, prompt), key=key) + for endpoint, prompt in zip(UNIFIED_ENDPOINTS, prompts, strict=True) + ) + assert [answer.status_code for answer in answers] == [200, 200, 200], [answer.text for answer in answers] + received: Final = _sent_for_callers(wire.drain()) + assert [request.target for request in received] == ["/v1/chat/completions", "/v1/responses", "/v1/responses"] + assert [json.loads(request.body)["model"] for request in received] == ["gpt-4o-mini"] * 3 + assert json.loads(received[0].body)["messages"] == [{"role": "user", "content": chat_prompt}] + assert messages_prompt in received[1].body.decode() + assert json.loads(received[2].body)["input"] == responses_prompt + _wait_for_requests(stored, owner, 3) + assert_key_owner_and_totals( + _activity_around_today(gateway, stored), + stored, + key_metadata(alias=alias, user=owner, email=email, exists=True), + _totals_of_requests(3), + ) + + +def test_key_purged_from_the_key_tables_is_reported_with_the_alias_its_spend_logs_name(gateway: Gateway) -> None: + prompts: Final = (_prompt(), _prompt(), _prompt()) + with wire_server(_provider) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire.url) + owner, email = user_with_an_email(scenario) + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + generated: Final = gateway.post("/key/generate", {"user_id": owner, "key_alias": alias, "models": [model]}) + key: Final = string_value(generated["key"]) + stored: Final = sha256(key.encode()).hexdigest() + try: + answers: Final = tuple( + gateway.request("POST", endpoint, _request_body(endpoint, model, prompt), key=key) + for endpoint, prompt in zip(UNIFIED_ENDPOINTS, prompts, strict=True) + ) + assert [answer.status_code for answer in answers] == [200, 200, 200], [answer.text for answer in answers] + received: Final = _sent_for_callers(wire.drain()) + assert [request.target for request in received] == [ + "/v1/chat/completions", + "/v1/responses", + "/v1/responses", + ] + _wait_for_requests(stored, owner, 3) + _wait_for_named_spend_logs(stored, 3) + finally: + purge_key_from_the_key_tables(stored) + assert_key_owner_and_totals( + _activity_around_today(gateway, stored), + stored, + key_metadata(alias=alias, user=owner, email=email, exists=False), + _totals_of_requests(3), + ) + + +def test_cli_session_spend_is_reported_with_the_user_and_team_of_the_session( + gateway: Gateway, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt")) + prompt: Final = _prompt() + with wire_server(_provider) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire.url) + owner, email = user_with_an_email(scenario) + team: Final = scenario.team(models=[model], members_with_roles=[{"role": "user", "user_id": owner}]) + answer: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": prompt}]}, + key=_cli_session_token(owner, team), + ) + assert answer.status_code == 200, answer.text + received: Final = _sent_for_callers(wire.drain()) + assert [request.target for request in received] == ["/v1/chat/completions"] + assert json.loads(received[0].body) == { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": prompt}], + } + stored: Final = f"cli-session-{owner}" + _wait_for_requests(stored, owner, 1) + assert_key_owner_and_totals( + _activity_around_today(gateway, stored), + stored, + key_metadata(alias=stored, team=team, user=owner, email=email), + _totals_of_requests(1), + ) + + +@pytest.mark.timeout(300) +def test_usage_ai_chat_hands_the_model_the_usage_summary_without_any_key_owner( + gateway: Gateway, tmp_path: Path +) -> None: + question: Final = f"what did we spend {uuid.uuid4().hex}" + owner: Final = f"integration-{uuid.uuid4().hex}" + ownerless_key: Final = f"integration-ownerless-{uuid.uuid4().hex}" + with ( + scratch_database() as scratch_url, + wire_server(_usage_analyst) as wire, + owned_proxy( + gateway, + tmp_path, + { + "DATABASE_URL": scratch_url, + "OPENAI_API_BASE": f"{wire.url}/v1", + "OPENAI_BASE_URL": f"{wire.url}/v1", + "OPENAI_API_KEY": "integration-provider-key", + }, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as candidate, + ): + candidate.post("/user/new", {"user_id": owner, "user_email": f"{owner}@example.com", "auto_create_key": False}) + with daily_rows((user_row(owner, ownerless_key, DAY),), database_url=scratch_url): + answer: Final = candidate.request( + "POST", + "/usage/ai/chat", + {"messages": [{"role": "user", "content": question}], "model": "openai/gpt-4o-mini"}, + ) + assert answer.status_code == 200, answer.text + tool_call: Final = { + "type": "tool_call", + "tool_name": "get_usage_data", + "tool_label": "global usage data", + "arguments": {"start_date": DAY, "end_date": DAY}, + } + events: Final = [ + json.loads(line.removeprefix("data: ")) for line in answer.text.splitlines() if line.startswith("data: ") + ] + assert events == [ + {"type": "status", "message": "Thinking..."}, + {**tool_call, "status": "running"}, + {**tool_call, "status": "complete"}, + {"type": "status", "message": "Analyzing results..."}, + {"type": "chunk", "content": ANSWER}, + {"type": "done"}, + ], answer.text + asked, analysed = wire.drain() + assert [asked.target, analysed.target] == ["/v1/chat/completions", "/v1/chat/completions"] + assert json.loads(asked.body)["messages"][-1] == {"role": "user", "content": question} + assert json.loads(analysed.body)["messages"][-1] == { + "role": "tool", + "tool_call_id": TOOL_CALL, + "content": SUMMARY_OF_ONE_SEEDED_ROW, + } + assert owner not in analysed.body.decode() + assert ownerless_key not in analysed.body.decode() + + +@pytest.mark.timeout(300) +def test_owner_is_reported_on_every_route_while_a_burst_of_requests_waits_on_the_provider(gateway: Gateway) -> None: + released: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + + def held_provider(request: Request) -> Reply: + if (request.method, request.target) == TOKEN_LIMIT_DISCOVERY: + return _provider(request) + held.put(request.target) + assert released.wait(timeout=120), "The burst was never released" + return _provider(request) + + api_key: Final = key_no_key_table_holds() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + prompts: Final = tuple(_prompt() for _ in range(REQUESTS_OF_A_BURST)) + entity_columns: Final = {route.table: route.entity_column for route in ROUTES if route.table != USER_SPEND} + with ( + wire_server(held_provider) as wire, + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.client.base_url, timeout=180, trust_env=False) as patient, + ThreadPoolExecutor(max_workers=REQUESTS_OF_A_BURST) as traffic, + ThreadPoolExecutor(max_workers=READS_DURING_A_BURST) as readers, + ): + model: Final = _priced_model(scenario, wire.url) + owner, email = user_with_an_email(scenario) + key: Final = scenario.key(models=[model]) + rows: Final = ( + user_row(owner, api_key, DAY), + *(seeded_row(table, column, entity, api_key, DAY) for table, column in entity_columns.items()), + ) + try: + with daily_rows(rows): + burst: Final = tuple( + traffic.submit( + patient.post, + UNIFIED_ENDPOINTS[index % len(UNIFIED_ENDPOINTS)], + json=_request_body(UNIFIED_ENDPOINTS[index % len(UNIFIED_ENDPOINTS)], model, prompt), + headers={"Authorization": f"Bearer {key}"}, + ) + for index, prompt in enumerate(prompts) + ) + eventually(held.qsize, lambda waiting: waiting >= REQUESTS_OF_A_BURST, seconds=60) + reads: Final = tuple( + readers.submit(_activity_on_route, gateway, ROUTES[index % len(ROUTES)], api_key, entity) + for index in range(READS_DURING_A_BURST) + ) + activity: Final = tuple(read.result() for read in reads) + still_waiting: Final = [call.done() for call in burst] + finally: + released.set() + answers: Final = tuple(call.result() for call in burst) + received: Final = tuple(request.body.decode() for request in _sent_for_callers(wire.drain())) + assert still_waiting == [False] * REQUESTS_OF_A_BURST + assert [answer.status_code for answer in answers] == [200] * REQUESTS_OF_A_BURST, [ + answer.text for answer in answers + ] + assert [sum(prompt in body for body in received) for prompt in prompts] == [1] * REQUESTS_OF_A_BURST + assert len(received) == REQUESTS_OF_A_BURST, len(received) + for response in activity: + assert_key_reported(response, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) diff --git a/tests/integration/spend/test_responses_bridge_spend_once.py b/tests/integration/spend/test_responses_bridge_spend_once.py new file mode 100644 index 00000000000..9eddd832ca0 --- /dev/null +++ b/tests/integration/spend/test_responses_bridge_spend_once.py @@ -0,0 +1,111 @@ +import json +from base64 import b64decode +from collections.abc import Callable +from hashlib import sha256 +from typing import Final +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +PROMPT_TOKENS: Final = 30 +COMPLETION_TOKENS: Final = 5 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE + + +def _responses_reply(prompt: str, response_id: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply( + body=json.dumps({"object": "list", "data": [{"id": "gpt-6.1-sol", "object": "model"}]}).encode() + ) + assert request.method == "POST", request.method + assert request.target == "/v1/responses", request.target + body: Final = object_value(json.loads(request.body)) + encoded_input: Final = json.dumps(body["input"]) + assert prompt in encoded_input, encoded_input + assert body.get("stream") is not True, body + response: Final = { + "id": response_id, + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": body["model"], + "output": [ + { + "type": "message", + "id": f"msg_{uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + } + return Reply(body=json.dumps(response).encode()) + + return respond + + +def _rows_for_key(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT litellm_call_id, request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +@pytest.mark.timeout(120) +def test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_once(gateway: Gateway) -> None: + prompt_prefix: Final = f"spend once {uuid4().hex[:8]} " + long_prompt: Final = (prompt_prefix + "lorem ipsum " * 30_000)[:300_000] + upstream_id: Final = f"resp_{uuid4().hex}" + with ( + wire_server(_responses_reply(long_prompt, upstream_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/responses/gpt-6.1-sol", + api_base=f"{wire.url}/v1", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + digest: Final = sha256(key.encode()).hexdigest() + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": long_prompt}], "stream": False}, + key=key, + ) + assert response.status_code == 200, response.text + assert "ok" in response.text, response.text + assert response.json()["id"] == upstream_id, response.text + call_id: Final = response.headers["x-litellm-call-id"] + + rows: Final = eventually( + lambda: _rows_for_key(key), + lambda values: any(row["litellm_call_id"] == call_id for row in values), + seconds=70, + ) + call_rows: Final = [row for row in rows if row["litellm_call_id"] == call_id] + assert len(call_rows) == 1, rows + assert float(str(call_rows[0]["spend"])) == pytest.approx(EXPECTED_SPEND), rows + logged_id: Final = b64decode(str(call_rows[0]["request_id"]).removeprefix("resp_")).decode() + assert logged_id.startswith("litellm:") and f";response_id:{upstream_id}" in logged_id, rows + + key_spend: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(str(values[0]["spend"])) >= EXPECTED_SPEND, + seconds=70, + ) + assert float(str(key_spend[0]["spend"])) == pytest.approx(EXPECTED_SPEND), key_spend diff --git a/tests/integration/spend/test_stream_alias_billing.py b/tests/integration/spend/test_stream_alias_billing.py new file mode 100644 index 00000000000..a8f7c2a4880 --- /dev/null +++ b/tests/integration/spend/test_stream_alias_billing.py @@ -0,0 +1,284 @@ +"""A streamed alias never replaces the deployment's model for pricing (LIT-9065). + +The proxy shows the client's alias on every streamed chunk, but the chunks kept for end-of-stream cost calculation +keep the deployment's model. "claude-opus-4.8-" is no cost-map key and only matches the claude capability +rules, whose model info carries no prices, so a stream through that alias must bill exactly what the plain alias +"integration-" bills at the same deployment rates, and the client must still see the alias it asked for. +Logging callbacks see that alias as the response model on streamed requests, the same as on non-streamed ones +""" + +import gzip +import json +import threading +from collections import deque +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from pydantic import JsonValue + + +def _sse_event(name: str, payload: dict[str, JsonValue]) -> bytes: + return f"event: {name}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _anthropic_reply(request: Request) -> Reply: + assert request.target.endswith("/v1/messages"), request.target + body: Final = json.loads(request.body) + assert body["model"] == "claude-opus-4-8", body + if body.get("stream") is not True: + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 30, "output_tokens": 40}, + } + ).encode() + ) + return Reply( + content_type="text/event-stream", + chunks=( + _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": f"msg_{uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 30, "output_tokens": 1}, + }, + }, + ), + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + ), + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 40}, + }, + ), + _sse_event("message_stop", {"type": "message_stop"}), + ), + ) + + +def _deployment( + scenario: Scenario, + model_name: str, + litellm_params: dict[str, JsonValue], + model_info: dict[str, JsonValue] | None = None, +) -> str: + created: Final = scenario.gateway.post( + "/model/new", {"model_name": model_name, "litellm_params": litellm_params, "model_info": model_info or {}} + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return model_name + + +def _streamed_spend(gateway: Gateway, scenario: Scenario, model: str, content: str) -> dict[str, JsonValue]: + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": content}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + key=key, + ) + assert response.status_code == 200, response.text + chunks: Final = tuple( + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert chunks and {chunk["model"] for chunk in chunks} == {model}, response.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _listed_deployments(gateway: Gateway, model_name: str) -> tuple[dict[str, JsonValue], ...]: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + return tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model_name) + + +def _deployment_pricing(gateway: Gateway, model_name: str) -> dict[str, JsonValue]: + listed: Final = eventually(lambda: _listed_deployments(gateway, model_name), lambda found: len(found) == 1) + return object_value(listed[0]["model_info"]) + + +_BACKENDS: Final = ( + pytest.param( + lambda _: {"model": "vertex_ai/claude-opus-4-8@default", "mock_response": "hi"}, + id="vertex-mock-response", + ), + pytest.param( + lambda wire_url: { + "model": "anthropic/claude-opus-4-8", + "api_key": "integration-provider-key", + "api_base": wire_url, + }, + id="anthropic-upstream", + ), +) + + +@pytest.mark.parametrize("litellm_params", _BACKENDS) +@pytest.mark.timeout(180) +def test_streamed_alias_matching_a_capability_rule_bills_the_deployment_price( + gateway: Gateway, litellm_params: Callable[[str], dict[str, JsonValue]] +) -> None: + with wire_server(_anthropic_reply) as wire, gateway.scenario() as scenario: + content: Final = f"alias billing {uuid4().hex}" + plain_alias: Final = f"integration-{uuid4().hex}" + rule_alias: Final = f"claude-opus-4.8-{uuid4().int % 10**8:08d}" + exact_row: Final = _streamed_spend( + gateway, scenario, _deployment(scenario, plain_alias, litellm_params(wire.url)), content + ) + alias_row: Final = _streamed_spend( + gateway, scenario, _deployment(scenario, rule_alias, litellm_params(wire.url)), content + ) + + for model_name, row in ((plain_alias, exact_row), (rule_alias, alias_row)): + pricing: Final = _deployment_pricing(gateway, model_name) + input_rate: Final = float(str(pricing["input_cost_per_token"])) + output_rate: Final = float(str(pricing["output_cost_per_token"])) + uplift: Final = float(str(pricing["regional_endpoint_uplift_multiplier"] or 1)) + assert input_rate > 0 and output_rate > 0, pricing + assert float(str(row["spend"])) == pytest.approx( + uplift + * (float(str(row["prompt_tokens"])) * input_rate + float(str(row["completion_tokens"])) * output_rate) + ), (model_name, row, pricing) + + +@dataclass(frozen=True, slots=True) +class _SpanSink: + exports: deque[Request] + guard: threading.Lock + + def receive(self, request: Request) -> Reply: + with self.guard: + self.exports.append(request) + return Reply(body=b"", content_type="application/x-protobuf") + + def span_attributes(self) -> tuple[dict[str, str], ...]: + with self.guard: + exports: Final = tuple(self.exports) + return tuple( + {attribute.key: attribute.value.string_value for attribute in span.attributes} + for span in _exported_spans(exports) + ) + + +def _exported_spans(exports: Sequence[Request]) -> Iterator[object]: + for export in exports: + body = gzip.decompress(export.body) if export.headers.get("content-encoding") == "gzip" else export.body + for resource in ExportTraceServiceRequest.FromString(body).resource_spans: + for scope in resource.scope_spans: + yield from scope.spans + + +@pytest.fixture(scope="module") +def otel_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, _SpanSink]]: + directory: Final = tmp_path_factory.mktemp("stream-alias-otel") + sink: Final = _SpanSink(deque(), threading.Lock()) + with wire_server(sink.receive) as collector, gateway_from_environment() as base: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"] = {**config["litellm_settings"], "callbacks": ["otel"]} + config["callback_settings"] = { + "otel": {"exporter": "http/json", "endpoint": collector.url, "use_simple_processor": True} + } + path: Final = directory / "otel.yaml" + path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": collector.url} + with owned_proxy(base, directory, overrides, config=path) as candidate: + yield candidate, sink + + +def _logged_response_models(sink: _SpanSink, call_ids: Mapping[str, str]) -> dict[str, JsonValue]: + return { + label: span["gen_ai.response.model"] + for span in sink.span_attributes() + for label, call_id in call_ids.items() + if span.get("litellm.call_id") == call_id and "gen_ai.response.model" in span + } + + +@pytest.mark.parametrize("litellm_params", _BACKENDS) +@pytest.mark.timeout(240) +def test_logged_response_model_is_the_client_alias_whether_or_not_the_request_streams( + otel_proxy: tuple[Gateway, _SpanSink], litellm_params: Callable[[str], dict[str, JsonValue]] +) -> None: + candidate, sink = otel_proxy + with wire_server(_anthropic_reply) as wire, candidate.scenario() as scenario: + alias: Final = f"claude-opus-4.8-{uuid4().int % 10**8:08d}" + key: Final = scenario.key(models=[_deployment(scenario, alias, litellm_params(wire.url))]) + call_ids: Final[dict[str, str]] = {} + for label, stream_fields in ( + ("non-streamed", {}), + ("streamed", {"stream": True, "stream_options": {"include_usage": True}}), + ): + response = candidate.request( + "POST", + "/v1/chat/completions", + {"model": alias, "messages": [{"role": "user", "content": f"logged alias {uuid4().hex}"}]} + | stream_fields, + key=key, + ) + assert response.status_code == 200, response.text + call_ids[label] = response.headers["x-litellm-call-id"] + logged: Final = eventually( + lambda: _logged_response_models(sink, call_ids), + lambda found: len(found) == 2, + seconds=60, + return_last_on_timeout=True, + ) + assert logged == {"non-streamed": alias, "streamed": alias}, call_ids diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index aebcd2b47ec..60f458eb1d1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -19,6 +20,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, update_metrics, ) +from litellm.proxy.utils import hash_token from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, SpendMetrics, @@ -503,6 +505,7 @@ async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_s mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma.db.query_raw = AsyncMock(return_value=[]) + recovery_query_raw = _recovery_transaction(mock_prisma) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -510,9 +513,10 @@ async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_s ) assert double_hashed not in result - issued_sql = [call.args[0] for call in mock_prisma.db.query_raw.call_args_list] - assert len(issued_sql) == 2 - assert not any("LiteLLM_SpendLogs" in sql for sql in issued_sql) + assert mock_prisma.db.query_raw.await_count == 2 + ((owner_sql, owner_keys),) = [call.args for call in recovery_query_raw.call_args_list] + assert _DAILY_USER_SPEND in owner_sql + assert owner_keys == [double_hashed] token_lookups = ( mock_prisma.db.litellm_verificationtoken.find_many.call_args_list + mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list @@ -520,14 +524,29 @@ async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_s assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups) -def _spend_log_transaction(mock_prisma: MagicMock, rows: list[dict[str, str | None]]) -> AsyncMock: +_DAILY_USER_SPEND: Final = '"LiteLLM_DailyUserSpend"' +_SPEND_LOGS: Final = '"LiteLLM_SpendLogs"' + + +def _recovery_transaction( + mock_prisma: MagicMock, + spend_log_rows: Sequence[dict[str, str | None]] = (), + daily_spend_owner_rows: Sequence[dict[str, str | None]] = (), +) -> AsyncMock: + async def query_raw(sql: str, *_: object) -> Sequence[dict[str, str | None]]: + return daily_spend_owner_rows if _DAILY_USER_SPEND in sql else spend_log_rows + transaction = MagicMock() transaction.execute_raw = AsyncMock(return_value=0) - transaction.query_raw = AsyncMock(return_value=rows) + transaction.query_raw = AsyncMock(side_effect=query_raw) mock_prisma.db.tx.return_value.__aenter__.return_value = transaction return transaction.query_raw +def _calls_reading(query_raw: AsyncMock, table: str) -> tuple[tuple[object, ...], ...]: + return tuple(call.args for call in query_raw.call_args_list if table in call.args[0]) + + def _spend_log_row(digest: str, key_alias: str, user_id: str) -> dict[str, str | None]: return { "digest": digest, @@ -551,15 +570,17 @@ async def test_get_api_key_metadata_permanent_miss_with_a_window_reads_spend_log mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma.db.query_raw = AsyncMock(return_value=[]) - spend_log_query_raw = _spend_log_transaction(mock_prisma, []) + recovery_query_raw = _recovery_transaction(mock_prisma) result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={double_hashed}, spend_logs_window=window) assert double_hashed not in result assert mock_prisma.db.query_raw.await_count == 2 - ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] + ((_, digests, start, end),) = _calls_reading(recovery_query_raw, _SPEND_LOGS) assert digests == [double_hashed] assert (start, end) == window + ((_, owner_keys),) = _calls_reading(recovery_query_raw, _DAILY_USER_SPEND) + assert owner_keys == [double_hashed] @pytest.mark.asyncio @@ -581,7 +602,7 @@ async def test_get_daily_activity_recovers_a_session_key_alias_from_spend_logs_a ) mock_prisma.db.query_raw = AsyncMock(return_value=[]) - spend_log_query_raw = _spend_log_transaction( + spend_log_query_raw = _recovery_transaction( mock_prisma, [_spend_log_row(session_digest, "cli-session-alias", "session-user")] ) @@ -2370,7 +2391,7 @@ async def test_get_api_key_metadata_resolves_session_key_via_spend_log_window(): ) mock_prisma.db.query_raw = AsyncMock(return_value=[]) - spend_log_query_raw = _spend_log_transaction( + spend_log_query_raw = _recovery_transaction( mock_prisma, [_spend_log_row(session_digest, "cli-session-user-42", "user-42")] ) @@ -2417,3 +2438,90 @@ async def test_get_api_key_metadata_resolves_cli_session_keys_from_the_key_itsel assert result["cli-session-alice"]["key_alias"] == "cli-session-alice" assert result["cli-session-alice"]["user_email"] == "alice@example.com" assert result["cli-session-alice"]["team_id"] == "team-a" + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_recovers_legacy_hashed_jwt_owner_from_daily_spend(): + api_key: Final = f"hashed-jwt-{hash_token('legacy-cli-session-daily-spend-owner')}" + user_id: Final = "legacy-owner" + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id=user_id, user_email="legacy-owner@example.com", teams=[])] + ) + recovery_query_raw: Final = _recovery_transaction( + mock_prisma, + daily_spend_owner_rows=[{"api_key": api_key, "first_owner": user_id, "last_owner": user_id}], + ) + + result: Final = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={api_key}, + spend_logs_window=(datetime(2026, 9, 7), datetime(2026, 9, 10)), + ) + + assert result.get(api_key, {}).get("user_id") == user_id + assert result.get(api_key, {}).get("user_email") == "legacy-owner@example.com" + assert len(_calls_reading(recovery_query_raw, _SPEND_LOGS)) == 1 + assert len(_calls_reading(recovery_query_raw, _DAILY_USER_SPEND)) == 1 + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_preserves_deleted_key_metadata_when_recovering_daily_spend_owner(): + api_key: Final = f"hashed-jwt-{hash_token('legacy-cli-session-daily-spend-metadata')}" + user_id: Final = "legacy-owner" + mock_prisma: Final = MagicMock() + deleted_key: Final = MagicMock() + deleted_key.token = api_key + deleted_key.key_alias = "legacy-cli-key" + deleted_key.team_id = "team-legacy" + deleted_key.user_id = None + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[deleted_key]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id=user_id, user_email="legacy-owner@example.com", teams=[])] + ) + recovery_query_raw: Final = _recovery_transaction( + mock_prisma, + daily_spend_owner_rows=[{"api_key": api_key, "first_owner": user_id, "last_owner": user_id}], + ) + + result: Final = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={api_key}) + + recovered_metadata: Final = result[api_key] + assert recovered_metadata.get("key_alias") == "legacy-cli-key" + assert recovered_metadata.get("team_id") == "team-legacy" + assert recovered_metadata.get("user_id") == user_id + assert recovered_metadata.get("user_email") == "legacy-owner@example.com" + recovery_query_raw.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_does_not_recover_daily_spend_owner_for_active_keys(): + api_key: Final = "active-token-value" + mock_prisma: Final = MagicMock() + active_key: Final = MagicMock() + active_key.token = api_key + active_key.key_alias = "active-key-alias" + active_key.team_id = "active-team" + active_key.user_id = "active-owner" + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[active_key]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="active-owner", user_email="active-owner@example.com", teams=[])] + ) + recovery_query_raw: Final = _recovery_transaction( + mock_prisma, + daily_spend_owner_rows=[{"api_key": api_key, "first_owner": "other-owner", "last_owner": "other-owner"}], + ) + + result: Final = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={api_key}) + + active_metadata: Final = result[api_key] + assert active_metadata.get("key_alias") == "active-key-alias" + assert active_metadata.get("team_id") == "active-team" + assert active_metadata.get("user_id") == "active-owner" + assert active_metadata.get("user_email") == "active-owner@example.com" + assert active_metadata.get("key_exists") is True + recovery_query_raw.assert_not_awaited() diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 86dd356e5f5..228f5a4c16d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -283,8 +283,9 @@ def test_restamp_streaming_chunk_model_overrides_model_on_basemodel(): "model": new_chunk.model, "logged": logged, "same_object": new_chunk is chunk, + "original_model": chunk.model, } - assert snapshot == {"model": "gpt-4", "logged": True, "same_object": True} + assert snapshot == {"model": "gpt-4", "logged": True, "same_object": False, "original_model": "openai/internal-x"} @pytest.mark.parametrize("return_raw_model_name", [False, True]) @@ -310,8 +311,7 @@ def test_restamp_streaming_chunk_model_overrides_model_on_dict(): request_data={}, model_mismatch_logged=True, ) - assert new_chunk["model"] == "gpt-4" - assert logged is True + assert (new_chunk["model"], chunk["model"], logged) == ("gpt-4", "internal", True) def test_restamp_streaming_chunk_model_uses_fallback_model_from_metadata(): @@ -443,7 +443,7 @@ def test_restamp_streaming_chunk_model_fastest_response_preserves_model(): assert logged is False -def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns(): +def test_restamp_streaming_chunk_model_restamps_a_frozen_chunk_through_a_copy(): from pydantic import ConfigDict class FrozenChunk(_simple_chunk().__class__): @@ -462,8 +462,30 @@ def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns(): request_data={"litellm_call_id": "test-id"}, model_mismatch_logged=False, ) - assert new_chunk.model == "openai/internal-x" - assert logged is True + assert (new_chunk.model, chunk.model, logged) == ("gpt-4", "openai/internal-x", True) + + +def test_restamp_streaming_chunk_model_records_the_client_model_on_the_logging_object(): + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="openai/internal-x", + messages=[], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="test-id", + function_id="test-id", + ) + _restamp_streaming_chunk_model( + chunk=_simple_chunk(model="openai/internal-x"), + requested_model_from_client="gpt-4", + request_data={"litellm_call_id": "test-id", "litellm_logging_obj": logging_obj}, + model_mismatch_logged=False, + ) + assert logging_obj.client_facing_stream_model == "gpt-4" def test_format_fallback_metadata_sse_event(): diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 1967d7b6aad..aeac491c959 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,18 +1,25 @@ import asyncio +import re import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock +import psycopg import pytest from prisma.errors import PrismaError +from psycopg.rows import dict_row +from psycopg.types.json import Jsonb +from pytest_postgresql import factories from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, ) from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, @@ -20,6 +27,7 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import ( recover_cli_session_key_metadata, recover_double_hashed_key_metadata, recover_key_metadata_from_spend_logs, + recover_key_owner_from_daily_spend, ) from litellm.proxy.utils import hash_token @@ -586,12 +594,311 @@ async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_state await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) - assert calls == [f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", "scan"] + assert calls == [ + f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", + "SET LOCAL enable_bitmapscan = off", + "scan", + ] assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS ) +_spend_logs_postgresql_proc: Final = factories.postgresql_proc() +_spend_logs_postgresql: Final = factories.postgresql("_spend_logs_postgresql_proc") + +_SPEND_LOGS_DDL: Final = """ + CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT PRIMARY KEY, + api_key TEXT NOT NULL DEFAULT '', + "startTime" TIMESTAMP(3) NOT NULL, + "user" TEXT DEFAULT '', + team_id TEXT, + metadata JSONB DEFAULT '{}' + ) +""" + +_API_KEY_START_TIME_INDEX_DDL: Final = ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime")' +) + +_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL: Final = """ + SELECT COALESCE(seq_tup_read, 0) + COALESCE(idx_tup_fetch, 0) AS rows_read + FROM pg_stat_xact_user_tables + WHERE relname = 'LiteLLM_SpendLogs' +""" + + +def _create_spend_logs_table(conn: psycopg.Connection) -> None: + conn.execute(_SPEND_LOGS_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_API_KEY_START_TIME_INDEX_DDL) # pyright: ignore[reportArgumentType] # DDL literal + + +def _psycopg_prisma(conn: psycopg.Connection) -> MagicMock: + async def query_raw(sql: str, *params: object) -> list[dict[str, object]]: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", sql), # pyright: ignore[reportArgumentType] # proxy SQL is not a literal + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + return cur.fetchall() + + async def execute_raw(sql: str) -> int: + conn.execute(sql) # pyright: ignore[reportArgumentType] # proxy SQL is not a literal + return 0 + + mock_prisma: Final = MagicMock() + transaction: Final = MagicMock() + transaction.query_raw = AsyncMock(side_effect=query_raw) + transaction.execute_raw = AsyncMock(side_effect=execute_raw) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return mock_prisma + + +def _commit_and_vacuum(conn: psycopg.Connection) -> None: + conn.commit() + conn.set_autocommit(True) + conn.execute('VACUUM (ANALYZE) "LiteLLM_SpendLogs"') + conn.set_autocommit(False) + + +def _insert_nameless_spend_logs(conn: psycopg.Connection, digest: str, rows: int) -> None: + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime") + SELECT %(digest)s || '-' || g, %(digest)s, %(start)s + g * interval '1 minute' + FROM generate_series(1, %(rows)s) g + """, + {"digest": digest, "start": datetime(2026, 9, 7), "rows": rows}, + ) + + +def _named_spend_log( + digest: str, logged_at: datetime, alias: str | None, user: str | None, team: str | None = None +) -> tuple[str, str, datetime, str, str | None, Jsonb]: + return ( + f"{digest}-{logged_at.isoformat()}", + digest, + logged_at, + user or "", + team, + Jsonb({"user_api_key_alias": alias} if alias else {}), + ) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_names_a_key_by_its_oldest_and_newest_named_rows_in_the_window( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + unnamed_edges, owner_logged_late, reowned, outside_window, never_named = ( + hash_token(f"cli-session-{name}") for name in ("edges", "late", "reowned", "window", "never") + ) + with conn.cursor() as cur: + cur.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", team_id, metadata)' + " VALUES (%s, %s, %s, %s, %s, %s)", + ( + _named_spend_log(unnamed_edges, datetime(2026, 9, 7, 1), None, None), + _named_spend_log(unnamed_edges, datetime(2026, 9, 8), "cli-a", "alice", "team-a"), + _named_spend_log(unnamed_edges, datetime(2026, 9, 9), "cli-a", "alice", "team-a"), + _named_spend_log(unnamed_edges, datetime(2026, 9, 9, 23), None, None), + _named_spend_log(owner_logged_late, datetime(2026, 9, 7, 1), "cli-b", None), + _named_spend_log(owner_logged_late, datetime(2026, 9, 9), "cli-b", "bob"), + _named_spend_log(reowned, datetime(2026, 9, 7, 1), "cli-c", "carol"), + _named_spend_log(reowned, datetime(2026, 9, 9), "cli-c", "dave"), + _named_spend_log(outside_window, datetime(2026, 9, 6), "stale-alias", "erin"), + _named_spend_log(outside_window, datetime(2026, 9, 8), "cli-d", "erin"), + _named_spend_log(outside_window, datetime(2026, 9, 10), "later-alias", "erin"), + _named_spend_log(never_named, datetime(2026, 9, 8), None, None), + ), + ) + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + {unnamed_edges, owner_logged_late, reowned, outside_window, never_named}, + (datetime(2026, 9, 7), datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert dict(result) == { + unnamed_edges: {"key_alias": "cli-a", "team_id": "team-a", "user_id": "alice"}, + owner_logged_late: {"key_alias": "cli-b", "team_id": None, "user_id": "bob"}, + reowned: {"key_alias": "cli-c", "team_id": None, "user_id": None}, + outside_window: {"key_alias": "cli-d", "team_id": None, "user_id": "erin"}, + } + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_reads_two_rows_per_key_however_many_the_key_logged( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + owners: Final[Mapping[str, str]] = {hash_token(f"cli-session-busy-{i}"): f"user-{i}" for i in range(5)} + for digest, owner in owners.items(): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", metadata) + SELECT %(digest)s || '-' || g, %(digest)s, %(start)s + g * interval '1 minute', %(owner)s, + jsonb_build_object('user_api_key_alias', 'cli-session-' || %(owner)s) + FROM generate_series(1, 2000) g + """, + {"digest": digest, "owner": owner, "start": datetime(2026, 9, 7)}, + ) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), frozenset(owners), (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert {digest: meta.get("user_id") for digest, meta in result.items()} == owners + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None and rows_read[0] <= 2 * len(owners) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_walks_a_bounded_number_of_nameless_rows_per_key( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + named_late: Final[Mapping[str, str]] = {hash_token(f"cli-session-late-{i}"): f"user-{i}" for i in range(3)} + never_named: Final = frozenset(hash_token(f"cli-session-never-{i}") for i in range(3)) + for digest in (*named_late, *never_named): + _insert_nameless_spend_logs(conn, digest, 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + for digest, owner in named_late.items(): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", metadata) + VALUES (%(digest)s || '-newest', %(digest)s, %(logged_at)s, %(owner)s, + jsonb_build_object('user_api_key_alias', 'cli-session-' || %(owner)s)) + """, + {"digest": digest, "owner": owner, "logged_at": datetime(2026, 9, 9)}, + ) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + frozenset(named_late) | never_named, + (datetime(2026, 9, 7), datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert {digest: meta.get("user_id") for digest, meta in result.items()} == named_late + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * (len(named_late) + len(never_named)) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_reads_a_short_nameless_key_once( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + rows_per_key: Final = SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE // 2 + never_named: Final = frozenset(hash_token(f"cli-session-short-{i}") for i in range(20)) + for digest in never_named: + _insert_nameless_spend_logs(conn, digest, rows_per_key) + _commit_and_vacuum(conn) + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), never_named, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert dict(result) == {} + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= rows_per_key * len(never_named) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_bounds_a_busy_nameless_key_among_short_keys_before_any_vacuum( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + busy: Final = frozenset(hash_token(f"cli-session-busy-nameless-{i}") for i in range(3)) + for short_key in range(200): + _insert_nameless_spend_logs( + conn, hash_token(f"cli-session-short-{short_key}"), SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE // 5 + ) + for digest in busy: + _insert_nameless_spend_logs(conn, digest, 30 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), busy, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert dict(result) == {} + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * len(busy) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_finds_a_name_logged_where_the_oldest_probe_stopped( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + start: Final = datetime(2026, 9, 7) + past_the_stop, tied_with_the_stop = (hash_token(f"cli-session-{name}") for name in ("past", "tied")) + same_millisecond: Final = tuple( + start + timedelta(minutes=SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, microseconds=n) for n in (100, 200, 300) + ) + with conn.cursor() as cur: + cur.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", team_id, metadata)' + " VALUES (%s, %s, %s, %s, %s, %s)", + ( + *( + _named_spend_log(past_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range(1, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 20) + ), + _named_spend_log( + past_the_stop, + start + timedelta(minutes=SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 20), + "cli-p", + "pat", + ), + *( + _named_spend_log(past_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range( + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 21, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 51 + ) + ), + *( + _named_spend_log(tied_with_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range(1, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + ), + _named_spend_log(tied_with_the_stop, same_millisecond[0], None, None), + _named_spend_log(tied_with_the_stop, same_millisecond[1], None, None), + _named_spend_log(tied_with_the_stop, same_millisecond[2], "cli-t", "tess"), + ), + ) + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + {past_the_stop, tied_with_the_stop}, + (start, datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert dict(result) == { + past_the_stop: {"key_alias": "cli-p", "team_id": None, "user_id": "pat"}, + tied_with_the_stop: {"key_alias": "cli-t", "team_id": None, "user_id": "tess"}, + } + + @pytest.mark.asyncio async def test_recover_cli_session_key_metadata_names_the_owner_only_when_the_suffix_is_a_real_user(): mock_prisma = MagicMock() @@ -664,3 +971,91 @@ async def test_attach_user_details_claims_no_team_for_a_multi_team_user_session_ assert "team_id" not in attached["cli-session-bob"] assert attached["cli-session-bob"]["user_email"] == "bob@example.com" + + +def _daily_spend_owner_row(api_key: str, first_owner: str, last_owner: str) -> dict[str, str]: + return {"api_key": api_key, "first_owner": first_owner, "last_owner": last_owner} + + +def _daily_spend_transaction(mock_prisma: MagicMock, query_raw: AsyncMock) -> MagicMock: + transaction: Final = MagicMock() + transaction.execute_raw = AsyncMock(return_value=0) + transaction.query_raw = query_raw + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return transaction + + +@pytest.mark.asyncio +async def test_recover_key_owner_from_daily_spend_keeps_a_unanimous_owner(): + key: Final = "hashed-jwt-digest-a" + mock_prisma: Final = MagicMock() + _daily_spend_transaction(mock_prisma, AsyncMock(return_value=[_daily_spend_owner_row(key, "owner-a", "owner-a")])) + + result: Final = await recover_key_owner_from_daily_spend(mock_prisma, {key}) + + assert dict(result) == {key: "owner-a"} + + +@pytest.mark.asyncio +async def test_recover_key_owner_from_daily_spend_drops_conflicting_owners(): + key: Final = "hashed-jwt-digest-b" + mock_prisma: Final = MagicMock() + _daily_spend_transaction(mock_prisma, AsyncMock(return_value=[_daily_spend_owner_row(key, "owner-a", "owner-b")])) + + result: Final = await recover_key_owner_from_daily_spend(mock_prisma, {key}) + + assert dict(result) == {} + + +@pytest.mark.asyncio +async def test_recover_key_owner_from_daily_spend_skips_empty_input(): + mock_prisma: Final = MagicMock() + transaction: Final = _daily_spend_transaction(mock_prisma, AsyncMock(return_value=[])) + + result: Final = await recover_key_owner_from_daily_spend(mock_prisma, frozenset()) + + assert dict(result) == {} + transaction.query_raw.assert_not_awaited() + mock_prisma.db.tx.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_owner_from_daily_spend_returns_empty_on_prisma_error(): + mock_prisma: Final = MagicMock() + _daily_spend_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("db down"))) + + result: Final = await recover_key_owner_from_daily_spend(mock_prisma, {"hashed-jwt-digest-c"}) + + assert dict(result) == {} + + +@pytest.mark.asyncio +async def test_recover_key_owner_from_daily_spend_names_no_owner_when_the_lookup_hits_the_statement_timeout(): + mock_prisma: Final = MagicMock() + _daily_spend_transaction( + mock_prisma, AsyncMock(side_effect=PrismaError("canceling statement due to statement timeout")) + ) + + result: Final = await recover_key_owner_from_daily_spend(mock_prisma, {"hashed-jwt-digest-d"}) + + assert dict(result) == {} + + +@pytest.mark.asyncio +async def test_recover_key_owner_from_daily_spend_bounds_the_lookup_with_a_statement_timeout(): + key: Final = "hashed-jwt-digest-e" + mock_prisma: Final = MagicMock() + transaction: Final = _daily_spend_transaction( + mock_prisma, AsyncMock(return_value=[_daily_spend_owner_row(key, "owner-a", "owner-a")]) + ) + + result: Final = await recover_key_owner_from_daily_spend(mock_prisma, {key}) + + assert dict(result) == {key: "owner-a"} + assert [name for name, _, _ in transaction.mock_calls] == ["execute_raw", "query_raw"] + transaction.execute_raw.assert_awaited_once_with( + f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" + ) + assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( + milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS + ) diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index d0f9bad795d..63b2a339282 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -4352,3 +4352,63 @@ def test_map_optional_params_verbosity_merges_into_text(): verbosity_only_request, ) assert verbosity_only_request["text"] == {"verbosity": "low"} + + +def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None: + from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage + + first: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_first", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "First message.", "annotations": []}], + } + ) + second_annotation: Final = { + "type": "url_citation", + "start_index": 0, + "end_index": 6, + "title": "Second", + "url": "https://example.com/second", + } + second: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Second message.", "annotations": [second_annotation]}], + } + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_1", + name="get_weather", + ) + + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly + (first, second, function_call) + ) + + assert len(choices) == 1 + assert choices[0].index == 0 + assert choices[0].finish_reason == "tool_calls" + assert choices[0].message.content == "First message.Second message." + assert choices[0].message.annotations == [ + {**second_annotation, "start_index": len("First message."), "end_index": len("First message.") + 6} + ] + tool_calls: Final = choices[0].message.tool_calls + assert tool_calls is not None + assert [call.model_dump(exclude_none=True) for call in tool_calls] == [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ] diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index f787d370f04..1d70a21af7b 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -2633,6 +2633,42 @@ class TestConfiguredInjectionPointsSurviveClientMarks: assert kwargs["cache_control"] is root_cache_control assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + @pytest.mark.parametrize( + "tools,kwargs,injected", + [ + ([MARKED_V1_TOOL], {"extra_body": {"tools": [UNMARKED_V1_TOOL]}}, False), + (None, {"cache_control": EPHEMERAL, "extra_body": {"cache_control": None}}, False), + ([UNMARKED_V1_TOOL], {"extra_body": {"tools": [UNMARKED_V1_TOOL]}}, True), + ], + ids=["extra_body_unmarks_direct_tool", "extra_body_nulls_root_cache_control", "no_client_mark_anywhere"], + ) + def test_v1_messages_automatic_defaults_stand_down_for_a_direct_mark_extra_body_hides( + self, monkeypatch, tools, kwargs, injected + ): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request_kwargs = {**copy.deepcopy(kwargs), "litellm_metadata": {}} + + result_messages, result_system = self._inject( + copy.deepcopy(self.V1_MESSAGES), request_kwargs, tools=copy.deepcopy(tools) + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == ( + 2 if injected else 0 + ) + assert ("litellm_gateway_injected_cache" in request_kwargs["litellm_metadata"]) is injected + + def test_chat_automatic_defaults_apply_when_extra_body_drops_the_only_client_mark(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + params = {"extra_body": {"tools": [self.UNMARKED_TOOL]}} + + self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES), tools=[self.MARKED_TOOL_TOP_LEVEL]) + affinity = AnthropicCacheControlHook.messages_with_default_injections( + copy.deepcopy(self.CLEAN_MESSAGES), ["claude-sonnet-4-5"], tools=[self.MARKED_TOOL_TOP_LEVEL], request_kwargs=params + ) + + assert [p["index"] for p in params["cache_control_injection_points"]] == [None, -1] + assert AnthropicCacheControlHook.count_request_cache_breakpoints(affinity) == 2 + @pytest.mark.parametrize( "marked_turns,expected_system", [(2, [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]), (3, "sys")], diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d9ddc3e950b..8cff0e04a8f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -321,6 +321,18 @@ async def test_mcp_direct_content_edit_invalidates_stale_structured_data(logging assert "SECRET-1234" not in result.model_dump_json() +def test_with_client_facing_stream_model_stamps_a_copy_of_the_priced_response(logging_obj): + response = ModelResponse(model="claude-opus-4-6@default") + logging_obj.client_facing_stream_model = "claude-opus-4.6" + logged = logging_obj._with_client_facing_stream_model(response) + assert (logged.model, response.model) == ("claude-opus-4.6", "claude-opus-4-6@default") + + +def test_with_client_facing_stream_model_keeps_the_response_when_the_proxy_set_no_model(logging_obj): + response = ModelResponse(model="claude-opus-4-6@default") + assert logging_obj._with_client_facing_stream_model(response) is response + + def test_get_combined_callback_list_preserves_insertion_order(logging_obj): assert logging_obj.get_combined_callback_list( dynamic_success_callbacks=["prometheus", "langfuse", "datadog", "otel", "s3"], @@ -5780,7 +5792,7 @@ def _interactions_logging_obj(stream: bool, call_type: str = "acreate"): messages=[], stream=stream, call_type=call_type, - start_time=time.time(), + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc).timestamp(), litellm_call_id="interactions-call-id", function_id="interactions-fn-id", ) @@ -8987,6 +8999,8 @@ async def test_prompt_management_with_unchanged_variables_replays_a_byte_identic assert json.dumps(messages_n_plus_one[: len(messages_n)], sort_keys=True) == json.dumps(messages_n, sort_keys=True) assert messages_n[0] == {"role": "system", "content": "You are a pirate. Answer in one sentence."} assert len(messages_n_plus_one) == len(messages_n) + 2 + + @pytest.mark.parametrize( ("input_cost_per_token", "output_cost_per_token"), [(1e-06, 2e-06), (0.0, 0.0)], @@ -9070,3 +9084,59 @@ def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_b litellm.model_cost.pop(deployment_id, None) +@pytest.mark.asyncio +async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch): + """ + The create's in_progress success handler can still be parked on awaited work when + the settlement arrives. The completed result is the only event carrying usage and + cost, so it has to reach the success callbacks whatever the one-success-log-per-request + dedupe does, or the settlement is marked billed with no spend row behind it. + """ + from litellm.litellm_core_utils import litellm_logging + from litellm.types.interactions import InteractionsAPIResponse + + class CountingLogger(CustomLogger): + def __init__(self): + super().__init__() + self.logged_results: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_results.append(response_obj) + + counting_logger: Final = CountingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger]) + + in_progress_parked: Final = asyncio.Event() + release_in_progress: Final = asyncio.Event() + real_truncate: Final = litellm_logging.truncate_base64_in_messages_async + calls: list[int] = [] + + async def parked_then_real_truncate(messages): + calls.append(1) + if len(calls) == 1: + in_progress_parked.set() + await release_in_progress.wait() + return await real_truncate(messages) + + monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", parked_then_real_truncate) + + logging_obj: Final = _interactions_logging_obj(stream=False) + in_progress: Final = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + completed: Final = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + first: Final = asyncio.create_task(logging_obj.async_success_handler(result=in_progress)) + await in_progress_parked.wait() + completion: Final = asyncio.create_task(logging_obj.async_log_background_interaction_completion(result=completed)) + await asyncio.sleep(0) + await asyncio.sleep(0) + release_in_progress.set() + await first + await completion + + assert counting_logger.logged_results == [completed, in_progress], counting_logger.logged_results diff --git a/tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 8043496f299..3e8b2842090 100644 --- a/tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -19,6 +19,7 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato _is_provider_error_chunk, anthropic_messages_response_as_sse_events, is_anthropic_content_delta_chunk, + is_anthropic_ping_chunk, parse_anthropic_error_event, ) @@ -171,6 +172,25 @@ def test_is_message_stop_chunk(): assert _is_message_stop_chunk("message_stop") is False +@pytest.mark.parametrize( + ("chunk", "expected"), + [ + (b'event: ping\ndata: {"type": "ping"}\n\n', True), + (b'event: ping\r\ndata: {"type": "ping"}\r\n\r\n', True), + (b'event: ping\ndata: {"type": "ping"}\n\nevent: ping\ndata: {"type": "ping"}\n\n', True), + ({"type": "ping"}, True), + (b'event: ping\ndata: {"ty', False), + (b'pe": "ping"}\n\n', False), + (b'pe": "message_start"}}\n\nevent: ping\ndata: {"type": "ping"}\n\n', False), + (b'event: ping\ndata: {"type": "ping"}\n\nevent: content_block_delta\ndata: {}\n\n', False), + ({"type": "message_start"}, False), + ("event: ping", False), + ], +) +def test_is_anthropic_ping_chunk_only_matches_whole_ping_frames(chunk: object, expected: bool): + assert is_anthropic_ping_chunk(chunk) is expected, chunk + + def test_is_message_stop_chunk_ignores_substring_in_payload(): """ Regression: a `content_block_delta` frame whose payload happens to contain diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index 8656a7564d2..1a6899f16ba 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -444,7 +444,7 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["thinking-binding-controls-2026-08-01"] - @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "bedrock_mantle", "vertex_ai"]) + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "bedrock", "bedrock_mantle", "vertex_ai"]) def test_dangerous_tool_use_forwarded(self, provider): """Claude Code's server-side auto-mode classifier sends `safeguards` together with dangerous-tool-use-2026-09-03. Bedrock Invoke, Bedrock Mantle, and Vertex rawPredict diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index c0266492b1e..c584f424f70 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -44,7 +44,6 @@ from litellm.router import ( _anthropic_stream_forwards_ping_live, _anthropic_stream_raised_error_status, _anthropic_stream_should_decline_fallback, - _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, ) from litellm.router_strategy import simple_shuffle @@ -4168,7 +4167,7 @@ def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"): class _InjectedFallbackRouter(Router): def __init__(self, fallback_response: object) -> None: - super().__init__(model_list=[]) + super().__init__(model_list=[], fallbacks=[{"primary": ["fallback"]}]) self._fallback_response: Final = fallback_response async def async_function_with_fallbacks_common_utils( @@ -13325,7 +13324,8 @@ def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream: return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) -def _anthropic_messages_make_router() -> Router: +def _anthropic_messages_make_router(**router_kwargs) -> Router: + router_kwargs.setdefault("fallbacks", [{"primary": ["fallback"]}]) return Router( model_list=[ { @@ -13341,7 +13341,8 @@ def _anthropic_messages_make_router() -> Router: "model": "bedrock/anthropic.claude-sonnet-4-5", }, }, - ] + ], + **router_kwargs, ) @@ -13529,24 +13530,286 @@ async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_c @pytest.mark.asyncio -async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_dropped(): - """Bugbot regression: a `ping` keepalive behind buffered lifecycle frames - carries no content and is dropped outright rather than buffered - - otherwise a slow-starting connection sending many pings could grow the - pre-content buffer without bound.""" - router = _anthropic_messages_make_router() +async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwarded_live(): + """A `ping` behind buffered lifecycle frames still reaches the client + live: it carries no lifecycle, so it cannot create overlapping + lifecycles, and it keeps the connection alive while a fallback-able + stream holds message_start back through a long thinking pass.""" + router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) + content_released = asyncio.Event() + + async def source(): + yield _anthropic_messages_message_start_chunk() + yield _anthropic_messages_ping_chunk() + await content_released.wait() + yield _anthropic_messages_content_chunk("hi") + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + + assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() + content_released.set() + assert [chunk async for chunk in wrapped] == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + ] + + +@pytest.mark.asyncio +async def test_anthropic_messages_split_ping_stays_in_order_behind_buffered_lifecycle_frame(): + """A ping the transport splits across two reads is not a whole frame, so + neither fragment may jump ahead of the buffered message_start: yielding + the head live and flushing the tail behind message_start would splice a + lifecycle frame into the middle of the ping on the wire.""" + router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) + ping_head, ping_tail = b'event: ping\ndata: {"ty', b'pe": "ping"}\n\n' source = _AnthropicMessagesFakeByteStream( - [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_ping_chunk(), - _anthropic_messages_content_chunk("hi"), - ] + [_anthropic_messages_message_start_chunk(), ping_head, ping_tail, _anthropic_messages_content_chunk("hi")] ) wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi")] + assert [chunk async for chunk in wrapped] == [ + _anthropic_messages_message_start_chunk(), + ping_head, + ping_tail, + _anthropic_messages_content_chunk("hi"), + ] + + +@pytest.mark.asyncio +async def test_anthropic_messages_no_fallback_message_start_reaches_client_before_content(): + """With no fallback able to take over, the stream is committed from the + first frame: message_start reaches the client live instead of waiting + behind the buffer for content that may be a whole thinking pass away.""" + router = _anthropic_messages_make_router(fallbacks=None) + content_released = asyncio.Event() + + async def source(): + yield _anthropic_messages_message_start_chunk() + await content_released.wait() + yield _anthropic_messages_content_chunk("hi") + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + + assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() + content_released.set() + assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] + + +@pytest.mark.asyncio +async def test_anthropic_messages_disabled_fallbacks_message_start_reaches_client_before_content(): + """A router with fallbacks configured cannot take over a request that + opted out with disable_fallbacks=True, so its lifecycle frames reach + the client live exactly like a no-fallback router's.""" + router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) + content_released = asyncio.Event() + + async def source(): + yield _anthropic_messages_message_start_chunk() + await content_released.wait() + yield _anthropic_messages_content_chunk("hi") + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary", "disable_fallbacks": True} + ) + + assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() + content_released.set() + assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] + + +@pytest.mark.asyncio +async def test_anthropic_messages_no_fallback_error_frame_reaches_client_verbatim(): + """With no fallback able to take over, a retriable provider error frame + is forwarded verbatim instead of triggering a fallback that does not + exist, and the frames already received stay in order ahead of it.""" + router = _anthropic_messages_make_router(fallbacks=None) + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecycle_frames(): + """A "*" default fallback can take over for any group, so lifecycle + frames are still held back until real content commits the primary.""" + router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}]) + content_released = asyncio.Event() + + async def source(): + yield _anthropic_messages_message_start_chunk() + await content_released.wait() + yield _anthropic_messages_content_chunk("hi") + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + + pending = asyncio.ensure_future(wrapped.__anext__()) + await asyncio.sleep(0.2) + assert not pending.done() + content_released.set() + assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() + assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] + + +def _anthropic_messages_two_order_primary_model_list() -> list: + return [ + { + "model_name": "primary", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": 1}, + }, + { + "model_name": "primary", + "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5", "order": 2}, + }, + { + "model_name": "fallback", + "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5"}, + }, + ] + + +@pytest.mark.parametrize( + "router_kwargs,request_kwargs,expected", + [ + pytest.param({"fallbacks": None}, {"model": "primary"}, False, id="no-fallbacks"), + pytest.param({"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary"}, True, id="group-fallback"), + pytest.param({"fallbacks": [{"other": ["fallback"]}]}, {"model": "primary"}, False, id="unrelated-group"), + pytest.param( + {"fallbacks": [{"*": ["fallback"]}]}, + {"model": "primary", "fallbacks": None}, + False, + id="wildcard-overridden-by-request-none", + ), + pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"), + pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"), + pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"), + pytest.param( + {"fallbacks": [{"primary": ["fallback"]}]}, + {"model": "primary", "disable_fallbacks": True}, + False, + id="disable-fallbacks", + ), + pytest.param( + {"fallbacks": None, "content_policy_fallbacks": [{"primary": ["fallback"]}]}, + {"model": "primary"}, + True, + id="content-policy-fallback", + ), + pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), + ], +) +def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): + router = _anthropic_messages_make_router(**router_kwargs) + assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected + + +@pytest.mark.parametrize( + "orders,expected", + [ + pytest.param([1, 2], True, id="distinct-orders-can-fall-back"), + pytest.param([1, 1], False, id="same-order-cannot-fall-back"), + ], +) +def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): + router = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order}, + } + for order in orders + ], + fallbacks=None, + ) + assert router._anthropic_messages_stream_can_fall_back("primary", {"model": "primary"}) is expected + + +@pytest.mark.parametrize( + "request_kwargs,expected", + [ + pytest.param({"model": "primary"}, True, id="no-target-order"), + pytest.param({"model": "primary", "_target_order": 1}, True, id="higher-order-remains"), + pytest.param({"model": "primary", "_target_order": 2}, False, id="top-order-no-order-fallback"), + pytest.param( + {"model": "primary", "_target_order": 2, "fallbacks": [{"primary": ["fallback"]}]}, + True, + id="top-order-external-fallback", + ), + ], +) +def test_anthropic_messages_stream_can_fall_back_order_target(request_kwargs, expected): + router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) + assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected + + +def test_anthropic_messages_order_levels_direct_call(): + router = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order}, + } + for order in (2, 1, None) + ], + fallbacks=None, + ) + assert router._anthropic_messages_order_levels("primary", {"model": "primary"}) == (1, 2) + + +@pytest.mark.asyncio +async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames(): + """Two order levels in one group are a real fallback target for the + dispatcher, so lifecycle frames stay buffered until content commits.""" + router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) + content_released = asyncio.Event() + + async def source(): + yield _anthropic_messages_message_start_chunk() + await content_released.wait() + yield _anthropic_messages_content_chunk("hi") + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + + pending = asyncio.ensure_future(wrapped.__anext__()) + await asyncio.sleep(0.2) + assert not pending.done() + content_released.set() + assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() + assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] + + +@pytest.mark.asyncio +async def test_anthropic_messages_request_fallbacks_none_forwards_message_start_live(): + """A per-request fallbacks=None override disables the router's wildcard + fallback, so lifecycle frames reach the client live before content.""" + router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}]) + content_released = asyncio.Event() + + async def source(): + yield _anthropic_messages_message_start_chunk() + await content_released.wait() + yield _anthropic_messages_content_chunk("hi") + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary", "fallbacks": None} + ) + + assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() + content_released.set() + assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] @pytest.mark.asyncio @@ -13949,21 +14212,12 @@ def test_merge_fallback_hidden_params_direct_call(): } -def test_anthropic_stream_should_drop_pre_content_ping_direct_call(): - ping = _anthropic_messages_ping_chunk() - content = _anthropic_messages_content_chunk("hi") - assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=False) is True - assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=True) is False - assert _anthropic_stream_should_drop_pre_content_ping(content, has_generated_content=False) is False - - def test_anthropic_stream_forwards_ping_live_direct_call(): ping = _anthropic_messages_ping_chunk() content = _anthropic_messages_content_chunk("hi") - assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=False, buffered_chunk_count=0) is True - assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=False, buffered_chunk_count=1) is False - assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=True, buffered_chunk_count=0) is False - assert _anthropic_stream_forwards_ping_live(content, has_generated_content=False, buffered_chunk_count=0) is False + assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=False) is True + assert _anthropic_stream_forwards_ping_live(ping, has_generated_content=True) is False + assert _anthropic_stream_forwards_ping_live(content, has_generated_content=False) is False def test_anthropic_stream_error_is_gateway_verdict_direct_call(): diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cb6a16e0c4c..bbcbda6a121 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -6484,3 +6484,67 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ ) assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} + + +@pytest.mark.asyncio +async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: + """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue + #44500: both exits enqueued a success handler, and on a large prompt the first one yielded to + the worker-thread base64 offload before it marked ``has_logged_async_success``, so the second + passed the check too and the request was logged and billed twice. The schedule step claims the + log for the object synchronously, so only the inner provider-shaped result is ever logged.""" + from litellm.litellm_core_utils import litellm_logging + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.utils import _dispatch_success_logging + + class CountingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_results: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + if kwargs["litellm_call_id"] == "bridge-call-id": + self.logged_results.append(response_obj) + + counting_logger: Final = CountingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger]) + real_truncate: Final = litellm_logging.truncate_base64_in_messages_async + + async def yielding_truncate(messages): + await asyncio.sleep(0) + return await real_truncate(messages) + + monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", yielding_truncate) + + messages: Final = [{"role": "user", "content": "hello"}] + logging_obj: Final = Logging( + model="gpt-5.6-luna", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp(), + litellm_call_id="bridge-call-id", + function_id="bridge-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="gpt-5.6-luna", custom_llm_provider="openai", input=messages + ) + inner_result: Final = ModelResponse(id="inner") + outer_result: Final = ModelResponse(id="outer") + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + for result in (inner_result, outer_result): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=now, + end_time=now, + is_completion_with_fallbacks=False, + is_litellm_internal_call=False, + ) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert len(counting_logger.logged_results) == 1, counting_logger.logged_results + assert counting_logger.logged_results[0] is inner_result diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index d53db68bb9f..6d0e0642ba2 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -58,6 +58,7 @@ export interface TopApiKeyData { api_key: string; key_alias: string | null; team_id: string | null; + user: string | null; spend: number; requests: number; tokens: number; diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index 1bd7655b5e3..569f69adb56 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -258,10 +258,20 @@ describe("ActivityMetrics", () => { api_key: "key-123", key_alias: "Test Key", team_id: "team1", + user: "owner@example.com", spend: 50.25, requests: 25, tokens: 12500, }, + { + api_key: "key-456", + key_alias: "Owner Alias", + team_id: null, + user: "Owner Alias", + spend: 40.25, + requests: 20, + tokens: 10000, + }, ], }, }; @@ -269,6 +279,9 @@ describe("ActivityMetrics", () => { render(); expect(screen.getByText("Top Virtual Keys by Spend")).toBeInTheDocument(); expect(screen.getByText("Test Key")).toBeInTheDocument(); + expect(screen.getByText("User: owner@example.com")).toBeInTheDocument(); + expect(screen.getByText("Owner Alias")).toBeInTheDocument(); + expect(screen.queryByText("User: Owner Alias")).not.toBeInTheDocument(); }); it("should display API key hash when alias is missing", () => { @@ -280,6 +293,7 @@ describe("ActivityMetrics", () => { api_key: "key-1234567890", key_alias: null, team_id: null, + user: null, spend: 50.25, requests: 25, tokens: 12500, @@ -301,6 +315,7 @@ describe("ActivityMetrics", () => { api_key: "key-123", key_alias: "Test Key", team_id: "team1", + user: null, spend: 50.25, requests: 25, tokens: 12500, @@ -1056,6 +1071,8 @@ describe("processActivityData", () => { metadata: { key_alias: "test-key-1", team_id: "team1", + user_id: "owner-id-1", + user_email: "owner-1@example.com", }, }, "key-2": { @@ -1073,6 +1090,7 @@ describe("processActivityData", () => { metadata: { key_alias: "test-key-2", team_id: "team2", + user_id: "owner-id-2", }, }, }, @@ -1094,6 +1112,10 @@ describe("processActivityData", () => { expect(result["gpt-4"].top_api_keys[0].spend).toBe(60.0); expect(result["gpt-4"].top_api_keys[0].api_key).toBe("key-1"); expect(result["gpt-4"].top_api_keys[1].spend).toBe(40.5); + expect(result["gpt-4"].top_api_keys.map(({ api_key, user }) => [api_key, user])).toEqual([ + ["key-1", "owner-1@example.com"], + ["key-2", "owner-id-2"], + ]); }); it("should limit top_api_keys to 5 entries", () => { diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index 95315f8ec27..5ddb89703ea 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -102,20 +102,26 @@ const ModelSection = ({

Top Virtual Keys by Spend

- {metrics.top_api_keys.map((keyData) => ( -
-
-

{keyData.key_alias || `${keyData.api_key.substring(0, 10)}...`}

- {keyData.team_id &&

Team: {keyData.team_id}

} + {metrics.top_api_keys.map((keyData) => { + const keyLabel = keyData.key_alias || `${keyData.api_key.substring(0, 10)}...`; + return ( +
+
+

{keyLabel}

+ {keyData.team_id &&

Team: {keyData.team_id}

} + {keyData.user && keyData.user !== keyLabel && ( +

User: {keyData.user}

+ )} +
+
+

${formatNumberWithCommas(keyData.spend, 2)}

+

+ {keyData.requests.toLocaleString()} requests | {keyData.tokens.toLocaleString()} tokens +

+
-
-

${formatNumberWithCommas(keyData.spend, 2)}

-

- {keyData.requests.toLocaleString()} requests | {keyData.tokens.toLocaleString()} tokens -

-
-
- ))} + ); + })}
@@ -585,6 +591,7 @@ export const processActivityData = ( api_key: apiKey, key_alias: keyActivityLabel(keyData.metadata, "") || null, team_id: keyData.metadata.team_id, + user: keyData.metadata.user_email ?? keyData.metadata.user_id ?? null, spend: 0, requests: 0, tokens: 0,