scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Merge pull request #30332 from BerriAI/litellm_backport_1_84_x_0612

Browse source 
chore(release): backport #27319, #27756, #27921, #28395, #29983, #29984, #29986, #30160, #30327 to stable/1.84.x and cut 1.84.8
This commit is contained in:
Mateo Wang 2026-06-12 18:20:57 -07:00 • committed by GitHub
commit 34e9d7d6f8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 2849 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -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

16
litellm/llms/anthropic/chat/handler.py
View file

@ -557,6 +557,10 @@ class ModelResponseIterator:
# 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 @@ class ModelResponseIterator:
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 @@ class ModelResponseIterator:
"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 ""),
)
]

15
litellm/llms/anthropic/chat/transformation.py
View file

@ -1856,8 +1856,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
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 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
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=(

41
litellm/llms/azure/containers/transformation.py
View file

@ -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 @@ class AzureContainerConfig(OpenAIContainerConfig):
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 @@ class AzureContainerConfig(OpenAIContainerConfig):
{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",
)

26
litellm/llms/custom_httpx/container_handler.py
View file

@ -257,14 +257,19 @@ class GenericContainerHandler:
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 @@ class GenericContainerHandler:
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 @@ class GenericContainerHandler:
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 @@ class GenericContainerHandler:
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}")

20
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -7816,7 +7816,7 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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 @@ class BaseLLMHTTPHandler:
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(

35
litellm/proxy/_types.py
View file

@ -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"
)

10
litellm/proxy/auth/auth_exception_handler.py
View file

@ -134,6 +134,16 @@ class UserAPIKeyAuthExceptionHandler:
)
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,

4
litellm/proxy/container_endpoints/endpoints.py
View file

@ -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,

14
litellm/proxy/container_endpoints/handler_factory.py
View file

@ -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,

75
litellm/proxy/container_endpoints/ownership.py
View file

@ -23,6 +23,13 @@ CONTAINER_OBJECT_PURPOSE = "container"
_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_<hex>``) 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,

86
litellm/proxy/db/exception_handler.py
View file

@ -109,6 +109,92 @@ class PrismaDBExceptionHandler:
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):
"""

64
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -100,6 +100,42 @@ class AnthropicPassthroughLoggingHandler:
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 @@ class AnthropicPassthroughLoggingHandler:
"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 @@ class AnthropicPassthroughLoggingHandler:
):
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 @@ class AnthropicPassthroughLoggingHandler:
# 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 @@ class AnthropicPassthroughLoggingHandler:
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,

95
litellm/proxy/proxy_cli.py
View file

@ -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:

79
litellm/proxy/utils.py
View file

@ -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 @@ class PrismaClient:
self, sql_query: str, *args
) -> Optional[dict]:
"""
Execute a query with automatic fallback for PostgreSQL cached plan errors.
Execute a query, recovering once from PostgreSQL's "cached plan must not
change result type" error.
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.
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.
Args:
sql_query: SQL query string to execute
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.
Returns:
Query result or None
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.
Raises:
Original exception if not a cached plan error
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 @@ class PrismaClient:
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 @@ class PrismaClient:
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:

10
litellm/router.py
View file

@ -5551,6 +5551,7 @@ class Router:
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 @@ class Router:
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(

4
pyproject.toml
View file

@ -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",
]

177
tests/proxy_unit_tests/test_deprecated_key_grace_period.py Normal file
View file

@ -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()

251
tests/test_litellm/containers/test_azure_container_transformation.py
View file

@ -109,6 +109,31 @@ class TestAzureContainerConfig:
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 @@ class TestAzureContainerKnownFailureRegressions:
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 @@ class TestAzureContainerKnownFailureRegressions:
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_<hex>``) 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"
)

287
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -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 = [

28
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -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",
[

151
tests/test_litellm/proxy/auth/test_auth_exception_handler.py
View file

@ -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()

112
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -28,6 +28,7 @@ from litellm.proxy._types import (
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()

89
tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py
View file

@ -13,7 +13,9 @@ Covers the critical gaps:
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 @@ from litellm.proxy._types import (
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 @@ class TestKeyRotationInitialization:
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()

195
tests/test_litellm/proxy/db/test_exception_handler.py
View file

@ -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",

329
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -321,6 +321,269 @@ class TestAzureAnthropicCostCalculation:
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 @@ class TestAnthropicBatchPassthroughCostTracking:
)
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

349
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -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 @@ class TestProxyInitializationHelpers:
), 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")

0
tests/test_litellm/proxy/utils/__init__.py Normal file
View file

0
tests/test_litellm/proxy/utils/prisma_and_spend/__init__.py Normal file
View file

180
tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py Normal file
View file

@ -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

192
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py Normal file
View file

@ -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

4
uv.lock generated
View file

@ -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" },