diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 89f4f3753b34..7caf74a1b58b 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -100,6 +100,7 @@ jobs: test-path: >- tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py + tests/proxy_unit_tests/test_deprecated_key_grace_period.py workers: 4 dist: loadscope timeout: 15 diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index f8e61d0166aa..2050c928cd8c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -557,6 +557,10 @@ def __init__( # Accumulate compaction blocks for multi-turn reconstruction self.compaction_blocks: List[Dict[str, Any]] = [] + # Accumulate streamed thinking text so final usage can split reasoning + # tokens from regular output tokens. + self.reasoning_content_chunks: List[str] = [] + # Track server tool use inputs and results for code_interpreter_results self._server_tool_inputs: Dict[str, Any] = {} self.tool_results: List[Dict[str, Any]] = [] @@ -587,9 +591,14 @@ def check_empty_tool_call_args(self) -> bool: return False def _handle_usage(self, anthropic_usage_chunk: Union[dict, UsageDelta]) -> Usage: + reasoning_content = ( + "".join(self.reasoning_content_chunks) + if self.reasoning_content_chunks + else None + ) return AnthropicConfig().calculate_usage( usage_object=cast(dict, anthropic_usage_chunk), - reasoning_content=None, + reasoning_content=reasoning_content, speed=self.speed, ) @@ -636,10 +645,13 @@ def _content_block_delta_helper(self, chunk: dict) -> Tuple[ "thinking" in content_block["delta"] or "signature" in content_block["delta"] ): + thinking_content = content_block["delta"].get("thinking") + if isinstance(thinking_content, str) and thinking_content: + self.reasoning_content_chunks.append(thinking_content) thinking_blocks = [ ChatCompletionThinkingBlock( type="thinking", - thinking=content_block["delta"].get("thinking") or "", + thinking=thinking_content or "", signature=str(content_block["delta"].get("signature") or ""), ) ] diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index b95bc8cb9b73..176bd6754545 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1856,8 +1856,16 @@ def calculate_usage( speed: Optional[str] = None, ) -> Usage: # NOTE: Sometimes the usage object has None set explicitly for token counts, meaning .get() & key access returns None, and we need to account for this - prompt_tokens = usage_object.get("input_tokens", 0) or 0 - completion_tokens = usage_object.get("output_tokens", 0) or 0 + raw_prompt_tokens = usage_object.get("input_tokens", 0) or 0 + prompt_tokens: int = ( + int(raw_prompt_tokens) if isinstance(raw_prompt_tokens, (int, float)) else 0 + ) + raw_completion_tokens = usage_object.get("output_tokens", 0) or 0 + completion_tokens: int = ( + int(raw_completion_tokens) + if isinstance(raw_completion_tokens, (int, float)) + else 0 + ) _usage = usage_object cache_creation_input_tokens: int = 0 cache_read_input_tokens: int = 0 @@ -1926,11 +1934,12 @@ def calculate_usage( text_tokens=raw_input_tokens, ) # Always populate completion_token_details, not just when there's reasoning_content - reasoning_tokens = ( + estimated_reasoning_tokens = ( token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) + reasoning_tokens = min(estimated_reasoning_tokens, completion_tokens) completion_token_details = CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else 0, text_tokens=( diff --git a/litellm/llms/azure/containers/transformation.py b/litellm/llms/azure/containers/transformation.py index 586b2e379a04..cd897511585f 100644 --- a/litellm/llms/azure/containers/transformation.py +++ b/litellm/llms/azure/containers/transformation.py @@ -1,9 +1,16 @@ from typing import Optional +from urllib.parse import parse_qs, urlparse, urlunparse from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.containers.transformation import OpenAIContainerConfig from litellm.types.router import GenericLiteLLMParams +# Endpoint-specific path suffixes that may appear in a deployment's api_base +# (e.g. the responses endpoint URL is stored as api_base for Azure models). +# Strip these before building the containers URL so we always start from the +# resource root (https://resource.cognitiveservices.azure.com). +_AZURE_ENDPOINT_PATHS = ("/openai/responses",) + class AzureContainerConfig(OpenAIContainerConfig): """ @@ -27,6 +34,27 @@ def validate_environment( litellm_params=GenericLiteLLMParams(api_key=api_key), ) + @staticmethod + def _normalize_api_base(api_base: Optional[str]) -> Optional[str]: + """Strip endpoint-specific path suffixes from api_base to get the resource root.""" + if not api_base: + return api_base + parsed = urlparse(api_base) + path = parsed.path.rstrip("/") + for ep in _AZURE_ENDPOINT_PATHS: + if path.endswith(ep): + return urlunparse( + (parsed.scheme, parsed.netloc, path[: -len(ep)], "", "", "") + ) + return api_base + + @staticmethod + def _extract_api_version(api_base: Optional[str]) -> Optional[str]: + """Return the api-version query param from api_base if present.""" + if not api_base: + return None + return parse_qs(urlparse(api_base).query).get("api-version", [None])[0] + def get_complete_url( self, api_base: Optional[str], @@ -39,10 +67,19 @@ def get_complete_url( {endpoint}/openai/v1/containers when api_version is 'v1', 'latest', or 'preview'; otherwise: {endpoint}/openai/containers + + The deployment's api_base may be the responses endpoint URL + (e.g. .../openai/responses?api-version=2025-04-01-preview). We + prefer the api-version embedded there over the deployment's + api_version field, which may point to an older chat API version. """ + effective_params = dict(litellm_params) + api_version_from_base = self._extract_api_version(api_base) + if api_version_from_base: + effective_params["api_version"] = api_version_from_base return BaseAzureLLM._get_base_azure_url( - api_base=api_base, - litellm_params=litellm_params, + api_base=self._normalize_api_base(api_base), + litellm_params=effective_params, route="/openai/containers", default_api_version="v1", ) diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 599cd705ebf4..501390d840b2 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -257,14 +257,19 @@ def _sync_handle( returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) + # An empty dict passed as `params` to httpx strips any existing query + # string from the URL (e.g. ?api-version=...). Use None instead so + # httpx leaves the URL's own query string intact. + effective_params = query_params or None + try: if method == "GET": response = http_client.get( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "DELETE": response = http_client.delete( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "POST": if is_multipart and "file" in kwargs: @@ -272,11 +277,11 @@ def _sync_handle( kwargs["file"], headers ) response = http_client.post( - url=url, headers=headers, params=query_params, files=files + url=url, headers=headers, params=effective_params, files=files ) else: response = http_client.post( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) else: raise ValueError(f"Unsupported HTTP method: {method}") @@ -376,14 +381,19 @@ async def _async_handle( returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) + # An empty dict passed as `params` to httpx strips any existing query + # string from the URL (e.g. ?api-version=...). Use None instead so + # httpx leaves the URL's own query string intact. + effective_params = query_params or None + try: if method == "GET": response = await http_client.get( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "DELETE": response = await http_client.delete( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "POST": if is_multipart and "file" in kwargs: @@ -391,11 +401,11 @@ async def _async_handle( kwargs["file"], headers ) response = await http_client.post( - url=url, headers=headers, params=query_params, files=files + url=url, headers=headers, params=effective_params, files=files ) else: response = await http_client.post( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) else: raise ValueError(f"Unsupported HTTP method: {method}") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2ffc7acbfb19..e6d2c437a340 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7816,7 +7816,7 @@ def container_list_handler( response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_list_response( @@ -7893,7 +7893,7 @@ async def async_container_list_handler( response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_list_response( @@ -7983,7 +7983,7 @@ def container_retrieve_handler( response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_retrieve_response( @@ -8060,7 +8060,7 @@ async def async_container_retrieve_handler( response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_retrieve_response( @@ -8150,7 +8150,7 @@ def container_delete_handler( response = sync_httpx_client.delete( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_delete_response( @@ -8227,7 +8227,7 @@ async def async_container_delete_handler( response = await async_httpx_client.delete( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_delete_response( @@ -8323,7 +8323,7 @@ def container_file_list_handler( response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_list_response( @@ -8402,7 +8402,7 @@ async def async_container_file_list_handler( response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_list_response( @@ -8490,7 +8490,7 @@ def container_file_content_handler( response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_content_response( @@ -8566,7 +8566,7 @@ async def async_container_file_content_handler( response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_content_response( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3a6d4e374e4b..6c3675ea06bc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2327,6 +2327,41 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): database_connection_timeout: Optional[float] = Field( 60, description="default timeout for a connection to the database" ) + database_connect_timeout: Optional[float] = Field( + None, + description=( + "Prisma `connect_timeout` URL param (seconds). Bounds how long the " + "engine waits to establish a new connection before failing. Defaults " + "to Prisma's built-in value when unset." + ), + ) + database_socket_timeout: Optional[float] = Field( + None, + description=( + "Prisma `socket_timeout` URL param (seconds). When set, an idle/slow " + "connection that has not produced data within this window is closed. " + "This is the main knob for capping idle DB connections from LiteLLM." + ), + ) + database_extra_connection_params: Optional[Dict[str, Any]] = Field( + None, + description=( + "Escape hatch: extra key/value pairs appended verbatim to the Prisma " + "DATABASE_URL / DIRECT_URL query string (e.g. `sslmode`, `pgbouncer`, " + "`statement_cache_size`). Keys here override any default LiteLLM sets." + ), + ) + database_disable_prepared_statements: Optional[bool] = Field( + None, + description=( + "Disable server-side prepared statements by setting Prisma's " + "`pgbouncer=true` URL param. Use this for pgbouncer transaction-pooling " + "deployments, or to prevent the 'cached plan must not change result " + "type' error that pooled connections hit during rolling schema " + "migrations. An explicit `pgbouncer` in `database_extra_connection_params` " + "takes precedence." + ), + ) database_type: Optional[Literal["dynamo_db"]] = Field( None, description="to use dynamodb instead of postgres db" ) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 5ded8136ef36..87dc2602d6e3 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -134,6 +134,16 @@ async def _handle_authentication_error( ) elif isinstance(e, ProxyException): raise e + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise ProxyException( + message=( + "Service Unavailable, the authentication database is " + "temporarily unreachable. Please retry shortly." + ), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) raise ProxyException( message="Authentication Error, " + str(e), type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 9650604bf81a..fc1f77bb684d 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -328,7 +328,7 @@ async def retrieve_container( custom_llm_provider=custom_llm_provider, ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id, original_container_id, custom_llm_provider, @@ -433,7 +433,7 @@ async def delete_container( custom_llm_provider=custom_llm_provider, ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id, original_container_id, custom_llm_provider, diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 4284cdd5d4a5..7eeb11fc3721 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -196,10 +196,12 @@ async def _process_binary_request( ) data: Dict[str, Any] = { "file_id": file_id, - **get_container_forwarding_params( - container_id=container_id, - original_container_id=original_container_id, - custom_llm_provider=resolved_provider, + **( + await get_container_forwarding_params( + container_id=container_id, + original_container_id=original_container_id, + custom_llm_provider=resolved_provider, + ) ), } processor = ProxyBaseLLMRequestProcessing(data=data) @@ -316,7 +318,7 @@ async def _process_multipart_upload_request( ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id=container_id, original_container_id=original_container_id, custom_llm_provider=resolved_provider, @@ -396,7 +398,7 @@ async def _process_request( ) ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id=path_params["container_id"], original_container_id=original_container_id, custom_llm_provider=resolved_provider, diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 568eca523ae7..57de6c4a63d6 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -23,6 +23,13 @@ _NEGATIVE_OWNER_SENTINEL = "__litellm_container_no_owner__" _CONTAINER_OWNER_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) +# Caches the stored ``unified_object_id`` (the encoded container ID +# captured at create time) so ``get_container_forwarding_params`` can +# recover the deployment ``model_id`` for native upstream IDs without +# re-hitting Prisma on every retrieve/delete. +_NEGATIVE_STORED_ID_SENTINEL = "__litellm_container_no_stored_id__" +_CONTAINER_STORED_ID_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) + # Per-caller-scope cache for ``GET /v1/containers`` list filtering. Without # this, every list call issues a fresh ``find_many`` against # ``litellm_managedobjecttable``. The cache key is the sorted owner-scope @@ -56,7 +63,7 @@ def decode_container_id_for_ownership( return original_container_id, custom_llm_provider -def get_container_forwarding_params( +async def get_container_forwarding_params( container_id: str, original_container_id: str, custom_llm_provider: str ) -> Dict[str, str]: params = { @@ -65,6 +72,20 @@ def get_container_forwarding_params( } decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) model_id = decoded.get("model_id") + if not (isinstance(model_id, str) and model_id): + # Native upstream IDs (e.g. Azure ``cntr_``) carry no LiteLLM + # routing payload, so decoding the user-supplied id yields no + # ``model_id``. Recover it from the encoded ``unified_object_id`` + # captured on the ownership row at create time — when the router + # selected a specific deployment that ID embeds the model_id. + stored_id = await _get_stored_container_id( + original_container_id, custom_llm_provider + ) + if stored_id and stored_id != container_id: + stored_decoded = ResponsesAPIRequestUtils._decode_container_id(stored_id) + stored_model_id = stored_decoded.get("model_id") + if isinstance(stored_model_id, str) and stored_model_id: + model_id = stored_model_id if isinstance(model_id, str) and model_id: params["model_id"] = model_id return params @@ -168,6 +189,7 @@ async def record_container_owner( ) _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner) + _CONTAINER_STORED_ID_CACHE.set_cache(model_object_id, container_id) # Drop the caller's own list-cache entry so the just-created container # shows up on their next ``GET /v1/containers``. Other callers with # disjoint scope tuples have their own entries; intersecting-scope @@ -207,9 +229,60 @@ async def _get_container_owner( _CONTAINER_OWNER_CACHE.set_cache( model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL ) + stored_id = getattr(row, "unified_object_id", None) if row is not None else None + _CONTAINER_STORED_ID_CACHE.set_cache( + model_object_id, + ( + stored_id + if isinstance(stored_id, str) and stored_id + else _NEGATIVE_STORED_ID_SENTINEL + ), + ) return owner +async def _get_stored_container_id( + original_container_id: str, custom_llm_provider: str +) -> Optional[str]: + """Return the ``unified_object_id`` stored at create time, if any. + + Used by :func:`get_container_forwarding_params` to recover the + deployment ``model_id`` for native upstream container IDs: the stored + value is the encoded form produced by ``encode_container_id_in_response`` + when the router selected a specific deployment. + """ + model_object_id = _container_model_object_id( + original_container_id, custom_llm_provider + ) + + cached = _CONTAINER_STORED_ID_CACHE.get_cache(model_object_id) + if cached == _NEGATIVE_STORED_ID_SENTINEL: + return None + if isinstance(cached, str) and cached: + return cached + + prisma_client = await _get_prisma_client() + if prisma_client is None: + return None + + row = await prisma_client.db.litellm_managedobjecttable.find_first( + where={ + "model_object_id": model_object_id, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + } + ) + stored_id = getattr(row, "unified_object_id", None) if row is not None else None + _CONTAINER_STORED_ID_CACHE.set_cache( + model_object_id, + ( + stored_id + if isinstance(stored_id, str) and stored_id + else _NEGATIVE_STORED_ID_SENTINEL + ), + ) + return stored_id if isinstance(stored_id, str) and stored_id else None + + async def assert_user_can_access_container( container_id: str, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index ab9d341aa510..c500e727595e 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -109,6 +109,92 @@ def is_database_transport_error(e: Exception) -> bool: return True return False + @staticmethod + def is_prisma_engine_internal_error(e: Exception) -> bool: + """True iff ``e`` is a non-``PrismaError`` exception raised from inside + prisma-client-py's query-engine layer. + + During the instant a DB connection is torn down, the query engine can + return a malformed error payload (``user_facing_error.meta`` is + ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` + before it can raise the proper P1001 "can't reach database server" + error. That AttributeError carries no connection keyword, so it can't + be matched by message; identify it by its ``prisma.engine`` origin + instead. + + Recognized ``PrismaError`` subclasses are excluded: connectivity ones + are already classified by type/keyword above, and data-layer ones + (the DB IS reachable) must stay 401. + """ + import prisma + + if isinstance(e, prisma.errors.PrismaError): + return False + tb = getattr(e, "__traceback__", None) + while tb is not None: + if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): + return True + tb = tb.tb_next + return False + + @staticmethod + def is_database_service_unavailable_error(e: Exception) -> bool: + """True iff the exception means the database could not answer at the + infrastructure level (connection refused, socket/interface failure, + timeout) rather than a genuine auth failure (key not found) or a + data-layer error (the DB IS reachable and rejected the data). + + Auth must answer 401 only for a key the DB confirms is invalid. When + the DB itself is unreachable, the request has to surface as 503 so + callers retry instead of treating valid keys as invalid during an + outage. + + Note: prisma-client-py mislabels the P1001 "can't reach database + server" connectivity failure as a ``DataError`` (a data-layer type), + so a type-only check misses real outages. ``is_database_transport_error`` + keyword-matches the connection message and catches that masquerade, + while genuine data errors (no connection keyword) correctly stay 401. + + The Postgres "cached plan must not change result type" error is matched + here, not in ``is_database_transport_error``: it is a transient stale-DB- + state condition (not an invalid key), but the connection is healthy so it + must not trigger a reconnect. + + A non-``PrismaError`` raised from inside the prisma query engine (e.g. + the ``AttributeError`` from ``handle_response_errors`` when the engine + returns a malformed error payload mid-tear-down) is also treated as + unavailable; see ``is_prisma_engine_internal_error``. + """ + import asyncio + + if PrismaDBExceptionHandler.is_database_connection_error(e): + return True + if PrismaDBExceptionHandler.is_database_transport_error(e): + return True + if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): + return True + if "cached plan must not change result type" in str(e).lower(): + return True + + # OSError already covers ConnectionError and (Py3.3+) TimeoutError. + # asyncio.TimeoutError is a distinct class before Py3.11. + if isinstance(e, (OSError, asyncio.TimeoutError)): + return True + + try: + import asyncpg + except ImportError: + return False + + return isinstance( + e, + ( + asyncpg.exceptions.PostgresConnectionError, + asyncpg.exceptions.InterfaceError, + ), + ) + @staticmethod def handle_db_exception(e: Exception): """ diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 82260eb2ae5e..6fd62e1a6ff3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -100,6 +100,42 @@ def _get_user_from_metadata( return get_end_user_id_from_request_body(request_body) return None + @staticmethod + def _resolve_costing_model(model: str, logging_obj: LiteLLMLoggingObj) -> str: + if model and model != "unknown": + return model + litellm_params = (getattr(logging_obj, "model_call_details", {}) or {}).get( + "litellm_params", {} + ) or {} + deployment_model = litellm_params.get("model") + if deployment_model and deployment_model != "unknown": + return deployment_model + model_group = (litellm_params.get("metadata", {}) or {}).get("model_group") + if model_group: + return model_group.removeprefix("passthrough/") + return model + + @staticmethod + def _extract_model_from_anthropic_chunks( + all_chunks: Sequence[Union[str, bytes]], + ) -> Optional[str]: + for raw in all_chunks: + text = raw.decode("utf-8") if isinstance(raw, bytes) else raw + for line in text.splitlines(): + if not line.startswith("data:"): + continue + try: + data = json.loads(line[len("data:") :].strip()) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(data, dict): + continue + if data.get("type") == "message_start": + model = (data.get("message") or {}).get("model") + if model: + return model + return None + @staticmethod def _create_anthropic_response_logging_payload( litellm_model_response: Union[ModelResponse, TextCompletionResponse], @@ -127,6 +163,10 @@ def _create_anthropic_response_logging_payload( "custom_llm_provider" ) + model = AnthropicPassthroughLoggingHandler._resolve_costing_model( + model, logging_obj + ) + # Prepend custom_llm_provider to model if not already present model_for_cost = model if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"): @@ -213,6 +253,15 @@ def _handle_logging_anthropic_collected_chunks( ): model = cast(str, litellm_logging_obj.model_call_details.get("model")) + if not model or model == "unknown": + chunk_model = ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( + all_chunks + ) + ) + if chunk_model: + model = chunk_model + complete_streaming_response = ( AnthropicPassthroughLoggingHandler._build_complete_streaming_response( all_chunks=all_chunks, @@ -301,6 +350,13 @@ def _build_complete_streaming_response( # Process each individual event for event_str in individual_events: try: + # Skip OpenAI-style [DONE] sentinels some Anthropic-compatible + # providers emit. Match the whole SSE line so a valid chunk whose + # text payload happens to contain "[DONE]" is not dropped. + if any( + line.strip() == "data: [DONE]" for line in event_str.split("\n") + ): + continue transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( chunk=event_str ) @@ -309,6 +365,14 @@ def _build_complete_streaming_response( except (StopIteration, StopAsyncIteration): break + except json.JSONDecodeError: + # Some upstreams emit non-JSON SSE lines; skip them so the + # logging pipeline is not broken by a single bad frame. + verbose_proxy_logger.debug( + "Skipping non-JSON SSE event: %s", + event_str[:200], + ) + continue complete_streaming_response = litellm.stream_chunk_builder( chunks=all_openai_chunks, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 71aeea67884a..ba2dc6971096 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -38,6 +38,41 @@ class LiteLLMDatabaseConnectionPool(Enum): database_connection_pool_timeout = 60 +def _build_db_connection_url_params( + connection_limit: int, + pool_timeout: Optional[Union[int, float]], + connect_timeout: Optional[Union[int, float]] = None, + socket_timeout: Optional[Union[int, float]] = None, + disable_prepared_statements: bool = False, + extra_params: Optional[dict] = None, +) -> dict: + """Build the Prisma DATABASE_URL query params controlling connection pool behavior. + + `connect_timeout` / `socket_timeout` map to the Prisma URL params of the same + name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are + omitted when None so Prisma's defaults apply. `disable_prepared_statements` + sets `pgbouncer=true`, which makes Prisma stop using server-side prepared + statements (pgbouncer transaction-pool compatible; also sidesteps the + "cached plan must not change result type" error during rolling migrations). + `extra_params` is an untyped passthrough — keys it provides win over the + named arguments above, so it can be used to override any default we set here. + """ + params: dict = { + "connection_limit": connection_limit, + } + if pool_timeout is not None: + params["pool_timeout"] = pool_timeout + if connect_timeout is not None: + params["connect_timeout"] = connect_timeout + if socket_timeout is not None: + params["socket_timeout"] = socket_timeout + if disable_prepared_statements: + params["pgbouncer"] = "true" + if extra_params: + params.update(extra_params) + return params + + def append_query_params(url: Optional[str], params: dict) -> str: from litellm._logging import verbose_proxy_logger @@ -745,7 +780,12 @@ def run_server( # noqa: PLR0915 ) db_connection_pool_limit = 100 - db_connection_timeout = 60 + # Starts optional due to config fallback checks; guaranteed non-None before use. + db_connection_timeout: Optional[Union[int, float]] = 60 + db_connect_timeout: Optional[Union[int, float]] = None + db_socket_timeout: Optional[Union[int, float]] = None + db_disable_prepared_statements: bool = False + db_extra_connection_params: Optional[dict] = None general_settings = {} ### GET DB TOKEN FOR IAM AUTH ### @@ -849,9 +889,30 @@ def run_server( # noqa: PLR0915 "database_connection_pool_limit", LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value, ) - db_connection_timeout = general_settings.get( - "database_connection_pool_timeout", - LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value, + db_connection_timeout = general_settings.get("database_connection_timeout") + if db_connection_timeout is None: + db_connection_timeout = general_settings.get( + "database_connection_pool_timeout" + ) + if db_connection_timeout is None: + db_connection_timeout = ( + LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value + ) + db_connect_timeout = general_settings.get("database_connect_timeout") + db_socket_timeout = general_settings.get("database_socket_timeout") + _disable_prepared_statements = general_settings.get( + "database_disable_prepared_statements", False + ) + if isinstance(_disable_prepared_statements, str): + from litellm.secret_managers.main import str_to_bool + + db_disable_prepared_statements = ( + str_to_bool(_disable_prepared_statements) is True + ) + else: + db_disable_prepared_statements = bool(_disable_prepared_statements) + db_extra_connection_params = general_settings.get( + "database_extra_connection_params" ) if database_url and database_url.startswith("os.environ/"): original_dir = os.getcwd() @@ -892,27 +953,27 @@ def run_server( # noqa: PLR0915 try: from litellm.secret_managers.main import get_secret + connection_url_params = _build_db_connection_url_params( + connection_limit=db_connection_pool_limit, + pool_timeout=db_connection_timeout, + connect_timeout=db_connect_timeout, + socket_timeout=db_socket_timeout, + disable_prepared_statements=db_disable_prepared_statements, + extra_params=db_extra_connection_params, + ) if os.getenv("DATABASE_URL", None) is not None: - ### add connection pool + pool timeout args - params = { - "connection_limit": db_connection_pool_limit, - "pool_timeout": db_connection_timeout, - } database_url = get_secret("DATABASE_URL", default_value=None) modified_url = append_query_params( - str(database_url) if database_url else None, params + str(database_url) if database_url else None, + connection_url_params, ) os.environ["DATABASE_URL"] = modified_url if os.getenv("DIRECT_URL", None) is not None: - ### add connection pool + pool timeout args - params = { - "connection_limit": db_connection_pool_limit, - "pool_timeout": db_connection_timeout, - } database_url = os.getenv("DIRECT_URL") - modified_url = append_query_params(database_url, params) + modified_url = append_query_params( + database_url, connection_url_params + ) os.environ["DIRECT_URL"] = modified_url - ### subprocess.run(["prisma"], capture_output=True) is_prisma_runnable = True except FileNotFoundError: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e7f5f4ee3964..9e7b4c74ffad 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2395,7 +2395,8 @@ def jsonify_object(data: dict) -> dict: return db_data -# In-memory cache for deprecated key lookups: maps old_token_hash -> (active_token_id, expires_at_ts) +# In-memory cache for deprecated key lookups: +# maps old_token_hash -> (active_token_id, cache_expires_at_ts, revoke_at_ts). # Avoids a DB query on every auth request for non-deprecated keys. # Bounded to prevent memory leaks from accumulated rotations. _deprecated_key_cache: LimitedSizeOrderedDict = LimitedSizeOrderedDict(max_size=1000) @@ -2417,26 +2418,25 @@ async def _lookup_deprecated_key( # Check cache first cached = _deprecated_key_cache.get(hashed_token) - cached = _deprecated_key_cache.get(hashed_token) if cached is not None: active_token_id, cache_expires_at_ts, revoke_at_ts = cached if now_ts < cache_expires_at_ts and now_ts < revoke_at_ts: return active_token_id - else: - _deprecated_key_cache.pop(hashed_token, None) + _deprecated_key_cache.pop(hashed_token, None) try: deprecated_row = await db.litellm_deprecatedverificationtoken.find_first( where={ "token": hashed_token, "revoke_at": {"gt": now}, - }, - select={"active_token_id": True}, + } ) if deprecated_row and deprecated_row.active_token_id: + revoke_at = deprecated_row.revoke_at _deprecated_key_cache[hashed_token] = ( deprecated_row.active_token_id, now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, + revoke_at.timestamp(), ) return deprecated_row.active_token_id # Only cache positive results; negative lookups are fast on indexed columns @@ -2841,40 +2841,49 @@ async def _query_first_with_cached_plan_fallback( self, sql_query: str, *args ) -> Optional[dict]: """ - Execute a query with automatic fallback for PostgreSQL cached plan errors. - - This handles the "cached plan must not change result type" error that occurs - during rolling deployments when schema changes are applied while old pods - still have cached query plans expecting the old schema. - - Args: - sql_query: SQL query string to execute - - Returns: - Query result or None - - Raises: - Original exception if not a cached plan error + Execute a query, recovering once from PostgreSQL's "cached plan must not + change result type" error. + + That error surfaces during rolling deployments when a schema change + invalidates the prepared-statement plans that pooled connections still + hold. Clearing only the server-side plans with DEALLOCATE ALL makes + things worse: Prisma's query engine keeps a per-connection client-side + cache of prepared-statement names, so once the server drops a plan the + engine re-sends a name PostgreSQL no longer recognizes and the + connection breaks with `prepared statement "sN" does not exist`. With a + small pool that connection stays poisoned and every auth lookup fails. + + Recreating the Prisma client kills the engine subprocess and drops the + server-side plans and the engine's client-side name cache together, so + the retried query is prepared fresh. We reconnect through + `attempt_db_reconnect`, which is singleflight: when a schema change + poisons every pooled connection at once, the first cached-plan error + recreates the client and the concurrent waiters reuse that single + recreate instead of racing to kill each other's fresh engine. We then + retry the identical query exactly once. + + The retry reuses the original query byte-for-byte. Mutating the SQL + (e.g. injecting a unique comment) would defeat PostgreSQL's plan cache, + forcing a fresh plan on every request and pegging the database CPU. + + If the reconnect is skipped because a recent reconnect is still within + its cooldown, the retry runs against the same connection and may fail + again; the get_data backoff decorator re-runs the lookup and a later + attempt reconnects once the cooldown elapses. """ try: return await self.db.query_first(sql_query, *args) except Exception as e: - error_str = str(e) - if "cached plan must not change result type" in error_str: - # Force PostgreSQL to re-plan by invalidating the cache - # Add a unique comment to make the query different - sql_query_retry = sql_query.replace( - "SELECT", - f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */", - ) - verbose_proxy_logger.warning( - "PostgreSQL cached plan error detected for token lookup, " - "retrying with fresh plan. This may occur during rolling deployments " - "when schema changes are applied." - ) - return await self.db.query_first(sql_query_retry, *args) - else: + if "cached plan must not change result type" not in str(e): raise + verbose_proxy_logger.warning( + "PostgreSQL cached plan error detected for token lookup; " + "recreating the database connection and retrying with the same " + "query. This may occur during rolling deployments when schema " + "changes are applied." + ) + await self.attempt_db_reconnect(reason="postgres_cached_plan_error") + return await self.db.query_first(sql_query, *args) @backoff.on_exception( backoff.expo, @@ -3230,7 +3239,10 @@ async def get_data( # noqa: PLR0915 db=self.db, hashed_token=hashed_token ) if active_token_id: - response = await self.get_data( + # The recursive call returns a finished + # LiteLLM_VerificationTokenView; the dict + # normalization below would crash subscripting it. + deprecated_response = await self.get_data( token=active_token_id, table_name="combined_view", query_type="find_unique", @@ -3238,10 +3250,11 @@ async def get_data( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, check_deprecated=False, ) - if response is not None: + if deprecated_response is not None: verbose_proxy_logger.debug( "Deprecated key used during grace period" ) + return deprecated_response if response is not None: if response["team_models"] is None: diff --git a/litellm/router.py b/litellm/router.py index 7512ee387dca..4e8c2152f371 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5551,6 +5551,7 @@ async def _init_containers_api_endpoints( from litellm.responses.utils import ResponsesAPIRequestUtils container_id = kwargs.get("container_id") + _forwarded_model_id = kwargs.get("model_id") if isinstance(container_id, str): decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) original_id = decoded.get("response_id", container_id) @@ -5559,7 +5560,14 @@ async def _init_containers_api_endpoints( decoded_provider = decoded.get("custom_llm_provider") if decoded_provider and kwargs.get("custom_llm_provider") == "openai": kwargs["custom_llm_provider"] = decoded_provider - model_id = decoded.get("model_id") + # Fall back to the model_id forwarded by the proxy when the container_id + # is a native upstream ID (e.g. Azure hex cntr_) that carries no LiteLLM + # routing payload, so deployment credentials (api_base, api_key) are applied. + model_id = decoded.get("model_id") or ( + _forwarded_model_id.strip() + if isinstance(_forwarded_model_id, str) and _forwarded_model_id.strip() + else None + ) if model_id: kwargs["model"] = model_id return await self._ageneric_api_call_with_fallbacks( diff --git a/pyproject.toml b/pyproject.toml index 4d2483894dc9..c7e205e6f3b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.84.7" +version = "1.84.8" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -242,7 +242,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.84.7" +version = "1.84.8" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/proxy_unit_tests/test_deprecated_key_grace_period.py b/tests/proxy_unit_tests/test_deprecated_key_grace_period.py new file mode 100644 index 000000000000..a91ecf95f329 --- /dev/null +++ b/tests/proxy_unit_tests/test_deprecated_key_grace_period.py @@ -0,0 +1,177 @@ +""" +Tests for the grace-period key-rotation feature (MLI-6358). + +Two bugs are confirmed in LiteLLM v1.83.7-stable (upstream BerriAI/litellm#27193). +Both live in _lookup_deprecated_key() (litellm/proxy/utils.py): + + Bug 1 — duplicate cache read (cosmetic, no functional impact on its own): + The cache is fetched twice in a row with no state change between the calls. + + Bug 2 — cache stores a 2-tuple but unpacks as a 3-tuple: + WRITE: _deprecated_key_cache[hash] = (active_token_id, cache_expires_at_ts) + READ: active_token_id, cache_expires_at_ts, revoke_at_ts = cached # ValueError! + The ValueError is NOT inside the try/except, so it propagates up through + PrismaClient.get_data() (which re-raises), killing the auth request. + +The local demo script confirmed +that all three requests with the old key returned HTTP 401 immediately after +rotation even though the grace-period window was still open. +""" + +from datetime import datetime, timedelta, timezone +from typing import Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest + + +# ── helpers ─────────────────────────────────────────────────────────────────── + +HASHED_TOKEN = "165efe575c98fe7e65d98cb2de71b68842049e286afd33a92d3491c340216880" +ACTIVE_TOKEN_HASH = "abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890ab" + + +def _make_db(active_token_id: Optional[str]) -> MagicMock: + """Prisma db mock whose deprecated-token find_first returns the given id.""" + row = MagicMock() + row.active_token_id = active_token_id + row.revoke_at = datetime.now(timezone.utc) + timedelta(minutes=5) + db = MagicMock() + db.litellm_deprecatedverificationtoken = MagicMock() + db.litellm_deprecatedverificationtoken.find_first = AsyncMock( + return_value=row if active_token_id else None + ) + return db + + +# ── Bug 1: first call (DB path) ─────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_db_miss_returns_none(): + """Token absent from deprecated table → returns None without error.""" + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=None) + + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + + assert result is None + db.litellm_deprecatedverificationtoken.find_first.assert_called_once() + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_db_hit_returns_active_token_id(): + """ + First call (cold cache): DB row exists within grace window → returns + active_token_id correctly. The DB path itself works; the bug is on the + second call when the result is read back from cache. + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + + assert result == ACTIVE_TOKEN_HASH + db.litellm_deprecatedverificationtoken.find_first.assert_called_once() + + +# ── Bug 2: second call (cache path) ────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_cache_hit_returns_on_second_call(): + """ + Regression guard: after first call warms the cache with a 3-tuple, + second call should return from cache without raising. + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + # First call: cold cache → DB hit → warms cache with 3-tuple → succeeds + r1 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r1 == ACTIVE_TOKEN_HASH, "First call (DB path) must succeed" + + # Second call: cache hit path should succeed without DB access + r2 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r2 == ACTIVE_TOKEN_HASH + + # DB is queried exactly once; the second call never reaches it + assert db.litellm_deprecatedverificationtoken.find_first.call_count == 1 + + +@pytest.mark.asyncio +async def test_lookup_deprecated_key_pre_warmed_cache_returns(): + """ + Pre-warmed 3-tuple cache entry should be served directly from cache. + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + now_ts = datetime.now(timezone.utc).timestamp() + _deprecated_key_cache[HASHED_TOKEN] = ( + ACTIVE_TOKEN_HASH, + now_ts + 60, + now_ts + 300, + ) + + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert result == ACTIVE_TOKEN_HASH + + db.litellm_deprecatedverificationtoken.find_first.assert_not_called() + + +# ── End-to-end reproduction of the demo ────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_grace_period_three_requests_mirrors_demo(): + """ + Reproduces Step 5 of the local demo script: + + Request 1 (cache miss — DB lookup) → succeeds + Request 2 (cache hit) → succeeds + Request 3 (cache hit) → succeeds + """ + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + db = _make_db(active_token_id=ACTIVE_TOKEN_HASH) + + r1 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r1 == ACTIVE_TOKEN_HASH, "Request 1 (DB path) should succeed" + + r2 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + r3 = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert r2 == ACTIVE_TOKEN_HASH + assert r3 == ACTIVE_TOKEN_HASH + + # DB hit only once; requests 2 and 3 never reach it + assert db.litellm_deprecatedverificationtoken.find_first.call_count == 1 + + +@pytest.mark.asyncio +async def test_cache_hit_respects_revoke_at_timestamp(): + """Cache entries should not remain valid past revoke_at even if cache TTL is still live.""" + from litellm.proxy.utils import _lookup_deprecated_key, _deprecated_key_cache + + _deprecated_key_cache.clear() + now_ts = datetime.now(timezone.utc).timestamp() + # cache_expires_at is in the future, but revoke_at is already past. + _deprecated_key_cache[HASHED_TOKEN] = ( + ACTIVE_TOKEN_HASH, + now_ts + 60, + now_ts - 1, + ) + + db = _make_db(active_token_id=None) + result = await _lookup_deprecated_key(db=db, hashed_token=HASHED_TOKEN) + assert result is None + db.litellm_deprecatedverificationtoken.find_first.assert_called_once() diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/test_litellm/containers/test_azure_container_transformation.py index 70181f6f03df..cdcccf7c04e7 100644 --- a/tests/test_litellm/containers/test_azure_container_transformation.py +++ b/tests/test_litellm/containers/test_azure_container_transformation.py @@ -109,6 +109,31 @@ def test_get_complete_url_with_latest_api_version(self): assert "/openai/v1/containers" in url + def test_get_complete_url_strips_responses_path_and_preserves_api_version(self): + """When api_base is the responses endpoint URL, get_complete_url must: + - strip /openai/responses (no double-path) + - use the api-version from api_base query string, NOT the deployment's + older api_version (e.g. 2024-08-01-preview → containers need 2025-04-01-preview) + """ + api_base = "https://my-resource.cognitiveservices.azure.com/openai/responses?api-version=2025-04-01-preview" + + url = self.config.get_complete_url( + api_base=api_base, + litellm_params={"api_version": "2024-08-01-preview"}, + ) + + assert ( + "/openai/responses/openai/containers" not in url + ), "path must not double /openai/responses" + assert "my-resource.cognitiveservices.azure.com" in url + assert "/openai/containers" in url or "/openai/v1/containers" in url + assert ( + "2025-04-01-preview" in url + ), "must use version from api_base, not litellm_params" + assert ( + "2024-08-01-preview" not in url + ), "must not fall back to older chat api_version" + def test_get_complete_url_raises_without_api_base(self, monkeypatch): monkeypatch.delenv("AZURE_API_BASE", raising=False) monkeypatch.setattr(litellm, "api_base", None) @@ -531,6 +556,92 @@ def test_regression_api_base_with_extra_query_params(self): assert qs.get("api-version") == ["v1"] assert qs.get("foo") == ["bar"] + @pytest.mark.asyncio + async def test_regression_no_container_id_does_not_use_user_supplied_model_id( + self, monkeypatch + ): + """Operations without container_id (create, list) must NOT route via + _ageneric_api_call_with_fallbacks using a caller-supplied model_id. + + Security boundary: only the path that holds a validated container_id + is trusted to fall back to the forwarded model_id. A caller setting + model_id without container_id on POST /v1/containers must not gain + access to an arbitrary deployment UUID. + """ + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "azure-model", + "litellm_params": { + "model": "azure/gpt-4", + "api_base": "https://my-resource.cognitiveservices.azure.com", + "api_key": "test-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "deployment-uuid-123"}, + } + ] + ) + + fallback_called = {"called": False} + + async def _mock_fallback(original_function, **kwargs): + fallback_called["called"] = True + return {} + + monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback) + + original_called = {"called": False} + + async def _noop(**kwargs): + original_called["called"] = True + return {} + + # No container_id — simulates create/list; caller injects a model_id + await router._init_containers_api_endpoints( + original_function=_noop, + model_id="deployment-uuid-123", + custom_llm_provider="azure", + ) + + assert not fallback_called["called"], ( + "_ageneric_api_call_with_fallbacks must NOT be called when " + "container_id is absent, even if model_id is supplied" + ) + assert original_called["called"], "original_function must be called directly" + + def test_regression_httpx_empty_params_strips_query_string(self): + """httpx erases the URL query-string when params={} (empty dict) is passed. + + Root cause of the Azure container 404s on POST/DELETE: + _build_query_params returns {} when the endpoint has no extra params; + passing that {} as params= to httpx wiped ?api-version=2025-04-01-preview. + + Fix: every container httpx call now uses `params or None` so an empty + dict falls back to None, which tells httpx to leave the URL untouched. + """ + url = ( + "https://resource.cognitiveservices.azure.com" + "/openai/containers/cntr_123?api-version=2025-04-01-preview" + ) + client = httpx.AsyncClient() + + req_none = client.build_request("DELETE", url, params=None) + assert "api-version=2025-04-01-preview" in str(req_none.url) + + req_empty = client.build_request("DELETE", url, params={}) + assert "api-version" not in str( + req_empty.url + ), "Documents root cause: params={} strips the query string" + + effective: dict = {} + req_guarded = client.build_request("DELETE", url, params=effective or None) + assert "api-version=2025-04-01-preview" in str( + req_guarded.url + ), "`params or None` must preserve ?api-version" + def test_regression_proxy_resolves_azure_text_same_as_azure(self): """Router/proxy treat azure_text like azure for container config.""" from litellm.proxy.container_endpoints.handler_factory import ( @@ -770,3 +881,143 @@ async def _mock_base_process_llm_request( assert captured["data"]["container_id"] == "cntr_123" assert captured["data"]["custom_llm_provider"] == "azure" assert captured["data"]["model_id"] == "model_abc123" + + @pytest.mark.asyncio + async def test_regression_get_container_forwarding_params_sets_model_id_for_managed_id( + self, + ): + """get_container_forwarding_params must extract model_id from a + LiteLLM-managed encoded container ID and include it in the forwarding + dict. This is the proxy-side half of the native-Azure-ID routing fix: + the router's _init_containers_api_endpoints reads kwargs["model_id"] + which is set here. + """ + from litellm.proxy.container_endpoints.ownership import ( + get_container_forwarding_params, + ) + + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="deployment-uuid-123", + container_id="cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df", + ) + + params = await get_container_forwarding_params( + container_id=encoded_id, + original_container_id="cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df", + custom_llm_provider="azure", + ) + + assert ( + params.get("model_id") == "deployment-uuid-123" + ), "model_id must be forwarded to the router for managed container IDs" + assert params.get("container_id") == ( + "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" + ) + assert params.get("custom_llm_provider") == "azure" + + @pytest.mark.asyncio + async def test_regression_get_container_forwarding_params_recovers_model_id_for_native_id( + self, monkeypatch + ): + """Native Azure IDs (``cntr_``) cannot be decoded, so model_id + must be recovered from the ownership row's ``unified_object_id`` — + the encoded form captured at create time when the router selected a + specific deployment. Without this, the router-side fallback for + native IDs in ``_init_containers_api_endpoints`` is dead code. + """ + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.container_endpoints import ownership + from litellm.proxy.container_endpoints.ownership import ( + get_container_forwarding_params, + ) + + native_id = "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" + encoded_stored_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="deployment-uuid-123", + container_id=native_id, + ) + + ownership._CONTAINER_STORED_ID_CACHE.flush_cache() + ownership._CONTAINER_OWNER_CACHE.flush_cache() + + table = AsyncMock() + table.find_first.return_value = SimpleNamespace( + created_by="user-1", + file_purpose=ownership.CONTAINER_OBJECT_PURPOSE, + unified_object_id=encoded_stored_id, + ) + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + params = await get_container_forwarding_params( + container_id=native_id, + original_container_id=native_id, + custom_llm_provider="azure", + ) + + assert params.get("model_id") == "deployment-uuid-123", ( + "model_id must be recovered from the stored unified_object_id " + "for native upstream container IDs" + ) + assert params.get("container_id") == native_id + assert params.get("custom_llm_provider") == "azure" + + @pytest.mark.asyncio + async def test_regression_native_azure_container_id_uses_forwarded_model_id( + self, monkeypatch + ): + """Native Azure container IDs (cntr_ + hex, no LiteLLM payload) must + still route through _ageneric_api_call_with_fallbacks using the + model_id forwarded from the proxy ownership check so that deployment + credentials (api_base) are applied.""" + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "azure-model", + "litellm_params": { + "model": "azure/gpt-4", + "api_base": "https://my-resource.cognitiveservices.azure.com", + "api_key": "test-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "deployment-uuid-123"}, + } + ] + ) + + called_with: dict = {} + + async def _mock_fallback(original_function, **kwargs): + called_with.update(kwargs) + return {} + + monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback) + + native_azure_id = "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" + + async def _noop(**kwargs): + return {} + + await router._init_containers_api_endpoints( + original_function=_noop, + container_id=native_azure_id, + model_id="deployment-uuid-123", + custom_llm_provider="azure", + ) + + assert called_with.get("model") == "deployment-uuid-123", ( + "_ageneric_api_call_with_fallbacks must be called with the forwarded " + "model_id when the container_id carries no LiteLLM routing payload" + ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index bf0461d89f1a..2fdd639e74d4 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -1,7 +1,11 @@ +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import AsyncMock, MagicMock import pytest +import litellm from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call from litellm.types.llms.openai import ( @@ -343,6 +347,289 @@ def test_text_only_streaming_has_index_zero(): ), f"Expected index=0, got {parsed.choices[0].index}" +def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage(): + """Anthropic streaming usage should account for emitted thinking deltas.""" + chunks = [ + { + "type": "message_start", + "message": { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "thinking_delta", + "thinking": "First I need to count the favorable outcomes. ", + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "thinking_delta", + "thinking": "Then I compare that count with all possible outcomes.", + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "sig_123"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "The probability is 3/8."}, + }, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 50}, + }, + ] + + iterator = ModelResponseIterator(None, sync_stream=True) + final_usage = None + reasoning_deltas = [] + + for chunk in chunks: + parsed = iterator.chunk_parser(chunk) + reasoning_content = getattr(parsed.choices[0].delta, "reasoning_content", None) + if reasoning_content: + reasoning_deltas.append(reasoning_content) + if parsed.usage is not None: + final_usage = parsed.usage + + assert reasoning_deltas == [ + "First I need to count the favorable outcomes. ", + "Then I compare that count with all possible outcomes.", + ] + assert final_usage is not None + completion_tokens_details = final_usage.completion_tokens_details + assert completion_tokens_details is not None + assert completion_tokens_details.reasoning_tokens > 0 + assert completion_tokens_details.text_tokens == ( + final_usage.completion_tokens - completion_tokens_details.reasoning_tokens + ) + + +def test_anthropic_completion_streaming_usage_matches_non_streaming_with_thinking(): + """The completion API should preserve Anthropic thinking usage in streaming mode.""" + thinking_parts = [ + "First I need to count the favorable outcomes. ", + "Then I compare that count with all possible outcomes.", + ] + thinking_text = "".join(thinking_parts) + answer_text = "The probability is 3/8." + requests_seen = [] + + class MockAnthropicHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, format, *args): # type: ignore[no-untyped-def] + return + + def do_POST(self): # type: ignore[no-untyped-def] + content_length = int(self.headers.get("content-length", "0")) + payload = json.loads(self.rfile.read(content_length).decode("utf-8")) + requests_seen.append( + { + "path": self.path, + "model": payload.get("model"), + "stream": payload.get("stream", False), + "thinking": payload.get("thinking"), + } + ) + + if payload.get("stream"): + events = [ + { + "type": "message_start", + "message": { + "id": "msg_mock", + "type": "message", + "role": "assistant", + "model": payload.get("model"), + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "thinking_delta", + "thinking": thinking_parts[0], + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "thinking_delta", + "thinking": thinking_parts[1], + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "signature_delta", + "signature": "sig_mock", + }, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": answer_text}, + }, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 50}, + }, + {"type": "message_stop"}, + ] + self._write_response( + content_type="text/event-stream", + body="".join( + f"data: {json.dumps(event)}\n\n" for event in events + ).encode("utf-8"), + ) + return + + self._write_response( + content_type="application/json", + body=json.dumps( + { + "id": "msg_mock", + "type": "message", + "role": "assistant", + "model": payload.get("model"), + "content": [ + { + "type": "thinking", + "thinking": thinking_text, + "signature": "sig_mock", + }, + {"type": "text", "text": answer_text}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 50}, + } + ).encode("utf-8"), + ) + + def _write_response(self, content_type: str, body: bytes) -> None: + self.send_response(200) + self.send_header("content-type", content_type) + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + server = ThreadingHTTPServer(("127.0.0.1", 0), MockAnthropicHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + try: + request_kwargs = { + "model": "anthropic/claude-sonnet-4-6", + "api_base": f"http://127.0.0.1:{server.server_port}", + "api_key": "test", + "messages": [ + { + "role": "user", + "content": "Solve a probability problem and show thinking.", + } + ], + "thinking": {"type": "adaptive"}, + "max_tokens": 128, + } + + non_stream_response = litellm.completion(**request_kwargs, stream=False) + non_stream_details = non_stream_response.usage.completion_tokens_details + assert non_stream_details is not None + assert non_stream_details.reasoning_tokens > 0 + + reasoning_chunks = [] + content_chunks = [] + stream_usage = None + for chunk in litellm.completion( + **request_kwargs, + stream=True, + stream_options={"include_usage": True}, + ): + chunk_dict = chunk.model_dump(exclude_none=True) + choices = chunk_dict.get("choices") or [] + if choices: + delta = choices[0].get("delta") or {} + if delta.get("reasoning_content"): + reasoning_chunks.append(delta["reasoning_content"]) + if delta.get("content"): + content_chunks.append(delta["content"]) + if chunk_dict.get("usage"): + stream_usage = chunk_dict["usage"] + + assert reasoning_chunks == thinking_parts + assert content_chunks == [answer_text] + assert stream_usage is not None + stream_completion_details = stream_usage["completion_tokens_details"] + assert ( + stream_completion_details["reasoning_tokens"] + == non_stream_details.reasoning_tokens + ) + assert stream_completion_details["text_tokens"] == ( + stream_usage["completion_tokens"] + - stream_completion_details["reasoning_tokens"] + ) + assert requests_seen == [ + { + "path": "/v1/messages", + "model": "claude-sonnet-4-6", + "stream": False, + "thinking": {"type": "adaptive"}, + }, + { + "path": "/v1/messages", + "model": "claude-sonnet-4-6", + "stream": True, + "thinking": {"type": "adaptive"}, + }, + ] + finally: + server.shutdown() + + def test_text_and_tool_streaming_has_index_zero(): """Test that mixed text and tool streaming responses have choice index=0""" chunks = [ diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index f96595ab8d8c..3323b9874723 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -97,6 +97,34 @@ def test_calculate_usage(): assert usage._cache_read_input_tokens == 0 +def test_calculate_usage_clamps_text_tokens_when_reasoning_estimate_exceeds_output(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={"input_tokens": 10, "output_tokens": 1}, + reasoning_content="This reasoning text intentionally tokenizes above one output token.", + ) + + assert usage.completion_tokens == 1 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == usage.completion_tokens + assert usage.completion_tokens_details.text_tokens == 0 + + +def test_calculate_usage_handles_mocked_output_tokens_with_reasoning_content(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={"input_tokens": 10, "output_tokens": MagicMock()}, + reasoning_content="mocked response reasoning", + ) + + assert usage.completion_tokens == 0 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 0 + + @pytest.mark.parametrize( "usage_object,expected_usage", [ diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 2d1586f0b171..01fd43ee032b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -112,6 +112,157 @@ async def test_handle_authentication_error_data_layer_errors_do_not_fall_back( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + asyncio.TimeoutError(), + OSError("network is unreachable"), + HTTPClientClosedError(), + PrismaError("can't reach database server"), + RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ), + ], +) +async def test_handle_authentication_error_db_infra_error_returns_503(db_error): + """Regression for the outage where valid keys got 401 for 4 hours: an + infrastructure-level DB failure during auth must surface as 503 (the DB + could not confirm the key), never as 401 ("Invalid API key").""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + db_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_handle_authentication_error_prisma_engine_teardown_returns_503(): + """Regression for the first-request-of-an-outage edge case: at the instant + the DB socket drops, the prisma query engine returns a malformed error + payload and prisma-client-py crashes with a bare + ``AttributeError: 'NoneType' object has no attribute 'get'`` before it can + raise P1001. That AttributeError reached auth and fell through to 401. It + must surface as 503 like every other infra failure during the outage.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + try: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + raise AssertionError("expected prisma to raise AttributeError") + except AttributeError as e: + teardown_error = e + + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + teardown_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error", + [ + # DB returned no row -> get_key_object raises this exact 401. + ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ), + # A bare auth failure raised as a plain Exception (e.g. master-key-only + # route) must keep returning 401, not get reclassified as 503. + Exception("Invalid proxy server token passed"), + ], +) +async def test_handle_authentication_error_genuine_auth_failure_stays_401(auth_error): + """Guard against the 503 conversion being too broad: a genuine auth + failure (missing key / wrong key) must still be 401.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + auth_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + @pytest.mark.asyncio async def test_handle_authentication_error_budget_exceeded(): handler = UserAPIKeyAuthExceptionHandler() diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 2e5eef2a0aa5..aeed0ad2825a 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -28,6 +28,7 @@ from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks +from fastapi import status from litellm.proxy.auth.user_api_key_auth import ( _route_requires_auth_despite_public, _reserve_budget_after_common_checks, @@ -2937,3 +2938,114 @@ async def test_master_key_auth_substitutes_alias_for_api_key(): finally: for k, v in _orig.items(): setattr(_proxy_server_mod, k, v) + + +def _proxy_attrs_for_db_lookup(): + """Minimal proxy_server attributes for driving the real + ``_user_api_key_auth_builder`` down to the DB key lookup.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + return { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {"allow_requests_on_db_unavailable": False}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + + +async def _run_builder_with_key_lookup(get_key_object_mock): + """Drive the real auth builder with ``get_key_object`` replaced by the + given mock. Returns the builder result.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs = _proxy_attrs_for_db_lookup() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", + get_key_object_mock, + ), + ): + return await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-db-lookup-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_builder_returns_503_when_db_lookup_raises_infra_error(): + """End-to-end: a DB infrastructure failure during the key lookup must + propagate past the ``except ProxyException`` guard and surface as 503, + not the 401 that masked the 4-hour outage. Killing the new 503 branch + flips this to 401 and fails the test.""" + get_key_object = AsyncMock(side_effect=ConnectionError("connection refused")) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_builder_returns_401_when_db_lookup_reports_missing_key(): + """Regression guard: a genuinely missing key (DB returned no row, which + ``get_key_object`` raises as a 401 ProxyException) must still be 401.""" + missing_key_error = ProxyException( + message="Authentication Error, Invalid proxy server token passed. key=..., not found in db.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + get_key_object = AsyncMock(side_effect=missing_key_error) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + +@pytest.mark.asyncio +async def test_builder_succeeds_when_db_lookup_returns_valid_token(): + """Regression guard: a valid key still authenticates. Proves the 503 + conversion only fires on the failure path and never intercepts success.""" + valid_token = UserAPIKeyAuth(api_key="sk-db-lookup-test", token="hashed-valid") + get_key_object = AsyncMock(return_value=valid_token) + + with patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=valid_token, + ) as mock_return: + result = await _run_builder_with_key_lookup(get_key_object) + + assert isinstance(result, UserAPIKeyAuth) + # Reaching the success-assembly return (never the exception handler) + # proves a valid key is unaffected by the 503 conversion. + mock_return.assert_awaited_once() diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py index f6ef02a86de7..d6e1d22fdde4 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py @@ -13,7 +13,9 @@ import os import sys from datetime import datetime, timedelta, timezone +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 import pytest @@ -24,6 +26,11 @@ LiteLLM_VerificationToken, ) from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager +from litellm.proxy.utils import ( + PrismaClient, + _deprecated_key_cache, + _lookup_deprecated_key, +) class TestMultiPodKeyRotation: @@ -557,3 +564,85 @@ async def test_lock_pattern_matches_spend_log_cleanup(self): assert acquire_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME assert release_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME + + +class TestDeprecatedKeyLookupDbE2E: + """DB-backed integration tests for deprecated key lookup behavior.""" + + @pytest.mark.asyncio + async def test_deprecated_key_grace_period_cache_hit_path(self): + """ + End-to-end validation against a real Prisma-backed DB: + - old key hash resolves through LiteLLM_DeprecatedVerificationToken + - repeated lookups hit the in-memory deprecated-key cache + - no ValueError/401 regression on subsequent requests + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + pytest.skip("DATABASE_URL not set; skipping DB-backed key-rotation E2E test.") + db_url = cast(str, database_url) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + prisma_client = PrismaClient( + database_url=db_url, proxy_logging_obj=proxy_logging_obj + ) + + old_token_hash = f"old-{uuid4().hex}" + active_token_hash = f"active-{uuid4().hex}" + _deprecated_key_cache.clear() + + await prisma_client.connect() + try: + await prisma_client.db.litellm_verificationtoken.create( + data={ + "token": active_token_hash, + "models": [], + } + ) + + await prisma_client.db.litellm_deprecatedverificationtoken.create( + data={ + "token": old_token_hash, + "active_token_id": active_token_hash, + "revoke_at": datetime.now(timezone.utc) + timedelta(minutes=5), + } + ) + + # Request 1 (DB path) + Request 2/3 (cache-hit path) + r1 = await _lookup_deprecated_key( + db=prisma_client.db, + hashed_token=old_token_hash, + ) + r2 = await _lookup_deprecated_key( + db=prisma_client.db, + hashed_token=old_token_hash, + ) + r3 = await _lookup_deprecated_key( + db=prisma_client.db, + hashed_token=old_token_hash, + ) + + assert r1 == active_token_hash + assert r2 == active_token_hash + assert r3 == active_token_hash + + cached = _deprecated_key_cache.get(old_token_hash) + assert isinstance(cached, tuple) + assert len(cached) == 3 + finally: + # Best-effort cleanup for idempotent reruns. + try: + await prisma_client.db.litellm_deprecatedverificationtoken.delete_many( + where={"token": old_token_hash} + ) + except Exception: + pass + try: + await prisma_client.db.litellm_verificationtoken.delete_many( + where={"token": active_token_hash} + ) + except Exception: + pass + _deprecated_key_cache.clear() + await prisma_client.disconnect() diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 9dcf5df4aebf..6021c2214267 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -107,6 +107,201 @@ def test_is_database_connection_generic_errors(): ) +@pytest.mark.parametrize( + "error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + OSError("network is unreachable"), + asyncio.TimeoutError(), + HTTPClientClosedError(), + ClientNotConnectedError(), + PrismaError("can't reach database server"), + PrismaError(), + ], +) +def test_is_database_service_unavailable_error_infra_failures(error): + """Infrastructure-level failures (socket/connection/timeout, prisma + transport, unknown PrismaError) mean the DB could not answer, so auth + must surface 503 instead of treating a valid key as invalid.""" + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is True + + +def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_dataerror(): + """Real-world regression: prisma-client-py raises the P1001 "can't reach + database server" connectivity failure as a DataError (a data-layer type). + A type-only check would miss it and return 401 during a genuine outage; + the message keyword must still classify it as service-unavailable -> 503.""" + p1001_as_dataerror = DataError( + data={ + "user_facing_error": { + "message": "Can't reach database server at `127.0.0.1`:`5499`", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + p1001_as_dataerror + ) + is True + ) + + +def test_is_database_service_unavailable_error_cached_plan_escapes_as_503(): + """Composes with the cached-plan retry: when that recovery fails and the + Postgres "cached plan must not change result type" error escapes (raised by + prisma as a data-layer RawQueryError), it is a transient stale-DB-state + condition, not an invalid key, so it must classify as service-unavailable + -> 503 rather than fall through to 401.""" + cached_plan_error = RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + cached_plan_error + ) + is True + ) + + +def test_is_database_service_unavailable_error_prisma_engine_malformed_payload(): + """Real-world regression: at the instant the DB socket drops, the prisma + query engine returns a malformed error payload (``user_facing_error.meta`` + is ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` before it + can raise the proper P1001 error. That bare AttributeError has no + connection keyword, so without the prisma-engine-origin check it falls + through to 401 on the first request of an outage. Reproduce the exact + prisma crash and assert it classifies as service-unavailable -> 503.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + with pytest.raises(AttributeError) as exc_info: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + + assert "no attribute 'get'" in str(exc_info.value) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is True + ) + + +def test_is_prisma_engine_internal_error_excludes_application_attributeerror(): + """The prisma-engine-origin check must stay narrow: a genuine AttributeError + raised by application code (a real bug) must NOT be classified as + service-unavailable, otherwise real bugs would silently become 503s.""" + + def application_bug(): + none_value = None + return none_value.get("oops") + + with pytest.raises(AttributeError) as exc_info: + application_bug() + + assert ( + PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc_info.value) + is False + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is False + ) + + +def test_is_prisma_engine_internal_error_excludes_data_layer_prisma_error(): + """A data-layer ``PrismaError`` (the DB IS reachable and rejected the data) + must stay 401. These are always raised from prisma internals, so the check + excludes any ``PrismaError`` by type before inspecting the traceback.""" + data_layer_error = UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "t"}}} + ) + try: + raise data_layer_error + except UniqueViolationError as e: + assert PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) is False + + +@pytest.mark.parametrize( + "error", + [ + DataError(data={"user_facing_error": {"meta": {"table": "t"}}}), + UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}), + Exception("some unrelated error"), + ValueError("bad value"), + ], +) +def test_is_database_service_unavailable_error_excludes_non_infra(error): + """Data-layer errors (the DB IS reachable and answered) and generic + non-DB errors must NOT be classified as service-unavailable, otherwise a + genuine 401 would be masked as a transient 503.""" + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is False + ) + + +def test_is_database_service_unavailable_error_asyncpg(monkeypatch): + """asyncpg connection/interface errors map to service-unavailable. asyncpg + is not a hard dependency, so inject a stand-in module to exercise the + branch deterministically regardless of the install environment.""" + import sys + import types + + fake_asyncpg = types.ModuleType("asyncpg") + fake_exceptions = types.ModuleType("asyncpg.exceptions") + + class PostgresConnectionError(Exception): + pass + + class InterfaceError(Exception): + pass + + class UniqueViolationError(Exception): # data-layer, must stay False + pass + + fake_exceptions.PostgresConnectionError = PostgresConnectionError + fake_exceptions.InterfaceError = InterfaceError + fake_exceptions.UniqueViolationError = UniqueViolationError + fake_asyncpg.exceptions = fake_exceptions + + monkeypatch.setitem(sys.modules, "asyncpg", fake_asyncpg) + monkeypatch.setitem(sys.modules, "asyncpg.exceptions", fake_exceptions) + + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + PostgresConnectionError("connection reset") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + InterfaceError("connection was closed") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + UniqueViolationError("duplicate key") + ) + is False + ) + + # Test should_allow_request_on_db_unavailable method @patch( "litellm.proxy.proxy_server.general_settings", diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 7593ea2f6dbf..7a9f73b70f3f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -321,6 +321,269 @@ def test_cost_calculation_does_not_duplicate_provider_prefix( assert call_kwargs["model"] == "azure_ai/claude-sonnet-4-5_gb_20250929" assert call_kwargs["custom_llm_provider"] == "azure_ai" + @patch("litellm.completion_cost") + def test_cost_calculation_resolves_unknown_model_from_litellm_params( + self, mock_completion_cost + ): + """When the body model is the "unknown" sentinel, the deployment model + from litellm_params must be used for costing, not "unknown" (which makes + completion_cost raise and the cost silently fall back to $0).""" + from datetime import datetime + + from litellm.types.utils import ModelResponse + + mock_completion_cost.return_value = 0.001 + + logging_obj = self._create_mock_logging_obj(model="unknown") + logging_obj.model_call_details["litellm_params"] = { + "model": "anthropic/claude-3-5-haiku-20241022", + "metadata": { + "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" + }, + } + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + + mock_response = MagicMock(spec=ModelResponse) + mock_response.id = "test-id" + mock_response.model = "unknown" + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=mock_response, + model="unknown", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + mock_completion_cost.assert_called_once() + assert ( + mock_completion_cost.call_args[1]["model"] + == "anthropic/claude-3-5-haiku-20241022" + ) + assert kwargs["response_cost"] == 0.001 + assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" + + @patch("litellm.completion_cost") + def test_cost_calculation_resolves_unknown_model_from_model_group( + self, mock_completion_cost + ): + """With only model_group available (no deployment litellm_params.model), + the leading passthrough/ prefix must be stripped so the cost map can + resolve the model.""" + from datetime import datetime + + from litellm.types.utils import ModelResponse + + mock_completion_cost.return_value = 0.002 + + logging_obj = self._create_mock_logging_obj(model="unknown") + logging_obj.model_call_details["litellm_params"] = { + "metadata": { + "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" + } + } + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + + mock_response = MagicMock(spec=ModelResponse) + mock_response.id = "test-id" + mock_response.model = "unknown" + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=mock_response, + model="unknown", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + mock_completion_cost.assert_called_once() + assert ( + mock_completion_cost.call_args[1]["model"] + == "anthropic/claude-3-5-haiku-20241022" + ) + assert kwargs["response_cost"] == 0.002 + + @patch("litellm.completion_cost") + def test_cost_calculation_skips_unknown_litellm_params_model_for_model_group( + self, mock_completion_cost + ): + """When litellm_params.model is itself the "unknown" sentinel, the + deployment-model branch must not short-circuit; resolution falls through + to model_group so costing still prices the real model instead of "unknown".""" + from datetime import datetime + + from litellm.types.utils import ModelResponse + + mock_completion_cost.return_value = 0.003 + + logging_obj = self._create_mock_logging_obj(model="unknown") + logging_obj.model_call_details["litellm_params"] = { + "model": "unknown", + "metadata": { + "model_group": "passthrough/anthropic/claude-3-5-haiku-20241022" + }, + } + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + + mock_response = MagicMock(spec=ModelResponse) + mock_response.id = "test-id" + mock_response.model = "unknown" + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=mock_response, + model="unknown", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + mock_completion_cost.assert_called_once() + assert ( + mock_completion_cost.call_args[1]["model"] + == "anthropic/claude-3-5-haiku-20241022" + ) + assert kwargs["response_cost"] == 0.003 + assert kwargs["model"] == "anthropic/claude-3-5-haiku-20241022" + + @patch("litellm.completion_cost") + def test_streaming_cost_calculation_resolves_model_from_message_start_chunk( + self, mock_completion_cost + ): + """On the bare /anthropic passthrough path litellm_params carries no model + or model_group and the body model is the "unknown" sentinel; the model + must be recovered from the message_start SSE event so completion_cost + prices the real model instead of failing on "unknown" and logging $0.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as RealLoggingObj, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + mock_completion_cost.return_value = 0.001 + + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + frames = [ + _sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ), + _sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + _sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hi"}, + }, + ), + _sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + _sse("message_stop", {"type": "message_stop"}), + ] + all_chunks = list( + PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) + ) + + logging_obj = RealLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1", + ) + logging_obj.model_call_details["model"] = "unknown" + logging_obj.model_call_details["stream"] = True + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + assert result["result"] is not None + mock_completion_cost.assert_called_once() + assert mock_completion_cost.call_args[1]["model"] == "claude-3-5-haiku-20241022" + assert result["kwargs"]["response_cost"] == 0.001 + assert result["kwargs"]["model"] == "claude-3-5-haiku-20241022" + + def test_extract_model_skips_non_dict_data_payload(self): + """A scalar data: payload (e.g. `data: null`) must be skipped, not crash + the streaming log handler with AttributeError, which would propagate out + and break spend logging for the whole request.""" + chunks = [ + "event: ping\ndata: null\n\n", + 'event: message_start\ndata: {"type": "message_start", "message": ' + '{"model": "claude-3-5-haiku-20241022"}}\n\n', + ] + + assert ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( + chunks + ) + == "claude-3-5-haiku-20241022" + ) + + def test_extract_model_parses_per_line_not_first_data_substring(self): + """A raw multi-line SSE event whose non-data line contains the substring + "data:" must not derail parsing: matching only lines that start with + "data:" recovers the message_start model, whereas a first-substring slice + would consume the wrong offset, fail to parse JSON, and return None.""" + raw_event = ( + "event: ping data: not-json\n" + 'data: {"type": "message_start", "message": ' + '{"model": "claude-3-5-haiku-20241022"}}\n\n' + ) + + assert ( + AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( + [raw_event] + ) + == "claude-3-5-haiku-20241022" + ) + class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" @@ -686,6 +949,72 @@ def test_store_batch_managed_object_success( ) +class TestBuildCompleteStreamingResponseRobustness: + """_build_complete_streaming_response must tolerate non-standard SSE frames.""" + + def _build(self, chunks: List[str]): + return AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=chunks, + litellm_logging_obj=MagicMock(), + model="claude-3-sonnet-20240229", + ) + + def test_done_frame_is_skipped(self): + """A bare 'data: [DONE]' control frame must not break reconstruction.""" + chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + "data: [DONE]", + ] + result = self._build(chunks) + assert result is not None + assert result.choices[0].message.content == "Hi" + + def test_non_json_sse_line_is_skipped(self): + """Non-JSON SSE lines (comments, keep-alive pings) must be skipped.""" + chunks = [ + ": ping", + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + "this is not json at all", + ] + # Must not raise; a malformed stream simply yields no usable response. + result = self._build(chunks) + assert result is None or hasattr(result, "choices") + + def test_mixed_valid_and_invalid_frames(self): + """Valid events are still collected when interleaved with invalid ones.""" + chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + "data: [DONE]", + ": keep-alive", + "not-json", + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ] + result = self._build(chunks) + assert result is not None + assert result.choices[0].message.content == "Hello" + + def test_done_in_text_payload_is_not_dropped(self): + """A valid event whose text content contains '[DONE]' must NOT be skipped.""" + chunks = [ + 'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-3-sonnet-20240229","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"The stream ends with [DONE]"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":8}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ] + result = self._build(chunks) + assert result is not None + assert result.choices[0].message.content == "The stream ends with [DONE]" class TestStreamFalseDeduplication: """ Regression tests for the duplicate-callback bug where a streaming pass-through diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6fbce4a54584..8723273b941d 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import fastapi import pytest @@ -308,6 +308,353 @@ def test_skip_server_startup( ), f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() + @pytest.mark.parametrize( + "timeout_config,expected_timeout", + [ + ({"database_connection_timeout": 30}, 30), + ({"database_connection_pool_timeout": 45}, 45), + ( + { + "database_connection_timeout": 30, + "database_connection_pool_timeout": 45, + }, + 30, + ), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_db_timeout_settings_are_forwarded_to_pool_timeout( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + timeout_config, + expected_timeout, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_connection_pool_limit": 5, + **timeout_config, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: ( + f"{url}?connection_limit={params['connection_limit']}&pool_timeout={params['pool_timeout']}" + ), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + assert appended_params["connection_limit"] == 5 + assert appended_params["pool_timeout"] == expected_timeout + + def test_build_db_connection_url_params_defaults(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60) + assert params == {"connection_limit": 10, "pool_timeout": 60} + + def test_build_db_connection_url_params_omits_none_timeouts(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + connect_timeout=None, + socket_timeout=None, + ) + assert "connect_timeout" not in params + assert "socket_timeout" not in params + + def test_build_db_connection_url_params_includes_optional_timeouts(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + connect_timeout=15, + socket_timeout=120, + ) + assert params["connect_timeout"] == 15 + assert params["socket_timeout"] == 120 + + def test_build_db_connection_url_params_extras_override_defaults(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + extra_params={ + "pgbouncer": "true", + "statement_cache_size": 0, + "pool_timeout": 5, + }, + ) + assert params["pgbouncer"] == "true" + assert params["statement_cache_size"] == 0 + assert params["pool_timeout"] == 5 + + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_db_connection_extra_params_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_connect_timeout": 15, + "database_socket_timeout": 120, + "database_extra_connection_params": { + "pgbouncer": "true", + "statement_cache_size": 0, + }, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: str(url), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + assert appended_params["connect_timeout"] == 15 + assert appended_params["socket_timeout"] == 120 + assert appended_params["pgbouncer"] == "true" + assert appended_params["statement_cache_size"] == 0 + + def test_build_db_connection_url_params_disable_prepared_statements(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + ) + assert params["pgbouncer"] == "true" + + def test_build_db_connection_url_params_no_pgbouncer_by_default(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + ) + assert "pgbouncer" not in params + + def test_build_db_connection_url_params_extra_pgbouncer_overrides_flag(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + extra_params={"pgbouncer": "false"}, + ) + assert params["pgbouncer"] == "false" + + @pytest.mark.parametrize( + "config_value, expect_pgbouncer", + [ + (True, True), + (False, False), + ("true", True), + ("false", False), + ("not-a-bool", False), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_disable_prepared_statements_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + config_value, + expect_pgbouncer, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_disable_prepared_statements": config_value, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: str(url), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + if expect_pgbouncer: + assert appended_params["pgbouncer"] == "true" + else: + assert "pgbouncer" not in appended_params + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") diff --git a/tests/test_litellm/proxy/utils/__init__.py b/tests/test_litellm/proxy/utils/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/__init__.py b/tests/test_litellm/proxy/utils/prisma_and_spend/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py new file mode 100644 index 000000000000..64611d1afcdc --- /dev/null +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -0,0 +1,180 @@ +"""Shared fixtures for tests/test_litellm/proxy/utils/prisma_and_spend/. + +All fixtures used by PR2 test files live here. Do NOT add fixtures inside +individual test files; if a fixture is missing, add it here and update the +Notion plan. + +The PrismaClient is exercised against a fully-mocked Prisma stack: the +``prisma.Prisma`` constructor and the writer/reader wrappers are patched +before PrismaClient.__init__ runs so the init code paths execute without +needing a generated Prisma client or a real database. +""" + +from __future__ import annotations + +import asyncio +import sys +from dataclasses import dataclass, field +from email.message import EmailMessage +from pathlib import Path +from typing import Any, Callable, Dict, Iterator, List, Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[5])) + + +_PRISMA_TABLES: List[str] = [ + "litellm_verificationtoken", + "litellm_teamtable", + "litellm_usertable", + "litellm_endusertable", + "litellm_organizationtable", + "litellm_proxymodeltable", + "litellm_modeltable", + "litellm_budgettable", + "litellm_spendlogs", + "litellm_config", + "litellm_usernotifications", + "litellm_healthchecktable", + "litellm_dailyuserspend", + "litellm_dailyteamspend", + "litellm_dailytagspend", + "litellm_managed_object_table", + "litellm_credentialstable", + "litellm_mcpservertable", + "litellm_audit_log", + "litellm_invitationlink", + "litellm_session_token_table", + "litellm_passthrough_endpoint_table", + "litellm_cron_job", + "litellm_passthrough_logs", + "litellm_promptstable", + "litellm_guardrailstable", + "litellm_managed_files", + "litellm_mcpusercredentials", + "litellm_objectpermissiontable", + "litellm_organizationmembership", +] + + +def _make_table_mock() -> MagicMock: + table = MagicMock() + table.find_unique = AsyncMock(return_value=None) + table.find_many = AsyncMock(return_value=[]) + table.find_first = AsyncMock(return_value=None) + table.create = AsyncMock() + table.create_many = AsyncMock() + table.update = AsyncMock() + table.update_many = AsyncMock() + table.upsert = AsyncMock() + table.delete = AsyncMock() + table.delete_many = AsyncMock() + table.count = AsyncMock(return_value=0) + table.group_by = AsyncMock(return_value=[]) + table.aggregate = AsyncMock(return_value={}) + return table + + + +@pytest.fixture +def mock_prisma_client() -> MagicMock: + """Bare ``db`` mock with all common LiteLLM_* tables stubbed. + + Override individual return values in a test:: + + mock_prisma_client.db.litellm_usertable.find_unique.return_value = user + """ + client = MagicMock(name="MockPrismaClient") + client.db = MagicMock(name="MockPrismaDB") + client.connect = AsyncMock() + client.disconnect = AsyncMock() + client.health_check = AsyncMock(return_value=[{"?column?": 1}]) + client.proxy_logging_obj = MagicMock() + client.proxy_logging_obj.failure_handler = AsyncMock() + client.spend_log_transactions = [] + client._spend_log_transactions_lock = asyncio.Lock() + client.jsonify_object = lambda data: dict(data) + client.db.is_connected = MagicMock(return_value=False) + client.db.connect = AsyncMock() + client.db.disconnect = AsyncMock() + client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) + client.db.execute_raw = AsyncMock() + client.db.tx = MagicMock() + client.db.batch_ = MagicMock() + for table_name in _PRISMA_TABLES: + setattr(client.db, table_name, _make_table_mock()) + return client + + +@pytest.fixture +def patched_prisma_import(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock]: + """Replace ``prisma.Prisma`` and ``PrismaWrapper`` so PrismaClient.__init__ + runs without a generated client. Yields the fake Prisma instance. + + ``prisma`` raises RuntimeError (not AttributeError) for the missing + ``Prisma`` attribute, so ``monkeypatch.setattr`` can't probe it; assign + directly and restore in teardown. + """ + import prisma as _prisma_pkg + import litellm.proxy.utils as _utils_mod + + fake_prisma = MagicMock(name="FakePrisma") + fake_prisma.is_connected = MagicMock(return_value=False) + fake_prisma.connect = AsyncMock() + fake_prisma.disconnect = AsyncMock() + + fake_prisma_factory = MagicMock(name="FakePrismaFactory", return_value=fake_prisma) + had_prisma_attr = "Prisma" in _prisma_pkg.__dict__ + previous_prisma_attr = _prisma_pkg.__dict__.get("Prisma") + _prisma_pkg.Prisma = fake_prisma_factory # type: ignore[attr-defined] + + fake_wrapper = MagicMock(name="FakePrismaWrapper") + fake_wrapper.is_connected = MagicMock(return_value=False) + fake_wrapper.connect = AsyncMock() + fake_wrapper.disconnect = AsyncMock() + fake_wrapper.query_raw = AsyncMock(return_value=[{"?column?": 1}]) + + def _fake_wrapper_ctor(*args: Any, **kwargs: Any) -> MagicMock: + return fake_wrapper + + monkeypatch.setattr(_utils_mod, "PrismaWrapper", _fake_wrapper_ctor) + fake_prisma.__wrapper__ = fake_wrapper + try: + yield fake_prisma + finally: + if had_prisma_attr: + _prisma_pkg.Prisma = previous_prisma_attr # type: ignore[attr-defined] + else: + try: + del _prisma_pkg.Prisma # type: ignore[attr-defined] + except AttributeError: + pass + + +@pytest.fixture +def prisma_client( + patched_prisma_import: MagicMock, + mock_prisma_client: MagicMock, + monkeypatch: pytest.MonkeyPatch, +) -> Any: + """Wired ``PrismaClient`` whose ``db`` attribute is the table mock. + + The init runs through the real code path (testing the constructor's + config-attribute setup) and is then snapped to the easier-to-assert + table mock for downstream behavior pinning. + """ + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False) + from litellm.proxy.utils import PrismaClient + + proxy_logging_obj = MagicMock(name="MockProxyLogging") + proxy_logging_obj.failure_handler = AsyncMock() + pc = PrismaClient( + database_url="postgresql://test:test@localhost:5432/test", + proxy_logging_obj=proxy_logging_obj, + ) + pc.db = mock_prisma_client.db + return pc + diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py new file mode 100644 index 000000000000..8c1bd6feeb8e --- /dev/null +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -0,0 +1,192 @@ +"""Pin ``PrismaClient`` read-side data operations. + +Symbols pinned here: + - ``PrismaClient.hash_token`` + - ``PrismaClient.jsonify_object`` + - ``PrismaClient.jsonify_team_object`` + - ``PrismaClient.check_view_exists`` + - ``PrismaClient.get_request_status`` + - ``PrismaClient.get_generic_data`` + - ``PrismaClient._query_first_with_cached_plan_fallback`` + - ``PrismaClient.get_data`` +""" + +from __future__ import annotations + +import hashlib +import json +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LiteLLM_VerificationTokenView +from litellm.proxy.utils import PrismaClient + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_happy_returns_row( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} + prisma_client.db.query_first = AsyncMock(return_value=expected) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + result = await prisma_client._query_first_with_cached_plan_fallback( + "SELECT * FROM x WHERE token = $1", "abc" + ) + actual = { + "result": result, + "call_count": prisma_client.db.query_first.await_count, + "args": prisma_client.db.query_first.await_args.args, + "matches": result == expected, + } + assert actual == { + "result": expected, + "call_count": 1, + "args": ("SELECT * FROM x WHERE token = $1", "abc"), + "matches": True, + } + prisma_client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_identical_query( + prisma_client: PrismaClient, +) -> None: + original_query = 'SELECT * FROM "LiteLLM_VerificationToken" WHERE v.token = $1' + expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} + manager = MagicMock() + query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + reconnect = AsyncMock(return_value=True) + manager.attach_mock(query_first, "query_first") + manager.attach_mock(reconnect, "attempt_db_reconnect") + prisma_client.db.query_first = query_first + prisma_client.attempt_db_reconnect = reconnect + + result = await prisma_client._query_first_with_cached_plan_fallback( + original_query, "abc" + ) + + assert result == expected + assert query_first.await_count == 2 + first_call, retry_call = query_first.await_args_list + assert retry_call.args == first_call.args == (original_query, "abc") + reconnect.assert_awaited_once() + assert reconnect.await_args.kwargs.get("force", False) is False + assert [name for name, *_ in manager.mock_calls] == [ + "query_first", + "attempt_db_reconnect", + "query_first", + ] + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_never_deallocates( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + prisma_client.db.execute_raw = AsyncMock(return_value=0) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + prisma_client.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_propagates_when_retry_also_fails( + prisma_client: PrismaClient, +) -> None: + plan_error = RuntimeError("cached plan must not change result type") + prisma_client.db.query_first = AsyncMock(side_effect=[plan_error, plan_error]) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with pytest.raises(RuntimeError, match="cached plan must not change result type"): + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + assert prisma_client.db.query_first.await_count == 2 + prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_retries_when_reconnect_returns_false( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + result = await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + assert result == expected + assert prisma_client.db.query_first.await_count == 2 + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors( + prisma_client: PrismaClient, +) -> None: + prisma_client.db.query_first = AsyncMock( + side_effect=RuntimeError("totally unrelated") + ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + with pytest.raises(RuntimeError, match="totally unrelated"): + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert prisma_client.db.query_first.await_count == 1 + prisma_client.attempt_db_reconnect.assert_not_awaited() + + + +@pytest.mark.asyncio +async def test_get_data_combined_view_returns_view_for_deprecated_key( + prisma_client: PrismaClient, +) -> None: + """Grace-period rotation, full get_data flow: the old hash misses the + combined view, the deprecated-key table resolves it to the active token, + and get_data must return the recursive lookup's finished view instead of + re-running dict normalization on it (which raised TypeError and turned + every grace-period request into a 401).""" + old_hash = "hashed-old-token-grace-e2e" + active_hash = "hashed-active-token-grace-e2e" + active_row = { + "token": active_hash, + "team_models": None, + "team_blocked": None, + "team_members_with_roles": None, + "user_id": None, + "expires": None, + } + prisma_client.db.query_first = AsyncMock(side_effect=[None, active_row]) + prisma_client.db.litellm_deprecatedverificationtoken = MagicMock() + prisma_client.db.litellm_deprecatedverificationtoken.find_first = AsyncMock( + return_value=SimpleNamespace( + active_token_id=active_hash, + revoke_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + ) + + response = await prisma_client.get_data( + token=old_hash, table_name="combined_view", query_type="find_unique" + ) + + assert isinstance(response, LiteLLM_VerificationTokenView) + assert response.token == active_hash diff --git a/uv.lock b/uv.lock index 4ca432e9260e..6f01d13e7847 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-07T23:51:30.443285Z" +exclude-newer = "2026-06-10T00:03:18.224478Z" exclude-newer-span = "P3D" [manifest] @@ -3083,7 +3083,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.84.7" +version = "1.84.8" source = { editable = "." } dependencies = [ { name = "aiohttp" },