Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_fix_v1_messages_midstream_timeout_failure_logging

# Conflicts:
#	litellm/litellm_core_utils/litellm_logging.py
This commit is contained in:
mateo-berri 2026-09-03 18:27:43 -07:00
commit 2f3d1ca575
60 changed files with 3338 additions and 568 deletions

View file

@ -79,6 +79,11 @@ test_paths:
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
locally rather than in pull-request jobs
paths:
- tests/load_tests/test_granian_admission_saturation.py
- reason: >-
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a

View file

@ -3,7 +3,7 @@
"limit": 14074
},
"reportArgumentType": {
"limit": 2214
"limit": 2206
},
"reportAssignmentType": {
"limit": 319
@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15287
"limit": 15285
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44362
"limit": 44360
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38323
"limit": 38311
},
"reportUnknownParameterType": {
"limit": 19624
},
"reportUnknownVariableType": {
"limit": 29861
"limit": 29847
},
"reportUnnecessaryCast": {
"limit": 111

View file

@ -1,8 +1,9 @@
"""Anthropic error format type definitions."""
from collections.abc import Mapping
from typing import Literal
from typing_extensions import Required, TypedDict
from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
# Known Anthropic error types
# Source: https://docs.anthropic.com/en/api/errors
@ -23,6 +24,7 @@ class AnthropicErrorDetail(TypedDict):
type: AnthropicErrorType
message: str
provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]]
class AnthropicErrorResponse(TypedDict, total=False):

View file

@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov
retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
topology changed.
redis-py 8.x fixed this upstream with gentler machinery than this override's
``node.disconnect()`` (which also kills connections other coroutines are mid-operation
on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
killed connection): it marks in-use connections for reconnect only after their current
operation completes, disconnects only the idle pooled ones, and defers reinitialization
to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
recovery API, the factory returns the base ``RedisCluster`` unmodified.
redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream
still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent
caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full
teardown. For those versions the factory returns a thin wrapper around upstream's
``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError,
a third consecutive timeout on the same node, or a concurrent request from any other command
or ``aclose()`` still reinits).
"""
import asyncio
@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol):
mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
own logic fully typed without a banned ``typing.cast``."""
name: str
async def execute_command(
self,
*args: object,
@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol):
#: this override can't see (Python won't error -- it'll just run our now-stale copy), so
#: construction logs a loud warning rather than silently trusting an unverified copy.
_VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3
def get_litellm_async_redis_cluster_class(
def get_litellm_async_redis_cluster_class( # noqa: C901 # supports redis-py version-specific cluster implementations
cluster_node_class: type | None = None,
base_cluster_class: type | None = None,
) -> type["_AsyncRedisClusterType"]:
"""Returns the base ``RedisCluster`` when the installed redis-py already recovers a
node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
subclass with the per-node isolation fix for older versions whose upstream branch
tears down the whole cluster client.
"""Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already
recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster``
subclass with the per-node isolation fix for older versions whose upstream branch tears
down the whole cluster client.
``cluster_node_class`` exists for dependency injection in tests; production callers
leave it unset and the installed ``ClusterNode`` is used.
``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests;
production callers leave them unset and the installed redis-py classes are used.
Imported lazily because this module is reachable from a base ``import litellm`` while
redis is not a base dependency. Cheap to call repeatedly: the underlying redis
@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class(
from redis.exceptions import TimeoutError as _RedisTimeoutError
node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster
if hasattr(node_class, "update_active_connections_for_reconnect"):
verbose_logger.debug(
"redis-py %s recovers a node-level connection error per-connection upstream; "
"using the base RedisCluster without litellm's node-isolation override.",
"redis-py %s recovers node connections per-connection upstream; using "
"LiteLLM's timeout-tolerant RedisCluster wrapper.",
redis.__version__,
)
return _BaseAsyncRedisCluster
class LiteLLMAsyncRedisClusterTimeoutTolerant(
base_class # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass] # the injected base class is selected at runtime
):
def __init__(
self,
*args: object,
**kwargs: object, # kwargs-ok: passes redis-py's constructor kwargs through untouched
) -> None:
self._litellm_initialize = False
self._litellm_reinit_requests = 0
self._litellm_tolerated_timeouts = 0
super().__init__(*args, **kwargs)
self._litellm_consecutive_timeouts: dict[ # mutable-ok: per-node counter updated on the command hot path
str, int
] = {}
@property
def _initialize(self) -> bool:
return self._litellm_initialize
@_initialize.setter
def _initialize(self, value: bool) -> None:
if value:
self._litellm_reinit_requests += 1
self._litellm_initialize = value
async def _execute_command(
self,
target_node: _ClusterNodeAttrs,
*args: object,
**kwargs: object, # kwargs-ok: matches redis-py's own command dispatch signature
) -> object:
outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts
pending_before: Final = self._litellm_initialize
try:
result: Final = await super()._execute_command(target_node, *args, **kwargs)
except _RedisTimeoutError:
timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1
if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT:
self._litellm_consecutive_timeouts.pop(target_node.name, None)
raise
self._litellm_consecutive_timeouts[target_node.name] = timeouts
self._litellm_tolerated_timeouts += 1
if (
not pending_before
and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before
):
self._initialize = False
raise
if self._litellm_consecutive_timeouts:
self._litellm_consecutive_timeouts.pop(target_node.name, None)
return result
return LiteLLMAsyncRedisClusterTimeoutTolerant
if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
verbose_logger.warning(

View file

@ -151,6 +151,7 @@ DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200"))
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600"))
MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60"))
# Default npm cache directory for STDIO MCP servers.
# npm/npx needs a writable cache dir; in containers the default (~/.npm)

View file

@ -10,7 +10,7 @@ import subprocess
import sys
import time
import traceback
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime as dt_object
from functools import lru_cache
from types import MappingProxyType, TracebackType
@ -5884,14 +5884,28 @@ def _get_status_fields(
#########################################################
# Map - guardrail_information.guardrail_status to guardrail_status
#########################################################
guardrail_status: GuardrailStatus = "not_run"
if guardrail_information and isinstance(guardrail_information, list):
for information in guardrail_information:
if isinstance(information, dict):
raw_status = information.get("guardrail_status", "not_run")
if raw_status != "not_run":
guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
break
# Severity order, least severe first. The status aggregates across ALL
# guardrail entries rather than taking the first non-"not_run" one: a
# pre_call guardrail that passed (e.g. a mask) records its entry before a
# later guardrail's block, and first-wins would report a blocked request
# as "success".
GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = (
"not_run",
"success",
"guardrail_failed_to_respond",
"guardrail_intervened",
)
entries: Final[Sequence[object]] = guardrail_information if isinstance(guardrail_information, list) else ()
raw_statuses: Final[Iterator[object]] = (
entry.get("guardrail_status", "not_run") for entry in entries if isinstance(entry, dict)
)
# A guardrail is free to write any value here, and an unhashable one would
# raise TypeError on the mapping lookup and drop the whole payload.
guardrail_status: Final[GuardrailStatus] = max(
(GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") for raw_status in raw_statuses if isinstance(raw_status, str)),
key=GUARDRAIL_STATUS_SEVERITY.index,
default="not_run",
)
return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status)

View file

@ -8,23 +8,31 @@ Routes to native Cortex REST API endpoints based on model:
Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
"""
import copy
import json
import re
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict
import httpx
from typing_extensions import ReadOnly
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.litellm_core_utils.prompt_templates.factory import (
anthropic_process_openai_file_message,
convert_to_anthropic_tool_result,
create_anthropic_image_param,
select_anthropic_content_block_type_for_file,
)
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolMessage
from litellm.types.utils import (
ChatCompletionMessageToolCall,
ChatCompletionUsageBlock,
Choices,
Function,
GenericStreamingChunk,
Message,
ModelResponse,
Usage,
ModelResponseStream,
)
from ...base_llm.base_model_iterator import BaseModelResponseIterator
@ -93,6 +101,103 @@ def _is_claude_model(model: str) -> bool:
return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES)
def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
"""One OpenAI ``image_url`` block in the native shape Cortex accepts.
Cortex documents base64 sources only, so remote URLs are inlined the way every
other base64-only Anthropic dialect (Bedrock invoke, Vertex) inlines them, and
pdf/text data URIs become document blocks rather than malformed image blocks.
"""
image_url: Final = block.get("image_url")
url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url")
if not url:
return block
converted: Final = (
anthropic_process_openai_file_message({"type": "file", "file": {"file_data": url}})
if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document"
else create_anthropic_image_param(
image_url if isinstance(image_url, dict) else url, # mutable-ok: caller's JSON block
format=_image_url_field(image_url, "format"),
is_bedrock_invoke=True,
)
)
cache_control: Final = block.get("cache_control")
if cache_control is None:
return converted
return {**converted, "cache_control": cache_control} # mutable-ok: JSON wire block
def _image_url_field(image_url: object, key: str) -> str | None:
value: Final = image_url.get(key) if isinstance(image_url, dict) else None
return value if isinstance(value, str) else None
def _data_uri_media_type(url: str) -> str:
match: Final = re.match(r"data:([^;,]+)", url)
return match.group(1) if match else ""
def _convert_image_url_blocks_to_anthropic(content: object) -> object:
if not isinstance(content, list):
return content
return [ # mutable-ok: JSON wire blocks
_convert_image_url_to_anthropic(block)
if isinstance(block, Mapping) and block.get("type") == "image_url"
else block
for block in content
]
def _convert_tool_result_to_anthropic(
content: object, tool_call_id: str, cache_control: object
) -> Mapping[str, object]:
"""The Anthropic ``tool_result`` block for one OpenAI tool message.
Delegating to the shared converter keeps image, document and per-block cache
breakpoints identical to every other Anthropic dialect; only the plain-string
and non-list shapes it does not model are handled here.
"""
if not isinstance(content, list):
plain: Final[dict[str, object]] = { # mutable-ok: JSON wire block
"type": "tool_result",
"tool_use_id": tool_call_id,
"content": content if isinstance(content, str) else json.dumps(content),
}
return {**plain, "cache_control": cache_control} if cache_control is not None else plain
converted: Final = convert_to_anthropic_tool_result(
ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content),
force_base64=True,
)
if cache_control is None:
return converted
return {**converted, "cache_control": cache_control} # mutable-ok: JSON wire block
def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]: # mutable-ok: JSON wire blocks
"""The assistant turn's thinking blocks that can legally be echoed back.
Only signed blocks round-trip: Cortex rejects a thinking block whose signature is
missing, which is what an unsigned block from a non-thinking turn would produce.
"""
blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None)
if not isinstance(blocks, list):
return [] # mutable-ok: JSON wire blocks
return [ # mutable-ok: JSON wire blocks
dict(block)
for block in blocks
if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking")
]
def _clean_input_schema(schema: object) -> object: # mutable-ok: JSON schema copy
return (
{key: value for key, value in schema.items() if key != "$schema"}
if isinstance(schema, Mapping)
else schema # mutable-ok: JSON schema copy
) # mutable-ok: JSON schema copy
class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
"""
Snowflake Cortex REST API unified provider.
@ -178,7 +283,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
if "description" in func:
anthropic_tool["description"] = func["description"]
if "parameters" in func:
anthropic_tool["input_schema"] = func["parameters"]
anthropic_tool["input_schema"] = _clean_input_schema(func["parameters"])
else:
anthropic_tool["input_schema"] = {
"type": "object",
@ -186,10 +291,16 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
}
anthropic_tools.append(anthropic_tool)
else:
anthropic_tools.append(tool)
anthropic_tools.append(
{**tool, "input_schema": _clean_input_schema(tool["input_schema"])} # mutable-ok: JSON wire tool
if "input_schema" in tool
else tool
)
return anthropic_tools
def _extract_system_and_messages(self, messages: list[AllMessageValues]) -> tuple[str | None, list[dict]]:
def _extract_system_and_messages( # mutable-ok: JSON wire messages
self, messages: list[AllMessageValues]
) -> tuple[list[dict] | None, list[dict]]:
"""
Split messages into system prompt and conversation turns for Anthropic format.
@ -197,26 +308,39 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
- assistant messages with tool_calls tool_use content blocks
- tool role messages user role with tool_result content blocks
"""
system_parts: Final[list[str]] = []
conversation: Final[list[dict]] = []
system_parts: Final[list[dict]] = [] # mutable-ok: JSON wire messages
conversation: Final[list[dict]] = [] # mutable-ok: JSON wire messages
for msg in messages:
if isinstance(msg, dict):
role = msg.get("role", "")
content: Any = msg.get("content", "")
msg_cache_control: object = msg.get("cache_control")
else:
role = getattr(msg, "role", "")
content = getattr(msg, "content", "")
msg_cache_control = getattr(msg, "cache_control", None)
if role == "system":
if isinstance(content, str) and content:
system_parts.append(content)
system_parts.append({"type": "text", "text": content}) # mutable-ok: JSON wire system block
elif isinstance(content, list):
system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text"))
system_parts.extend(
{ # mutable-ok: JSON wire system block
"type": "text",
"text": block.get("text", ""),
**(
{"cache_control": block["cache_control"]} if "cache_control" in block else {}
), # mutable-ok: JSON wire block
}
for block in content
if isinstance(block, Mapping) and block.get("type") == "text"
)
elif role == "assistant":
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
thinking_blocks = _signed_thinking_blocks(msg)
if tool_calls:
content_blocks: list[dict[str, object]] = []
content_blocks: list[dict[str, object]] = list(thinking_blocks) # mutable-ok: JSON wire blocks
if content:
content_blocks.append({"type": "text", "text": content})
for tc in tool_calls:
@ -239,18 +363,26 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
}
)
conversation.append({"role": "assistant", "content": content_blocks})
elif thinking_blocks:
thinking_content = (
[
*thinking_blocks,
*copy.deepcopy(content),
]
if isinstance(content, list)
else [*thinking_blocks, *([{"type": "text", "text": content}] if content else [])]
) # rebind-ok: loop-local normalized content
conversation.append({"role": "assistant", "content": thinking_content})
else:
conversation.append({"role": "assistant", "content": content})
elif role == "tool":
tool_call_id = (
tool_call_id_value = (
msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "")
)
tool_content = content if isinstance(content, str) else json.dumps(content)
tool_result_block = {
"type": "tool_result",
"tool_use_id": tool_call_id,
"content": tool_content,
}
tool_call_id = (
tool_call_id_value if isinstance(tool_call_id_value, str) else ""
) # rebind-ok: normalized loop value
tool_result_block = _convert_tool_result_to_anthropic(content, tool_call_id, msg_cache_control)
if (
conversation
and conversation[-1]["role"] == "user"
@ -260,11 +392,18 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
):
conversation[-1]["content"].append(tool_result_block)
else:
conversation.append({"role": "user", "content": [tool_result_block]})
conversation.append(
{"role": "user", "content": [tool_result_block]} # mutable-ok: JSON wire message
) # mutable-ok: JSON wire message
else:
conversation.append({"role": role, "content": content})
conversation.append( # mutable-ok: JSON wire message
{ # mutable-ok: JSON wire message
"role": role,
"content": _convert_image_url_blocks_to_anthropic(content),
} # mutable-ok: JSON wire message
)
system: Final[str | None] = "\n\n".join(system_parts) if system_parts else None
system: Final[list[dict] | None] = system_parts if system_parts else None # mutable-ok: JSON wire messages
return system, conversation
def transform_request(
@ -339,7 +478,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
extra_body: dict,
) -> dict:
"""Anthropic Messages format for /messages endpoint."""
system, conversation = self._extract_system_and_messages(messages)
passthrough_system: Final = optional_params.pop("system", None)
extracted_system, conversation = self._extract_system_and_messages(messages)
system: Final = passthrough_system if passthrough_system is not None else extracted_system
if "tools" in optional_params:
optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"])
@ -353,16 +494,19 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
model_name: Final = model.removeprefix("snowflake/")
body: Final[dict[str, object]] = {
"model": model_name,
"messages": conversation,
"stream": stream,
**optional_params,
**extra_body,
}
body: Final[dict[str, object]] = normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire body
{ # mutable-ok: JSON wire body
"model": model_name,
"messages": conversation,
"stream": stream,
**optional_params,
**extra_body, # mutable-ok: JSON wire body
}
)
if system is not None:
body["system"] = system
body["system"] = normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire payload
{"system": system} # mutable-ok: JSON wire payload
)["system"]
if "max_tokens" not in body:
body["max_tokens"] = 4096 # reasonable default; Anthropic API max varies by model
@ -435,23 +579,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
additional_args={"complete_input_dict": request_data},
)
text_content = ""
tool_calls: Final = []
for block in response_json.get("content", []):
if block.get("type") == "text":
text_content += block.get("text", "")
elif block.get("type") == "tool_use":
tool_calls.append(
ChatCompletionMessageToolCall(
id=block.get("id", ""),
type="function",
function=Function(
name=block.get("name", ""),
arguments=json.dumps(block.get("input", {})),
),
)
)
anthropic_config: Final = AnthropicConfig()
text_content, _, thinking_blocks, reasoning_content, tool_calls, _, _, _ = (
anthropic_config.extract_response_content(completion_response=dict(response_json))
)
_stop_reason_map: Final = {
"end_turn": "stop",
@ -461,9 +592,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
}
finish_reason: Final = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop")
message: Final = Message(content=text_content or None, role="assistant")
if tool_calls:
message.tool_calls = tool_calls
message: Final = Message(
content=text_content or None,
role="assistant",
tool_calls=tool_calls or None,
thinking_blocks=thinking_blocks,
reasoning_content=reasoning_content,
)
choice: Final = Choices(
finish_reason=finish_reason,
@ -471,11 +606,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
message=message,
)
usage_data: Final = response_json.get("usage", {})
usage: Final = Usage(
prompt_tokens=usage_data.get("input_tokens", 0),
completion_tokens=usage_data.get("output_tokens", 0),
total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0),
# Cortex reports prompt-cache creation/read counts alongside input_tokens; the
# shared calculator folds them into prompt_tokens_details so cached input is
# visible and billed at its own rate.
usage: Final = anthropic_config.calculate_usage(
usage_object=response_json.get("usage", {}),
reasoning_content=reasoning_content,
completion_response=dict(response_json),
)
model_response.choices = [choice]
@ -516,15 +653,19 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
json_mode: bool | None = False,
):
super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
self._tool_index = 0
self._tool_id = ""
self._tool_name = ""
self._input_tokens = 0
# Cortex streams the Anthropic SSE dialect on /messages, so its events are parsed
# by Anthropic's own parser: thinking deltas, signatures and prompt-cache usage
# all arrive the way they do on every other Anthropic-dialect provider.
self._anthropic_parser: Final = AnthropicStreamParser(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
if "choices" in chunk:
return self._parse_openai_chunk(chunk)
return self._parse_anthropic_chunk(chunk)
return self._anthropic_parser.chunk_parser(chunk)
def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk:
choices: Final = chunk.get("choices", [])
@ -566,117 +707,3 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
index=choice.get("index", 0),
tool_use=tool_use,
)
def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk:
event_type: Final = chunk.get("type", "")
if event_type == "message_start":
message: Final = chunk.get("message", {})
usage_data = message.get("usage", {})
self._input_tokens = usage_data.get("input_tokens", 0)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
elif event_type == "content_block_delta":
delta = chunk.get("delta", {})
delta_type: Final = delta.get("type", "")
if delta_type == "text_delta":
return GenericStreamingChunk(
text=delta.get("text", ""),
is_finished=False,
finish_reason="",
usage=None,
index=chunk.get("index", 0),
tool_use=None,
)
elif delta_type == "input_json_delta":
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=chunk.get("index", 0),
tool_use=ChatCompletionToolCallChunk(
id=self._tool_id,
type="function",
function={
"name": self._tool_name,
"arguments": delta.get("partial_json", ""),
},
index=self._tool_index,
),
)
elif event_type == "content_block_start":
content_block: Final = chunk.get("content_block", {})
if content_block.get("type") == "tool_use":
self._tool_id = content_block.get("id", "")
self._tool_name = content_block.get("name", "")
self._tool_index = chunk.get("index", 0)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=chunk.get("index", 0),
tool_use=ChatCompletionToolCallChunk(
id=self._tool_id,
type="function",
function={"name": self._tool_name, "arguments": ""},
index=self._tool_index,
),
)
elif event_type == "message_delta":
delta = chunk.get("delta", {})
stop_reason: Final = delta.get("stop_reason", "")
usage_data = chunk.get("usage", {})
_stop_map: Final = {
"end_turn": "stop",
"max_tokens": "length",
"tool_use": "tool_calls",
"stop_sequence": "stop",
}
usage = None
if usage_data or self._input_tokens:
output_t: Final = usage_data.get("output_tokens", 0)
input_t: Final = self._input_tokens or usage_data.get("input_tokens", 0)
usage = ChatCompletionUsageBlock(
prompt_tokens=input_t,
completion_tokens=output_t,
total_tokens=input_t + output_t,
)
return GenericStreamingChunk(
text="",
is_finished=True,
finish_reason=_stop_map.get(stop_reason, "stop"),
usage=usage,
index=0,
tool_use=None,
)
elif event_type == "message_stop":
return GenericStreamingChunk(
text="",
is_finished=True,
finish_reason="stop",
usage=None,
index=0,
tool_use=None,
)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)

View file

@ -11,7 +11,8 @@ being registered, so a gateway with no EMA upstream never stores bearer material
The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the
id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an
expired assertion with a refresh token is still renewable, and the DB row is the source of
truth, the same contract as the per-user OAuth credential store.
truth, the same contract as the per-user OAuth credential store. Reads use a per-process cache with
TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against stale in-flight reads.
"""
from __future__ import annotations
@ -24,6 +25,8 @@ import jwt
from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
@ -45,6 +48,46 @@ class SSOIdentityAssertion(BaseModel):
expires_at: datetime | None = None
class SSOAssertionCache:
"""Process-local read cache. ``invalidate`` bumps a process-wide epoch so a fetch that started
before a login cannot repopulate the old assertion after it."""
def __init__(self, ttl_seconds: int = MCP_SSO_ASSERTION_CACHE_TTL_SECONDS) -> None:
self._entries = InMemoryCache(
max_size_in_memory=MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
default_ttl=ttl_seconds,
)
self._epoch: int = 0
def epoch(self) -> int:
return self._epoch
def get(self, user_id: str) -> SSOIdentityAssertion | None:
cached: Final = self._entries.get_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
user_id
)
return cached if isinstance(cached, SSOIdentityAssertion) else None
def set_if_unchanged(self, user_id: str, assertion: SSOIdentityAssertion, seen_epoch: int) -> None:
if self._epoch != seen_epoch:
return
self._entries.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
user_id, assertion
)
def invalidate(self, user_id: str) -> None:
self._epoch += 1
self._entries.delete_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
user_id
)
def flush(self) -> None:
self._entries.flush_cache() # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
_ASSERTION_CACHE: Final = SSOAssertionCache()
class _IdTokenClaims(BaseModel):
exp: float | None = None
iss: str | None = None
@ -107,7 +150,9 @@ async def ema_assertion_retention_enabled() -> bool:
return row is not None
async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAssertion) -> None:
async def persist_sso_identity_assertion(
user_id: str, assertion: SSOIdentityAssertion, cache: SSOAssertionCache = _ASSERTION_CACHE
) -> None:
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper # noqa: PLC0415 # runtime global
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global
@ -127,11 +172,10 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss
"update": {"assertion_b64": encoded},
},
)
cache.invalidate(user_id)
async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None:
"""The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None:
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper # noqa: PLC0415 # runtime global
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global
@ -160,6 +204,21 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N
)
async def fetch_sso_identity_assertion(
user_id: str, cache: SSOAssertionCache = _ASSERTION_CACHE
) -> SSOIdentityAssertion | None:
"""The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
cached: Final = cache.get(user_id)
if cached is not None:
return cached
seen_epoch: Final = cache.epoch()
assertion: Final = await _read_assertion_from_db(user_id)
if assertion is not None:
cache.set_if_unchanged(user_id, assertion, seen_epoch)
return assertion
class AssertionStoreUnavailable(Exception):
"""Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down).
@ -189,9 +248,12 @@ class DbSSOAssertionStore:
from credential resolution and from the upstream-401 retry.
"""
def __init__(self, cache: SSOAssertionCache = _ASSERTION_CACHE) -> None:
self._cache = cache
async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
try:
return await fetch_sso_identity_assertion(user_id)
return await fetch_sso_identity_assertion(user_id, cache=self._cache)
except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence
raise AssertionStoreUnavailable(str(exc)) from exc

View file

@ -2404,6 +2404,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"""
completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls")
max_in_flight_requests_per_worker: int | None = Field(
None, gt=0, description="maximum concurrent requests handled by each worker"
)
max_queued_requests_per_worker: int | None = Field(
None, ge=0, description="maximum requests waiting for a worker slot"
)
admission_queue_timeout_seconds: float = Field(
1.0, gt=0, description="maximum time a request waits for a worker slot"
)
plugins: list[PluginConfig] | None = Field(
None, description="external services registered as embeddable UI plugins"
)

View file

@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping
from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.anthropic.experimental_pass_through.context_management import (
AnthropicContextManagementError,
@ -30,6 +30,27 @@ from litellm.types.utils import TokenCountResponse
router: Final = APIRouter()
def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse:
from litellm.proxy.proxy_server import (
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
)
status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500
_close_dangling_otel_server_span(request, status_code, exc=exc)
envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error(
status_code=status_code,
raw_message=exc.message,
request_id=request.headers.get("x-request-id"),
)
if not exc.provider_specific_fields:
return JSONResponse(status_code=status_code, content=envelope, headers=exc.headers)
content: Final[AnthropicErrorResponse] = {
**envelope,
"error": {**envelope["error"], "provider_specific_fields": exc.provider_specific_fields},
}
return JSONResponse(status_code=status_code, content=content, headers=exc.headers)
def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
"""Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM
injects into Anthropic /v1/messages responses.
@ -195,7 +216,7 @@ async def anthropic_response(
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
if isinstance(e, ProxyException):
raise
return _anthropic_error_json_response(e, request)
# Extract model_id from request metadata (same as success path)
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@ -216,15 +237,18 @@ async def anthropic_response(
)
if isinstance(e, HTTPException):
raise proxy_exception_from_http_exception(e, headers)
return _anthropic_error_json_response(proxy_exception_from_http_exception(e, headers), request)
error_msg: Final = f"{e}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", 500),
headers=headers,
return _anthropic_error_json_response(
ProxyException(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", 500),
headers=headers,
),
request,
)

View file

@ -50,6 +50,9 @@ from litellm.proxy.health_check import (
perform_health_check,
run_with_timeout,
)
from litellm.proxy.middleware.admission_control_middleware import (
get_admission_control_stats,
)
from litellm.proxy.middleware.in_flight_requests_middleware import (
get_in_flight_requests,
)
@ -63,6 +66,13 @@ from litellm.secret_managers.main import get_secret_bool
#### Health ENDPOINTS ####
class _HealthBacklogResponse(TypedDict):
in_flight_requests: ReadOnly[int]
admitted_requests: ReadOnly[int]
queued_requests: ReadOnly[int]
rejected_requests: ReadOnly[int]
def _reject_os_environ_references(params: dict) -> None:
"""
Validate that the provided params do not contain any ``os.environ/``
@ -1759,7 +1769,14 @@ async def health_backlog():
for the event loop to get to them, adding latency before LiteLLM even starts
its own timer.
"""
return {"in_flight_requests": get_in_flight_requests()}
stats: Final = get_admission_control_stats()
response: Final[_HealthBacklogResponse] = {
"in_flight_requests": get_in_flight_requests(),
"admitted_requests": stats.admitted,
"queued_requests": stats.queued,
"rejected_requests": stats.rejected_total,
}
return response
@router.get(

View file

@ -0,0 +1,315 @@
import asyncio
import os
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from functools import lru_cache
from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable
from pydantic import Field, TypeAdapter, ValidationError
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Receive, Scope, Send
from litellm._logging import verbose_proxy_logger
_EXEMPT_PATHS: Final[frozenset[str]] = frozenset(
{
"/health/liveliness",
"/health/liveness",
"/health/readiness",
"/health/readiness/details",
"/health/backlog",
"/health/drain",
"/metrics",
"/metrics/",
}
)
@dataclass(frozen=True, slots=True)
class AdmissionControlSettings:
max_in_flight_requests: int
max_queued_requests: int
queue_timeout_seconds: float
AdmissionControlSettingsGetter: TypeAlias = Callable[[], AdmissionControlSettings | None] # mutable-ok: Callable params
@dataclass(frozen=True, slots=True)
class AdmissionControlStats:
admitted: int
queued: int
rejected_total: int
@runtime_checkable
class _Gauge(Protocol):
def inc(self, amount: float = 1) -> None: ...
def dec(self, amount: float = 1) -> None: ...
@runtime_checkable
class _CounterChild(Protocol):
def inc(self, amount: float = 1) -> None: ...
@runtime_checkable
class _Counter(Protocol):
def labels(self, reason: str) -> _CounterChild: ...
@dataclass(frozen=True, slots=True)
class AdmissionControlMetrics:
admitted_gauge: _Gauge
queued_gauge: _Gauge
rejected_counter: _Counter
AdmissionControlMetricsFactory: TypeAlias = Callable[[], AdmissionControlMetrics | None] # mutable-ok: Callable params
class AdmissionControlState:
"""Per-process admission counters and the in-flight semaphore shared by one worker's requests."""
def __init__(self, metrics_factory: AdmissionControlMetricsFactory) -> None:
self._metrics_factory = metrics_factory
self._metrics: AdmissionControlMetrics | None = None
self._metrics_init_attempted = False
self._admitted = 0
self._queued = 0
self._rejected_total = 0
self._semaphore: asyncio.Semaphore | None = None
self._semaphore_loop: asyncio.AbstractEventLoop | None = None
def get_stats(self) -> AdmissionControlStats:
return AdmissionControlStats(
admitted=self._admitted,
queued=self._queued,
rejected_total=self._rejected_total,
)
def get_semaphore(self, max_in_flight_requests: int) -> asyncio.Semaphore:
loop: Final = asyncio.get_running_loop()
if self._semaphore_loop is not loop:
self._semaphore = asyncio.Semaphore(max_in_flight_requests)
self._semaphore_loop = loop
semaphore: Final = self._semaphore
if semaphore is None:
raise RuntimeError("Admission control semaphore was not initialized")
return semaphore
def record_admission(self) -> None:
self._admitted += 1
metrics: Final = self._get_metrics()
if metrics is not None:
metrics.admitted_gauge.inc()
def record_release(self) -> None:
self._admitted -= 1
metrics: Final = self._get_metrics()
if metrics is not None:
metrics.admitted_gauge.dec()
def record_queue(self) -> None:
self._queued += 1
metrics: Final = self._get_metrics()
if metrics is not None:
metrics.queued_gauge.inc()
def record_dequeue(self) -> None:
self._queued -= 1
metrics: Final = self._get_metrics()
if metrics is not None:
metrics.queued_gauge.dec()
def record_rejection(self, reason: str) -> None:
self._rejected_total += 1
metrics: Final = self._get_metrics()
if metrics is not None:
metrics.rejected_counter.labels(reason=reason).inc()
def _get_metrics(self) -> AdmissionControlMetrics | None:
if not self._metrics_init_attempted:
self._metrics_init_attempted = True
self._metrics = self._metrics_factory()
return self._metrics
class AdmissionControlMiddleware:
def __init__(
self,
app: ASGIApp,
get_settings: AdmissionControlSettingsGetter,
state: AdmissionControlState,
) -> None:
self.app = app
self.get_settings = get_settings
self.state = state
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
settings: Final = self.get_settings()
if settings is None or _get_route_path(scope) in _EXEMPT_PATHS:
await self.app(scope, receive, send)
return
state: Final = self.state
semaphore: Final = state.get_semaphore(settings.max_in_flight_requests)
if not semaphore.locked():
await semaphore.acquire()
state.record_admission()
elif state.get_stats().queued >= settings.max_queued_requests:
state.record_rejection("queue_full")
await _overloaded_response(state)(scope, receive, send)
return
else:
state.record_queue()
try:
await asyncio.wait_for(
semaphore.acquire(),
timeout=settings.queue_timeout_seconds,
)
except asyncio.TimeoutError:
state.record_dequeue()
state.record_rejection("queue_timeout")
await _overloaded_response(state)(scope, receive, send)
return
except asyncio.CancelledError:
state.record_dequeue()
raise
state.record_dequeue()
state.record_admission()
try:
await self.app(scope, receive, send)
finally:
semaphore.release()
state.record_release()
def _get_route_path(scope: Scope) -> str:
"""Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching."""
path: Final[str] = scope["path"]
root_path: Final[str] = scope.get("root_path", "")
if not root_path or not path.startswith(root_path):
return path
if path == root_path:
return ""
if path[len(root_path)] == "/":
return path[len(root_path) :]
return path
def _create_gauge(gauge_type: Callable[..., object], name: str, description: str) -> _Gauge:
metric: Final = (
gauge_type(name, description, multiprocess_mode="livesum")
if "PROMETHEUS_MULTIPROC_DIR" in os.environ
else gauge_type(name, description)
)
if not isinstance(metric, _Gauge):
raise TypeError("Admission gauge has an unexpected type")
return metric
def create_prometheus_admission_metrics() -> AdmissionControlMetrics | None:
try:
from prometheus_client import Counter, Gauge
return AdmissionControlMetrics(
admitted_gauge=_create_gauge(
Gauge,
"litellm_admission_admitted_requests",
"Number of requests admitted by this worker",
),
queued_gauge=_create_gauge(
Gauge,
"litellm_admission_queued_requests",
"Number of requests queued by this worker",
),
rejected_counter=Counter( # mutable-ok: Prometheus requires runtime Counter construction
"litellm_admission_rejected_requests_total",
"Number of requests rejected by this worker",
labelnames=("reason",),
),
)
except (ImportError, ValueError):
return None
admission_control_state: Final = AdmissionControlState(create_prometheus_admission_metrics)
def get_admission_control_stats() -> AdmissionControlStats:
return admission_control_state.get_stats()
_PositiveInt: TypeAlias = Annotated[int, Field(gt=0)]
_NonNegativeInt: TypeAlias = Annotated[int, Field(ge=0)]
_PositiveFloat: TypeAlias = Annotated[float, Field(gt=0)]
_AdmissionControlRaw: TypeAlias = int | float | str | None
def _hashable(value: object) -> _AdmissionControlRaw:
return value if value is None or isinstance(value, (int, float, str)) else repr(value)
_POSITIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_PositiveInt)
_NON_NEGATIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_NonNegativeInt)
_POSITIVE_FLOAT_ADAPTER: Final[TypeAdapter[float]] = TypeAdapter(_PositiveFloat)
@lru_cache(maxsize=16)
def _parse_admission_control_settings(
max_in_flight_raw: _AdmissionControlRaw,
max_queued_raw: _AdmissionControlRaw,
queue_timeout_raw: _AdmissionControlRaw,
) -> AdmissionControlSettings | None:
try:
max_in_flight: Final = _POSITIVE_INT_ADAPTER.validate_python(max_in_flight_raw)
max_queued: Final = (
max_in_flight if max_queued_raw is None else _NON_NEGATIVE_INT_ADAPTER.validate_python(max_queued_raw)
)
queue_timeout: Final = _POSITIVE_FLOAT_ADAPTER.validate_python(queue_timeout_raw)
except ValidationError as exc:
verbose_proxy_logger.error(
"Ignoring invalid admission control settings, per-worker admission control is disabled: %s",
exc,
)
return None
return AdmissionControlSettings(
max_in_flight_requests=max_in_flight,
max_queued_requests=max_queued,
queue_timeout_seconds=queue_timeout,
)
def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionControlSettings | None:
max_in_flight_raw: Final = settings.get("max_in_flight_requests_per_worker")
if max_in_flight_raw is None:
return None
return _parse_admission_control_settings(
_hashable(max_in_flight_raw),
_hashable(settings.get("max_queued_requests_per_worker")),
_hashable(settings.get("admission_queue_timeout_seconds", 1.0)),
)
def _overloaded_response(state: AdmissionControlState) -> JSONResponse:
stats: Final = state.get_stats()
return JSONResponse(
status_code=503,
headers={"retry-after": "1"}, # mutable-ok: Starlette expects a plain headers mapping
content={ # mutable-ok: Starlette serializes a plain response mapping
"error": { # mutable-ok: nested response mapping
"message": (
f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later."
),
"type": "overloaded_error",
"code": "503",
}
},
)

View file

@ -1659,6 +1659,12 @@ async def azure_proxy_route(
from abc import ABC, abstractmethod
_VERTEX_LOCATION_REQUIRED_DETAIL: Final = (
"No Vertex AI location for this request. Include /projects/<project>/locations/<location>/ in the "
"route, set vertex_location in default_vertex_config (or DEFAULT_VERTEXAI_LOCATION), or add the "
"model to model_list with use_in_pass_through: true."
)
class BaseVertexAIPassThroughHandler(ABC):
@staticmethod
@ -1666,29 +1672,18 @@ class BaseVertexAIPassThroughHandler(ABC):
def get_default_base_target_url(vertex_location: str | None) -> str:
pass
@staticmethod
@abstractmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
pass
class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
@staticmethod
def get_default_base_target_url(vertex_location: str | None) -> str:
return "https://discoveryengine.googleapis.com/"
@staticmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
return base_target_url
class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
@staticmethod
def get_default_base_target_url(vertex_location: str | None) -> str:
return get_vertex_base_url(vertex_location)
@staticmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
if vertex_location is None:
raise HTTPException(status_code=400, detail=_VERTEX_LOCATION_REQUIRED_DETAIL)
return get_vertex_base_url(vertex_location)
@ -1911,10 +1906,8 @@ async def _prepare_vertex_auth_headers(
router_credentials: LiteLLM_ManagedVectorStore | None,
vertex_project: str | None,
vertex_location: str | None,
base_target_url: str | None,
get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]:
) -> tuple[Mapping[str, str], bool, str | None, str | None]:
"""
Prepare authentication headers for Vertex AI pass-through requests.
@ -1924,15 +1917,12 @@ async def _prepare_vertex_auth_headers(
router_credentials: Optional vector store credentials from registry
vertex_project: Vertex project ID
vertex_location: Vertex location
base_target_url: Base URL for the Vertex AI service
get_vertex_pass_through_handler: Handler for the specific Vertex AI service
user_api_key_dict: The caller's resolved authentication, so only the secret that
authenticated them is stripped on the credential-less branch
Returns:
tuple containing:
- headers: dict - Authentication headers to use
- base_target_url: str | None - Updated base target URL
- headers_passed_through: bool - Whether headers were passed through from request
- vertex_project: str | None - Updated vertex project ID
- vertex_location: str | None - Updated vertex location
@ -1985,14 +1975,8 @@ async def _prepare_vertex_auth_headers(
# Add the Authorization header with vendor credentials
headers["Authorization"] = f"Bearer {auth_header}"
if base_target_url is not None:
base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location(
base_target_url, vertex_location
)
return (
headers,
base_target_url,
headers_passed_through,
vertex_project,
vertex_location,
@ -2085,12 +2069,9 @@ async def _base_vertex_proxy_route(
location=vertex_location,
)
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
# Prepare authentication headers
(
headers,
base_target_url,
headers_passed_through,
vertex_project,
vertex_location,
@ -2100,13 +2081,10 @@ async def _base_vertex_proxy_route(
router_credentials=router_credentials,
vertex_project=vertex_project,
vertex_location=vertex_location,
base_target_url=base_target_url,
get_vertex_pass_through_handler=get_vertex_pass_through_handler,
user_api_key_dict=user_api_key_dict,
)
if base_target_url is None:
base_target_url = get_vertex_base_url(vertex_location)
base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
request_route: Final = encoded_endpoint
verbose_proxy_logger.debug("request_route %s", request_route)

View file

@ -583,6 +583,11 @@ try:
except ImportError:
build_billing_metrics_recorder = None
shutdown_billing_metrics_recorder = None
from litellm.proxy.middleware.admission_control_middleware import (
AdmissionControlMiddleware,
admission_control_state,
get_admission_control_settings,
)
from litellm.proxy.middleware.in_flight_requests_middleware import (
InFlightRequestsMiddleware,
)
@ -16502,6 +16507,9 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
{
"max_parallel_requests": "Integer",
"global_max_parallel_requests": "Integer",
"max_in_flight_requests_per_worker": "Integer",
"max_queued_requests_per_worker": "Integer",
"admission_queue_timeout_seconds": "Float",
"max_request_size_mb": "Integer",
"max_batch_file_size_mb": "Integer",
"max_file_size_mb": "Integer",
@ -18177,6 +18185,11 @@ app.add_middleware(
get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
is_request_size_limit_enabled=lambda: premium_user is True,
)
app.add_middleware(
AdmissionControlMiddleware,
get_settings=lambda: get_admission_control_settings(general_settings),
state=admission_control_state,
)
async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":

View file

@ -3,7 +3,8 @@ import collections
import json
import os
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from datetime import date, datetime, timedelta, timezone
from itertools import groupby
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -16,7 +17,6 @@ from typing import (
TypeAlias,
TypedDict,
TypeVar,
cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings
)
import fastapi
@ -201,16 +201,12 @@ class _SessionSpendStats(NamedTuple):
_SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats]
class _SpendSumAggregate(TypedDict, total=False):
spend: ReadOnly[float]
class _SpendGroupByRow(TypedDict):
class _SpendDailySummaryRow(TypedDict):
day: ReadOnly[str]
api_key: ReadOnly[str]
user: ReadOnly[str | None]
model: ReadOnly[str]
startTime: ReadOnly[object]
_sum: ReadOnly[_SpendSumAggregate]
spend: ReadOnly[float]
async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
@ -251,6 +247,66 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken
return VerificationTokenRepository(prisma_client).table
def _spend_logs_daily_summary_sql(
*,
start_date_iso: str,
end_date_iso: str,
api_key: str | None,
request_id: str | None,
user_id: str | None,
) -> tuple[str, tuple[object, ...]]:
filter_params: Final[tuple[tuple[str, object], ...]] = tuple(
(column, value)
for column, value in (
("api_key", api_key),
("request_id", request_id),
('"user"', user_id),
)
if value is not None
)
filter_clauses: Final[tuple[str, ...]] = tuple(
f"AND {column} = ${index}" for index, (column, _) in enumerate(filter_params, start=3)
)
filter_sql: Final = "\n".join(filter_clauses)
sql_query: Final = f"""
SELECT
to_char(date_trunc('day', "startTime"), 'YYYY-MM-DD') AS day,
api_key,
"user",
model,
SUM(spend) AS spend
FROM "LiteLLM_SpendLogs"
WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND "startTime" <= ($2::timestamptz AT TIME ZONE 'UTC')
{filter_sql}
GROUP BY 1, 2, 3, 4
ORDER BY 1
"""
params: Final[tuple[object, ...]] = (
start_date_iso,
end_date_iso,
*(value for _, value in filter_params),
)
return sql_query, params
def _sum_spend_by(
rows: Sequence[_SpendDailySummaryRow], column: Literal["api_key", "user", "model"]
) -> Mapping[str | None, float]:
keys: Final = frozenset(row[column] for row in rows)
return {key: sum(float(row["spend"]) for row in rows if row[column] == key) for key in keys}
def _daily_summary_item(summary_date: date, rows: Sequence[_SpendDailySummaryRow]) -> Mapping[str, object]:
api_key_spend: Final = {key: value for key, value in _sum_spend_by(rows, "api_key").items() if key is not None}
return {
**api_key_spend,
"startTime": summary_date,
"spend": sum(float(row["spend"]) for row in rows),
"users": _sum_spend_by(rows, "user"),
"models": _sum_spend_by(rows, "model"),
}
async def _find_spend_logs(
prisma_client: PrismaClient,
where: Mapping[str, object],
@ -3266,18 +3322,22 @@ async def view_spend_logs(
start_date_iso: Final = start_date_obj.isoformat()
end_date_iso: Final = end_date_obj.isoformat()
filter_query: Final = {
filter_query: Final[
dict[str, object]
] = { # mutable-ok: legacy filters are extended for optional parameters
"startTime": {
"gte": start_date_iso, # Greater than or equal to Start Date
"lte": end_date_iso, # Less than or equal to End Date
}
}
summary_api_key: Final[str | None] = (
prisma_client.hash_token(token=api_key)
if api_key is not None and api_key.startswith("sk-")
else api_key
)
if api_key is not None and isinstance(api_key, str):
if api_key.startswith("sk-"):
filter_query["api_key"] = prisma_client.hash_token(token=api_key)
else:
filter_query["api_key"] = api_key
filter_query["api_key"] = summary_api_key
if request_id is not None and isinstance(request_id, str):
filter_query["request_id"] = request_id
if user_id is not None and isinstance(user_id, str):
@ -3296,58 +3356,34 @@ async def view_spend_logs(
return data
# Legacy behavior: return summarized data (when summarize=true)
# SQL query
response: Final = await SpendLogsRepository(prisma_client).table.group_by(
by=["api_key", "user", "model", "startTime"],
where=filter_query,
sum={
"spend": True,
},
summary_sql_and_params: Final = _spend_logs_daily_summary_sql(
start_date_iso=start_date_iso,
end_date_iso=end_date_iso,
api_key=summary_api_key,
request_id=request_id,
user_id=user_id,
)
sql_query, params = summary_sql_and_params
rows: Final[Sequence[_SpendDailySummaryRow]] = await _query_raw(prisma_client, sql_query, *params)
if len(rows) == 0:
return [] # pyright: ignore[reportUnknownVariableType] # empty summary has no element type
if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
spend_rows: Final = cast(Sequence[_SpendGroupByRow], response) # cast-ok: by/sum fix the shape
result: Final[dict] = {}
for record in spend_rows:
dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
date = dt_object.date()
if date not in result:
result[date] = {"users": {}, "models": {}}
api_key = record["api_key"]
user_id = record["user"]
model = record["model"]
result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0)
result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0)
result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get(
"spend", 0
)
result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get(
"spend", 0
)
return_list: Final = []
final_date = None
for k, v in sorted(result.items()):
return_list.append({**v, "startTime": k})
final_date = k
end_date_date: Final = end_date_obj.date()
if final_date is not None and final_date < end_date_date:
current_date = final_date + timedelta(days=1)
while current_date <= end_date_date:
# Represent current_date as string because original response has it this way
return_list.append(
{
"startTime": current_date,
"spend": 0,
"users": {},
"models": {},
}
) # If no data, will stay as zero
current_date += timedelta(days=1) # Move on to the next day
return return_list
return response
summary_items: Final = tuple(
_daily_summary_item(date.fromisoformat(day), tuple(day_rows))
for day, day_rows in groupby(rows, key=lambda row: row["day"])
)
final_date: Final = date.fromisoformat(rows[-1]["day"])
end_date_date: Final = end_date_obj.date()
padding: Final[tuple[Mapping[str, object], ...]] = tuple(
{
"startTime": final_date + timedelta(days=offset),
"spend": 0,
"users": {},
"models": {},
}
for offset in range(1, (end_date_date - final_date).days + 1)
)
return [*summary_items, *padding]
else:
scoped_filter: Final[dict[str, str]] = {}

View file

@ -9546,6 +9546,7 @@ class Router:
public_model_name for _, public_model_name in self.team_model_to_deployment_indices
)
self.pattern_router.remove_deployment(model_id)
for team_id in list(self.team_pattern_routers.keys()):
team_pattern_router = self.team_pattern_routers[team_id]
team_pattern_router.remove_deployment(model_id)

View file

@ -190,6 +190,9 @@ model_list:
# Let that replacement also override a kept session pin, for image turns only (default: false)
modality_pin_override: true
# Refreshes on every pin reuse, so this is idle time rather than total session length (default: 3600)
session_affinity_ttl_seconds: 300
```
## Usage
@ -240,6 +243,10 @@ affinity write happens upstream of the gate and stores the session's own model,
turn replays the original pin and the override is never pinned in its place. It does nothing
unless `modality_routing` is also on.
### Session pin retention
`session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds.
### Heuristic-first chaining
`classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM

View file

@ -96,7 +96,7 @@
"limit": 10
},
"DTZ007": {
"limit": 17
"limit": 6
},
"DTZ011": {
"limit": 3

View file

@ -0,0 +1,218 @@
import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage } from "../../helpers/navigation";
import { CHAT_MODEL_A, masterKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
const VIEWPORT = { width: 1280, height: 720 };
const SEED_ROWS = 40;
const LOG_ROWS = 20;
const BODY_SCROLL_PX = 500;
const MAX_FOOTER_GAP_PX = 40;
interface GeneratedKey {
key: string;
}
interface CreatedTeam {
team_id: string;
}
interface CreatedModel {
model_info: { id: string };
}
interface BoxMetrics {
top: number;
bottom: number;
scrollHeight: number;
clientHeight: number;
scrollWidth: number;
clientWidth: number;
}
const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
const adminHeaders = (): Record<string, string> => ({ Authorization: `Bearer ${masterKey()}` });
const appShellMain = (page: PlaywrightPage): Locator => page.locator("main").first();
const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first();
const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr");
const metrics = (locator: Locator): Promise<BoxMetrics> =>
locator.evaluate((el) => {
const rect = el.getBoundingClientRect();
return {
top: rect.top,
bottom: rect.bottom,
scrollHeight: el.scrollHeight,
clientHeight: el.clientHeight,
scrollWidth: el.scrollWidth,
clientWidth: el.clientWidth,
};
});
async function postOk<T>(request: APIRequestContext, path: string, data: Record<string, unknown>): Promise<T> {
const res = await request.post(path, { headers: adminHeaders(), data });
expect(res.ok(), `POST ${path} failed (${res.status()}): ${await res.text()}`).toBe(true);
return (await res.json()) as T;
}
const oneAtATime = <T>(count: number, call: (index: number) => Promise<T>): Promise<readonly T[]> =>
Array.from({ length: count }, (_, i) => i).reduce<Promise<readonly T[]>>(
async (previous, i) => [...(await previous), await call(i)],
Promise.resolve([]),
);
async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise<void> {
await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count);
}
async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise<void> {
await visibleTestId(page, "pagination-page-size").click();
await page.getByRole("option", { name: size, exact: true }).click();
}
async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void> {
const scroller = await metrics(appShellMain(page));
const body = visibleTestId(page, "data-table-scroller");
const bodyBefore = await metrics(body);
const headBefore = await metrics(visibleTestId(page, "data-table-head"));
const footer = await metrics(visibleDataTable(page));
expect(scroller.scrollHeight, "page scroller must not overflow vertically").toBe(scroller.clientHeight);
expect(scroller.scrollWidth, "page scroller must not overflow horizontally").toBe(scroller.clientWidth);
expect(bodyBefore.scrollHeight, "table body must be the element that scrolls").toBeGreaterThan(
bodyBefore.clientHeight,
);
expect(footer.bottom, "pagination footer must be inside the page").toBeLessThanOrEqual(scroller.bottom);
expect(scroller.bottom - footer.bottom, "pagination footer must sit at the bottom of the page").toBeLessThanOrEqual(
MAX_FOOTER_GAP_PX,
);
await body.evaluate((el, px) => {
el.scrollTop = px;
}, BODY_SCROLL_PX);
await expect.poll(() => body.evaluate((el) => el.scrollTop)).toBeGreaterThan(0);
const headAfter = await metrics(visibleTestId(page, "data-table-head"));
expect(Math.round(headAfter.top), "header must stay put while the body scrolls").toBe(Math.round(headBefore.top));
}
const rowsPaintingPastAnAncestor = (page: PlaywrightPage): Promise<string[]> =>
visibleDataTable(page)
.locator("table")
.evaluate((table) => {
const scrollsVertically = (el: Element): boolean =>
/auto|scroll/.test(getComputedStyle(el).overflowY) && el.scrollHeight > el.clientHeight + 1;
const boxesUpToTheScroller = (el: Element | null): Element[] =>
el === null || el === document.body || scrollsVertically(el)
? []
: [el, ...boxesUpToTheScroller(el.parentElement)];
const describe = (el: Element): string =>
`<${el.tagName.toLowerCase()} class="${el.getAttribute("class") ?? ""}">`;
return Array.from(table.querySelectorAll("tbody tr")).flatMap((row, index) => {
const rowBottom = row.getBoundingClientRect().bottom;
return boxesUpToTheScroller(row.parentElement)
.filter((box) => rowBottom > box.getBoundingClientRect().bottom + 1)
.map(
(box) =>
`row ${index} bottom ${Math.round(rowBottom)} past ${describe(box)} bottom ${Math.round(box.getBoundingClientRect().bottom)}`,
);
});
});
test.describe("Admin tables scroll inside the page", () => {
test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT });
test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({
page,
request,
}) => {
const suffix = uniqueSuffix();
const keys = await oneAtATime(SEED_ROWS, (i) =>
postOk<GeneratedKey>(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
);
try {
await navigateToPage(page, Page.ApiKeys);
await expectRowsAtLeast(page, SEED_ROWS);
await expectBodyIsTheOnlyScroller(page);
} finally {
await request.post("/key/delete", { headers: adminHeaders(), data: { keys: keys.map((k) => k.key) } });
}
});
test("Teams: rows scroll under a sticky header and the page itself never scrolls", async ({ page, request }) => {
const suffix = uniqueSuffix();
const teams = await oneAtATime(SEED_ROWS, (i) =>
postOk<CreatedTeam>(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
);
try {
await navigateToPage(page, Page.Teams);
await expectRowsAtLeast(page, SEED_ROWS);
await expectBodyIsTheOnlyScroller(page);
} finally {
await request.post("/team/delete", { headers: adminHeaders(), data: { team_ids: teams.map((t) => t.team_id) } });
}
});
test("Request Logs: rows scroll under a sticky header and the page itself never scrolls", async ({
page,
request,
}) => {
const suffix = uniqueSuffix();
const ids = await oneAtATime(LOG_ROWS, (i) =>
sendChatCompletion(request, { model: CHAT_MODEL_A, prompt: `scroll ${suffix} ${i}` }),
);
await waitForSpendLog(request, ids[ids.length - 1]);
await navigateToPage(page, Page.Logs);
await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
await setRowsPerPage(page, "25");
await expectRowsAtLeast(page, LOG_ROWS);
await expectBodyIsTheOnlyScroller(page);
});
test("Tags: no row paints past the box it lives in", async ({ page, request }) => {
const suffix = uniqueSuffix();
const names = Array.from({ length: SEED_ROWS }, (_, i) => `e2e-scroll-tag-${suffix}-${i}`);
await oneAtATime(SEED_ROWS, (i) =>
postOk<unknown>(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }),
);
try {
await navigateToPage(page, Page.TagManagement);
await expectRowsAtLeast(page, SEED_ROWS);
expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
} finally {
await oneAtATime(SEED_ROWS, (i) =>
request.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } }),
);
}
});
test("Model Hub: no row paints past the box it lives in", async ({ page, request }) => {
const suffix = uniqueSuffix();
const models = await oneAtATime(SEED_ROWS, (i) =>
postOk<CreatedModel>(request, "/model/new", {
model_name: `e2e-scroll-model-${suffix}-${i}`,
litellm_params: {
model: "openai/fake-gpt-4",
api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`,
api_key: "fake-key",
},
}),
);
try {
await navigateToPage(page, Page.ModelHubTable);
await expectRowsAtLeast(page, SEED_ROWS);
expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
} finally {
await oneAtATime(SEED_ROWS, (i) =>
request.post("/model/delete", { headers: adminHeaders(), data: { id: models[i].model_info.id } }),
);
}
});
});

View file

@ -17,7 +17,7 @@ test.describe("Internal Users Search", () => {
test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => {
await goToInternalUsers(page);
const search = page.getByPlaceholder("Search by email…");
const search = page.getByPlaceholder("Search by email or ID…");
await expect(search).toBeVisible();
await search.fill("noteam@");

View file

@ -806,3 +806,91 @@ def test_guardrail_status_fields_computation():
)
assert status_fields_no_guardrail.get("llm_api_status") == "success"
assert status_fields_no_guardrail.get("guardrail_status") == "not_run"
@pytest.mark.parametrize(
"status, guardrail_information, expected_guardrail_status",
[
pytest.param(
"failure",
[
{"guardrail_status": "success"},
{"guardrail_status": "guardrail_intervened"},
],
"guardrail_intervened",
id="pre_call_success_before_blocker",
),
pytest.param(
"failure",
[
{"guardrail_status": "guardrail_intervened"},
{"guardrail_status": "success"},
],
"guardrail_intervened",
id="blocker_before_success",
),
pytest.param(
"failure",
[
{"guardrail_status": "success"},
{"guardrail_status": "guardrail_failed_to_respond"},
],
"guardrail_failed_to_respond",
id="failure_outranks_success",
),
pytest.param(
"failure",
[
{"guardrail_status": "guardrail_failed_to_respond"},
{"guardrail_status": "guardrail_intervened"},
],
"guardrail_intervened",
id="intervention_outranks_failure",
),
pytest.param(
"success",
[
{"guardrail_status": "success"},
{"guardrail_status": "success"},
],
"success",
id="all_success_stays_success",
),
pytest.param(
"failure",
[
{"guardrail_status": "some_new_status"},
{"guardrail_status": "blocked"},
],
"guardrail_intervened",
id="unknown_status_does_not_mask_blocker",
),
pytest.param(
"failure",
[
{"guardrail_status": {"unhashable": True}},
{"guardrail_status": "guardrail_intervened"},
],
"guardrail_intervened",
id="unhashable_status_is_skipped",
),
],
)
def test_guardrail_status_fields_severity_across_entries(
status, guardrail_information, expected_guardrail_status
):
"""
A blocked request must never be reported as a guardrail success.
With multiple guardrails on one request (e.g. a pre_call mask that passes,
then a post_call guardrail that blocks), entries are recorded in execution
order, so the earlier "success" entry must not shadow the later
"guardrail_intervened" entry: the aggregate takes the most severe status,
regardless of entry order.
"""
from litellm.litellm_core_utils.litellm_logging import _get_status_fields
fields = _get_status_fields(
status=status, guardrail_information=guardrail_information, error_str=None
)
assert fields.get("guardrail_status") == expected_guardrail_status

View file

@ -0,0 +1,150 @@
import asyncio
import os
import socket
import subprocess
import sys
import time
from pathlib import Path
from typing import Final
import httpx
import pytest
pytestmark = pytest.mark.skipif(
os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1",
reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark",
)
def _free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
listener.bind(("127.0.0.1", 0))
return int(listener.getsockname()[1])
def _percentile(values: list[float], percentile: float) -> float:
return sorted(values)[min(int(len(values) * percentile), len(values) - 1)]
@pytest.mark.asyncio
async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
fake_port: Final = _free_port()
proxy_port: Final = _free_port()
fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py"
config_path: Final = tmp_path / "saturation_config.yaml"
config_path.write_text(
f"""model_list:
- model_name: slow-endpoint
litellm_params:
model: openai/slow-endpoint
api_base: http://127.0.0.1:{fake_port}/v1
general_settings:
master_key: sk-saturation
max_in_flight_requests_per_worker: 8
max_queued_requests_per_worker: 8
admission_queue_timeout_seconds: 0.5
"""
)
fake_process: Final = subprocess.Popen(
[sys.executable, str(fake_script), "--host", "127.0.0.1", "--port", str(fake_port)],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
try:
proxy_process: Final = subprocess.Popen(
[
sys.executable,
"-m",
"litellm.proxy.proxy_cli",
"--config",
str(config_path),
"--run_granian",
"--num_workers",
"1",
"--port",
str(proxy_port),
],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
try:
async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{proxy_port}") as client:
deadline: Final = time.monotonic() + 60
while time.monotonic() < deadline:
try:
response: Final = await client.get("/health/liveliness", timeout=2)
if response.status_code == 200:
break
except httpx.HTTPError:
pass
await asyncio.sleep(0.25)
else:
raise AssertionError("Granian proxy did not become healthy")
liveness_latencies: Final[list[float]] = []
stop_sampling: Final = asyncio.Event()
async def sample_liveness() -> None:
while not stop_sampling.is_set():
start: Final = time.perf_counter()
try:
response = await client.get("/health/liveliness", timeout=2)
response.raise_for_status()
liveness_latencies.append(time.perf_counter() - start)
except httpx.HTTPError:
pass
await asyncio.sleep(0.05)
async def send_completion() -> tuple[int, float, bool]:
start: Final = time.perf_counter()
response = await client.post(
"/chat/completions",
headers={"Authorization": "Bearer sk-saturation"},
json={
"model": "slow-endpoint",
"messages": [{"role": "user", "content": "hello"}],
},
timeout=10,
)
return response.status_code, time.perf_counter() - start, "retry-after" in response.headers
sampler: Final = asyncio.create_task(sample_liveness())
results: Final = await asyncio.gather(*(send_completion() for _ in range(200)))
stop_sampling.set()
await sampler
statuses: Final = [result[0] for result in results]
latencies: Final = [result[1] for result in results]
rejected: Final = [result for result in results if result[0] == 503]
assert set(statuses) <= {200, 503}
assert rejected
assert all(result[2] for result in rejected)
assert _percentile(latencies, 0.99) < 5
assert liveness_latencies
assert _percentile(liveness_latencies, 0.95) < 0.5
duration: Final = max(latencies)
print(
"\nmetric value\n"
f"rps {len(results) / duration:.2f}\n"
f"200 count {statuses.count(200)}\n"
f"503 count {statuses.count(503)}\n"
f"p50 {_percentile(latencies, 0.50):.3f}s\n"
f"p95 {_percentile(latencies, 0.95):.3f}s\n"
f"p99 {_percentile(latencies, 0.99):.3f}s\n"
f"liveness p95 {_percentile(liveness_latencies, 0.95):.3f}s"
)
finally:
proxy_process.terminate()
try:
proxy_process.wait(timeout=10)
except subprocess.TimeoutExpired:
proxy_process.kill()
proxy_process.wait()
finally:
fake_process.terminate()
try:
fake_process.wait(timeout=10)
except subprocess.TimeoutExpired:
fake_process.kill()
fake_process.wait()

View file

@ -5,20 +5,26 @@ CLIENT PAUSE) showed 100% of concurrent commands to the other two, untouched nod
stalling for the full pause duration before this fix, and zero after -- these tests pin
the same behavior at the unit level so it can run without a live Redis Cluster."""
import asyncio
from typing import TYPE_CHECKING
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, Mock, patch
import pytest
from redis.exceptions import (
AskError,
BusyLoadingError,
ClusterDownError,
ClusterError,
MaxConnectionsError,
MovedError,
TryAgainError,
)
from redis.exceptions import (
ConnectionError as RedisConnectionError,
)
from redis.exceptions import TimeoutError as RedisTimeoutError
from redis.exceptions import (
TimeoutError as RedisTimeoutError,
)
from litellm.caching.redis_cluster_node_isolation import (
get_litellm_async_redis_cluster_class,
@ -39,10 +45,35 @@ class _NodeClassWithoutPerConnectionRecovery:
class _FakeClusterNode:
def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None:
self.name = name
self.execute_command = AsyncMock(side_effect=raises, return_value=response)
async def execute_command(*args: object, **kwargs: object) -> object:
await asyncio.sleep(0)
if raises is not None:
raise raises
return response
self.execute_command = AsyncMock(side_effect=execute_command)
self.disconnect = AsyncMock()
class _Fake8xRedisCluster:
def __init__(self) -> None:
self._initialize = False
async def _execute_command(
self, target_node: _FakeClusterNode, *args: object, **kwargs: object
) -> object:
try:
return await target_node.execute_command(*args, **kwargs)
except (RedisConnectionError, RedisTimeoutError):
self._initialize = True
await asyncio.sleep(0)
raise
async def aclose(self) -> None:
self._initialize = True
class _FakeNodesManager:
def __init__(self, node_to_return: _FakeClusterNode) -> None:
self._moved_exception: object = None
@ -68,31 +99,198 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType":
return instance
def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None:
"""Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level
connection error per-connection, the factory must NOT install the copied override,
whose node.disconnect() also kills connections other coroutines are mid-operation on."""
from redis.asyncio.cluster import RedisCluster
def _build_8x_cluster_instance() -> _Fake8xRedisCluster:
cluster_cls = get_litellm_async_redis_cluster_class(
cluster_node_class=_NodeClassWithPerConnectionRecovery
cluster_node_class=_NodeClassWithPerConnectionRecovery,
base_cluster_class=_Fake8xRedisCluster,
)
return cluster_cls()
def test_unverified_redis_version_logs_warning(caplog: pytest.LogCaptureFixture) -> None:
import redis
with patch.object(redis, "__version__", "8.0.1"):
get_litellm_async_redis_cluster_class(cluster_node_class=_NodeClassWithoutPerConnectionRecovery)
assert "not in the set this cluster-teardown-storm fix was verified against" in caplog.text
@pytest.mark.asyncio
async def test_single_timeout_does_not_request_topology_reinit() -> None:
error = RedisTimeoutError("timeout")
target_node = _FakeClusterNode("node-a")
target_node.execute_command.side_effect = error
instance = _build_8x_cluster_instance()
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is error
assert instance._initialize is False
@pytest.mark.asyncio
async def test_connection_error_preserves_upstream_topology_reinit() -> None:
error = RedisConnectionError("connection error")
target_node = _FakeClusterNode("node-a")
target_node.execute_command.side_effect = error
instance = _build_8x_cluster_instance()
with pytest.raises(RedisConnectionError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is error
assert instance._initialize is True
@pytest.mark.asyncio
async def test_three_consecutive_timeouts_request_topology_reinit_and_reset_counter() -> None:
errors = [
RedisTimeoutError("timeout-1"),
RedisTimeoutError("timeout-2"),
RedisTimeoutError("timeout-3"),
]
fourth_error = RedisTimeoutError("timeout-4")
target_node = _FakeClusterNode("node-a")
target_node.execute_command.side_effect = [*errors, fourth_error]
instance = _build_8x_cluster_instance()
for error in errors:
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is error
assert instance._initialize is True
instance._initialize = False
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is fourth_error
assert instance._initialize is False
@pytest.mark.asyncio
async def test_success_resets_consecutive_timeout_counter() -> None:
errors = [RedisTimeoutError("timeout-1"), RedisTimeoutError("timeout-2")]
final_error = RedisTimeoutError("timeout-3")
target_node = _FakeClusterNode("node-a")
target_node.execute_command.side_effect = [*errors, b"value", final_error]
instance = _build_8x_cluster_instance()
for error in errors:
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is error
assert instance._initialize is False
result = await instance._execute_command(target_node, "GET", "k")
assert result == b"value"
assert instance._initialize is False
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is final_error
assert instance._initialize is False
@pytest.mark.asyncio
async def test_timeout_counters_are_per_node() -> None:
node_a_errors = [RedisTimeoutError("node-a-1"), RedisTimeoutError("node-a-2")]
node_b_error = RedisTimeoutError("node-b-1")
node_a = _FakeClusterNode("node-a")
node_b = _FakeClusterNode("node-b")
node_a.execute_command.side_effect = node_a_errors
node_b.execute_command.side_effect = node_b_error
instance = _build_8x_cluster_instance()
for target_node, error in (
(node_a, node_a_errors[0]),
(node_b, node_b_error),
(node_a, node_a_errors[1]),
):
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is error
assert instance._initialize is False
@pytest.mark.asyncio
async def test_timeout_does_not_clear_concurrent_topology_reinit_request() -> None:
error = RedisTimeoutError("timeout")
instance = _build_8x_cluster_instance()
async def request_reinit(*args: object, **kwargs: object) -> object:
await instance.aclose()
raise error
target_node = _FakeClusterNode("node-a")
target_node.execute_command.side_effect = request_reinit
with pytest.raises(RedisTimeoutError) as exc_info:
await instance._execute_command(target_node, "GET", "k")
assert exc_info.value is error
assert instance._initialize is True
@pytest.mark.asyncio
async def test_tolerated_timeout_does_not_erase_concurrent_connection_error_reinit() -> None:
instance = _build_8x_cluster_instance()
failing_node = _FakeClusterNode("node-a", raises=RedisConnectionError("gone"))
slow_node = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow"))
results = await asyncio.gather(
instance._execute_command(failing_node, "GET", "a"),
instance._execute_command(slow_node, "GET", "b"),
return_exceptions=True,
)
assert cluster_cls is RedisCluster
assert isinstance(results[0], RedisConnectionError)
assert isinstance(results[1], RedisTimeoutError)
assert instance._initialize is True
def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None:
"""Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(),
so those versions must keep litellm's per-node isolation override."""
from redis.asyncio.cluster import RedisCluster
@pytest.mark.asyncio
async def test_overlapping_tolerated_timeouts_do_not_request_topology_reinit() -> None:
instance = _build_8x_cluster_instance()
node_a = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow-a"))
node_b = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow-b"))
cluster_cls = get_litellm_async_redis_cluster_class(
cluster_node_class=_NodeClassWithoutPerConnectionRecovery
results = await asyncio.gather(
instance._execute_command(node_a, "GET", "a"),
instance._execute_command(node_b, "GET", "b"),
return_exceptions=True,
)
assert cluster_cls is not RedisCluster
assert issubclass(cluster_cls, RedisCluster)
assert "_execute_command" in cluster_cls.__dict__
assert all(isinstance(result, RedisTimeoutError) for result in results)
assert instance._initialize is False
@pytest.mark.asyncio
async def test_tolerated_timeout_does_not_clear_pending_reinit() -> None:
instance = _build_8x_cluster_instance()
instance._initialize = True
target_node = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow"))
with pytest.raises(RedisTimeoutError):
await instance._execute_command(target_node, "GET", "k")
assert instance._initialize is True
@pytest.mark.asyncio
async def test_success_returns_value_without_topology_reinit() -> None:
target_node = _FakeClusterNode("node-a", response=b"value")
instance = _build_8x_cluster_instance()
result = await instance._execute_command(target_node, "GET", "k")
assert result == b"value"
assert instance._initialize is False
@pytest.mark.asyncio
@ -110,6 +308,48 @@ async def test_node_level_error_resets_only_that_node_not_the_whole_client(error
instance.aclose.assert_not_awaited()
@pytest.mark.asyncio
async def test_moved_error_retries_without_full_reinit_before_threshold() -> None:
moved_error = MovedError("1 127.0.0.1:7001")
target_node = _FakeClusterNode("node-a")
target_node.execute_command = AsyncMock(side_effect=[moved_error, b"value"])
instance = _build_cluster_instance()
instance.RedisClusterRequestTTL = 2
instance.nodes_manager = _FakeNodesManager(node_to_return=target_node)
instance._determine_slot = AsyncMock(return_value=0)
result = await instance._execute_command(target_node, "GET", "k")
assert result == b"value"
assert instance.nodes_manager._moved_exception is moved_error
instance.aclose.assert_not_awaited()
@pytest.mark.asyncio
async def test_ask_error_sends_asking_and_retries_on_redirected_node() -> None:
ask_error = AskError("0 127.0.0.1:7001")
target_node = _FakeClusterNode("node-a")
target_node.execute_command = AsyncMock(side_effect=[ask_error, None, b"value"])
instance = _build_cluster_instance()
instance.RedisClusterRequestTTL = 2
instance.get_node = Mock(return_value=target_node)
result = await instance._execute_command(target_node, "GET", "k")
assert result == b"value"
instance.get_node.assert_called_once_with(node_name="127.0.0.1:7001")
@pytest.mark.asyncio
async def test_try_again_error_exhausts_ttl() -> None:
target_node = _FakeClusterNode("node-a", raises=TryAgainError("try again"))
instance = _build_cluster_instance()
instance.RedisClusterRequestTTL = 2
with pytest.raises(ClusterError):
await instance._execute_command(target_node, "GET", "k")
@pytest.mark.asyncio
async def test_successful_command_touches_neither_disconnect_nor_aclose() -> None:
target_node = _FakeClusterNode("node-a", response=b"v")

View file

@ -17,7 +17,7 @@ import pytest
import litellm
from litellm import completion, acompletion
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.snowflake.chat.transformation import SnowflakeConfig
from litellm.llms.snowflake.chat.transformation import SnowflakeConfig, SnowflakeStreamingHandler
from litellm.types.utils import ModelResponse
@ -114,8 +114,7 @@ class TestSnowflakeToolTransformation:
)
assert transformed_request["tool_choice"] == value, (
f"tool_choice='{value}' should pass through unchanged, "
f"got {transformed_request['tool_choice']}"
f"tool_choice='{value}' should pass through unchanged, got {transformed_request['tool_choice']}"
)
def test_transform_response_with_tool_calls(self):
@ -159,9 +158,7 @@ class TestSnowflakeToolTransformation:
headers={"Content-Type": "application/json"},
)
model_response = ModelResponse(
choices=[litellm.Choices(index=0, message=litellm.Message())]
)
model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
logging_obj = MagicMock()
@ -232,9 +229,7 @@ class TestSnowflakeToolTransformation:
headers={"Content-Type": "application/json"},
)
model_response = ModelResponse(
choices=[litellm.Choices(index=0, message=litellm.Message())]
)
model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
logging_obj = MagicMock()
@ -280,9 +275,7 @@ class TestSnowflakeToolTransformation:
headers={"Content-Type": "application/json"},
)
model_response = ModelResponse(
choices=[litellm.Choices(index=0, message=litellm.Message())]
)
model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
logging_obj = MagicMock()
@ -300,10 +293,7 @@ class TestSnowflakeToolTransformation:
# Verify standard response works
assert isinstance(result, ModelResponse)
assert (
result.choices[0].message.content
== "Hello! I'm doing well, thank you for asking."
)
assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking."
def test_get_supported_openai_params_includes_tools(self):
"""
@ -318,6 +308,385 @@ class TestSnowflakeToolTransformation:
assert "max_tokens" in supported_params
class TestSnowflakeCortexClaudeFixes:
def setup_method(self):
self.config = SnowflakeConfig()
@staticmethod
def _transform(messages, optional_params=None):
return SnowflakeConfig().transform_request(
model="snowflake/claude-sonnet-4-6",
messages=messages,
optional_params=optional_params or {},
litellm_params={},
headers={},
)
def test_thinking_is_offered_on_every_claude_model(self):
"""Cortex documents extended thinking (budget_tokens) for Claude generally, so a
4.6-only gate would silently drop it on the models that do support it."""
for model in (
"snowflake/claude-sonnet-4-6",
"snowflake/claude-sonnet-4-5",
"snowflake/claude-3-7-sonnet",
"snowflake/claude-4-opus",
):
assert "thinking" in self.config.get_supported_openai_params(model), model
assert "thinking" not in self.config.get_supported_openai_params("snowflake/llama3.1-70b")
def test_system_blocks_preserve_cache_control_and_strip_ttl(self):
body = self._transform(
[
{
"role": "system",
"content": [
{
"type": "text",
"text": "You are helpful",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
},
{"role": "user", "content": "hi"},
]
)
assert body["system"] == [{"type": "text", "text": "You are helpful", "cache_control": {"type": "ephemeral"}}]
def test_direct_system_param_is_normalized(self):
body = self._transform(
[{"role": "user", "content": "hi"}],
{"system": [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral", "ttl": "1h"}}]},
)
assert body["system"] == [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral"}}]
def test_message_and_tool_cache_control_are_normalized(self):
body = self._transform(
[
{
"role": "user",
"content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}],
}
],
{
"tools": [
{
"name": "f",
"input_schema": {"type": "object", "properties": {}},
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
]
},
)
assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"}
assert body["tools"][0]["cache_control"] == {"type": "ephemeral"}
def test_extra_body_message_override_is_normalized(self):
body = self._transform(
[{"role": "user", "content": "original"}],
{
"extra_body": {
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": "override",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
}
]
}
},
)
assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"}
def test_image_blocks_are_converted_to_anthropic_source(self):
body = self._transform(
[
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,ZmFrZQ==", "format": "image/jpeg"},
}
],
}
]
)
assert body["messages"][0]["content"] == [
{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "ZmFrZQ=="}}
]
def test_tool_result_image_list_is_converted(self):
body = self._transform(
[
{"role": "user", "content": "look"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
],
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}}],
},
]
)
assert body["messages"][2]["content"][0]["content"] == [
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}}
]
def test_tool_result_preserves_cache_control(self):
"""A cache breakpoint the bridge puts on a tool message must survive onto the tool_result."""
for tool_content in ("done", [{"type": "text", "text": "done"}]):
body = self._transform(
[
{"role": "user", "content": "look"},
{
"role": "tool",
"tool_call_id": "call_1",
"content": tool_content,
"cache_control": {"type": "ephemeral", "ttl": "1h"},
},
]
)
tool_result = body["messages"][1]["content"][0]
assert tool_result["cache_control"] == {"type": "ephemeral"}, tool_content
def test_pdf_data_uri_becomes_a_document_block(self):
"""A bridged pdf data URI is a document block; forwarding it as an image is malformed."""
body = self._transform(
[
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": "data:application/pdf;base64,ZmFrZQ=="}},
],
}
]
)
assert body["messages"][0]["content"] == [
{
"type": "document",
"source": {"type": "base64", "media_type": "application/pdf", "data": "ZmFrZQ=="},
}
]
def test_multipart_tool_result_preserves_text_and_converts_image(self):
body = self._transform(
[
{"role": "user", "content": "look"},
{
"role": "tool",
"tool_call_id": "call_1",
"content": [
{"type": "text", "text": "first"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}},
{"type": "text", "text": "last"},
],
},
]
)
assert body["messages"][1]["content"][0]["content"] == [
{"type": "text", "text": "first"},
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}},
{"type": "text", "text": "last"},
]
def test_plain_text_tool_result_remains_string(self):
body = self._transform(
[{"role": "user", "content": "look"}, {"role": "tool", "tool_call_id": "call_1", "content": "done"}]
)
assert body["messages"][1]["content"][0]["content"] == "done"
def test_anthropic_tool_schema_strips_only_top_level_schema_key(self):
tools = [
{
"name": "f",
"input_schema": {"$schema": "schema", "type": "object", "properties": {"$schema": {"type": "string"}}},
}
]
body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools})
schema = body["tools"][0]["input_schema"]
assert "$schema" not in schema
assert "$schema" in schema["properties"]
def test_tool_schema_strips_only_top_level_schema_key(self):
tools = [
{
"type": "function",
"function": {
"name": "f",
"parameters": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {"$schema": {"type": "string"}},
},
},
}
]
body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools})
schema = body["tools"][0]["input_schema"]
assert "$schema" not in schema
assert "$schema" in schema["properties"]
def test_streaming_tool_identity_is_emitted_only_on_start(self):
handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True)
start = handler.chunk_parser(
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "tool_use", "id": "tool_1", "name": "read"},
}
)
first_delta = handler.chunk_parser(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"path":'},
}
)
second_delta = handler.chunk_parser(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '"/tmp"}'},
}
)
def _tool_call(chunk):
return chunk.choices[0].delta.tool_calls[0]
assert _tool_call(start).id == "tool_1"
assert _tool_call(start).function.name == "read"
assert _tool_call(first_delta).id is None
assert _tool_call(first_delta).function.name is None
assert _tool_call(second_delta).id is None
assert _tool_call(second_delta).function.name is None
assert _tool_call(first_delta).function.arguments == '{"path":'
assert _tool_call(second_delta).function.arguments == '"/tmp"}'
def test_signed_thinking_blocks_lead_the_assistant_turn(self):
"""Multi-turn tool use with thinking only works if the signed block is echoed back first."""
body = self._transform(
[
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"thinking_blocks": [
{"type": "thinking", "thinking": "391", "signature": "Eto"},
{"type": "thinking", "thinking": "unsigned"},
],
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
],
},
]
)
blocks = body["messages"][1]["content"]
assert blocks[0] == {"type": "thinking", "thinking": "391", "signature": "Eto"}
assert [b["type"] for b in blocks] == ["thinking", "tool_use"]
def test_signed_thinking_blocks_lead_a_plain_text_assistant_turn(self):
"""A thinking response without a tool call must also round-trip on the next request."""
body = self._transform(
[
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": "391",
"thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
},
{"role": "user", "content": "continue"},
]
)
assert body["messages"][1] == {
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "391", "signature": "Eto"},
{"type": "text", "text": "391"},
],
}
def test_signed_thinking_blocks_preserve_list_content(self):
"""Cached assistant text reaches this transform as a content list, not a string."""
body = self._transform(
[
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [{"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}}],
"thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
},
{"role": "user", "content": "continue"},
]
)
assert body["messages"][1] == {
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "391", "signature": "Eto"},
{"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}},
],
}
def test_thinking_only_assistant_turn_sends_no_empty_text_block(self):
"""Anthropic-shaped APIs reject empty text blocks, so a content-less thinking turn is thinking only."""
body = self._transform(
[
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
},
{"role": "user", "content": "continue"},
]
)
assert body["messages"][1]["content"] == [{"type": "thinking", "thinking": "391", "signature": "Eto"}]
def test_streaming_surfaces_thinking_and_prompt_cache_usage(self):
"""Cortex streams thinking deltas, signatures and cache counts; all must reach the caller."""
handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True)
handler.chunk_parser(
{
"type": "message_start",
"message": {"usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323}},
}
)
thinking = handler.chunk_parser(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "391"},
}
)
signature = handler.chunk_parser(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "signature_delta", "signature": "Eto"},
}
)
final = handler.chunk_parser(
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 8, "cache_read_input_tokens": 1323},
}
)
assert thinking.choices[0].delta.reasoning_content == "391"
assert signature.choices[0].delta.thinking_blocks[0]["signature"] == "Eto"
assert final.usage.prompt_tokens_details.cached_tokens == 1323
class TestSnowFlakeCompletion:
model_name = "mistral"
@ -380,10 +749,7 @@ class TestSnowFlakeCompletion:
# PAT key was used
post_kwargs = mock_post.call_args_list[-1][1]
assert "xxxxx" in post_kwargs["headers"]["Authorization"]
assert (
post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"]
== "PROGRAMMATIC_ACCESS_TOKEN"
)
assert post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN"
# account id was used
assert "AAAA-BBBB" in post_kwargs["url"]
@ -495,9 +861,7 @@ class TestSnowflakeChatCompletion:
)
mock_post.assert_called_once()
else:
with patch.object(
AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp
) as mock_post:
with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp) as mock_post:
response = asyncio.run(
acompletion(
model="snowflake/mistral-7b",
@ -580,8 +944,4 @@ class TestSnowflakeChatCompletion:
chunks_received = asyncio.run(_run())
assert len(chunks_received) > 0
content = "".join(
c.choices[0].delta.content
for c in chunks_received
if c.choices[0].delta.content
)
content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content)

View file

@ -338,7 +338,7 @@ class TestAnthropicConfigRequest:
litellm_params={},
headers={},
)
assert body["system"] == "You are helpful."
assert body["system"] == [{"type": "text", "text": "You are helpful."}]
assert all(m["role"] != "system" for m in body["messages"])
assert body["messages"][0] == {"role": "user", "content": "Hello"}
@ -422,6 +422,64 @@ class TestAnthropicConfigResponse:
assert result.usage.completion_tokens == 5
assert result.usage.total_tokens == 15
def test_prompt_cache_usage_is_surfaced(self):
"""Cortex reports cache creation/read counts; dropping them hides caching and bills cached input at full price."""
raw = httpx.Response(
200,
json={
"id": "msg_1",
"model": "claude-sonnet-4-6",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323, "cache_read_input_tokens": 0},
},
)
result = self.cfg.transform_response(
model="snowflake/claude-sonnet-4-6",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.usage.prompt_tokens == 1341
assert result.usage.prompt_tokens_details.cache_creation_tokens == 1323
assert result.usage.prompt_tokens_details.cached_tokens == 0
def test_thinking_block_and_signature_are_preserved(self):
"""The signature must survive so a client can echo the thinking block on the next turn."""
raw = httpx.Response(
200,
json={
"id": "msg_1",
"model": "claude-sonnet-4-6",
"content": [
{"type": "thinking", "thinking": "391", "signature": "Eto"},
{"type": "text", "text": "391"},
],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
},
)
result = self.cfg.transform_response(
model="snowflake/claude-sonnet-4-6",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
message = result.choices[0].message
assert message.content == "391"
assert message.reasoning_content == "391"
assert message.thinking_blocks[0]["signature"] == "Eto"
def test_stop_reason_end_turn_maps_to_stop(self):
raw = _make_anthropic_response()
result = self.cfg.transform_response(

View file

@ -7,6 +7,7 @@ round-trips exactly, a store failure never escapes into the login path, and a sa
rotation re-encrypts stored rows like the sibling per-user credential tables.
"""
import asyncio
import json
import os
import time
@ -16,8 +17,10 @@ import jwt as pyjwt
import pytest
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
_ASSERTION_CACHE,
AssertionStoreUnavailable,
DbSSOAssertionStore,
SSOAssertionCache,
assertion_from_sso_login,
ema_assertion_retention_enabled,
fetch_sso_identity_assertion,
@ -25,7 +28,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s
retain_sso_identity_assertion_for_ema,
rotate_sso_identity_assertions_master_key,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper
from litellm.types.mcp import MCPAuth
SALT_KEY = "test-salt-key-for-sso-assertion-tests-1234"
@ -38,6 +41,11 @@ def _set_salt_key(monkeypatch):
monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY)
@pytest.fixture(autouse=True)
def _flush_assertion_cache():
_ASSERTION_CACHE.flush()
def _make_id_token(exp_offset: int = 3600, iss: str = ISSUER) -> str:
return pyjwt.encode(
{"iss": iss, "sub": "u1", "exp": int(time.time()) + exp_offset},
@ -52,9 +60,7 @@ def _make_prisma(stored: dict, db_has_id_jag_server: bool = False):
``db_has_id_jag_server`` drives the retention gate's authoritative DB fallback;
it is wired explicitly so the gate never reads a truthy bare MagicMock."""
prisma = MagicMock()
prisma.db.litellm_mcpservertable.find_first = AsyncMock(
return_value=MagicMock() if db_has_id_jag_server else None
)
prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=MagicMock() if db_has_id_jag_server else None)
async def _upsert(where, data):
stored[where["user_id"]] = data["update"]["assertion_b64"]
@ -235,6 +241,78 @@ async def test_persist_overwrites_previous_login():
assert fetched.refresh_token is not None
@pytest.mark.asyncio
async def test_fetch_serves_second_read_from_cache_without_db_read():
stored = {}
prisma = _make_prisma(stored)
cache = SSOAssertionCache()
token = _make_id_token()
assertion = assertion_from_sso_login(token, "rt_1")
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion, cache=cache)
first = await fetch_sso_identity_assertion("user-a", cache=cache)
second = await fetch_sso_identity_assertion("user-a", cache=cache)
assert first is not None
assert second is not None
assert first.id_token.get_secret_value() == token
assert second.id_token.get_secret_value() == token
prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once()
@pytest.mark.asyncio
async def test_persist_busts_cache_so_relogin_is_visible_immediately():
stored = {}
prisma = _make_prisma(stored)
cache = SSOAssertionCache()
first_token = _make_id_token(exp_offset=100)
second_token = _make_id_token(exp_offset=7200)
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
first = await fetch_sso_identity_assertion("user-a", cache=cache)
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache)
second = await fetch_sso_identity_assertion("user-a", cache=cache)
assert first is not None
assert second is not None
assert first.id_token.get_secret_value() == first_token
assert second.id_token.get_secret_value() == second_token
@pytest.mark.asyncio
async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion():
stored = {}
prisma = _make_prisma(stored)
cache = SSOAssertionCache()
first_token = _make_id_token(exp_offset=100)
second_token = _make_id_token(exp_offset=7200)
db_read_started = asyncio.Event()
relogin_done = asyncio.Event()
unpaused_find_unique = prisma.db.litellm_ssoidentityassertion.find_unique
async def _paused_find_unique(where):
row = await unpaused_find_unique(where=where)
db_read_started.set()
await relogin_done.wait()
return row
prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache))
await db_read_started.wait()
await persist_sso_identity_assertion(
"user-a",
assertion_from_sso_login(second_token, "rt_new"),
cache=cache,
)
relogin_done.set()
raced = await racing_fetch
after = await fetch_sso_identity_assertion("user-a", cache=cache)
assert raced is not None
assert after is not None
assert raced.id_token.get_secret_value() == first_token
assert after.id_token.get_secret_value() == second_token
@pytest.mark.asyncio
async def test_fetch_missing_row_returns_none():
prisma = _make_prisma({})
@ -242,6 +320,18 @@ async def test_fetch_missing_row_returns_none():
assert await fetch_sso_identity_assertion("nobody") is None
@pytest.mark.asyncio
async def test_fetch_does_not_cache_a_missing_row():
prisma = _make_prisma({})
cache = SSOAssertionCache()
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
first = await fetch_sso_identity_assertion("nobody", cache=cache)
second = await fetch_sso_identity_assertion("nobody", cache=cache)
assert first is None
assert second is None
assert prisma.db.litellm_ssoidentityassertion.find_unique.await_count == 2
@pytest.mark.asyncio
async def test_fetch_undecryptable_row_returns_none():
prisma = _make_prisma({"user-a": "not-an-encrypted-blob"})
@ -251,13 +341,37 @@ async def test_fetch_undecryptable_row_returns_none():
@pytest.mark.asyncio
async def test_fetch_unparseable_payload_returns_none():
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
prisma = _make_prisma({"user-a": encrypt_value_helper("]]not json")})
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
assert await fetch_sso_identity_assertion("user-a") is None
@pytest.mark.asyncio
async def test_cached_assertion_expires_after_ttl():
stored = {}
prisma = _make_prisma(stored)
cache = SSOAssertionCache(ttl_seconds=1)
first_token = _make_id_token(exp_offset=100)
second_token = _make_id_token(exp_offset=7200)
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
first = await fetch_sso_identity_assertion("user-a", cache=cache)
await persist_sso_identity_assertion(
"user-a",
assertion_from_sso_login(second_token, "rt_new"),
cache=SSOAssertionCache(),
)
cached = await fetch_sso_identity_assertion("user-a", cache=cache)
time.sleep(1.1)
expired = await fetch_sso_identity_assertion("user-a", cache=cache)
assert first is not None
assert cached is not None
assert expired is not None
assert first.id_token.get_secret_value() == first_token
assert cached.id_token.get_secret_value() == first_token
assert expired.id_token.get_secret_value() == second_token
@pytest.mark.asyncio
async def test_retain_noop_when_no_id_jag_server():
stored = {}
@ -357,6 +471,24 @@ async def test_db_store_converts_a_driver_failure_into_assertion_store_unavailab
await DbSSOAssertionStore().fetch("alice")
@pytest.mark.asyncio
async def test_db_store_uses_injected_cache():
stored = {}
prisma = _make_prisma(stored)
cache = SSOAssertionCache()
token = _make_id_token()
with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary
await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache)
store = DbSSOAssertionStore(cache=cache)
first = await store.fetch("alice")
second = await store.fetch("alice")
assert first is not None
assert second is not None
assert first.id_token.get_secret_value() == token
assert second.id_token.get_secret_value() == token
prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once()
@pytest.mark.asyncio
async def test_db_store_returns_none_for_a_user_with_no_stored_assertion():
"""An absent row stays an absence, not an outage, so a user who never signed in still gets the

View file

@ -125,11 +125,12 @@ class TestBlockedResponseUsage:
mock_logging.post_call_failure_hook.assert_awaited_once()
class TestProxyExceptionPassthrough:
class TestProxyExceptionAnthropicEnvelope:
@pytest.mark.asyncio
async def test_anthropic_response_reraises_proxy_exception_unwrapped(self):
"""A 400 ProxyException from request validation must surface as-is,
not be re-wrapped into a code-500 ProxyException."""
async def test_anthropic_response_maps_proxy_exception_to_anthropic_envelope(self):
"""LIT-6468: a 400 ProxyException from request validation must surface as
Anthropic's documented {"type": "error", "error": {...}} envelope with the
original status and message, not the OpenAI {"error": {...}} envelope."""
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyErrorTypes, ProxyException
@ -140,6 +141,8 @@ class TestProxyExceptionPassthrough:
param="metadata",
code=400,
)
request = MagicMock()
request.headers = {"x-request-id": "req_test_6468"}
with (
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})),
@ -151,30 +154,61 @@ class TestProxyExceptionPassthrough:
patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
):
mock_logging.post_call_failure_hook = AsyncMock()
with pytest.raises(ProxyException) as exc_info:
await ep.anthropic_response(
fastapi_response=MagicMock(),
request=MagicMock(),
user_api_key_dict=MagicMock(),
)
response = await ep.anthropic_response(
fastapi_response=MagicMock(),
request=request,
user_api_key_dict=MagicMock(),
)
assert exc_info.value is exc
assert exc_info.value.code == "400"
assert exc_info.value.param == "metadata"
assert response.status_code == 400
body = json.loads(response.body)
assert body == {
"type": "error",
"error": {
"type": "invalid_request_error",
"message": "Invalid type for 'metadata': expected an object, but got a string instead.",
},
"request_id": "req_test_6468",
}
mock_logging.post_call_failure_hook.assert_awaited_once()
@pytest.mark.asyncio
async def test_anthropic_response_maps_429_to_rate_limit_error(self):
"""The Anthropic error type follows the status code (429 -> rate_limit_error),
and a code-less exception falls back to 500 api_error."""
import litellm.proxy.anthropic_endpoints.endpoints as ep
from litellm.proxy._types import ProxyException
request = MagicMock()
request.headers = {}
response = ep._anthropic_error_json_response(
ProxyException(message="Rate limit exceeded", type="rate_limit_error", param=None, code=429),
request,
)
assert response.status_code == 429
assert json.loads(response.body)["error"]["type"] == "rate_limit_error"
fallback = ep._anthropic_error_json_response(
ProxyException(message="boom", type="None", param=None, code=None),
request,
)
assert fallback.status_code == 500
assert json.loads(fallback.body)["error"]["type"] == "api_error"
class TestHttpExceptionDictDetail:
@pytest.mark.asyncio
async def test_anthropic_response_serializes_dict_detail_http_exception(self):
"""LIT-6466: a post_call guardrail's HTTPException(detail=<dict>) must
surface with a clean message plus provider_specific_fields, matching
/v1/chat/completions and /v1/responses, not the str() of the exception."""
"""LIT-6466 + LIT-6468: a post_call guardrail's HTTPException(detail=<dict>)
must surface as Anthropic's {"type": "error", "error": {...}} envelope with
the guardrail's clean message plus provider_specific_fields, not the str()
of the exception and not the OpenAI envelope."""
from fastapi import HTTPException
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy._types import UserAPIKeyAuth
detail = {
"error": "Content blocked: keyword 'kumquat' detected",
@ -182,6 +216,8 @@ class TestHttpExceptionDictDetail:
"guardrail": "keyword-block",
}
exc = HTTPException(status_code=400, detail=detail)
request = MagicMock()
request.headers = {}
with (
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
@ -193,17 +229,19 @@ class TestHttpExceptionDictDetail:
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
):
mock_logging.post_call_failure_hook = AsyncMock()
with pytest.raises(ProxyException) as exc_info:
await ep.anthropic_response(
fastapi_response=MagicMock(),
request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(),
)
response = await ep.anthropic_response(
fastapi_response=MagicMock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
assert exc_info.value.message == "Content blocked: keyword 'kumquat' detected"
assert "{'error'" not in exc_info.value.message
assert exc_info.value.provider_specific_fields == detail
assert exc_info.value.code == "400"
assert response.status_code == 400
body = json.loads(response.body)
assert body["type"] == "error"
assert body["error"]["type"] == "invalid_request_error"
assert body["error"]["message"] == "Content blocked: keyword 'kumquat' detected"
assert "{'error'" not in body["error"]["message"]
assert body["error"]["provider_specific_fields"] == detail
mock_logging.post_call_failure_hook.assert_awaited_once()
@ -215,7 +253,7 @@ class TestFailureHookRequestData:
handler must pass that replaced dict, not the raw request body dict."""
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy._types import UserAPIKeyAuth
captured = {}
@ -224,18 +262,23 @@ class TestFailureHookRequestData:
captured["processor_data"] = self.data
raise RuntimeError("provider timeout")
request = MagicMock()
request.headers = {}
with (
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})),
patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process),
patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
):
mock_logging.post_call_failure_hook = AsyncMock()
with pytest.raises(ProxyException):
await ep.anthropic_response(
fastapi_response=MagicMock(),
request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(),
)
response = await ep.anthropic_response(
fastapi_response=MagicMock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
assert response.status_code == 500
assert json.loads(response.body)["error"]["message"] == "provider timeout"
hook_request_data = mock_logging.post_call_failure_hook.await_args.kwargs["request_data"]
assert hook_request_data is captured["processor_data"]

View file

@ -1238,6 +1238,18 @@ def test_health_liveness_endpoint(proxy_client):
print(f"\n/health/liveness response time: {duration_ms:.2f}ms")
def test_health_backlog_includes_admission_control_stats(proxy_client):
response = proxy_client.get("/health/backlog")
assert response.status_code == 200, response.text
assert set(response.json()) == {
"in_flight_requests",
"admitted_requests",
"queued_requests",
"rejected_requests",
}
def test_health_readiness(proxy_client):
"""
Test /health/readiness endpoint.

View file

@ -0,0 +1,402 @@
import asyncio
import json
from typing import Final
import pytest
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from litellm.proxy.middleware.admission_control_middleware import (
AdmissionControlMetrics,
AdmissionControlMiddleware,
AdmissionControlSettings,
AdmissionControlState,
AdmissionControlStats,
_parse_admission_control_settings,
create_prometheus_admission_metrics,
get_admission_control_settings,
)
@pytest.fixture
def state() -> AdmissionControlState:
return AdmissionControlState(lambda: None)
async def _call(
middleware: AdmissionControlMiddleware,
path: str = "/",
root_path: str = "",
) -> tuple[Message, ...]:
messages: Final[list[Message]] = []
async def receive() -> Message:
return {"type": "http.request", "body": b"", "more_body": False}
async def send(message: Message) -> None:
messages.append(message)
scope: Final[Scope] = {
"type": "http",
"path": path,
"root_path": root_path,
"method": "GET",
"headers": [],
}
await middleware(scope, receive, send)
return tuple(messages)
def _handler_with_release(
started: asyncio.Event,
release: asyncio.Event,
) -> ASGIApp:
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
started.set()
await release.wait()
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
return handler
def test_is_not_base_http_middleware() -> None:
assert not issubclass(AdmissionControlMiddleware, BaseHTTPMiddleware)
@pytest.mark.asyncio
async def test_capacity_rejects_excess_and_releases_queued_request(state: AdmissionControlState) -> None:
started: Final = asyncio.Event()
release: Final = asyncio.Event()
middleware: Final = AdmissionControlMiddleware(
_handler_with_release(started, release),
lambda: AdmissionControlSettings(1, 1, 1.0),
state,
)
first: Final = asyncio.create_task(_call(middleware))
await started.wait()
second: Final = asyncio.create_task(_call(middleware))
await asyncio.sleep(0)
assert state.get_stats().queued == 1
third: Final = await _call(middleware)
assert third[0]["status"] == 503
headers: Final = dict(third[0]["headers"])
assert headers[b"retry-after"] == b"1"
assert headers[b"content-type"] == b"application/json"
assert json.loads(third[1]["body"])["error"] == {
"message": "Worker at capacity: 1 in-flight, 1 queued requests. Retry later.",
"type": "overloaded_error",
"code": "503",
}
assert state.get_stats().rejected_total == 1
release.set()
assert (await first)[0]["status"] == 200
assert (await second)[0]["status"] == 200
assert state.get_stats() == AdmissionControlStats(0, 0, 1)
@pytest.mark.asyncio
async def test_pending_waiter_is_not_skipped_after_admission_is_released(state: AdmissionControlState) -> None:
started: Final = asyncio.Event()
release: Final = asyncio.Event()
third_trigger: Final = asyncio.Event()
middleware: Final = AdmissionControlMiddleware(
_handler_with_release(started, release),
lambda: AdmissionControlSettings(1, 2, 1.0),
state,
)
first: Final = asyncio.create_task(_call(middleware))
await started.wait()
second: Final = asyncio.create_task(_call(middleware))
await asyncio.sleep(0)
async def call_third() -> tuple[Message, ...]:
await third_trigger.wait()
return await _call(middleware)
third: Final = asyncio.create_task(call_third())
await asyncio.sleep(0)
release.set()
third_trigger.set()
await asyncio.sleep(0)
assert state.get_stats().queued == 2
await asyncio.gather(first, second, third)
@pytest.mark.asyncio
async def test_queue_timeout_rejects_and_decrements_queue(state: AdmissionControlState) -> None:
started: Final = asyncio.Event()
release: Final = asyncio.Event()
middleware: Final = AdmissionControlMiddleware(
_handler_with_release(started, release),
lambda: AdmissionControlSettings(1, 1, 0.05),
state,
)
first: Final = asyncio.create_task(_call(middleware))
await started.wait()
start_time: Final = asyncio.get_running_loop().time()
second: Final = await _call(middleware)
elapsed: Final = asyncio.get_running_loop().time() - start_time
assert second[0]["status"] == 503
assert elapsed < 0.5
assert state.get_stats().queued == 0
assert state.get_stats().rejected_total == 1
release.set()
await first
@pytest.mark.asyncio
@pytest.mark.parametrize(
("root_path", "probe_path"),
(
("", "/health/liveliness"),
("/proxy", "/proxy/health/liveliness"),
("/proxy", "/proxy/metrics"),
),
)
async def test_exempt_path_passes_through_when_saturated(
state: AdmissionControlState,
root_path: str,
probe_path: str,
) -> None:
started: Final = asyncio.Event()
release: Final = asyncio.Event()
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
if scope["path"] == "/":
started.set()
await release.wait()
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
first: Final = asyncio.create_task(_call(middleware))
await started.wait()
health: Final = await _call(middleware, probe_path, root_path)
assert health[0]["status"] == 200
blocked: Final = await _call(middleware, "/proxy/v1/chat/completions", root_path)
assert blocked[0]["status"] == 503
lookalike: Final = await _call(middleware, "/proxyhealth/liveliness", "/proxy")
assert lookalike[0]["status"] == 503
release.set()
await first
@pytest.mark.asyncio
async def test_non_http_scope_passes_through_when_saturated(state: AdmissionControlState) -> None:
seen: Final[list[str]] = []
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
seen.append(scope["type"])
middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
state.record_admission()
async def receive() -> Message:
return {"type": "lifespan.startup"}
async def send(message: Message) -> None:
return None
await middleware({"type": "lifespan"}, receive, send)
assert seen == ["lifespan"]
@pytest.mark.asyncio
async def test_none_settings_does_not_limit_concurrency() -> None:
active: Final = [0]
peak: Final = [0]
all_started: Final = asyncio.Event()
release: Final = asyncio.Event()
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
active[0] += 1
peak[0] = max(peak[0], active[0])
if active[0] == 3:
all_started.set()
await release.wait()
active[0] -= 1
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok", "more_body": False})
middleware: Final = AdmissionControlMiddleware(handler, lambda: None, AdmissionControlState(lambda: None))
requests: Final = tuple(asyncio.create_task(_call(middleware)) for _ in range(3))
await all_started.wait()
assert peak[0] == 3
release.set()
results: Final = await asyncio.gather(*requests)
assert tuple(result[0]["status"] for result in results) == (200, 200, 200)
@pytest.mark.asyncio
async def test_cancelling_queued_request_does_not_leak_counter(state: AdmissionControlState) -> None:
started: Final = asyncio.Event()
release: Final = asyncio.Event()
middleware: Final = AdmissionControlMiddleware(
_handler_with_release(started, release),
lambda: AdmissionControlSettings(1, 1, 1.0),
state,
)
first: Final = asyncio.create_task(_call(middleware))
await started.wait()
queued: Final = asyncio.create_task(_call(middleware))
await asyncio.sleep(0)
queued.cancel()
with pytest.raises(asyncio.CancelledError):
await queued
assert state.get_stats().queued == 0
release.set()
await first
@pytest.mark.asyncio
async def test_streaming_response_holds_admission_until_final_body(state: AdmissionControlState) -> None:
first_chunk_sent: Final = asyncio.Event()
finish_stream: Final = asyncio.Event()
async def handler(scope: Scope, receive: Receive, send: Send) -> None:
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"first", "more_body": True})
first_chunk_sent.set()
await finish_stream.wait()
await send({"type": "http.response.body", "body": b"last", "more_body": False})
middleware: Final = AdmissionControlMiddleware(
handler,
lambda: AdmissionControlSettings(1, 1, 1.0),
state,
)
first: Final = asyncio.create_task(_call(middleware))
await first_chunk_sent.wait()
second: Final = asyncio.create_task(_call(middleware))
await asyncio.sleep(0)
assert not second.done()
assert state.get_stats().queued == 1
finish_stream.set()
assert (await first)[0]["status"] == 200
assert (await second)[0]["status"] == 200
assert state.get_stats().admitted == 0
assert state.get_stats().queued == 0
class _FakeGauge:
def __init__(self) -> None:
self.value = 0.0
def inc(self, amount: float = 1) -> None:
self.value += amount
def dec(self, amount: float = 1) -> None:
self.value -= amount
class _FakeCounter:
def __init__(self) -> None:
self.by_reason: Final[dict[str, _FakeGauge]] = {}
def labels(self, reason: str) -> _FakeGauge:
return self.by_reason.setdefault(reason, _FakeGauge())
@pytest.mark.asyncio
async def test_metrics_track_admitted_queued_and_rejected() -> None:
admitted: Final = _FakeGauge()
queued: Final = _FakeGauge()
rejected: Final = _FakeCounter()
state: Final = AdmissionControlState(
lambda: AdmissionControlMetrics(admitted_gauge=admitted, queued_gauge=queued, rejected_counter=rejected)
)
started: Final = asyncio.Event()
release: Final = asyncio.Event()
middleware: Final = AdmissionControlMiddleware(
_handler_with_release(started, release),
lambda: AdmissionControlSettings(1, 1, 0.05),
state,
)
first: Final = asyncio.create_task(_call(middleware))
await started.wait()
second: Final = asyncio.create_task(_call(middleware))
await asyncio.sleep(0)
assert (admitted.value, queued.value) == (1.0, 1.0)
await _call(middleware)
assert rejected.by_reason["queue_full"].value == 1.0
await second
assert rejected.by_reason["queue_timeout"].value == 1.0
release.set()
await first
assert (admitted.value, queued.value) == (0.0, 0.0)
def test_create_prometheus_admission_metrics_registers_named_metrics() -> None:
from prometheus_client import REGISTRY
metrics: Final = create_prometheus_admission_metrics()
if metrics is not None:
metrics.admitted_gauge.inc()
metrics.queued_gauge.inc()
metrics.rejected_counter.labels(reason="queue_full").inc()
assert REGISTRY.get_sample_value("litellm_admission_admitted_requests") == 1.0
assert REGISTRY.get_sample_value("litellm_admission_queued_requests") == 1.0
assert REGISTRY.get_sample_value("litellm_admission_rejected_requests_total", {"reason": "queue_full"}) is not None
assert create_prometheus_admission_metrics() is None
@pytest.mark.parametrize(
("settings", "expected"),
(
({}, None),
({"max_in_flight_requests_per_worker": None}, None),
({"max_in_flight_requests_per_worker": 0}, None),
({"max_in_flight_requests_per_worker": "many"}, None),
({"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": -1}, None),
({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": 0}, None),
({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": -0.5}, None),
(
{"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": 0},
AdmissionControlSettings(3, 0, 1.0),
),
(
{"max_in_flight_requests_per_worker": 3},
AdmissionControlSettings(3, 3, 1.0),
),
(
{
"max_in_flight_requests_per_worker": 3,
"max_queued_requests_per_worker": 5,
"admission_queue_timeout_seconds": 0.25,
},
AdmissionControlSettings(3, 5, 0.25),
),
),
)
def test_get_admission_control_settings(
settings: dict[str, object],
expected: AdmissionControlSettings | None,
) -> None:
assert get_admission_control_settings(settings) == expected
def test_invalid_admission_control_settings_logs_once(caplog: pytest.LogCaptureFixture) -> None:
_parse_admission_control_settings.cache_clear()
caplog.set_level("ERROR")
settings: Final = {"max_in_flight_requests_per_worker": [1]}
assert get_admission_control_settings(settings) is None
assert get_admission_control_settings(settings) is None
messages: Final = tuple(
record.message
for record in caplog.records
if record.message.startswith("Ignoring invalid admission control settings")
)
assert len(messages) == 1

View file

@ -319,9 +319,6 @@ class TestVertexAIPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
f"https://{test_location}-aiplatform.googleapis.com/"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value=f"https://{test_location}-aiplatform.googleapis.com/"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -427,9 +424,6 @@ class TestVertexAIPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
"https://aiplatform.googleapis.com/"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value="https://aiplatform.googleapis.com/"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -530,9 +524,6 @@ class TestVertexAIPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
f"https://{default_location}-aiplatform.googleapis.com/"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value=f"https://{default_location}-aiplatform.googleapis.com/"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -1308,9 +1299,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
"https://discoveryengine.googleapis.com"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value="https://discoveryengine.googleapis.com"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -3650,7 +3638,6 @@ class TestVertexRawPredictStreamingClassification:
base_url = "https://us-east5-aiplatform.googleapis.com/"
mock_handler = Mock()
mock_handler.get_default_base_target_url.return_value = base_url
mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url)
module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
with (
@ -4234,6 +4221,126 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items())
class TestVertexPassthroughDefaultLocationOnShortRoutes:
PROJECT = "test-project"
SHORT_ROUTE = "publishers/google/models/gemini-2.5-flash:generateContent"
@staticmethod
def _forwarder() -> Mock:
return Mock(return_value=AsyncMock(return_value={"status": "success"}))
async def _forward(
self,
monkeypatch,
endpoint: str,
default_config: dict | None,
headers: list[tuple[bytes, bytes]],
forwarder: Mock,
) -> None:
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
PassthroughEndpointRouter,
)
async def receive():
return {"type": "http.request", "body": b"{}", "more_body": False}
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": f"/vertex_ai/{endpoint}",
"headers": headers,
"query_string": b"",
},
receive=receive,
)
router: Final = PassthroughEndpointRouter()
if default_config is not None:
router.set_default_vertex_config(dict(default_config))
module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
monkeypatch.setattr(f"{module}.passthrough_endpoint_router", router)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
mock_credentials: Final = Mock()
mock_credentials.token = "test-token"
caller: Final = UserAPIKeyAuth(api_key="test-key")
with (
mock.patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
return_value=(mock_credentials, self.PROJECT),
),
mock.patch(f"{module}.create_pass_through_route", new=forwarder),
mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)),
):
await vertex_proxy_route(
endpoint=endpoint,
request=request,
fastapi_response=Response(),
user_api_key_dict=caller,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("endpoint", "location", "expected_target"),
[
(
SHORT_ROUTE,
"global",
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
),
(
f"v1/{SHORT_ROUTE}",
"global",
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
),
(
f"v1beta1/{SHORT_ROUTE}",
"global",
"https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/" + SHORT_ROUTE,
),
(
SHORT_ROUTE,
"us-central1",
"https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/"
+ SHORT_ROUTE,
),
],
)
async def test_default_vertex_config_location_fills_routes_without_project_and_location(
self, monkeypatch, endpoint, location, expected_target
):
forwarder: Final = self._forwarder()
await self._forward(
monkeypatch,
endpoint,
{"vertex_project": self.PROJECT, "vertex_location": location, "vertex_credentials": "test-creds"},
[(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
forwarder,
)
forwarded: Final = forwarder.call_args.kwargs
assert str(forwarded["target"]) == expected_target
assert forwarded["custom_headers"]["Authorization"] == "Bearer test-token"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("default_config", "headers"),
[
(None, [(b"content-type", b"application/json"), (b"authorization", b"Bearer ya29.byo-google-oauth")]),
(
{"vertex_project": PROJECT, "vertex_credentials": "test-creds"},
[(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
),
],
)
async def test_no_location_anywhere_is_a_400_not_a_500(self, monkeypatch, default_config, headers):
forwarder: Final = self._forwarder()
with pytest.raises(HTTPException) as raised:
await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers, forwarder)
forwarder.assert_not_called()
assert raised.value.status_code == 400
assert "/projects/<project>/locations/<location>/" in str(raised.value.detail)
assert "default_vertex_config" in str(raised.value.detail)
class TestGetAzureAISearchIndexFromEndpoint:
"""The operable index is only the segment right after ``indexes``.

View file

@ -4,6 +4,7 @@ import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
VertexAIPassThroughHandler,
_base_vertex_proxy_route,
_upstream_headers_for_vertex_route,
)
@ -20,6 +21,7 @@ async def test_vertex_passthrough_load_balancing():
mock_request = MagicMock()
mock_response = MagicMock()
mock_handler = MagicMock()
mock_handler.get_default_base_target_url.return_value = "https://test.url"
# Mock the router
mock_router = MagicMock()
@ -68,7 +70,6 @@ async def test_vertex_passthrough_load_balancing():
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
mock_prep_headers.return_value = (
{},
"https://test.url",
False,
"test-project-lb",
"us-central1-lb",
@ -290,12 +291,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
mock_vertex_credentials.vertex_location = "us-central1"
mock_vertex_credentials.vertex_credentials = "test-credentials"
# Create mock handler
mock_handler = MagicMock()
mock_handler.update_base_target_url_with_credential_location.return_value = (
"https://us-central1-aiplatform.googleapis.com"
)
with (
patch.object(
VertexBase,
@ -313,7 +308,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
# Call the function
(
headers,
base_target_url,
headers_passed_through,
vertex_project,
vertex_location,
@ -323,8 +317,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
router_credentials=None,
vertex_project="test-project",
vertex_location="us-central1",
base_target_url="https://us-central1-aiplatform.googleapis.com",
get_vertex_pass_through_handler=mock_handler,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
)
@ -394,7 +386,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
"content-type": "application/json",
"Authorization": "Bearer vertex-access-token",
},
"https://aiplatform.googleapis.com",
False,
"test-project",
"global",
@ -406,7 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
request=MagicMock(),
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=MagicMock(),
get_vertex_pass_through_handler=VertexAIPassThroughHandler(),
)
upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
@ -473,12 +464,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
mock_vertex_credentials.vertex_location = "us-central1"
mock_vertex_credentials.vertex_credentials = "test-credentials"
# Create mock handler
mock_handler = MagicMock()
mock_handler.update_base_target_url_with_credential_location.return_value = (
"https://us-central1-aiplatform.googleapis.com"
)
with (
patch.object(
VertexBase,
@ -495,7 +480,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
(
headers,
_base_target_url,
_headers_passed_through,
_vertex_project,
_vertex_location,
@ -505,8 +489,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
router_credentials=None,
vertex_project="test-project",
vertex_location="us-central1",
base_target_url="https://us-central1-aiplatform.googleapis.com",
get_vertex_pass_through_handler=mock_handler,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
)
@ -742,7 +724,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
mock_prep_headers.return_value = (
{},
"https://global-aiplatform.googleapis.com",
False,
"nv-gcpllmgwit-20250411173346",
"global",

View file

@ -5,14 +5,12 @@ import hashlib
import json
import re
from datetime import timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
from unittest.mock import AsyncMock, MagicMock, patch
import litellm
import litellm.proxy.proxy_server as ps
@ -3325,7 +3323,7 @@ def _compare_nested_dicts(
return differences
# Check for keys in actual but not in expected
for key in actual.keys():
for key in actual:
current_path = f"{path}.{key}" if path else key
if current_path not in ignore_keys and key not in expected:
differences.append(f"Extra key in actual: {current_path}")
@ -3495,24 +3493,22 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch):
# Return individual log entries when summarize=false
return mock_spend_logs
async def group_by(self, *args, **kwargs):
# Return grouped data when summarize=true
# Simplified mock response for grouped data
async def query_raw(self, sql_query, *params):
yesterday = datetime.datetime.now(timezone.utc) - timedelta(days=1)
return [
{
"api_key": "sk-test-key",
"user": "test_user_1",
"model": "gpt-3.5-turbo",
"startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
"_sum": {"spend": 0.05},
"day": yesterday.date().isoformat(),
"spend": 0.05,
},
{
"api_key": "sk-test-key",
"user": "test_user_1",
"model": "gpt-4",
"startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
"_sum": {"spend": 0.10},
"day": yesterday.date().isoformat(),
"spend": 0.10,
},
]
@ -3850,47 +3846,30 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
"""
from datetime import datetime, timedelta, timezone
# This simulates the summarized data that Prisma's `group_by` would return.
mock_summarized_response = [
{
"api_key": "sk-test-key",
"user": "test_user_1",
"model": "gpt-4",
"startTime": (datetime.now(timezone.utc) - timedelta(days=1)).strftime(
"%Y-%m-%dT%H:%M:%S.%fZ"
),
"_sum": {"spend": 0.15},
"day": (datetime.now(timezone.utc) - timedelta(days=1)).date().isoformat(),
"spend": 0.15,
}
]
# This mock class will replace the real Prisma client.
class MockDB:
def __init__(self):
self.litellm_spendlogs = self
async def group_by(self, *args, **kwargs):
# We assert that the `gte` and `lte` values are strings in ISO format.
# If they were datetime objects, this test would fail.
where_clause = kwargs.get("where", {})
start_time_filter = where_clause.get("startTime", {})
assert "gte" in start_time_filter
assert "lte" in start_time_filter
assert isinstance(start_time_filter["gte"], str)
assert isinstance(start_time_filter["lte"], str)
assert "T" in start_time_filter["gte"] # Check for ISO format 'T' separator
# If the assertions pass, return the mock response.
async def query_raw(self, sql_query, *params):
assert isinstance(params[0], str)
assert isinstance(params[1], str)
assert "T" in params[0]
assert "T" in params[1]
return mock_summarized_response
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
# Apply the monkeypatch to replace the real prisma_client with our mock.
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
# Define a date range for the test.
start_date = (datetime.now(timezone.utc) - timedelta(days=2)).strftime("%Y-%m-%d")
end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
@ -3898,8 +3877,6 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
# Call the endpoint with both start and end dates.
# We don't need `summarize=true` as it's the default.
response = client.get(
"/spend/logs",
params={
@ -3909,11 +3886,9 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
headers={"Authorization": "Bearer sk-test"},
)
# ASSERTIONS
assert response.status_code == 200
data = response.json()
# Check that the response is not empty and has the summarized structure.
assert isinstance(data, list)
assert len(data) > 0
assert "startTime" in data[0]
@ -3924,6 +3899,183 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatch):
mock_rows = [
{
"day": "2024-01-01",
"api_key": "hashed::sk-abc",
"user": "u1",
"model": "gpt-4",
"spend": 0.1,
},
{
"day": "2024-01-01",
"api_key": "hashed::sk-abc",
"user": "u1",
"model": "gpt-4o",
"spend": 0.2,
},
]
class MockDB:
def __init__(self):
self.captured_sql = None
self.captured_params = None
async def query_raw(self, sql_query, *params):
self.captured_sql = sql_query
self.captured_params = params
return mock_rows
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
def hash_token(self, token):
return "hashed::" + token
mock_prisma_client = MockPrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
response = client.get(
"/spend/logs",
params={
"start_date": "2024-01-01",
"end_date": "2024-01-03",
"api_key": "sk-abc",
"request_id": "req-123",
"user_id": "u1",
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
sql = mock_prisma_client.db.captured_sql
assert "date_trunc('day'" in sql
assert "GROUP BY" in sql
assert "find_many" not in sql
assert not hasattr(mock_prisma_client.db, "group_by")
assert mock_prisma_client.db.captured_params == (
"2024-01-01T00:00:00+00:00",
"2024-01-03T00:00:00+00:00",
"hashed::sk-abc",
"req-123",
"u1",
)
assert len(data) == 3
assert data[0]["startTime"] == "2024-01-01"
assert data[0]["spend"] == pytest.approx(0.3)
assert data[0]["models"] == {"gpt-4": 0.1, "gpt-4o": 0.2}
assert data[0]["users"] == {"u1": pytest.approx(0.3)}
assert data[0]["hashed::sk-abc"] == pytest.approx(0.3)
assert data[1] == {
"startTime": "2024-01-02",
"spend": 0,
"users": {},
"models": {},
}
assert data[2] == {
"startTime": "2024-01-03",
"spend": 0,
"users": {},
"models": {},
}
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch):
class MockDB:
async def query_raw(self, sql_query, *params):
return []
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
response = client.get(
"/spend/logs",
params={"start_date": "2024-01-01", "end_date": "2024-01-01"},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
assert response.json() == []
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client, monkeypatch):
mock_rows = [
{
"day": "2024-01-01",
"api_key": "plain-key",
"user": "u1",
"model": "gpt-4",
"spend": 0.4,
}
]
class MockDB:
def __init__(self):
self.captured_params = None
async def query_raw(self, sql_query, *params):
self.captured_params = params
return mock_rows
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
mock_prisma_client = MockPrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
response = client.get(
"/spend/logs",
params={
"start_date": "2024-01-01",
"end_date": "2024-01-01",
"api_key": "plain-key",
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert mock_prisma_client.db.captured_params == (
"2024-01-01T00:00:00+00:00",
"2024-01-01T00:00:00+00:00",
"plain-key",
)
assert data == [
{
"startTime": "2024-01-01",
"spend": pytest.approx(0.4),
"plain-key": pytest.approx(0.4),
"users": {"u1": pytest.approx(0.4)},
"models": {"gpt-4": pytest.approx(0.4)},
}
]
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_with_error_code(client):
"""Test filtering spend logs by error code"""
@ -4832,13 +4984,14 @@ class _CaptureFilterDB:
def __init__(self):
self.litellm_spendlogs = self
self.captured_where = None
self.captured_params = None
async def find_many(self, *args, **kwargs):
self.captured_where = kwargs.get("where")
return []
async def group_by(self, *args, **kwargs):
self.captured_where = kwargs.get("where")
async def query_raw(self, sql_query, *params):
self.captured_params = params
return []

View file

@ -5101,6 +5101,40 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment():
)
def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete():
"""
Regression for #29064: upsert_deployment removed the old deployment from
model_list but left it in the global pattern_router, so wildcard requests
round-robined between the stale and the corrected deployment.
"""
from litellm.types.router import Deployment, LiteLLM_Params
router = litellm.Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/openai/*", "api_key": "sk-old"},
"model_info": {"id": "global-wildcard"},
}
]
)
router.upsert_deployment(
Deployment(
model_name="openai/*",
litellm_params=LiteLLM_Params(model="openai/*", api_key="sk-new"),
model_info={"id": "global-wildcard"},
)
)
matches = router.pattern_router.route("openai/gpt-5.2")
assert matches is not None
assert [m["litellm_params"]["api_key"] for m in matches] == ["sk-new"]
router.delete_deployment(id="global-wildcard")
assert router.pattern_router.patterns == {}
def test_pattern_match_router_remove_deployment():
"""
remove_deployment must drop only the deployment with the given model id and

View file

@ -3,7 +3,7 @@
"limit": 22328
},
"LIT002": {
"limit": 26758
"limit": 26750
},
"LIT003": {
"limit": 261
@ -27,10 +27,10 @@
"limit": 0
},
"LIT010": {
"limit": 16470
"limit": 16468
},
"LIT011": {
"limit": 5516
"limit": 5514
},
"LIT012": {
"limit": 4489

View file

@ -41,6 +41,7 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
data={data}
columns={columns}
getRowId={(tag, index) => tag.name || String(index)}
fillHeight
sortingMode="client"
sorting={sorting}
onSortingChange={setSorting}

View file

@ -126,7 +126,7 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
}, [accessToken]);
return (
<div className="mx-4 h-[75vh]">
<div className="mx-4 h-full">
{selectedTagId ? (
<TagInfoView
tagId={selectedTagId}
@ -139,7 +139,7 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
editTag={editTag}
/>
) : (
<div className="mt-2 h-[75vh] w-full gap-2 p-8">
<div className="flex h-full w-full flex-col p-8 pt-10">
<div className="mt-2 mb-4 flex w-full items-center justify-between">
<h1>Tag Management</h1>
<div className="flex items-center space-x-2">
@ -162,23 +162,21 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
</p>
</div>
<Button className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
<Button className="mb-4 self-start" onClick={() => setIsCreateModalVisible(true)}>
+ Create New Tag
</Button>
<div className="mt-2 grid h-[75vh] w-full grid-cols-1 gap-2 pt-2 pb-2">
<div>
<TagTable
data={tags}
isLoading={isLoadingTags}
onEdit={(tag) => {
setSelectedTagId(tag.name);
setEditTag(true);
}}
onDelete={handleDelete}
onSelectTag={setSelectedTagId}
/>
</div>
<div className="mt-2 flex min-h-0 flex-1 flex-col">
<TagTable
data={tags}
isLoading={isLoadingTags}
onEdit={(tag) => {
setSelectedTagId(tag.name);
setEditTag(true);
}}
onDelete={handleDelete}
onSelectTag={setSelectedTagId}
/>
</div>
{/* Create Tag Modal */}

View file

@ -137,8 +137,8 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
/>
</div>
) : (
<div className="mx-4 h-[75vh]">
<div className="gap-2 p-8 h-[75vh] w-full mt-2">
<div className="mx-4">
<div className="gap-2 p-8 w-full mt-2">
<div className="flex justify-between mt-2 w-full items-center mb-4">
<h1 className="text-xl font-semibold tracking-tight text-foreground">Vector Store Management</h1>
<div className="flex items-center space-x-2">

View file

@ -400,7 +400,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
}
return (
<div className="mx-4 h-[75vh]">
<div className="mx-4">
{publicPage == false ? (
<div className="w-full m-2 mt-2 p-8">
{/* Header with Title, Description and URL */}

View file

@ -559,6 +559,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
{
key: "your-teams",
label: "Your Teams",
className: "flex min-h-0 flex-1 flex-col",
children: (
<>
<TeamsTable
@ -608,6 +609,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
{
key: "available-teams",
label: "Available Teams",
className: "min-h-0 flex-1 overflow-y-auto",
children: <AvailableTeamsPanel accessToken={accessToken} userID={userID} />,
},
...(isProxyAdminRole(userRole || "")
@ -615,6 +617,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
{
key: "default-settings",
label: "Default Team Settings",
className: "min-h-0 flex-1 overflow-y-auto",
children: <TeamSSOSettings accessToken={accessToken} userID={userID || ""} userRole={userRole || ""} />,
},
]
@ -622,7 +625,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
];
return (
<main className={selectedTeamId ? "px-12 py-6" : "p-8"}>
<main className={selectedTeamId ? "px-12 py-6" : "flex h-full flex-col p-8"}>
{selectedTeamId ? (
<TeamInfoView
teamId={selectedTeamId}
@ -642,7 +645,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
premiumUser={premiumUser}
/>
) : (
<Tabs defaultValue={tabItems[0].key} className="gap-6">
<Tabs defaultValue={tabItems[0].key} className="min-h-0 flex-1 gap-6">
<PageHeader
icon={<Users />}
title="Teams"
@ -674,7 +677,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
)}
/>
{tabItems.map((item) => (
<TabsContent key={item.key} value={item.key}>
<TabsContent key={item.key} value={item.key} className={item.className}>
{item.children}
</TabsContent>
))}

View file

@ -164,7 +164,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
isLoading={isLoading}
loadingMessage="Loading teams..."
noDataMessage="No teams found"
maxBodyHeight="calc(75vh - 210px)"
fillHeight
size="compact"
toolbar={(table) => (
<>

View file

@ -256,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
}
return (
<div className="flex h-full flex-col gap-6 overflow-hidden">
<div className="flex min-h-0 flex-1 flex-col gap-6">
<PageHeader
icon={<KeyRound />}
title="Virtual Keys"
@ -283,7 +283,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
isLoading={isLoading}
loadingMessage="Loading keys..."
noDataMessage="No keys found"
maxBodyHeight="calc(75vh - 210px)"
fillHeight
size="compact"
toolbar={(table) => (
<>

View file

@ -1,26 +1,58 @@
import React from "react";
import { Input } from "@/components/ui/input";
import { Switch } from "@/components/ui/switch";
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
import { DEFAULT_DEPLOYMENT_AFFINITY } from "./ComplexityRouterConfig";
import { DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_SESSION_AFFINITY_TTL_SECONDS } from "./ComplexityRouterConfig";
export const AffinityControls: React.FC<{
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
}> = ({ value, onChange }) => (
<>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
aria-label="Pin a session to one deployment per model group"
/>
<strong className="font-semibold">Pin a session to one deployment per model group</strong>
</div>
<span className="block text-xs text-muted-foreground">
Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
load-balance every turn.
</span>
</>
);
}> = ({ value, onChange }) => {
const [ttlDraft, setTtlDraft] = React.useState<string | null>(null);
const commitTtl = (raw: string) => {
setTtlDraft(null);
if (raw.trim() === "") {
onChange({ ...value, session_affinity_ttl_seconds: undefined });
return;
}
const parsed = Number(raw);
if (!Number.isFinite(parsed)) return;
onChange({ ...value, session_affinity_ttl_seconds: Math.max(1, Math.round(parsed)) });
};
return (
<>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
aria-label="Pin a session to one deployment per model group"
/>
<strong className="font-semibold">Pin a session to one deployment per model group</strong>
</div>
<span className="block text-xs mb-3 text-muted-foreground">
Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
load-balance every turn.
</span>
<div style={{ maxWidth: 320 }}>
<label className="block text-sm font-medium mb-1" htmlFor="session-affinity-ttl">
How long a pin survives idle (seconds)
</label>
<Input
id="session-affinity-ttl"
inputMode="numeric"
value={ttlDraft ?? value.session_affinity_ttl_seconds ?? ""}
placeholder={String(DEFAULT_SESSION_AFFINITY_TTL_SECONDS)}
onChange={(event) => setTtlDraft(event.target.value)}
onBlur={(event) => commitTtl(event.target.value)}
/>
<span className="block text-xs mt-1 text-muted-foreground">
Refreshes after every request that reuses a pin. Empty tracks the backend default of{" "}
{DEFAULT_SESSION_AFFINITY_TTL_SECONDS} seconds.
</span>
</div>
</>
);
};

View file

@ -947,6 +947,46 @@ describe("ComplexityRouterConfig affinity panel", () => {
expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked();
});
it("writes an idle TTL on blur and keeps the partial input as a draft while typing", () => {
const onChange = vi.fn();
renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Affinity"));
const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
expect(ttl).toHaveAttribute("placeholder", "3600");
fireEvent.change(ttl, { target: { value: "300" } });
expect(onChange).not.toHaveBeenCalled();
fireEvent.blur(ttl);
expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 300 });
});
it("clearing the idle TTL returns the router to its backend default", () => {
const onChange = vi.fn();
const value = { ...defaultValue, session_affinity_ttl_seconds: 300 };
renderWithProviders(<ComplexityRouterConfig {...baseProps} value={value} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Affinity"));
const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
expect(ttl).toHaveValue("300");
fireEvent.change(ttl, { target: { value: "" } });
fireEvent.blur(ttl);
expect(onChange).toHaveBeenCalledWith({ ...value, session_affinity_ttl_seconds: undefined });
});
it("clamps a non-positive idle TTL to the backend's minimum", () => {
const onChange = vi.fn();
renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
fireEvent.click(screen.getByText("Advanced: Affinity"));
const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
fireEvent.change(ttl, { target: { value: "0" } });
fireEvent.blur(ttl);
expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 1 });
});
});
describe("ComplexityRouterConfig default model", () => {

View file

@ -58,6 +58,7 @@ export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3;
export const DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS = 8000;
export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120;
export const DEFAULT_SESSION_AFFINITY = false;
export const DEFAULT_SESSION_AFFINITY_TTL_SECONDS = 3600;
export const DEFAULT_DEPLOYMENT_AFFINITY = true;
export type ClassificationMode = "every_request" | "user_turn";
@ -411,6 +412,7 @@ export interface ComplexityRouterConfigValue {
hybrid_boundary_margin?: number;
classification_mode?: ClassificationMode;
session_affinity?: boolean;
session_affinity_ttl_seconds?: number;
modality_routing?: boolean;
modality_pin_override?: boolean;
deployment_affinity?: boolean;

View file

@ -493,7 +493,7 @@ describe("AddAutoRouterTab", () => {
});
});
it("carries session affinity turned on through to the create payload", async () => {
it("carries session affinity turned on and its idle window through to the create payload", async () => {
const user = userEvent.setup();
vi.mocked(getMissingTiersError).mockReturnValue(null);
@ -503,12 +503,17 @@ describe("AddAutoRouterTab", () => {
expandDetailedConfiguration();
await user.click(screen.getByText("Advanced: Classification Method"));
await user.click(await screen.findByRole("radio", { name: /Once per session/ }));
await user.click(screen.getByText("Advanced: Affinity"));
const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
fireEvent.change(ttl, { target: { value: "300" } });
fireEvent.blur(ttl);
await user.click(screen.getByRole("button", { name: /add auto router/i }));
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
session_affinity: true,
session_affinity_ttl_seconds: 300,
});
});

View file

@ -388,6 +388,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score,
enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation,
contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer,
sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds,
};
const submitRecommendedRouter = async (name: string) => {

View file

@ -73,6 +73,15 @@ describe("buildComplexityRouterConfig", () => {
expect(config.context_window_escalation_buffer).toBe(0.9);
});
it("omits session_affinity_ttl_seconds when untouched, so the router tracks the backend default", () => {
expect(buildComplexityRouterConfig(baseParams)).not.toHaveProperty("session_affinity_ttl_seconds");
});
it("emits an explicit session affinity idle window", () => {
const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinityTtlSeconds: 300 });
expect(config.session_affinity_ttl_seconds).toBe(300);
});
it("trims escalation keywords and drops blank entries", () => {
const config = buildComplexityRouterConfig({
...baseParams,

View file

@ -138,6 +138,7 @@ export interface BuildComplexityRouterConfigParams {
tierModelParams?: TierModelParamsByTier;
enableContextWindowEscalation?: boolean;
contextWindowEscalationBuffer?: number;
sessionAffinityTtlSeconds?: number;
}
/**
@ -175,6 +176,7 @@ export interface ComplexityRouterConfigPayload {
hybrid_boundary_margin?: number;
classification_mode: ClassificationMode;
session_affinity: boolean;
session_affinity_ttl_seconds?: number;
deployment_affinity: boolean;
modality_routing: boolean;
modality_pin_override: boolean;
@ -456,6 +458,7 @@ export const buildComplexityRouterConfig = ({
tierModelParams,
enableContextWindowEscalation,
contextWindowEscalationBuffer,
sessionAffinityTtlSeconds,
}: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => {
const serializedTierModelConfigs = customTierSet
? serializeTierModelConfigs(
@ -522,6 +525,9 @@ export const buildComplexityRouterConfig = ({
...(contextWindowEscalationBuffer !== undefined && {
context_window_escalation_buffer: contextWindowEscalationBuffer,
}),
...(sessionAffinityTtlSeconds !== undefined && {
session_affinity_ttl_seconds: sessionAffinityTtlSeconds,
}),
...scorerKnobs,
};
if (!customTierSet) return payload;

View file

@ -257,6 +257,43 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => {
});
});
describe("buildUpdatedComplexityRouterConfig session affinity ttl", () => {
it("writes an edited idle window", () => {
const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, session_affinity_ttl_seconds: 300 });
expect(result.session_affinity_ttl_seconds).toBe(300);
});
it("carries a stored idle window through an untouched open-and-save", () => {
const stored = { ...STORED, session_affinity_ttl_seconds: 900 };
const result = buildUpdatedComplexityRouterConfig(stored, hydrateComplexityRouterConfig(stored, undefined));
expect(result.session_affinity_ttl_seconds).toBe(900);
});
it("drops the key when the field is cleared, so the router goes back to tracking the backend default", () => {
const result = buildUpdatedComplexityRouterConfig(
{ ...STORED, session_affinity_ttl_seconds: 900 },
{ ...FORM_VALUE, session_affinity_ttl_seconds: undefined },
);
expect(result).not.toHaveProperty("session_affinity_ttl_seconds");
});
it("keeps the idle window on a custom tier set, whose deployment pin still uses it", () => {
const result = buildUpdatedComplexityRouterConfig(STORED, {
...FORM_VALUE,
session_affinity_ttl_seconds: 300,
custom_tier_set: {
tiers: [
{ id: "a", name: "CASUAL", definition: "small talk", models: ["gpt-4o-mini"] },
{ id: "b", name: "AUDIT", definition: "security review", models: ["o1"] },
],
fallback_tier_id: "a",
},
});
expect(result.session_affinity).toBe(false);
expect(result.session_affinity_ttl_seconds).toBe(300);
});
});
describe("buildUpdatedComplexityRouterConfig modality pin override", () => {
it("writes modality_pin_override explicitly both ways", () => {
expect(
@ -529,6 +566,7 @@ describe("managed keys survive an untouched open-and-save", () => {
classifier_fallback: "default_model",
classification_mode: "user_turn",
session_affinity: true,
session_affinity_ttl_seconds: 300,
modality_routing: true,
modality_pin_override: true,
deployment_affinity: false,

View file

@ -550,6 +550,49 @@ describe("EditAutoRouterModal deployment affinity", () => {
expect(savedConfig().deployment_affinity).toBe(false);
});
it("preserves an idle TTL through an untouched save", async () => {
const user = userEvent.setup();
renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 });
await user.click(await screen.findByText("Advanced: Affinity"));
expect(await screen.findByLabelText("How long a pin survives idle (seconds)")).toHaveValue("300");
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().session_affinity_ttl_seconds).toBe(300);
});
it("persists an edited idle TTL", async () => {
const user = userEvent.setup();
renderWithStoredConfig(STORED_CONFIG);
await user.click(await screen.findByText("Advanced: Affinity"));
const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
fireEvent.change(ttl, { target: { value: "300" } });
fireEvent.blur(ttl);
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig().session_affinity_ttl_seconds).toBe(300);
});
it("removes the idle TTL when cleared", async () => {
const user = userEvent.setup();
renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 });
await user.click(await screen.findByText("Advanced: Affinity"));
const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
fireEvent.change(ttl, { target: { value: "" } });
fireEvent.blur(ttl);
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(savedConfig()).not.toHaveProperty("session_affinity_ttl_seconds");
});
// modality_pin_override is a managed key, so the modal rewrites it from form state on save. A
// hydration gap would silently turn a stored override off on the next untouched save.
it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => {

View file

@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig {
dimension_weights?: unknown;
reasoning_override_min_score?: unknown;
session_affinity?: unknown;
session_affinity_ttl_seconds?: unknown;
modality_routing?: unknown;
modality_pin_override?: unknown;
deployment_affinity?: unknown;
@ -182,6 +183,11 @@ export const hydrateComplexityRouterConfig = (
reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score),
session_affinity:
typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY,
session_affinity_ttl_seconds:
typeof parsedConfig.session_affinity_ttl_seconds === "number" &&
Number.isFinite(parsedConfig.session_affinity_ttl_seconds)
? parsedConfig.session_affinity_ttl_seconds
: undefined,
modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false,
modality_pin_override:
typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false,
@ -224,6 +230,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
"hybrid_boundary_margin",
"classification_mode",
"session_affinity",
"session_affinity_ttl_seconds",
"modality_routing",
"modality_pin_override",
"deployment_affinity",
@ -322,6 +329,7 @@ export const buildUpdatedComplexityRouterConfig = (
classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns,
classifierFallback: value.classifier_fallback,
sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY,
sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds,
modalityRouting: value.modality_routing ?? false,
modalityPinOverride: value.modality_pin_override ?? false,
deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,

View file

@ -613,8 +613,11 @@ describe("DataTable layout", () => {
it("makes the header sticky and constrains body height when maxBodyHeight is set", () => {
render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameEmailColumns} maxBodyHeight={240} />);
expect(screen.getByTestId("data-table-head")).toHaveClass("sticky");
expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" });
const scroller = screen.getByTestId("data-table-scroller");
expect(scroller).toHaveStyle({ maxHeight: "240px" });
expect(scroller).toHaveClass("overflow-auto");
expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible");
expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background");
});
it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => {

View file

@ -59,18 +59,22 @@ const noop = () => {};
/**
* Height-filling mode. The table still sizes to its rows; the parent's height is only a ceiling, so
* a short table keeps its footer under the last row and a long one scrolls its rows instead of the
* page. `table-container` is the Table primitive's own overflow-x wrapper; left as a scroll box it
* captures the sticky header and the header scrolls away with the rows. And rows pass under that
* header, which the semi-transparent header row tint alone would not hide.
* page.
*/
const FILL_CLASSES = {
outer: "flex max-h-full min-h-0 flex-col",
frame: "flex min-h-0 flex-col",
body: "min-h-0 [&_[data-slot=table-container]]:overflow-visible",
body: "min-h-0",
} as const;
const NO_FILL_CLASSES = { outer: "", frame: "", body: "" } as const;
const STICKY_CLASSES = {
body: "[&_[data-slot=table-container]]:overflow-visible",
header: "bg-background",
} as const;
const NO_FILL_CLASSES = { outer: "", frame: "", body: "", header: "" } as const;
const NO_STICKY_CLASSES = { body: "", header: "" } as const;
function columnDefId<TData, TValue>(column: ColumnDef<TData, TValue>): string | undefined {
if ("id" in column && typeof column.id === "string") {
@ -533,6 +537,7 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
const visibleColumnCount = table.getVisibleLeafColumns().length;
const stickyHeader = maxBodyHeight !== undefined || fillHeight;
const fill = fillHeight ? FILL_CLASSES : NO_FILL_CLASSES;
const sticky = stickyHeader ? STICKY_CLASSES : NO_STICKY_CLASSES;
const tableStyle = enableColumnResizing ? { width: table.getTotalSize(), minWidth: "100%" } : undefined;
const renderPagination = (): React.ReactNode => {
@ -593,13 +598,13 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
{toolbar !== undefined && <div className="shrink-0 border-b border-border px-4 py-3">{toolbar(table)}</div>}
<div
data-testid="data-table-scroller"
className={cn(stickyHeader ? "overflow-auto" : "overflow-x-auto", fill.body)}
className={cn(stickyHeader ? "overflow-auto" : "overflow-x-auto", sticky.body, fill.body)}
style={maxBodyHeight !== undefined ? { maxHeight: maxBodyHeight } : undefined}
>
<TableRoot className={enableColumnResizing ? "table-fixed" : ""} style={tableStyle}>
<TableHeader
data-testid="data-table-head"
className={cn(stickyHeader ? "sticky top-0 z-sticky" : "", fill.header)}
className={cn(stickyHeader ? "sticky top-0 z-sticky" : "", sticky.header)}
>
{table.getHeaderGroups().map((headerGroup) => (
<TableRow key={headerGroup.id} className="bg-muted/50">

View file

@ -443,7 +443,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
}, []);
return (
<div className="w-full h-full overflow-hidden">
<div className="w-full">
{selectedKey ? (
<KeyInfoView
keyId={selectedKey.token}
@ -453,7 +453,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
onDelete={refetch}
/>
) : (
<div className="py-4 flex-1 overflow-hidden">
<div className="py-4">
<DataTable
data={displayKeys}
columns={columns}
@ -471,7 +471,6 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
columnResizeMode="onChange"
isLoading={isLoading || isFetching}
loadingMessage="Loading keys..."
maxBodyHeight="75vh"
size="compact"
toolbar={(table) => (
<>

View file

@ -216,24 +216,22 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer";
return (
<main className="h-[75vh] p-8">
<div className="flex h-full flex-col">
<VirtualKeysTable
headerActions={
canCreateKey ? (
<CreateKey
key={selectedTeam ? selectedTeam.team_id : null}
team={selectedTeam as Team | null}
teams={teams as Team[]}
data={keys}
addKey={addKey}
autoOpenCreate={autoOpenCreate}
prefillData={prefillData}
/>
) : undefined
}
/>
</div>
<main className="flex h-full flex-col p-8">
<VirtualKeysTable
headerActions={
canCreateKey ? (
<CreateKey
key={selectedTeam ? selectedTeam.team_id : null}
team={selectedTeam as Team | null}
teams={teams as Team[]}
data={keys}
addKey={addKey}
autoOpenCreate={autoOpenCreate}
prefillData={prefillData}
/>
) : undefined
}
/>
</main>
);
};

View file

@ -86,6 +86,7 @@ export function RequestLogsTable({
data={data}
columns={columns}
getRowId={(row) => row.request_id}
fillHeight
sortingMode="server"
sorting={sorting}
onSortingChange={onSortingChange}

View file

@ -27,6 +27,9 @@ const AUDIT_LOGS_TAB: LogsTab = { id: "audit logs", label: "Audit Logs" };
const DELETED_KEYS_TAB: LogsTab = { id: "deleted keys", label: "Deleted Keys" };
const DELETED_TEAMS_TAB: LogsTab = { id: "deleted teams", label: "Deleted Teams" };
const tabContentClassName = (tabId: LogsTabId): string =>
tabId === REQUEST_LOGS_TAB.id ? "flex min-h-0 flex-1 flex-col" : "min-h-0 flex-1 overflow-y-auto";
export default function SpendLogsTable({ accessToken, token, userRole, userID, premiumUser }: SpendLogsTableProps) {
const [activeTab, setActiveTab] = useState<LogsTabId>(REQUEST_LOGS_TAB.id);
const canViewAuditLogs = useCan("viewAuditLogs");
@ -78,8 +81,8 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
};
return (
<div className="box-border w-full overflow-x-hidden p-6">
<Tabs value={activeTab} onValueChange={(value) => setActiveTab(value as LogsTabId)}>
<div className="flex h-full w-full flex-col p-6">
<Tabs value={activeTab} onValueChange={(value) => setActiveTab(value as LogsTabId)} className="min-h-0 flex-1">
<TabsList variant="line">
{tabs.map((tab) => (
<TabsTrigger key={tab.id} value={tab.id} className="flex-none">
@ -88,7 +91,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
))}
</TabsList>
{tabs.map((tab) => (
<TabsContent key={tab.id} value={tab.id} keepMounted>
<TabsContent key={tab.id} value={tab.id} keepMounted className={tabContentClassName(tab.id)}>
{renderPanel(tab.id)}
</TabsContent>
))}

View file

@ -83,9 +83,19 @@ describe("autorouter_presets", () => {
expect(config.tier_boundaries).toBeUndefined();
expect(config.token_thresholds).toBeUndefined();
expect(config.dimension_weights).toBeUndefined();
expect(config.session_affinity_ttl_seconds).toBeUndefined();
}
});
it("carries a preset's session affinity idle window into the prefilled form state", () => {
const config = getPresetByKey("anthropic_family")!.complexity_router_config;
const prefill = buildPresetPrefill({ ...config, session_affinity_ttl_seconds: 300 }, groupsOnly([]));
expect(prefill.complexityRouterConfig.session_affinity_ttl_seconds).toBe(300);
expect(
buildPresetPrefill(config, groupsOnly([])).complexityRouterConfig.session_affinity_ttl_seconds,
).toBeUndefined();
});
it("keeps the model-family presets on the heuristic classifier", () => {
for (const key of ["anthropic_family", "gemini_family", "openai_family"]) {
expect(getPresetByKey(key)!.complexity_router_config.classifier_type).toBe("heuristic");

View file

@ -284,6 +284,7 @@ export const buildPresetPrefill = (
classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns,
classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE,
session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
session_affinity_ttl_seconds: config.session_affinity_ttl_seconds,
deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
modality_routing: config.modality_routing ?? false,
modality_pin_override: config.modality_pin_override ?? false,

View file

@ -25501,6 +25501,12 @@ export interface components {
* @description Documents all the fields supported by `general_settings` in config.yaml
*/
ConfigGeneralSettings: {
/**
* Admission Queue Timeout Seconds
* @description maximum time a request waits for a worker slot
* @default 1
*/
admission_queue_timeout_seconds: number;
/**
* Alert To Webhook Url
* @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`
@ -25709,11 +25715,21 @@ export interface components {
* @description max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider
*/
max_file_size_mb?: number | null;
/**
* Max In Flight Requests Per Worker
* @description maximum concurrent requests handled by each worker
*/
max_in_flight_requests_per_worker?: number | null;
/**
* Max Parallel Requests
* @description maximum parallel requests for each api key
*/
max_parallel_requests?: number | null;
/**
* Max Queued Requests Per Worker
* @description maximum requests waiting for a worker slot
*/
max_queued_requests_per_worker?: number | null;
/**
* Max Request Size Mb
* @description max request size in MB, if a request is larger than this size it will be rejected