diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 37b4d2d3937..2ff523571f2 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -1,7 +1,8 @@ import asyncio import os import time -from collections.abc import Callable +from collections.abc import Callable, Mapping +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from functools import cache from typing import Final @@ -15,6 +16,7 @@ from litellm.constants import ( AZURE_STORAGE_MSFT_VERSION, ) from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id from litellm.llms.custom_httpx.http_handler import ( @@ -47,6 +49,20 @@ def adls_safe_file_name(payload_id: str | None) -> str: return f"{name.translate(_ADLS_SAFE_NAME)}.json" +@dataclass(frozen=True, slots=True) +class AzureStorageLogQueueItem: + payload: StandardLoggingPayload + file_name: str + + +def _queued_log(kwargs: Mapping[str, object], payload: StandardLoggingPayload) -> AzureStorageLogQueueItem: + payload_id: Final = payload.get("id") + return AzureStorageLogQueueItem( + payload=payload, + file_name=adls_safe_file_name(payload_id and logged_event_key(kwargs, payload_id)), + ) + + @cache def _cached_credential_chain_token_provider() -> Callable[[], str]: return get_azure_ad_token_provider( @@ -99,7 +115,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: list[StandardLoggingPayload] = [] + self.log_queue: list[AzureStorageLogQueueItem] = [] super().__init__(**kwargs, flush_lock=self.flush_lock) except Exception as e: verbose_logger.exception( @@ -115,7 +131,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): def azure_storage_blob_endpoint(self) -> str: return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}" - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime): """ Async Log success events to Azure Blob Storage @@ -133,12 +149,12 @@ class AzureBlobStorageLogger(CustomBatchLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") - self.log_queue.append(standard_logging_payload) + self.log_queue.append(_queued_log(kwargs, standard_logging_payload)) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime): """ Async Log failure events to Azure Blob Storage @@ -156,7 +172,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") - self.log_queue.append(standard_logging_payload) + self.log_queue.append(_queued_log(kwargs, standard_logging_payload)) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) @@ -195,13 +211,15 @@ class AzureBlobStorageLogger(CustomBatchLogger): len(self.log_queue), ) - for payload in self.log_queue: - await self.async_upload_payload_to_azure_blob_storage(payload=payload) + for item in self.log_queue: + await self.async_upload_payload_to_azure_blob_storage(payload=item.payload, file_name=item.file_name) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Error sending batch API - %s", e) - async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload): + async def async_upload_payload_to_azure_blob_storage( + self, payload: StandardLoggingPayload, file_name: str | None = None + ): """ Uploads the payload to Azure Blob Storage using a 3-step process: 1. Create file resource @@ -210,14 +228,14 @@ class AzureBlobStorageLogger(CustomBatchLogger): """ try: if self.azure_storage_account_key: - await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload) + await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload, file_name=file_name) else: # Get a valid token instead of always requesting a new one await self.set_valid_azure_ad_token() async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry payload_bytes: Final = json_payload.encode("utf-8") - filename: Final = adls_safe_file_name(payload.get("id")) + filename: Final = file_name or adls_safe_file_name(payload.get("id")) base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}" # Execute the 3-step upload process @@ -377,7 +395,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client_timeout = self._clock() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS return self._service_client - async def upload_to_azure_data_lake_with_azure_account_key(self, payload: StandardLoggingPayload): + async def upload_to_azure_data_lake_with_azure_account_key( + self, payload: StandardLoggingPayload, file_name: str | None = None + ): """ Uploads the payload to Azure Data Lake using the Azure SDK @@ -386,11 +406,11 @@ class AzureBlobStorageLogger(CustomBatchLogger): try: today: Final = datetime.now().strftime("%Y-%m-%d") - file_name: Final = adls_safe_file_name(payload.get("id")) + resolved_file_name: Final = file_name or adls_safe_file_name(payload.get("id")) await self._upload_bytes_with_account_key( - file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8") + file_path=f"{today}/{resolved_file_name}", content=safe_dumps(payload).encode("utf-8") ) - verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name) + verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, resolved_file_name) except Exception as e: verbose_logger.exception("Error occurred: %s", e) diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 1d01965cd2c..c8453456052 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -19,6 +19,7 @@ from litellm._logging import verbose_logger from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.tracing.types import SpendLogPayload, SpendLogRecord @@ -134,7 +135,7 @@ def spend_log_row_from_payload( usage: Mapping[str, Any] = metadata.get("usage_object") or hidden_params.get("usage_object") or MappingProxyType({}) cache_read_tokens, cache_write_tokens = _cache_tokens(usage) trace_id, span_id = _find_traceparent(metadata, kwargs) - request_id = str(payload.get("id") or "") + response_id: Final = str(payload.get("id") or "") redact = litellm.turn_off_message_logging is True track_end_user: Final = not litellm.disable_end_user_cost_tracking completion_start_ms = _to_ms(payload.get("completionStartTime")) @@ -157,8 +158,8 @@ def spend_log_row_from_payload( ) ) return SpendLogRecord( - request_id=request_id, - response_id=strip_cache_hit_suffix(request_id), + request_id=logged_event_key(kwargs, response_id), + response_id=strip_cache_hit_suffix(response_id), provider_request_id=( get_provider_request_id(kwargs.get("response_headers")) or get_provider_request_id(hidden_params.get("additional_headers")) diff --git a/litellm/integrations/deepeval/deepeval.py b/litellm/integrations/deepeval/deepeval.py index 71d149e4014..e38b987d46f 100644 --- a/litellm/integrations/deepeval/deepeval.py +++ b/litellm/integrations/deepeval/deepeval.py @@ -15,6 +15,7 @@ from litellm.integrations.deepeval.utils import ( to_zod_compatible_iso, validate_environment, ) +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key # This file includes the custom callbacks for LiteLLM Proxy @@ -109,8 +110,10 @@ class DeepEvalLogger(CustomLogger): ) else: output = str(standard_logging_object.get("error_string", "")) + logged_id: Final = standard_logging_object.get("id", uuid.uuid4()) + span_id: Final = logged_event_key(kwargs, logged_id) if isinstance(logged_id, str) else logged_id return BaseApiSpan( - uuid=standard_logging_object.get("id", uuid.uuid4()), + uuid=span_id, name=("litellm_success_callback" if is_success else "litellm_failure_callback"), status=(TraceSpanApiStatus.SUCCESS if is_success else TraceSpanApiStatus.ERRORED), type=SpanApiType.LLM, diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index 6569e70329b..8672eb00ce5 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -9,6 +9,8 @@ from typing import Final, Protocol import litellm from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key class _DynamoTable(Protocol): @@ -22,11 +24,13 @@ class _DynamoResource(Protocol): class DyanmoDBLogger: # Class variables or attributes - def __init__(self): + def __init__(self, dynamodb: _DynamoResource | None = None): # Instance variables import boto3 - self.dynamodb: Final[_DynamoResource] = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"]) + self.dynamodb: Final[_DynamoResource] = dynamodb or boto3.resource( + "dynamodb", region_name=os.environ["AWS_REGION_NAME"] + ) if litellm.dynamodb_table_name is None: raise ValueError( "LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=`" @@ -45,23 +49,23 @@ class DyanmoDBLogger: _async_log_event = async_log_event - def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): + def log_event(self, kwargs: Mapping[str, object], response_obj, start_time, end_time, print_verbose): try: print_verbose(f"DynamoDB Logging - Enters logging function for model {kwargs}") # construct payload to send to DynamoDB # follows the same params as langfuse.py - litellm_params: Final = kwargs.get("litellm_params", {}) - metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} messages: Final = kwargs.get("messages") optional_params: Final = kwargs.get("optional_params", {}) call_type: Final = kwargs.get("call_type", "litellm.completion") usage: Final = response_obj["usage"] - id: Final = response_obj.get("id", str(uuid.uuid4())) + response_id: Final = response_obj.get("id", str(uuid.uuid4())) # Build the initial payload payload: Final[dict[str, object]] = { - "id": id, + "id": logged_event_key(kwargs, response_id) if isinstance(response_id, str) else response_id, "call_type": call_type, "startTime": start_time, "endTime": end_time, diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 99cd8b31b63..7057e719764 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload @@ -296,30 +297,28 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): else: await self._send_individual_logs(items_to_process) - def _get_object_name(self, kwargs: dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str: + def _get_object_name( + self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any + ) -> str: """ Get the object name to use for the current payload """ current_date: Final = self._get_object_date_from_datetime(datetime.now(timezone.utc)) + metadata: Final = as_str_mapping((as_str_mapping(kwargs.get("litellm_params")) or {}).get("metadata")) or {} + gcs_log_id: Final = metadata.get("gcs_log_id") + safe_log_id: Final = sanitize_cloud_object_component( + gcs_log_id if isinstance(gcs_log_id, str) else None, fallback="" + ) + if safe_log_id: + return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" if logging_payload.get("error_str", None) is not None: - object_name = self._generate_failure_object_name( - request_date_str=current_date, - ) - else: - object_name = self._generate_success_object_name( - request_date_str=current_date, - response_id=response_obj.get("id", ""), - ) - - # used for testing - _litellm_params: Final = kwargs.get("litellm_params", None) or {} - _metadata: Final = _litellm_params.get("metadata", None) or {} - if "gcs_log_id" in _metadata: - safe_log_id: Final = sanitize_cloud_object_component(_metadata.get("gcs_log_id"), fallback="") - if safe_log_id: - object_name = f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" - - return object_name + return self._generate_failure_object_name(request_date_str=current_date) + raw_response_id: Final[object] = response_obj.get("id", "") + response_id: Final = raw_response_id if isinstance(raw_response_id, str) else "" + return self._generate_success_object_name( + request_date_str=current_date, + response_id=logged_event_key(kwargs, response_id), + ) async def get_request_response_payload( self, diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index db2fe386dec..918e849a086 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -3,12 +3,15 @@ import json import os +from collections.abc import Mapping from typing import Final import httpx import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -48,7 +51,7 @@ class OpenMeterLogger(CustomLogger): if len(missing_keys) > 0: raise Exception(f"Missing keys={missing_keys} in environment.") - def _common_logic(self, kwargs: dict, response_obj): + def _common_logic(self, kwargs: Mapping[str, object], response_obj): call_id: Final = response_obj.get("id", kwargs.get("litellm_call_id")) dt: Final = get_utc_datetime().isoformat() cost: Final = kwargs.get("response_cost", None) @@ -74,8 +77,8 @@ class OpenMeterLogger(CustomLogger): # If no user provided directly, try to get it from token user_id if user_param is None: # Check if user_id is available from the API key metadata - litellm_params: Final = kwargs.get("litellm_params", {}) - metadata: Final = litellm_params.get("metadata", {}) + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} user_api_key_user_id: Final = metadata.get("user_api_key_user_id", None) if user_api_key_user_id is not None: @@ -89,7 +92,7 @@ class OpenMeterLogger(CustomLogger): return { "specversion": "1.0", "type": os.getenv("OPENMETER_EVENT_TYPE", "litellm_tokens"), - "id": call_id, + "id": logged_event_key(kwargs, call_id) if isinstance(call_id, str) else call_id, "time": dt, "subject": subject, "source": "litellm-proxy", diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 453d4358b35..138a01c6524 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -26,7 +26,12 @@ from litellm.constants import ( MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, NON_INFERENCE_CALL_TYPES, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, @@ -35,6 +40,8 @@ from litellm.types.utils import ( BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" + ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" @@ -151,6 +158,26 @@ def is_unbilled_non_inference_call_from_params( return is_unbilled_non_inference_call(call_type, metadata, response) +def settles_background_response(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> bool: + return kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES or ( + call_type in NON_INFERENCE_CALL_TYPES + and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + ) + + +def logged_event_key(kwargs: Mapping[str, object], event_id: str) -> str: + call_type: Final = kwargs.get("call_type") + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = ( + as_str_mapping(litellm_params.get("litellm_metadata")) or as_str_mapping(litellm_params.get("metadata")) or {} + ) + if not event_id or not settles_background_response( + call_type if isinstance(call_type, str) else "", metadata, kwargs + ): + return event_id + return f"{event_id}{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}" + + def sanitize_user_api_key_auth(auth: object) -> object: """Copy of the auth object with its budget reservation removed; the cost callback falls back to reading the reservation from inside the auth object.""" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5f909f639c4..37657a94f16 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -872,6 +872,7 @@ class Logging(LiteLLMLoggingBaseClass): # 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.client_facing_response_id: str | None = None self.zero_cost_warned: bool = False self.llm_caching_handler: LLMCachingHandler | None = None @@ -3238,6 +3239,23 @@ class Logging(LiteLLMLoggingBaseClass): return response return response.model_copy(update={"id": response_id}) + def _with_client_facing_response_id( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + response_id: Final = self.client_facing_response_id + if response_id is None or getattr(response, "id", None) == response_id: + return response + return response.model_copy(update={"id": response_id}) + + def _client_facing_streaming_response( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + return self._with_client_facing_response_id( + self._with_continued_response_id(self._with_client_facing_stream_model(response)) + ) + def _success_handler_helper_fn( self, result: object = None, @@ -3641,9 +3659,7 @@ class Logging(LiteLLMLoggingBaseClass): ) self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) - logged_streaming_response: Final = self._with_continued_response_id( - self._with_client_facing_stream_model(complete_streaming_response) - ) + logged_streaming_response: Final = self._client_facing_streaming_response(complete_streaming_response) self.model_call_details["complete_streaming_response"] = logged_streaming_response self.capture_usage_before_logging(logged_streaming_response, start_time, end_time) ## STANDARDIZED LOGGING PAYLOAD @@ -4240,9 +4256,7 @@ class Logging(LiteLLMLoggingBaseClass): await self._prepare_baseline_cache_estimate(complete_streaming_response) - logged_streaming_response: Final = self._with_continued_response_id( - self._with_client_facing_stream_model(complete_streaming_response) - ) + logged_streaming_response: Final = self._client_facing_streaming_response(complete_streaming_response) self.model_call_details["async_complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 771d68a0ea0..61ab1eb4771 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -7,22 +7,25 @@ instead of writing immediately on each request. from collections.abc import AsyncGenerator, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import LitellmUserRoles from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, ResponsesAPIResponse, ) +from litellm.types.responses.main import DeleteResponseResult from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LLMResponseTypes, SpecialEnums if TYPE_CHECKING: @@ -67,6 +70,8 @@ def _proxy_signing_key() -> str | None: _RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_ResponseObjectT = TypeVar("_ResponseObjectT", bound=BaseLiteLLMOpenAIResponseObject) +_AddressedResultT = TypeVar("_AddressedResultT", bound=ResponsesAPIResponse | DeleteResponseResult) def _response_payload(response_obj: object) -> Mapping[str, object] | None: @@ -77,9 +82,9 @@ def _response_payload(response_obj: object) -> Mapping[str, object] | None: def _rewrite_advertised_id( - event: BaseLiteLLMOpenAIResponseObject, + event: _ResponseObjectT, rewrite: Callable[[str], str], -) -> BaseLiteLLMOpenAIResponseObject: +) -> _ResponseObjectT: event_id: Final = getattr(event, "id", None) if isinstance(event_id, str) and event_id.startswith("resp_"): setattr(event, "id", rewrite(event_id)) @@ -116,6 +121,35 @@ def _streams_response_ids(request_route: str | None) -> bool: return canonical in _RESPONSES_API_ROUTES or (parent in _RESPONSES_API_ROUTES and response_id != "") +_NO_HELD_IDS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _upstream_response_id(response_id: str) -> str: + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id) + + +def _names_the_same_response(response_id: str | None, original_response_id: str) -> bool: + return response_id is not None and _upstream_response_id(response_id) == _upstream_response_id(original_response_id) + + +def _held_id_naming(response_id: str | None, held_ids: Mapping[str, str]) -> str | None: + return next( + ( + held_id + for original_response_id, held_id in held_ids.items() + if _names_the_same_response(response_id, original_response_id) + ), + None, + ) + + +def _id_sent_back(response: ResponsesAPIResponse | DeleteResponseResult, held_ids: Mapping[str, str]) -> str | None: + held_id: Final = _held_id_naming(response.id, held_ids) + if held_id is not None: + return held_id + return response.id if isinstance(response, DeleteResponseResult) else None + + class ResponsesIDSecurity(CustomLogger): def __init__( self, @@ -275,10 +309,18 @@ class ResponsesIDSecurity(CustomLogger): def _encrypt_response_id( self, - response: BaseLiteLLMOpenAIResponseObject, + response: _ResponseObjectT, user_api_key_dict: "UserAPIKeyAuth", request_cache: dict[str, str] | None = None, - ) -> BaseLiteLLMOpenAIResponseObject: + ) -> _ResponseObjectT: + encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache) + return response if encrypt is None else _rewrite_advertised_id(response, encrypt) + + def _response_id_encrypter( + self, + user_api_key_dict: "UserAPIKeyAuth", + request_cache: dict[str, str] | None, + ) -> Callable[[str], str] | None: signing_key: Final = self._get_signing_key() if signing_key is None: verbose_proxy_logger.debug( @@ -287,7 +329,7 @@ class ResponsesIDSecurity(CustomLogger): "Skipping response ID encryption. " "See: https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" ) - return response + return None def encrypt(original_id: str) -> str: cached: Final = request_cache.get(original_id) if request_cache is not None else None @@ -304,47 +346,73 @@ class ResponsesIDSecurity(CustomLogger): request_cache[original_id] = encrypted_id return encrypted_id - return _rewrite_advertised_id(response, encrypt) + return encrypt + + def _streamed_response_id_advertiser( + self, + user_api_key_dict: "UserAPIKeyAuth", + request_data: Mapping[str, object], + ) -> Callable[[str], str] | None: + if not _streams_response_ids(user_api_key_dict.request_route): + return None + if self._general_settings_reader().get("disable_responses_id_security", False): + return None + encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache={}) + if encrypt is None: + return None + logging_obj: Final = request_data.get("litellm_logging_obj") + held_ids: Final = self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD)) + + def advertise(original_id: str) -> str: + logged_id: Final = encrypt(original_id) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.client_facing_response_id = logged_id + return _held_id_naming(original_id, held_ids) or logged_id + + return advertise + + def _ids_the_caller_holds(self, addressed_id: object) -> Mapping[str, str]: + if not isinstance(addressed_id, str) or not self._is_encrypted_response_id(addressed_id): + return _NO_HELD_IDS + original_response_id, _, _ = self._decrypt_response_id(addressed_id) + return MappingProxyType({original_response_id: addressed_id}) + + def _advertised_response( + self, + response: _AddressedResultT, + user_api_key_dict: "UserAPIKeyAuth", + request_data: Mapping[str, object], + ) -> _AddressedResultT: + sent_back_id: Final = _id_sent_back( + response, self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD)) + ) + logged_response: Final = self._encrypt_response_id(response, user_api_key_dict, request_cache=None) + return logged_response if sent_back_id is None else logged_response.model_copy(update={"id": sent_back_id}) async def async_post_call_success_hook( self, - data: dict, + data: dict[str, object], user_api_key_dict: "UserAPIKeyAuth", response: LLMResponseTypes, ) -> LLMResponseTypes: - """ - Queue response IDs for batch processing instead of writing directly to DB. - - This method adds response IDs to an in-memory queue, which are then - batch-processed by the DBSpendUpdateWriter during regular database update cycles. - """ general_settings: Final = self._general_settings_reader() if general_settings.get("disable_responses_id_security", False): return response - if isinstance(response, ResponsesAPIResponse): - response = cast( - ResponsesAPIResponse, - self._encrypt_response_id(response, user_api_key_dict, request_cache=None), - ) + if isinstance(response, ResponsesAPIResponse | DeleteResponseResult): + return self._advertised_response(response, user_api_key_dict, data) return response async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict + self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict[str, object] ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: - general_settings: Final = self._general_settings_reader() - - # Create a request-scoped cache for consistent encryption across streaming chunks. - request_encryption_cache: Final[dict[str, str]] = {} - + advertise: Final = self._streamed_response_id_advertiser(user_api_key_dict, request_data) async for chunk in response: - if ( - isinstance(chunk, BaseLiteLLMOpenAIResponseObject) - and _streams_response_ids(user_api_key_dict.request_route) - and not general_settings.get("disable_responses_id_security", False) - ): - chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) - yield chunk + yield ( + _rewrite_advertised_id(chunk, advertise) + if advertise is not None and isinstance(chunk, BaseLiteLLMOpenAIResponseObject) + else chunk + ) def decrypted_response_id(response_id: str) -> str: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ba5f873c4e9..8b0ff718fec 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -20,7 +20,6 @@ from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, - INTERNAL_CALL_ORIGIN_METADATA_KEY, KUBERNETES_POD_ROUTING_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, @@ -30,7 +29,6 @@ from litellm.constants import ( MAX_SPEND_LOG_MODEL_NAME_LENGTH, MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, - NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL, @@ -53,7 +51,11 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider from litellm.litellm_core_utils.hidden_params import served_from_cache -from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.internal_call_metadata import ( + BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX, + is_unbilled_non_inference_call, + settles_background_response, +) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, UsageObservation, @@ -79,13 +81,8 @@ from litellm.responses.litellm_completion_transformation.compaction_items import summary_only_compaction_content, with_compaction_summary, ) -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - SETTLES_BACKGROUND_RESPONSE_KWARG, - BackgroundResponseSettlement, -) from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, SilentExperimentFlag from litellm.types.utils import ( - BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, PROMPT_CARRYING_GUARDRAIL_FIELDS, AzureSpillover, CallTypes, @@ -347,24 +344,21 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" -BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" -def _cost_row_request_id_suffix( - call_type: str, metadata: Mapping[str, object], settles_background_response: bool -) -> str: +def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX - if settles_background_response or ( - call_type in NON_INFERENCE_CALL_TYPES - and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - ): + if settles_background_response(call_type, metadata, kwargs): return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX return "" def get_spend_logs_id( - call_type: str, response_obj: dict, kwargs: Mapping[str, object], metadata: Mapping[str, object] = EMPTY_MAPPING + call_type: str, + response_obj: Mapping[str, object], + kwargs: Mapping[str, object], + metadata: Mapping[str, object] = EMPTY_MAPPING, ) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( @@ -377,10 +371,7 @@ def get_spend_logs_id( ) if resolved_id is None: return None - settles_background_response: Final = ( - kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES - ) - return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, settles_background_response)}" + return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" _MISSING_ATTRIBUTE: Final = object() diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 7cf163e7ce8..7e6ee84e2b2 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,14 +1,16 @@ import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from itertools import chain +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.responses.litellm_completion_transformation.client_tool_search import ( bridge_client_tool_search, client_tool_search_loaded_tools, @@ -69,6 +71,90 @@ def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool: return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES and not _is_a_failed_turn(spend_log) +_RESPONSE_READ_CALL_TYPES: Final = frozenset(("aget_responses", "acancel_responses")) + + +def _output_item_id(item: object) -> str | None: + match item: + case {"id": str() as item_id}: + return item_id + case _: + return None + + +def _is_a_response_read(spend_log: "SpendLogsPayload") -> bool: + return "call_type" in spend_log and spend_log["call_type"] in _RESPONSE_READ_CALL_TYPES + + +def _stored_response(spend_log: "SpendLogsPayload") -> object: + return spend_log.get("response") + + +def _output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str | None, object]]: + output: Final = _stored_responses_api_output(_stored_response(spend_log)) or [] + return zip(map(_output_item_id, output), output) + + +def _identified_output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str, object]]: + return ((item_id, item) for item_id, item in _output_items(spend_log) if item_id is not None) + + +def _output_item_ids(spend_log: "SpendLogsPayload") -> frozenset[str]: + return frozenset(item_id for item_id, _ in _identified_output_items(spend_log)) + + +def _first_index_holding_each_item(indexed_turns: Iterable[tuple[int, "SpendLogsPayload"]]) -> Mapping[str, int]: + return dict( + chain.from_iterable( + ((item_id, index) for item_id in _output_item_ids(turn)) for index, turn in reversed(tuple(indexed_turns)) + ) + ) + + +def _turn_replaying_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, int]: + indexed_turns: Final = tuple(enumerate(turns)) + first_read_holding: Final = _first_index_holding_each_item( + (index, turn) for index, turn in indexed_turns if _is_a_response_read(turn) + ) + first_other_turn_holding: Final = _first_index_holding_each_item( + (index, turn) for index, turn in indexed_turns if not _is_a_response_read(turn) + ) + return MappingProxyType({**first_read_holding, **first_other_turn_holding}) + + +def _latest_read_copy_of_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, object]: + return MappingProxyType( + dict(chain.from_iterable(_identified_output_items(turn) for turn in turns if _is_a_response_read(turn))) + ) + + +def _with_only_new_output( + index: int, turn: "SpendLogsPayload", turn_replaying: Mapping[str, int], latest_copy: Mapping[str, object] +) -> "SpendLogsPayload | None": + if not _is_a_response_read(turn): + return turn + stored_response: Final = as_str_mapping(_stored_response(turn)) + new_output: Final = [ + item if item_id is None else latest_copy[item_id] + for item_id, item in _output_items(turn) + if item_id is None or turn_replaying.get(item_id) == index + ] + if stored_response is None or not new_output: + return None + replayed_turn: Final[SpendLogsPayload] = {**turn, "response": {**stored_response, "output": new_output}} + return replayed_turn + + +def _conversation_turns(spend_logs: Sequence["SpendLogsPayload"]) -> tuple["SpendLogsPayload", ...]: + turns: Final = tuple(filter(_is_a_conversation_turn, spend_logs)) + turn_replaying: Final = _turn_replaying_each_item(turns) + latest_copy: Final = _latest_read_copy_of_each_item(turns) + replayed: Final = ( + _with_only_new_output(index, turn, turn_replaying, latest_copy) for index, turn in enumerate(turns) + ) + return tuple(turn for turn in replayed if turn is not None) + + _MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results")) @@ -101,7 +187,7 @@ class StoredTurns: def _compaction_point( index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None ) -> _ReplayStart | None: - if response_holds_compaction_block(spend_log.get("response")): + if response_holds_compaction_block(_stored_response(spend_log)): return _ReplayStart(index=index, at_output=True) if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")): return _ReplayStart(index=index, at_output=False) @@ -271,12 +357,10 @@ class ResponsesSessionHandler: previous_response_id: str, client_response_ids: tuple[str, ...] = (), ) -> StoredTurns: - spend_logs: Final = tuple( - spend_log - for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + spend_logs: Final = _conversation_turns( + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( previous_response_id, client_response_ids ) - if _is_a_conversation_turn(spend_log) ) return StoredTurns( spend_logs=spend_logs, @@ -378,7 +462,7 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final[object] = spend_log.get("response", "{}") + _response_output: Final = _stored_response(spend_log) if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( diff --git a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py index 4a08259c028..0e954af0952 100644 --- a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py +++ b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py @@ -69,9 +69,6 @@ def _retrieve_response(proxy: ProxyClient, key: str, response_id: str) -> Respon class TestResponsesRetrieve: - @pytest.mark.skip( - reason="stage red: product gap (LIT-5446), retrieve returns a different id than the stored response (non-idempotent response-id re-encryption)" - ) @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") @meta( Subject( diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths.py b/tests/integration/compatibility/test_openai_sdk_client_paths.py index f8d7449038a..d6b6d1d460d 100644 --- a/tests/integration/compatibility/test_openai_sdk_client_paths.py +++ b/tests/integration/compatibility/test_openai_sdk_client_paths.py @@ -630,19 +630,6 @@ def _spend_row(identity: str) -> Mapping[str, JsonValue]: return rows[0] -def _spend_row_keyed_by_any(identities: frozenset[str]) -> Mapping[str, JsonValue]: - keys: Final = tuple(sorted(identities)) - placeholders: Final = ", ".join("%s" for _ in keys) - rows: Final = eventually( - lambda: read_rows( - f'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', keys - ), - lambda rows: len(rows) == 1, - seconds=70, - ) - return rows[0] - - def _spend_request_ids(identities: tuple[str, str, str]) -> frozenset[str]: rows: Final = eventually( lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', identities), @@ -815,7 +802,7 @@ def test_r07_responses_stream_through_the_async_openai_sdk(gateway: Gateway) -> served: Final = _served(wire) assert _routes(served) == (("POST", "/v1/responses"),), served assert _body(served[0])["stream"] is True, served[0].body - assert _spend_row_keyed_by_any(response_identities(last.response.id))["model"] == DEPLOYMENT_MODEL + assert _spend_row(last.response.id)["model"] == DEPLOYMENT_MODEL def test_r08_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: diff --git a/tests/integration/observability/test_langfuse_otel_metadata_chaos.py b/tests/integration/observability/test_langfuse_otel_metadata_chaos.py index bb386f6d41e..15bc076b8ba 100644 --- a/tests/integration/observability/test_langfuse_otel_metadata_chaos.py +++ b/tests/integration/observability/test_langfuse_otel_metadata_chaos.py @@ -107,9 +107,11 @@ def test_langfuse_otel_sink_outage_mid_burst_lands_every_response_id_once(gatewa ) payloads: Final = tuple(json.loads(line) for line in events) found: Final = next( - payload.get("response_id") or (payload.get("message") or {}).get("id") or payload.get("id") + (payload.get("response") or {}).get("id") + or (payload.get("message") or {}).get("id") + or payload.get("id") for payload in payloads - if payload.get("response_id") or payload.get("message") or payload.get("id") + if payload.get("response") or payload.get("message") or payload.get("id") ) return str(found) response: Final = rig.candidate.client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) diff --git a/tests/integration/observability/test_langfuse_otel_metadata_clients.py b/tests/integration/observability/test_langfuse_otel_metadata_clients.py index c2e0573b48a..fd2b9758449 100644 --- a/tests/integration/observability/test_langfuse_otel_metadata_clients.py +++ b/tests/integration/observability/test_langfuse_otel_metadata_clients.py @@ -162,8 +162,13 @@ def test_langfuse_otel_metadata_reaches_langfuse_for_responses_sse_stream(gatewa for line in stream.iter_lines() if line.startswith("data: ") and line != "data: [DONE]" ) - response_id: Final = next( - event["response_id"] for event in events if event.get("type") == "response.output_text.delta" + advertised: Final = frozenset( + event["response"]["id"] + for event in events + if event.get("type") in ("response.created", "response.completed") ) + assert len(advertised) == 1, events + response_id: Final = next(iter(advertised)) attrs: Final = exported_llm_span(rig.collector, response_id) + assert attrs["gen_ai.response.id"] == response_id, attrs assert_identity(attrs, chat_identity_fields(identity)) diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index d7f4de2961b..cd095991e53 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -1,5 +1,4 @@ import asyncio -import base64 import json import re import signal @@ -118,17 +117,6 @@ def _responses_reply(identity: str, stream: bool) -> Reply: ) -def _decoded_responses_id(identity: str) -> str: - try: - return base64.b64decode(identity.removeprefix("resp_").encode()).decode() - except (ValueError, UnicodeDecodeError): - return identity - - -def _canonical_id(identity: str) -> str: - return _decoded_responses_id(identity).rpartition("response_id:")[2] - - def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]" @@ -170,8 +158,7 @@ class Collector: return tuple( attributes for attributes in self.attributes() - if isinstance(logged := attributes.get("gen_ai.response.id", {}).get("stringValue"), str) - and _canonical_id(logged) == _canonical_id(response_id) + if attributes.get("gen_ai.response.id", {}).get("stringValue") == response_id ) def conversation_ids(self, response_id: str) -> tuple[str | None, ...]: @@ -183,10 +170,6 @@ class Collector: def single_span(self, response_id: str) -> str | None: return eventually(lambda: self.conversation_ids(response_id), lambda values: len(values) == 1, seconds=30)[0] - def logged_id(self, response_id: str) -> str: - spans: Final = eventually(lambda: self.spans(response_id), lambda values: len(values) == 1, seconds=30) - return str(spans[0]["gen_ai.response.id"]["stringValue"]) - @pytest.fixture(scope="session") def collector() -> Iterator[Collector]: @@ -437,12 +420,13 @@ def test_responses_stream_raw_metadata_session_id_lands_as_conversation_id(rig: completed: Final = tuple(event for event in events if event["type"] == "response.completed") assert len(completed) == 1, events assert completed[0]["response"]["output"][0]["id"] == f"msg_resp_{marker}", completed - assert str(completed[0]["response"]["id"]).startswith("resp_"), completed + client_id: Final = str(completed[0]["response"]["id"]) + assert client_id.startswith("resp_"), completed assert rig.upstream_bodies(marker) == ( {"model": "gpt-4o-mini", "input": marker, "metadata": {"session_id": session}, "stream": True}, ) - assert rig.sink.single_span(f"resp_{marker}") == session - assert rig.spend_session(rig.sink.logged_id(f"resp_{marker}")) == session + assert rig.sink.single_span(client_id) == session + assert rig.spend_session(client_id) == session @pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_chat_raw") @@ -689,15 +673,11 @@ def _burst(rig: Rig, count: int, session_for: Mapping[int, str]) -> tuple[tuple[ return tuple(pool.map(one, range(count))) -def _is_encrypted_responses_id(identity: str) -> bool: - return identity.startswith("resp_") and _decoded_responses_id(identity) == identity - - def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]: spans: Final = rig.sink.attributes() return { session: tuple( - _canonical_id(str(attributes["gen_ai.response.id"]["stringValue"])) + str(attributes["gen_ai.response.id"]["stringValue"]) for attributes in spans if attributes.get(CONVERSATION, {}).get("stringValue") == session and "gen_ai.response.id" in attributes ) @@ -707,21 +687,13 @@ def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]] def _assert_exactly_once(rig: Rig, expected: Mapping[str, str], landed: Mapping[str, tuple[str, ...]]) -> None: spend: Final = eventually( - lambda: { - session: tuple(_canonical_id(identity) for identity in rig.spend_request_ids(session)) - for session in expected.values() - }, + lambda: {session: rig.spend_request_ids(session) for session in expected.values()}, lambda rows: all(len(values) >= 1 for values in rows.values()), seconds=70, ) assert landed == spend, (landed, spend) assert all(len(values) == 1 for values in landed.values()), landed - caller_visible: Final = { - session: (_canonical_id(identity),) - for identity, session in expected.items() - if not _is_encrypted_responses_id(identity) - } - assert {session: landed[session] for session in caller_visible} == caller_visible, landed + assert landed == {session: (identity,) for identity, session in expected.items()}, landed @pytest.mark.covers("other.observability.otel.conversation_id_sink_outage_recovers_exactly_once") diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index e25126e4fa8..93539d7641b 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -1,5 +1,4 @@ import asyncio -import base64 import json import re import signal @@ -9,7 +8,6 @@ from collections import deque from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from functools import partial from pathlib import Path from typing import Final @@ -137,17 +135,6 @@ def _upstream(request: Request) -> Reply: return _chat_reply(f"chatcmpl-{marker}", stream) -def _decoded_responses_id(identity: str) -> str: - try: - return base64.b64decode(identity.removeprefix("resp_").encode()).decode() - except (ValueError, UnicodeDecodeError): - return identity - - -def _canonical_id(identity: str) -> str: - return _decoded_responses_id(identity).rpartition("response_id:")[2] - - def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( object_value(JSON.validate_json(line[6:])) @@ -232,8 +219,7 @@ class Collector: return tuple( span for span in self.spans() - if RESPONSE_ID in span.attributes - and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(response_id) + if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == response_id ) def single_span(self, response_id: str, *, elsewhere: "Collector | None" = None) -> Span: @@ -249,11 +235,8 @@ class Collector: def landed(self, response_ids: Sequence[str]) -> dict[str, int]: spans: Final = self.spans() return { - _canonical_id(identity): sum( - 1 - for span in spans - if RESPONSE_ID in span.attributes - and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(identity) + identity: sum( + 1 for span in spans if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == identity ) for identity in response_ids } @@ -448,7 +431,7 @@ def _assert_operator_span(rig: Rig, response_id: str, marker: str) -> Span: span: Final = rig.sink.single_span(response_id) assert span.target == "/v1/traces", span assert span.ingestion_key == OPERATOR_KEY, span - assert rig.tenant_sink.landed((response_id,)) == {_canonical_id(response_id): 0} + assert rig.tenant_sink.landed((response_id,)) == {response_id: 0} bodies: Final = rig.upstream_bodies(marker) assert len(bodies) == 1, bodies assert "signoz" not in json.dumps(bodies[0]).replace(marker, ""), bodies[0] @@ -522,7 +505,7 @@ def test_responses_stream_raw_httpx_span_lands_once_after_the_stream_is_consumed marker: Final = _marker() response: Final = rig.proxy.request("POST", "/v1/responses", {"model": rig.model, "input": marker, "stream": True}) assert response.status_code == 200, response.text - _assert_operator_span(rig, _responses_id(response, marker), marker) + _assert_operator_span(rig, _responses_id(response), marker) def test_v2_flag_on_still_delivers_the_operator_span_with_the_ingestion_key(v2_rig: Rig) -> None: @@ -821,7 +804,7 @@ def _chat_id(response: httpx.Response) -> str: return next(iter(identities)) -def _responses_id(response: httpx.Response, marker: str) -> str: +def _responses_id(response: httpx.Response) -> str: if not _is_event_stream(response): return _body_id(response) completed: Final = tuple( @@ -830,7 +813,7 @@ def _responses_id(response: httpx.Response, marker: str) -> str: if event.get("type") == "response.completed" ) assert len(completed) == 1 and completed[0].startswith("resp_"), response.text - return f"resp_{marker}" + return completed[0] def _message_id(response: httpx.Response) -> str: @@ -852,7 +835,7 @@ def _burst(rig: Rig, count: int) -> tuple[tuple[int, str, str | None], ...]: stream: Final = index % 2 == 0 path, body, identity_of = ( ("/v1/chat/completions", {"model": rig.model, "messages": [{"role": "user", "content": marker}]}, _chat_id), - ("/v1/responses", {"model": rig.model, "input": marker}, partial(_responses_id, marker=marker)), + ("/v1/responses", {"model": rig.model, "input": marker}, _responses_id), ( "/v1/messages", {"model": rig.model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, @@ -874,8 +857,8 @@ def _assert_exactly_once(rig: Rig, identities: Sequence[str]) -> None: landed: Final = eventually( lambda: rig.sink.landed(identities), lambda seen: all(count >= 1 for count in seen.values()), seconds=80 ) - assert landed == {_canonical_id(identity): 1 for identity in identities}, landed - charged: Final = frozenset(identity for identity in identities if not identity.startswith("resp_")) + assert landed == {identity: 1 for identity in identities}, landed + charged: Final = frozenset(identities) spend: Final = eventually( lambda: read_rows( 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', diff --git a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py index 3d6eb7fcaff..0788b9ac85a 100644 --- a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py +++ b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py @@ -62,12 +62,12 @@ def _body(request: Request) -> dict[str, JsonValue]: # ciphertext the caller received, because the spend row id is read from response_obj["id"] before the # ResponsesIDSecurity hook rewrites it in place; the row is looked up under both ids until that ordering is fixed on # main -def _spend_row(client_id: str, issued_id: str) -> dict[str, JsonValue]: +def _spend_row(*request_ids: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows( 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' "WHERE request_id = ANY(%s)", - ([client_id, issued_id],), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array + (list(request_ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array ), lambda found: len(found) == 1, seconds=70, @@ -120,4 +120,4 @@ async def test_async_openai_sdk_responses_stream_is_served_by_the_native_respons assert issued.upstream == f"resp_upstream_{marker}", completed.response.id request: Final = _native_request(wire) assert _body(request) == {"model": GPT, "input": _prompt(marker), "stream": True}, request.body - assert _spend_row(completed.response.id, issued.issued) == _success_row(model) + assert _spend_row(completed.response.id) == _success_row(model) diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 89ad226761e..5aede070896 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -516,7 +516,7 @@ def test_responses_max_output_tokens_is_forwarded_natively_on_both_legs(gateway: for attempt in attempts: assert attempt.body["max_output_tokens"] == 1000 and _limits(attempt.body) == (None, None), attempt.body _assert_logged(issued, model, "success") - _assert_logged(_managed_response_id(streamed_issued), model, "success") + _assert_logged(streamed_issued, model, "success") def test_region_prefixed_deployment_on_a_v1_base_sends_the_bare_model_with_max_completion_tokens( diff --git a/tests/integration/providers/test_bedrock_model_discovery_wire.py b/tests/integration/providers/test_bedrock_model_discovery_wire.py index 3bf4583e736..2871e51de5e 100644 --- a/tests/integration/providers/test_bedrock_model_discovery_wire.py +++ b/tests/integration/providers/test_bedrock_model_discovery_wire.py @@ -301,11 +301,7 @@ def _spend_row_id(rows: Sequence[Mapping[str, JsonValue]], tag: str) -> str: def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ...]) -> None: - """One spend row per tagged call, carrying the id the client saw. - - FIXME: a streamed /v1/responses row carries the managed `resp_` id LiteLLM built before the proxy encrypted the - advertised one, so that row is matched by tag alone until the spend log reads the id the client received. - """ + """One spend row per tagged call, carrying the id the client saw.""" rows: Final = eventually( lambda: read_rows( 'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_tags ?| %s', @@ -316,9 +312,6 @@ def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ... ) for (identity, path, stream), tag in zip(served, tags, strict=True): row_id: Final = _spend_row_id(rows, tag) - if path == "/v1/responses" and stream: - assert row_id.startswith("resp_"), (row_id, identity) - continue assert row_id == identity, (path, stream, row_id, identity) diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 3cadc1b5ccf..42950fb18c3 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -172,7 +172,11 @@ def _upstream_id_inside(row_id: str) -> str | None: def _row_belongs_to(row_id: str, served: _Served) -> bool: if row_id == _response_id(served): return True - return served.call.endpoint == "responses" and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}" + return ( + served.call.endpoint == "responses" + and not served.call.stream + and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}" + ) def _assert_each_success_landed_once(rows: list[dict[str, JsonValue]], served: tuple[_Served, ...]) -> None: diff --git a/tests/integration/providers/test_responses_lifecycle_wire.py b/tests/integration/providers/test_responses_lifecycle_wire.py index 02269885b44..ac5b4dc6d5c 100644 --- a/tests/integration/providers/test_responses_lifecycle_wire.py +++ b/tests/integration/providers/test_responses_lifecycle_wire.py @@ -203,9 +203,6 @@ def test_retrieve_routes_to_the_deployment_that_created_the_response( @pytest.mark.parametrize("prefix", _ALIASES) def test_retrieve_returns_the_id_the_client_holds(gateway: Gateway, prefix: str) -> None: - pytest.skip( - "BUG: GET /v1/responses/{id} returns a re-encrypted id that differs byte-for-byte from the id the client sent" - ) raw_id: Final = f"resp_retrieve_id_{uuid.uuid4().hex}" def respond(request: Request) -> Reply: @@ -697,7 +694,6 @@ def test_delete_returns_the_client_held_id_not_the_provider_id( gateway: Gateway, prefix: str, ) -> None: - pytest.skip("BUG: DELETE /v1/responses/{id} returns the raw provider id instead of the proxy-issued id") raw_id: Final = f"resp_delete_id_{uuid.uuid4().hex}" def respond(request: Request) -> Reply: diff --git a/tests/integration/spend/test_responses_stream_spend_row_id.py b/tests/integration/spend/test_responses_stream_spend_row_id.py new file mode 100644 index 00000000000..58ed9f7785f --- /dev/null +++ b/tests/integration/spend/test_responses_stream_spend_row_id.py @@ -0,0 +1,465 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value +from integration._support.process import owned_proxy +from integration._support.responses_stream import frame, response_object +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_NATIVE_MODEL: Final = "openai/gpt-5.6" +_BRIDGED_MODEL: Final = "deepseek/deepseek-v4-pro" +_ANSWER: Final = "streamed answer" +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35} +_JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _sse(payload: dict[str, JsonValue]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _native_completed(upstream_id: str) -> dict[str, JsonValue]: + message: Final[dict[str, JsonValue]] = { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": _ANSWER, "annotations": []}], + } + return response_object(upstream_id, "completed", output=[message], usage=_USAGE) + + +def _native_stream(upstream_id: str) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=( + frame( + { + "type": "response.created", + "sequence_number": 0, + "response": response_object(upstream_id, "in_progress"), + } + ), + frame( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": _ANSWER, + } + ), + frame({"type": "response.completed", "sequence_number": 2, "response": _native_completed(upstream_id)}), + ), + ) + + +def _chat_stream(upstream_id: str, model: str) -> Reply: + def chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> dict[str, JsonValue]: + return { + "id": upstream_id, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + + return Reply( + content_type="text/event-stream", + chunks=( + _sse(chunk({"role": "assistant", "content": _ANSWER}, None)), + _sse( + { + **chunk({}, "stop"), + "usage": {"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35}, + } + ), + b"data: [DONE]\n\n", + ), + ) + + +def _respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + if request.method == "GET": + stored_id: Final = request.target.removeprefix("/v1/responses/") + assert stored_id.startswith("resp_upstream-"), request.target + return Reply(body=json.dumps(_native_completed(stored_id)).encode()) + if request.target.endswith("/cancel"): + cancelled_id: Final = request.target.removeprefix("/v1/responses/").removesuffix("/cancel") + assert cancelled_id.startswith("resp_upstream-"), request.target + return Reply(body=json.dumps(response_object(cancelled_id, "cancelled")).encode()) + body: Final = JSON_OBJECT.validate_json(request.body) + upstream_id: Final = f"upstream-{uuid4().hex}" + if request.target == "/v1/responses" and body.get("stream") is not True: + return Reply(body=json.dumps(_native_completed(f"resp_{upstream_id}")).encode()) + assert body.get("stream") is True, body + if request.target == "/v1/responses": + return _native_stream(f"resp_{upstream_id}") + assert request.target == "/v1/chat/completions", request.target + return _chat_stream(f"chatcmpl-{upstream_id}", string_value(body["model"])) + + +def _terminal_response_id( + gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> str: + response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers) + assert response.status_code == 200, response.text + frames: Final = tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert frames and frames[-1]["type"] == "response.completed", response.text + advertised: Final = tuple( + dict.fromkeys(string_value(object_value(event["response"])["id"]) for event in frames if "response" in event) + ) + assert len(advertised) == 1, advertised + return advertised[0] + + +def _created_response_id( + gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> str: + if body.get("stream") is True: + return _terminal_response_id(gateway, key, body, headers) + response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _read_and_cancel( + gateway: Gateway, key: str, response_id: str, headers: Mapping[str, str] | None = None +) -> tuple[httpx.Response, ...]: + responses: Final = tuple( + gateway.request(method, path, key=key, headers=headers) + for method, path in ( + ("GET", f"/v1/responses/{response_id}"), + ("GET", f"/v1/responses/{response_id}"), + ("POST", f"/v1/responses/{response_id}/cancel"), + ) + ) + assert [response.status_code for response in responses] == [200, 200, 200], [r.text for r in responses] + return responses + + +def _spend_logs_lookup(gateway: Gateway, request_id: str) -> list[JsonValue]: + response: Final = gateway.request("GET", "/spend/logs", params={"request_id": request_id}) + assert response.status_code == 200, response.text + return [row["request_id"] for row in _JSON_LIST.validate_json(response.content)] + + +def _request_ids_for_key(gateway: Gateway, key: str, count: int) -> list[str]: + rows: Final = eventually( + lambda: gateway.rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "endTime"', + (sha256(key.encode()).hexdigest(),), + ), + lambda found: len(found) >= count, + seconds=70, + ) + return [string_value(row["request_id"]) for row in rows] + + +def _logs_page_search(gateway: Gateway, client_id: str) -> list[str]: + now: Final = datetime.now(timezone.utc) + page: Final = gateway.get( + "/spend/logs/ui", + params={ + "search": client_id, + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + }, + ) + return [string_value(object_value(row)["request_id"]) for row in list_value(page["data"])] + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("litellm_model", [_NATIVE_MODEL, _BRIDGED_MODEL], ids=["native", "bridge"]) +def test_streamed_responses_spend_row_is_filed_under_the_id_the_client_saw( + gateway: Gateway, litellm_model: str +) -> None: + with wire_server(_respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=litellm_model, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + + client_id: Final = _terminal_response_id( + gateway, key, {"model": model, "input": f"marker-{uuid4().hex}", "stream": True} + ) + + assert client_id.startswith("resp_"), client_id + assert _request_ids_for_key(gateway, key, 1) == [client_id] + assert _logs_page_search(gateway, client_id) == [client_id] + assert _spend_logs_lookup(gateway, client_id) == [client_id] + + +def _proxy_config( + directory: Path, + general_settings: Mapping[str, JsonValue], + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "responses-spend.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "general_settings": {**object_value(stock["general_settings"]), **general_settings}, + "litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings}, + } + ) + ) + return path + + +def _prompt_storing_config(directory: Path) -> Path: + return _proxy_config(directory, {"store_prompts_in_spend_logs": True}) + + +def test_a_follow_up_turn_on_a_streamed_bridged_response_replays_the_first_turn( + gateway: Gateway, tmp_path: Path +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + model: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + first_id: Final = _terminal_response_id(storing, key, {"model": model, "input": first_prompt, "stream": True}) + assert _request_ids_for_key(storing, key, 1) == [first_id] + wire.drain() + + second_id: Final = _terminal_response_id( + storing, + key, + {"model": model, "input": "second turn", "previous_response_id": first_id, "stream": True}, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed + assert _request_ids_for_key(storing, key, 2) == [first_id, second_id] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("stream_first_turn", [True, False], ids=["streamed-first-turn", "non-streamed-first-turn"]) +def test_a_follow_up_through_the_id_a_retrieve_returned_replays_the_first_turn_once( + gateway: Gateway, tmp_path: Path, stream_first_turn: bool +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id( + storing, key, {"model": native, "input": first_prompt, "stream": stream_first_turn} + ) + + retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=key) + + assert retrieved.status_code == 200, retrieved.text + retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"]) + assert retrieved_id == first_id + stored: Final = _request_ids_for_key(storing, key, 2) + assert stored.count(first_id) == 1, stored + assert len(set(stored)) == 2, stored + assert _spend_logs_lookup(storing, first_id) == [first_id] + wire.drain() + + _terminal_response_id( + storing, + key, + {"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True}, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed + + +def _replayed_turns(wire_body: bytes) -> list[tuple[JsonValue, JsonValue]]: + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire_body)["messages"]) + return [(message["role"], message["content"]) for message in replayed] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("reader", ["proxy-admin", "master-key"]) +def test_a_reader_other_than_the_creator_continues_through_the_id_its_retrieve_returned( + gateway: Gateway, tmp_path: Path, reader: str +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + creator: Final = scenario.key(user_id=scenario.user(), models=[native, bridged]) + reader_key: Final = ( + scenario.key(user_id=scenario.user(user_role="proxy_admin"), models=[native, bridged]) + if reader == "proxy-admin" + else storing.key + ) + first_id: Final = _terminal_response_id( + storing, creator, {"model": native, "input": first_prompt, "stream": True} + ) + assert _request_ids_for_key(storing, creator, 1) == [first_id] + + retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=reader_key) + + assert retrieved.status_code == 200, retrieved.text + retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"]) + wire.drain() + _terminal_response_id( + storing, + reader_key, + {"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True}, + ) + assert _replayed_turns(wire.drain()[-1].body) == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ] + assert retrieved_id == first_id + + +@pytest.mark.timeout(180) +def test_reads_and_cancels_sent_with_a_session_header_add_no_replayed_turns(gateway: Gateway, tmp_path: Path) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"} + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id(storing, key, {"model": native, "input": first_prompt}, session) + + _read_and_cancel(storing, key, first_id, session) + + assert len(_request_ids_for_key(storing, key, 4)) == 4 + wire.drain() + _terminal_response_id( + storing, + key, + {"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True}, + session, + ) + assert _replayed_turns(wire.drain()[-1].body) == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ] + + +@pytest.mark.timeout(180) +def test_each_read_and_cancel_reaches_callbacks_and_spend_logs_under_an_id_of_its_own( + gateway: Gateway, tmp_path: Path +) -> None: + logged_call_types: Final = frozenset({"aresponses", "aget_responses", "acancel_responses"}) + with ( + wire_server(_respond) as wire, + wire_server(lambda _: Reply()) as sink, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url}, + config=_proxy_config(tmp_path, {}, {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}), + ) as logging_proxy, + logging_proxy.scenario() as scenario, + ): + model: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + first_id: Final = _created_response_id(logging_proxy, key, {"model": model, "input": f"marker-{uuid4().hex}"}) + + reads: Final = _read_and_cancel(logging_proxy, key, first_id) + + assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3 + batches: Final[list[bytes]] = [] # mutable-ok: drain() consumes, later polls keep earlier batches + + def delivered() -> tuple[tuple[str, str], ...]: + batches.extend(batch.body for batch in sink.drain()) + return tuple( + (string_value(event["call_type"]), string_value(event["id"])) + for event in (object_value(item) for batch in batches for item in _JSON_LIST.validate_json(batch)) + if event["call_type"] in logged_call_types + ) + + events: Final = eventually(delivered, lambda found: len(found) >= 4, seconds=30) + read_ids: Final = [logged_id for call_type, logged_id in events if call_type != "aresponses"] + assert [logged_id for call_type, logged_id in events if call_type == "aresponses"] == [first_id], events + assert len(read_ids) == 3, events + assert first_id not in read_ids, events + assert len(set(read_ids)) == 3, events + assert sorted(_request_ids_for_key(logging_proxy, key, 4)) == sorted([first_id, *read_ids]) + + +@pytest.mark.timeout(240) +def test_with_id_security_off_reads_and_cancels_add_no_rows_and_no_replayed_turns( + gateway: Gateway, tmp_path: Path +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"} + with ( + wire_server(_respond) as wire, + owned_proxy( + gateway, + tmp_path, + {}, + config=_proxy_config( + tmp_path, {"store_prompts_in_spend_logs": True, "disable_responses_id_security": True} + ), + ) as insecure, + insecure.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id(insecure, key, {"model": native, "input": first_prompt}, session) + + reads: Final = _read_and_cancel(insecure, key, first_id, session) + + assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3 + later_id: Final = _created_response_id(insecure, key, {"model": native, "input": "a later request"}) + assert _request_ids_for_key(insecure, key, 2) == [first_id, later_id] + wire.drain() + + _terminal_response_id( + insecure, + key, + {"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True}, + session, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed diff --git a/tests/unit/integrations/azure_storage/test_azure_storage.py b/tests/unit/integrations/azure_storage/test_azure_storage.py index 0f0ce24ff84..ee12291e7a0 100644 --- a/tests/unit/integrations/azure_storage/test_azure_storage.py +++ b/tests/unit/integrations/azure_storage/test_azure_storage.py @@ -5,19 +5,28 @@ import re import sys import threading from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS +from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.azure_storage import azure_storage as azure_storage_module from litellm.integrations.azure_storage.azure_storage import ( AzureBlobStorageLogger, _cached_credential_chain_token_provider, adls_safe_file_name, ) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType -from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + StandardAuditLogPayload, + StandardLoggingPayload, +) @pytest.fixture @@ -648,3 +657,171 @@ async def test_audit_log_is_uploaded_to_a_dated_audit_logs_path_with_an_account_ assert json.loads(uploaded) == audit_log assert file_client.append_data.await_args.kwargs == {"data": uploaded, "offset": 0, "length": len(uploaded)} file_client.flush_data.assert_awaited_once_with(position=len(uploaded), offset=0) + + +BACKGROUND_RESPONSE_ID: Final = "resp_YWJjZA==" +BACKGROUND_COST_POLL_KWARGS: Final = { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}, + }, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} + + +def _background_create_cost_and_other_events( + cost_kwargs: dict[str, object], +) -> tuple[tuple[dict[str, object], StandardLoggingPayload], ...]: + create: Final = StandardLoggingPayload( + id=BACKGROUND_RESPONSE_ID, + call_type="aresponses", + messages=[{"role": "user", "content": "background create input"}], + ) + cost: Final = StandardLoggingPayload(id=BACKGROUND_RESPONSE_ID, call_type="aget_responses", response_cost=0.42) + ordinary: Final = StandardLoggingPayload(id="chatcmpl-ordinary", call_type="acompletion") + batch_retrieve: Final = StandardLoggingPayload(id="batch_abc123", call_type="aretrieve_batch", response_cost=0.07) + return ( + ({"call_type": "aresponses", "litellm_params": {"metadata": {}}, "standard_logging_object": create}, create), + ({**cost_kwargs, "standard_logging_object": cost}, cost), + ( + {"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": ordinary}, + ordinary, + ), + ( + { + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": {}}, + "standard_logging_object": batch_retrieve, + }, + batch_retrieve, + ), + ) + + +def _expected_files( + events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...], +) -> dict[str, StandardLoggingPayload]: + create, cost, ordinary, batch_retrieve = (payload for _, payload in events) + return { + "resp_YWJjZA.json": create, + "resp_YWJjZA_background_cost.json": cost, + "chatcmpl-ordinary.json": ordinary, + "batch_abc123.json": batch_retrieve, + } + + +async def _log_and_flush( + logger: AzureBlobStorageLogger, events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...] +) -> None: + for kwargs, payload in events: + await logger.async_log_success_event(kwargs, {"id": payload["id"]}, None, None) + await logger.async_send_batch() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), +) +async def test_background_cost_event_keeps_its_own_file_beside_the_create_on_the_entra_path( + mock_env_vars, monkeypatch, cost_kwargs +): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + events: Final = _background_create_cost_and_other_events(cost_kwargs) + with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + await _log_and_flush(logger, events) + + base_url: Final = "https://test-account.dfs.core.windows.net/test-container/" + appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]] + uploaded: Final = { + call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends + } + assert uploaded == _expected_files(events), ( + "the create, its background cost event and every other event must each keep their own file" + ) + + +@pytest.mark.asyncio +async def test_failure_event_is_uploaded_to_a_file_named_after_its_id(mock_env_vars, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + failure: Final = StandardLoggingPayload(id="chatcmpl-failed", call_type="acompletion", error_str="provider failure") + with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + await logger.async_log_failure_event( + {"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": failure}, + None, + None, + None, + ) + await logger.async_send_batch() + + base_url: Final = "https://test-account.dfs.core.windows.net/test-container/" + appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]] + assert {call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends} == { + "chatcmpl-failed.json": failure + } + + +@pytest.mark.asyncio +async def test_background_cost_event_keeps_its_own_file_beside_the_create_with_an_account_key( + workload_identity_env_vars, monkeypatch +): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr(azure_storage_module, "datetime", _FrozenDatetime) + events: Final = _background_create_cost_and_other_events(BACKGROUND_COST_POLL_KWARGS) + + uploaded: Final[dict[str, StandardLoggingPayload]] = {} + + def file_client_for(file_path: str) -> MagicMock: + file_client = MagicMock() + file_client.create_file = AsyncMock() + file_client.flush_data = AsyncMock() + + async def append_data(data: bytes, offset: int, length: int) -> None: + uploaded[file_path] = json.loads(data) + + file_client.append_data = append_data + return file_client + + file_system_client = MagicMock() + file_system_client.get_file_client = MagicMock(side_effect=file_client_for) + service_client = MagicMock() + service_client.get_file_system_client = MagicMock(return_value=file_system_client) + fake_aio_module = MagicMock() + fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client) + + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): + logger = AzureBlobStorageLogger() + await _log_and_flush(logger, events) + + today: Final = _FrozenDatetime.now().strftime("%Y-%m-%d") + assert uploaded == {f"{today}/{name}": payload for name, payload in _expected_files(events).items()}, ( + "the create, its background cost event and every other event must each keep their own file" + ) diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index f94d7bbee37..42e69f95bb4 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -20,10 +20,16 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import ( strip_cache_hit_suffix, ) from litellm.integrations.clickhouse.context import lens_analysis +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.secret_redaction import REDACTED from litellm.tracing.types import SpendLogRecord -from litellm.types.utils import StandardLoggingPayload +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, StandardLoggingPayload _JSON_OBJECT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) @@ -349,6 +355,41 @@ def test_cache_hit_id_is_stripped_for_response_id(): assert strip_cache_hit_suffix("chatcmpl-xyz") == "chatcmpl-xyz" +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + }, + id="cost-poll", + ), + pytest.param( + { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, + }, + id="passthrough-settlement", + ), + ), +) +def test_background_cost_row_keeps_its_own_request_id_beside_its_create(cost_kwargs: Mapping[str, object]) -> None: + payload: Final = cast(StandardLoggingPayload, _payload(id="resp_YWJjZA", call_type="aresponses")) + create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}} + + rows: Final = [spend_log_row_from_payload(payload, kwargs) for kwargs in (create_kwargs, cost_kwargs)] + + assert [row["request_id"] for row in rows] == [ + "resp_YWJjZA", + f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}", + ] + assert [row["response_id"] for row in rows] == ["resp_YWJjZA", "resp_YWJjZA"] + + def test_parse_traceparent_valid_missing_malformed(): assert parse_traceparent(TRACEPARENT) == (TRACE_ID, SPAN_ID) assert parse_traceparent(None) == ("", "") diff --git a/tests/unit/integrations/conftest.py b/tests/unit/integrations/conftest.py index 48a01ed8d48..9278a280fe7 100644 --- a/tests/unit/integrations/conftest.py +++ b/tests/unit/integrations/conftest.py @@ -1,14 +1,18 @@ import functools import http.server +import importlib.abc import ipaddress import os import queue import ssl +import sys import threading -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from importlib.machinery import ModuleSpec from pathlib import Path +from types import ModuleType from typing import Final import pytest @@ -75,6 +79,36 @@ def write_self_signed_cert(directory: Path, stem: str) -> tuple[Path, Path]: return certificate_path, key_path +PROXY_ONLY_PACKAGES: Final = frozenset({"fastapi", "starlette"}) + + +class _ProxyOnlyPackageBlocker(importlib.abc.MetaPathFinder): + def find_spec( + self, + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> ModuleSpec | None: + if fullname.partition(".")[0] in PROXY_ONLY_PACKAGES: + raise ModuleNotFoundError(f"No module named {fullname!r}") + return None + + +def _needs_proxy_extras(module_name: str) -> bool: + return ( + module_name.partition(".")[0] in PROXY_ONLY_PACKAGES + or module_name == "litellm.proxy" + or module_name.startswith("litellm.proxy.") + ) + + +@pytest.fixture +def sdk_only_install(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in tuple(filter(_needs_proxy_extras, sys.modules)): + monkeypatch.delitem(sys.modules, module_name) + monkeypatch.setattr(sys, "meta_path", [_ProxyOnlyPackageBlocker(), *sys.meta_path]) + + @pytest.fixture(autouse=True) def restore_process_environment() -> Iterator[None]: original: Final = dict(os.environ) diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index 6f143a244f6..c29985b1e41 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -13,9 +13,11 @@ import litellm import pytest import respx +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN mock_response_data: Final = { @@ -172,9 +174,7 @@ async def test_get_payload_not_found(monkeypatch): assert payload is None -def _gcs_logger_with_upload_status( - monkeypatch: pytest.MonkeyPatch, status_code: int -) -> tuple[GCSBucketLogger, Mock]: +def _gcs_logger_with_upload_status(monkeypatch: pytest.MonkeyPatch, status_code: int) -> tuple[GCSBucketLogger, Mock]: monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr(google.auth, "default", _google_default_credentials) monkeypatch.setenv("GCS_FLUSH_INTERVAL", "3600") @@ -394,3 +394,98 @@ async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch: p (thread,) = threads assert thread is not threading.main_thread() + + +@pytest.mark.asyncio +async def test_unbatched_background_cost_event_keeps_its_own_object_readable_by_its_spend_log_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false") + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + object_prefix: Final = "/storage/v1/b/test-bucket/o/" + stored: Final[dict[str, object]] = {} + + def bucket(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + stored[request.url.params["name"]] = json.loads(request.content) + return httpx.Response(200, json={"name": request.url.params["name"]}) + name: Final = request.url.path.removeprefix(object_prefix) + if name not in stored: + return httpx.Response(404, text="No such object") + return httpx.Response(200, content=json.dumps(stored[name]).encode("utf-8")) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket)) + response_id: Final = "resp_YmFja2dyb3VuZA" + create: Final = { + "id": response_id, + "messages": [{"role": "user", "content": "background create input"}], + "error_str": None, + } + cost: Final = {"id": response_id, "response_cost": 0.42, "error_str": None} + ordinary: Final = {"id": "chatcmpl-ordinary", "error_str": None} + batch_retrieve: Final = {"id": "batch_abc123", "response_cost": 0.07, "error_str": None} + events: Final = ( + ({"call_type": "aresponses", "litellm_params": {"metadata": {}}}, create), + ( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + }, + cost, + ), + ({"call_type": "acompletion", "litellm_params": {"metadata": {}}}, ordinary), + ({"call_type": "aretrieve_batch", "litellm_params": {"metadata": {}}}, batch_retrieve), + ) + + for kwargs, payload in events: + await logger.async_log_success_event( + {**kwargs, "standard_logging_object": payload}, {"id": payload["id"]}, None, None + ) + await logger.flush_queue() + + assert {name.split("/", maxsplit=1)[1]: payload for name, payload in stored.items()} == { + response_id: create, + f"{response_id}_background_cost": cost, + "chatcmpl-ordinary": ordinary, + "batch_abc123": batch_retrieve, + }, "the create, its background cost event and every other event must each keep their own object" + now: Final = datetime.now(timezone.utc) + assert await logger.get_request_response_payload(f"{response_id}_background_cost", now, None) == cost + assert await logger.get_request_response_payload(response_id, now, None) == create + + +@pytest.mark.asyncio +async def test_unbatched_gcs_log_id_still_names_a_background_cost_event_object(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false") + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + stored: Final[dict[str, object]] = {} + + def bucket(request: httpx.Request) -> httpx.Response: + stored[request.url.params["name"]] = json.loads(request.content) + return httpx.Response(200, json={"name": request.url.params["name"]}) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket)) + cost: Final = {"id": "resp_YmFja2dyb3VuZA", "response_cost": 0.42, "error_str": None} + await logger.async_log_success_event( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": { + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + "gcs_log_id": "qa-run-7", + } + }, + "standard_logging_object": cost, + }, + {"id": cost["id"]}, + None, + None, + ) + await logger.flush_queue() + + ((name, payload),) = stored.items() + object_name: Final = name.split("/", maxsplit=1)[1] + assert object_name.startswith("custom-") and object_name.endswith("-qa-run-7"), object_name + assert payload == cost diff --git a/tests/unit/integrations/test_deepeval.py b/tests/unit/integrations/test_deepeval.py index 6dd3aef0a03..4523a9c61d2 100644 --- a/tests/unit/integrations/test_deepeval.py +++ b/tests/unit/integrations/test_deepeval.py @@ -2,11 +2,15 @@ import os import unittest from litellm._uuid import uuid from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch from litellm.integrations.deepeval.api import Endpoints, HttpMethods from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.deepeval.types import SpanApiType, TraceSpanApiStatus +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN class TestDeepEvalLogger(unittest.TestCase): @@ -94,6 +98,37 @@ class TestDeepEvalLogger(unittest.TestCase): self.assertEqual(llm_span["inputTokenCount"], 10) self.assertEqual(llm_span["outputTokenCount"], 20) + def test_background_cost_span_keeps_its_own_uuid_beside_its_create(self) -> None: + standard_logging_object: Final = { + "id": "resp_YWJjZA", + "trace_id": self.trace_id, + "model": self.model, + "response": {"usage": {}, "choices": [{"message": {"content": "OK"}}]}, + } + create_kwargs: Final = { + "input": self.input_str, + "call_type": "aresponses", + "litellm_params": {"metadata": {}}, + "standard_logging_object": standard_logging_object, + } + cost_kwargs: Final = { + "input": self.input_str, + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + "standard_logging_object": standard_logging_object, + } + + for kwargs in (create_kwargs, cost_kwargs): + self.logger.log_success_event(kwargs, self.mock_response_obj, self.start_time, self.end_time) + + span_uuids: Final = [ + call.kwargs["body"]["llmSpans"][0]["uuid"] + for call in self.mock_api_instance.send_request.call_args_list + ] + self.assertEqual(span_uuids, ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"]) + def test_log_failure_event(self): error_message = "This is an error." kwargs = { diff --git a/tests/unit/integrations/test_dynamodb.py b/tests/unit/integrations/test_dynamodb.py new file mode 100644 index 00000000000..d139c49a49e --- /dev/null +++ b/tests/unit/integrations/test_dynamodb.py @@ -0,0 +1,52 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.integrations.dynamodb import DyanmoDBLogger +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = { + "call_type": "aget_responses", + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} + + +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), +) +@pytest.mark.usefixtures("sdk_only_install") +def test_background_cost_event_keeps_its_own_item_beside_its_create( + monkeypatch: pytest.MonkeyPatch, cost_kwargs: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "dynamodb_table_name", "litellm-logs") + dynamodb: Final = MagicMock() + logger: Final = DyanmoDBLogger(dynamodb=dynamodb) + response_obj: Final = {"id": "resp_YWJjZA", "usage": {"total_tokens": 3}} + now: Final = datetime.now(timezone.utc) + create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}} + + for kwargs in (create_kwargs, cost_kwargs): + logger.log_event(kwargs, response_obj, now, now, print_verbose=print) + + dynamodb.Table.assert_called_with("litellm-logs") + written_ids: Final = [call.kwargs["Item"]["id"] for call in dynamodb.Table.return_value.put_item.call_args_list] + assert written_ids == ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"] diff --git a/tests/unit/integrations/test_openmeter.py b/tests/unit/integrations/test_openmeter.py index 2d09e1572db..fc4e2fcd7a3 100644 --- a/tests/unit/integrations/test_openmeter.py +++ b/tests/unit/integrations/test_openmeter.py @@ -1,11 +1,30 @@ import json import os +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.openmeter import OpenMeterLogger +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = { + "call_type": "aget_responses", + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} class TestOpenMeterIntegration: @@ -36,6 +55,31 @@ class TestOpenMeterIntegration: with pytest.raises(Exception, match=r"Missing keys.*OPENMETER_API_KEY"): OpenMeterLogger() + @pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), + ) + @pytest.mark.usefixtures("sdk_only_install") + def test_background_cost_event_gets_its_own_event_id_beside_its_create( + self, cost_kwargs: Mapping[str, object] + ) -> None: + logger: Final = OpenMeterLogger() + response_obj: Final = {"id": "resp_YWJjZA"} + create_kwargs: Final[Mapping[str, object]] = { + "user": "user-1", + "call_type": "aresponses", + "litellm_params": {"metadata": {}}, + } + + create_event: Final = logger._common_logic(create_kwargs, response_obj) + cost_event: Final = logger._common_logic({**cost_kwargs, "user": "user-1", "response_cost": 0.42}, response_obj) + + assert create_event["id"] == "resp_YWJjZA" + assert cost_event["id"] == f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}" + def test_common_logic_with_string_user(self): """Test that _common_logic correctly handles string user parameter""" logger = OpenMeterLogger() diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f1854563cd5..13e181b2b92 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -11090,6 +11090,40 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() assert logging_obj.response_cost_calculator(result=assembled) == 0.0042 +@pytest.mark.parametrize("continued_response_id", ["resp_continuedafteracutstream", None]) +@pytest.mark.parametrize("client_facing_response_id", ["resp_encryptedidtheclientsaw", None]) +def test_sync_success_handler_logs_a_stream_under_the_id_the_client_saw( + client_facing_response_id, continued_response_id +): + class Recorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged = [] + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged.append((response_obj.id, kwargs["standard_logging_object"]["id"])) + + recorder: Final = Recorder() + logging_obj: Final = _responses_stream_logging_obj() + logging_obj.dynamic_success_callbacks = [recorder] + logging_obj.client_facing_response_id = client_facing_response_id + if continued_response_id is not None: + logging_obj.record_continued_partial( + Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), 0.0, continued_response_id + ) + now: Final = datetime.datetime.now() + + logging_obj.success_handler( + result=_completed_responses_event(ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5)), + start_time=now, + end_time=now, + cache_hit=False, + ) + + expected_id: Final = client_facing_response_id or continued_response_id or "resp-1" + assert recorder.logged == [(expected_id, expected_id)] + + def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): logging_obj: Final = _responses_stream_logging_obj() inner_response: Final = ResponsesAPIResponse( diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index f52f40ad55d..f73d33cc93e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5150,6 +5150,23 @@ def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_th assert request_ids == ("resp_shared", "resp_shared_background_cost") +@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) +def test_a_read_or_cancel_row_carries_the_id_its_callbacks_logged(call_type: str): + row: Final = get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-5.6", + "litellm_call_id": "call-read", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj={"id": "resp_logged_by_callbacks", "object": "response", "status": "completed"}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert row["request_id"] == "resp_logged_by_callbacks" + + @pytest.mark.parametrize("forged_marker", [True, BackgroundResponseSettlement.SETTLES.value]) def test_a_request_body_cannot_move_its_spend_row_onto_a_background_settlement_id(forged_marker): """A failed call logs its raw request body as kwargs, so a body field named like the settlement diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index f1d841a8e5e..9a8fa19e2e4 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1083,6 +1083,128 @@ async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored ] +def _stored_response_turn( + request_id: str, call_type: str, status: str, output: list[dict[str, object]], prompt: str | None = None +) -> dict[str, object]: + return { + "request_id": request_id, + "call_type": call_type, + "session_id": "session-with-reads", + "proxy_server_request": json.dumps({} if prompt is None else {"input": [{"role": "user", "content": prompt}]}), + "response": {"id": request_id, "object": "response", "status": status, "output": output}, + } + + +def _identified_output_message(item_id: str, text: str) -> dict[str, object]: + return {**_output_message(text), "id": item_id} + + +def _replayed_roles_and_texts(session: ChatCompletionSession) -> list[tuple[object, object]]: + return [ + (role, _text_of(content)) for role, content in (_role_and_content(message) for message in session["messages"]) + ] + + +@pytest.mark.asyncio +async def test_reads_and_cancels_of_a_response_the_session_already_holds_replay_no_turns(): + answer: Final = [_identified_output_message("msg_answer", "OK")] + spend_logs: Final = ( + _stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Reply with OK"), + _stored_response_turn("resp_read", "aget_responses", "completed", answer), + _stored_response_turn("resp_streamed_read", "aget_responses", "completed", answer), + _stored_response_turn("resp_cancel", "acancel_responses", "completed", answer), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_create") + + assert _replayed_roles_and_texts(session) == [("user", "Reply with OK"), ("assistant", "OK")] + + +@pytest.mark.asyncio +async def test_a_polled_background_response_replays_its_finished_output_once(): + answer: Final = [_identified_output_message("msg_answer", "One, two, three")] + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"), + _stored_response_turn("resp_poll_running", "aget_responses", "in_progress", []), + _stored_response_turn("resp_poll_done", "aget_responses", "completed", answer), + _stored_response_turn("resp_poll_again", "aget_responses", "completed", answer), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to three"), ("assistant", "One, two, three")] + + +@pytest.mark.asyncio +async def test_a_background_response_that_grows_between_polls_replays_each_output_item_once(): + first: Final = _identified_output_message("msg_first", "One") + second: Final = _identified_output_message("msg_second", "Two") + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"), + _stored_response_turn("resp_poll_partial", "aget_responses", "in_progress", [first]), + _stored_response_turn("resp_poll_done", "aget_responses", "completed", [first, second]), + _stored_response_turn("resp_poll_again", "aget_responses", "completed", [first, second]), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")] + + +@pytest.mark.asyncio +async def test_an_item_a_poll_caught_half_written_replays_its_finished_text_once(): + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"), + _stored_response_turn( + "resp_poll_partial", "aget_responses", "in_progress", [_identified_output_message("msg_answer", "One, t")] + ), + _stored_response_turn( + "resp_poll_done", + "aget_responses", + "completed", + [_identified_output_message("msg_answer", "One, two, three"), _identified_output_message("msg_more", "Done")], + ), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [ + ("user", "Count to three"), + ("assistant", "One, two, three"), + ("assistant", "Done"), + ] + + +@pytest.mark.asyncio +async def test_a_polled_output_item_without_an_id_still_replays_in_order(): + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"), + _stored_response_turn( + "resp_poll_done", + "aget_responses", + "completed", + [_output_message("One"), _identified_output_message("msg_two", "Two")], + ), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")] + + +@pytest.mark.asyncio +async def test_a_read_that_ends_before_its_streamed_create_replays_the_answer_once_after_the_prompt(): + answer: Final = [_identified_output_message("msg_answer", "OK")] + spend_logs: Final = ( + _stored_response_turn("resp_read_mid_stream", "aget_responses", "completed", answer), + _stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Say OK"), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_create") + + assert _replayed_roles_and_texts(session) == [("user", "Say OK"), ("assistant", "OK")] + + def _tool_call_ids(message: object) -> tuple[str, ...]: tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ()) diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py index 7e74bfd0f17..8dc50fa33c7 100644 --- a/tests/unit/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -4,23 +4,33 @@ Tests for ResponsesIDSecurity hook. Tests the security hook that prevents user B from seeing response from user A. """ +import datetime +from collections.abc import Callable, Mapping from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.responses_id_security import ( ResponsesIDSecurity, _streams_response_ids, ) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( GenericEvent, + ResponseAPIUsage, ResponseCompletedEvent, ResponseCreatedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import SpecialEnums +from litellm.types.responses.main import DeleteResponseResult +from litellm.types.utils import CallTypesLiteral, Choices, Message, ModelResponse, SpecialEnums, Usage @pytest.fixture @@ -55,12 +65,12 @@ class TestIsEncryptedResponseId: import litellm.proxy.hooks.responses_id_security as responses_module with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: - mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456" - - result = responses_id_security._is_encrypted_response_id( - "resp_encrypted_value" + mock_decrypt.return_value = ( + f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456" ) + result = responses_id_security._is_encrypted_response_id("resp_encrypted_value") + assert result is True mock_decrypt.assert_called_once() @@ -88,9 +98,7 @@ class TestDecryptResponseId: with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_original_123;user_id:user-456;team_id:team-789" - original_id, user_id, team_id = responses_id_security._decrypt_response_id( - "resp_encrypted_value" - ) + original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_encrypted_value") assert original_id == "resp_original_123" assert user_id == "user-456" @@ -104,9 +112,7 @@ class TestDecryptResponseId: with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: mock_decrypt.return_value = None - original_id, user_id, team_id = responses_id_security._decrypt_response_id( - "resp_plain_value" - ) + original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_plain_value") assert original_id == "resp_plain_value" assert user_id is None @@ -116,9 +122,7 @@ class TestDecryptResponseId: class TestCheckUserAccessToResponseId: """Test check_user_access_to_response_id function""" - def test_check_user_access_same_user( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_same_user(self, responses_id_security, mock_user_api_key_dict): """Test that same user has access to their response ID""" result = responses_id_security.check_user_access_to_response_id( response_id_user_id="test-user-123", @@ -128,9 +132,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_different_user_raises_exception( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_different_user_raises_exception(self, responses_id_security, mock_user_api_key_dict): """Test that different user is denied access to response ID""" with patch("litellm.proxy.proxy_server.general_settings", {}): with pytest.raises(HTTPException) as exc_info: @@ -143,9 +145,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "Forbidden" in exc_info.value.detail - def test_check_user_access_different_team_raises_exception( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_different_team_raises_exception(self, responses_id_security, mock_user_api_key_dict): """Test that different team is denied access to response ID""" with patch("litellm.proxy.proxy_server.general_settings", {}): with pytest.raises(HTTPException) as exc_info: @@ -158,9 +158,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "Forbidden" in exc_info.value.detail - def test_check_user_access_team_a_to_team_b_without_user_id( - self, responses_id_security - ): + def test_check_user_access_team_a_to_team_b_without_user_id(self, responses_id_security): """Test that key from team A (without user_id) cannot access response from team B (without user_id)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -179,9 +177,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "team" in exc_info.value.detail.lower() - def test_check_user_access_team_a_to_team_b_with_user_id( - self, responses_id_security - ): + def test_check_user_access_team_a_to_team_b_with_user_id(self, responses_id_security): """Test that key from team A (without user_id) cannot access response from team B (with user_id)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -217,9 +213,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_admin_can_access_any_response( - self, responses_id_security - ): + def test_check_user_access_admin_can_access_any_response(self, responses_id_security): """Test that proxy admin can access any response ID""" from litellm.proxy._types import LitellmUserRoles @@ -238,9 +232,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_security_disabled( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_security_disabled(self, responses_id_security, mock_user_api_key_dict): """Test that when security is disabled, any user can access any response""" with patch( "litellm.proxy.proxy_server.general_settings", @@ -266,9 +258,7 @@ class TestAsyncPreCallHook: """Test pre-call hook decrypts previous_response_id for aresponses call""" data = {"previous_response_id": "resp_encrypted_value"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -284,15 +274,11 @@ class TestAsyncPreCallHook: assert result["previous_response_id"] == "resp_original_123" @pytest.mark.asyncio - async def test_async_pre_call_hook_aget_responses( - self, responses_id_security, mock_user_api_key_dict, mock_cache - ): + async def test_async_pre_call_hook_aget_responses(self, responses_id_security, mock_user_api_key_dict, mock_cache): """Test pre-call hook decrypts response_id for aget_responses call""" data = {"response_id": "resp_encrypted_456"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -308,9 +294,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_456" @pytest.mark.asyncio - async def test_async_pre_call_hook_team_a_accessing_team_b_response( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_team_a_accessing_team_b_response(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from accessing team B response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -320,9 +304,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -341,9 +323,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_team_a_accessing_team_b_with_user( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_team_a_accessing_team_b_with_user(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A (no user) from accessing team B response (with user)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -353,9 +333,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b_with_user"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -375,9 +353,7 @@ class TestAsyncPreCallHook: assert "forbidden" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_same_team_access( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_same_team_access(self, responses_id_security, mock_cache): """Test pre-call hook allows team A to access their own team's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -387,9 +363,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_a"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -405,9 +379,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_team_a" @pytest.mark.asyncio - async def test_async_pre_call_hook_adelete_responses_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_adelete_responses_team_security(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from deleting team B's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -417,9 +389,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -438,9 +408,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_acancel_responses_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_acancel_responses_team_security(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from canceling team B's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -450,9 +418,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -470,16 +436,13 @@ class TestAsyncPreCallHook: assert exc_info.value.status_code == 403 assert "team" in exc_info.value.detail.lower() - @pytest.mark.asyncio async def test_async_pre_call_hook_alist_input_items_decrypts_response_id( self, responses_id_security, mock_user_api_key_dict, mock_cache ): data = {"response_id": "resp_encrypted_789"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -496,9 +459,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_789" @pytest.mark.asyncio - async def test_async_pre_call_hook_alist_input_items_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_alist_input_items_team_security(self, responses_id_security, mock_cache): mock_auth_team_a = MagicMock() mock_auth_team_a.user_id = None mock_auth_team_a.team_id = "team-a" @@ -506,9 +467,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -620,9 +579,7 @@ class TestAsyncPostCallStreamingIteratorHook: "/openai/v1/responses/resp_rawprovider123", ], ) - async def test_streamed_id_encrypted_on_all_responses_routes( - self, responses_id_security, route, monkeypatch - ): + async def test_streamed_id_encrypted_on_all_responses_routes(self, responses_id_security, route, monkeypatch): streamed_id = await self._drain_streamed_id(responses_id_security, route, monkeypatch) assert streamed_id != "resp_rawprovider123" @@ -634,12 +591,8 @@ class TestAsyncPostCallStreamingIteratorHook: ) @pytest.mark.asyncio - async def test_streamed_id_untouched_on_non_responses_route( - self, responses_id_security, monkeypatch - ): - streamed_id = await self._drain_streamed_id( - responses_id_security, "/chat/completions", monkeypatch - ) + async def test_streamed_id_untouched_on_non_responses_route(self, responses_id_security, monkeypatch): + streamed_id = await self._drain_streamed_id(responses_id_security, "/chat/completions", monkeypatch) assert streamed_id == "resp_rawprovider123" assert not responses_id_security._is_encrypted_response_id(streamed_id) @@ -698,12 +651,8 @@ class TestStreamedGenericEventIdEncryption: def _advertised_ids(events): nested = (getattr(event, "response", None) for event in events) return [ - payload["id"] if isinstance(payload, dict) else payload.id - for payload in nested - if payload is not None - ] + [ - event.id for event in events if isinstance(getattr(event, "id", None), str) - ] + payload["id"] if isinstance(payload, dict) else payload.id for payload in nested if payload is not None + ] + [event.id for event in events if isinstance(getattr(event, "id", None), str)] async def _drain(self, responses_id_security, monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij") @@ -723,9 +672,7 @@ class TestStreamedGenericEventIdEncryption: ] @pytest.mark.asyncio - async def test_every_event_advertises_the_same_encrypted_id( - self, responses_id_security, monkeypatch - ): + async def test_every_event_advertises_the_same_encrypted_id(self, responses_id_security, monkeypatch): events = await self._drain(responses_id_security, monkeypatch) advertised = self._advertised_ids(events) @@ -742,22 +689,14 @@ class TestStreamedGenericEventIdEncryption: ) @pytest.mark.asyncio - async def test_raw_provider_id_never_reaches_the_client( - self, responses_id_security, monkeypatch - ): + async def test_raw_provider_id_never_reaches_the_client(self, responses_id_security, monkeypatch): events = await self._drain(responses_id_security, monkeypatch) - assert [self.RAW_ID in event.model_dump_json() for event in events] == [ - False - ] * len(events) + assert [self.RAW_ID in event.model_dump_json() for event in events] == [False] * len(events) @pytest.mark.asyncio - async def test_sibling_fields_survive_the_rewrite( - self, responses_id_security, monkeypatch - ): - _, queued, keepalive, later, _ = await self._drain( - responses_id_security, monkeypatch - ) + async def test_sibling_fields_survive_the_rewrite(self, responses_id_security, monkeypatch): + _, queued, keepalive, later, _ = await self._drain(responses_id_security, monkeypatch) assert queued.response["status"] == "queued" assert later.response["status"] == "in_progress" @@ -769,27 +708,19 @@ class TestAsyncPostCallSuccessHook: """Test async_post_call_success_hook function""" @pytest.mark.asyncio - async def test_async_post_call_success_hook_encrypts_response( - self, responses_id_security, mock_user_api_key_dict - ): + async def test_async_post_call_success_hook_encrypts_response(self, responses_id_security, mock_user_api_key_dict): """Test post-call hook encrypts ResponsesAPIResponse""" - mock_response = ResponsesAPIResponse( - id="resp_789", created_at=1234567890, output=[], status="completed" - ) + mock_response = ResponsesAPIResponse(id="resp_789", created_at=1234567890, output=[], status="completed") data = {} - with patch.object( - responses_id_security, "_encrypt_response_id", return_value=mock_response - ) as mock_encrypt: + with patch.object(responses_id_security, "_encrypt_response_id", return_value=mock_response) as mock_encrypt: result = await responses_id_security.async_post_call_success_hook( data=data, user_api_key_dict=mock_user_api_key_dict, response=mock_response, ) - mock_encrypt.assert_called_once_with( - mock_response, mock_user_api_key_dict, request_cache=None - ) + mock_encrypt.assert_called_once_with(mock_response, mock_user_api_key_dict, request_cache=None) assert result == mock_response @pytest.mark.asyncio @@ -846,9 +777,7 @@ def _auth(user_id="owner-user", team_id="owner-team", user_role=None): def _issue_managed_id(hook, owner, provider_response_id=_FABRICATED_PROVIDER_RESPONSE_ID): """Mint an id exactly the way the proxy hands one to a client on create.""" issued = hook._encrypt_response_id( - ResponsesAPIResponse( - id=provider_response_id, created_at=1234567890, output=[], status="completed" - ), + ResponsesAPIResponse(id=provider_response_id, created_at=1234567890, output=[], status="completed"), owner, ) return issued.id @@ -1080,9 +1009,7 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert exc_info.value.status_code == 403 @pytest.mark.asyncio - async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response( - self, mock_cache, salt_key_env - ): + async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response(self, mock_cache, salt_key_env): hook = _hook() stranger = _auth(user_id="stranger-user", team_id="stranger-team") stranger_id = _issue_managed_id(hook, stranger, provider_response_id="resp_strangerownprovideridcccccccc") @@ -1098,3 +1025,436 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert result["response_id"] == "resp_strangerownprovideridcccccccc" assert result["response_id"] != victim_provider_id + + +def _admin(): + from litellm.proxy._types import LitellmUserRoles + + return _auth(user_id="admin-user", team_id="admin-team", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _provider_answer(response_id: str) -> ResponsesAPIResponse: + return ResponsesAPIResponse(id=response_id, created_at=1234567890, output=[], status="completed") + + +class TestReadsAnswerWithTheIdTheClientHolds: + """A follow-up looks its history up by the id the client holds, so a read must not hand out a new one.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) + @pytest.mark.parametrize( + "issuer, reader", + [ + (_auth(), _auth()), + (_auth(team_id=None), _auth(team_id="another-team")), + (_auth(user_id=None), _auth(user_id=None)), + (_auth(user_id=None), _auth(user_id="teammate-user")), + (_auth(), _admin()), + (_auth(user_id=None), _auth(user_id=None, team_id="admin-team", user_role=_admin().user_role)), + ], + ids=[ + "same-key", + "same-user-new-team", + "same-team-no-user", + "teammate-reading-a-team-key-id", + "admin-user", + "admin-without-user", + ], + ) + async def test_a_caller_with_access_gets_back_the_exact_id_they_sent( + self, mock_cache, salt_key_env, call_type, issuer, reader + ): + hook = _hook() + issued_id = _issue_managed_id(hook, issuer) + data = await hook.async_pre_call_hook( + user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type=call_type + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=reader, response=_provider_answer(data["response_id"]) + ) + + assert result.id == issued_id + + @pytest.mark.asyncio + async def test_a_follow_up_gets_an_id_for_its_own_response_never_the_previous_one(self, mock_cache, salt_key_env): + hook = _hook() + owner = _auth() + previous_id = _issue_managed_id(hook, owner) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, + cache=mock_cache, + data={"model": "gpt-fake", "previous_response_id": previous_id}, + call_type="aresponses", + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=owner, response=_provider_answer("resp_followupprovideridzzzzzzzzzzzzzzz") + ) + + assert result.id != previous_id + assert hook._decrypt_response_id(result.id) == ( + "resp_followupprovideridzzzzzzzzzzzzzzz", + "owner-user", + "owner-team", + ) + + @pytest.mark.asyncio + async def test_an_admin_read_by_a_raw_provider_id_gets_a_managed_id_back(self, mock_cache, salt_key_env): + hook = _hook() + admin = _admin() + data = await hook.async_pre_call_hook( + user_api_key_dict=admin, + cache=mock_cache, + data={"response_id": _FABRICATED_UNMANAGED_ID}, + call_type="aget_responses", + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=admin, response=_provider_answer(_FABRICATED_UNMANAGED_ID) + ) + + assert hook._is_encrypted_response_id(result.id) + assert hook._decrypt_response_id(result.id) == (_FABRICATED_UNMANAGED_ID, "admin-user", "admin-team") + + +_LOGGED_MANAGED_ID = "resp_managedidlitellmbuiltbeforethehookran" + + +def _streamed_response(response_id): + return ResponsesAPIResponse( + id=response_id, + created_at=0, + model="gpt-5.1", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + usage=ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5), + ) + + +def _upstream_frames(): + return ( + ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=_streamed_response(_LOGGED_MANAGED_ID) + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID) + ), + ) + + +def _native_logged_result(): + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID) + ) + + +def _bridged_logged_result(): + return ModelResponse( + id="chatcmpl-bridgedchatcompletionid", + model="gpt-5.1", + choices=[Choices(message=Message(role="assistant", content="hi"), finish_reason="stop", index=0)], + usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +class _SpendRowRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.spend_row_ids = [] + self.standard_logging_ids = [] + self.callback_response_ids = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.spend_row_ids.append(get_logging_payload(kwargs, response_obj, start_time, end_time)["request_id"]) + self.callback_response_ids.append(response_obj.id) + if "standard_logging_object" in kwargs: + self.standard_logging_ids.append(kwargs["standard_logging_object"]["id"]) + + +async def _stream_logged_mid_stream( + hook: ResponsesIDSecurity, + route: str, + logged_result: ResponseCompletedEvent | ModelResponse, + caller: UserAPIKeyAuth | None = None, + request_data: Mapping[str, object] | None = None, +) -> tuple[str, _SpendRowRecorder]: + streaming_caller = caller if caller is not None else _auth() + recorder = _SpendRowRecorder() + logging_obj = LiteLLMLoggingObj( + model="gpt-5.1", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime.datetime.now(), + litellm_call_id="streamed-responses-call", + function_id="streamed-responses-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""} + ) + upstream = _upstream_frames() + + async def frames(): + for frame in upstream: + yield frame + + client_stream = hook.async_post_call_streaming_iterator_hook( + user_api_key_dict=streaming_caller.model_copy(update={"request_route": route}), + response=frames(), + request_data={**(request_data or {}), "litellm_logging_obj": logging_obj}, + ) + first_frame = await anext(client_stream) + now = datetime.datetime.now() + await logging_obj.async_success_handler(result=logged_result, start_time=now, end_time=now, cache_hit=False) + client_frames = [first_frame, *[frame async for frame in client_stream]] + return client_frames[-1].response.id, recorder + + +class TestStreamedSpendRowCarriesTheIdTheClientSaw: + @pytest.mark.asyncio + @pytest.mark.parametrize("logged_result", [_native_logged_result, _bridged_logged_result], ids=["native", "bridge"]) + @pytest.mark.parametrize("route", ["/v1/responses", "/responses", "/openai/v1/responses"]) + async def test_spend_row_and_callbacks_log_the_encrypted_id_of_the_terminal_frame( + self, salt_key_env, route, logged_result + ): + hook = _hook() + + client_id, recorder = await _stream_logged_mid_stream(hook, route, logged_result()) + + assert hook._decrypt_response_id(client_id) == (_LOGGED_MANAGED_ID, "owner-user", "owner-team") + assert recorder.spend_row_ids == [client_id] + assert recorder.standard_logging_ids == [client_id] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "hook, route", + [ + (_hook(), "/chat/completions"), + (_hook(general_settings={"disable_responses_id_security": True}), "/v1/responses"), + (_hook(signing_key=None), "/v1/responses"), + ], + ids=["not-a-create-route", "security-disabled", "no-signing-key"], + ) + async def test_stream_the_hook_leaves_alone_keeps_its_logged_id(self, salt_key_env, hook, route): + client_id, recorder = await _stream_logged_mid_stream(hook, route, _native_logged_result()) + + assert client_id == _LOGGED_MANAGED_ID + assert recorder.spend_row_ids == [_LOGGED_MANAGED_ID] + assert recorder.standard_logging_ids == [_LOGGED_MANAGED_ID] + + +class TestStreamedReadAnswersWithTheIdTheClientHolds: + @pytest.mark.asyncio + @pytest.mark.parametrize("route_prefix", ["/v1/responses", "/responses", "/openai/v1/responses"]) + async def test_the_owner_sees_their_id_while_callbacks_log_one_of_the_reads_own( + self, mock_cache, salt_key_env, route_prefix + ): + hook = _hook() + owner = _auth() + issued_id = _issue_managed_id(hook, owner, provider_response_id=_LOGGED_MANAGED_ID) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, recorder = await _stream_logged_mid_stream( + hook, f"{route_prefix}/{issued_id}", _native_logged_result(), caller=owner, request_data=data + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.standard_logging_ids + assert len(recorder.standard_logging_ids) == 1 + assert recorder.standard_logging_ids[0] != issued_id + assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == ( + _LOGGED_MANAGED_ID, + "owner-user", + "owner-team", + ) + + @pytest.mark.asyncio + async def test_an_admin_streamed_read_of_someone_elses_response_sees_the_id_sent_and_logs_one_of_its_own( + self, mock_cache, salt_key_env + ): + hook = _hook() + admin = _admin() + issued_id = _issue_managed_id(hook, _auth(), provider_response_id=_LOGGED_MANAGED_ID) + data = await hook.async_pre_call_hook( + user_api_key_dict=admin, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, recorder = await _stream_logged_mid_stream( + hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=admin, request_data=data + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.standard_logging_ids + assert len(recorder.standard_logging_ids) == 1 + assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == ( + _LOGGED_MANAGED_ID, + "admin-user", + "admin-team", + ) + + + @pytest.mark.asyncio + async def test_a_frame_naming_the_held_response_under_another_wrapping_still_gets_the_id_sent( + self, mock_cache, salt_key_env + ): + hook = _hook() + owner = _auth() + issued_id = _issue_managed_id(hook, owner, provider_response_id=_routed_response_id(_LOGGED_MANAGED_ID)) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, _ = await _stream_logged_mid_stream( + hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=owner, request_data=data + ) + + assert client_id == issued_id + + +async def _read_logged_by_callbacks( + hook: ResponsesIDSecurity, + caller: UserAPIKeyAuth, + issued_id: str, + call_type: CallTypesLiteral, + cache: DualCache, + answer: Callable[[str], ResponsesAPIResponse | DeleteResponseResult] = _provider_answer, +) -> tuple[str, _SpendRowRecorder]: + recorder = _SpendRowRecorder() + now = datetime.datetime.now() + logging_obj = LiteLLMLoggingObj( + model="gpt-5.1", + messages=[], + stream=False, + call_type=call_type, + start_time=now, + litellm_call_id=f"{call_type}-call", + function_id=f"{call_type}-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""} + ) + data = await hook.async_pre_call_hook( + user_api_key_dict=caller, cache=cache, data={"response_id": issued_id}, call_type=call_type + ) + provider_answer = answer(data["response_id"]) + + client_response = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=caller, response=provider_answer + ) + await logging_obj.async_success_handler(result=provider_answer, start_time=now, end_time=now, cache_hit=False) + return client_response.id, recorder + + +class TestEveryReadLogsAnIdOfItsOwn: + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + async def test_callbacks_and_the_spend_row_log_a_new_id_per_read_never_the_one_the_client_holds( + self, mock_cache, salt_key_env, call_type, reader + ): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth()) + + reads = [await _read_logged_by_callbacks(hook, reader, issued_id, call_type, mock_cache) for _ in range(2)] + + assert [client_id for client_id, _ in reads] == [issued_id, issued_id] + logged_ids = [id_ for _, recorder in reads for id_ in recorder.standard_logging_ids] + assert [recorder.spend_row_ids for _, recorder in reads] == [[logged_id] for logged_id in logged_ids] + assert len(set(logged_ids)) == 2 + assert issued_id not in logged_ids + assert {hook._decrypt_response_id(logged_id) for logged_id in logged_ids} == { + (_FABRICATED_PROVIDER_RESPONSE_ID, reader.user_id, reader.team_id) + } + + +_RAW_UPSTREAM_RESPONSE_ID = "resp_06de59f3upstreamiddddddddddddddddd" + + +def _routed_response_id(upstream_response_id: str) -> str: + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + _provider_answer(upstream_response_id), "openai", {"model_info": {"id": "deployment-1"}} + ).id + + +def _deletion_of(response_id: str) -> DeleteResponseResult: + return DeleteResponseResult( + id=ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id), + object="response.deleted", + deleted=True, + ) + + +class TestDeleteAnswersWithTheIdTheClientHolds: + @pytest.mark.asyncio + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + @pytest.mark.parametrize( + "stored_id", + [_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID], + ids=["routed-deployment", "plain-id"], + ) + async def test_a_caller_with_access_gets_back_the_exact_id_they_sent( + self, mock_cache, salt_key_env, reader, stored_id + ): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth(), provider_response_id=stored_id) + data = await hook.async_pre_call_hook( + user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type="adelete_responses" + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=reader, response=_deletion_of(data["response_id"]) + ) + + assert (result.id, result.object, result.deleted) == (issued_id, "response.deleted", True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + async def test_callbacks_and_the_spend_row_log_an_id_of_its_own_never_the_one_the_client_holds( + self, mock_cache, salt_key_env, reader + ): + hook = _hook() + issued_id = _issue_managed_id( + hook, _auth(), provider_response_id=_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID) + ) + + client_id, recorder = await _read_logged_by_callbacks( + hook, reader, issued_id, "adelete_responses", mock_cache, answer=_deletion_of + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.callback_response_ids + assert len(recorder.callback_response_ids) == 1 + assert recorder.callback_response_ids[0] != issued_id + assert hook._decrypt_response_id(recorder.callback_response_ids[0]) == ( + _RAW_UPSTREAM_RESPONSE_ID, + reader.user_id, + reader.team_id, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "sent_id", + [_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID], + ids=["routed-deployment", "plain-id"], + ) + async def test_an_admin_delete_by_a_raw_provider_id_gets_the_provider_id_back_and_logs_an_id_of_its_own( + self, mock_cache, salt_key_env, sent_id + ): + hook = _hook() + + client_id, recorder = await _read_logged_by_callbacks( + hook, _admin(), sent_id, "adelete_responses", mock_cache, answer=_deletion_of + ) + + assert client_id == _RAW_UPSTREAM_RESPONSE_ID + assert [hook._decrypt_response_id(logged_id) for logged_id in recorder.callback_response_ids] == [ + (_RAW_UPSTREAM_RESPONSE_ID, "admin-user", "admin-team") + ]