diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 32232de381c..0da07038152 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -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 diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 9a9b1138a7d..cb3756575f7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index b6f8bf2dc5b..7df14565c3f 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -48,6 +48,18 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec } +def _build_search_condition(search: str) -> dict[str, object]: + """Match a row whose id, changed_by, object_id, or changed_by_api_key equals the search value.""" + return { + "OR": ( + {"id": search}, + {"changed_by": search}, + {"object_id": search}, + {"changed_by_api_key": search}, + ) + } + + @router.get( "/audit", tags=["Audit Logging"], @@ -83,6 +95,10 @@ async def get_audit_logs( None, description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)", ), + search: str | None = Query( + None, + description="Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value", + ), # Sorting parameters sort_by: str | None = Query( None, @@ -118,6 +134,11 @@ async def get_audit_logs( *([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []), ] + and_conditions: Final[tuple[dict[str, object], ...]] = ( + *json_field_conditions, + *((_build_search_condition(search),) if search else ()), + ) + # Build filter conditions where_conditions: Final[dict[str, object]] = { **({"changed_by": changed_by} if changed_by else {}), @@ -126,14 +147,14 @@ async def get_audit_logs( **({"table_name": table_name} if table_name else {}), **({"object_id": object_id} if object_id else {}), **({"updated_at": date_filter} if start_date or end_date else {}), - **({"AND": json_field_conditions} if json_field_conditions else {}), + **({"AND": and_conditions} if and_conditions else {}), } order_by: Final[dict[str, str]] = ( {sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order} ) - audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table + audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table # Get paginated results audit_logs: Final = await audit_log_table.find_many( @@ -195,7 +216,7 @@ async def get_audit_log_by_id( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) - audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table + audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table # Get the audit log by ID audit_log: Final = await audit_log_table.find_unique(where={"id": id}) diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index ae333d1f4ad..91bcf82f455 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -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): diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 58733b384b9..aaee7188d86 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -27,6 +27,7 @@ from litellm.constants import ( REDIS_CIRCUIT_BREAKER_ENABLED, REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, + REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, ) from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker @@ -41,6 +42,8 @@ from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prometheus_client import Counter as _PromCounter + from prometheus_client import Gauge as _PromGauge from redis.asyncio import Redis, RedisCluster from redis.asyncio.client import Pipeline from redis.asyncio.cluster import ClusterPipeline @@ -135,10 +138,20 @@ class RedisCircuitBreaker: HALF_OPEN - recovery probe: allow one request through Transitions: - CLOSED -> OPEN after failure_threshold consecutive failures + CLOSED -> OPEN after failure_threshold consecutive hard connectivity + failures, or after an unbroken run of timeout failures + (no success or hard failure in between) that reaches + failure_threshold and spans timeout_min_duration seconds OPEN -> HALF_OPEN after recovery_timeout seconds HALF_OPEN -> CLOSED on success HALF_OPEN -> OPEN on failure (resets timer) + + Timeouts are accounted separately from hard connectivity failures because the async + Redis timeout includes time waiting for the worker event loop to resume: one loop + stall makes every in-flight operation time out together, which satisfies a purely + consecutive threshold instantly even though Redis is healthy. Requiring a + timeout-only streak to also span timeout_min_duration filters such bursts while a + real outage that surfaces as timeouts still opens the breaker after that duration. """ CLOSED = "closed" @@ -150,13 +163,19 @@ class RedisCircuitBreaker: failure_threshold: int, recovery_timeout: int, enabled: bool = True, + timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, ) -> None: self.failure_threshold = failure_threshold self.recovery_timeout = recovery_timeout self.enabled = enabled + self.timeout_min_duration = timeout_min_duration self._failure_count = 0 + self._hard_failure_count = 0 + self._timeout_count = 0 + self._timeout_streak_started_at: float | None = None self._opened_at: float | None = None self._state = self.CLOSED + _breaker_metrics().record_state_change(None, self._state) def is_open(self) -> bool: """Returns True if Redis calls should be skipped.""" @@ -169,24 +188,45 @@ class RedisCircuitBreaker: return True if self._state == self.OPEN: if time.time() - (self._opened_at or 0) > self.recovery_timeout: - self._state = self.HALF_OPEN + self._set_state(self.HALF_OPEN) return False # this caller is the designated probe return True return False - def record_failure(self) -> None: + def _should_open(self, now: float) -> bool: + if self._state == self.HALF_OPEN: + return True + if self._hard_failure_count >= self.failure_threshold: + return True + if self._timeout_count < self.failure_threshold: + return False + return now - (self._timeout_streak_started_at or now) >= self.timeout_min_duration + + def record_failure(self, is_timeout: bool = False) -> None: if not self.enabled: return + now: Final = time.time() self._failure_count += 1 - self._opened_at = time.time() - if self._failure_count >= self.failure_threshold: + if is_timeout: + self._timeout_count += 1 + if self._timeout_streak_started_at is None: + self._timeout_streak_started_at = now + else: + self._hard_failure_count += 1 + self._timeout_count = 0 + self._timeout_streak_started_at = None + self._opened_at = now + _breaker_metrics().record_failure("timeout" if is_timeout else "connectivity") + if self._should_open(now): if self._state != self.OPEN: verbose_logger.warning( - "Redis circuit breaker OPENED after %d consecutive failures — fast-failing Redis calls for %ds", + "Redis circuit breaker OPENED after %d consecutive failures" + " (%d hard connectivity) — fast-failing Redis calls for %ds", self._failure_count, + self._hard_failure_count, self.recovery_timeout, ) - self._state = self.OPEN + self._set_state(self.OPEN) def record_success(self) -> None: if not self.enabled: @@ -194,7 +234,17 @@ class RedisCircuitBreaker: if self._state == self.HALF_OPEN: verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered") self._failure_count = 0 - self._state = self.CLOSED + self._hard_failure_count = 0 + self._timeout_count = 0 + self._timeout_streak_started_at = None + self._set_state(self.CLOSED) + + def _set_state(self, state: str) -> None: + if state == self._state: + return + _breaker_metrics().record_transition(state) + _breaker_metrics().record_state_change(self._state, state) + self._state = state _RedisCallResult = TypeVar("_RedisCallResult") @@ -234,6 +284,78 @@ def _is_redis_health_failure(exc: BaseException) -> bool: return True +@functools.lru_cache(maxsize=1) +def _redis_timeout_error_types() -> tuple[type, ...]: + """Health failures that are timeouts rather than unambiguous connectivity errors. + + ``builtins.TimeoutError`` covers ``asyncio.TimeoutError`` and ``socket.timeout`` + (aliases since py3.11 / py3.10). ``redis.exceptions.TimeoutError`` does not subclass + either, so it is listed explicitly. + """ + try: + from redis.exceptions import TimeoutError as RedisTimeoutError + except ImportError: + return (TimeoutError,) + return (RedisTimeoutError, TimeoutError) + + +def _is_redis_timeout_failure(exc: BaseException) -> bool: + return isinstance(exc, _redis_timeout_error_types()) + + +class _BreakerMetrics: + """Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent. + + Registered lazily on the default registry (which /metrics serves) via the module-level + ``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers. + """ + + def __init__(self) -> None: + self._state_gauge: _PromGauge | None = None + self._transitions: _PromCounter | None = None + self._failures: _PromCounter | None = None + try: + from prometheus_client import Counter as PromCounter + from prometheus_client import Gauge + except ImportError: + return + self._state_gauge = Gauge( + "litellm_redis_circuit_breaker_state", + "Number of Redis circuit breakers currently in each state", + labelnames=("state",), + ) + self._transitions = PromCounter( + "litellm_redis_circuit_breaker_transitions", + "Redis circuit breaker state transitions", + labelnames=("state",), + ) + self._failures = PromCounter( + "litellm_redis_circuit_breaker_failures", + "Redis health failures counted by the circuit breaker", + labelnames=("failure_class",), + ) + + def record_state_change(self, old_state: str | None, new_state: str) -> None: + if self._state_gauge is None: + return + if old_state is not None: + self._state_gauge.labels(old_state).dec() + self._state_gauge.labels(new_state).inc() + + def record_transition(self, state: str) -> None: + if self._transitions is not None: + self._transitions.labels(state).inc() + + def record_failure(self, failure_class: str) -> None: + if self._failures is not None: + self._failures.labels(failure_class).inc() + + +@functools.lru_cache(maxsize=1) +def _breaker_metrics() -> _BreakerMetrics: + return _BreakerMetrics() + + def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None: """Record a Redis failure that the calling method is about to swallow. @@ -245,7 +367,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep """ if not _is_redis_health_failure(exc): return - breaker.record_failure() + breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc)) _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) @@ -281,7 +403,7 @@ async def _run_under_circuit_breaker( result: Final = await call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure() + breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) raise _exit_circuit_breaker(breaker, swallowed_before) return result diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py index ae8c78709d9..0035801018b 100644 --- a/litellm/caching/redis_cluster_node_isolation.py +++ b/litellm/caching/redis_cluster_node_isolation.py @@ -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( diff --git a/litellm/constants.py b/litellm/constants.py index f5acadc32ab..da731cb5eb2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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) @@ -432,6 +433,9 @@ REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIME REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)) REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)) REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" +# minimum seconds a timeout-only failure streak must span before it can open the breaker, +# so one event-loop stall timing out many queued calls at once does not trip it +REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)) # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart # (e.g. ElastiCache Serverless maintenance) is not reused while broken diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 16ef6920114..13058bf4f22 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -36,7 +36,7 @@ AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default" def _cached_credential_chain_token_provider() -> Callable[[], str]: return get_azure_ad_token_provider( azure_scope=AZURE_STORAGE_TOKEN_SCOPE, - azure_credential=AzureCredentialType.DefaultAzureCredential, + azure_credential=AzureCredentialType.DeploymentIdentityCredential, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 51d6858b7dc..17a19f05fa3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,7 +10,7 @@ import subprocess import sys import time import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType @@ -5878,14 +5878,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) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 21587af73aa..8e24302b440 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -415,40 +415,64 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = return False -def _coerce_off_peak_rate(value: object, default: float) -> float: +@dataclass(frozen=True, slots=True) +class TokenRates: + input_rate: float + output_rate: float + cache_read_rate: float + cache_creation_rate: float + reasoning_rate: float | None + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + +def _parse_off_peak_rate(value: object) -> float | None: if isinstance(value, bool): - return default + return None if isinstance(value, (int, float)): return float(value) if isinstance(value, str): try: return float(value) except ValueError: - return default - return default + return None + return None -def apply_off_peak_pricing( - model_info: ModelInfo, - current_time: datetime | None, - prompt_base_cost: float, - completion_base_cost: float, - cache_read_cost: float, -) -> tuple[float, float, float]: +def _off_peak_rate(off_peak: Mapping[str, object], key: str, standard_rate: float) -> float: + parsed: Final = _parse_off_peak_rate(off_peak.get(key)) + return standard_rate if parsed is None else parsed + + +def _open_off_peak_block(model_info: ModelInfo, current_time: datetime | None) -> Mapping[str, object] | None: + off_peak: Final = model_info.get("off_peak_pricing") + if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time): + return None + return off_peak + + +def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: """Swap in off-peak per-token rates when the current UTC time is inside one of the model's off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in windows. An off-peak rate replaces the rate that would otherwise apply rather than discounting it, so a model that also has tiered or above-threshold pricing bills the flat off-peak rate for the whole request while the window is open. Any rate left unset in - off_peak_pricing falls back to the standard rate. + off_peak_pricing falls back to the standard rate, so a block without + output_cost_per_reasoning_token keeps the model's own reasoning rate, or its off-peak output + rate when reasoning has no dedicated rate at all. """ - off_peak: Final = model_info.get("off_peak_pricing") - if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time): - return prompt_base_cost, completion_base_cost, cache_read_cost - return ( - _coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost), - _coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost), - _coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost), + off_peak: Final = _open_off_peak_block(model_info, current_time) + if off_peak is None: + return rates + off_peak_reasoning_rate: Final = _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token")) + return TokenRates( + input_rate=_off_peak_rate(off_peak, "input_cost_per_token", rates.input_rate), + output_rate=_off_peak_rate(off_peak, "output_cost_per_token", rates.output_rate), + cache_read_rate=_off_peak_rate(off_peak, "cache_read_input_token_cost", rates.cache_read_rate), + cache_creation_rate=_off_peak_rate(off_peak, "cache_creation_input_token_cost", rates.cache_creation_rate), + reasoning_rate=rates.reasoning_rate if off_peak_reasoning_rate is None else off_peak_reasoning_rate, ) @@ -458,14 +482,28 @@ def _apply_off_peak_to_base_costs( base_costs: tuple[float, float, float, float, float], ) -> tuple[float, float, float, float, float]: """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path - produced them. Cache-creation rates are passed through untouched, since off_peak_pricing - has no field for them. + produced them. The one-hour cache-creation rate passes through untouched, since + off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate. """ prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs - off_peak_prompt, off_peak_completion, off_peak_cache_read = apply_off_peak_pricing( - model_info, current_time, prompt, completion, cache_read + rates: Final = apply_off_peak_pricing( + model_info, + current_time, + TokenRates( + input_rate=prompt, + output_rate=completion, + cache_read_rate=cache_read, + cache_creation_rate=cache_creation, + reasoning_rate=None, + ), + ) + return ( + rates.input_rate, + rates.output_rate, + rates.cache_creation_rate, + cache_creation_above_1hr, + rates.cache_read_rate, ) - return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read) def _get_token_base_cost( @@ -1029,6 +1067,29 @@ def _resolve_reasoning_token_cost( return standard_reasoning_cost if standard_reasoning_cost is not None else completion_base_cost +def _resolve_billed_reasoning_rate( + model_info: ModelInfo, + usage: Usage, + service_tier: str | None, + completion_base_cost: float, + current_time: datetime | None, +) -> float: + off_peak: Final = _open_off_peak_block(model_info, current_time) + off_peak_reasoning_rate: Final = ( + None if off_peak is None else _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token")) + ) + if off_peak_reasoning_rate is not None: + return off_peak_reasoning_rate + tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) + if tiered_reasoning_rate is not None: + return tiered_reasoning_rate + return _resolve_reasoning_token_cost( + model_info=model_info, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + ) + + def generic_cost_per_token( model: str, usage: Usage, @@ -1037,6 +1098,7 @@ def generic_cost_per_token( data_residency: str | None = None, model_info: ModelInfo | None = None, vertex_location: str | None = None, + current_time: datetime | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -1051,6 +1113,7 @@ def generic_cost_per_token( - vertex_location: optional Vertex AI location the request was served from (e.g. "us-east5", "global"), used to apply the per-model regional-endpoint uplift multiplier when non-global. + - current_time: the moment the request is billed at, for off_peak_pricing; defaults to now, UTC Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -1117,6 +1180,7 @@ def generic_cost_per_token( usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0 ) + billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) ( prompt_base_cost, completion_base_cost, @@ -1127,6 +1191,7 @@ def generic_cost_per_token( model_info=model_info, usage=usage, service_tier=service_tier, + current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) @@ -1185,17 +1250,13 @@ def generic_cost_per_token( ## REASONING COST if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0: - tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) - _output_cost_per_reasoning_token = ( - tiered_reasoning_rate - if tiered_reasoning_rate is not None - else _resolve_reasoning_token_cost( - model_info=model_info, - service_tier=service_tier, - completion_base_cost=completion_base_cost, - ) + completion_cost += float(reasoning_tokens) * _resolve_billed_reasoning_rate( + model_info=model_info, + usage=usage, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + current_time=billing_time, ) - completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token ## IMAGE COST if not is_text_tokens_total and image_tokens and image_tokens > 0: @@ -1247,6 +1308,7 @@ def get_token_type_cost_breakdown( service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + current_time: datetime | None = None, ) -> TokenTypeCostBreakdown: """ Provider-agnostic cost of reasoning and cache tokens, derived from the usage @@ -1265,6 +1327,7 @@ def get_token_type_cost_breakdown( except Exception: return TokenTypeCostBreakdown(0.0, 0.0, 0.0) + billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) ( _prompt_base_cost, completion_base_cost, @@ -1275,6 +1338,7 @@ def get_token_type_cost_breakdown( model_info=model_info, usage=usage, service_tier=service_tier, + current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) @@ -1284,18 +1348,12 @@ def get_token_type_cost_breakdown( if not reasoning_tokens: reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - # Reasoning is billed at the selected tier's reasoning rate for tiered models, - # else at the service-tier-aware per-reasoning-token rate - this mirrors how the - # total completion cost is computed, so the breakdown can never diverge from it. - tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) - reasoning_rate: Final = ( - tiered_reasoning_rate - if tiered_reasoning_rate is not None - else _resolve_reasoning_token_cost( - model_info=model_info, - service_tier=service_tier, - completion_base_cost=completion_base_cost, - ) + reasoning_rate: Final = _resolve_billed_reasoning_rate( + model_info=model_info, + usage=usage, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + current_time=billing_time, ) reasoning_cost = float(reasoning_tokens) * reasoning_rate diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 9deff950724..242300c7b6d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -6,6 +6,7 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +from litellm.types.llms.openai import ChatCompletionSystemMessage if TYPE_CHECKING: from litellm.exceptions import ContentPolicyViolationError @@ -36,6 +37,16 @@ def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> " ) +def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMessage | None: + """ + Return the Anthropic Messages top-level ``system`` (a string or a list of text + blocks) as an OpenAI-style system message, or None when the request has none. + """ + if not isinstance(system, (str, list)) or not system: + return None + return ChatCompletionSystemMessage(role="system", content=system) + + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: """ diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index d8eb1f9f8d7..17f70ec5db7 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,12 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ -from dataclasses import dataclass, replace +from dataclasses import dataclass from datetime import datetime from typing import Final from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.litellm_core_utils.llm_cost_calc.utils import ( + TokenRates, apply_off_peak_pricing, parse_completion_tokens_details, parse_prompt_tokens_details, @@ -34,19 +35,6 @@ class TokenBreakdown: return self.text_tokens + self.cached_tokens + self.cache_creation_tokens -@dataclass(frozen=True, slots=True) -class TokenRates: - input_rate: float - cache_read_rate: float - cache_creation_rate: float - output_rate: float - reasoning_rate: float | None - - @property - def billed_reasoning_rate(self) -> float: - return self.output_rate if self.reasoning_rate is None else self.reasoning_rate - - def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: prompt_details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = prompt_details["cache_hit_tokens"] @@ -105,13 +93,6 @@ def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: ) -def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: - input_rate, output_rate, cache_read_rate = apply_off_peak_pricing( - model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate - ) - return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate) - - def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]: prompt_cost: Final = ( (breakdown.text_tokens * rates.input_rate) @@ -155,6 +136,6 @@ def cost_per_token( else None ) standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier) - rates: Final = _off_peak_rates(model_info, current_time, standard_rates) + rates: Final = apply_off_peak_pricing(model_info, current_time, standard_rates) return _bill(breakdown, rates) diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index ecec161ed46..283fdfb92c2 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -8,7 +8,7 @@ from urllib.parse import urlparse import litellm from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str -from .common_utils import OpenAIError +from .common_utils import OpenAIError, is_openai_backed_api_base if TYPE_CHECKING: from collections.abc import Callable @@ -16,7 +16,6 @@ if TYPE_CHECKING: from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth OPENAI_WIF_CLIENT_ID: Final = "litellm" -_OPENAI_API_HOST: Final = "api.openai.com" _SDK_UPGRADE_MESSAGE: Final = ( "OpenAI workload identity federation requires openai>=2.32.0. " "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / " @@ -75,7 +74,7 @@ def _targets_openai_api(api_base: str | None) -> bool: if api_base is None: return True parsed: Final = urlparse(api_base) - return parsed.scheme == "https" and parsed.hostname == _OPENAI_API_HOST + return parsed.scheme == "https" and is_openai_backed_api_base(api_base) @lru_cache(maxsize=16) diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 0968185b084..c64fc583edc 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -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, - ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 31cea8d8ce2..ea0a9755caa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3656,30 +3656,48 @@ class MCPServerManager: user_api_key_auth: UserAPIKeyAuth | None, raw_headers: Mapping[str, str] | None = None, ) -> None: - """Run the OBO exchange for a caller-supplied subject at the transport edge. + """Mint an exchange-backed server's upstream credential at the transport edge. Single-server routes call this before the MCP session opens, where an HTTP status and ``WWW-Authenticate`` still reach the client. A rejected subject raises the RFC 9728 challenge and any other ``CredError`` maps onto its public HTTP status, so an exchange failure surfaces as a failure instead of the session continuing into an empty tool list. A successful exchange is cached by the exchanger, so the session's list/call reuses it. + + Each mode pre-flights only where it would resolve the subject the session goes on to use, + which is what keeps the pre-flight from reaching a verdict the session would contradict. + ``oauth2_token_exchange`` mints from the caller's inbound bearer, so without one there is + nothing to exchange and the missing-subject case stays the preemptive challenge's job. + ``oauth2_id_jag`` is the mirror image: tool listing resolves it from the identity assertion + captured for this user at SSO login and never from the inbound bearer, so the pre-flight is + faithful exactly when no identity bearer was sent (a LiteLLM key in ``Authorization`` is not one), + and a caller that did send one is passed through + untouched rather than judged against a subject the listing will not use. That store-sourced + case is the one whose missing-assertion 412 and store-outage 503 the session cannot report. + Only OBO has a discovery challenge to raise; ID-JAG's failures are plain statuses whose body + already names what the user has to do, so they map through ``raise_public`` as at egress. """ - if server.auth_type != MCPAuth.oauth2_token_exchange: - return - if not self._extract_bearer_token(oauth2_headers, None): - return - resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) - spec: Final = to_server_spec(resolved_server) - if spec is None or not isinstance(spec.config, TokenExchangeConfig): - return subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) - if subject_token is None: + match server.auth_type: + case MCPAuth.oauth2_token_exchange: + if not self._extract_bearer_token(oauth2_headers, None): + return + case MCPAuth.oauth2_id_jag: + if subject_token is not None: + return + case _: + return + resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) + spec: Final = _to_server_spec_fail_closed(resolved_server) + if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): + return + if subject_token is None and isinstance(spec.config, TokenExchangeConfig): raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path()) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return case Error(err): - if err.tag == "unauthorized": + if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): raise_token_exchange_challenge( resolved_server, root_path=get_server_root_path(), diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 8a0d41584f9..6552008ca54 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index bb075220530..3d7c947a913 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3851,15 +3851,15 @@ if MCP_AVAILABLE: raise_token_exchange_challenge(server, root_path=get_server_root_path()) - # token_exchange (OBO) with a subject present: run the exchange here at the transport - # edge, so a rejected subject raises the RFC 9728 challenge (and a gateway fault its - # public status) instead of the session opening and list_tools masking the failure as - # an empty tool list. Gated to single-server routes; the multi-server aggregate keeps - # absorbing per-server auth failures so one bad server cannot 401 the whole connect. + # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run + # the exchange here at the transport edge, so a rejected subject raises the RFC 9728 + # challenge and any other failure its public status, instead of the session opening and + # list_tools masking it as an empty tool list. The manager owns which modes pre-flight + # and what each mints from. Gated to single-server routes the key may reach; the + # multi-server aggregate keeps absorbing per-server auth failures so one bad server + # cannot 401 the whole connect. if ( server - and server.auth_type == MCPAuth.oauth2_token_exchange - and oauth2_headers and len(mcp_servers or []) == 1 and server.server_id in frozenset( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5d5a25e7cd6..1a807fd39bb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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" ) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 7f0045c1d93..b243b737b0a 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -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, ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 1625198892f..de2ca4762f1 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -387,33 +387,22 @@ def _get_wildcard_models( all_wildcard_models: Final = [] for model in unique_models: if _check_wildcard_routing(model=model): - if return_wildcard_routes: # will add the wildcard route to the list eg: anthropic/*. + if return_wildcard_routes: all_wildcard_models.append(model) - ## get litellm params from model - if llm_router is not None: - model_list = llm_router.get_model_list(model_name=model, team_id=team_id) - if model_list: - for router_model in model_list: - wildcard_models = get_known_models_from_wildcard( + models_to_remove.add(model) + + model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router else None + if model_list: + for router_model in model_list: + all_wildcard_models.extend( + get_known_models_from_wildcard( wildcard_model=model, litellm_params=LiteLLM_Params(**router_model["litellm_params"]), ) - all_wildcard_models.extend(wildcard_models) - else: - # Router has no deployment for this wildcard (e.g., BYOK team models) - # Fall back to expanding from known provider models - wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None) - if wildcard_models: - models_to_remove.add(model) - all_wildcard_models.extend(wildcard_models) + ) else: - # get all known provider models - wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None) - - if wildcard_models: - models_to_remove.add(model) - all_wildcard_models.extend(wildcard_models) + all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)) for model in models_to_remove: unique_models.remove(model) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8b57bdca2fe..65d0ec8c0dc 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -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( diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 0357bc7dbc6..1f91eeedf64 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -1,4 +1,5 @@ from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, status @@ -20,7 +21,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache from litellm.proxy.utils import get_prisma_client_or_throw -from litellm.repositories.table_repositories import AccessGroupRepository +from litellm.repositories.table_repositories import AccessGroupRepository, TeamRepository from litellm.types.access_group import ( AccessGroupCreateRequest, AccessGroupResponse, @@ -74,11 +75,11 @@ class _AccessGroupTable(Protocol): class _TeamTable(Protocol): - async def find_unique(self, where: Mapping[str, object]) -> _TeamRecord | None: ... + async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ... - async def find_many(self, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ... - async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... class _KeyTable(Protocol): @@ -119,8 +120,56 @@ def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None: ) -def _record_to_response(record: _AccessGroupRecord) -> AccessGroupResponse: - return AccessGroupResponse.model_validate(record.dict()) +def _record_to_response( + record: _AccessGroupRecord, *, assigned_team_ids: Sequence[str] | None = None +) -> AccessGroupResponse: + stored: Final = record.dict() + payload: Final = ( + stored if assigned_team_ids is None else MappingProxyType({**stored, "assigned_team_ids": assigned_team_ids}) + ) + return AccessGroupResponse.model_validate(payload) + + +def _attached_team_ids_by_group( + records: Sequence[_AccessGroupRecord], teams: Sequence[_TeamRecord] +) -> Mapping[str, tuple[str, ...]]: + """Teams really attached to each group: the stored column minus ghosts, plus teams the mirror missed.""" + real_team_ids: Final = frozenset(team.team_id for team in teams) + + def attached(record: _AccessGroupRecord) -> tuple[str, ...]: + stored: Final = (team_id for team_id in (record.assigned_team_ids or ()) if team_id in real_team_ids) + carrying: Final = (team.team_id for team in teams if record.access_group_id in (team.access_group_ids or ())) + return tuple(dict.fromkeys((*stored, *carrying))) + + return MappingProxyType({record.access_group_id: attached(record) for record in records}) + + +async def _attached_team_ids_for( + team_table: _TeamTable, records: Sequence[_AccessGroupRecord] +) -> Mapping[str, tuple[str, ...]]: + if not records: + return MappingProxyType({}) + group_ids: Final = tuple(record.access_group_id for record in records) + stored_team_ids: Final = tuple( + dict.fromkeys(team_id for record in records for team_id in (record.assigned_team_ids or ())) + ) + carrying: Final = {"access_group_ids": {"hasSome": group_ids}} # mutable-ok: prisma where is a dict + listed: Final = {"team_id": {"in": stored_team_ids}} # mutable-ok: prisma where is a dict + teams: Final = await team_table.find_many(where={"OR": (carrying, listed)}) # mutable-ok: prisma where is a dict + return _attached_team_ids_by_group(records, teams) + + +async def _require_teams_exist(tx: _AccessGroupTx, team_ids: Sequence[str]) -> None: + if not team_ids: + return + where: Final = {"team_id": {"in": team_ids}} # mutable-ok: prisma where is a dict + found: Final = await tx.litellm_teamtable.find_many(where=where) + missing: Final = frozenset(team_ids) - frozenset(team.team_id for team in found) + if missing: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Unknown team ids: {', '.join(sorted(missing))}", + ) def _record_to_access_group_table(record: _AccessGroupRecord) -> LiteLLM_AccessGroupTable: @@ -330,6 +379,7 @@ async def create_access_group( status_code=status.HTTP_409_CONFLICT, detail=f"Access group '{data.access_group_name}' already exists", ) + await _require_teams_exist(tx, data.assigned_team_ids or ()) record: Final = await tx.litellm_accessgrouptable.create( data={ @@ -390,7 +440,8 @@ async def list_access_groups( table: Final = AccessGroupRepository(prisma_client).table records: Final = await table.find_many(order={"created_at": "desc"}) - return [_record_to_response(r) for r in records] + attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, records) + return [_record_to_response(r, assigned_team_ids=attached[r.access_group_id]) for r in records] @router.get( @@ -411,7 +462,8 @@ async def get_access_group( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", ) - return _record_to_response(record) + attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, (record,)) + return _record_to_response(record, assigned_team_ids=attached[record.access_group_id]) @router.put( @@ -461,8 +513,10 @@ async def update_access_group( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", ) + await _require_teams_exist(tx, data.assigned_team_ids or ()) - old_team_ids: Final[set[str]] = set(existing.assigned_team_ids or []) + attached: Final = await _attached_team_ids_for(tx.litellm_teamtable, (existing,)) + old_team_ids: Final[set[str]] = set(attached[access_group_id]) old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or []) new_team_ids: Final[set[str]] = ( set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d7d20d168b5..324d380b85b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -147,6 +147,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyResponse, BulkUpdateTeamKeysRequest, FailedKeyUpdate, + KeySearchWhere, SuccessfulKeyUpdate, ) from litellm.types.router import Deployment @@ -5800,6 +5801,10 @@ async def list_keys( None, description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.", ), + search: str | None = Query( + None, + description="Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive).", + ), return_full_object: bool = Query(False, description="Return full key object"), include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."), include_created_by_keys: bool = Query(False, description="Include keys created by the user"), @@ -5943,6 +5948,7 @@ async def list_keys( agent_id=agent_id, use_substring_matching=use_substring_matching, expires_filter=expires if isinstance(expires, str) else None, + search=search, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -6162,6 +6168,16 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, return {"OR": [{"expires": None}, {"expires": {"gte": now}}]} +def _build_key_search_where(search: str) -> KeySearchWhere: + search_where: Final[KeySearchWhere] = { + "OR": ( + {"token": search}, + {"key_alias": {"contains": search, "mode": "insensitive"}}, + ) + } + return search_where + + def _build_key_filter_conditions( user_id: str | None, team_id: str | None, @@ -6177,6 +6193,7 @@ def _build_key_filter_conditions( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, + search: str | None = None, ) -> Mapping[str, object]: """Build filter conditions for key listing. @@ -6266,7 +6283,7 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) - global_filters: Final[tuple[dict[str, object], ...]] = ( + global_filters: Final[tuple[Mapping[str, object], ...]] = ( *( ( {"key_alias": {"contains": key_alias, "mode": "insensitive"}} @@ -6277,6 +6294,7 @@ def _build_key_filter_conditions( else () ), *(({"token": key_hash},) if key_hash and isinstance(key_hash, str) else ()), + *((_build_key_search_where(search),) if isinstance(search, str) and search else ()), *(({"team_id": team_id},) if team_id and isinstance(team_id, str) else ()), *(({"project_id": project_id},) if project_id else ()), *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()), @@ -6316,6 +6334,7 @@ async def _list_key_helper( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, + search: str | None = None, ) -> KeyListResponseObject: """ Helper function to list keys @@ -6354,6 +6373,7 @@ async def _list_key_helper( agent_id=agent_id, use_substring_matching=use_substring_matching, expires_filter=expires_filter, + search=search, ) # Calculate skip for pagination diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index f58f3722741..69356922ea1 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -12,7 +12,7 @@ All /policy management endpoints import copy import json import os -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Request @@ -20,6 +20,7 @@ from fastapi.responses import Response, StreamingResponse from pydantic import BaseModel, Field from typing_extensions import TypedDict +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( COMPETITOR_LLM_TEMPERATURE, @@ -32,6 +33,10 @@ from litellm.llms.openai.chat.guardrail_translation.handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.sse_keepalive import ( + SSE_COMMENT_PING, + wrap_sse_stream_with_keepalive_pings, +) from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail, @@ -811,7 +816,7 @@ async def _stream_competitor_events( llm_enrichment: dict, brand_name: str, model: str, -) -> AsyncIterator[str]: +) -> AsyncGenerator[str, None]: """Stream competitor names as SSE events, then emit a final 'done' event.""" competitors: Final[list[str]] = list(data.competitors or []) @@ -883,7 +888,11 @@ async def enrich_policy_template_stream( model: Final = data.model or DEFAULT_COMPETITOR_DISCOVERY_MODEL return StreamingResponse( - _stream_competitor_events(data, template, llm_enrichment, brand_name, model), + wrap_sse_stream_with_keepalive_pings( + _stream_competitor_events(data, template, llm_enrichment, brand_name, model), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + ping_chunk=SSE_COMMENT_PING, + ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 069f86c852c..ceb67e3eee8 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1123,13 +1123,15 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou user_id=user_id, user_email=user_id, # We don't have email from group membership user_alias=None, - teams=[], # Teams will be added separately metadata={"created_via": created_via}, auto_create_key=False, user_role=default_role, ) - created_user: Final = await new_user(data=new_user_request) + created_user: Final = await new_user( + data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) verbose_proxy_logger.info("Created user %s via %s", user_id, created_via) return created_user @@ -1699,7 +1701,7 @@ async def create_user( user_id=user_id, user_email=user_data["user_email"], user_alias=user_data["user_alias"], - teams=user_data["teams"], + teams=user_data["teams"] or None, metadata=metadata, auto_create_key=False, user_role=resolved_role if admin_group is not None else default_role, @@ -1717,6 +1719,7 @@ async def create_user( created_user: Final = await new_user( data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) scim_user: Final = await ScimTransformations.transform_litellm_user_to_scim_user(user=created_user) @@ -1771,22 +1774,25 @@ async def update_user( roles=user_data["roles"], ) + # SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send + # no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team" + target_teams: Final = user_data["teams"] or existing_user.teams await _handle_team_membership_changes( user_id=user_id, - existing_teams=existing_user.teams or [], - new_teams=user_data["teams"], + existing_teams=existing_user.teams, + new_teams=target_teams, ) update_data: Final = { "user_email": user_data["user_email"], "user_alias": user_data["user_alias"], "sso_user_id": user_data["sso_user_id"], - "teams": user_data["teams"], + "teams": target_teams, "metadata": safe_dumps(metadata), } admin_group: Final = await _get_scim_admin_group() - if admin_group is not None: + if admin_group is not None and user_data["teams"]: update_data["user_role"] = _resolve_scim_user_role( user.groups or [], admin_group, _default_scim_user_role() ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 9d5ddda017a..7ef496f94e3 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -4,7 +4,7 @@ usage/spend data by querying the aggregated daily activity endpoints. """ import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import date from typing import Any, Final, Literal, Protocol, cast, overload @@ -543,7 +543,7 @@ async def stream_usage_ai_chat( model: str | None = None, user_id: str | None = None, is_admin: bool = False, -) -> AsyncIterator[str]: +) -> AsyncGenerator[str, None]: """Stream SSE events: status → tool_call → chunk → done.""" resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py index b7b0ae2d8e5..d1e92d0c7e4 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -10,8 +10,13 @@ from fastapi import APIRouter, Depends, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field +import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.sse_keepalive import ( + SSE_COMMENT_PING, + wrap_sse_stream_with_keepalive_pings, +) router: Final = APIRouter() @@ -56,11 +61,15 @@ async def usage_ai_chat( messages: Final = [{"role": m.role, "content": m.content} for m in data.messages] return StreamingResponse( - stream_usage_ai_chat( - messages=messages, - model=data.model, - user_id=user_id, - is_admin=is_admin, + wrap_sse_stream_with_keepalive_pings( + stream_usage_ai_chat( + messages=messages, + model=data.model, + user_id=user_id, + is_admin=is_admin, + ), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + ping_chunk=SSE_COMMENT_PING, ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 98c5fdd198c..d8f72d200c7 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -22,6 +22,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final from fastapi import APIRouter, Depends, HTTPException, Query +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -91,6 +92,36 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object return {"OR": ors} +class _StartsWith(TypedDict): + startsWith: ReadOnly[str] + + +class _MemoryKeyWhere(TypedDict): + key: ReadOnly[str | _StartsWith] + + +class _MemoryIdWhere(TypedDict): + memory_id: ReadOnly[str] + + +class _MemorySearchWhere(TypedDict): + OR: ReadOnly[tuple[_MemoryKeyWhere, _MemoryIdWhere]] + + +def _key_filter(search: str | None, key_prefix: str | None, key: str | None) -> Mapping[str, object] | None: + """`search` matches a key prefix or an exact memory_id; otherwise `key_prefix` wins over `key`.""" + if search is not None: + search_where: Final[_MemorySearchWhere] = {"OR": ({"key": {"startsWith": search}}, {"memory_id": search})} + return search_where + if key_prefix is not None: + prefix_where: Final[_MemoryKeyWhere] = {"key": {"startsWith": key_prefix}} + return prefix_where + if key is not None: + exact_where: Final[_MemoryKeyWhere] = {"key": key} + return exact_where + return None + + def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow: return LiteLLM_MemoryRow( memory_id=row.memory_id, @@ -326,6 +357,13 @@ async def list_memory( "Mutually exclusive with `key`; if both are provided, `key_prefix` wins." ), ), + search: str | None = Query( + None, + description=( + "Match entries whose key starts with this value or whose memory_id equals it. " + "Takes precedence over `key_prefix` and `key` when provided." + ), + ), page: int = Query(1, ge=1), page_size: int = Query(50, ge=1, le=500), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -333,22 +371,16 @@ async def list_memory( """List memory entries visible to the caller.""" prisma_client: Final = _require_prisma() - # Build the key filter first (prefix wins if both `key` and `key_prefix` - # are passed). Then AND it with the visibility filter via an explicit - # top-level "AND" — safer than `dict.update` since future visibility - # filters could grow an "OR" key that would clobber this one if merged - # by key. - key_filter: Final[dict[str, object]] = {} - if key_prefix is not None: - key_filter["key"] = {"startsWith": key_prefix} - elif key is not None: - key_filter["key"] = key + # AND the key filter with the visibility filter via an explicit top-level + # "AND": both sides can carry an "OR" key (`search`, non-admin visibility), + # so merging them by key would let one clobber the other and leak rows. + key_filter: Final = _key_filter(search=search, key_prefix=key_prefix, key=key) vis: Final = _visibility_filter(user_api_key_dict) - where: Mapping[str, object] + where: Mapping[str, object] | None if vis is None: where = key_filter - elif not key_filter: + elif key_filter is None: where = vis else: where = {"AND": [key_filter, vis]} diff --git a/litellm/proxy/middleware/admission_control_middleware.py b/litellm/proxy/middleware/admission_control_middleware.py new file mode 100644 index 00000000000..aa62ef9e3bf --- /dev/null +++ b/litellm/proxy/middleware/admission_control_middleware.py @@ -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", + } + }, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 29f216fd450..6b1d6405a6a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -9,10 +9,11 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. from __future__ import annotations import hmac +import inspect import json import os import re -from collections.abc import Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, cast @@ -32,6 +33,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.passthrough.main import AsyncPassthroughStreamingResponse from litellm.proxy._types import * from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks @@ -40,6 +42,7 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, user_api_key_auth_websocket, ) +from litellm.proxy.common_request_processing import open_sse_before_first_byte from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -47,6 +50,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( get_form_data, get_request_body, ) +from litellm.proxy.common_utils.sse_keepalive import ( + wrap_passthrough_sse_bytes_with_keepalive_pings, +) from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, @@ -1478,6 +1484,74 @@ def is_azure_ai_search_service_level_index_create(method: str, endpoint: str) -> return path == "indexes" or path.endswith("/indexes") +async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> AsyncGenerator[bytes, None]: + try: + async for chunk in upstream: + yield chunk + finally: + await upstream.aclose() + + +async def _relay_azure_router_model( + llm_router: litellm.Router, + model: str, + endpoint: str, + request: Request, + request_body: Mapping[str, object], + is_streaming_request: bool, + user_api_key_dict: UserAPIKeyAuth, +) -> Response: + result: Final = await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=_safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if request.headers.get("content-type") == "application/json" else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + ) + + if not is_streaming_request: + upstream: Final = cast(httpx.Response, result) + return Response( + content=await upstream.aread(), + status_code=upstream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), + ) + + if inspect.isasyncgen(result): + sse_headers: Final = {"content-type": "text/event-stream"} + return StreamingResponse( + content=wrap_passthrough_sse_bytes_with_keepalive_pings( + stream=_relay_upstream_bytes(result), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + upstream_headers=sse_headers, + ), + status_code=200, + headers=sse_headers, + ) + + upstream_stream: Final = cast(AsyncPassthroughStreamingResponse, result) + return StreamingResponse( + content=wrap_passthrough_sse_bytes_with_keepalive_pings( + stream=_relay_upstream_bytes(upstream_stream), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + upstream_headers=upstream_stream.headers, + ), + status_code=upstream_stream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers( + headers=upstream_stream.headers, custom_headers=None + ), + ) + + @router.api_route( "/azure_ai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -1528,55 +1602,18 @@ async def azure_proxy_route( if is_router_model: request_body = await get_request_body(request) is_streaming_request = is_passthrough_request_streaming(request_body) - result = await llm_router.allm_passthrough_route( - model=part, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=_safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), - ) - - if is_streaming_request: - # Check if result is an async generator (from _async_streaming) - import inspect - - if inspect.isasyncgen(result): - # Result is already an async generator, use it directly - return StreamingResponse( - content=result, - status_code=200, - headers={"content-type": "text/event-stream"}, - ) - else: - # Result is an httpx.Response, use aiter_bytes() - result = cast(httpx.Response, result) - return StreamingResponse( - content=result.aiter_bytes(), - status_code=result.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=result.headers, - custom_headers=None, - ), - ) - - # Non-streaming response - result = cast(httpx.Response, result) - content = await result.aread() - return Response( - content=content, - status_code=result.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=result.headers, - custom_headers=None, + return await open_sse_before_first_byte( + _relay_azure_router_model( + llm_router=llm_router, + model=part, + endpoint=endpoint, + request=request, + request_body=request_body, + is_streaming_request=is_streaming_request, + user_api_key_dict=user_api_key_dict, + ), + ping_interval_seconds=( + litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None ), ) elif is_vector_store_index: @@ -1659,6 +1696,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//locations// 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 +1709,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 +1943,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 +1954,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 +2012,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 +2106,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 +2118,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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c43cc510990..1f39a78e12a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, ) @@ -15233,20 +15238,33 @@ async def async_queue_request( if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - - response: Final = await llm_router.schedule_acompletion(**data) + router: Final = llm_router if "stream" in data and data["stream"] is True: # use generate_responses to stream responses - return StreamingResponse( - async_data_generator( - user_api_key_dict=user_api_key_dict, - response=response, - request_data=data, - request=request, - ), - media_type="text/event-stream", + + async def produce_queue_stream() -> StreamingResponse: + return StreamingResponse( + async_data_generator( + user_api_key_dict=user_api_key_dict, + response=await router.schedule_acompletion(**data), + request_data=data, + request=request, + ), + media_type="text/event-stream", + ) + + async def audit_late_failure(exc: Exception) -> HTTPException | None: + return await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=data + ) + + return await open_sse_before_first_byte( + produce_queue_stream(), + ping_interval_seconds=ttft_keepalive_interval(data, router), + on_late_failure=audit_late_failure, ) + response: Final = await router.schedule_acompletion(**data) fastapi_response.headers.update({"x-litellm-priority": str(data["priority"])}) return response except Exception as e: @@ -16502,6 +16520,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 +18198,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": diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index c8c6c505375..d6a402e1860 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -23,11 +23,14 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import * from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + open_sse_before_first_byte, + ttft_keepalive_interval, +) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -48,6 +51,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) from litellm.repositories.table_repositories import ManagedVectorStoresRepository +from litellm.types.utils import ModelResponse if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -756,43 +760,53 @@ async def rag_query( merged_retrieval_config.get("custom_llm_provider"), ) - # Call query - response: Final = await litellm.aquery( - model=model, - messages=messages, - retrieval_config=merged_retrieval_config, - vector_store_params=store_data, - rerank=rerank, - stream=stream, - router=llm_router, - **request_data, - ) - - hidden_params: Final = getattr(response, "_hidden_params", {}) or {} - custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=hidden_params.get("litellm_call_id", None) or "", - model_id=hidden_params.get("model_id", None) or "", - cache_key=hidden_params.get("cache_key", None) or "", - api_base=hidden_params.get("api_base", None) or "", - version=version, - response_cost=hidden_params.get("response_cost", None), - request_data=request_data, - ) - - if isinstance(response, CustomStreamWrapper): - return StreamingResponse( - select_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=request_data, - request=request, - ), - media_type="text/event-stream", - headers=custom_headers, + async def query() -> ModelResponse: + return await litellm.aquery( + model=model, + messages=messages, + retrieval_config=merged_retrieval_config, + vector_store_params=store_data, + rerank=rerank, + stream=stream, + router=llm_router, + **request_data, ) - fastapi_response.headers.update(custom_headers) + def custom_headers_for(response: ModelResponse) -> Mapping[str, str]: + hidden_params: Final = getattr(response, "_hidden_params", {}) or {} + return ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=hidden_params.get("litellm_call_id", None) or "", + model_id=hidden_params.get("model_id", None) or "", + cache_key=hidden_params.get("cache_key", None) or "", + api_base=hidden_params.get("api_base", None) or "", + version=version, + response_cost=hidden_params.get("response_cost", None), + request_data=request_data, + ) + + if stream: + + async def produce_stream() -> StreamingResponse: + response: Final = await query() + return StreamingResponse( + select_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + request=request, + ), + media_type="text/event-stream", + headers=custom_headers_for(response), + ) + + return await open_sse_before_first_byte( + produce_stream(), + ping_interval_seconds=ttft_keepalive_interval(data, llm_router), + ) + + response: Final = await query() + fastapi_response.headers.update(custom_headers_for(response)) return response except HTTPException: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 2a50d5170f0..dff100bdea7 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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], @@ -2229,6 +2285,31 @@ async def calculate_spend(request: SpendCalculateRequest): ) +class _SpendLogSearchCondition(NamedTuple): + sql: str + params: tuple[object, ...] + + +def _build_spend_log_search_condition( + search: str, + start_date: datetime, + end_date: datetime, + next_param_index: int, +) -> _SpendLogSearchCondition: + """request_id (indexed) matches across all time; the unindexed id columns only inside the window.""" + raw: Final = f"${next_param_index}" + window_start: Final = f"${next_param_index + 1}" + window_end: Final = f"${next_param_index + 2}" + sql: Final = ( + f"(request_id = {raw} OR (" + f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') " + f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') " + f'AND (api_key = {raw} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} ' + f"OR session_id = {raw} OR model_id = {raw})))" + ) + return _SpendLogSearchCondition(sql=sql, params=(search, start_date, end_date)) + + @router.get( "/spend/logs/v2", tags=["Budget & Spend Tracking"], @@ -2329,6 +2410,14 @@ async def ui_view_spend_logs( "UI route only, honored when sorting by startTime" ), ), + search: str | None = fastapi.Query( + default=None, + description=( + "Match a log whose request_id, api_key (hash), team_id, user, end_user, " + "session_id, or model_id equals this value. request_id matches across all time; the other columns " + "match inside start_date/end_date, which stay required" + ), + ), ): """ View spend logs with pagination support. @@ -2392,8 +2481,10 @@ async def ui_view_spend_logs( try: is_admin_view: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) is_request_id_lookup: Final = request_id is not None and not is_v2 + is_search_lookup: Final = search is not None + search_owns_window: Final = is_search_lookup and not is_v2 - if is_request_id_lookup: + if is_request_id_lookup and not is_search_lookup: # request_id is the @id primary key: it identifies a single row, so a # time window is meaningless. The dashboard always sends a default 24h # window, which hid ids copied from an older page (LIT-3981). Drop the @@ -2576,7 +2667,7 @@ async def ui_view_spend_logs( # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison # against the plain `timestamp` column does not depend on the DB session # timezone (see #22529). Absent for a request_id-only lookup (see above). - if start_date_obj is not None and end_date_obj is not None: + if start_date_obj is not None and end_date_obj is not None and not search_owns_window: sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") sql_params.append(start_date_obj) p += 1 @@ -2584,6 +2675,17 @@ async def ui_view_spend_logs( sql_params.append(end_date_obj) p += 1 + if search is not None and start_date_obj is not None and end_date_obj is not None: + search_condition: Final = _build_spend_log_search_condition( + search=search, + start_date=start_date_obj, + end_date=end_date_obj, + next_param_index=p, + ) + sql_conditions.append(search_condition.sql) + sql_params.extend(search_condition.params) + p += len(search_condition.params) # rebind-ok: advances the file's shared $N placeholder counter + # Equality filters - read effective values from where_conditions (post-authorization) for sql_col, wc_key in [ ("team_id", "team_id"), @@ -2662,7 +2764,13 @@ async def ui_view_spend_logs( sql_params.append(f"%{error_message}%") p += 1 - if group_by_session is True and not is_v2 and not is_request_id_lookup and sort_by == "startTime": + if ( + group_by_session is True + and not is_v2 + and not is_request_id_lookup + and not is_search_lookup + and sort_by == "startTime" + ): return await _ui_session_grouped_spend_logs( prisma_client=prisma_client, sql_conditions=sql_conditions, @@ -2696,7 +2804,7 @@ async def ui_view_spend_logs( _order_expr = order_column joined_conditions: Final = " AND ".join(sql_conditions) - session_grouping: Final = group_by_session is True + session_grouping: Final = group_by_session is True and not is_search_lookup count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else "" count_query: Final = f""" SELECT COUNT(*) AS total_count @@ -3214,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): @@ -3244,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]] = {} diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 18cf884f267..739d6ada71c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -176,6 +176,10 @@ class PolicyAttachmentRepository(PrismaTableRepository["prisma_models.LiteLLM_Po table_name = "litellm_policyattachmenttable" +class TeamRepository(PrismaTableRepository["prisma_models.LiteLLM_TeamTable"]): + table_name = "litellm_teamtable" + + class DeletedTeamRepository(PrismaTableRepository["prisma_models.LiteLLM_DeletedTeamTable"]): table_name = "litellm_deletedteamtable" diff --git a/litellm/router.py b/litellm/router.py index f33dfbba7bf..6da201725b6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -197,6 +197,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import ( from litellm.scheduler import FlowItem, Scheduler from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionToolParam, FileTypes, OpenAIFileObject, OpenAIFilesPurpose, @@ -9545,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) @@ -11762,7 +11764,7 @@ class Router: self, messages: list[dict[str, str]] | None, input: str | list | None, - instructions: str | None = None, + request_kwargs: Mapping[str, object] | None = None, ) -> int: """ Count input tokens for context-window pre-call checks. @@ -11772,9 +11774,28 @@ class Router: The Responses payload is normalized to chat messages via the shared LiteLLMCompletionResponsesConfig transform so the same token_counter path covers both API surfaces and `instructions` tokens are included in the count. + + Prompt content the message list never carries is read from `request_kwargs`: + `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the + Anthropic Messages top-level `system` block. """ + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + anthropic_system_to_openai_message, + ) + + extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({}) + raw_instructions: Final = extras.get("instructions") + instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None + raw_tools: Final = extras.get("tools") + tools: Final = ( + cast(list[ChatCompletionToolParam], raw_tools) # cast-ok: token_counter formats any tool dict shape + if isinstance(raw_tools, list) and raw_tools + else None + ) + system_message: Final = anthropic_system_to_openai_message(extras.get("system")) if messages is not None: - return litellm.token_counter(messages=messages) + counted_messages: Final = (system_message, *messages) if system_message is not None else messages + return litellm.token_counter(messages=counted_messages, tools=tools) if input is not None: from openai.types.responses.response_create_params import ResponseInputParam @@ -11787,7 +11808,10 @@ class Router: input=typed_input, responses_api_request={"instructions": instructions} if instructions is not None else {}, ) - return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages + return litellm.token_counter( + messages=cast(list, input_messages), # cast-ok: transformed chat messages + tools=tools, + ) raise ValueError("Either messages or input must be provided to count tokens") def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None: @@ -11833,14 +11857,13 @@ class Router: """ if messages is None and input is None: return None - raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None try: if not self._pre_call_checks_need_token_count(model, healthy_deployments): return None return await asyncify(self._count_pre_call_check_tokens)( messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter - instructions=raw_instructions if isinstance(raw_instructions, str) else None, + request_kwargs=request_kwargs, ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( @@ -11887,8 +11910,6 @@ class Router: _rate_limit_error = False parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs) - raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None - instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None has_countable_input: Final = messages is not None or input is not None ## get model group RPM ## @@ -11919,7 +11940,7 @@ class Router: return _returned_deployments try: input_tokens = self._count_pre_call_check_tokens( - messages=messages, input=input, instructions=instructions + messages=messages, input=input, request_kwargs=request_kwargs ) except Exception as e: verbose_router_logger.error( diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index e3da70f50fe..afa27719064 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -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 diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index c2dc09bc65d..5d056ea3fe0 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -57,9 +57,11 @@ def get_azure_ad_token_provider( from azure import identity from azure.identity import ( CertificateCredential, + ChainedTokenCredential, ClientSecretCredential, DefaultAzureCredential, ManagedIdentityCredential, + WorkloadIdentityCredential, get_bearer_token_provider, ) @@ -101,6 +103,28 @@ def get_azure_ad_token_provider( # DefaultAzureCredential doesn't require explicit environment variables # It automatically discovers credentials from the environment (managed identity, CLI, etc.) credential = DefaultAzureCredential() + elif cred == AzureCredentialType.DeploymentIdentityCredential: + # DefaultAzureCredential cannot express this: excluding its developer credentials still + # leaves one managed identity link, which AZURE_CLIENT_ID pins to a user assigned identity, + # so a host running as a system assigned identity never gets asked + workload_client_id: Final = os.environ.get("AZURE_CLIENT_ID") + workload_tenant_id: Final = os.environ.get("AZURE_TENANT_ID") + workload_token_file: Final = os.environ.get("AZURE_FEDERATED_TOKEN_FILE") + credential = ChainedTokenCredential( + *( + ( + WorkloadIdentityCredential( + client_id=workload_client_id, + tenant_id=workload_tenant_id, + token_file_path=workload_token_file, + ), + ) + if workload_client_id and workload_tenant_id and workload_token_file + else () + ), + *((ManagedIdentityCredential(client_id=workload_client_id),) if workload_client_id else ()), + ManagedIdentityCredential(), + ) else: cred_cls: Final = getattr(identity, cred) credential = cred_cls() diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 0f17f2f23ab..9fb5bea81e3 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -2,6 +2,23 @@ from datetime import datetime from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, model_validator +from typing_extensions import ReadOnly, TypedDict + +from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains + + +class KeyTokenWhere(TypedDict): + token: ReadOnly[str] + + +class KeyAliasContainsWhere(TypedDict): + key_alias: ReadOnly[InsensitiveContains] + + +class KeySearchWhere(TypedDict): + """Prisma filter behind `/key/list?search=`: exact token or case-insensitive alias substring.""" + + OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]] class BulkUpdateKeyRequestItem(BaseModel): diff --git a/litellm/types/secret_managers/get_azure_ad_token_provider.py b/litellm/types/secret_managers/get_azure_ad_token_provider.py index 5d2f7409f95..6b4700d081d 100644 --- a/litellm/types/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/types/secret_managers/get_azure_ad_token_provider.py @@ -6,3 +6,4 @@ class AzureCredentialType(str, Enum): ManagedIdentityCredential = "ManagedIdentityCredential" CertificateCredential = "CertificateCredential" DefaultAzureCredential = "DefaultAzureCredential" + DeploymentIdentityCredential = "DeploymentIdentityCredential" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 34c47d3f201..c1d90694f14 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -222,7 +222,9 @@ class OffPeakPricing(TypedDict, total=False): weekday_timezone: ReadOnly[str] input_cost_per_token: ReadOnly[float] output_cost_per_token: ReadOnly[float] + output_cost_per_reasoning_token: ReadOnly[float] cache_read_input_token_cost: ReadOnly[float] + cache_creation_input_token_cost: ReadOnly[float] class ModelInfoBase(ProviderSpecificModelInfo, total=False): diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 24c0ff6b181..8763318b4eb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -96,7 +96,7 @@ "limit": 10 }, "DTZ007": { - "limit": 17 + "limit": 6 }, "DTZ011": { "limit": 3 diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts new file mode 100644 index 00000000000..5c7438cfd11 --- /dev/null +++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts @@ -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 => ({ 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 => + 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(request: APIRequestContext, path: string, data: Record): Promise { + 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 = (count: number, call: (index: number) => Promise): Promise => + Array.from({ length: count }, (_, i) => i).reduce>( + async (previous, i) => [...(await previous), await call(i)], + Promise.resolve([]), + ); + +async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise { + await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count); +} + +async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise { + await visibleTestId(page, "pagination-page-size").click(); + await page.getByRole("option", { name: size, exact: true }).click(); +} + +async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise { + 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 => + 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(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(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(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(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 } }), + ); + } + }); +}); diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index fa8f32764e8..64fce24bf2a 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -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@"); diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py index a0a26c089eb..fd1b05ff060 100644 --- a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py +++ b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py @@ -1,5 +1,6 @@ from datetime import datetime, timedelta -from unittest.mock import AsyncMock, patch +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI @@ -8,10 +9,12 @@ from litellm_enterprise.proxy.audit_logging_endpoints import router as audit_rou from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth # Create an app with just the audit router for testing app = FastAPI() app.include_router(audit_router) +app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin") client = TestClient(app) # Mock data for testing @@ -130,3 +133,45 @@ async def test_get_audit_log_by_id_not_found(mock_prisma_client): data = response.json() assert "message" in data["detail"] assert "not found" in data["detail"]["message"].lower() + + +def _list_audit_logs_where(mock_prisma_client: MagicMock, query: str) -> dict[str, object]: + mock_prisma_client.db.litellm_auditlog.find_many.return_value = [] + mock_prisma_client.db.litellm_auditlog.count.return_value = 0 + + response: Final = client.get(f"/audit?{query}") + + assert response.status_code == 200, response.text + find_many_where: Final = mock_prisma_client.db.litellm_auditlog.find_many.call_args.kwargs["where"] + assert mock_prisma_client.db.litellm_auditlog.count.call_args.kwargs["where"] == find_many_where + return find_many_where + + +def test_search_matches_any_id_column_alongside_the_other_filters(mock_prisma_client): + where: Final = _list_audit_logs_where(mock_prisma_client, "search=abc-123&action=create&object_team_id=team-1") + + assert where == { + "action": "create", + "AND": ( + { + "OR": [ + {"before_value": {"path": ["team_id"], "string_contains": "team-1"}}, + {"updated_values": {"path": ["team_id"], "string_contains": "team-1"}}, + ] + }, + { + "OR": ( + {"id": "abc-123"}, + {"changed_by": "abc-123"}, + {"object_id": "abc-123"}, + {"changed_by_api_key": "abc-123"}, + ) + }, + ), + } + + +def test_an_empty_search_leaves_the_where_clause_unchanged(mock_prisma_client): + where: Final = _list_audit_logs_where(mock_prisma_client, "action=create&search=") + + assert where == {"action": "create"} diff --git a/tests/guardrails_tests/test_tracing_guardrails.py b/tests/guardrails_tests/test_tracing_guardrails.py index bd8b7bad33f..ac85803ba39 100644 --- a/tests/guardrails_tests/test_tracing_guardrails.py +++ b/tests/guardrails_tests/test_tracing_guardrails.py @@ -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 diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 9ffb57924b6..3fab20a28ad 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -763,7 +763,7 @@ class _MigrateDeployHarness: "_resolve_specific_migration", staticmethod(self.resolved.append), ) - monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) self.baseline_succeeds = True diff --git a/tests/load_tests/test_granian_admission_saturation.py b/tests/load_tests/test_granian_admission_saturation.py new file mode 100644 index 00000000000..b42c06037c2 --- /dev/null +++ b/tests/load_tests/test_granian_admission_saturation.py @@ -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() diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 71be8730df1..e4724ff8705 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -779,7 +779,7 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): "error, opens_breaker", [ pytest.param("ConnectionError", True, id="connection_refused_is_unhealthy"), - pytest.param("TimeoutError", True, id="timeout_is_unhealthy"), + pytest.param("TimeoutError", False, id="timeout_burst_is_ambiguous"), pytest.param("BusyLoadingError", True, id="loading_is_unhealthy"), pytest.param("ResponseError", False, id="wrong_type_command_is_not"), pytest.param("DataError", False, id="bad_data_is_not"), @@ -791,6 +791,10 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker) They say nothing about connectivity, and a caller able to provoke them (an INCR against a non-numeric value, say) could otherwise trip the shared breaker on demand and drop rate limiting to per-process counters, which spreading traffic across replicas outruns. + + A rapid burst of timeouts is ambiguous too: the async timeout includes event-loop + scheduling delay, so a loop stall times out every queued call at once against a + healthy Redis. It must not open the breaker until the streak spans a minimum duration. """ import redis.exceptions @@ -810,3 +814,147 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker) await _run_under_circuit_breaker(breaker, "op", failing_call) assert breaker.is_open() is opens_breaker + + +@pytest.mark.asyncio +async def test_event_loop_stall_timeout_burst_keeps_breaker_closed(): + """One blocking stall of the worker event loop must not trip the breaker. + + Every operation already waiting on the loop times out together when the loop resumes, + so a purely consecutive threshold is satisfied instantly even though the Redis on the + other end (here an in-process fake that answers immediately) is healthy. + """ + import time as time_mod + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + async def healthy_redis_call_with_client_timeout(): + return await asyncio.wait_for(asyncio.sleep(0.001, result="ok"), timeout=0.05) + + async def stall_the_loop(): + await asyncio.sleep(0) + time_mod.sleep(0.2) + + results = await asyncio.gather( + *(_run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) for _ in range(8)), + stall_the_loop(), + return_exceptions=True, + ) + timeouts = [r for r in results if isinstance(r, asyncio.TimeoutError)] + assert len(timeouts) >= breaker.failure_threshold, "the stall must time out a full burst" + + assert breaker.is_open() is False, "a healthy Redis behind one loop stall must stay in the pool" + assert await _run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) == "ok" + + +@pytest.mark.asyncio +async def test_persistent_timeouts_still_open_the_breaker(): + """A real outage that surfaces only as timeouts must still open the breaker. + + Once the timeout-only streak spans the minimum duration with no success in between, + Redis is genuinely unusable from this worker and protection has to kick in. + """ + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.1) + + async def timing_out_call(): + raise RedisTimeoutError("read timed out") + + for _ in range(breaker.failure_threshold): + with pytest.raises(RedisTimeoutError): + await _run_under_circuit_breaker(breaker, "op", timing_out_call) + assert breaker.is_open() is False, "the burst has not spanned the minimum duration yet" + + await asyncio.sleep(0.12) + with pytest.raises(RedisTimeoutError): + await _run_under_circuit_breaker(breaker, "op", timing_out_call) + + assert breaker.is_open() is True + + +@pytest.mark.asyncio +async def test_stale_timeout_does_not_let_sub_threshold_hard_failures_open_the_breaker(): + """Hard connectivity failures below the threshold must not open the breaker just + because an old timeout already started the streak and the duration has elapsed. + + Each class has to earn the open on its own terms: hard failures by reaching the + threshold, timeouts by reaching the threshold and spanning the minimum duration. + """ + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) + + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + await asyncio.sleep(0.06) + for _ in range(breaker.failure_threshold - 1): + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + assert breaker.is_open() is False, "2 hard failures and 1 stale timeout are below both thresholds" + + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + assert breaker.is_open() is True, "the threshold-th hard failure must still open it" + + +@pytest.mark.asyncio +async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its_own_duration(): + """A stale timeout followed by hard failures must not pre-age the duration gate: + a later short timeout burst has to span timeout_min_duration on its own. + """ + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) + + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + await asyncio.sleep(0.06) + for _ in range(breaker.failure_threshold): + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + assert breaker.is_open() is False, "the burst is instantaneous, so the duration gate must hold it closed" + + await asyncio.sleep(0.06) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + assert breaker.is_open() is True, "the same run of timeouts persisting past the duration must open it" + + +@pytest.mark.asyncio +async def test_breaker_metrics_track_state_and_failure_class(): + """Breaker accounting must be observable: failure class, transitions, and state.""" + from prometheus_client import REGISTRY + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + + def sample(name, labels=None): + return REGISTRY.get_sample_value(name, labels) or 0.0 + + timeout_before = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) + hard_before = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) + opened_before = sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) + open_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "open"}) + closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) + + breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("t"))) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + + assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) == timeout_before + 1 + assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) == hard_before + 2 + assert sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) == opened_before + 1 + assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + 1 + assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + + breaker.record_success() + assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1 diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py index c16ceec8c31..9ebc3673ab6 100644 --- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py +++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py @@ -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") diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py index a96eae0f9c3..6e1dab4a71a 100644 --- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py +++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py @@ -42,6 +42,7 @@ def workload_identity_env_vars(monkeypatch): "AZURE_STORAGE_ENDPOINT_SUFFIX", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", + "AZURE_TOKEN_CREDENTIALS", "AZURE_SCOPE", ): monkeypatch.delenv(unset, raising=False) @@ -206,10 +207,28 @@ def test_default_chain_provider_is_storage_scoped_and_built_once_per_process(): assert first() == "chain-token" mock_builder.assert_called_once_with( azure_scope="https://storage.azure.com/.default", - azure_credential=AzureCredentialType.DefaultAzureCredential, + azure_credential=AzureCredentialType.DeploymentIdentityCredential, ) +def test_storage_chain_reaches_only_the_identities_a_deployment_carries(workload_identity_env_vars): + """ + The chain runs on a server, where a developer sign-in is a person and not the deployment, so + the storage token must come from workload identity or managed identity or from nothing + """ + _cached_credential_chain_token_provider.cache_clear() + with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "chain-token") as bearer: + _cached_credential_chain_token_provider() + _cached_credential_chain_token_provider.cache_clear() + + bearer.assert_called_once() + with bearer.call_args.args[0] as chain: + assert {type(link).__name__ for link in chain.credentials} == { + "WorkloadIdentityCredential", + "ManagedIdentityCredential", + } + + @pytest.mark.asyncio async def test_chain_tokens_are_read_from_the_provider_on_every_refresh( workload_identity_env_vars, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 0b6832d4bef..311ba7aebc0 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -29,11 +29,13 @@ from litellm.types.utils import ( from litellm.litellm_core_utils.llm_cost_calc.utils import ( CostCalculatorUtils, PromptTokensDetailsResult, + TokenRates, TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, _is_off_peak, _is_within_off_peak_window, + apply_off_peak_pricing, calculate_cache_writing_cost, generic_cost_per_token, get_token_type_cost_breakdown, @@ -782,6 +784,258 @@ def test_get_token_base_cost_off_peak_wins_over_tiered_pricing(): assert outside[:2] == (3e-6, 6e-6) +def _register_off_peak_reasoning_model( + model_name: str, off_peak_pricing: dict, reasoning_rate: float | None = 4e-6, **service_tier_rates: float +) -> None: + reasoning_entry = {} if reasoning_rate is None else {"output_cost_per_reasoning_token": reasoning_rate} + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 1.25e-6, + "off_peak_pricing": off_peak_pricing, + **reasoning_entry, + **service_tier_rates, + } + } + ) + + +def _off_peak_reasoning_usage() -> Usage: + return Usage( + prompt_tokens=100, + completion_tokens=80, + total_tokens=180, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=30, text_tokens=50), + ) + + +def test_generic_cost_per_token_off_peak_reasoning_rate(): + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-reasoning" + _register_off_peak_reasoning_model( + model_name, + {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6, "output_cost_per_reasoning_token": 5e-7}, + ) + + _, inside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7) + + _, outside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside == pytest.approx(50 * 2e-6 + 30 * 4e-6) + + +def test_generic_cost_per_token_off_peak_block_without_reasoning_rate(): + from datetime import datetime, timezone + + inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + block = {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6} + + _register_off_peak_reasoning_model("litellm-test-off-peak-model-reasoning-rate", block) + _, with_model_rate = generic_cost_per_token( + model="litellm-test-off-peak-model-reasoning-rate", + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=inside_window, + ) + assert with_model_rate == pytest.approx(50 * 1e-6 + 30 * 4e-6) + + _register_off_peak_reasoning_model("litellm-test-off-peak-no-reasoning-rate", block, reasoning_rate=None) + _, without_model_rate = generic_cost_per_token( + model="litellm-test-off-peak-no-reasoning-rate", + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=inside_window, + ) + assert without_model_rate == pytest.approx(80 * 1e-6) + + +def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier(): + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-tiered-reasoning" + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 128000], + "input_cost_per_token": 3e-6, + "output_cost_per_token": 6e-6, + "output_cost_per_reasoning_token": 8e-6, + }, + ], + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "output_cost_per_token": 1e-6, + "output_cost_per_reasoning_token": 5e-7, + }, + } + } + ) + + _, inside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7) + + _, outside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside == pytest.approx(50 * 6e-6 + 30 * 8e-6) + + +def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_tier(): + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-reasoning-service-tier" + _register_off_peak_reasoning_model( + model_name, + {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6, "output_cost_per_reasoning_token": 5e-7}, + output_cost_per_token_priority=3e-6, + output_cost_per_reasoning_token_priority=6e-6, + ) + + _, inside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + service_tier="priority", + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7) + + _, outside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + service_tier="priority", + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside == pytest.approx(50 * 3e-6 + 30 * 6e-6) + + +def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings(): + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-odd-values" + _register_off_peak_reasoning_model( + model_name, + { + "hours_utc": "16:30-00:30", + "cache_creation_input_token_cost": True, + "output_cost_per_reasoning_token": "5e-7", + }, + ) + standard = TokenRates( + input_rate=1e-6, output_rate=2e-6, cache_read_rate=1e-7, cache_creation_rate=1.25e-6, reasoning_rate=4e-6 + ) + + rates = apply_off_peak_pricing( + litellm.get_model_info(model_name, custom_llm_provider="openai"), + datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + standard, + ) + assert rates.cache_creation_rate == 1.25e-6 + assert rates.reasoning_rate == 5e-7 + + +def test_get_token_base_cost_off_peak_cache_creation_rate(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 1.25e-6, + "cache_creation_input_token_cost_above_1hr": 2e-6, + "off_peak_pricing": {"hours_utc": "16:30-00:30", "cache_creation_input_token_cost": 5e-7}, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + + inside = _get_token_base_cost(model_info, usage, current_time=inside_window) + assert inside[2] == 5e-7 + assert inside[3] == 2e-6 + + outside = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert outside[2] == 1.25e-6 + + without_key = cast( + ModelInfo, + {**model_info, "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}}, + ) + assert _get_token_base_cost(without_key, usage, current_time=inside_window)[2] == 1.25e-6 + + +def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates(): + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-breakdown" + _register_off_peak_reasoning_model( + model_name, + { + "hours_utc": "16:30-00:30", + "output_cost_per_reasoning_token": 5e-7, + "cache_creation_input_token_cost": 5e-7, + }, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=80, + total_tokens=1080, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=30, text_tokens=50), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, cache_creation_tokens=400, text_tokens=500), + ) + + inside = get_token_type_cost_breakdown( + model=model_name, + custom_llm_provider="openai", + usage=usage, + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside.reasoning_cost == pytest.approx(30 * 5e-7) + assert inside.cache_creation_cost == pytest.approx(400 * 5e-7) + assert inside.cache_read_cost == pytest.approx(100 * 1e-7) + + outside = get_token_type_cost_breakdown( + model=model_name, + custom_llm_provider="openai", + usage=usage, + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside.reasoning_cost == pytest.approx(30 * 4e-6) + assert outside.cache_creation_cost == pytest.approx(400 * 1.25e-6) + + def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index b6281834f24..a30d35d46f2 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -649,6 +649,90 @@ class TestDashscopeCostCalculator: assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + def test_dashscope_off_peak_reasoning_rate_replaces_the_dedicated_reasoning_rate(self): + self._register_off_peak_flat_model( + "dashscope/qwen-reasoning-rate-off-peak-test", + { + "hours_utc": self.OFF_PEAK_WINDOW, + "output_cost_per_token": 2.4e-06, + "output_cost_per_reasoning_token": 4.5e-06, + }, + ) + litellm.model_cost["dashscope/qwen-reasoning-rate-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06 + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + _, completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-rate-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + assert math.isclose(completion_cost, (150 * 2.4e-06) + (50 * 4.5e-06), rel_tol=1e-10) + + _, peak_completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-rate-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + assert math.isclose(peak_completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_cache_creation_rate_replaces_the_standard_rate(self): + self._register_off_peak_flat_model( + "dashscope/qwen-cache-creation-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "cache_creation_input_token_cost": 1.5e-06}, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100), + ) + + prompt_cost, _ = dashscope_cost_per_token( + model="qwen-cache-creation-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + assert math.isclose(prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 1.5e-06), rel_tol=1e-10) + + peak_prompt_cost, _ = dashscope_cost_per_token( + model="qwen-cache-creation-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_reasoning_and_cache_creation_rates_override_the_selected_tier(self): + self._register_tiered_model( + "dashscope/qwen-tiered-reasoning-off-peak-test", + [ + { + "range": [0, 1000], + "input_cost_per_token": 4e-07, + "cache_creation_input_token_cost": 3e-07, + "output_cost_per_token": 1.6e-06, + "output_cost_per_reasoning_token": 3.2e-06, + }, + ], + ) + litellm.model_cost["dashscope/qwen-tiered-reasoning-off-peak-test"]["off_peak_pricing"] = { + "hours_utc": self.OFF_PEAK_WINDOW, + "cache_creation_input_token_cost": 1e-07, + "output_cost_per_reasoning_token": 8e-07, + } + usage = Usage( + prompt_tokens=500, + completion_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_creation_tokens=200), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=40), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-tiered-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + assert math.isclose(prompt_cost, (300 * 4e-07) + (200 * 1e-07), rel_tol=1e-10) + assert math.isclose(completion_cost, (60 * 1.6e-06) + (40 * 8e-07), rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="qwen-tiered-reasoning-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + assert math.isclose(peak_prompt_cost, (300 * 4e-07) + (200 * 3e-07), rel_tol=1e-10) + assert math.isclose(peak_completion_cost, (60 * 1.6e-06) + (40 * 3.2e-06), rel_tol=1e-10) + def test_dashscope_off_peak_defaults_to_the_current_time(self): """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the default current time.""" diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index d8d9936e9a1..db107e00df0 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -85,6 +85,31 @@ class TestResolveConfig: def test_plaintext_http_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="http://api.openai.com/v1") is None + @pytest.mark.parametrize( + "api_base", + ( + "https://southcentralus.privatelink.api.openai.com/v1", + "https://eu.api.openai.com/v1", + "https://us.api.openai.com/v1", + ), + ) + def test_openai_backed_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig, api_base: str) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base=api_base) == wif_env + + @pytest.mark.parametrize( + "api_base", + ( + "https://api.openai.com.evil.example/v1", + "https://openai.com/v1", + "https://euapi.openai.com/v1", + "http://southcentralus.privatelink.api.openai.com/v1", + ), + ) + def test_lookalike_or_plaintext_api_base_disables( + self, wif_env: OpenAIWorkloadIdentityConfig, api_base: str + ) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base=api_base) is None + def test_foreign_env_base_url_disables( self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -158,6 +183,14 @@ class TestClientConstruction: assert client.api_key == "workload-identity-auth" assert client._workload_identity_auth is not None + def test_privatelink_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client( + is_async=False, api_key=None, api_base="https://southcentralus.privatelink.api.openai.com/v1" + ) + assert isinstance(client, OpenAI) + assert client.api_key == "workload-identity-auth" + assert client._workload_identity_auth is not None + def test_static_key_client_unaffected(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key="sk-static", api_base=None) assert isinstance(client, OpenAI) @@ -231,6 +264,16 @@ class TestResponsesValidateEnvironment: ) assert headers["Authorization"] == "Bearer None" + @respx.mock + def test_privatelink_api_base_mints_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + mock_token_exchange() + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, + model="gpt-4o-mini", + litellm_params=GenericLiteLLMParams(api_base="https://southcentralus.privatelink.api.openai.com/v1"), + ) + assert headers["Authorization"] == "Bearer exchanged-bearer-token" + def test_litellm_proxy_subclass_never_mints_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: headers: Final = LiteLLMProxyResponsesAPIConfig().validate_environment( headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index a182656e4a8..25a961c3413 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -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) diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py index fb21e2e6f6b..7970f7771fc 100644 --- a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py @@ -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( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index 7b82e004f37..5d6d47b8c38 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 0fd35e674b7..9a6815a61e5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8201,6 +8201,129 @@ class TestPreemptive401ModeAware: await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) +class TestSingleServerPreflightReachesIdJag: + """The connect-time preflight is what turns a credential failure into an HTTP status the client + can read. An oauth2_id_jag server has to reach it: its subject comes from the assertion stored at + SSO login, so the failure is decided before any IdP call and there is nothing later in the session + that can report it (tools/list degrades to an empty list, tools/call to 'tool not found').""" + + def _id_jag_server(self) -> MCPServer: + return MCPServer( + server_id="id-idjag", + name="idjag", + alias="idjag", + server_name="idjag", + url="https://idjag.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.test/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.test/oauth2/token", + mcp_info={"server_name": "idjag"}, + ) + + async def _run(self, server: MCPServer, mcp_servers: list[str], preflight: AsyncMock) -> None: + from litellm.proxy._experimental.mcp_server import server as server_module + + with ( + patch.object( # test-quality-ok: route wiring must use the manager's configured server + server_module.global_mcp_server_manager, + "get_mcp_server_by_name", + return_value=server, + ), + patch.object( # test-quality-ok: route wiring must invoke the manager preflight + server_module.global_mcp_server_manager, + "preflight_token_exchange", + preflight, + ), + patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer + server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[server]) + ), + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": "/mcp/idjag", "headers": []}, + mcp_servers=mcp_servers, + oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"}, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + client_ip=None, + ) + + @pytest.mark.asyncio + async def test_id_jag_single_server_route_surfaces_the_preflight_status(self): + """The 412 the preflight raises must propagate out of connect, not be swallowed.""" + server = self._id_jag_server() + preflight = AsyncMock(side_effect=HTTPException(status_code=412, detail="no stored assertion")) + + with pytest.raises(HTTPException) as exc: + await self._run(server, ["idjag"], preflight) + + assert exc.value.status_code == 412 + assert preflight.await_args.kwargs["server"] is server + + @pytest.mark.asyncio + async def test_token_exchange_without_a_bearer_still_challenges_and_never_pre_flights(self): + """The already-shipped OBO path must be untouched by the call site dropping its mode test. + A token_exchange server with no inbound bearer has nothing to exchange, so it still gets the + RFC 9728 discovery challenge from the block above and the preflight is never reached; pushing + a subject-less exchange through the resolver would turn that challenge into some other status + and strand a client that only had to SSO and retry.""" + from litellm.proxy._experimental.mcp_server import server as server_module + + token_exchange = MCPServer( + server_id="id-obo", + name="obo", + alias="obo", + server_name="obo", + url="https://obo.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.test/oauth2/token", + client_id="cid", + client_secret="csec", + mcp_info={"server_name": "obo"}, + ) + preflight = AsyncMock() + + with ( + patch.object( # test-quality-ok: route wiring must use the manager's configured server + server_module.global_mcp_server_manager, + "get_mcp_server_by_name", + return_value=token_exchange, + ), + patch.object( # test-quality-ok: route wiring must invoke the manager preflight + server_module.global_mcp_server_manager, + "preflight_token_exchange", + preflight, + ), + pytest.raises(HTTPException) as exc, + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": "/mcp/obo", "headers": []}, + mcp_servers=["obo"], + oauth2_headers=None, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + client_ip=None, + ) + + assert exc.value.status_code == 401 + headers = exc.value.headers or {} + assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "") + preflight.assert_not_awaited() + + @pytest.mark.asyncio + async def test_id_jag_multi_server_route_still_absorbs_the_failure(self): + """The aggregate contract is unchanged: with more than one target the preflight does not run, + so one server with no stored assertion cannot fail the whole connect.""" + preflight = AsyncMock(side_effect=HTTPException(status_code=412, detail="no stored assertion")) + + await self._run(self._id_jag_server(), ["idjag", "other"], preflight) + + preflight.assert_not_awaited() + + def _make_obo_server(alias: str) -> MCPServer: return MCPServer( server_id=f"id-{alias}", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index cc41315c001..b84c132cdbe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2706,6 +2706,165 @@ class TestMCPServerManager: await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None) assert resolved == ["good-subject"] + def _id_jag_server(self, server_id: str) -> "MCPServer": + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + @pytest.mark.asyncio + async def test_preflight_id_jag_surfaces_missing_assertion_as_a_plain_412(self): + """ID-JAG's missing/expired-assertion precondition must reach the client as a 412 whose body + names the fix, at the transport edge. Without the preflight the session opens and the caller + gets a 200 with an empty tool list and then 'tool not found', which is not what happened. + 412 is a precondition, not an RFC 9728 discovery challenge, so it carries no + WWW-Authenticate: there is nothing for the client to discover and retry against.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + summary = ( + "ID-JAG requires an IdP identity assertion for this user and none is stored. " + "Sign in through LiteLLM SSO so the gateway captures one." + ) + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_precondition_required(summary)) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-412"), + oauth2_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + ) + assert exc_info.value.status_code == 412 + assert summary in exc_info.value.detail + assert not (exc_info.value.headers or {}) + + @pytest.mark.asyncio + async def test_preflight_id_jag_surfaces_assertion_store_outage_as_503(self): + """A store outage is the other failure the session would swallow, and it is a different + answer than 412: the user has nothing to fix by signing in again.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_upstream_unavailable("assertion store unreachable")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-503"), + oauth2_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + ) + assert exc_info.value.status_code == 503 + + @pytest.mark.asyncio + async def test_preflight_id_jag_preflights_litellm_key_and_skips_identity_bearer(self): + """ID-JAG preflights when Authorization carries a LiteLLM key, but skips a caller identity + bearer that the session passes through unchanged.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + subjects = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + subjects.append( + ( + subject.subject_id, + subject.inbound_token.get_secret_value() if subject.inbound_token else None, + ) + ) + return Ok(StaticHeaderAuth("Bearer minted-id-jag", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + caller = UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1") + + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-key"), + oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"}, + raw_headers={"authorization": "Bearer sk-litellm-virtual-key"}, + user_api_key_auth=caller, + ) + assert subjects == [("u-1", None)] + + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-identity"), + oauth2_headers={"Authorization": "Bearer caller-idp-id-token"}, + raw_headers={ + "x-litellm-api-key": "Bearer sk-admission-key", + "authorization": "Bearer caller-idp-id-token", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="u-1"), + ) + assert subjects == [("u-1", None)] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "server_fields", + [ + {"auth_type": MCPAuth.none}, + {"auth_type": MCPAuth.api_key, "authentication_token": "static-upstream-key"}, + {"auth_type": MCPAuth.bearer_token, "authentication_token": "static-upstream-key"}, + { + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "client_credentials", + "client_id": "cid", + "client_secret": "csec", + "token_url": "https://idp.example.com/token", + }, + {"auth_type": MCPAuth.true_passthrough}, + ], + ) + async def test_preflight_resolves_nothing_for_a_mode_that_does_not_pre_flight(self, server_fields): + """The manager is the only thing deciding which modes pre-flight, so it has to reject every + other mode itself. The single-server call site no longer tests the mode before calling, so a + mode that falls through here would start resolving its credential a second time, at connect, + for flows that never had a connect-time resolution at all.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + calls = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + calls.append(server.server_id) + return Error(CredError.of_misconfigured("the preflight must never get here")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = MCPServer( + server_id="not-pre-flighted", + name="not-pre-flighted-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + **server_fields, + ) + + assert ( + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + ) + is None + ) + assert calls == [] + @pytest.mark.asyncio @pytest.mark.parametrize( "authorization", @@ -2728,7 +2887,9 @@ class TestMCPServerManager: resolved: Final[list[str | None]] = [] class _FakeProvider: - async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]: + async def resolve_credentials( + self, subject: Subject, server: ServerSpec + ) -> Ok[StaticHeaderAuth, CredError]: resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) @@ -2754,7 +2915,9 @@ class TestMCPServerManager: resolved: Final[list[str | None]] = [] class _FakeProvider: - async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]: + async def resolve_credentials( + self, subject: Subject, server: ServerSpec + ) -> Ok[StaticHeaderAuth, CredError]: resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index c83ba142011..f809fadc879 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -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=) 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=) + 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"] diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 62073f4bf51..d58683fd1e5 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -754,6 +754,84 @@ def test_expand_wildcard_invalid_litellm_params_passthrough(): assert result == [deployment] +def test_get_complete_model_list_excludes_wildcard_routes_by_default(): + """Regression (LIT-4108): a wildcard with a matching router deployment leaked into /v1/models.""" + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + router = Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + }, + ] + ) + + result = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=["bedrock/*", "gpt-4"], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=False, + llm_router=router, + ) + + assert "bedrock/*" not in result + assert "gpt-4" in result + assert any(m.startswith("bedrock/") for m in result) + + +def test_get_complete_model_list_excludes_wildcard_routes_without_router(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=["bedrock/*", "gpt-4"], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=False, + llm_router=None, + ) + + assert "bedrock/*" not in result + assert "gpt-4" in result + assert any(m.startswith("bedrock/") for m in result) + + +def test_get_complete_model_list_includes_wildcard_routes_when_requested(): + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + router = Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + }, + ] + ) + + result = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=["bedrock/*"], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=True, + llm_router=router, + ) + + assert result.count("bedrock/*") == 1 + assert any(m.startswith("bedrock/") and m != "bedrock/*" for m in result) + + def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): """models_by_provider was a frozen import-time snapshot of set unions, so cost map reloads (which call add_known_models) never reached wildcard expansion until a diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e3f71692c78..0e90c107865 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -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. diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py index 4e063dd0c5b..86f9aeafc08 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py @@ -225,3 +225,71 @@ def test_compute_overall_action_all_passed(): def test_compute_overall_action_empty(): assert _compute_overall_action([]) == "passed" + + +class TestEnrichPolicyTemplateStreamKeepalive: + async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]: + import asyncio + from unittest.mock import MagicMock + + import litellm + import litellm.proxy.management_endpoints.policy_endpoints.endpoints as policy_endpoints + import litellm.proxy.proxy_server as proxy_server + from fastapi.responses import StreamingResponse + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.policy_endpoints.endpoints import ( + EnrichTemplateRequest, + enrich_policy_template_stream, + ) + + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + async def _name_chunks(): + await asyncio.sleep(delay) + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Rival Air\n" + yield chunk + + class SlowRouter: + async def acompletion(self, **kwargs): + return _name_chunks() + + async def _no_variations(competitors, model): + return {} + + monkeypatch.setattr(proxy_server, "llm_router", SlowRouter()) + monkeypatch.setattr(policy_endpoints, "_generate_competitor_variations", _no_variations) + + response = await enrich_policy_template_stream( + data=EnrichTemplateRequest( + template_id="competitor-mention-detection", + parameters={"brand_name": "Acme"}, + model="gpt-5.4-mini", + ), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert isinstance(response, StreamingResponse) + chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + return chunks, dict(response.headers) + + @pytest.mark.asyncio + async def test_endpoint_pings_while_competitor_discovery_is_still_running(self, monkeypatch): + chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05) + + assert headers["content-type"].startswith("text/event-stream") + assert headers["cache-control"] == "no-cache" + assert headers["x-accel-buffering"] == "no" + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b'data: {"type": "competitor", "name": "Rival Air"}\n\n' in chunks + assert chunks[-1].startswith(b'data: {"type": "done"') + + @pytest.mark.asyncio + async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch): + chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15) + + assert b": ping\n\n" not in chunks + assert chunks[0] == b'data: {"type": "competitor", "name": "Rival Air"}\n\n' + assert chunks[-1].startswith(b'data: {"type": "done"') diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 1697b77b99a..60f9a1a55e2 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -19,11 +19,13 @@ from litellm.proxy._types import ( NewUserResponse, ProxyErrorTypes, ProxyException, + UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( SCIMRosterSyncError, UserProvisionerHelpers, _apply_group_patch_updates, + _create_user_if_not_exists, _extract_group_member_ids, _extract_ids_from_path_filter, _handle_group_membership_changes, @@ -37,8 +39,8 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( delete_group, delete_user, get_groups, - get_users, get_service_provider_config, + get_users, merge_placeholder, patch_group, patch_team_membership, @@ -304,6 +306,85 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN +def _mock_scim_create_user_deps(mocker: MockerFixture, scim_user: SCIMUser) -> AsyncMock: + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + return mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id=scim_user.userName)), + ) + + +@pytest.mark.asyncio +async def test_create_user_without_groups_defers_to_default_team(mocker: MockerFixture, monkeypatch): + """IdPs omit groups on POST /Users; teams must stay unset so new_user applies default_internal_user_params.teams""" + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-user", + emails=[SCIMUserEmail(value="new@example.com")], + ) + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = _mock_scim_create_user_deps(mocker, scim_user) + + await create_user(user=scim_user) + + assert new_user_mock.call_args.kwargs["data"].teams is None + assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch): + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-user", + emails=[SCIMUserEmail(value="new@example.com")], + groups=[SCIMUserGroup(value="idp-team")], + ) + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = _mock_scim_create_user_deps(mocker, scim_user) + + await create_user(user=scim_user) + + assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"] + + +@pytest.mark.asyncio +async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch): + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = mocker.patch( # test-quality-ok: new_user is imported inside the helper, not injectable + "litellm.proxy.management_endpoints.internal_user_endpoints.new_user", + AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")), + ) + + created = await _create_user_if_not_exists(user_id="group-user") + + assert created is not None + assert new_user_mock.call_args.kwargs["data"].teams is None + assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + + @pytest.mark.asyncio async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch): """ @@ -1176,6 +1257,67 @@ async def test_update_user_success(mocker): assert call_args[1]["data"]["teams"] == ["new-team"] +@pytest.mark.asyncio +@pytest.mark.parametrize("groups", [None, []], ids=["groups-omitted", "groups-empty"]) +async def test_update_user_without_groups_preserves_memberships_and_role(mocker, monkeypatch, groups): + """Okta profile PUTs carry no `groups` or `groups: []`; neither may drop teams (and their keys) or recompute role""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins", "engineering"] + existing_user.metadata = {} + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="okta-user", + name=SCIMUserName(familyName="Renamed", givenName="Okta"), + emails=[SCIMUserEmail(value="okta@example.com")], + **({} if groups is None else {"groups": groups}), + ) + response_scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + id="okta-user", + userName="okta-user", + emails=[SCIMUserEmail(value="okta@example.com")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "okta-user"}) + + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=response_scim_user), + ) + patch_membership = mocker.patch( # test-quality-ok: roster writes are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + + result = await update_user(user_id="okta-user", user=scim_user) + + assert result == response_scim_user + patch_membership.assert_not_awaited() + update_data = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"] + assert update_data["teams"] == ["litellm-admins", "engineering"] + assert "user_role" not in update_data + + @pytest.mark.asyncio async def test_update_user_not_found(mocker): """Should raise 404 when user doesn't exist""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index e8f768c14ef..81816e21c10 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -12,13 +12,12 @@ from fastapi.testclient import TestClient from prisma.errors import PrismaError import litellm.proxy.proxy_server as ps -from litellm.proxy.proxy_server import app from litellm.proxy._types import ( CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth, ) - +from litellm.proxy.proxy_server import app def _make_access_group_record( @@ -58,6 +57,10 @@ def _make_access_group_record( return record +def _make_team_record(team_id: str, access_group_ids: list[str] | None = None): + return types.SimpleNamespace(team_id=team_id, access_group_ids=access_group_ids or []) + + @pytest.fixture def client_and_mocks(monkeypatch): """Setup mock prisma and admin auth for access group endpoints.""" @@ -185,7 +188,8 @@ ACCESS_GROUP_PATHS = ["/v1/access_group", "/v1/unified_access_group"] ) def test_create_access_group_success(client_and_mocks, base_path, payload): """Create access group with various payloads returns 201.""" - client, _, mock_table, *_ = client_and_mocks + client, mock_prisma, mock_table, *_ = client_and_mocks + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[_make_team_record("team-1")]) resp = client.post(base_path, json=payload) assert resp.status_code == 201 @@ -277,13 +281,45 @@ def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) def test_list_access_groups_success_empty(client_and_mocks, base_path): - """List access groups returns empty list when none exist.""" - client, _, mock_table, *_ = client_and_mocks + """List access groups returns empty list when none exist, without querying teams.""" + client, mock_prisma, mock_table, *_ = client_and_mocks resp = client.get(base_path) assert resp.status_code == 200 assert resp.json() == [] mock_table.find_many.assert_awaited_once() + mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited() + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_list_access_groups_attributes_teams_per_group_with_one_query(client_and_mocks, base_path): + """List derives each group's teams from the team table in a single query, attributed per group.""" + client, mock_prisma, mock_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + records = [ + _make_access_group_record(access_group_id="ag-1", access_group_name="group-1"), + _make_access_group_record(access_group_id="ag-2", access_group_name="group-2"), + ] + mock_table.find_many = AsyncMock(return_value=records) + mock_team_table.find_many = AsyncMock( + return_value=[ + _make_team_record("team-x", ["ag-1"]), + _make_team_record("team-y", ["ag-2"]), + _make_team_record("team-z", ["ag-1", "ag-2"]), + ] + ) + + resp = client.get(base_path) + assert resp.status_code == 200 + body = resp.json() + assert body[0]["assigned_team_ids"] == ["team-x", "team-z"] + assert body[1]["assigned_team_ids"] == ["team-y", "team-z"] + + mock_team_table.find_many.assert_awaited_once() + carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"] + assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-1", "ag-2"] + assert list(listed["team_id"]["in"]) == [] @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) @@ -373,6 +409,43 @@ def test_get_access_group_success(client_and_mocks, base_path, access_group_id): assert resp.json()["access_group_id"] == access_group_id +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_get_access_group_derives_assigned_teams_from_team_table(client_and_mocks, base_path): + """Get drops ghost ids from the stored column and adds teams that carry the group but were never mirrored.""" + client, mock_prisma, mock_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + record = _make_access_group_record(access_group_id="ag-123", assigned_team_ids=["team-a", "ghost-team"]) + mock_table.find_unique = AsyncMock(return_value=record) + mock_team_table.find_many = AsyncMock( + return_value=[ + _make_team_record("team-a", ["ag-123"]), + _make_team_record("team-b", ["ag-123"]), + _make_team_record("team-c", ["ag-123"]), + ] + ) + + resp = client.get(f"{base_path}/ag-123") + assert resp.status_code == 200 + assert resp.json()["assigned_team_ids"] == ["team-a", "team-b", "team-c"] + + carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"] + assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-123"] + assert list(listed["team_id"]["in"]) == ["team-a", "ghost-team"] + + +def test_get_access_group_empty_column_and_no_teams_returns_empty(client_and_mocks): + """Get returns [] when the column is empty and no team carries the group.""" + client, mock_prisma, mock_table, *_ = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-123")) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/access_group/ag-123") + assert resp.status_code == 200 + assert resp.json()["assigned_team_ids"] == [] + + def test_get_access_group_not_found(client_and_mocks): """Get access group returns 404 when not found.""" client, _, mock_table, *_ = client_and_mocks @@ -985,6 +1058,28 @@ def test_record_to_access_group_table(): assert result.access_agent_ids == ["agent-1"] +def test_attached_team_ids_by_group_keeps_column_order_then_appends_unmirrored_teams(): + """Stored ids that resolve keep their order, ghosts drop, carriers the mirror missed append once, per group.""" + from litellm.proxy.management_endpoints.access_group_endpoints import ( + _attached_team_ids_by_group, + ) + + records = [ + _make_access_group_record(access_group_id="ag-1", assigned_team_ids=["team-b", "ghost", "team-a"]), + _make_access_group_record(access_group_id="ag-2", assigned_team_ids=[]), + ] + teams = [ + _make_team_record("team-a", ["ag-1"]), + _make_team_record("team-b", []), + _make_team_record("team-c", ["ag-1"]), + _make_team_record("team-d", ["ag-2"]), + ] + + result = _attached_team_ids_by_group(records, teams) + + assert dict(result) == {"ag-1": ("team-b", "team-a", "team-c"), "ag-2": ("team-d",)} + + # --------------------------------------------------------------------------- # Sync tests: CREATE # --------------------------------------------------------------------------- @@ -997,9 +1092,8 @@ def test_create_access_group_syncs_assigned_teams(client_and_mocks): ) mock_team_table = mock_prisma.db.litellm_teamtable - team_record = MagicMock() - team_record.team_id = "team-1" - team_record.access_group_ids = [] + team_record = _make_team_record("team-1") + mock_team_table.find_many = AsyncMock(return_value=[team_record]) mock_team_table.find_unique = AsyncMock(return_value=team_record) resp = client.post( @@ -1043,20 +1137,22 @@ def test_create_access_group_syncs_assigned_keys(client_and_mocks): assert "ag-new" in call_kwargs["data"]["access_group_ids"] -def test_create_access_group_skips_sync_for_nonexistent_team(client_and_mocks): - """Create skips updating a team that doesn't exist in DB.""" - client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks +def test_create_access_group_rejects_nonexistent_team(client_and_mocks): + """Create refuses to store a team id that does not resolve to a team row.""" + client, mock_prisma, mock_access_group_table, *_ = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - mock_team_table.find_unique = AsyncMock(return_value=None) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-real")]) resp = client.post( "/v1/access_group", json={ "access_group_name": "new-group", - "assigned_team_ids": ["nonexistent-team"], + "assigned_team_ids": ["team-real", "nonexistent-team", "also-missing"], }, ) - assert resp.status_code == 201 + assert resp.status_code == 400 + assert resp.json()["detail"] == "Unknown team ids: also-missing, nonexistent-team" + mock_access_group_table.create.assert_not_awaited() mock_team_table.update.assert_not_awaited() @@ -1065,9 +1161,8 @@ def test_create_access_group_idempotent_team_sync(client_and_mocks): client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - team_record = MagicMock() - team_record.team_id = "team-1" - team_record.access_group_ids = ["ag-new"] # already synced + team_record = _make_team_record("team-1", ["ag-new"]) + mock_team_table.find_many = AsyncMock(return_value=[team_record]) mock_team_table.find_unique = AsyncMock(return_value=team_record) resp = client.post( @@ -1095,9 +1190,8 @@ def test_update_access_group_syncs_added_teams(client_and_mocks): ) mock_access_group_table.find_unique = AsyncMock(return_value=existing) - team_record = MagicMock() - team_record.team_id = "team-new" - team_record.access_group_ids = [] + team_record = _make_team_record("team-new") + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-existing", ["ag-update"]), team_record]) mock_team_table.find_unique = AsyncMock(return_value=team_record) resp = client.put( @@ -1113,6 +1207,25 @@ def test_update_access_group_syncs_added_teams(client_and_mocks): assert "ag-update" in call_kwargs["data"]["access_group_ids"] +def test_update_access_group_rejects_nonexistent_team(client_and_mocks): + """Update refuses to store a team id that does not resolve to a team row and leaves the group untouched.""" + client, mock_prisma, mock_access_group_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-existing"]) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-existing", ["ag-update"])]) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": ["team-existing", "team-ghost"]}, + ) + assert resp.status_code == 400 + assert resp.json()["detail"] == "Unknown team ids: team-ghost" + mock_access_group_table.update.assert_not_awaited() + mock_team_table.update.assert_not_awaited() + + def test_update_access_group_syncs_removed_teams(client_and_mocks): """Update removes access_group_id from de-assigned teams.""" client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( @@ -1125,9 +1238,8 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): ) mock_access_group_table.find_unique = AsyncMock(return_value=existing) - team_to_remove = MagicMock() - team_to_remove.team_id = "team-remove" - team_to_remove.access_group_ids = ["ag-update"] + team_to_remove = _make_team_record("team-remove", ["ag-update"]) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), team_to_remove]) mock_team_table.find_unique = AsyncMock(return_value=team_to_remove) resp = client.put( @@ -1145,6 +1257,28 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): assert "ag-update" not in call_kwargs["data"]["access_group_ids"] +def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks): + """Update removes the group from a team that carries it but was never written to the stored column.""" + client, mock_prisma, mock_access_group_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep"]) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + unmirrored = _make_team_record("team-unmirrored", ["ag-update", "ag-other"]) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), unmirrored]) + mock_team_table.find_unique = AsyncMock(return_value=unmirrored) + + resp = client.put("/v1/access_group/ag-update", json={"assigned_team_ids": ["team-keep"]}) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-unmirrored"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-unmirrored"} + assert call_kwargs["data"]["access_group_ids"] == ["ag-other"] + + def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks): """Update does not sync teams when assigned_team_ids is absent from the payload.""" client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 7954a4693cc..0e4af9f75a5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6347,6 +6347,93 @@ def test_build_key_filter_conditions_key_hash_narrows_team_admin_visibility(): assert {"token": "hashed-token-123"} in where["AND"], f"key_hash not ANDed: {where}" +def _search_clause(search: str, token: str) -> dict: + return {"OR": [{"token": token}, {"key_alias": {"contains": search, "mode": "insensitive"}}]} + + +def test_build_key_filter_conditions_search_ors_token_and_alias_contains(): + """ + LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by + its ID (the token column), with the pasted value used verbatim. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + hashed_where = json.loads( + json.dumps( + _build_key_filter_conditions( + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + search="already-hashed-token", + ) + ) + ) + assert _search_clause("already-hashed-token", "already-hashed-token") in hashed_where["AND"], ( + f"hashed search not used verbatim: {hashed_where}" + ) + + +def test_build_key_filter_conditions_search_narrows_team_admin_visibility(): + """ + LIT-4741, same class as LIT-3243: `search` must be a top-level AND so it + narrows a team admin's admin-team branch instead of being bypassed by it. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + where = json.loads( + json.dumps( + _build_key_filter_conditions( + user_id="team-admin-user", + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=["team-a"], + member_team_ids=["team-a"], + include_created_by_keys=False, + search="member-key-id", + ) + ) + ) + + assert where.get("AND"), f"expected top-level AND, got: {where}" + assert _search_clause("member-key-id", "member-key-id") in where["AND"], f"search not ANDed: {where}" + assert json.dumps({"team_id": {"in": ["team-a"]}}) in json.dumps(where) + + +@pytest.mark.asyncio +async def test_list_key_helper_applies_search_to_prisma_where(): + """LIT-4741: `search` given to _list_key_helper must reach the Prisma where clause.""" + mock_prisma_client = AsyncMock() + mock_find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + + await _list_key_helper( + prisma_client=mock_prisma_client, + page=1, + size=50, + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + search="key-id-123", + ) + + where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"])) + assert _search_clause("key-id-123", "key-id-123") in where["AND"], f"search not in Prisma where: {where}" + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ @@ -14870,6 +14957,16 @@ async def test_list_keys_non_admin_cannot_opt_into_substring(): assert kwargs["user_id"] == "alice" +@pytest.mark.asyncio +async def test_list_keys_search_is_honored_for_non_admin(): + """LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's + search reaches the helper while their own-user scoping stays in place.""" + user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice") + kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="key-id-123") + assert kwargs["search"] == "key-id-123" + assert kwargs["user_id"] == "alice" + + @pytest.mark.asyncio async def test_cli_session_token_delegation_ceiling_blocked_by_team_budget(): team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0) diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index 3a32b3cc128..e5616a1975d 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -466,3 +466,61 @@ class TestUsageAiChatServiceAccountGuard: is_admin=False, ) assert "Endpoint-level guard missing" in str(exc_info.value) + + +class TestUsageAiChatKeepalive: + async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]: + import asyncio + + import litellm + from fastapi.responses import StreamingResponse + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( + ChatMessage, + UsageAIChatRequest, + usage_ai_chat, + ) + + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + async def slow_acompletion(**kwargs): + await asyncio.sleep(delay) + response = MagicMock() + response.choices = [MagicMock()] + response.choices[0].message.tool_calls = None + response.choices[0].message.content = "Total spend is $50.25" + return response + + with patch( # test-quality-ok: the stream calls the module-level litellm.acompletion directly; no injection seam + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm.acompletion", + new=AsyncMock(side_effect=slow_acompletion), + ): + response = await usage_ai_chat( + data=UsageAIChatRequest(messages=[ChatMessage(role="user", content="hi")], model="gpt-4o-mini"), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert isinstance(response, StreamingResponse) + chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + return chunks, dict(response.headers) + + @pytest.mark.asyncio + async def test_endpoint_pings_while_the_planning_completion_is_still_running(self, monkeypatch): + chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05) + + assert headers["content-type"].startswith("text/event-stream") + assert headers["cache-control"] == "no-cache" + assert headers["x-accel-buffering"] == "no" + assert chunks[0].startswith(b'data: {"type": "status"') + assert chunks[1] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b'"content": "Total spend is $50.25"' in b"".join(chunks) + assert chunks[-1] == b'data: {"type": "done"}\n\n' + + @pytest.mark.asyncio + async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch): + chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15) + + assert b": ping\n\n" not in chunks + assert chunks[0].startswith(b'data: {"type": "status"') + assert chunks[-1] == b'data: {"type": "done"}\n\n' diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index be75d980d9d..dff0e80fa77 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -615,6 +615,96 @@ class TestMemoryEndpoints: assert keys == {"user:profile"} assert body["total"] == 1 + def test_list_memory_search_matches_key_prefix_or_memory_id_within_scope(self): + """ + `search` matches a key prefix OR an exact memory_id, and stays ANDed + with the visibility filter so a pasted foreign id cannot leak a row. + """ + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None), + _make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None), + _make_row(memory_id="mem-foreign", key="user:secret", user_id="user-b", team_id=None), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + by_id = client.get("/v1/memory?search=mem-target") + by_prefix = client.get("/v1/memory?search=user:") + foreign_id = client.get("/v1/memory?search=mem-foreign") + + assert by_id.status_code == 200, by_id.text + assert [m["memory_id"] for m in by_id.json()["memories"]] == ["mem-target"] + assert by_id.json()["total"] == 1 + + assert by_prefix.status_code == 200, by_prefix.text + assert {m["key"] for m in by_prefix.json()["memories"]} == {"user:profile"} + assert by_prefix.json()["total"] == 1 + + assert foreign_id.status_code == 200, foreign_id.text + assert foreign_id.json()["memories"] == [] + assert foreign_id.json()["total"] == 0 + + def test_list_memory_search_by_memory_id_for_admin_sees_any_scope(self): + """Admins have no visibility filter, so an id search returns the row whoever owns it.""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="mem-a", key="a", user_id="user-a", team_id=None), + _make_row(memory_id="mem-b", key="b", user_id="user-b", team_id=None), + ] + ) + client = _make_client(_admin_auth()) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?search=mem-b") + assert resp.status_code == 200, resp.text + assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-b"] + assert resp.json()["total"] == 1 + + def test_list_memory_search_wins_over_key_prefix(self): + """When both are sent, `search` decides the match and `key_prefix` is ignored.""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None), + _make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?search=mem-target&key_prefix=user:") + assert resp.status_code == 200, resp.text + assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-target"] + assert resp.json()["total"] == 1 + + def test_list_memory_key_prefix_never_matches_memory_id(self): + """`key_prefix` stays a pure key-prefix match; only `search` consults memory_id.""" + table = self.prisma.db.litellm_memorytable + table.rows.append(_make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None)) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?key_prefix=mem-target") + assert resp.status_code == 200, resp.text + assert resp.json()["memories"] == [] + assert resp.json()["total"] == 0 + + def test_list_memory_key_exact_filter(self): + """`key` is an exact match, never a prefix.""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="m1", key="user:profile", user_id="user-a", team_id=None), + _make_row(memory_id="m2", key="user:profile:archived", user_id="user-a", team_id=None), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?key=user:profile") + assert resp.status_code == 200, resp.text + assert [m["memory_id"] for m in resp.json()["memories"]] == ["m1"] + assert resp.json()["total"] == 1 + def test_list_memory_admin_sees_all(self): table = self.prisma.db.litellm_memorytable table.rows.extend( diff --git a/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py b/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py new file mode 100644 index 00000000000..f1ca13daa03 --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 1d4b0264879..acb45038df0 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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//locations//" 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``. @@ -5009,3 +5116,93 @@ class TestAzureRouterModelStreamingDispatch: assert result.status_code == 200 body = b"".join([chunk async for chunk in result.body_iterator]) assert body == upstream_body + + +class TestAzureRouterModelStreamingKeepalive: + async def _dispatch(self, monkeypatch, interval, headers_delay=0.0, body_delay=0.0) -> StreamingResponse: + import asyncio + + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + class _StallingBody(httpx.AsyncByteStream): + async def __aiter__(self): + await asyncio.sleep(body_delay) + yield b"data: hello\n\n" + + async def _upstream_response() -> httpx.Response: + await asyncio.sleep(headers_delay) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream", "x-upstream": "kept"}, + stream=_StallingBody(), + request=httpx.Request("POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/x"), + ) + + logging_obj = MagicMock() + logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() + + class StreamingRouter: + async def allm_passthrough_route(self, **kwargs): + return await AsyncPassthroughStreamingResponse( + response=_upstream_response(), + litellm_logging_obj=logging_obj, + provider_config=MagicMock(), + ) + + async def fake_get_request_body(_request): + return {"model": "gpt-5", "stream": True} + + monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + assert isinstance(result, StreamingResponse) + return result + + @pytest.mark.asyncio + async def test_pings_while_upstream_headers_are_still_pending(self, monkeypatch): + result = await self._dispatch(monkeypatch, interval=0.05, headers_delay=0.3) + + chunks = [chunk async for chunk in result.body_iterator] + + assert result.status_code == 200 + assert result.headers["x-accel-buffering"] == "no" + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b"".join(chunks).endswith(b"data: hello\n\n") + + @pytest.mark.asyncio + async def test_pings_while_upstream_body_is_still_pending(self, monkeypatch): + result = await self._dispatch(monkeypatch, interval=0.05, body_delay=0.3) + + chunks = [chunk async for chunk in result.body_iterator] + + assert result.status_code == 200 + assert result.headers["x-upstream"] == "kept" + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert chunks[-1] == b"data: hello\n\n" + + @pytest.mark.asyncio + async def test_relays_upstream_bytes_untouched_while_keepalives_are_unconfigured(self, monkeypatch): + result = await self._dispatch(monkeypatch, interval=None, headers_delay=0.15, body_delay=0.15) + + chunks = [chunk async for chunk in result.body_iterator] + + assert result.headers["x-upstream"] == "kept" + assert chunks == [b"data: hello\n\n"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 6735f2a3780..e8fd5579631 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -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", diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index fdaad567f95..87e10ce7e8d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -1819,3 +1819,92 @@ async def test_run_thread_stream_is_untouched_while_keepalives_are_unconfigured( assert not any(chunk.startswith(": ping") for chunk in chunks) assert chunks[-1] == "data: [DONE]\n\n" + + +# --------------------------------------------------------------------------- +# async_queue_request: SSE keepalives during the time-to-first-token +# --------------------------------------------------------------------------- + + +async def _queue_streaming(monkeypatch, interval, delay=0.3, fails_with=None): + _patch_logging_flags(monkeypatch) + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + router = MagicMock() + router.get_model_list.return_value = [] + + async def _schedule_after_the_scheduler_queue_drains(**kwargs): + await asyncio.sleep(delay) + if fails_with is not None: + raise fails_with + return _async_iter([_simple_chunk(content="queued reply")]) + + router.schedule_acompletion = _schedule_after_the_scheduler_queue_drains + monkeypatch.setattr(ps, "llm_router", router) + + request = MagicMock() + request.url = "http://testserver/queue/chat/completions" + request.method = "POST" + request.headers = {} + request.json = AsyncMock( + return_value={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "priority": 0, + "stream": True, + } + ) + request.is_disconnected = AsyncMock(return_value=False) + + return await ps.async_queue_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=_user_auth(), + ) + + +@pytest.mark.asyncio +async def test_queue_request_pings_while_the_scheduler_is_still_waiting(monkeypatch): + response = await _queue_streaming(monkeypatch, interval=0.05) + + assert isinstance(response, StreamingResponse) + assert response.headers["x-accel-buffering"] == "no" + chunks = [chunk async for chunk in response.body_iterator] + + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b'"content":"queued reply"' in chunks[-2] + assert chunks[-1] == b"data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_queue_request_audits_a_failure_that_arrives_after_the_first_ping(monkeypatch): + audited = [] + + async def _record_failure(*, user_api_key_dict, original_exception, request_data, **kwargs): + audited.append(original_exception) + return None + + monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _record_failure) + + boom = RuntimeError("scheduler died after the wire was already open") + response = await _queue_streaming(monkeypatch, interval=0.05, fails_with=boom) + + assert isinstance(response, StreamingResponse) + chunks = [chunk async for chunk in response.body_iterator] + + assert chunks[0] == b": ping\n\n" + assert audited == [boom] + assert json.loads(chunks[-2].removeprefix(b"data: "))["error"]["code"] == "500" + assert chunks[-1] == b"data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigured(monkeypatch): + response = await _queue_streaming(monkeypatch, interval=None, delay=0.15) + + assert isinstance(response, StreamingResponse) + chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + + assert not any(chunk.startswith(b": ping") for chunk in chunks) + assert chunks[-1] == b"data: [DONE]\n\n" diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index a176e91eaa4..832435711c6 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -11,7 +11,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -324,6 +323,100 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): assert "data: [DONE]" in response.text +def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal_user, monkeypatch): + import asyncio + + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 0.05) + + async def slow_aquery(**kwargs): + await asyncio.sleep(0.3) + return await litellm_module.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the codename?"}], + mock_response="The codename is AZURE-FALCON-42.", + stream=True, + api_key="test-key", + ) + + with ( + patch( # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=slow_aquery), + ), + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + }, + "stream": True, + }, + ) + + assert response.status_code == 200, response.text + assert response.headers.get("content-type", "").startswith("text/event-stream") + assert response.headers["x-accel-buffering"] == "no" + assert response.text.startswith(": ping\n\n") + assert response.text.count(": ping\n\n") >= 3 + assert '"object":"chat.completion.chunk"' in response.text + assert response.text.endswith("data: [DONE]\n\n") + + +def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive( + client_internal_user, monkeypatch +): + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 5) + + async def fast_aquery(**kwargs): + response = await litellm_module.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the codename?"}], + mock_response="The codename is AZURE-FALCON-42.", + stream=True, + api_key="test-key", + ) + response._hidden_params["response_cost"] = 3.45e-06 + return response + + with ( + patch( # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fast_aquery), + ), + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + }, + "stream": True, + }, + ) + + assert response.status_code == 200, response.text + assert response.headers.get("content-type", "").startswith("text/event-stream") + assert response.headers.get("x-litellm-response-cost") == "3.45e-06" + assert not response.text.startswith(": ping") + assert '"object":"chat.completion.chunk"' in response.text + assert response.text.endswith("data: [DONE]\n\n") + + def test_rag_query_merges_managed_store_params(client_internal_user): """ Regression: /v1/rag/query must consult the managed vector store registry diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index f4cd8814bc1..329a33eb440 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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 @@ -58,6 +56,24 @@ def _filter_logs_by_date_range(logs, where): return filtered +_SEARCH_CLAUSE_RE = re.compile( + r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' + r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' + r'AND \(api_key = \$\1 OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 ' + r"OR session_id = \$\1 OR model_id = \$\1\)\)\)" +) + + +def _matches_spend_log_search(log, search): + """Mirror the search clause: request_id across all time, the other id columns inside the window.""" + if log.get("request_id") == search["value"]: + return True + if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}): + return False + columns = ("api_key", "team_id", "user", "end_user", "session_id", "model_id") + return any(log.get(col) == search["value"] for col in columns) + + def _reconstruct_ui_where_from_sql(sql_query, params): """ Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the @@ -77,6 +93,16 @@ def _reconstruct_ui_where_from_sql(sql_query, params): def _iso(value): return value.isoformat() if hasattr(value, "isoformat") else str(value) + search_clause = _SEARCH_CLAUSE_RE.search(clause.group(1)) + if search_clause: + raw_index, start_index, end_index = (int(g) for g in search_clause.groups()) + where["search"] = { + "value": params[raw_index - 1], + "gte": _iso(params[start_index - 1]), + "lte": _iso(params[end_index - 1]), + } + remaining = clause.group(1) if search_clause is None else clause.group(1).replace(search_clause.group(0), "") + eq_cols = { "team_id": "team_id", '"user"': "user", @@ -89,7 +115,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): } date_bounds: dict = {} metadata_conds: list = [] - for cond in (c.strip() for c in clause.group(1).split(" AND ")): + for cond in (c.strip() for c in remaining.split(" AND ")): gte = re.search(r'"startTime" >= \(\$(\d+)', cond) lte = re.search(r'"startTime" <= \(\$(\d+)', cond) alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) @@ -2352,6 +2378,208 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( app.dependency_overrides.pop(ps.user_api_key_auth, None) +def test_build_spend_log_search_condition_windows_every_branch_except_request_id(): + """LIT-4741: request_id matches across all time; the six other id columns only inside the window, + all comparing the pasted value verbatim.""" + start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc) + end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc) + + condition = spend_management_endpoints._build_spend_log_search_condition( + search="key-hash-7", start_date=start, end_date=end, next_param_index=3 + ) + + assert condition.sql == ( + "(request_id = $3 OR (\"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC') " + "AND \"startTime\" <= ($5::timestamptz AT TIME ZONE 'UTC') " + 'AND (api_key = $3 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))' + ) + assert condition.params == ("key-hash-7", start, end) + + +def _search_fixture_logs(today): + recent = (today - datetime.timedelta(days=1)).isoformat() + old = (today - datetime.timedelta(days=90)).isoformat() + base = { + "api_key": "hashed-other", + "user": "user-x", + "team_id": "team-x", + "end_user": "cust-x", + "session_id": "sess-x", + "model_id": "mdl-x", + "spend": 0.01, + "model": "gpt-4", + } + return [ + {**base, "request_id": "req-session", "session_id": "sess-42", "startTime": recent}, + {**base, "request_id": "req-session-old", "session_id": "sess-42", "startTime": old}, + {**base, "request_id": "req-key", "api_key": "hashed-7", "startTime": recent}, + {**base, "request_id": "req-team", "team_id": "team-7", "startTime": recent}, + {**base, "request_id": "req-user", "user": "user-7", "startTime": recent}, + {**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent}, + {**base, "request_id": "req-model", "model_id": "mdl-7", "startTime": recent}, + ] + + +def _search_filter_fn(logs, captured): + def filter_fn(where): + captured["where"] = where + rows = _filter_logs_by_date_range(logs, where) + if "user" in where: + rows = [row for row in rows if row["user"] == where["user"]] + if "search" in where: + rows = [row for row in rows if _matches_spend_log_search(row, where["search"])] + return rows + + return filter_fn + + +def _five_day_window(today): + return { + "start_date": (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": today.strftime("%Y-%m-%d %H:%M:%S"), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "search,expected_request_ids", + [ + ("req-session-old", {"req-session-old"}), + ("sess-42", {"req-session"}), + ("hashed-7", {"req-key"}), + ("team-7", {"req-team"}), + ("user-7", {"req-user"}), + ("cust-7", {"req-end-user"}), + ("mdl-7", {"req-model"}), + ("no-such-id", set()), + ], +) +async def test_ui_view_spend_logs_search_matches_any_id(client, monkeypatch, search, expected_request_ids): + """LIT-4741: one box matches any id column. A request_id is found across all time (the 5-day + window excludes the 90-day-old row), every other column only inside the window, and a raw + sk- key is hashed before it is compared with api_key. The window is not applied globally.""" + today = datetime.datetime.now(timezone.utc) + logs = _search_fixture_logs(today) + captured = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + try: + response = client.get( + "/spend/logs/ui", + params={"search": search, **_five_day_window(today)}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert {row["request_id"] for row in data["data"]} == expected_request_ids + assert data["total"] == len(expected_request_ids) + assert "startTime" not in captured["where"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_logs_v2_search_keeps_global_window(client, monkeypatch): + """The public route keeps the caller's window on the whole query, so a search only finds rows + inside it even by request_id; the windowless request_id branch is a dashboard-only relaxation.""" + today = datetime.datetime.now(timezone.utc) + logs = _search_fixture_logs(today) + captured = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + try: + response = client.get( + "/spend/logs/v2", + params={"search": "req-session-old", **_five_day_window(today)}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["data"] == [] + assert data["total"] == 0 + assert "startTime" in captured["where"] + assert captured["where"]["search"]["value"] == "req-session-old" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "params", + [ + {"search": "req-old"}, + {"search": "req-old", "request_id": "req-old"}, + ], +) +async def test_ui_view_spend_logs_search_requires_dates(client, monkeypatch, params): + """A search needs the window for its non-request_id branches, so it stays required even + alongside a request_id, which on its own may drop the window.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([], lambda where: []), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + try: + response = client.get("/spend/logs/ui", params=params, headers={"Authorization": "Bearer sk-test"}) + assert response.status_code == 400 + assert "date" in response.text.lower() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "search,expected_request_ids", + [("sess-9", {"req-own"}), ("req-foreign", set())], +) +async def test_ui_view_spend_logs_search_keeps_non_admin_scope(client, monkeypatch, search, expected_request_ids): + """A search is scoped like any other listing: an internal user only sees their own rows even + when the id is on someone else's row, and the request_id ownership shortcut is not used.""" + yesterday = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=1)).isoformat() + base = {"api_key": "hashed-key", "team_id": None, "spend": 0.01, "startTime": yesterday, "model": "gpt-4"} + logs = [ + {**base, "request_id": "req-own", "user": "internal_user_1", "session_id": "sess-9"}, + {**base, "request_id": "req-own-other", "user": "internal_user_1", "session_id": "sess-other"}, + {**base, "request_id": "req-foreign", "user": "internal_user_2", "session_id": "sess-9"}, + ] + captured = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=[]), + ) + ownership_check = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._assert_user_can_view_request_id", + ownership_check, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"search": search, "start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert {row["request_id"] for row in response.json()["data"]} == expected_request_ids + assert captured["where"]["user"] == "internal_user_1" + ownership_check.assert_not_awaited() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_unauthorized(client): # Test without authorization header @@ -3095,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}") @@ -3265,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, }, ] @@ -3620,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") @@ -3668,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={ @@ -3679,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] @@ -3694,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""" @@ -4602,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 [] @@ -6351,3 +6734,46 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort assert "OFFSET" in emitted_sql[1] finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_session(client, monkeypatch): + """The dashboard lists sessions by default; a search for an id lists every matching row instead, + so both calls of a session show up rather than one representative, and no session cursor is returned.""" + rows = [_session_representative_row("req-1", "sess-1"), _session_representative_row("req-2", "sess-1")] + + async def mock_query_raw(sql_query, *params): + if "mcp_tool_call_count" in sql_query: + return [] + grouped = "DISTINCT ON" in sql_query or "GROUP BY" in sql_query + visible = rows[:1] if grouped else rows + if "COUNT(*)" in sql_query: + return [{"total_count": len(visible)}] + return visible + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "search": "sess-1", + "group_by_session": "true", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert [row["request_id"] for row in data["data"]] == ["req-1", "req-2"] + assert data["total"] == 2 + assert "next_session_cursor" not in data + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index 9ae932ff01f..a7de3f1d8d6 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -184,6 +184,7 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -247,6 +248,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -314,6 +316,7 @@ async def test_spend_logs_ui_caps_total_for_large_result_sets(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -359,6 +362,7 @@ async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -406,6 +410,7 @@ async def test_spend_logs_ui_out_of_range_page_keeps_total(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=99, @@ -552,6 +557,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -616,6 +622,7 @@ async def test_spend_logs_ui_group_by_session_offset_pages_for_other_sorts(monke api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=2, @@ -664,6 +671,7 @@ async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(m api_key=None, user_id=None, request_id="req-deep-link", + search=None, start_date=None, end_date=None, page=1, diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py index c9ec22ab0df..4bc7c21d8d2 100644 --- a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules import pytest +from azure.core.exceptions import ClientAuthenticationError from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -16,6 +17,143 @@ from litellm.types.secret_managers.get_azure_ad_token_provider import ( ) +class TestDeploymentIdentityCredential: + @staticmethod + def _chain_for(credential_type): + with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "token") as bearer: + get_azure_ad_token_provider( + azure_scope="https://storage.azure.com/.default", + azure_credential=credential_type, + ) + bearer.assert_called_once() + with bearer.call_args.args[0] as chain: + return {type(link).__name__ for link in chain.credentials} + + @staticmethod + def _managed_identity_client_ids(credential_type): + with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "token") as bearer: + get_azure_ad_token_provider( + azure_scope="https://storage.azure.com/.default", + azure_credential=credential_type, + ) + with bearer.call_args.args[0] as chain: + return [ + (link._credential._settings or {}).get("client_id") + for link in chain.credentials + if type(link).__name__ == "ManagedIdentityCredential" + ] + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "workload-identity-client-id", + "AZURE_TENANT_ID": "workload-identity-tenant-id", + "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token", + }, + clear=True, + ) + def test_deployment_identity_reaches_workload_and_managed_identity_only(self): + assert self._chain_for(AzureCredentialType.DeploymentIdentityCredential) == { + "WorkloadIdentityCredential", + "ManagedIdentityCredential", + } + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "workload-identity-client-id", + "AZURE_TENANT_ID": "workload-identity-tenant-id", + "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token", + "AZURE_TOKEN_CREDENTIALS": "dev", + }, + clear=True, + ) + def test_deployment_identity_survives_a_developer_only_token_credentials_setting(self): + """AZURE_TOKEN_CREDENTIALS=dev asks the SDK for developer credentials only, which is every + credential this chain drops, so the deployment's own identity has to win over it""" + assert self._chain_for(AzureCredentialType.DeploymentIdentityCredential) == { + "WorkloadIdentityCredential", + "ManagedIdentityCredential", + } + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "azure-openai-client-id", + "AZURE_CLIENT_SECRET": "azure-openai-client-secret", + "AZURE_TENANT_ID": "azure-openai-tenant-id", + }, + clear=True, + ) + def test_default_azure_credential_keeps_its_full_chain(self): + """Azure OpenAI callers pass DefaultAzureCredential and must be unaffected by the + narrowing that the storage callback asks for""" + full_chain = self._chain_for(AzureCredentialType.DefaultAzureCredential) + + assert "EnvironmentCredential" in full_chain + assert "AzureCliCredential" in full_chain + assert "EnvironmentCredential" not in self._chain_for( + AzureCredentialType.DeploymentIdentityCredential + ) + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "azure-openai-client-id", + "AZURE_CLIENT_SECRET": "azure-openai-client-secret", + "AZURE_TENANT_ID": "azure-openai-tenant-id", + }, + clear=True, + ) + def test_deployment_identity_refuses_to_mint_a_token_for_a_configured_service_principal(self): + """A host carrying only an Azure OpenAI client secret must get no token at all, and the + refusal must name the identities that were actually tried""" + provider = get_azure_ad_token_provider( + azure_scope="https://storage.azure.com/.default", + azure_credential=AzureCredentialType.DeploymentIdentityCredential, + ) + + with pytest.raises(ClientAuthenticationError) as refusal: + provider() + + assert "ManagedIdentityCredential" in str(refusal.value) + assert "EnvironmentCredential" not in str(refusal.value) + assert "AzureCliCredential" not in str(refusal.value) + assert "azure-openai-client-secret" not in str(refusal.value) + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "azure-openai-client-id", + "AZURE_CLIENT_SECRET": "azure-openai-client-secret", + "AZURE_TENANT_ID": "azure-openai-tenant-id", + }, + clear=True, + ) + def test_deployment_identity_still_reaches_a_system_assigned_managed_identity(self): + """AZURE_CLIENT_ID names one identity for the whole proxy, and pointing it at Azure OpenAI + must not hide the system assigned identity the host runs as""" + client_ids = self._managed_identity_client_ids(AzureCredentialType.DeploymentIdentityCredential) + + assert "azure-openai-client-id" in client_ids + assert None in client_ids + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "user-assigned-identity-client-id", + "AZURE_TOKEN_CREDENTIALS": "dev", + }, + clear=True, + ) + def test_deployment_identity_keeps_the_user_assigned_identity_under_a_dev_only_setting(self): + """AZURE_TOKEN_CREDENTIALS=dev asks the SDK for developer credentials only, and the + identity a host actually runs as has to survive that""" + assert "user-assigned-identity-client-id" in self._managed_identity_client_ids( + AzureCredentialType.DeploymentIdentityCredential + ) + + class TestGetAzureAdTokenProvider: @patch.dict( os.environ, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index e4236afc586..5fc96bcfbb1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3855,7 +3855,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input) with_instructions_tokens = router._count_pre_call_check_tokens( - messages=None, input=short_input, instructions=long_instructions + messages=None, input=short_input, request_kwargs={"instructions": long_instructions} ) assert with_instructions_tokens > input_only_tokens @@ -3871,6 +3871,164 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) +_OVERSIZED_TOOL_DESCRIPTION = "look up the answer in the knowledge base. " * 40 + + +@pytest.mark.parametrize( + "prompt_kwargs, tool", + [ + pytest.param( + {"messages": [{"role": "user", "content": "hi"}]}, + { + "type": "function", + "function": { + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + }, + id="chat_completions_tool", + ), + pytest.param( + {"input": "hi"}, + { + "type": "function", + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + id="responses_tool", + ), + pytest.param( + {"messages": [{"role": "user", "content": "hi"}]}, + { + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + id="anthropic_messages_tool", + ), + ], +) +def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwargs, tool): + """ + Tool definitions are sent to the model as prompt tokens but never appear in + `messages` or `input`. A request whose prompt alone fits the context window but + whose prompt plus `tools` exceeds it must be rejected before dispatch, for the + Chat Completions, Responses and Anthropic Messages tool shapes alike. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + + prompt_only_tokens = router._count_pre_call_check_tokens( + messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} + ) + + assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + request_kwargs={"tools": [tool]}, + **prompt_kwargs, + ) + + +@pytest.mark.parametrize( + "system", + [ + pytest.param("You are a meticulous assistant. " * 40, id="system_string"), + pytest.param( + [{"type": "text", "text": "You are a meticulous assistant. " * 40}], + id="system_blocks", + ), + ], +) +def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): + """ + The Anthropic Messages API carries the system prompt as a top-level `system` field, + not as a message. Its tokens reach the model, so a request whose `messages` fit but + whose `messages` plus `system` exceed the context window must be rejected. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + messages = [{"role": "user", "content": "hi"}] + + messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) + + assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + messages=messages, + request_kwargs={"system": system}, + ) + + +@pytest.mark.asyncio +async def test_aanthropic_messages_enforces_context_window_with_system_and_tools(): + """ + End-to-end router regression for /v1/messages: a request whose only oversized + content lives in the top-level `system` field or in `tools` must trip the pre-call + context-window check instead of being dispatched (the deployment uses mock_response, + so reaching the provider handler would return a response rather than raise). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "small-ctx", + "litellm_params": {"model": "anthropic/claude-3-5-haiku-20241022", "mock_response": "hi"}, + "model_info": {"max_input_tokens": 20}, + } + ], + enable_pre_call_checks=True, + ) + messages = [{"role": "user", "content": "hi"}] + + response = await router.aanthropic_messages(model="small-ctx", messages=messages, max_tokens=5) + assert response is not None + + with pytest.raises(litellm.ContextWindowExceededError): + await router.aanthropic_messages( + model="small-ctx", + messages=messages, + max_tokens=5, + system="You are a meticulous assistant. " * 40, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await router.aanthropic_messages( + model="small-ctx", + messages=messages, + max_tokens=5, + tools=[ + { + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ], + ) + + def test_count_pre_call_check_tokens_across_api_surfaces(): """ _count_pre_call_check_tokens must count tokens from chat `messages`, a Responses @@ -4943,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 @@ -7362,14 +7554,14 @@ def test_get_configured_mode_reads_deployment_model_info(): router = litellm.Router( model_list=[ { - "model_name": "chat-model", - "litellm_params": {"model": "openai/some-unmapped-model"}, - "model_info": {"mode": "chat"}, + "model_name": "tts-model", + "litellm_params": {"model": "openai/some-unmapped-tts-model"}, + "model_info": {"mode": "audio_speech"}, } ] ) - assert router.get_configured_mode("chat-model") == "chat" + assert router.get_configured_mode("tts-model") == "audio_speech" def test_get_configured_mode_returns_none_for_unset_or_unknown(): @@ -12701,33 +12893,3 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo bucket = captured.get("litellm_metadata") or captured["metadata"] assert captured["model_info"]["id"] == "provisional-dep" assert bucket["litellm_gateway_injected_cache"] == "" - - -def test_get_configured_mode_reads_deployment_model_info(): - router = Router( - model_list=[ - { - "model_name": "my-tts", - "litellm_params": {"model": "openai/some-unmapped-mode-model"}, - "model_info": {"mode": "audio_speech"}, - } - ] - ) - - assert router.get_configured_mode("my-tts") == "audio_speech" - - -@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": " "}, {"mode": 123}]) -def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info): - router = Router( - model_list=[ - { - "model_name": "plain-model", - "litellm_params": {"model": "openai/some-unmapped-mode-model"}, - "model_info": model_info, - } - ] - ) - - assert router.get_configured_mode("plain-model") is None - assert router.get_configured_mode("unknown-model") is None diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2f85128b4b6..ab1a793e09d 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22328 }, "LIT002": { - "limit": 26760 + "limit": 26750 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16470 + "limit": 16468 }, "LIT011": { - "limit": 5516 + "limit": 5514 }, "LIT012": { "limit": 4489 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx index 4d18ec2ef5f..bef938cd31c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx @@ -90,7 +90,7 @@ describe("AgentsTable", () => { />, ); - const search = screen.getByPlaceholderText("Search agent names or descriptions..."); + const search = screen.getByPlaceholderText("Search agents by name, ID, or description..."); await user.type(search, "billing"); expect(screen.getByText("Billing Router")).toBeInTheDocument(); expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); @@ -101,11 +101,36 @@ describe("AgentsTable", () => { expect(screen.queryByText("Billing Router")).not.toBeInTheDocument(); }); + it("filters agents by a pasted agent_id so only that agent's row survives", async () => { + const user = userEvent.setup(); + render( + , + ); + + const search = screen.getByPlaceholderText("Search agents by name, ID, or description..."); + await user.click(search); + await user.paste("5f3c2a1b-9d8e-4f7a-b6c5-d4e3f2a1b0c9"); + expect(screen.getByText("Billing Router")).toBeInTheDocument(); + expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); + + await user.clear(search); + await user.paste("ffffffff-0000-4000-8000-000000000000"); + expect(screen.queryByText("Billing Router")).not.toBeInTheDocument(); + expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); + expect(screen.getByText("No matching agents")).toBeInTheDocument(); + }); + it("shows the no-match empty state when the search matches nothing", async () => { const user = userEvent.setup(); render(); - await user.type(screen.getByPlaceholderText("Search agent names or descriptions..."), "zzzz"); + await user.type(screen.getByPlaceholderText("Search agents by name, ID, or description..."), "zzzz"); expect(screen.queryByText("Test Agent")).not.toBeInTheDocument(); expect(screen.getByText("No matching agents")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx index 35ed6b66425..aceb07e2e9a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -55,7 +55,12 @@ const AgentsTable: React.FC = ({ const [sorting, setSorting] = useState(DEFAULT_SORTING); const [searchTerm, setSearchTerm] = useState(""); const filteredAgents = useMemo( - () => filterBySearchTerm(agents, searchTerm, (agent) => [agent.agent_name, agent.agent_card_params?.description]), + () => + filterBySearchTerm(agents, searchTerm, (agent) => [ + agent.agent_name, + agent.agent_id, + agent.agent_card_params?.description, + ]), [agents, searchTerm], ); @@ -83,7 +88,7 @@ const AgentsTable: React.FC = ({ setSearchTerm(e.target.value)} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts index 8c9b33f2c3e..84be7e2ef49 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -518,6 +518,24 @@ describe("useKeys", () => { const callUrl = mockFetch.mock.calls[0][0]; expect(callUrl).not.toContain("agent_id"); }); + + it("sends the combined alias-or-ID search as the search param, separate from key_alias and key_hash", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const { result } = renderHook(() => useKeys(1, 10, { search: "pasted-key-id" }), { wrapper }); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + const callUrl = new URL(mockFetch.mock.calls[0][0], "http://localhost"); + expect(callUrl.searchParams.get("search")).toBe("pasted-key-id"); + expect(callUrl.searchParams.has("key_alias")).toBe(false); + expect(callUrl.searchParams.has("key_hash")).toBe(false); + }); }); describe("useDeletedKeys", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 94ded01679d..7e7089e685f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -40,6 +40,7 @@ export interface KeyListCallOptions { selectedKeyAlias?: string | null; userID?: string | null; keyHash?: string | null; + search?: string | null; sortBy?: string | null; sortOrder?: string | null; expand?: string | null; @@ -61,6 +62,7 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number, organization_id: options.organizationID, key_alias: options.selectedKeyAlias, key_hash: options.keyHash, + search: options.search, user_id: options.userID, page, size: pageSize, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx index 5100b998b80..984b8135466 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx @@ -95,6 +95,7 @@ describe("MemoryTable", () => { it("shows the filtered-empty copy when a search is active", () => { render(); expect(screen.getByText("No matching memories")).toBeInTheDocument(); + expect(screen.getByText("No memories match your search.")).toBeInTheDocument(); expect(screen.queryByText("No memories stored yet")).not.toBeInTheDocument(); }); @@ -128,6 +129,7 @@ describe("MemoryTable", () => { const onRefresh = vi.fn(); render(); + expect(screen.getByPlaceholderText("Search by key prefix or memory ID…")).toBeInTheDocument(); fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "u" } }); expect(onSearchChange).toHaveBeenCalledWith("u"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx index 50dd04ee14c..3e37faafe15 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx @@ -36,7 +36,7 @@ function MemoryEmptyState({ hasActiveSearch }: { hasActiveSearch: boolean }) {
{hasActiveSearch - ? "No memories have keys starting with your search." + ? "No memories match your search." : "Memories your agents store under /v1/memory will appear here."}
@@ -81,7 +81,7 @@ export function MemoryTable({ table={table} searchValue={searchValue} onSearchChange={onSearchChange} - searchPlaceholder='Filter by key prefix, e.g. "user:"' + searchPlaceholder="Search by key prefix or memory ID…" onRefresh={onRefresh} isRefreshing={isRefreshing} showViewOptions={false} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx index 9ccef5357b9..b703df652c2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx @@ -1,8 +1,9 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, render, screen } from "@testing-library/react"; +import type { PaginationState } from "@tanstack/react-table"; +import { act, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { MemoryRow } from "@/components/networking"; @@ -13,10 +14,13 @@ interface CapturedTableProps { rowCount: number; data: MemoryRow[]; hasActiveSearch: boolean; + onSearchChange: (value: string) => void; + onPaginationChange: (state: PaginationState) => void; onViewClick: (row: MemoryRow) => void; } const captured = vi.hoisted(() => ({ current: null as CapturedTableProps | null })); +const fetchMemoryListMock = vi.hoisted(() => vi.fn()); vi.mock("./MemoryTable", () => ({ MemoryTable: function MemoryTableMock(props: CapturedTableProps) { @@ -25,6 +29,15 @@ vi.mock("./MemoryTable", () => ({ }, })); +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + fetchMemoryList: fetchMemoryListMock, +})); + +vi.mock("@tanstack/react-pacer/debouncer", () => ({ + useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }], +})); + const renderView = (accessToken: string | null) => { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); return render( @@ -35,6 +48,28 @@ const renderView = (accessToken: string | null) => { }; describe("MemoryView", () => { + beforeEach(() => { + fetchMemoryListMock.mockReset(); + fetchMemoryListMock.mockResolvedValue({ memories: [], total: 0 }); + }); + + it("queries the server with the search box value as `search` and resets to page 1", async () => { + renderView("token"); + await waitFor(() => expect(fetchMemoryListMock).toHaveBeenCalled()); + + act(() => captured.current?.onPaginationChange({ pageIndex: 2, pageSize: 50 })); + await waitFor(() => + expect(fetchMemoryListMock).toHaveBeenLastCalledWith("token", expect.objectContaining({ page: 3 })), + ); + + act(() => captured.current?.onSearchChange("mem-abc123")); + + await waitFor(() => + expect(fetchMemoryListMock).toHaveBeenLastCalledWith("token", { search: "mem-abc123", page: 1, pageSize: 50 }), + ); + expect(captured.current?.hasActiveSearch).toBe(true); + }); + it("keeps the table out of the skeleton state when the token is null (disabled query)", () => { renderView(null); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx index 1d2e5150a62..58d4e42aa96 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx @@ -43,10 +43,8 @@ export const MemoryView: React.FC = ({ accessToken }) => { queryKey: [MEMORY_LIST_KEY, debouncedSearch, pagination.pageIndex, pagination.pageSize], queryFn: () => { if (!accessToken) throw new Error("Access token required"); - // Prefix search matches the Redis-style mental model (namespace scan): - // typing "user:" finds "user:profile", "user:prefs", etc. return fetchMemoryList(accessToken, { - keyPrefix: debouncedSearch || undefined, + search: debouncedSearch || undefined, page: pagination.pageIndex + 1, pageSize: pagination.pageSize, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx index 488190a0fdf..feea01f19ca 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx @@ -41,6 +41,7 @@ const TagTable: React.FC = ({ data, onEdit, onDelete, onSelectTag data={data} columns={columns} getRowId={(tag, index) => tag.name || String(index)} + fillHeight sortingMode="client" sorting={sorting} onSortingChange={setSorting} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx index a04afcd6d45..583492bf837 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx @@ -126,7 +126,7 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => }, [accessToken]); return ( -
+
{selectedTagId ? ( = ({ accessToken, userID, userRole }) => editTag={editTag} /> ) : ( -
+

Tag Management

@@ -162,23 +162,21 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) =>

- -
-
- { - setSelectedTagId(tag.name); - setEditTag(true); - }} - onDelete={handleDelete} - onSelectTag={setSelectedTagId} - /> -
+
+ { + setSelectedTagId(tag.name); + setEditTag(true); + }} + onDelete={handleDelete} + onSelectTag={setSelectedTagId} + />
{/* Create Tag Modal */} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index f74f444ce26..1745c710c51 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -137,8 +137,8 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID />
) : ( -
-
+
+

Vector Store Management

diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 475e3dcd70b..74f44db5cb9 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -400,7 +400,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, } return ( -
+
{publicPage == false ? (
{/* Header with Title, Description and URL */} diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index ef58237a6aa..4be91f22339 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -559,6 +559,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser { key: "your-teams", label: "Your Teams", + className: "flex min-h-0 flex-1 flex-col", children: ( <> = ({ accessToken, userID, userRole, premiumUser { key: "available-teams", label: "Available Teams", + className: "min-h-0 flex-1 overflow-y-auto", children: , }, ...(isProxyAdminRole(userRole || "") @@ -615,6 +617,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser { key: "default-settings", label: "Default Team Settings", + className: "min-h-0 flex-1 overflow-y-auto", children: , }, ] @@ -622,7 +625,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ]; return ( -
+
{selectedTeamId ? ( = ({ accessToken, userID, userRole, premiumUser premiumUser={premiumUser} /> ) : ( - + } title="Teams" @@ -674,7 +677,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser )} /> {tabItems.map((item) => ( - + {item.children} ))} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index 0b38d51e36a..3fb19e522f3 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -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) => ( <> diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index c69662b69dc..881b0b93ff9 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -542,6 +542,19 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { expect((lastCall[2] ?? {}).userID).toBeUndefined(); }); }); + + it("sends the search box as the combined alias-or-ID search rather than the key-alias filter", async () => { + renderWithProviders(); + + fireEvent.change(screen.getByPlaceholderText(/Search by key alias or ID/), { target: { value: mockKey.token } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: mockKey.token })); + }); + const lastOptions = mockUseKeys.mock.calls.at(-1)?.[2]; + expect(lastOptions?.selectedKeyAlias).toBeUndefined(); + expect(lastOptions?.keyHash).toBeUndefined(); + }); }); describe("pagination display – total count comes from useKeys", () => { @@ -663,7 +676,7 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(mockUseKeys).toHaveBeenLastCalledWith( 3, 25, - expect.objectContaining({ selectedKeyAlias: "prod", sortBy: "spend", sortOrder: "asc" }), + expect.objectContaining({ search: "prod", sortBy: "spend", sortOrder: "asc" }), ); }); expect(screen.getByPlaceholderText(/Search by key alias/)).toHaveValue("prod"); @@ -736,7 +749,7 @@ describe("table state lives in the URL so it survives leaving and returning to t fireEvent.change(screen.getByPlaceholderText(/Search by key alias/), { target: { value: "prod" } }); await waitFor(() => { - expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "prod" })); + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: "prod" })); }); await waitFor(() => { expect(lastSearchParam(onUrlUpdate, "page")).toBeNull(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index c424966a0a3..f28ae27b6bc 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -118,7 +118,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const keyListOptions = { teamID: appliedFilters.team_id || undefined, organizationID: appliedFilters.org_id || undefined, - selectedKeyAlias: searchQuery.trim() || undefined, + search: searchQuery.trim() || undefined, userID: appliedFilters.user_id || undefined, keyHash: appliedFilters.key_hash || undefined, sortBy, @@ -256,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { } return ( -
+
} 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) => ( <> @@ -291,7 +291,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { table={table} searchValue={searchInput} onSearchChange={handleSearchChange} - searchPlaceholder="Search by key alias…" + searchPlaceholder="Search by key alias or ID…" onRefresh={() => refetch?.()} isRefreshing={isFetching} onOpenFilters={() => setFiltersOpen(true)} diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx index 4d9e2122739..9022d424369 100644 --- a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx @@ -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 }) => ( - <> -
- onChange({ ...value, deployment_affinity: deploymentAffinity })} - aria-label="Pin a session to one deployment per model group" - /> - Pin a session to one deployment per model group -
- - Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to - load-balance every turn. - - -); +}> = ({ value, onChange }) => { + const [ttlDraft, setTtlDraft] = React.useState(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 ( + <> +
+ onChange({ ...value, deployment_affinity: deploymentAffinity })} + aria-label="Pin a session to one deployment per model group" + /> + Pin a session to one deployment per model group +
+ + Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to + load-balance every turn. + +
+ + setTtlDraft(event.target.value)} + onBlur={(event) => commitTtl(event.target.value)} + /> + + Refreshes after every request that reuses a pin. Empty tracks the backend default of{" "} + {DEFAULT_SESSION_AFFINITY_TTL_SECONDS} seconds. + +
+ + ); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 02cb0543e17..bdca5205b2e 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -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(); + 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(); + 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(); + 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", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 791fd1fcfbb..06363830d64 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 836b4c6ff3d..dfef8171c51 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -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, }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index c335beade91..1a1725b8dd0 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -388,6 +388,7 @@ const AddAutoRouterTab: React.FC = ({ 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) => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 53ab859baa0..9ee555f5dd2 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index b8af55e8c8a..956e593a234 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -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; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index f24f3033901..877b35199a4 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 6067e72e547..970bcaa545f 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -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 () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 54766df93ec..ea1e5cba6a3 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -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, diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 7a220dd4711..9df9e9a9209 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -854,3 +854,48 @@ describe("userListCall search serialization", () => { expect(lastParams(mockFetch).get("user_email")).toBe("ada@example.com"); }); }); + +describe("fetchMemoryList search serialization", () => { + const originalFetch = global.fetch; + + afterEach(() => { + global.fetch = originalFetch; + }); + + const mockOkFetch = () => { + const emptyPage = { memories: [], total: 0 }; + const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue(emptyPage) } as any); + global.fetch = mockFetch as any; + return mockFetch; + }; + + const lastParams = (mockFetch: ReturnType) => { + const [url] = mockFetch.mock.calls.at(-1) ?? []; + return new URL(url as string, "http://example.com").searchParams; + }; + + it("sends the search box value as search and omits key_prefix and key", async () => { + const mockFetch = mockOkFetch(); + + await Networking.fetchMemoryList("token", { search: "mem-abc123", page: 1, pageSize: 50 }); + + const params = lastParams(mockFetch); + expect(params.get("search")).toBe("mem-abc123"); + expect(params.has("key_prefix")).toBe(false); + expect(params.has("key")).toBe(false); + expect(params.get("page")).toBe("1"); + expect(params.get("page_size")).toBe("50"); + }); + + it("keeps key_prefix and key working when no search is given", async () => { + const mockFetch = mockOkFetch(); + + await Networking.fetchMemoryList("token", { keyPrefix: "user:" }); + expect(lastParams(mockFetch).get("key_prefix")).toBe("user:"); + expect(lastParams(mockFetch).has("search")).toBe(false); + + await Networking.fetchMemoryList("token", { key: "user:profile" }); + expect(lastParams(mockFetch).get("key")).toBe("user:profile"); + expect(lastParams(mockFetch).has("search")).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index d8762565e08..1384679a88a 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2056,6 +2056,7 @@ interface UiSpendLogsParams { exclude_internal_health_checks?: boolean; group_by_session?: boolean; session_cursor?: string; + search?: string; } interface UiSpendLogsCallOptions { @@ -6563,6 +6564,7 @@ interface UiAuditLogsParams { changed_by_api_key?: string; object_team_id?: string; object_key_hash?: string; + search?: string | null; sort_by?: string; sort_order?: "asc" | "desc"; } @@ -8061,15 +8063,18 @@ export const fetchMemoryList = async ( options: { key?: string; keyPrefix?: string; + search?: string; page?: number; pageSize?: number; } = {}, ): Promise => { const base = proxyBaseUrl ? `${proxyBaseUrl}/v1/memory` : `/v1/memory`; const params = new URLSearchParams(); - // keyPrefix takes precedence — backend also does, but we omit `key` + // Backend precedence is search > key_prefix > key; only the winner is sent // to keep the URL clean and intent obvious. - if (options.keyPrefix) { + if (options.search) { + params.append("search", options.search); + } else if (options.keyPrefix) { params.append("key_prefix", options.keyPrefix); } else if (options.key) { params.append("key", options.key); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index f502ea77127..149554a3ac3 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -613,8 +613,11 @@ describe("DataTable layout", () => { it("makes the header sticky and constrains body height when maxBodyHeight is set", () => { render(); - 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", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 60267606951..17a5fe42d1c 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -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(column: ColumnDef): string | undefined { if ("id" in column && typeof column.id === "string") { @@ -533,6 +537,7 @@ export function DataTable(props: DataTableProps { @@ -593,13 +598,13 @@ export function DataTable(props: DataTableProps{toolbar(table)}
}
{table.getHeaderGroups().map((headerGroup) => ( diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index 8bf5d639d6c..0d9d0988aa1 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -30,6 +30,8 @@ vi.mock("@tanstack/react-pacer/debouncer", () => ({ const mockUseKeys = useKeys as MockedFunction; +const KEY_HASH = "88a145505dd6e87e2ea166fcef1e4b53948dbdb32af6431dfd05ec06b571ee52"; + const createMockKey = (overrides: Partial = {}): KeyResponse => ({ token: "sk-test123", @@ -277,7 +279,7 @@ describe("TeamVirtualKeysTable", () => { ); }); - it("maps the search box to a server-side key-alias query", async () => { + it("maps the Key ID drawer filter to a server-side useKeys query and clears it", async () => { const user = userEvent.setup(); mockUseKeys.mockReturnValue({ data: { keys: [createMockKey()], total_count: 1, current_page: 1, total_pages: 1 }, @@ -288,11 +290,42 @@ describe("TeamVirtualKeysTable", () => { renderWithProviders(); - fireEvent.change(await screen.findByTestId("datatable-search"), { target: { value: "check-002" } }); + await user.click(await screen.findByTestId("datatable-filters-trigger")); + const drawerBody = await screen.findByTestId("filter-drawer-body"); + fireEvent.change(within(drawerBody).getByPlaceholderText("Enter Key ID…"), { target: { value: KEY_HASH } }); + await user.click(screen.getByTestId("filter-drawer-apply")); await waitFor(() => - expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "check-002" })), + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ keyHash: KEY_HASH })), ); + expect(screen.getByTestId("filter-chip-key_hash")).toHaveTextContent("Key ID"); + + await user.click(screen.getByTestId("datatable-clear-filters")); + await waitFor(() => + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ keyHash: undefined })), + ); + }); + + it("maps the search box to the combined alias-or-ID search rather than the key-alias filter", async () => { + mockUseKeys.mockReturnValue({ + data: { keys: [createMockKey()], total_count: 1, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as unknown as ReturnType); + + renderWithProviders(); + + const searchBox = await screen.findByTestId("datatable-search"); + expect(searchBox).toHaveAttribute("placeholder", "Search by key alias or ID…"); + fireEvent.change(searchBox, { target: { value: KEY_HASH } }); + + await waitFor(() => + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: KEY_HASH })), + ); + const lastOptions = mockUseKeys.mock.calls.at(-1)?.[2]; + expect(lastOptions?.selectedKeyAlias).toBeUndefined(); + expect(lastOptions?.keyHash).toBeUndefined(); }); it("should show Loading keys when isPending", async () => { diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index bd7faa41ee9..b0380255b95 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -68,19 +68,17 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const pageIndex = tablePagination.pageIndex; const pageSize = tablePagination.pageSize; - const { - data: keys, - isPending: isLoading, - isFetching, - refetch, - } = useKeys(pageIndex + 1, pageSize, { + const keyListOptions = { teamID: teamId, - selectedKeyAlias: searchQuery.trim() || undefined, + search: searchQuery.trim() || undefined, userID: getFilterValue("user_id"), + keyHash: getFilterValue("key_hash"), sortBy: sortBy || undefined, sortOrder: sortOrder || undefined, expand: "user", - }); + }; + + const { data: keys, isPending: isLoading, isFetching, refetch } = useKeys(pageIndex + 1, pageSize, keyListOptions); const displayKeys = useMemo(() => { const kList = keys?.keys || []; @@ -445,7 +443,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, []); return ( -
+
{selectedKey ? ( ) : ( -
+
( <> @@ -481,11 +478,11 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi table={table} searchValue={searchInput} onSearchChange={handleSearchChange} - searchPlaceholder="Search by key alias…" + searchPlaceholder="Search by key alias or ID…" onRefresh={() => refetch?.()} isRefreshing={isFetching} onOpenFilters={() => setFiltersOpen(true)} - filterLabels={{ user_id: "User ID" }} + filterLabels={{ user_id: "User ID", key_hash: "Key ID" }} /> {({ get, set }) => ( - - set("user_id", event.target.value)} - placeholder="Filter by user ID…" - /> - + <> + + set("user_id", event.target.value)} + placeholder="Filter by user ID…" + /> + + + set("key_hash", event.target.value)} + placeholder="Enter Key ID…" + /> + + )} diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 614b3b7af62..dcadc141103 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -216,24 +216,22 @@ const UserDashboard: React.FC = ({ const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer"; return ( -
-
- - ) : undefined - } - /> -
+
+ + ) : undefined + } + />
); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx new file mode 100644 index 00000000000..3b27f663b8e --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx @@ -0,0 +1,145 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { chooseSelectOption } from "../../../tests/test-utils"; +import AuditLogsPanel from "./AuditLogsPanel"; + +vi.mock("../networking", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, uiAuditLogsCall: vi.fn() }; +}); + +// Resolve the debounced search synchronously so typed input reaches the query within the test tick. +vi.mock("@tanstack/react-pacer/debouncer", () => ({ + useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }], +})); + +import { uiAuditLogsCall } from "../networking"; + +type AuditLogsParams = NonNullable[0]["params"]>; + +const PAGE_SIZE = 50; + +const ID_PARAM_KEYS = [ + "search", + "object_id", + "changed_by", + "object_team_id", + "object_key_hash", + "action", + "table_name", +] as const satisfies readonly (keyof AuditLogsParams)[]; + +const respondWith = (total: number) => { + const response = { audit_logs: [], total, page: 1, page_size: PAGE_SIZE, total_pages: Math.ceil(total / PAGE_SIZE) }; + return vi.mocked(uiAuditLogsCall).mockResolvedValue(response); +}; + +const lastCall = () => vi.mocked(uiAuditLogsCall).mock.calls.at(-1)?.[0]; +const sentIdParams = () => ID_PARAM_KEYS.filter((key) => lastCall()?.params?.[key] !== undefined); + +const defaultProps = { + accessToken: "sk-test", + token: "jwt-test", + userRole: "Admin", + userID: "user-1", + isActive: true, + premiumUser: true, +}; + +const renderPanel = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; + +const TEXT_FILTERS: { filterId: string; placeholder: string; paramKey: keyof AuditLogsParams }[] = [ + { filterId: "object_id", placeholder: "Enter object ID…", paramKey: "object_id" }, + { filterId: "changed_by", placeholder: "Enter user ID…", paramKey: "changed_by" }, + { filterId: "team_id", placeholder: "Enter team ID…", paramKey: "object_team_id" }, + { filterId: "key_hash", placeholder: "Enter key hash…", paramKey: "object_key_hash" }, +]; + +const SELECT_FILTERS: { + label: string; + comboboxIndex: number; + option: string; + paramKey: keyof AuditLogsParams; + value: string; +}[] = [ + { label: "Action", comboboxIndex: 0, option: "Created", paramKey: "action", value: "created" }, + { label: "Table", comboboxIndex: 1, option: "Teams", paramKey: "table_name", value: "LiteLLM_TeamTable" }, +]; + +describe("AuditLogsPanel", () => { + beforeEach(() => { + vi.clearAllMocks(); + respondWith(0); + }); + + it("sends the typed search as params.search and returns to the first page", async () => { + const user = userEvent.setup(); + respondWith(120); + renderPanel(); + await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled()); + expect(lastCall()?.params?.search).toBeUndefined(); + + await user.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastCall()?.page).toBe(2)); + + await user.type(screen.getByTestId("datatable-search"), "team-abc"); + + await waitFor(() => expect(lastCall()?.params?.search).toBe("team-abc")); + expect(lastCall()?.page).toBe(1); + expect(sentIdParams()).toEqual(["search"]); + }); + + it("trims the search and drops params.search once the box is cleared", async () => { + const user = userEvent.setup(); + renderPanel(); + const input = await screen.findByTestId("datatable-search"); + + await user.type(input, " abc"); + await waitFor(() => expect(lastCall()?.params?.search).toBe("abc")); + + await user.clear(input); + + await waitFor(() => expect(lastCall()?.params?.search).toBeUndefined()); + expect(sentIdParams()).toEqual([]); + }); + + it.each(TEXT_FILTERS)("maps the $filterId drawer filter to params.$paramKey", async ({ placeholder, paramKey }) => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled()); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + fireEvent.change(await screen.findByPlaceholderText(placeholder), { target: { value: "val-1" } }); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastCall()?.params?.[paramKey]).toBe("val-1")); + expect(sentIdParams()).toEqual([paramKey]); + }); + + it.each(SELECT_FILTERS)( + "maps the $label drawer select to params.$paramKey", + async ({ comboboxIndex, option, paramKey, value }) => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled()); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + const triggers = await screen.findAllByRole("combobox"); + await chooseSelectOption(user, triggers[comboboxIndex], option); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastCall()?.params?.[paramKey]).toBe(value)); + expect(sentIdParams()).toEqual([paramKey]); + }, + ); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx index 81bd4a19f76..5ce4ca6053f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx @@ -1,7 +1,9 @@ import { useCallback, useState } from "react"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { useQuery, keepPreviousData } from "@tanstack/react-query"; import { ColumnFiltersState, OnChangeFn, PaginationState } from "@tanstack/react-table"; import { resolveLogoSrc } from "@/lib/assetPaths"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { uiAuditLogsCall } from "../networking"; import { AuditLogEntry } from "./AuditLogsTableColumns"; import { AuditLogsTable } from "./AuditLogsTable"; @@ -39,9 +41,13 @@ export default function AuditLogsPanel({ }: AuditLogsProps) { const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); const [columnFilters, setColumnFilters] = useState([]); + const [searchInput, setSearchInput] = useState(""); + const [debouncedSearch] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); const [selectedLog, setSelectedLog] = useState(null); const [drawerOpen, setDrawerOpen] = useState(false); + const searchTerm = debouncedSearch.trim(); + const getFilterValue = (columnId: string): string | undefined => { const entry = columnFilters.find((filter) => filter.id === columnId); return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; @@ -50,7 +56,7 @@ export default function AuditLogsPanel({ const canQueryAuditLogs = !!accessToken && !!token && !!userRole && !!userID && isActive && premiumUser; const query = useQuery({ - queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters], + queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters, searchTerm], queryFn: async () => { if (!accessToken) { return { audit_logs: [], total: 0, page: 1, page_size: pagination.pageSize, total_pages: 0 }; @@ -60,6 +66,7 @@ export default function AuditLogsPanel({ page: pagination.pageIndex + 1, page_size: pagination.pageSize, params: { + search: searchTerm || undefined, object_id: getFilterValue("object_id"), changed_by: getFilterValue("changed_by"), object_key_hash: getFilterValue("key_hash"), @@ -80,6 +87,11 @@ export default function AuditLogsPanel({ setPagination((prev) => ({ ...prev, pageIndex: 0 })); }, []); + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); + setPagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + const handleViewLog = useCallback((log: AuditLogEntry) => { setSelectedLog(log); setDrawerOpen(true); @@ -128,6 +140,8 @@ export default function AuditLogsPanel({ onPaginationChange={setPagination} columnFilters={columnFilters} onColumnFiltersChange={handleColumnFiltersChange} + searchValue={searchInput} + onSearchChange={handleSearchChange} onRefresh={() => query.refetch()} onViewLog={handleViewLog} /> diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx index 7349c3019ae..d8e549715af 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx @@ -120,6 +120,24 @@ describe("AuditLogsTable", () => { expect(screen.getByText("No matching audit logs")).toBeInTheDocument(); }); + it("renders the toolbar search box from the search props and forwards typed input", () => { + const onSearchChange = vi.fn(); + renderTable({ searchValue: "team-", onSearchChange }); + + const input = screen.getByPlaceholderText("Search audit logs by ID…"); + expect(input).toHaveValue("team-"); + + fireEvent.change(input, { target: { value: "team-7" } }); + expect(onSearchChange).toHaveBeenCalledWith("team-7"); + }); + + it("treats an active search as a filter for the empty state", () => { + const emptySearchResult = { data: [], rowCount: 0, searchValue: "zzz", onSearchChange: vi.fn() }; + renderTable(emptySearchResult); + + expect(screen.getByText("No matching audit logs")).toBeInTheDocument(); + }); + it("renders active filter chips with human-readable labels", () => { const filters: ColumnFiltersState = [{ id: "action", value: "created" }]; renderTable({ columnFilters: filters }); diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx index 799505ed07c..cef828838e2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx @@ -24,6 +24,8 @@ interface AuditLogsTableProps { onPaginationChange: OnChangeFn; columnFilters: ColumnFiltersState; onColumnFiltersChange: OnChangeFn; + searchValue?: string; + onSearchChange?: (value: string) => void; onRefresh: () => void; onViewLog: (log: AuditLogEntry) => void; } @@ -102,11 +104,14 @@ export function AuditLogsTable({ onPaginationChange, columnFilters, onColumnFiltersChange, + searchValue, + onSearchChange, onRefresh, onViewLog, }: AuditLogsTableProps) { const [filtersOpen, setFiltersOpen] = useState(false); const columns = useMemo(() => getAuditLogsTableColumns({ onViewLog }), [onViewLog]); + const hasActiveSearch = Boolean(searchValue?.trim()); return ( 0} />} + noDataMessage={ 0 || hasActiveSearch} />} size="compact" toolbar={(table) => ( <> setFiltersOpen(true)} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 93066c106ee..be0d0049c13 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -54,6 +54,14 @@ vi.mock("./LogDetailsDrawer", () => ({ }, })); +const debounce = vi.hoisted(() => ({ settled: null as string | null })); + +vi.mock("@tanstack/react-pacer/debouncer", () => ({ + useDebouncedValue: vi.fn((value: unknown) => [debounce.settled ?? value, { cancel: vi.fn(), flush: vi.fn() }]), +})); + +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { uiSpendLogsCall } from "../networking"; const logEntry = (overrides: Partial): LogEntry => ({ @@ -136,6 +144,7 @@ describe("RequestLogsPanel", () => { sessionStorage.clear(); testQueryClient.clear(); respondWith([]); + debounce.settled = null; }); describe("server-grouped session pagination (#38060)", () => { @@ -322,9 +331,8 @@ describe("RequestLogsPanel", () => { }); }); - describe("search by request id (LIT-3981)", () => { - it("sends the typed request id to the server on the first page instead of filtering the loaded rows", async () => { - const user = userEvent.setup(); + describe("search by any id (LIT-3981, LIT-4741)", () => { + it("sends the typed id to the server as search on the first page instead of filtering the loaded rows", async () => { renderPanel(); await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); @@ -334,9 +342,54 @@ describe("RequestLogsPanel", () => { await waitFor(() => { const call = lastCall(); if (!call) throw new Error("uiSpendLogsCall was not called"); - expect(call.params?.request_id).toBe("req-on-another-page"); + expect(call.params?.search).toBe("req-on-another-page"); expect(call.page).toBe(1); }); + expect(lastCall()?.params?.request_id).toBeUndefined(); + expect(lastCall()?.params?.session_cursor).toBeUndefined(); + }); + + it("sends the debounced value to the server while the box shows what is being typed", async () => { + debounce.settled = "settled-id"; + renderPanel(); + + await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); + + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "still-typing" } }); + + expect(screen.getByTestId("datatable-search")).toHaveValue("still-typing"); + await waitFor(() => + expect(useDebouncedValue).toHaveBeenLastCalledWith("still-typing", { wait: DEBOUNCE_WAIT_MS }), + ); + await waitFor(() => expect(lastCall()?.params?.search).toBe("settled-id")); + const sentLiveValue = vi + .mocked(uiSpendLogsCall) + .mock.calls.some(([options]) => options.params?.search === "still-typing"); + expect(sentLiveValue).toBe(false); + }); + + it("shows a Search chip whose remove button clears the box and restores the unsearched listing", async () => { + const user = userEvent.setup(); + vi.mocked(uiSpendLogsCall).mockImplementation(async ({ params }) => { + const data = + params?.search === "sess-42" + ? [logEntry({ request_id: "req-sess", session_id: "sess-42" })] + : [logEntry({ request_id: "req-initial" })]; + return { data, total: data.length, page: 1, page_size: 50, total_pages: 1 }; + }); + renderPanel(); + + await waitFor(() => expect(row("req-initial")).not.toBeNull()); + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "sess-42" } }); + await waitFor(() => expect(row("req-sess")).not.toBeNull()); + expect(row("req-initial")).toBeNull(); + expect(screen.getByTestId("filter-chip-search")).toHaveTextContent("Search:sess-42"); + + await user.click(screen.getByRole("button", { name: "Remove Search filter" })); + + expect(screen.getByTestId("datatable-search")).toHaveValue(""); + await waitFor(() => expect(row("req-initial")).not.toBeNull()); + expect(row("req-sess")).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 96cf2bd5dde..9b99c6af923 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -1,11 +1,13 @@ "use client"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { useQuery, type UseQueryOptions } from "@tanstack/react-query"; import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import moment from "moment"; import { useCallback, useEffect, useMemo, useState } from "react"; import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import type { KeyResponse } from "../key_team_helpers/key_list"; import { keyInfoV1Call, uiSpendLogsCall } from "../networking"; import KeyInfoView from "../templates/key_info_view"; @@ -75,12 +77,22 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, sessionStorage.setItem("excludeInternalHealthChecks", JSON.stringify(excludeInternalHealthChecks)); }, [excludeInternalHealthChecks]); + const searchTerm = useMemo(() => { + const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.SEARCH); + return typeof entry?.value === "string" ? entry.value : ""; + }, [columnFilters]); + const [debouncedSearch] = useDebouncedValue(searchTerm, { wait: DEBOUNCE_WAIT_MS }); + const queryColumnFilters = useMemo(() => { + const others = columnFilters.filter((filter) => filter.id !== LOG_FILTER_IDS.SEARCH); + return debouncedSearch === "" ? others : [...others, { id: LOG_FILTER_IDS.SEARCH, value: debouncedSearch }]; + }, [columnFilters, debouncedSearch]); + const { logsQuery, filteredLogs, allTeams, usesSessionCursor } = useLogFilterLogic({ accessToken, token, userRole, userID, - columnFilters, + columnFilters: queryColumnFilters, activeTab: isActive ? "request logs" : "inactive", isLiveTail, excludeInternalHealthChecks, @@ -155,15 +167,10 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const rows: LogEntry[] = filteredLogs.data; - const searchTerm = useMemo(() => { - const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.REQUEST_ID); - return typeof entry?.value === "string" ? entry.value : ""; - }, [columnFilters]); - const handleSearchChange = useCallback((value: string) => { setColumnFilters((previous) => { - const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.REQUEST_ID); - return value === "" ? others : [...others, { id: LOG_FILTER_IDS.REQUEST_ID, value }]; + const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.SEARCH); + return value === "" ? others : [...others, { id: LOG_FILTER_IDS.SEARCH, value }]; }); setSessionCursors({}); setPagination((previous) => ({ ...previous, pageIndex: 0 })); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index 4159b3b699b..db108ee681c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -86,6 +86,7 @@ export function RequestLogsTable({ data={data} columns={columns} getRowId={(row) => row.request_id} + fillHeight sortingMode="server" sorting={sorting} onSortingChange={onSortingChange} @@ -108,7 +109,7 @@ export function RequestLogsTable({ table={table} searchValue={searchValue} onSearchChange={onSearchChange} - searchPlaceholder="Search by Request ID" + searchPlaceholder="Search logs by ID…" onRefresh={onRefresh} isRefreshing={isRefreshing} onOpenFilters={() => setFiltersOpen(true)} diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index f20285dffa0..aadf90fad6c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -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(REQUEST_LOGS_TAB.id); const canViewAuditLogs = useCan("viewAuditLogs"); @@ -78,8 +81,8 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p }; return ( -
- setActiveTab(value as LogsTabId)}> +
+ setActiveTab(value as LogsTabId)} className="min-h-0 flex-1"> {tabs.map((tab) => ( @@ -88,7 +91,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p ))} {tabs.map((tab) => ( - + {renderPanel(tab.id)} ))} diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx index 1b738db097d..6af791cc50e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx @@ -91,6 +91,7 @@ describe("useLogFilterLogic", () => { { id: LOG_FILTER_IDS.ERROR_CODE, value: "429", param: "error_code" }, { id: LOG_FILTER_IDS.ERROR_MESSAGE, value: "rate limited", param: "error_message" }, { id: LOG_FILTER_IDS.USER_ID, value: "user-9", param: "user_id" }, + { id: LOG_FILTER_IDS.SEARCH, value: "any-id", param: "search" }, ]; it.each(cases)("sends $id as $param", async ({ id, value, param }) => { diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 90f0f0a60f1..3d368527ad9 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -33,6 +33,7 @@ export const LOG_FILTER_IDS = { PUBLIC_MODEL_OR_SEARCH_TOOL: "model", REQUEST_ID: "request_id", USER_ID: "user_id", + SEARCH: "search", } as const; export const LOG_FILTER_LABELS: Record = { @@ -48,6 +49,7 @@ export const LOG_FILTER_LABELS: Record = { [LOG_FILTER_IDS.SESSION_ID]: "Session ID", [LOG_FILTER_IDS.MODEL_ID]: "Model", [LOG_FILTER_IDS.PUBLIC_MODEL_OR_SEARCH_TOOL]: "Public model / search tool", + [LOG_FILTER_IDS.SEARCH]: "Search", }; export interface LogsWindow { @@ -175,6 +177,7 @@ export function useLogFilterLogic({ api_key: getFilterValue(columnFilters, LOG_FILTER_IDS.KEY_HASH), team_id: getFilterValue(columnFilters, LOG_FILTER_IDS.TEAM_ID), request_id: getFilterValue(columnFilters, LOG_FILTER_IDS.REQUEST_ID), + search: getFilterValue(columnFilters, LOG_FILTER_IDS.SEARCH), session_id: getFilterValue(columnFilters, LOG_FILTER_IDS.SESSION_ID), user_id: userIdFilter, end_user: getFilterValue(columnFilters, LOG_FILTER_IDS.END_USER), diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 3dd9911c794..ffede7b2a6b 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -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"); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index ff482bd23b8..b01108f1631 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -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, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 93bc1de8e7e..3dcfeb64866 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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 @@ -40904,6 +40920,8 @@ export interface operations { object_team_id?: string | null; /** @description Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only) */ object_key_hash?: string | null; + /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value */ + search?: string | null; /** @description Column to sort by (e.g. 'updated_at', 'action', 'table_name') */ sort_by?: string | null; /** @description Sort order ('asc' or 'desc') */ @@ -49622,6 +49640,8 @@ export interface operations { key_hash?: string | null; /** @description Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */ key_alias?: string | null; + /** @description Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive). */ + search?: string | null; /** @description Return full key object */ return_full_object?: boolean; /** @description Include all keys for teams that user is an admin of. */ @@ -56880,6 +56900,8 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; + /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + search?: string | null; }; header?: never; path?: never; @@ -56996,6 +57018,8 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; + /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + search?: string | null; }; header?: never; path?: never; @@ -63308,6 +63332,8 @@ export interface operations { key?: string | null; /** @description Filter by key prefix (Redis-style namespace scan). Mutually exclusive with `key`; if both are provided, `key_prefix` wins. */ key_prefix?: string | null; + /** @description Match entries whose key starts with this value or whose memory_id equals it. Takes precedence over `key_prefix` and `key` when provided. */ + search?: string | null; page?: number; page_size?: number; };