diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index baabfad6852..be761e1258b 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -23,7 +23,7 @@ from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from .base_cache import BaseCache from .in_memory_cache import InMemoryCache -from .redis_cache import RedisCache, log_redis_failure +from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -206,9 +206,12 @@ class DualCache(BaseCache): redis_result: Final = self.redis_cache.batch_get_cache( key_list=sublist_keys, parent_otel_span=parent_otel_span ) - except Exception: + except Exception as e: # Do not throttle subsequent callers if the Redis read fails. self._rollback_redis_batch_key_reservations(previous_access_times) + if isinstance(e, RedisCircuitBreakerOpenError): + verbose_logger.debug("LiteLLM Cache: batch_get_cache served from memory only: %s", e) + return result raise if self.in_memory_cache is not None: @@ -325,9 +328,12 @@ class DualCache(BaseCache): redis_result: Final = await self.redis_cache.async_batch_get_cache( sublist_keys, parent_otel_span=parent_otel_span ) - except Exception: + except Exception as e: # Do not throttle subsequent callers if the Redis read fails. self._rollback_redis_batch_key_reservations(previous_access_times) + if isinstance(e, RedisCircuitBreakerOpenError): + verbose_logger.debug("LiteLLM Cache: async_batch_get_cache served from memory only: %s", e) + return result raise # Short-circuit if redis_result is None or contains only None values diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 6b93529e456..2c36995c4f8 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,6 +18,7 @@ import logging import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar +from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast @@ -196,8 +197,14 @@ class RedisCircuitBreaker: self._timeout_streak_started_at: float | None = None self._opened_at: float | None = None self._state = self.CLOSED + self._generation = 0 _breaker_metrics().record_state_change(None, self._state) + @property + def generation(self) -> int: + """Counts state transitions, so a call can tell whether the breaker moved while it ran.""" + return self._generation + def is_open(self) -> bool: """Returns True if Redis calls should be skipped.""" if not self.enabled: @@ -250,7 +257,7 @@ class RedisCircuitBreaker: self._set_state(self.OPEN) def record_success(self) -> None: - if not self.enabled: + if not self.enabled or self._state == self.OPEN: return if self._state == self.HALF_OPEN: verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered") @@ -266,6 +273,7 @@ class RedisCircuitBreaker: _breaker_metrics().record_transition(state) _breaker_metrics().record_state_change(self._state, state) self._state = state + self._generation += 1 _RedisCallResult = TypeVar("_RedisCallResult") @@ -405,21 +413,33 @@ def log_redis_failure( logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None) -def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int: - """Reject the call if the breaker is open, else return the swallowed-failure count to compare against.""" +@dataclass(frozen=True, slots=True) +class _BreakerAdmission: + swallowed_before: int + generation: int + + +def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> _BreakerAdmission: + """Reject the call if the breaker is open, else record what its success may later prove.""" if breaker.is_open(): raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open — skipping {name}") - return _swallowed_redis_failures.get() + return _BreakerAdmission(swallowed_before=_swallowed_redis_failures.get(), generation=breaker.generation) -def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None: - """Record success only when nothing failed while the call ran. +def _exit_circuit_breaker(breaker: RedisCircuitBreaker, admission: _BreakerAdmission) -> None: + """Record success only when nothing failed while the call ran and the breaker has not moved since. Several Redis methods catch their own connection errors and return a default, so a - method that returned is not on its own proof of a healthy Redis. + method that returned is not on its own proof of a healthy Redis. A success also vouches + only for the breaker state that admitted the call: a call admitted before the breaker + opened, or a probe admitted before a later failure reopened it, finishes knowing nothing + about whether Redis has recovered since, so only the current probe may close the breaker. """ - if _swallowed_redis_failures.get() == swallowed_before: - breaker.record_success() + if _swallowed_redis_failures.get() != admission.swallowed_before: + return + if breaker.generation != admission.generation: + return + breaker.record_success() async def _run_under_circuit_breaker( @@ -432,14 +452,14 @@ async def _run_under_circuit_breaker( Shared by the method decorator and the Lua script executor so both feed the same health signal. """ - swallowed_before: Final = _enter_circuit_breaker(breaker, name) + admission: Final = _enter_circuit_breaker(breaker, name) try: result: Final = await call() except Exception as e: if _is_redis_health_failure(e): breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) raise - _exit_circuit_breaker(breaker, swallowed_before) + _exit_circuit_breaker(breaker, admission) return result @@ -449,14 +469,14 @@ def _run_under_circuit_breaker_sync( call: Callable[[], _RedisCallResult], ) -> _RedisCallResult: """Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path.""" - swallowed_before: Final = _enter_circuit_breaker(breaker, name) + admission: Final = _enter_circuit_breaker(breaker, name) try: result: Final = call() except Exception as e: if _is_redis_health_failure(e): breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) raise - _exit_circuit_breaker(breaker, swallowed_before) + _exit_circuit_breaker(breaker, admission) return result @@ -1337,6 +1357,7 @@ class RedisCache(BaseCache): except Exception: return ast.literal_eval(decoded) + @_redis_circuit_breaker_guard_sync def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): try: key = self.check_and_fix_namespace(key=key) @@ -1356,8 +1377,8 @@ class RedisCache(BaseCache): print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}") return self._get_cache_logic(cached_response=cached_response) except Exception as e: - # NON blocking - notify users Redis is throwing an exception - verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e) + verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ @@ -1394,12 +1415,12 @@ class RedisCache(BaseCache): key_value_dict = {} _key_list: Final = [key for key in key_list if key is not None] start_time: Final = time.time() + admission: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") try: - swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") _keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list] results: Final = self._run_redis_mget_operation(keys=_keys) - _exit_circuit_breaker(self._circuit_breaker, swallowed_before) + _exit_circuit_breaker(self._circuit_breaker, admission) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 12e515e19a8..f4883b57fbc 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,6 +9,7 @@ import threading import time from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy +from io import BytesIO from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict, TypeVar @@ -505,6 +506,10 @@ async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> N raise MaskedHTTPStatusError(e, message=_text, text=_text) from None +class HTTPResponseLimitError(ValueError): + pass + + class MaskedHTTPStatusError(httpx.HTTPStatusError): def __init__(self, original_error, message: str | None = None, text: str | None = None): # Create a new error with the masked URL @@ -654,6 +659,7 @@ class AsyncHTTPHandler: headers: dict | None = None, follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT @@ -661,6 +667,16 @@ class AsyncHTTPHandler: params = params or {} params.update(HTTPHandler.extract_query_params(url)) + if max_response_bytes is not None: + return await self._get_with_response_limit( + url, + params=httpx.QueryParams(params), + headers=httpx.Headers(headers), + max_bytes=max_response_bytes, + follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects, + timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout), + ) + response: Final = await self.client.get( url, params=params, @@ -670,6 +686,57 @@ class AsyncHTTPHandler: ) return response + async def _get_with_response_limit( + self, + url: str, + *, + params: httpx.QueryParams, + headers: httpx.Headers, + timeout: httpx.Timeout, + max_bytes: int, + follow_redirects: bool, + ) -> httpx.Response: + request: Final = self.client.build_request( + "GET", + url, + headers=MappingProxyType({**headers, "accept-encoding": "identity"}), + params=params, + timeout=timeout, + ) + response: Final = await self.client.send(request, stream=True, follow_redirects=False) + return await self._read_with_response_limit(response, max_bytes=max_bytes, follow_redirects=follow_redirects) + + async def _read_with_response_limit( + self, response: httpx.Response, *, max_bytes: int, follow_redirects: bool, redirects_remaining: int = 10 + ) -> httpx.Response: + try: + if response.next_request is not None and follow_redirects: + if redirects_remaining == 0: + raise ValueError("Too many redirects") + await response.aclose() + following: Final = await self.client.send( + response.next_request, auth=None, stream=True, follow_redirects=False + ) + return await self._read_with_response_limit( + following, max_bytes=max_bytes, follow_redirects=True, redirects_remaining=redirects_remaining - 1 + ) + if response.is_redirect or response.is_error: + return httpx.Response(response.status_code, headers=response.headers, request=response.request) + if response.headers.get("content-encoding", "identity").lower() != "identity": + raise HTTPResponseLimitError("Response size limits require an uncompressed response") + if int(response.headers.get("content-length", "0")) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + with BytesIO() as body: + async for chunk in response.aiter_bytes(chunk_size=65536): + if body.tell() + len(chunk) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + body.write(chunk) + return httpx.Response( + response.status_code, headers=response.headers, content=body.getvalue(), request=response.request + ) + finally: + await response.aclose() + @track_llm_api_timing() async def post( self, diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f4889008e94..f3fdd54b39d 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -122,13 +122,13 @@ _USED_CODE_CACHE_PREFIX: Final = "mcp_gateway_dcr_code_used:" _USED_FLOW_CACHE_PREFIX: Final = "mcp_gateway_dcr_flow_used:" _USED_REFRESH_CACHE_PREFIX: Final = "mcp_gateway_dcr_refresh_used:" -MAX_REDIRECT_URIS: Final = 3 +MAX_REDIRECT_URIS: Final = 4 MAX_REDIRECT_URI_LENGTH: Final = 256 MAX_CLIENT_ID_LENGTH: Final = 2048 """Registration bounds. They exist to bound the sealed client_id, which rides inside -every session-token claim set: 3 URIs of 256 bytes seal to roughly 1.2KB, comfortably -under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP -Inspector register one or two redirect URIs.""" +every session-token claim set. Four 256-character ASCII URIs seal to roughly 1.5KB; +the encoded client_id is checked against its own cap before registration succeeds. +VS Code registers four callbacks for its web and desktop environments.""" MAX_STATE_LENGTH: Final = 1024 """Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index aab7f93164f..7ee4daed477 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -25,6 +25,7 @@ from collections.abc import ( ) from contextlib import asynccontextmanager from dataclasses import dataclass, replace +from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -883,6 +884,53 @@ def _sanitized_error_text(exc: Exception) -> str: return re.sub(r"https?://\S+", "", str(exc))[:200] +async def _openapi_spec_health( + spec_path: str, *, timeout: float +) -> tuple[Literal["healthy", "unhealthy", "unknown"], str | None]: + """Check specification availability, not upstream operations or user credentials.""" + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + if not spec_path.startswith(("http://", "https://")): + return "unknown", "OpenAPI servers have no protocol-level health probe" + try: + await asyncio.wait_for(load_openapi_spec_async(spec_path, max_bytes=10 * 1024 * 1024), timeout=timeout) + except asyncio.TimeoutError: + return "unhealthy", f"OpenAPI specification check timed out after {timeout} seconds" + except HTTPStatusError as exc: + return "unhealthy", f"OpenAPI specification request failed (HTTP {exc.response.status_code})" + except HTTPResponseLimitError as exc: + return "unknown", f"OpenAPI specification probe refused: {exc}" + except (httpx.RequestError, ValueError, OSError) as exc: + return "unhealthy", f"OpenAPI specification could not be loaded ({type(exc).__name__})" + return "healthy", None + + +class _OpenAPIHealthProbe: + def __init__(self, spec_path: str, clock: Callable[[], float] = time.monotonic) -> None: + self.spec_path = spec_path + self.clock = clock + self.lock = asyncio.Lock() + self.checked_at = float("-inf") + self.result: tuple[Literal["healthy", "unhealthy", "unknown"], str | None, datetime.datetime] | None = None + + async def check(self) -> tuple[Literal["healthy", "unhealthy", "unknown"], str | None, datetime.datetime]: + async with self.lock: + if self.result is not None and self.clock() - self.checked_at < 30.0: + return self.result + try: + status, error = await _openapi_spec_health(self.spec_path, timeout=MCP_HEALTH_CHECK_TIMEOUT) + except asyncio.CancelledError: + return ( + "unknown", + "OpenAPI specification check was cancelled", + datetime.datetime.now(datetime.timezone.utc), + ) + self.result = (status, error, datetime.datetime.now(datetime.timezone.utc)) + self.checked_at = self.clock() + return self.result + + def _discovery_failure_leaves_needs_unresolved( *, needs_authorization_url: bool, @@ -1749,6 +1797,7 @@ class MCPServerManager: token_exchanger=build_token_exchanger(), ) self.registry: dict[str, MCPServer] = {} + self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} """ eg. @@ -6698,6 +6747,18 @@ class MCPServerManager: last_health_check=datetime.now(), ) + if server.spec_path: + spec_status, spec_error, spec_checked_at = await self._openapi_health_probes(server.spec_path).check() + return self._build_mcp_server_table(server).model_copy( + update=MappingProxyType( + { + "status": spec_status, + "health_check_error": spec_error, + "last_health_check": spec_checked_at, + } + ) + ) + status: Literal["healthy", "unhealthy", "unknown"] = "unknown" health_check_error = None diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 16f58ef5b76..d115eb8b3c1 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -163,10 +163,14 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]: return asyncio.run(load_openapi_spec_async(filepath)) -async def load_openapi_spec_async(filepath: str) -> dict[str, Any]: +async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> dict[str, Any]: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - r: Final[httpx.Response] = await async_safe_get(client, filepath) + r: Final[httpx.Response] = ( + await async_safe_get(client, filepath) + if max_bytes is None + else await async_safe_get(client, filepath, max_response_bytes=max_bytes) + ) r.raise_for_status() return r.json() diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 4be1331e955..bc67617e444 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -1,10 +1,11 @@ import asyncio import json +import logging from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cache import RedisCache, log_redis_failure from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.types.services import ServiceTypes @@ -109,7 +110,7 @@ end ) return False except Exception as e: - verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e) + log_redis_failure(verbose_proxy_logger, logging.ERROR, f"Error acquiring Redis lock for {cronjob_id}", e) return False async def release_lock( @@ -151,7 +152,7 @@ end cronjob_id, ) except Exception as e: - verbose_proxy_logger.error("Error releasing Redis lock for %s: %s", cronjob_id, e) + log_redis_failure(verbose_proxy_logger, logging.ERROR, f"Error releasing Redis lock for {cronjob_id}", e) async def _compare_and_delete_lock(self, lock_key: str) -> int: """ diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c60888e298f..1ba90725eff 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -47,6 +47,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import ( CLI_SSO_CLAIM_MAP, CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, @@ -336,6 +337,16 @@ def _check_cli_sso_start_rate_limit( ) +def _read_cli_sso_flow(cache: DualCache, cache_key: str) -> object: + redis_cache: Final = cache.redis_cache + if redis_cache is None: + return cache.get_cache(key=cache_key) + try: + return redis_cache.get_cache(key=cache_key) + except RedisCircuitBreakerOpenError: + return None + + def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: if isinstance(login_id, str) and login_id.startswith("sk-"): raise HTTPException( @@ -348,12 +359,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: if not _is_valid_cli_sso_login_id(login_id): raise HTTPException(status_code=400, detail="Invalid CLI login session id") - cache_key: Final = _get_cli_sso_flow_cache_key(cast(str, login_id)) - redis_cache: Final = cache.redis_cache - if redis_cache is not None: - flow = redis_cache.get_cache(key=cache_key) - else: - flow = cache.get_cache(key=cache_key) + flow = _read_cli_sso_flow(cache, _get_cli_sso_flow_cache_key(cast(str, login_id))) if isinstance(flow, str): try: flow = _as_object(json.loads(flow)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0f24acb8bb4..94e74b20297 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -251,6 +251,7 @@ import litellm._redis from litellm import Router from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, @@ -3412,8 +3413,10 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme ] try: results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) - except Exception: + except Exception as e: await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending)) + if isinstance(e, RedisCircuitBreakerOpenError): + return raise for item, current_value in zip(pending, results or ()): spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 8235761ca98..686d57e2b77 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -3,12 +3,13 @@ Base class across routing strategies to abstract commmon functions like batch in """ import asyncio +import logging from abc import ABC from typing import Final from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisPipelineIncrementOperation, log_redis_failure from litellm.constants import DEFAULT_REDIS_SYNC_INTERVAL @@ -147,7 +148,7 @@ class BaseRoutingStrategy(ABC): return return_result except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) self.redis_increment_operation_queue = [] def add_to_in_memory_keys_to_update(self, key: str): diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index a8d51f95e45..3e094df7ac8 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -20,6 +20,7 @@ anthropic: import asyncio import builtins +import logging from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import Any, Final @@ -27,7 +28,7 @@ from typing import Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -92,6 +93,13 @@ class _LiteLLMParamsDictView: return dict(self._params) +async def _push_increments_to_redis(redis_cache: RedisCache, queued: list[RedisPipelineIncrementOperation]) -> None: + try: + await redis_cache.async_increment_pipeline(increment_list=queued) + except Exception as e: + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) + + class RouterBudgetLimiting(CustomLogger): def __init__( self, @@ -536,17 +544,13 @@ class RouterBudgetLimiting(CustomLogger): "Pushing Redis Increment Pipeline for queue: %s", self.redis_increment_operation_queue, ) - if len(self.redis_increment_operation_queue) > 0: - asyncio.create_task( - self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=self.redis_increment_operation_queue, - ) - ) - + queued: Final = self.redis_increment_operation_queue self.redis_increment_operation_queue = [] + if queued: + asyncio.create_task(_push_increments_to_redis(self.dual_cache.redis_cache, queued)) except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) async def _sync_in_memory_spend_with_redis(self): """ @@ -601,7 +605,7 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value) except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) def _get_budget_config_for_deployment( self, diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 38dfd143cbc..1a4764c291e 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -455,6 +455,13 @@ If 2+ reasoning markers are detected in the user message, the request is promote Reasoning markers in the system prompt do **not** trigger the reasoning override. This prevents system prompts like "Think step by step before answering" from forcing all requests to the reasoning tier. +For requests identified by a `claude-cli/` or `claude-code/` user agent, the LLM classifier omits caller system +text to avoid classifying environment, agent, and skill catalogs. The current ask, configured prior-turn context, +and trajectory signal remain unchanged. The routed completion still receives the original system text. This +also excludes genuine task constraints supplied only in Claude Code system messages. Other clients keep the +existing system-context behavior. The browser routing preview has no client-identity field and retains that +generic behavior; use the real client when checking Claude Code routing. + ### Harness Reminder Blocks Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a748bb5df27..7519f4c5156 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -48,6 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( request_contains_image_content, ) from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload +from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -1993,8 +1994,7 @@ class ComplexityRouter(CustomLogger): Args: prompt: The current user ask text (already extracted as the real human ask, not tool results) - system_prompt: The caller's system prompt (task constraints), always included so later - turns never lose it + system_prompt: Caller task constraints, omitted from classification for Claude Code requests request_kwargs: Request metadata for spend attribution messages: Full message history for extracting prior turns and the trajectory signal """ @@ -2027,9 +2027,18 @@ class ComplexityRouter(CustomLogger): ) encrypted_task: Final = _encrypted_classifier_task(request_kwargs, marker_pairs) + caller_system_prompt: Final = ( + None + if any( + is_claude_code_user_agent(user_agent) + for metadata in (self._iter_metadata_dicts(request_kwargs) if request_kwargs is not None else ()) + if isinstance(user_agent := metadata.get("user_agent"), str) + ) + else system_prompt + ) user_payload: Final = self._build_classifier_user_payload( prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt, - system_prompt=system_prompt, + system_prompt=caller_system_prompt, prior_turns=prior_turns, messages=messages, has_prior_conversation=has_prior_conversation, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index fb7003c887c..1f1b5a5cc4b 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -971,9 +971,11 @@ class ComplexityRouterConfig(BaseModel): "classified against what it refers to. Counts turns of both roles when " "classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier " "model, which may " - "be a different deployment or provider than the routed completion model; that call already " - "carries the current user ask and the caller's system prompt in full. Set to 0 to send neither " - "prior turns nor any conversation context beyond the current ask. Only applies when " + "be a different deployment or provider than the routed completion model; that call carries " + "the current user ask and, except for Claude Code requests, the extracted system-role text in full. " + "Claude Code system text is omitted to avoid classifying harness instructions; the routed " + "completion still receives it. Set to 0 to send neither prior turns nor " + "any conversation context beyond the current ask. Only applies when " "classifier_type is 'llm'." ), ) @@ -985,9 +987,9 @@ class ComplexityRouterConfig(BaseModel): "context window, per classification call. Turns are taken newest first and quoted whole " "while they fit, so a conversation small enough to quote entirely is never cut; once the " "budget runs out the older turns are dropped whole and only the turn straddling the " - "boundary is truncated, into whatever space is left. The current ask and the caller's " - "system prompt sit outside this budget and are always sent in full, as does the numbering " - "each quoted turn carries. A budget under 120 leaves no room to quote a turn and " + "boundary is truncated, into whatever space is left. The current ask and, except for Claude " + "Code requests, the extracted system-role text sit outside this budget and are sent in full, as does " + "the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and " "suppresses the block; set classifier_context_window_size to 0 to turn context off " "deliberately. Only applies when classifier_type is 'llm'." ), diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 22d816e13e9..c8ca7105392 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from litellm import verbose_logger from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -27,6 +28,16 @@ class DeploymentHealthStateValue(TypedDict): reason: str +def _read_shared_health_snapshot(cache: DualCache, key: str) -> object: + redis_cache: Final = cache.redis_cache + if redis_cache is None: + return None + try: + return redis_cache.get_cache(key) + except RedisCircuitBreakerOpenError: + return None + + class DeploymentHealthCache: """ Cache for deployment health states produced by background health checks. @@ -50,13 +61,12 @@ class DeploymentHealthCache: coexist on the one shared entry without erasing each other's results. The snapshot is read from Redis when available, since a pod-local read would only ever see this writer's own previous merge. When the Redis - read comes back empty (a miss, or a swallowed connection error), the - pod-local copy of the last merge is used so peers are not erased. + read comes back empty (a miss, a swallowed connection error, or a read + refused by the open circuit breaker), the pod-local copy of the last + merge is used so peers are not erased. """ try: - redis_raw: Final = ( - self.cache.redis_cache.get_cache(self.CACHE_KEY) if self.cache.redis_cache is not None else None - ) + redis_raw: Final = _read_shared_health_snapshot(self.cache, self.CACHE_KEY) raw: Final = redis_raw if isinstance(redis_raw, dict) else self.cache.get_cache(key=self.CACHE_KEY) existing: Final = raw if isinstance(raw, dict) else {} expiry_seconds: Final = self.staleness_threshold * 1.5 diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index eae8bbfdaff..4c9068722b8 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -677,3 +677,30 @@ async def test_a_real_redis_failure_still_logs_an_error(caplog): errors = [record for record in caplog.records if record.levelno == logging.ERROR] assert [record.getMessage() for record in errors] == ["LiteLLM Cache: exception in async_get_cache: redis is down"] assert errors[0].exc_info is not None + + +def _dual_cache_with_open_breaker_and_a_memory_hit() -> DualCache: + in_memory = InMemoryCache() + in_memory.set_cache("k1", "v1") + return DualCache(in_memory_cache=in_memory, redis_cache=_OpenBreakerRedis(), default_redis_batch_cache_expiry=10) # pyright: ignore[reportArgumentType] # duck-typed Redis double + + +def test_open_breaker_keeps_sync_batch_read_memory_hits_and_releases_reservations(): + """A refused Redis batch read must still answer with the in-memory hits and hold no reservation. + + The refusal was logged and turned into a bare None, so a caller lost its in-memory hits + for as long as the breaker stayed open, and the reserved keys stayed throttled until + the batch expiry passed even though nothing was ever read for them. + """ + cache = _dual_cache_with_open_breaker_and_a_memory_hit() + + assert list(cache.batch_get_cache(["k1", "k2"])) == ["v1", None] + assert "k2" not in cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_open_breaker_keeps_async_batch_read_memory_hits_and_releases_reservations(): + cache = _dual_cache_with_open_breaker_and_a_memory_hit() + + assert list(await cache.async_batch_get_cache(["k1", "k2"])) == ["v1", None] + assert "k2" not in cache.last_redis_batch_access_time diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index bcaa58c9c40..bcae33b976e 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1,11 +1,12 @@ import asyncio +import time from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm._service_logger import ServiceLogging -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError @pytest.fixture @@ -515,14 +516,46 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(call_met await call_method(cache) -def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_redis_cache): - """An open breaker must preserve the sync batch read's dictionary fallback.""" +def test_circuit_breaker_open_makes_sync_batch_get_cache_fast_fail(sync_batch_redis_cache, caplog): + """Once the breaker is open the sync batch read refuses with the typed error instead of a miss. + + Swallowing the refusal into `{}` made every sync batch read on an open breaker emit an ERROR + log and a service failure event per call, and the DualCache caller could not tell the + refusal from a dead Redis, so it dropped its in-memory hits too. + """ from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} - assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} + caplog.clear() + with caplog.at_level("INFO"): + with pytest.raises(RedisCircuitBreakerOpenError): + sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) + sync_batch_redis_cache.redis_client.mget.assert_called() + assert caplog.records == [] + + +def test_sync_get_cache_failure_feeds_the_breaker_and_logs_a_well_formed_record(sync_batch_redis_cache, caplog): + """The sync get path swallowed its Redis error without recording it, and its log call was malformed. + + `verbose_logger.error("...: ", e)` passes the exception as a format argument to a message + with no placeholder, so the record carried no error text. Nothing fed the breaker either, + so a dead Redis read through this path never opened it. + """ + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + + sync_batch_redis_cache.redis_client.get.side_effect = OSError("redis unavailable") + + with caplog.at_level("ERROR"): + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert sync_batch_redis_cache.get_cache("lit7468") is None + + assert all("redis unavailable" in record.getMessage() for record in caplog.records) + assert len(caplog.records) == REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + assert sync_batch_redis_cache._circuit_breaker.is_open() is True + with pytest.raises(RedisCircuitBreakerOpenError): + sync_batch_redis_cache.get_cache("lit7468") def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): @@ -661,7 +694,8 @@ def test_sync_batch_get_cache_survives_a_service_callback_that_raises( with ThreadPoolExecutor(max_workers=1) as pool: assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {} - assert cache.batch_get_cache(key_list=["lit6729"]) == {} + with pytest.raises(RedisCircuitBreakerOpenError): + cache.batch_get_cache(key_list=["lit6729"]) def test_call_stack_info_skips_breaker_guard_frames(): @@ -1010,6 +1044,8 @@ async def test_breaker_metrics_track_state_and_failure_class(): 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._opened_at = time.time() - 9999 + assert breaker.is_open() is False 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 @@ -1030,3 +1066,139 @@ def test_sync_guard_counts_a_timeout_as_a_timeout(): _run_under_circuit_breaker_sync(breaker, "op", timing_out_call) assert breaker.is_open() is False + + +def test_success_admitted_before_the_breaker_opened_cannot_close_it(): + """A stale in-flight success must not close a breaker that opened while it ran. + + Calls admitted while the breaker was still closed finish after later failures opened it. + Recording their success unconditionally closed the breaker again, skipping the recovery + timeout and the single half-open probe, so the breaker flapped between open and closed + on every straggler while Redis was still down. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + breaker.record_failure() + assert breaker._state == breaker.OPEN + + breaker.record_success() + + assert breaker._state == breaker.OPEN + assert breaker.is_open() is True + + +def test_recovery_probe_still_closes_the_breaker(): + from litellm.caching.redis_cache import RedisCircuitBreaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + breaker.record_failure() + breaker._opened_at = time.time() - 9999 + assert breaker.is_open() is False + assert breaker._state == breaker.HALF_OPEN + + breaker.record_success() + + assert breaker._state == breaker.CLOSED + assert breaker.is_open() is False + + +@pytest.mark.asyncio +async def test_stale_success_during_the_recovery_probe_leaves_the_breaker_to_the_probe(): + """A call admitted before the trip that finishes while HALF_OPEN must not close the breaker. + + Only the one call designated as the recovery probe has actually reached Redis after the + outage, so closing on the straggler's success resumed full Redis traffic before the probe + had proven anything. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + stale_admitted = asyncio.Event() + stale_release = asyncio.Event() + probe_admitted = asyncio.Event() + probe_release = asyncio.Event() + + async def stale_call() -> str: + stale_admitted.set() + await stale_release.wait() + return "stale" + + async def probe_call() -> str: + probe_admitted.set() + await probe_release.wait() + return "probe" + + stale = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", stale_call)) + await stale_admitted.wait() + for _ in range(3): + breaker.record_failure() + assert breaker._state == breaker.OPEN + breaker._opened_at = time.time() - 9999 + probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", probe_call)) + await probe_admitted.wait() + assert breaker._state == breaker.HALF_OPEN + + stale_release.set() + assert await stale == "stale" + + assert breaker._state == breaker.HALF_OPEN, "the straggler must not close the breaker for the probe" + assert breaker.is_open() is True + + probe_release.set() + assert await probe == "probe" + + assert breaker._state == breaker.CLOSED + assert breaker.is_open() is False + + +@pytest.mark.asyncio +async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new_probe(): + """A probe still in flight when a late failure reopens the breaker must not close it for the next probe. + + Once the breaker has reopened, only the probe admitted after that outage has reached + Redis, so the older probe's success no longer says anything about whether Redis recovered. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + old_probe_admitted = asyncio.Event() + old_probe_release = asyncio.Event() + new_probe_admitted = asyncio.Event() + new_probe_release = asyncio.Event() + + async def old_probe_call() -> str: + old_probe_admitted.set() + await old_probe_release.wait() + return "old probe" + + async def new_probe_call() -> str: + new_probe_admitted.set() + await new_probe_release.wait() + return "new probe" + + for _ in range(3): + breaker.record_failure() + breaker._opened_at = time.time() - 9999 + old_probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", old_probe_call)) + await old_probe_admitted.wait() + assert breaker._state == breaker.HALF_OPEN + + breaker.record_failure() + assert breaker._state == breaker.OPEN + breaker._opened_at = time.time() - 9999 + new_probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", new_probe_call)) + await new_probe_admitted.wait() + assert breaker._state == breaker.HALF_OPEN + + old_probe_release.set() + assert await old_probe == "old probe" + + assert breaker._state == breaker.HALF_OPEN, "the overtaken probe must not close the breaker for the new probe" + assert breaker.is_open() is True + + new_probe_release.set() + assert await new_probe == "new probe" + assert breaker._state == breaker.CLOSED diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 3d4ba264c1c..f8868cfaf83 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1612,3 +1612,66 @@ async def test_a_retried_put_stays_a_put_and_still_refuses_redirects(): assert attempts == [("PUT", "/first"), ("PUT", "/first")] finally: await handler.client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", ["https://example.com/final.json?next=1", "https://other.example/final.json?next=1"]) +async def test_bounded_get_preserves_sdk_redirect_auth_and_query_handling(respx_mock, monkeypatch, target): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + respx_mock.get("https://example.com/spec.json?original=1").respond(302, headers={"location": target}) + destination = respx_mock.get(target).respond(200, json={"paths": {}}) + handler = AsyncHTTPHandler() + try: + response = await handler.get( + "https://example.com/spec.json?original=1", max_response_bytes=100, follow_redirects=True, + headers={"Authorization": "Bearer sentinel", "Accept-Encoding": "gzip"}, timeout=2.0, + ) + finally: + await handler.close() + assert response.json() == {"paths": {}} + request = destination.calls[0].request + assert request.headers.get("authorization") == (None if "other.example" in target else "Bearer sentinel") + assert request.headers["accept-encoding"] == "identity" + assert str(request.url) == target + assert request.extensions["timeout"]["read"] == 2.0 + + +@pytest.mark.asyncio +async def test_bounded_get_stops_redirect_loops(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route = respx_mock.get("https://example.com/spec.json").respond(302, headers={"location": "/spec.json"}) + handler = AsyncHTTPHandler() + try: + with pytest.raises(ValueError, match="Too many redirects"): + await handler.get("https://example.com/spec.json", max_response_bytes=100, follow_redirects=True) + finally: + await handler.close() + assert route.call_count == 11 + + +@pytest.mark.asyncio +async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + started = asyncio.Event() + closed = asyncio.Event() + + class SlowStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"x" + started.set() + await asyncio.Event().wait() + + async def aclose(self): + closed.set() + + respx_mock.get("https://example.com/slow.json").respond(200, stream=SlowStream()) + handler = AsyncHTTPHandler() + try: + task = asyncio.create_task(handler.get("https://example.com/slow.json", max_response_bytes=100)) + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + await handler.close() + assert closed.is_set() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 8aa4cfc5619..7c80ee77cd7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -6,6 +6,7 @@ import re from base64 import urlsafe_b64encode from datetime import datetime, timedelta, timezone from http.cookies import SimpleCookie +from typing import Final from urllib.parse import parse_qs, urlparse import pytest @@ -18,6 +19,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, + MAX_CLIENT_ID_LENGTH, ConsentTeam, MintedProxyCredential, _GatewayAuthCode, @@ -53,6 +55,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" +VSCODE_REDIRECT_URIS: Final = ( + "https://insiders.vscode.dev/redirect", + "https://vscode.dev/redirect", + "http://127.0.0.1/", + "http://127.0.0.1:33418/", +) +MAX_LENGTH_REDIRECT_URIS: Final = tuple(f"https://client.example/{index}/".ljust(256, "a") for index in range(4)) CODE_VERIFIER = "verifier-" + "v" * 43 CODE_CHALLENGE = urlsafe_b64encode(hashlib.sha256(CODE_VERIFIER.encode("ascii")).digest()).rstrip(b"=").decode("ascii") @@ -104,6 +113,58 @@ async def test_register_mints_stateless_public_client(): assert record.redirect_uris == (REDIRECT_URI,) +@pytest.mark.asyncio +@pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) +async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={ + "client_name": "Visual Studio Code", + "client_uri": "https://code.visualstudio.com", + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "redirect_uris": list(redirect_uris), + "token_endpoint_auth_method": "none", + "application_type": "native", + }, + ) + assert response.status_code == 201 + body: Final = json.loads(response.body) + assert body["redirect_uris"] == list(redirect_uris) + assert body["token_endpoint_auth_method"] == "none" + assert "client_secret" not in body + assert len(body["client_id"]) <= MAX_CLIENT_ID_LENGTH + record: Final = open_gateway_dcr_client(body["client_id"]) + assert record is not None + assert record.redirect_uris == redirect_uris + + +@pytest.mark.asyncio +async def test_register_rejects_five_valid_callbacks() -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]}, + ) + assert response.status_code == 400 + assert json.loads(response.body) == { + "error": "invalid_redirect_uri", + "error_description": "redirect_uris must be a list of 1 to 4 URIs", + } + + +@pytest.mark.asyncio +async def test_register_four_callbacks_preserves_encoded_size_guard() -> None: + response: Final = await register_aggregate_client( + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]}, + ) + assert response.status_code == 400 + assert json.loads(response.body) == { + "error": "invalid_client_metadata", + "error_description": "registered metadata is too large", + } + + @pytest.mark.asyncio async def test_register_allows_loopback_http_for_dev_clients(): body = await _register(["http://localhost:6274/oauth/callback"]) @@ -162,7 +223,6 @@ async def test_register_rejects_userinfo_spoofed_origin(): ["https://claude.ai/cb#fragment"], ["ftp://claude.ai/cb"], ["https://a.example.com/" + "p" * 300], - ["https://a.example.com/1", "https://a.example.com/2", "https://a.example.com/3", "https://a.example.com/4"], [12345], ], ) @@ -248,12 +308,14 @@ def _flow_cookie_from(response) -> tuple: @pytest.mark.asyncio -async def test_full_walk_register_authorize_complete_token_and_replay(): +@pytest.mark.parametrize("redirect_uris", [(REDIRECT_URI,), VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) +async def test_full_walk_register_authorize_complete_token_and_replay(redirect_uris: tuple[str, ...]): """The whole front door on one deterministic walk: register -> authorize -> complete -> token, then the security edges on the same artifacts (user mismatch, PKCE mismatch, single-use replay, refresh rotation, cross-client refresh).""" - client_id = (await _register([REDIRECT_URI]))["client_id"] - authorize_response = _authorize(client_id, session_user_id="u1") + redirect_uri: Final = redirect_uris[-1] + client_id = (await _register(list(redirect_uris)))["client_id"] + authorize_response = _authorize(client_id, session_user_id="u1", redirect_uri=redirect_uri) handle, cookies = _flow_cookie_from(authorize_response) denied = await complete_connect_flow( @@ -280,7 +342,7 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): ) assert completed.status_code == 303 redirect = urlparse(completed.headers["location"]) - assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == REDIRECT_URI + assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == redirect_uri params = parse_qs(redirect.query) assert params["state"] == ["client-state-123"] code = params["code"][0] @@ -293,7 +355,7 @@ async def test_full_walk_register_authorize_complete_token_and_replay(): "request": _request("/token", method="POST"), "grant_type": "authorization_code", "code": code, - "redirect_uri": REDIRECT_URI, + "redirect_uri": redirect_uri, "client_id": client_id, "code_verifier": CODE_VERIFIER, "refresh_token": None, 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 f1675b72f77..105c78867ad 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 @@ -4473,6 +4473,116 @@ class TestMCPServerManager: assert len(result) == 1 assert result[0].name == "github_tool_1" + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2]) + @pytest.mark.parametrize("is_byok", [False, True]) + @pytest.mark.parametrize("scheme", ["http", "https"]) + async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="openapi-health", + name="openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path=f"{scheme}://93.184.216.34/openapi.json", + auth_type=auth_type, + is_byok=is_byok, + authentication_token=None if is_byok else "shared-secret", + static_headers={"Authorization": "Bearer static-secret"}, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).respond(200, json={"openapi": "3.0.0", "paths": {}}) + result = await manager.health_check_server(server.server_id, mcp_auth_header="caller-secret") + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None + assert result.spec_path == server.spec_path + assert route.call_count == 1 + assert "authorization" not in route.calls[0].request.headers + assert "x-api-key" not in route.calls[0].request.headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token]) + @pytest.mark.parametrize("spec_path", ["/config/openapi.json", "relative/openapi.json"]) + async def test_openapi_local_spec_health_is_unknown(self, respx_mock, auth_type, spec_path): + manager = MCPServerManager() + server = MCPServer( + server_id="local-openapi-health", + name="local-openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path=spec_path, + auth_type=auth_type, + is_byok=True, + ) + manager.registry = {server.server_id: server} + result = await manager.health_check_server(server.server_id) + assert result.status == "unknown" + assert result.health_check_error == "OpenAPI servers have no protocol-level health probe" + assert result.last_health_check is not None + assert not respx_mock.calls + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("failure", "expected_status", "expected_error"), + [ + (httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"), + (httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"), + (httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"), + (httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"), + (httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"), + ], + ) + async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="failed-openapi-health", + name="failed-openapi-health", + transport=MCPTransport.http, + url="https://rest.example.com", + spec_path="https://93.184.216.34/key-secret?token=query-secret", + auth_type=MCPAuth.bearer_token, + is_byok=True, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).mock(side_effect=[failure]) + result = await manager.health_check_server(server.server_id) + assert result.status == expected_status + assert result.health_check_error == expected_error + assert result.last_health_check is not None + assert route.call_count == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancel", [False, True]) + async def test_openapi_health_timeout_and_cancellation_cleanup(self, respx_mock, monkeypatch, cancel): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_spec_health + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + started = asyncio.Event() + cancelled = asyncio.Event() + + async def slow_load(request): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + respx_mock.get("https://93.184.216.34/slow.json").mock(side_effect=slow_load) + task = asyncio.create_task(_openapi_spec_health("https://93.184.216.34/slow.json", timeout=0.1)) + await asyncio.wait_for(started.wait(), timeout=1) + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + status, error = await task + assert status == "unhealthy" + assert error == "OpenAPI specification check timed out after 0.1 seconds" + assert cancelled.is_set() + @pytest.mark.asyncio async def test_health_check_server_healthy(self): """Test health check for a healthy server""" @@ -12777,3 +12887,118 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li assert "Credential=AKIDEXAMPLE/" in request.headers["Authorization"] finally: request_ctx.reset(token) + +@pytest.mark.asyncio +async def test_openapi_health_coalesces_concurrent_checks_and_reuses_results(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="coalesced", + name="coalesced", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/coalesced.json", + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + started = asyncio.Event() + release = asyncio.Event() + + async def serve(request): + started.set() + await release.wait() + return httpx.Response(200, json={"paths": {}}) + + route = respx_mock.get(server.spec_path).mock(side_effect=serve) + tasks = [asyncio.create_task(manager.health_check_server(server.server_id)) for _ in range(4)] + await asyncio.wait_for(started.wait(), timeout=1) + release.set() + results = await asyncio.gather(*tasks) + cached = await manager.health_check_server(server.server_id) + assert [result.status for result in results] == ["healthy"] * 4 + assert cached.status == "healthy" + assert {result.last_health_check for result in [*results, cached]} == {results[0].last_health_check} + assert route.call_count == 1 + + +@pytest.mark.asyncio +async def test_openapi_health_cache_expires_at_thirty_seconds(respx_mock, monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _OpenAPIHealthProbe + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + clock = iter([0.0, 29.0, 30.0, 30.0]) + probe = _OpenAPIHealthProbe("https://93.184.216.34/expiry.json", clock=clock.__next__) + route = respx_mock.get(probe.spec_path).mock( + side_effect=[ + httpx.Response(200, json={"paths": {}}), + httpx.Response(503), + ] + ) + first = await probe.check() + assert first[0] == "healthy" + assert await probe.check() == first + refreshed = await probe.check() + assert refreshed[0] == "unhealthy" + assert refreshed[1] == "OpenAPI specification request failed (HTTP 503)" + assert refreshed[2] >= first[2] + assert route.call_count == 2 + + +@pytest.mark.asyncio +async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="oversized", + name="oversized", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/large.json", + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + route = respx_mock.get(server.spec_path).respond(200, headers={"content-length": str(12 * 1024 * 1024)}) + result = await manager.health_check_server(server.server_id) + cached = await manager.health_check_server(server.server_id) + assert result.status == "unknown" + assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" + assert cached.health_check_error == result.health_check_error + assert cached.last_health_check == result.last_health_check + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("already_waiting", [False, True]) +async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, monkeypatch, already_waiting): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager = MCPServerManager() + server = MCPServer( + server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http, + spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} + started = asyncio.Event() + attempts = [] + + async def serve(request): + attempts.append(request.url) + if not started.is_set(): + started.set() + await asyncio.Event().wait() + return httpx.Response(200, json={"paths": {}}) + + route = respx_mock.get(server.spec_path).mock(side_effect=serve) + leader = asyncio.create_task(manager.health_check_server(server.server_id)) + await asyncio.wait_for(started.wait(), timeout=1) + follower = asyncio.create_task(manager.health_check_server(server.server_id)) if already_waiting else None + await asyncio.sleep(0) + leader.cancel() + cancelled = await leader + assert cancelled.status == "unknown" + assert cancelled.health_check_error == "OpenAPI specification check was cancelled" + recovered = await follower if follower is not None else await manager.health_check_server(server.server_id) + assert recovered.status == "healthy" + assert recovered.health_check_error is None + cached = await manager.health_check_server(server.server_id) + assert cached.last_health_check == recovered.last_health_check + assert cached.status == "healthy" + assert len(attempts) == 2 + assert route.call_count == 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index e59616e53c1..5fa202224e3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1378,3 +1378,83 @@ class TestUpstreamStatusIsClassified: assert exc.value.status_code == status_code assert secret_body not in str(exc.value) assert str(exc.value) == f"upstream returned HTTP {status_code}" + +class TestBoundedOpenAPISpecLoading: + @pytest.mark.asyncio + @pytest.mark.parametrize("max_bytes", [12, 13]) + async def test_exact_size_and_smaller_specs_load(self, respx_mock, monkeypatch, max_bytes): + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route = respx_mock.get("https://93.184.216.34/spec.json").respond(200, content=b'{"paths":{}}') + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=max_bytes) == {"paths": {}} + assert route.calls[0].request.headers["accept-encoding"] == "identity" + + @pytest.mark.asyncio + @pytest.mark.parametrize("headers", [{"content-length": "1000000"}, {"content-encoding": "gzip"}]) + async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers): + import httpx + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + load_openapi_spec_async, + ) + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + closed = [] + + class UnreadableStream(httpx.AsyncByteStream): + async def __aiter__(self): + pytest.fail("Oversized or compressed response must not be consumed") + yield b"" + + async def aclose(self): + closed.append(True) + + respx_mock.get("https://93.184.216.34/spec.json").respond(200, headers=headers, stream=UnreadableStream()) + with pytest.raises(HTTPResponseLimitError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12) + assert closed == [True] + + @pytest.mark.asyncio + async def test_chunked_response_is_bounded_and_closed(self, respx_mock, monkeypatch): + import httpx + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + load_openapi_spec_async, + ) + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + consumed = [] + closed = [] + + class ChunkedStream(httpx.AsyncByteStream): + async def __aiter__(self): + for index in range(10): + consumed.append(index) + yield b"x" * 65536 + + async def aclose(self): + closed.append(True) + + respx_mock.get("https://93.184.216.34/spec.json").respond(200, stream=ChunkedStream()) + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=65536) + assert consumed == [0, 1] + assert closed == [True] + + @pytest.mark.asyncio + @pytest.mark.parametrize("target", ["https://93.184.216.35/final.json", "http://127.0.0.1/private.json"]) + async def test_bounded_spec_redirects_preserve_ssrf_protection(self, respx_mock, monkeypatch, target): + from litellm.litellm_core_utils.url_utils import SSRFError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + respx_mock.get("https://93.184.216.34/spec.json").respond(302, headers={"location": target}) + destination = respx_mock.get(target).respond(200, json={"paths": {}}) + if "127.0.0.1" in target: + with pytest.raises(SSRFError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) + assert not destination.called + else: + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}} + assert destination.call_count == 1 diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index ecd5c5f50c0..4684c3213d6 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -1,4 +1,5 @@ import json +import logging from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -6,6 +7,7 @@ import pytest from fastapi.testclient import TestClient +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager @@ -215,6 +217,22 @@ async def test_redis_error_handling(pod_lock_manager, mock_redis): ) +@pytest.mark.asyncio +async def test_lock_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(pod_lock_manager, mock_redis, caplog): + """Every cron job retries its lock on a timer, so an open breaker must not add an error line per cycle.""" + refused = RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping async_set_cache") + mock_redis.async_set_cache.side_effect = refused + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id + mock_redis.async_delete_cache.side_effect = refused + + with caplog.at_level(logging.ERROR): + acquired = await pod_lock_manager.acquire_lock(cronjob_id="test_job") + await pod_lock_manager.release_lock(cronjob_id="test_job") + + assert acquired is False + assert caplog.records == [] + + @pytest.mark.asyncio async def test_bytes_handling(pod_lock_manager, mock_redis): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 8d8bc15f9be..2050e65d2a1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2616,6 +2616,28 @@ class TestCLIKeyRegenerationFlow: ) cache.set_cache.assert_not_called() + def test_cli_sso_flow_lookup_treats_an_open_redis_breaker_as_a_miss(self): + """A Redis read refused by the open circuit breaker is a missing session, not a server error. + + The direct Redis read is what keeps the flow authoritative across workers, so the + refusal must not fall back to a possibly stale in-memory copy either. + """ + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError + from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_or_raise + + redis_cache = MagicMock() + redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError("Redis circuit breaker is open") + cache = MagicMock() + cache.redis_cache = redis_cache + cache.get_cache.return_value = {"poll_secret_hash": "stale", "sso_complete": False} + + with pytest.raises(HTTPException) as exc_info: + _get_cli_sso_flow_or_raise(login_id="cli-breaker_open_1234567890", cache=cache) + + assert exc_info.value.status_code == 400 + assert "not found or expired" in exc_info.value.detail + cache.get_cache.assert_not_called() + def test_cli_sso_flow_with_enum_survives_redis_round_trip(self): """ RedisCache stores values via str(value) and reads them back through diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index c343652efd9..2f47736a398 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1146,6 +1146,54 @@ async def test_prepare_window_spend_counter_increment_missing_window_start_inval assert fake_cache.redis_cache.async_increment.called is False +# --------------------------------------------------------------------------- +# _apply_spend_counter_increments +# --------------------------------------------------------------------------- + + +def _two_pending_increments() -> tuple[ps._PendingSpendIncrement, ...]: + return ( + ps._PendingSpendIncrement(counter_key="spend:key:k", increment=1.5), + ps._PendingSpendIncrement(counter_key="spend:team:t", increment=1.5), + ) + + +@pytest.mark.asyncio +async def test_apply_spend_counter_increments_open_breaker_invalidates_and_returns(monkeypatch): + """An open Redis circuit breaker is a known, already-logged state, not a per-request tracking failure. + + Re-raising the refusal sent every request through the cost callback's error path, which + logged an ERROR and fired the failed-tracking alert once per request for the whole outage. + """ + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_increment_pipeline = AsyncMock( + side_effect=RedisCircuitBreakerOpenError("Redis circuit breaker is open") + ) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + await ps._apply_spend_counter_increments(_two_pending_increments()) + + deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) + assert deleted_keys == ["spend:key:k", "spend:team:t"] + fake_cache.in_memory_cache.set_cache.assert_not_called() + + +@pytest.mark.asyncio +async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_raises(monkeypatch): + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + with pytest.raises(ConnectionError, match="redis down"): + await ps._apply_spend_counter_increments(_two_pending_increments()) + + deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) + assert deleted_keys == ["spend:key:k", "spend:team:t"] + fake_cache.in_memory_cache.set_cache.assert_not_called() + + # --------------------------------------------------------------------------- # _ensure_spend_counter_initialized # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_base_routing_strategy.py b/tests/test_litellm/router_strategy/test_base_routing_strategy.py index 154042692d0..dc75c3d2916 100644 --- a/tests/test_litellm/router_strategy/test_base_routing_strategy.py +++ b/tests/test_litellm/router_strategy/test_base_routing_strategy.py @@ -1,4 +1,5 @@ import json +import logging from typing import Any, Dict, List, Optional, Set, Union import pytest @@ -9,7 +10,7 @@ from unittest.mock import MagicMock, patch from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -146,3 +147,18 @@ async def test_cache_keys_management(base_strategy): # Test resetting cache keys base_strategy.reset_in_memory_keys_to_update() assert len(base_strategy.get_in_memory_keys_to_update()) == 0 + + +@pytest.mark.asyncio +async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(base_strategy, mock_dual_cache, caplog): + """The sync loop pushes every 100 ms under usage-based routing, so an open breaker must not add an error line per cycle.""" + mock_dual_cache.redis_cache.async_increment_pipeline.side_effect = RedisCircuitBreakerOpenError( + "Redis circuit breaker is open - skipping async_increment_pipeline" + ) + base_strategy.redis_increment_operation_queue = [{"key": "k", "increment_value": 1.0, "ttl": 60}] + + with caplog.at_level(logging.ERROR): + await base_strategy._push_in_memory_increments_to_redis() + + assert caplog.records == [] + assert base_strategy.redis_increment_operation_queue == [] diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 36fa38bacb5..4cc8fe78811 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -1,7 +1,13 @@ +import asyncio +import gc +import logging +from unittest.mock import AsyncMock, MagicMock + import pytest import litellm from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.router import LiteLLM_Params from litellm.types.utils import BudgetConfig @@ -303,3 +309,90 @@ def test_router_add_deployment_registers_deployment_budget( ) assert config is not None assert config.max_budget == 0.000000000001 + + +@pytest.mark.asyncio +async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_task(disable_budget_sync, caplog): + """The budget sync runs every second, so an open breaker must not add an error line or an unretrieved task exception per cycle.""" + refused = RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping async_increment_pipeline") + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=refused) + redis_cache.async_batch_get_cache = AsyncMock(side_effect=refused) + limiter = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": BudgetConfig(max_budget=1.0, budget_duration="1d")}, + ) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + limiter.redis_increment_operation_queue = [{"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60}] + loop = asyncio.get_running_loop() + unretrieved = MagicMock() + loop.set_exception_handler(unretrieved) + + try: + with caplog.at_level(logging.ERROR): + await limiter._sync_in_memory_spend_with_redis() + await asyncio.sleep(0) + gc.collect() + finally: + loop.set_exception_handler(None) + + assert caplog.records == [] + unretrieved.assert_not_called() + assert limiter.redis_increment_operation_queue == [] + assert redis_cache.async_increment_pipeline.await_count == 1 + + +async def _limiter_with_redis(redis_cache: MagicMock) -> RouterBudgetLimiting: + limiter = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": BudgetConfig(max_budget=1.0, budget_duration="1d")}, + ) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + limiter.redis_increment_operation_queue = [{"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60}] + return limiter + + +@pytest.mark.asyncio +async def test_push_returns_before_redis_answers(disable_budget_sync): + """The push runs inside the request success callback, so it must hand the Redis round trip to a task instead of waiting on it.""" + redis_answered = asyncio.Event() + + async def wait_for_redis(**_: object) -> None: + await redis_answered.wait() + + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=wait_for_redis) + limiter = await _limiter_with_redis(redis_cache) + + await asyncio.wait_for(limiter._push_in_memory_increments_to_redis(), timeout=1) + await asyncio.sleep(0) + + assert not redis_answered.is_set() + assert redis_cache.async_increment_pipeline.await_count == 1 + assert limiter.redis_increment_operation_queue == [] + redis_answered.set() + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + + +@pytest.mark.asyncio +async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sync, caplog): + """A real Redis failure on the background push must surface as one error line, never as an unretrieved task exception.""" + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("Error 61 connecting to 127.0.0.1:6379")) + limiter = await _limiter_with_redis(redis_cache) + loop = asyncio.get_running_loop() + unretrieved = MagicMock() + loop.set_exception_handler(unretrieved) + + try: + with caplog.at_level(logging.ERROR): + await limiter._push_in_memory_increments_to_redis() + await asyncio.sleep(0) + gc.collect() + finally: + loop.set_exception_handler(None) + + assert [record.getMessage() for record in caplog.records] == [ + "Error syncing in-memory cache with Redis: Error 61 connecting to 127.0.0.1:6379" + ] + unretrieved.assert_not_called() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index b8e728c44a9..dddaee71a63 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2676,6 +2676,26 @@ class TestEncryptedTaskClassifier: assert "source-secret" not in json.dumps(call) assert "originating_request_masked" not in call["proxy_server_request"]["body"] + @pytest.mark.asyncio + async def test_claude_code_encrypted_task_omits_caller_instructions(self): + router, dependency = _native_classifier_router() + task: Final = _encrypted_agent_task() + request: Final = { + "input": [task], + "instructions": "CLAUDE_CODE_SYSTEM", + "litellm_metadata": {"user_agent": "claude-cli/2.1.233"}, + } + original: Final = deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert result.routing_decision["cause"] == "llm_classifier" + assert request == original + call: Final = dependency.aresponses.call_args.kwargs + assert call["instructions"] == classification_system_prompt(router.config.classifier_context_window_size) + assert "CLAUDE_CODE_SYSTEM" not in json.dumps(call["input"][:-1]) + assert call["input"][-1] == task + @pytest.mark.asyncio @pytest.mark.parametrize( "items", @@ -7994,6 +8014,113 @@ _CODEX_ENVELOPES: Final = ( class TestContextAwareClassifier: """Test the new classifier context window and trajectory signals.""" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_metadata,forwards_system", + [ + ({"metadata": {"user_agent": "claude-cli/2.1.233"}}, False), + ({"litellm_metadata": {"user_agent": "claude-code/2.1.233"}}, False), + ({"metadata": {"user_agent": "curl/8.7.1"}}, True), + ({"litellm_metadata": {}}, True), + ( + {"metadata": {"user_agent": "claude-cli/2.1.233"}, "litellm_metadata": {"user_agent": "curl/8.7.1"}}, + False, + ), + ({"metadata": {"user_agent": "Claude-Code/2.1.233"}}, True), + ], + ) + async def test_claude_code_classifier_omits_harness_system_prompt( + self, + llm_classifier_config: dict[str, object], + request_metadata: dict[str, object], + forwards_system: bool, + ) -> None: + dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}'))) + router: Final = ComplexityRouter( + "test-complexity-router", + dependency, + { + **llm_classifier_config, + "classifier_context_include_assistant_turns": True, + }, + ) + messages: Final = [ + {"role": "user", "content": "Design the retry state machine"}, + {"role": "assistant", "content": "The design needs a lease and fencing token"}, + {"role": "user", "content": "Now prove it cannot livelock"}, + { + "role": "system", + "content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}], + }, + ] + top_level_system: Final = [{"type": "text", "text": "TOP_LEVEL_HARNESS_SYSTEM"}] + claude_kwargs: Final = { + "metadata": {"user_agent": "claude-cli/2.1.233"}, + "system": top_level_system, + "proxy_server_request": {"body": {"system": top_level_system}}, + } + compared_kwargs: Final = { + **request_metadata, + "system": top_level_system, + "proxy_server_request": {"body": {"system": top_level_system}}, + } + original_messages: Final = deepcopy(messages) + original_kwargs: Final = deepcopy((claude_kwargs, compared_kwargs)) + results: Final = ( + await router.async_pre_routing_hook("test-complexity-router", claude_kwargs, messages), + await router.async_pre_routing_hook("test-complexity-router", compared_kwargs, messages), + ) + + assert all(result is not None and result.routing_decision["cause"] == "llm_classifier" for result in results) + assert all(result is not None and result.messages == original_messages for result in results) + assert messages == original_messages + assert (claude_kwargs, compared_kwargs) == original_kwargs + calls: Final = tuple(call.kwargs["messages"] for call in dependency.acompletion.await_args_list) + assert calls[0][0]["content"] == calls[1][0]["content"] == classification_system_prompt( + router.config.classifier_context_window_size + ) + payloads: Final = (calls[0][1]["content"], calls[1][1]["content"]) + for payload, expected_system in zip(payloads, (False, forwards_system)): + assert payload.endswith("Classify this message:\nNow prove it cannot livelock") + assert ("ENVIRONMENT_CATALOG" in payload) is expected_system + assert ("AGENT_CATALOG" in payload) is expected_system + assert ("SKILL_CATALOG" in payload) is expected_system + assert "Design the retry state machine" in payload + assert "lease and fencing token" in payload + assert "TOP_LEVEL_HARNESS_SYSTEM" not in payload + assert "Conversation so far: ~35 tokens across the request" in payload + + @pytest.mark.asyncio + async def test_claude_code_first_turn_without_context_omits_harness_system_prompt( + self, llm_classifier_config: dict[str, object] + ) -> None: + dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}'))) + router: Final = ComplexityRouter( + "test-complexity-router", + dependency, + {**llm_classifier_config, "classifier_context_window_size": 0}, + ) + messages: Final = [ + {"role": "user", "content": "What is two plus two?"}, + { + "role": "system", + "content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}], + }, + ] + request_kwargs: Final = {"litellm_metadata": {"user_agent": "claude-code/2.1.233"}} + original: Final = deepcopy((messages, request_kwargs)) + + result: Final = await router.async_pre_routing_hook("test-complexity-router", request_kwargs, messages) + + assert result is not None and result.routing_decision["cause"] == "llm_classifier" + assert result.messages == messages == original[0] + assert request_kwargs == original[1] + classifier_messages: Final = dependency.acompletion.call_args.kwargs["messages"] + assert classifier_messages[0]["content"] == classification_system_prompt( + router.config.classifier_context_window_size + ) + assert classifier_messages[1]["content"].strip() == "Classify this message:\nWhat is two plus two?" + @pytest.mark.parametrize( "tail,expected", ( diff --git a/tests/test_litellm/router_utils/test_health_state_cache.py b/tests/test_litellm/router_utils/test_health_state_cache.py index ffd031f9b7d..aa976bb1002 100644 --- a/tests/test_litellm/router_utils/test_health_state_cache.py +++ b/tests/test_litellm/router_utils/test_health_state_cache.py @@ -7,6 +7,7 @@ import time import pytest from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.router_utils.health_state_cache import DeploymentHealthCache @@ -145,8 +146,11 @@ class _SharedRedisFake: def __init__(self): self.store = {} self.fail_get = False + self.breaker_open = False def get_cache(self, key, parent_otel_span=None, **kwargs): + if self.breaker_open: + raise RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping get_cache") if self.fail_get: return None # RedisCache.get_cache swallows connection errors and returns None return self.store.get(key) @@ -192,3 +196,26 @@ def test_failed_redis_read_falls_back_to_local_copy(): {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} ) assert set(redis_fake.store[DeploymentHealthCache.CACHE_KEY]) == {"prod-bad", "internal-bad"} + + +def test_open_circuit_breaker_read_still_merges_into_local_copy(caplog): + """A read refused by the open breaker is a miss, so the merge and local write still happen quietly.""" + redis_fake = _SharedRedisFake() + pod_a = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) + pod_b = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) + pod_a.set_deployment_health_states( + {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} + ) + pod_b.set_deployment_health_states( + {"internal-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "timeout"}} + ) + pod_a.set_deployment_health_states( + {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} + ) + redis_fake.breaker_open = True + with caplog.at_level("ERROR"): + pod_a.set_deployment_health_states( + {"prod-new-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} + ) + assert caplog.records == [] + assert pod_a.get_unhealthy_deployment_ids() == {"prod-bad", "internal-bad", "prod-new-bad"} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8d8317c2fd1..4811aad012b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35089,7 +35089,7 @@ export interface components { classification_prompt?: string | null; /** * Classifier Context Budget Chars - * @description Maximum characters of prior-turn text quoted to the LLM classifier, across the whole context window, per classification call. Turns are taken newest first and quoted whole while they fit, so a conversation small enough to quote entirely is never cut; once the budget runs out the older turns are dropped whole and only the turn straddling the boundary is truncated, into whatever space is left. The current ask and the caller's system prompt sit outside this budget and are always sent in full, as does the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and suppresses the block; set classifier_context_window_size to 0 to turn context off deliberately. Only applies when classifier_type is 'llm'. + * @description Maximum characters of prior-turn text quoted to the LLM classifier, across the whole context window, per classification call. Turns are taken newest first and quoted whole while they fit, so a conversation small enough to quote entirely is never cut; once the budget runs out the older turns are dropped whole and only the turn straddling the boundary is truncated, into whatever space is left. The current ask and, except for Claude Code requests, the extracted system-role text sit outside this budget and are sent in full, as does the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and suppresses the block; set classifier_context_window_size to 0 to turn context off deliberately. Only applies when classifier_type is 'llm'. * @default 8000 */ classifier_context_budget_chars: number; @@ -35106,7 +35106,7 @@ export interface components { classifier_context_per_turn_chars?: number | null; /** * Classifier Context Window Size - * @description Number of prior user turns (tool output and harness reminders excluded) to include as context in the LLM classifier prompt, so a follow-up like 'now do the same for the streaming path' is classified against what it refers to. Counts turns of both roles when classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier model, which may be a different deployment or provider than the routed completion model; that call already carries the current user ask and the caller's system prompt in full. Set to 0 to send neither prior turns nor any conversation context beyond the current ask. Only applies when classifier_type is 'llm'. + * @description Number of prior user turns (tool output and harness reminders excluded) to include as context in the LLM classifier prompt, so a follow-up like 'now do the same for the streaming path' is classified against what it refers to. Counts turns of both roles when classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier model, which may be a different deployment or provider than the routed completion model; that call carries the current user ask and, except for Claude Code requests, the extracted system-role text in full. Claude Code system text is omitted to avoid classifying harness instructions; the routed completion still receives it. Set to 0 to send neither prior turns nor any conversation context beyond the current ask. Only applies when classifier_type is 'llm'. * @default 3 */ classifier_context_window_size: number;