diff --git a/.github/workflows/guard-main-branch.yml b/.github/workflows/guard-main-branch.yml index 21aad18d298..aa4968f0c1e 100644 --- a/.github/workflows/guard-main-branch.yml +++ b/.github/workflows/guard-main-branch.yml @@ -31,12 +31,12 @@ jobs: echo "PR head repo: $HEAD_REPO" echo "PR head branch: $HEAD_REF" if [ "$HEAD_REPO" != "$BASE_REPO" ]; then - echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the 'litellm_oss_staging' branch instead." + echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead." exit 1 fi if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then echo "Allowed source branch." exit 0 fi - echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_oss_staging' instead." + echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead." exit 1 diff --git a/CLAUDE.md b/CLAUDE.md index 5affa7748d7..78da2c65d96 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,7 +19,7 @@ Same thing for bug fixes. The tests should make it so that this specific bug can End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` -When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose +When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose for internal contributors; external / OSS contributions target the current daily OSS branch instead, named `litellm_oss_daily_YYYY_MM_DD` (a fresh one is cut each weekday, so use the most recent) When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1080579d0fa..0202965ec4b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -322,7 +322,7 @@ npm run build ## Submitting Your PR 1. **Push your branch**: `git push origin your-feature-branch` -2. **Create a PR**: Go to GitHub and create a pull request +2. **Create a PR**: Go to GitHub and open a pull request against the current daily OSS branch, named `litellm_oss_daily_YYYY_MM_DD`. A fresh one is cut each weekday, so pick the most recent from the [branch list](https://github.com/BerriAI/litellm/branches/all?query=litellm_oss_daily). Do not target `main`. 3. **Fill out the PR template**: Provide clear description of changes 4. **Wait for review**: Maintainers will review and provide feedback 5. **Address feedback**: Make requested changes and push updates diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index b67f7d42127..02574ca505d 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -46,6 +46,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/fallback", "/fallbacks", "/cache_settings", + "/coordination_redis/", "/cost_tracking", "/cost/", "/credentials", diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 25b02dd5f37..469d52c03a7 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -76,10 +76,13 @@ so fall back to "default" (or an explicit override) to avoid a cyclic dependency {{- end }} {{/* -Get redis service name +Get redis service name. +The bundled Redis subchart only serves sentinel in "replication" architecture +(it rejects standalone + sentinel outright), and in that mode the sentinel +Service is named "-redis", not "-redis-master". */}} {{- define "litellm.redis.serviceName" -}} -{{- if and (eq .Values.redis.architecture "standalone") .Values.redis.sentinel.enabled -}} +{{- if .Values.redis.sentinel.enabled -}} {{- printf "%s-%s" .Release.Name (default "redis" .Values.redis.nameOverride | trunc 63 | trimSuffix "-") -}} {{- else -}} {{- printf "%s-%s-master" .Release.Name (default "redis" .Values.redis.nameOverride | trunc 63 | trimSuffix "-") -}} diff --git a/helm/litellm-helm/templates/configmap-litellm.yaml b/helm/litellm-helm/templates/configmap-litellm.yaml index acbe4e3a4b5..03e4f620206 100644 --- a/helm/litellm-helm/templates/configmap-litellm.yaml +++ b/helm/litellm-helm/templates/configmap-litellm.yaml @@ -1,9 +1,22 @@ {{- if .Values.proxyConfigMap.create }} +{{- $config := deepCopy .Values.proxy_config }} +{{- if and .Values.redis.enabled (dig "coordination" "enabled" true .Values.redis) }} +{{- $generalSettings := (get $config "general_settings") | default dict }} +{{- if not (hasKey $generalSettings "coordination_redis") }} +{{- $coordinationRedis := dict "host" "os.environ/REDIS_HOST" "port" "os.environ/REDIS_PORT" "password" "os.environ/REDIS_PASSWORD" }} +{{- if .Values.redis.sentinel.enabled }} +{{- $sentinelNode := list (include "litellm.redis.serviceName" .) (include "litellm.redis.port" . | int) }} +{{- $coordinationRedis = dict "sentinel_nodes" (list $sentinelNode) "service_name" (default "mymaster" .Values.redis.sentinel.masterSet) "password" "os.environ/REDIS_PASSWORD" }} +{{- end }} +{{- $_ := set $generalSettings "coordination_redis" $coordinationRedis }} +{{- $_ := set $config "general_settings" $generalSettings }} +{{- end }} +{{- end }} apiVersion: v1 kind: ConfigMap metadata: name: {{ include "litellm.fullname" . }}-config data: config.yaml: | -{{ .Values.proxy_config | toYaml | indent 6 }} +{{ $config | toYaml | indent 6 }} {{- end }} diff --git a/helm/litellm-helm/tests/coordination_redis_tests.yaml b/helm/litellm-helm/tests/coordination_redis_tests.yaml new file mode 100644 index 00000000000..0b58b1e6bc8 --- /dev/null +++ b/helm/litellm-helm/tests/coordination_redis_tests.yaml @@ -0,0 +1,143 @@ +suite: test coordination redis +templates: + - configmap-litellm.yaml + - deployment.yaml +tests: + - it: should not render coordination_redis when redis is disabled + template: configmap-litellm.yaml + set: + redis.enabled: false + asserts: + - notMatchRegex: + path: data["config.yaml"] + pattern: coordination_redis + + - it: should not emit redis env vars when redis is disabled + template: deployment.yaml + set: + redis.enabled: false + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: RELEASE-NAME-redis-master + any: true + + - it: should render coordination_redis pointing at the bundled redis when enabled + template: configmap-litellm.yaml + set: + redis.enabled: true + asserts: + - matchRegex: + path: data["config.yaml"] + pattern: "coordination_redis:\n host: os.environ/REDIS_HOST\n password: os.environ/REDIS_PASSWORD\n port: os.environ/REDIS_PORT\n" + - matchRegex: + path: data["config.yaml"] + pattern: "master_key: os.environ/PROXY_MASTER_KEY" + + - it: should emit redis env vars backing the coordination_redis os.environ refs + template: deployment.yaml + set: + redis.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: RELEASE-NAME-redis-master + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PORT + value: "6379" + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: RELEASE-NAME-redis + key: redis-password + + - it: should not render coordination_redis when coordination is opted out + template: configmap-litellm.yaml + set: + redis.enabled: true + redis.coordination.enabled: false + asserts: + - notMatchRegex: + path: data["config.yaml"] + pattern: coordination_redis + + - it: should keep emitting redis env vars when coordination is opted out + template: deployment.yaml + set: + redis.enabled: true + redis.coordination.enabled: false + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: RELEASE-NAME-redis-master + + - it: should not clobber a user supplied coordination_redis block + template: configmap-litellm.yaml + set: + redis.enabled: true + proxy_config.general_settings.coordination_redis: + url: os.environ/COORDINATION_REDIS_URL + asserts: + - matchRegex: + path: data["config.yaml"] + pattern: "coordination_redis:\n url: os.environ/COORDINATION_REDIS_URL\n" + - notMatchRegex: + path: data["config.yaml"] + pattern: "host: os.environ/REDIS_HOST" + + - it: should render sentinel_nodes and service_name in sentinel mode + template: configmap-litellm.yaml + set: + redis.enabled: true + redis.architecture: replication + redis.sentinel.enabled: true + asserts: + # The sentinel Service the redis subchart renders is "-redis", and a + # plain client cannot speak the sentinel protocol, so host/port must not appear + - matchRegex: + path: data["config.yaml"] + pattern: "coordination_redis:\n password: os.environ/REDIS_PASSWORD\n sentinel_nodes:\n - - RELEASE-NAME-redis\n - 26379\n service_name: mymaster\n" + - notMatchRegex: + path: data["config.yaml"] + pattern: "host: os.environ/REDIS_HOST" + + - it: should carry a custom sentinel masterSet into service_name + template: configmap-litellm.yaml + set: + redis.enabled: true + redis.architecture: replication + redis.sentinel.enabled: true + redis.sentinel.masterSet: litellm-master + asserts: + - matchRegex: + path: data["config.yaml"] + pattern: "service_name: litellm-master" + + - it: should point REDIS_HOST at the sentinel service in sentinel mode + template: deployment.yaml + set: + redis.enabled: true + redis.architecture: replication + redis.sentinel.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: RELEASE-NAME-redis + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PORT + value: "26379" diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 6e30a6af444..d3821a547e5 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -331,12 +331,28 @@ postgresql: # secretKeys: # userPasswordKey: password -# requires cache: true in config file -# either enable this or pass a secret for REDIS_HOST, REDIS_PORT, REDIS_PASSWORD or REDIS_URL -# with cache: true to use existing redis instance +# Redis is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend +# tracking, and the pod lock manager. Enabling this deploys the bundled Redis +# subchart, wires REDIS_HOST / REDIS_PORT / REDIS_PASSWORD into the proxy, and +# renders a `general_settings.coordination_redis` block into the proxy config. +# +# To point at an existing Redis instead, leave `enabled: false` and pass a +# secret for REDIS_HOST, REDIS_PORT, REDIS_PASSWORD or REDIS_URL; the proxy +# falls back to those env vars for coordination. Set `cache: true` in the proxy +# config only if you also want LLM response caching, which is independent of +# coordination +# +# When `redis.sentinel.enabled` is set, the coordination block is rendered with +# `sentinel_nodes` and `service_name` (from `redis.sentinel.masterSet`) instead +# of host/port, because a plain Redis client cannot talk to the sentinel port redis: enabled: false architecture: standalone + coordination: + # Set to false to keep the bundled Redis for response caching only and leave + # `general_settings.coordination_redis` out of the rendered config. A + # `coordination_redis` block you define yourself in `proxy_config` always wins + enabled: true # Prisma migration job settings migrationJob: diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 4319907883e..7c281aa158b 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -213,6 +213,10 @@ harmless no-op for the Job and authoritative for the app pods. */}} - name: DISABLE_SCHEMA_UPDATE value: "true" +{{/* These feed the proxy's coordination Redis (cross-pod rate limits, spend + tracking, pod lock manager) via its REDIS_* env fallback. An explicit + `general_settings.coordination_redis` block in proxy_config takes + precedence over anything emitted here. */}} {{- if $root.Values.redis.host }} - name: REDIS_HOST value: {{ $root.Values.redis.host | quote }} @@ -226,10 +230,11 @@ harmless no-op for the Job and authoritative for the app pods. key: {{ $root.Values.redis.passwordSecret.passwordKey | default "password" }} {{- end }} {{- if $root.Values.redis.cluster }} -{{/* The proxy's Cache() reads REDIS_CLUSTER_NODES as JSON and constructs a - RedisClusterCache when it's set (litellm/caching/caching.py:169-192). - We seed with the single configured endpoint — the cluster client - discovers the remaining nodes from CLUSTER SLOTS at startup. */}} +{{/* The proxy falls back to REDIS_CLUSTER_NODES (JSON) to build a cluster-mode + coordination client when `general_settings.coordination_redis` is absent + and no plain-Redis response cache is configured. We seed with the single + configured endpoint; the cluster client discovers the remaining nodes from + CLUSTER SLOTS at startup. */}} - name: REDIS_CLUSTER_NODES value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }} {{- end }} diff --git a/helm/litellm/tests/redis_env_tests.yaml b/helm/litellm/tests/redis_env_tests.yaml new file mode 100644 index 00000000000..684d7071b35 --- /dev/null +++ b/helm/litellm/tests/redis_env_tests.yaml @@ -0,0 +1,109 @@ +suite: test redis coordination env vars +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: gateway omits redis env vars when no host is configured + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: redis.example.com + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_CLUSTER_NODES + any: true + + - it: gateway emits host, port and password when redis is configured + template: gateway/deployment.yaml + set: + redis.host: redis.example.com + redis.port: 6380 + redis.passwordSecret.name: redis-secret + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: redis.example.com + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PORT + value: "6380" + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: redis-secret + key: password + + - it: backend emits the same redis env vars so both pods coordinate on one redis + template: backend/deployment.yaml + set: + redis.host: redis.example.com + redis.passwordSecret.name: redis-secret + redis.passwordSecret.passwordKey: redis-password + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: redis.example.com + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: redis-secret + key: redis-password + + - it: gateway omits REDIS_PASSWORD for an auth-less redis + template: gateway/deployment.yaml + set: + redis.host: redis.example.com + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + any: true + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + value: redis.example.com + + - it: gateway seeds REDIS_CLUSTER_NODES from host and port in cluster mode + template: gateway/deployment.yaml + set: + redis.host: redis.example.com + redis.port: 6380 + redis.cluster: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_CLUSTER_NODES + value: '[{"host":"redis.example.com","port":6380}]' + + - it: gateway omits REDIS_CLUSTER_NODES when cluster mode is off + template: gateway/deployment.yaml + set: + redis.host: redis.example.com + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_CLUSTER_NODES + any: true diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 6aa5dd39cd0..a8f2d39663e 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -100,7 +100,18 @@ database: usernameKey: username passwordKey: password -# Optional Redis (caching, rate limiting). Leave host empty to disable. +# Optional Redis. Leave host empty to disable. +# +# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend +# tracking, and the pod lock manager. The chart emits REDIS_HOST / REDIS_PORT / +# REDIS_PASSWORD, which the proxy picks up through its coordination Redis env +# fallback. Response caching is separate and off unless you enable it in +# `proxy_config.litellm_settings.cache`. +# +# For full control, define `general_settings.coordination_redis` in +# `proxy_config` (host/port/password/username/url/ssl/startup_nodes/ +# sentinel_nodes/sentinel_password/service_name, each accepting os.environ/VAR +# refs). An explicit block overrides these env vars. # # Set `cluster: true` for Redis Cluster mode (e.g. AWS ElastiCache Cluster, # self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from diff --git a/litellm/_redis.py b/litellm/_redis.py index bb3a0974241..0b91cdabffc 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -325,8 +325,19 @@ def _get_redis_client_logic(**env_overrides): value = get_secret(v) # type: ignore env_overrides[k] = value + environment_kwargs = _redis_kwargs_from_environment() + + # An explicitly configured connection target outranks REDIS_URL from the + # environment. Without this, the url branch below strips the caller's + # host/port/password and silently connects to whatever REDIS_URL names. + caller_named_a_target = any( + env_overrides.get(key) is not None for key in ("host", "startup_nodes", "sentinel_nodes") + ) + if caller_named_a_target and env_overrides.get("url") is None: + environment_kwargs.pop("url", None) + redis_kwargs = { - **_redis_kwargs_from_environment(), + **environment_kwargs, **env_overrides, } @@ -678,9 +689,8 @@ def get_redis_connection_pool( redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) connection_class = async_redis.Connection - if "ssl" in redis_kwargs: + if redis_kwargs.pop("ssl", False): connection_class = async_redis.SSLConnection - redis_kwargs.pop("ssl", None) redis_kwargs["connection_class"] = connection_class return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs) diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index bd62d1e303a..20239d831cc 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -19,7 +19,7 @@ import os import time import traceback from datetime import datetime as datetimeObj -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Sequence, Union import httpx from httpx import Response @@ -50,6 +50,7 @@ from litellm.types.integrations.base_health_check import IntegrationHealthCheckS from litellm.types.integrations.datadog import ( DD_ERRORS, DD_MAX_BATCH_SIZE, + DD_MAX_PAYLOAD_SIZE_BYTES, DataDogStatus, DatadogInitParams, DatadogPayload, @@ -384,8 +385,10 @@ class DataDogLogger( async def _send_with_413_split(self, batch: List) -> List: """ - Send a batch, halving any sub-batch that 413s (payload too large) and retrying the - halves, since Datadog enforces a 5MB uncompressed limit per request. + Send a batch, halving any sub-batch that exceeds Datadog's intake limits before + sending, and halving again on a 413 (payload too large) response, since Datadog + enforces a 5MB uncompressed limit per request. The proactive split avoids paying + a serialize + gzip + round trip for a payload the intake is guaranteed to reject. A 413 surfaces as a raised MaskedHTTPStatusError (httpx raise_for_status), not a returned response, so both paths are handled. A lone event that still 413s is @@ -398,6 +401,11 @@ class DataDogLogger( chunk = pending.pop() if not chunk: continue + if len(chunk) > 1 and self._exceeds_intake_limits(chunk): + mid = len(chunk) // 2 + pending.append(chunk[mid:]) + pending.append(chunk[:mid]) + continue try: response = await self.async_send_compressed_data(chunk) except Exception as e: @@ -436,6 +444,21 @@ class DataDogLogger( def _undelivered(chunk: List, pending: List[List]) -> List: return chunk + [event for remaining in reversed(pending) for event in remaining] + @staticmethod + def _exceeds_intake_limits(chunk: Sequence[DatadogPayload]) -> bool: + """ + True when a chunk would breach Datadog's log intake limits: more than + DD_MAX_BATCH_SIZE events per payload, or a serialized size above + DD_MAX_PAYLOAD_SIZE_BYTES (held under Datadog's 5MB uncompressed cap so + the batch is split before the intake rejects it with a 413). + """ + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + if len(chunk) > DD_MAX_BATCH_SIZE: + return True + payload_size_bytes = len(safe_dumps(chunk).encode("utf-8")) + return payload_size_bytes > DD_MAX_PAYLOAD_SIZE_BYTES + async def flush_queue(self): if self.flush_lock is None: return diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index 17011bb8db7..3038bdb90b2 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -223,6 +223,15 @@ lives in [`plumbing/`](./plumbing): readers/exporters receive them alongside the server metrics, and one is built and registered as the global only when none is set (mirroring how V2 owns trace export). +- [`events.py`](./plumbing/events.py) — GenAI client events. Gated on + `enable_events` (`LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS`), a failed LLM call + records the semconv `gen_ai.client.operation.exception` log event at severity + WARN, carrying `exception.type` / `exception.message` / `exception.stacktrace` + and correlated to the failed span through the trace and span ids. The + `LoggerProvider` is resolved like the meter provider, except that an explicit + `NoOpLoggerProvider` global is an operator opt-out that builds no recorder at + all. The deprecated `error.*` span attributes and the `exception` span event + are still stamped by the emitter for backwards compatibility. ### Adapter diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index 46aa166a8bb..f97f8b8394c 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -18,6 +18,7 @@ from litellm.integrations.otel.model.payloads import ( ServiceSpanData, SpanError, ) +from litellm.integrations.otel.plumbing.events import GenAIEventRecorder from litellm.integrations.otel.plumbing.providers import to_otel_span_kind from litellm.integrations.otel.model.semconv import Error, ExceptionEvent, LiteLLMError from litellm.integrations.otel.model.spans import ( @@ -77,9 +78,11 @@ class SpanEmitter: tracer: Tracer, config: OpenTelemetryV2Config, mappers: Sequence[AttributeMapper] | None = None, + event_recorder: GenAIEventRecorder | None = None, ) -> None: self._tracer = tracer self._config = config + self._event_recorder = event_recorder # The mapper chain is the sole source of span attributes. When not # passed in, resolve it from the config so there's one source of truth. self._mappers: list[AttributeMapper] = ( @@ -223,6 +226,14 @@ class SpanEmitter: ExceptionEvent.NAME, {ExceptionEvent.TYPE: error_type, ExceptionEvent.MESSAGE: message}, ) + if self._event_recorder is not None and role is SpanRole.LLM_CALL: + self._event_recorder.record_operation_exception( + span_context=span.get_span_context(), + error_type=error_type, + message=message, + stack_trace=error.stack_trace, + timestamp_ns=end_time_ns, + ) # On success leave the status UNSET (the semconv default) rather than # forcing OK — that matches the FastAPI server span and avoids implying a # span-level health signal litellm doesn't actually evaluate. Only a diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e258b239d93..be72fabd387 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -6,6 +6,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, Sequence, cast from opentelemetry.context import Context, attach, get_current +from opentelemetry.sdk._logs import LoggerProvider from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import Span, Tracer, get_current_span, use_span @@ -40,14 +41,17 @@ from litellm.integrations.otel.model.payloads import ( is_mcp_list_tools, is_mcp_tool_call, ) +from litellm.integrations.otel.plumbing.events import GenAIEventRecorder from litellm.integrations.otel.plumbing.metrics import ( GenAIMetricRecorder, create_genai_metrics, ) from litellm.integrations.otel.plumbing.providers import ( build_tracer_provider, + get_event_logger, get_meter, get_tracer, + resolve_logger_provider, resolve_meter_provider, ) from litellm.integrations.otel.plumbing.routing import TenantTracerCache @@ -104,7 +108,7 @@ class OpenTelemetryV2(CustomLogger): config: OpenTelemetryV2Config | None = None, callback_name: str | None = None, tracer_provider: TracerProvider | None = None, - logger_provider: Any | None = None, # reserved for OTel logs + logger_provider: LoggerProvider | None = None, meter_provider: Any | None = None, **kwargs: Any, ) -> None: @@ -117,7 +121,12 @@ class OpenTelemetryV2(CustomLogger): self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME) self._metrics_recorder = self._init_metrics(meter_provider) self._metric_filter_error_logged = False - self._emitter = SpanEmitter(self.tracer, self.config, mappers=resolve_mappers(self.config.mapper_names)) + self._emitter = SpanEmitter( + self.tracer, + self.config, + mappers=resolve_mappers(self.config.mapper_names), + event_recorder=self._init_events(logger_provider), + ) self._tenant_tracers = TenantTracerCache(self.config, callback_name, LITELLM_TRACER_NAME) self._open_llm_calls: "OrderedDict[str, _LLMCallSpan]" = OrderedDict() self._init_otel_logger_on_litellm_proxy() @@ -136,6 +145,22 @@ class OpenTelemetryV2(CustomLogger): meter = get_meter(provider, LITELLM_TRACER_NAME) return GenAIMetricRecorder(create_genai_metrics(meter), self.callback_name) + def _init_events(self, logger_provider: LoggerProvider | None) -> "GenAIEventRecorder | None": + """Create the GenAI event recorder when events are enabled, else ``None``. + + ``logger_provider`` is an explicit override (tests inject one); otherwise the + provider is resolved from the OTel global so an operator-configured logs + pipeline receives the events, building and registering one only when no + global provider is set. A ``None`` resolution means the operator opted out + of the logs signal, so no recorder is built. + """ + if not self.config.enable_events: + return None + provider = resolve_logger_provider(self.config, logger_provider) + if provider is None: + return None + return GenAIEventRecorder(get_event_logger(provider, LITELLM_TRACER_NAME)) + # ====================================================================== # # Proxy global registration # ====================================================================== # diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index aab80c7e5c4..44b2f7e0488 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -177,6 +177,19 @@ class ExceptionEvent: NAME: Final = "exception" TYPE: Final = "exception.type" MESSAGE: Final = "exception.message" + STACKTRACE: Final = "exception.stacktrace" + + +class GenAIEvent: + """GenAI semconv event names, from the GenAI registry's *events* section. + + ``gen_ai.client.operation.exception`` is defined as a log-based event + (severity WARN) carrying the ``exception.*`` trio, correlated to the failed + span via the trace/span ids — the semconv-compliant home for GenAI failure + details, unlike the deprecated ``error.message`` span attribute. + """ + + OPERATION_EXCEPTION: Final = "gen_ai.client.operation.exception" class Server: diff --git a/litellm/integrations/otel/plumbing/events.py b/litellm/integrations/otel/plumbing/events.py new file mode 100644 index 00000000000..f674526d04f --- /dev/null +++ b/litellm/integrations/otel/plumbing/events.py @@ -0,0 +1,52 @@ +"""GenAI client events: the ``gen_ai.client.operation.exception`` log event. + +The GenAI semantic conventions define exception recording for client +operations as a log-based event (severity WARN) carrying the ``exception.*`` +attribute trio, correlated to the failed span through the trace/span ids — +not as a span attribute or span event. This module owns building and +emitting that event; the exporter pipeline it rides is built in +:mod:`litellm.integrations.otel.plumbing.providers`. +""" + +from dataclasses import dataclass + +from opentelemetry._events import Event, EventLogger +from opentelemetry._logs.severity import SeverityNumber +from opentelemetry.trace import SpanContext + +from litellm.integrations.otel.model.semconv import ExceptionEvent, GenAIEvent + + +@dataclass(frozen=True, slots=True) +class GenAIEventRecorder: + event_logger: EventLogger + + def record_operation_exception( + self, + span_context: SpanContext, + error_type: str, + message: str, + stack_trace: str | None, + timestamp_ns: int | None, + ) -> None: + # ``exception.type`` and ``exception.message`` are the semconv-required + # pair and always ride the event; only the recommended stacktrace is + # conditional on the payload carrying one. + stacktrace = ((ExceptionEvent.STACKTRACE, stack_trace),) if stack_trace else () + self.event_logger.emit( + Event( + name=GenAIEvent.OPERATION_EXCEPTION, + timestamp=timestamp_ns, + trace_id=span_context.trace_id, + span_id=span_context.span_id, + trace_flags=span_context.trace_flags, + severity_number=SeverityNumber.WARN, + attributes=dict( + ( + (ExceptionEvent.TYPE, error_type), + (ExceptionEvent.MESSAGE, message), + *stacktrace, + ) + ), + ) + ) diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index ac971c6daa8..ced65aa1ec3 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -2,9 +2,20 @@ from typing import TYPE_CHECKING, Any, Callable, Iterable -from opentelemetry import baggage, metrics +from opentelemetry import _logs, baggage, metrics +from opentelemetry._events import EventLogger +from opentelemetry._logs import LoggerProvider, NoOpLoggerProvider from opentelemetry.context import Context from opentelemetry.metrics import MeterProvider, NoOpMeterProvider +from opentelemetry.sdk._events import EventLoggerProvider +from opentelemetry.sdk._logs import LoggerProvider as SDKLoggerProvider +from opentelemetry.sdk._logs.export import ( + BatchLogRecordProcessor, + ConsoleLogExporter, + InMemoryLogExporter, + LogExporter, + SimpleLogRecordProcessor, +) from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider @@ -224,6 +235,112 @@ def build_metric_reader(config: OpenTelemetryV2Config) -> "MetricReader": return PeriodicExportingMetricReader(exporter, export_interval_millis=5000) +def _otlp_logs_endpoint(endpoint: str | None) -> str | None: + """Point an OTLP/HTTP base endpoint at the ``/v1/logs`` signal path. + + The OTLP/HTTP exporter only appends ``/v1/logs`` when it reads + ``OTEL_EXPORTER_OTLP_ENDPOINT`` itself; an explicitly passed endpoint is used + verbatim, so a base URL would POST to the root. Mirror ``_otlp_traces_endpoint`` + for the logs signal (rewriting a sibling signal path when present). + """ + if not endpoint: + return endpoint + endpoint = endpoint.rstrip("/") + if endpoint.endswith("/v1/logs"): + return endpoint + for other_signal in ("/v1/traces", "/v1/metrics"): + if endpoint.endswith(other_signal): + return endpoint[: -len(other_signal)] + "/v1/logs" + return endpoint + "/v1/logs" + + +def build_log_exporter(config: OpenTelemetryV2Config) -> LogExporter: + """Build a log exporter mirroring the exporter selection of the other signals. + + ``console`` (and any unrecognized kind) exports to the console; ``otlp_http`` + and ``otlp_grpc`` export over OTLP with the configured endpoint/headers; + ``in_memory`` buffers for tests. Like GenAI metrics, events ride the + single-destination shorthand fields, not the multi-exporter ``exporters`` list. + """ + kind = (config.exporter or "console").lower() + if kind in ("in_memory", "inmemory", "memory"): + return InMemoryLogExporter() + if kind in ("otlp_http", "http", "http/protobuf", "http/json"): + from opentelemetry.exporter.otlp.proto.http._log_exporter import ( + OTLPLogExporter as HTTPLogExporter, + ) + + return HTTPLogExporter( + endpoint=_otlp_logs_endpoint(config.endpoint), + headers=parse_headers(config.headers), + ) + if kind in ("otlp_grpc", "grpc"): + try: + from opentelemetry.exporter.otlp.proto.grpc._log_exporter import ( + OTLPLogExporter as GRPCLogExporter, + ) + except ImportError as exc: + raise ImportError( + "OpenTelemetry OTLP gRPC log exporter is not available. Install " + "`opentelemetry-exporter-otlp` and `grpcio` (or `litellm[grpc]`)." + ) from exc + + return GRPCLogExporter(endpoint=config.endpoint, headers=parse_headers(config.headers)) + return ConsoleLogExporter() + + +def build_logger_provider( + config: OpenTelemetryV2Config, + log_exporter: LogExporter | None = None, +) -> SDKLoggerProvider: + """Build the :class:`LoggerProvider` GenAI events export through. + + ``log_exporter`` is an explicit override (tests inject an + ``InMemoryLogExporter``); otherwise the exporter is selected from the config's + exporter kind via :func:`build_log_exporter`. Console and in-memory exporters + get a Simple processor (synchronous export, which tests rely on), everything + else a Batch processor — the same split as span processing. + """ + exporter = log_exporter if log_exporter is not None else build_log_exporter(config) + provider = SDKLoggerProvider(resource=build_resource(config)) + use_simple = isinstance(exporter, (ConsoleLogExporter, InMemoryLogExporter)) + provider.add_log_record_processor( + SimpleLogRecordProcessor(exporter) if use_simple else BatchLogRecordProcessor(exporter) + ) + return provider + + +def resolve_logger_provider( + config: OpenTelemetryV2Config, + logger_provider: SDKLoggerProvider | None = None, +) -> SDKLoggerProvider | None: + """Resolve the :class:`LoggerProvider` GenAI events record through, or ``None`` + when the operator has opted out of the logs signal. + + Same resolution order as :func:`resolve_meter_provider`: an injected provider + wins (DI/tests); an operator-configured SDK global is reused so events ride + their pipeline; an explicit ``NoOpLoggerProvider`` global is an opt-out and + yields ``None``, so no event is ever built. Only the default placeholder + global makes V2 build a provider from the config and publish it as the global. + """ + if logger_provider is not None: + return logger_provider + + existing: LoggerProvider = _logs.get_logger_provider() + if isinstance(existing, SDKLoggerProvider): + return existing + if isinstance(existing, NoOpLoggerProvider): + return None + + provider = build_logger_provider(config) + _logs.set_logger_provider(provider) + return provider + + +def get_event_logger(provider: SDKLoggerProvider, name: str = "litellm") -> EventLogger: + return EventLoggerProvider(logger_provider=provider).get_event_logger(name, litellm_version) + + def build_meter_provider( config: OpenTelemetryV2Config, metric_reader: "MetricReader | None" = None, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 60fc021a6a8..f575372fc3d 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1618,6 +1618,14 @@ class PrometheusLogger(CustomLogger): user_id: Optional[str] = None, user_api_key_org_id: Optional[str] = None, ): + if ( + isinstance(self.litellm_remaining_team_budget_metric, NoOpMetric) + and isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric) + and isinstance(self.litellm_remaining_user_budget_metric, NoOpMetric) + and isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric) + ): + return + _metadata = litellm_params.get("metadata") or {} _team_spend = _metadata.get("user_api_key_team_spend", None) _team_max_budget = _metadata.get("user_api_key_team_max_budget", None) @@ -3332,6 +3340,9 @@ class PrometheusLogger(CustomLogger): - looks up team info from db if not available in metadata - Set team budget metrics """ + if isinstance(self.litellm_remaining_team_budget_metric, NoOpMetric): + return + if user_api_team: team_object = await self._assemble_team_object( team_id=user_api_team, @@ -3453,6 +3464,9 @@ class PrometheusLogger(CustomLogger): - Fetches org info via cache (get_org_object) - Sets org budget metrics """ + if isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric): + return + if not org_id: return @@ -3582,6 +3596,9 @@ class PrometheusLogger(CustomLogger): key_max_budget: Optional[float], key_spend: Optional[float], ): + if isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric): + return + if user_api_key: user_api_key_dict = await self._assemble_key_object( user_api_key=user_api_key, @@ -3642,6 +3659,9 @@ class PrometheusLogger(CustomLogger): - looks up user info from db if not available in metadata - Set user budget metrics """ + if isinstance(self.litellm_remaining_user_budget_metric, NoOpMetric): + return + if user_id: user_object = await self._assemble_user_object( user_id=user_id, diff --git a/litellm/litellm_core_utils/fallback_generalizations.py b/litellm/litellm_core_utils/fallback_generalizations.py index abc171f900a..8b5a797cbd7 100644 --- a/litellm/litellm_core_utils/fallback_generalizations.py +++ b/litellm/litellm_core_utils/fallback_generalizations.py @@ -8,9 +8,12 @@ The metadata is a partial cost-map entry: ``litellm_provider`` drives provider routing, and the remaining fields (``mode``, ``supports_*``, context window, pricing, ...) drive ``get_model_info`` / ``supports_*``. -Precedence: rules are evaluated in file order and the first match wins. They are -consulted only after exact and case-insensitive lookups miss, so an exact entry -always takes precedence over a rule. +Precedence: rules are evaluated in file order and the first match wins. Callers +with extra constraints (model-info resolution checks the provider) use +``match_all_fallback_generalizations`` to skip inapplicable earlier rules instead +of discarding the model name. Rules are consulted only after exact and +case-insensitive lookups miss, so an exact entry always takes precedence over a +rule. Patterns are matched case-insensitively with ``re.search`` and are not implicitly anchored: a rule must include ``^`` and ``$`` (as the shipped rules do) to bind to @@ -105,15 +108,15 @@ class _FallbackGeneralizations: ) return compiled - def match(self, model: str) -> Optional[dict]: + def matches(self, model: str) -> list[dict]: if not model: - return None + return [] if self._compiled is None: self._compiled = self._compile() - for pattern, model_info in self._compiled: - if pattern.search(model) is not None: - return dict(model_info) - return None + return [dict(model_info) for pattern, model_info in self._compiled if pattern.search(model) is not None] + + def match(self, model: str) -> Optional[dict]: + return next(iter(self.matches(model)), None) _registry = _FallbackGeneralizations() @@ -139,3 +142,12 @@ def match_fallback_generalization(model: str) -> Optional[dict]: O(number of rules). Only call this once exact lookups have missed. """ return _registry.match(model) + + +def match_all_fallback_generalizations(model: str) -> list[dict]: + """Return the ``model_info`` of every rule whose regex matches ``model``, in rule order. + + Lets a caller with extra constraints (e.g. a provider match) skip an + inapplicable earlier rule instead of discarding the whole candidate. + """ + return _registry.matches(model) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index db540e5441d..964186fd76a 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -289,6 +289,13 @@ class AnthropicModelInfo(BaseLLMModelInfo): status_code=400, ) + @staticmethod + def _strip_version_suffix(model: str) -> str: + at = model.rfind("@") + if at > 0: + return model[:at] + return model + @staticmethod def _model_map_lookup_candidates(model: str) -> List[str]: """Model-map keys to try for ``model``: the id itself, the same id with a @@ -324,6 +331,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): _DATED_RELEASE_SUFFIX_RE.sub("", cand), _DOTTED_VERSION_RE.sub(r"\1-\2", cand), _strip_bedrock_id_suffixes(cand), + AnthropicModelInfo._strip_version_suffix(cand), ) ) return list(dict.fromkeys((*primary, *normalized))) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index e78802a1587..b74bbda0b66 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -3,6 +3,7 @@ from typing import Any, AsyncIterator, Dict, List, Optional, Tuple import httpx from litellm.constants import ( + ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, @@ -32,6 +33,12 @@ from ...common_utils import ( DEFAULT_ANTHROPIC_API_VERSION = "2023-06-01" +DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING = ( + "Dropping adaptive `thinking`/`output_config.effort` for model=%s: the model " + "does not support extended thinking, or max_tokens is too small to fit the " + "minimum thinking budget." +) + class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def get_supported_anthropic_messages_params(self, model: str) -> list: @@ -253,6 +260,111 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): existing_output_config.setdefault("effort", effort) optional_params["output_config"] = existing_output_config + @staticmethod + def _translate_adaptive_effort_for_non_adaptive_model( + model: str, optional_params: Dict, max_tokens: Optional[int] + ) -> None: + """Translate the 4.6+ adaptive-thinking interface (``thinking.type=adaptive`` + and/or ``output_config.effort``) down to what an older Anthropic model + supports. Clients like Claude Code send this interface unconditionally, so + without translation it reaches a pre-4.6 model and Anthropic rejects it with + "This model does not support the effort parameter". + + The reshape is silent, matching how the messages path already strips + unsupported ``output_config`` for older models (bedrock invoke, issue + #22797): the goal is to keep the request working, not to fail it. + + ``thinking.type=adaptive`` and ``output_config.effort`` are independent + capabilities. Adaptive thinking needs ``supports_adaptive_thinking`` (4.6+); + ``output_config.effort`` needs ``supports_output_config``, which some + non-adaptive models (e.g. Claude Opus 4.5) advertise on its own. So the two + are handled separately: + + - Adaptive-thinking models (4.6+): both are native, left untouched. + - ``supports_output_config`` but non-adaptive (Opus 4.5): keep + ``output_config.effort`` (native), only drop the unsupported adaptive + ``thinking`` block. When adaptive thinking is being dropped and the + effort level itself isn't supported by the model (e.g. ``xhigh``/``max`` + on Opus 4.5, which only accepts low/medium/high, while ``xhigh`` is + Claude Code's default), fall through to the legacy translation below + instead of forwarding a level Anthropic would reject. Effort-only + requests are always left untouched: provider subclasses own their level + normalization (bedrock clamps ``xhigh`` to the model's ceiling after + this base transform runs). + - Thinking-capable but neither (``supports_reasoning``, e.g. Haiku/Sonnet + 4.5): map effort to legacy ``thinking={type: enabled, budget_tokens}`` via + ``AnthropicConfig._map_reasoning_effort``, capped below ``max_tokens`` + (Anthropic requires ``max_tokens > budget_tokens``) and dropped when + ``max_tokens`` can't fit even the minimum budget. + - No reasoning support: ``thinking`` is dropped. + + For the last two, only the consumed ``effort`` key is removed from + ``output_config``; any residual (e.g. ``format``) is left for provider + subclasses to handle. + """ + from litellm.exceptions import BadRequestError as _BadRequestError + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + if AnthropicConfig._is_adaptive_thinking_model(model): + return + + output_config = optional_params.get("output_config") + thinking = optional_params.get("thinking") + effort = output_config.get("effort") if isinstance(output_config, dict) else None + adaptive_thinking = isinstance(thinking, dict) and thinking.get("type") == "adaptive" + if effort is None and not adaptive_thinking: + return + + if AnthropicConfig._model_supports_effort_param(model) and ( + not adaptive_thinking or AnthropicConfig._validate_effort_for_model(model, effort) is None + ): + if adaptive_thinking: + optional_params.pop("thinking", None) + return + + supports_thinking = AnthropicModelInfo._supports_model_capability(model, "supports_reasoning") + try: + legacy_thinking = ( + AnthropicConfig._map_reasoning_effort(reasoning_effort=effort or "medium", model=model) + if supports_thinking + else None + ) + except _BadRequestError as e: + raise AnthropicError(message=str(e.message), status_code=400) + capped_thinking = ( + AnthropicMessagesConfig._cap_thinking_budget_to_max_tokens(legacy_thinking, max_tokens) + if legacy_thinking is not None + else None + ) + + if capped_thinking is not None: + optional_params["thinking"] = capped_thinking + else: + verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_EFFORT_WARNING, model) + optional_params.pop("thinking", None) + + if isinstance(output_config, dict) and "effort" in output_config: + residual = {k: v for k, v in output_config.items() if k != "effort"} + if residual: + optional_params["output_config"] = residual + else: + optional_params.pop("output_config", None) + + @staticmethod + def _cap_thinking_budget_to_max_tokens(thinking: Dict, max_tokens: Optional[int]) -> Optional[Dict]: + """Cap a legacy ``thinking.budget_tokens`` below ``max_tokens`` (Anthropic + requires ``max_tokens > budget_tokens``). Returns the (possibly capped) + thinking dict, or ``None`` when ``max_tokens`` is too small to fit even the + minimum thinking budget and thinking should be dropped.""" + budget = thinking.get("budget_tokens") + if max_tokens is None or not isinstance(budget, int): + return thinking + if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: + return None + if budget < max_tokens: + return thinking + return {**thinking, "budget_tokens": max_tokens - 1} + def transform_anthropic_messages_request( self, model: str, @@ -284,6 +396,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params=anthropic_messages_optional_request_params, ) + self._translate_adaptive_effort_for_non_adaptive_model( + model=model, + optional_params=anthropic_messages_optional_request_params, + max_tokens=max_tokens, + ) + system_param = anthropic_messages_optional_request_params.get("system") if self.should_strip_billing_metadata() and system_param is not None: filtered_system = self._filter_billing_headers_from_system(system_param) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index daee3369a3c..56a47432997 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -93,26 +93,48 @@ class AmazonAnthropicClaudeMessagesConfig( return [{"type": "text", "text": value}] return [value] - def _normalize_system_role_messages_for_bedrock(self, anthropic_messages_request: dict) -> None: - """Bedrock Invoke rejects ``role: "system"`` entries inside ``messages`` on - some Claude aliases; Anthropic Messages carries that content in the - top-level ``system`` field. Move any such entries into ``system`` before - the Invoke request is built.""" + @staticmethod + def _is_system_role_message(message: Any) -> bool: + return isinstance(message, dict) and message.get("role") == "system" + + def _normalize_system_role_messages_for_bedrock(self, anthropic_messages_request: dict, model: str) -> None: + """Bedrock Invoke validates ``role: "system"`` entries inside ``messages`` + per model. Models carrying ``supports_mid_conversation_system`` in the + cost map (the Opus 4.8 family) only reject a leading run ("messages.0: + use the top-level 'system' parameter for the initial system prompt") and + accept mid-conversation entries (e.g. Claude Code's + ``mid-conversation-system-2026-04-07`` reminders) in place, where they + MUST stay: hoisting one mutates the ``system`` prefix and invalidates the + prompt cache for the entire message history. Older Claude models (Opus + 4.7, Sonnet 4.6, Haiku 4.5, ...) reject the role in every position + ("role 'system' is not supported on this model"), so without the flag + every system entry is hoisted into the top-level ``system`` field. + Billing-header system blocks are stripped from the top-level ``system`` + field regardless of whether anything was hoisted.""" messages = anthropic_messages_request.get("messages") if not isinstance(messages, list): return - system_role_messages = [m for m in messages if isinstance(m, dict) and m.get("role") == "system"] - if not system_role_messages: - return - - anthropic_messages_request["messages"] = [ - m for m in messages if not (isinstance(m, dict) and m.get("role") == "system") - ] + if _supports_factory( + model=model, + custom_llm_provider="bedrock", + key="supports_mid_conversation_system", + ): + leading_count = next( + (i for i, m in enumerate(messages) if not self._is_system_role_message(m)), + len(messages), + ) + hoisted = messages[:leading_count] + remaining = messages[leading_count:] + else: + hoisted = [m for m in messages if self._is_system_role_message(m)] + remaining = [m for m in messages if not self._is_system_role_message(m)] + if hoisted: + anthropic_messages_request["messages"] = remaining system_content = [ block for source in ( anthropic_messages_request.get("system"), - *(m.get("content") for m in system_role_messages), + *(m.get("content") for m in hoisted), ) for block in self._as_system_content_blocks(source) ] @@ -669,7 +691,7 @@ class AmazonAnthropicClaudeMessagesConfig( litellm_params=litellm_params, headers=headers, ) - self._normalize_system_role_messages_for_bedrock(anthropic_messages_request) + self._normalize_system_role_messages_for_bedrock(anthropic_messages_request, model=model) ######################################################### ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### ######################################################### diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2a269d693ec..6ecadaf1fe6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1359,6 +1359,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1393,6 +1394,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1427,6 +1429,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1461,6 +1464,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1481,6 +1485,7 @@ "anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1516,6 +1521,7 @@ "global.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1551,6 +1557,7 @@ "us.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1586,6 +1593,7 @@ "eu.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1621,6 +1629,43 @@ "au.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true + }, + "jp.anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1703,6 +1748,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1737,6 +1783,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1771,6 +1818,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1805,6 +1853,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1839,6 +1888,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1873,6 +1923,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -44998,6 +45049,17 @@ }, "fallback_generalizations": { "rules": [ + { + "name": "bedrock-anthropic-claude-mid-conversation-system", + "pattern": "anthropic\\.claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "description": "Bedrock Invoke ids for Claude 4.8 or higher: anthropic.claude- with minor 4.8 through 4.99, any 5.x or later major-minor, or a bare 5+ major, which also admits new families such as fable. These models accept mid-conversation role system messages in place (verified live on Opus 4.8, Sonnet 5 and Fable 5), so unmapped future Bedrock Claudes keep the cache-preserving in-place handling instead of the hoist-all default. Listed first so bare-id provider inference, which takes the first pattern hit, resolves these Bedrock ids to bedrock; model-info resolution skips provider-mismatched rules either way.", + "extends": "anthropic-claude", + "model_info": { + "litellm_provider": "bedrock", + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true + } + }, { "name": "anthropic-claude-adaptive-thinking", "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index c7b32614c61..97cefb3f2cb 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -76,6 +76,34 @@ _TOKEN_EXCHANGE_COLUMN_FIELDS: frozenset = frozenset( } ) +# The client-forwarded token modes share one stored-credential shape: the admin-declared upstream +# OAuth app (client_id/client_secret) plus the same authorize relay, and neither mints anything the +# gateway keeps. So a switch WITHIN this class must preserve the stored app, unlike a cross-class +# switch (e.g. an oauth2 row whose client may be DCR-minted and is not reusable elsewhere). +_CLIENT_FORWARDED_AUTH_TYPES: frozenset = frozenset({"true_passthrough", "oauth_delegate"}) + +# Minted token material that must never survive a client rotation on a persisted row. +_MINTED_TOKEN_CREDENTIAL_FIELDS: frozenset = frozenset({"access_token", "refresh_token", "expires_in"}) + + +def _credential_auth_class(auth_type: Optional[str]) -> Optional[str]: + """Collapse the client-forwarded modes to one credential class; every other auth_type is its own + class. Used so credential handling keys off whether the stored-credential shape actually changed, + not off a raw auth_type inequality that treats true_passthrough<->oauth_delegate as a full reset.""" + if auth_type in _CLIENT_FORWARDED_AUTH_TYPES: + return "client_forwarded" + return auth_type + + +def _drop_stale_minted_on_client_rotation(merged: Dict[str, Any], new_creds: Dict[str, Any]) -> Dict[str, Any]: + """When the update rotates the client, drop stale minted token keys it did not itself set, so an old + app's access/refresh token never rides forward under the new client. A no-op when no client key changed.""" + if "client_id" not in new_creds and "client_secret" not in new_creds: + return merged + return { + key: value for key, value in merged.items() if key not in _MINTED_TOKEN_CREDENTIAL_FIELDS or key in new_creds + } + def _is_global_env_var_scope(scope: Any) -> bool: """``scope="user"`` entries are placeholders the user fills in; everything @@ -679,7 +707,9 @@ async def update_mcp_server( existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) auth_type_changed = bool( - data.auth_type and existing and existing.auth_type is not None and existing.auth_type != data.auth_type + data.auth_type + and existing + and _credential_auth_class(existing.auth_type) != _credential_auth_class(data.auth_type) ) # Clear stale credentials when auth_type changes but no new credentials provided @@ -712,11 +742,12 @@ async def update_mcp_server( # would wipe encrypted secrets that the UI cannot display back. if "credentials" in data_dict and data_dict["credentials"] is not None: if existing and existing.credentials: - # Only merge when auth_type is unchanged. Switching auth types - # (e.g. oauth2 → api_key) should replace credentials entirely - # to avoid stale secrets from the previous auth type lingering. - auth_type_unchanged = data.auth_type is None or data.auth_type == existing.auth_type - if auth_type_unchanged: + # Only merge when the credential CLASS is unchanged. A cross-class switch + # (e.g. oauth2 → api_key, or oauth2 → true_passthrough) replaces credentials + # entirely to avoid stale secrets from the previous class lingering; a switch + # within the client-forwarded class (true_passthrough ↔ oauth_delegate) keeps + # the same declared app and so must merge, not replace. + if not auth_type_changed: existing_creds = ( json.loads(existing.credentials) if isinstance(existing.credentials, str) @@ -727,8 +758,9 @@ async def update_mcp_server( if isinstance(data_dict["credentials"], str) else dict(data_dict["credentials"]) ) - # New values override existing; existing keys not in update are preserved - merged = {**existing_creds, **new_creds} + # New values override existing; existing keys not in update are preserved. A client + # rotation additionally drops the previous app's stale minted token keys. + merged = _drop_stale_minted_on_client_rotation({**existing_creds, **new_creds}, new_creds) # Migrate-on-write for legacy rows: token-exchange settings the # old blob shape carried move to their dedicated columns (unless # the caller set the column this update, or the row already has @@ -747,6 +779,14 @@ async def update_mcp_server( # Add audit fields data_dict["updated_by"] = touched_by + # prisma-python rejects a raw ``None`` for a ``Json?`` field ("value is required but not set"); the + # clear paths above use ``None`` as the merge-skip sentinel, so translate it here to ``Json(None)``, + # which writes SQL null and reads back as ``None``. Done at the edge so the merge guards stay simple. + if "credentials" in data_dict and data_dict["credentials"] is None: + from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools + + data_dict["credentials"] = Json(None) + updated_mcp_server = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, data=data_dict, # type: ignore diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py new file mode 100644 index 00000000000..5530fbc46fd --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -0,0 +1,169 @@ +"""Producer and consumer helpers for the DCR-bridge ``oauth_delegate`` envelope. + +A DCR-bridge ``oauth_delegate`` client presents ONE bearer that is a litellm-signed +envelope (see :mod:`.envelope`) carrying both a litellm identity and the upstream OAuth +token. The gateway token endpoint mints it (producer) at OAuth issuance, and at the MCP +admission edge the gateway derives the envelope keys from the proxy ``master_key``, opens +it, admits the request under the recovered identity, and forwards the inner upstream token +to the upstream MCP server (consumer). This module is the pure surface for both sides; the +token-endpoint and admission wiring live in their respective call sites. +""" + +import hashlib +from datetime import datetime +from functools import lru_cache +from typing import Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeIdentity, + EnvelopeKeys, + EnvelopeMintError, + OpenedEnvelope, + SealedEnvelope, + UpstreamTokenGrant, + is_envelope, + mint_envelope, + open_envelope, +) + +_SIGNING_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-signing:" +_ENCRYPTION_KEY_DOMAIN = b"litellm-mcp-bridge:envelope-encryption:" + +# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation, which +# makes offline guessing of a candidate master key memory-hard rather than a bare hash comparison. +_SCRYPT_N = 2**15 +_SCRYPT_R = 8 +_SCRYPT_P = 1 +# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales +# with every work factor and a future p or r bump does not trip "memory limit exceeded". +_SCRYPT_MAXMEM = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 +_DERIVED_KEY_BYTES = 32 + + +@lru_cache(maxsize=8) +def envelope_keys_from_master_key(master_key: str) -> EnvelopeKeys: + """Derive the envelope signing and encryption keys from the proxy master key. + + A memory-hard scrypt KDF (RFC 7914) over two distinct domain-label salts yields two + independent 256-bit subkeys from the one secret, so the producer (mint) and consumer + (open) agree on keys without persisting any. scrypt is used rather than a bare hash or + HMAC so that a captured envelope is not a cheap offline oracle for the master key: each + candidate guess costs a full memory-hard derivation, which is what protects a deployment + whose master key is weaker than it should be. The result is cached (the master key is + fixed for a process), so the KDF runs once per key and adds nothing to the per-request + admission path. The derivation is deterministic; rotating ``master_key`` invalidates + every outstanding envelope, which is the intended behavior for a signing-key change. + """ + signing = hashlib.scrypt( + master_key.encode(), + salt=_SIGNING_KEY_DOMAIN, + n=_SCRYPT_N, + r=_SCRYPT_R, + p=_SCRYPT_P, + maxmem=_SCRYPT_MAXMEM, + dklen=_DERIVED_KEY_BYTES, + ).hex() + encryption = hashlib.scrypt( + master_key.encode(), + salt=_ENCRYPTION_KEY_DOMAIN, + n=_SCRYPT_N, + r=_SCRYPT_R, + p=_SCRYPT_P, + maxmem=_SCRYPT_MAXMEM, + dklen=_DERIVED_KEY_BYTES, + ).hex() + return EnvelopeKeys(signing_key=SecretStr(signing), encryption_key=SecretStr(encryption)) + + +def build_bridge_token_response( + identity: EnvelopeIdentity, + grant: UpstreamTokenGrant, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``grant`` for ``identity`` into the client-held bearer the token endpoint returns. + + The producer mirror of :func:`resolve_bridge_envelope`: a thin, pure wrapper over + :func:`mint_envelope` that returns the sealed envelope, or the mint error as a value + (an oversized grant) for the caller to map onto an OAuth error response. + """ + return mint_envelope(identity, grant, keys, now) + + +class NotBridgeEnvelope(BaseModel): + """The bearer is not an envelope; admission continues on its normal path.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_bridge_envelope"] = "not_bridge_envelope" + + +class BridgeEnvelopeAdmitted(BaseModel): + """A valid envelope: the identity to admit under and the full upstream ``Authorization`` + value (``token_type access_token``) to forward to the upstream MCP server.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["admitted"] = "admitted" + identity: EnvelopeIdentity + upstream_authorization: SecretStr + + +class BridgeEnvelopeInvalid(BaseModel): + """The bearer is envelope-shaped but did not open (expired, tampered, wrong key); + admission must fail closed rather than fall through to normal validation.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["invalid"] = "invalid" + + +BridgeEnvelopeResult: TypeAlias = NotBridgeEnvelope | BridgeEnvelopeAdmitted | BridgeEnvelopeInvalid + + +def _strip_bearer(value: str) -> str: + parts = value.split(None, 1) + if len(parts) == 2 and parts[0].lower() == "bearer": + return parts[1] + return value + + +def is_bridge_envelope_shaped(authorization_value: str) -> bool: + """Cheap, keyless test that an ``Authorization`` value carries an envelope (optional + ``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an + envelope, so a plain upstream bearer falls through to normal oauth2 admission.""" + return is_envelope(_strip_bearer(authorization_value)) + + +def resolve_bridge_envelope( + authorization_value: str, + keys: EnvelopeKeys, + now: datetime, + expected_server_id: str, +) -> BridgeEnvelopeResult: + """Classify an ``Authorization`` value presented to a bridge ``oauth_delegate`` server. + + Strips an optional ``Bearer`` scheme, then returns ``NotBridgeEnvelope`` for a + non-envelope bearer (normal admission continues), ``BridgeEnvelopeAdmitted`` with the + recovered identity and the upstream ``Authorization`` value to forward for a valid + envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not + open. Never raises: it is total over hostile input via :func:`open_envelope`. + + ``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an + opened envelope whose sealed ``server_id`` does not match is rejected as + ``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents + replaying an envelope minted for one server against another, which would forward the + first server's upstream credential across a server boundary. ``server_id`` is not a + secret (the caller targets that server), so a plain equality check is sufficient and, + unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id. + """ + candidate = _strip_bearer(authorization_value) + if not is_envelope(candidate): + return NotBridgeEnvelope() + opened = open_envelope(candidate, keys, now) + if not isinstance(opened, OpenedEnvelope): + return BridgeEnvelopeInvalid() + if opened.identity.server_id != expected_server_id: + return BridgeEnvelopeInvalid() + grant = opened.grant + upstream_authorization = f"{grant.token_type} {grant.access_token.get_secret_value()}" + return BridgeEnvelopeAdmitted(identity=opened.identity, upstream_authorization=SecretStr(upstream_authorization)) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py new file mode 100644 index 00000000000..298bc8d98cc --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -0,0 +1,358 @@ +"""Client-held sealed envelope for the oauth_delegate DCR bridge. + +A DCR-bridge client holds ONE bearer that must carry BOTH a litellm identity and the +upstream OAuth grant, with zero server-side storage. The gateway token endpoint mints a +litellm-signed envelope (:func:`mint_envelope`); the MCP edge validates it, recovers the +identity claims and the inner upstream grant (:func:`open_envelope`), and forwards the +inner access token upstream. This module is pure and unwired: it imports nothing from +endpoint or edge code, reads no proxy globals, and takes all key material and the clock +as explicit parameters. + +Wire shape: ``llm_env_`` + an HS256 JWT (same signing approach as the BYOK session +bearer in ``byok_oauth_endpoints.py``). Registered claims are ``iss``/``iat``/``exp``; +custom claims are ``user_id``, ``server_id``, and ``grant``, where ``grant`` is the +upstream token grant serialized to JSON, encrypted with the repo's symmetric +encryption helpers (``encrypt_value``/``decrypt_value`` from +``encrypt_decrypt_utils`` — the same family ``encrypt_value_helper`` applies to +persisted DCR credentials), and base64url-encoded, so the inner token never appears +in plaintext anywhere in the envelope. + +Failures are values: :func:`open_envelope` returns one of the frozen +``EnvelopeOpenError`` variants (discriminated on ``tag``) for invalid, expired, +tampered, or undecryptable input, and :func:`mint_envelope` returns +``EnvelopeTooLarge`` for oversized grants. Error values carry tags and sizes only, +never token material. + +The pydantic input models reject programmer errors at construction (e.g. a +non-positive ``expires_in`` or an empty required field). :func:`open_envelope` is +additionally total over hostile, attacker-controlled input: it never raises, only +returns an ``EnvelopeOpenError``. :func:`mint_envelope` operates on a +gateway-supplied grant (an upstream IdP's UTF-8 JSON token response), so it does not +defend against non-UTF-8 field content that cannot survive JSON parsing; its only +value-typed failure is ``EnvelopeTooLarge``. +""" + +from __future__ import annotations + +import base64 +from datetime import datetime, timedelta +from typing import Literal, TypeAlias + +import jwt +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value + +ENVELOPE_PREFIX = "llm_env_" +"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope +from a raw upstream token before doing any cryptography.""" + +ENVELOPE_ISSUER = "litellm-mcp-bridge" +"""``iss`` claim stamped into every envelope and required back on open.""" + +MAX_ENVELOPE_TTL_SECONDS = 3600 +"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)`` +(the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the +BYOK session bearer this module's signing approach is borrowed from: a client-held +credential should never outlive a bounded window even when the upstream token does.""" + +MAX_ENVELOPE_BYTES = 12288 +"""Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs +commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the +envelope, and common proxy/server header limits sit around 16KB total. 12288 leaves +comfortable headroom for a large upstream token while keeping the envelope safely +transmittable as a single Authorization header. Oversized grants are rejected with a +typed error, never truncated.""" + +_ENVELOPE_JWT_ALGORITHM = "HS256" + + +class EnvelopeIdentity(BaseModel): + """The litellm identity the envelope binds the inner grant to.""" + + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + server_id: str = Field(min_length=1) + + +class UpstreamTokenGrant(BaseModel): + """The upstream OAuth token response fields sealed inside the envelope. + + ``expires_in`` must be positive when present; a non-positive value is a programmer + error rejected at construction. Token fields are ``SecretStr`` so reprs never leak + them. + """ + + model_config = ConfigDict(frozen=True) + access_token: SecretStr = Field(min_length=1) + token_type: str = Field(min_length=1) + refresh_token: SecretStr | None = None + scope: str | None = None + expires_in: int | None = Field(default=None, gt=0) + + +class EnvelopeKeys(BaseModel): + """Injected key material: the HS256 signing key and the symmetric encryption key. + + ``signing_key`` must be at least 32 bytes: HS256's HMAC-SHA256 has a 256-bit + security level, RFC 7518 requires a key of at least that size, and a shorter key + makes PyJWT emit ``InsecureKeyLengthWarning``. + """ + + model_config = ConfigDict(frozen=True) + signing_key: SecretStr = Field(min_length=32) + encryption_key: SecretStr = Field(min_length=1) + + +class SealedEnvelope(BaseModel): + """A minted envelope: the client-held bearer value and when it expires.""" + + model_config = ConfigDict(frozen=True) + token: SecretStr + expires_at: datetime + + +class OpenedEnvelope(BaseModel): + """A validated envelope: the identity it was minted for and the recovered grant.""" + + model_config = ConfigDict(frozen=True) + identity: EnvelopeIdentity + grant: UpstreamTokenGrant + + +class EnvelopeTooLarge(BaseModel): + """The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["envelope_too_large"] = "envelope_too_large" + size_bytes: int + max_bytes: int + + +EnvelopeMintError: TypeAlias = EnvelopeTooLarge + + +class NotAnEnvelope(BaseModel): + """The candidate does not carry the envelope prefix.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_an_envelope"] = "not_an_envelope" + + +class BadSignature(BaseModel): + """The JWT signature does not verify under the provided signing key.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["bad_signature"] = "bad_signature" + + +class Expired(BaseModel): + """The envelope's ``exp`` is not in the future relative to the provided ``now``.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["expired"] = "expired" + + +class MalformedPayload(BaseModel): + """The token is not a well-formed envelope: undecodable JWT, wrong issuer, missing + or mistyped claims, or a decrypted grant that fails validation.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["malformed_payload"] = "malformed_payload" + + +class DecryptFailed(BaseModel): + """The signed ``grant`` blob could not be decrypted under the provided key.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["decrypt_failed"] = "decrypt_failed" + + +EnvelopeOpenError: TypeAlias = NotAnEnvelope | BadSignature | Expired | MalformedPayload | DecryptFailed + + +class _EnvelopeClaims(BaseModel): + """Decoded-claims boundary that pins the exact shape :func:`mint_envelope` emits. + + ``user_id``/``server_id`` mirror the ``min_length`` constraints of + :class:`EnvelopeIdentity` so any claim set that validates here also constructs an + identity, keeping :func:`open_envelope` raise-free: a correctly signed JWT with an + empty identity claim fails here and maps to ``MalformedPayload``. + + ``strict`` rejects coerced types (``exp: "123"``, ``exp: 123.0``) rather than opening + on them, and ``extra="forbid"`` rejects any claim the gateway never mints (a hostile + ``nbf``/``aud``/... rides along on a re-signed token). Since PyJWT's own ``iat``/ + ``nbf``/``exp`` validators are disabled at decode (they raise on hostile claim types + and, for ``iat``/``nbf``, compare against the wall clock rather than the injected + ``now``), this model is the sole, total type gate for every registered claim. + """ + + model_config = ConfigDict(frozen=True, strict=True, extra="forbid") + iss: str + iat: int + exp: int + user_id: str = Field(min_length=1) + server_id: str = Field(min_length=1) + grant: str = Field(min_length=1) + + +class _GrantWire(BaseModel): + model_config = ConfigDict(frozen=True) + access_token: str + token_type: str + refresh_token: str | None = None + scope: str | None = None + expires_in: int | None = None + + +def is_envelope(candidate: str) -> bool: + """Cheap prefix check so the edge can route envelopes vs raw tokens without crypto.""" + return candidate.startswith(ENVELOPE_PREFIX) + + +def mint_envelope( + identity: EnvelopeIdentity, + grant: UpstreamTokenGrant, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``grant`` for ``identity`` into a client-held envelope. + + ``exp`` is ``min(grant.expires_in, MAX_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` + (the cap alone when ``expires_in`` is absent). Returns ``EnvelopeTooLarge`` when the + serialized envelope exceeds ``MAX_ENVELOPE_BYTES``. + """ + expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in)) + claims = _EnvelopeClaims( + iss=ENVELOPE_ISSUER, + iat=int(now.timestamp()), + exp=int(expires_at.timestamp()), + user_id=identity.user_id, + server_id=identity.server_id, + grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), + ) + token = ENVELOPE_PREFIX + jwt.encode( + claims.model_dump(), + keys.signing_key.get_secret_value(), + algorithm=_ENVELOPE_JWT_ALGORITHM, + ) + size_bytes = len(token.encode("utf-8")) + if size_bytes > MAX_ENVELOPE_BYTES: + return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES) + return SealedEnvelope(token=SecretStr(token), expires_at=expires_at) + + +def open_envelope( + candidate: str, + keys: EnvelopeKeys, + now: datetime, +) -> OpenedEnvelope | EnvelopeOpenError: + """Validate ``candidate`` and recover the identity and inner grant. + + Never raises for bad input: every invalid, expired, tampered, or undecryptable + candidate maps to a distinct ``EnvelopeOpenError`` variant. The recovered + ``grant.expires_in`` is the value the upstream reported at mint time and is not + re-derived, so it is stale by up to the envelope's lifetime; callers that need a + live remaining lifetime should use ``now`` against the upstream, not this field. + """ + if not is_envelope(candidate): + return NotAnEnvelope() + # UTF-8 byte length is never below character length, so a character count already over the + # cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then + # runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters. + if len(candidate) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key) + if not isinstance(claims, _EnvelopeClaims): + return claims + if now.timestamp() >= claims.exp: + return Expired() + grant = _decrypt_grant(claims.grant, keys.encryption_key) + if not isinstance(grant, UpstreamTokenGrant): + return grant + return OpenedEnvelope( + identity=EnvelopeIdentity(user_id=claims.user_id, server_id=claims.server_id), + grant=grant, + ) + + +def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int: + if upstream_expires_in is None: + return MAX_ENVELOPE_TTL_SECONDS + return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS) + + +def _grant_plaintext(grant: UpstreamTokenGrant) -> str: + wire = _GrantWire( + access_token=grant.access_token.get_secret_value(), + token_type=grant.token_type, + refresh_token=None if grant.refresh_token is None else grant.refresh_token.get_secret_value(), + scope=grant.scope, + expires_in=grant.expires_in, + ) + return wire.model_dump_json(exclude_none=True) + + +def _decode_claims( + compact: str, + signing_key: SecretStr, +) -> _EnvelopeClaims | BadSignature | MalformedPayload: + """Verify the HS256 signature and shape of an attacker-controlled compact JWT. + + ``compact`` is fully hostile and bounded to ``MAX_ENVELOPE_BYTES`` by the caller. + PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim + types and, for ``iat``/``nbf``, compare against the wall clock rather than the + injected ``now`` (``exp`` is checked by the caller against ``now``). Apart from a + signature mismatch (``BadSignature``), every decode failure is ``MalformedPayload``: + a non-UTF-8 candidate surfaces as ``UnicodeEncodeError`` (a ``ValueError``), a + non-string registered claim such as ``iss`` as a ``TypeError`` from PyJWT's claim + validators, and a wrong issuer or structurally invalid token as an + ``InvalidTokenError``. ``_EnvelopeClaims`` is the total type gate for the payload. + """ + try: + payload = jwt.decode( + compact, + signing_key.get_secret_value(), + algorithms=[_ENVELOPE_JWT_ALGORITHM], + issuer=ENVELOPE_ISSUER, + options={ + "verify_exp": False, + "verify_iat": False, + "verify_nbf": False, + "require": ["iss", "iat", "exp"], + }, + ) + except jwt.InvalidSignatureError: + return BadSignature() + except (jwt.InvalidTokenError, ValueError, TypeError): + return MalformedPayload() + try: + return _EnvelopeClaims.model_validate(payload) + except ValidationError: + return MalformedPayload() + + +def _encrypt_grant_blob(plaintext: str, encryption_key: SecretStr) -> str: + ciphertext = bytes(encrypt_value(value=plaintext, signing_key=encryption_key.get_secret_value())) + return base64.urlsafe_b64encode(ciphertext).decode("ascii") + + +def _decrypt_grant( + blob: str, + encryption_key: SecretStr, +) -> UpstreamTokenGrant | DecryptFailed | MalformedPayload: + from nacl.exceptions import CryptoError + + try: + plaintext = decrypt_value( + value=base64.urlsafe_b64decode(blob), + signing_key=encryption_key.get_secret_value(), + ) + except (CryptoError, ValueError): + return DecryptFailed() + try: + return UpstreamTokenGrant.model_validate_json(plaintext) + except ValidationError: + return MalformedPayload() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 43ddf302692..f441a1d3f84 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2151,6 +2151,41 @@ class PluginConfig(LiteLLMPydanticObjectBase): ) +class CoordinationRedisNode(LiteLLMPydanticObjectBase): + """A single startup node of a cluster-mode Redis used for proxy coordination.""" + + host: str = Field(description="hostname of the cluster node") + port: int = Field(description="port of the cluster node") + + +class CoordinationRedisParams(LiteLLMPydanticObjectBase): + """ + Connection params for the proxy's coordination Redis (cross-pod tpm/rpm rate + limits, spend tracking, pod lock manager, shared health checks), configured + independently of the response-cache backend in `litellm_settings.cache_params`. + """ + + model_config = ConfigDict(extra="allow", protected_namespaces=()) + + host: Optional[str] = Field(None, description="Redis hostname") + port: Optional[int] = Field(None, description="Redis port") + password: Optional[str] = Field(None, description="Redis password") + username: Optional[str] = Field(None, description="Redis username") + url: Optional[str] = Field(None, description="full Redis connection url, e.g. redis://:pass@host:6379") + ssl: Optional[bool] = Field(None, description="connect over TLS") + startup_nodes: Optional[List[CoordinationRedisNode]] = Field( + None, description="cluster-mode startup nodes; when set a cluster client is used" + ) + sentinel_nodes: Optional[List[List[Union[str, int]]]] = Field( + None, description="sentinel [host, port] pairs; when set a sentinel-managed client is used" + ) + sentinel_password: Optional[str] = Field(None, description="password for the sentinel nodes") + service_name: Optional[str] = Field(None, description="sentinel service name") + + def has_connection_target(self) -> bool: + return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes)) + + class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ Documents all the fields supported by `general_settings` in config.yaml @@ -2166,6 +2201,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): use_google_kms: Optional[bool] = Field(None, description="decrypt keys with google kms") use_azure_key_vault: Optional[bool] = Field(None, description="load keys from azure key vault") master_key: Optional[str] = Field(None, description="require a key for all calls to proxy") + coordination_redis: Optional[CoordinationRedisParams] = Field( + None, + description=( + "standalone Redis for cross-pod coordination (tpm/rpm rate limits, " + "spend tracking, pod lock manager, shared health checks), configured " + "independently of the response-cache backend; takes precedence over " + "borrowing the `cache_params` Redis and over the REDIS_* env fallback" + ), + ) allow_cli_sso_verification_uri_complete: bool | None = Field( None, description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine", diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 43b64aebd3b..aadaff61238 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -76,6 +76,10 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None: """LiteLLM Proxy CLI - Manage your LiteLLM proxy server""" ctx.ensure_object(dict) + # Normalize once here so every downstream command (login, agents, http, ...) can safely + # do f"{base_url}/some/path" without producing a double slash. + base_url = base_url.rstrip("/") + # If no API key provided via flag or environment variable, try to load from saved token. # Pass base_url so we only use the stored key when it was issued for this server. if api_key is None: diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index 8cde91e32b2..445257a1e01 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -35,15 +35,49 @@ class ComplianceChecker: If a guardrail doesn't have a mode specified, it's treated as pre-call (the most common case). """ - result = [] - for g in self.guardrails: - g_mode = g.get("guardrail_mode") - # If no mode specified, default to pre_call - if g_mode is None and mode == "pre_call": - result.append(g) - elif g_mode == mode: - result.append(g) - return result + return [g for g in self.guardrails if self._mode_matches(g.get("guardrail_mode"), mode)] + + @staticmethod + def _mode_matches(g_mode: object, mode: str) -> bool: + """ + Return True only when a guardrail with logged ``guardrail_mode`` of + ``g_mode`` is guaranteed to have run in ``mode`` for the audited request. + + ``guardrail_mode`` in a spend log can take several shapes because + ``LitellmParams.mode`` is typed ``Union[str, List[str], Mode]``, and + when the event type cannot be inferred at write time the raw config is + logged verbatim. The spend log records the configured mode(s), not the + concrete hook that fired for a given request; a match reports a mode + satisfied only when every configured branch runs in that mode, so True + never claims a hook the guardrail may not have actually executed. + + Fails safe: if the guarantee cannot be established (missing default, + divergent per-tag override, or a list that runs in more than one mode), + the guardrail counts for no mode. The precise fix is to log the + resolved event mode and match on it; this is the safe interim. + """ + if g_mode is None: + return mode == "pre_call" + if isinstance(g_mode, str): + return g_mode == mode + if isinstance(g_mode, (list, tuple)): + return bool(g_mode) and all(m == mode for m in g_mode) + if isinstance(g_mode, dict): + default = g_mode.get("default") + if default is None: + return False + tags = g_mode.get("tags") + tag_branches = list(tags.values()) if isinstance(tags, dict) else [] + + def _branch_runs_in_mode(branch: object) -> bool: + if isinstance(branch, str): + return branch == mode + if isinstance(branch, (list, tuple)): + return bool(branch) and all(m == mode for m in branch) + return False + + return all(_branch_runs_in_mode(branch) for branch in [default, *tag_branches]) + return False def _has_guardrail_intervention(self, guardrails: List[Dict]) -> bool: """Check if any guardrail intervened (blocked/masked content).""" diff --git a/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml b/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml index b353924f000..eb091cc72c5 100644 --- a/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml +++ b/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml @@ -5,8 +5,9 @@ model_list: api_key: my-fake-key api_base: os.environ/FAKE_OPENAI_API_BASE -litellm_settings: - cache: True - cache_params: - type: redis +general_settings: + coordination_redis: + host: os.environ/REDIS_HOST + port: os.environ/REDIS_PORT + password: os.environ/REDIS_PASSWORD diff --git a/litellm/proxy/example_config_yaml/spend_tracking_config.yaml b/litellm/proxy/example_config_yaml/spend_tracking_config.yaml index 60adadbd8d4..d66fd5fa601 100644 --- a/litellm/proxy/example_config_yaml/spend_tracking_config.yaml +++ b/litellm/proxy/example_config_yaml/spend_tracking_config.yaml @@ -7,9 +7,6 @@ model_list: general_settings: use_redis_transaction_buffer: true - -litellm_settings: - cache: True - cache_params: - type: redis - supported_call_types: [] \ No newline at end of file + coordination_redis: + host: os.environ/REDIS_HOST + port: os.environ/REDIS_PORT \ No newline at end of file diff --git a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py new file mode 100644 index 00000000000..7ab4e3019c3 --- /dev/null +++ b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py @@ -0,0 +1,431 @@ +""" +COORDINATION REDIS SETTINGS MANAGEMENT + +Endpoints for managing `general_settings.coordination_redis` - the standalone +Redis the proxy uses for cross-pod coordination (tpm/rpm rate limits, spend +tracking, pod lock manager, shared health checks), configured independently of +the response-cache backend. + +GET /coordination_redis/settings - Get the coordination Redis settings, field metadata, and which source is active +POST /coordination_redis/settings - Save coordination Redis settings to the database +POST /coordination_redis/settings/test - Test a coordination Redis connection with the provided credentials +""" + +import asyncio +import json +from collections.abc import Mapping +from contextlib import suppress +from datetime import datetime, timezone +from typing import Optional + +from fastapi import APIRouter, Depends, Header, HTTPException +from pydantic import BaseModel, Field, TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.caching.caching import RedisCache +from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.proxy._types import ( + AUDIT_ACTIONS, + CoordinationRedisParams, + LiteLLM_AuditLogs, + LitellmTableNames, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import invalidate_config_param +from litellm.repositories.config_repository import ConfigRepository +from litellm.secret_managers.main import get_secret_str +from litellm.types.management_endpoints import ( + COORDINATION_REDIS_SETTINGS_FIELDS, + CoordinationRedisSettingsField, + CoordinationRedisSource, +) + +router = APIRouter() + +_GENERAL_SETTINGS_PARAM_NAME = "general_settings" +_COORDINATION_REDIS_KEY = "coordination_redis" + +# Fields that carry credentials. Redacted on read so a plaintext Redis / +# Sentinel password never leaves the server, and scrubbed out of connection-test +# error strings. `url` is here because a Redis url can embed a password inline +# (e.g. redis://:secret@host:6379/1). +_SENSITIVE_FIELDS: frozenset[str] = frozenset({"password", "sentinel_password", "url"}) + +_REDACTED_VALUE = "***REDACTED***" + +_ENV_REF_PREFIX = "os.environ/" + +_PING_TIMEOUT_SECONDS = 5.0 + +_SETTINGS_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +def _enforce_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins can manage coordination Redis settings"}, + ) + + +def _resolve_env_ref(value: object) -> object: + """Resolve an `os.environ/VAR` reference to its value, passing anything else through.""" + if isinstance(value, str) and value.startswith(_ENV_REF_PREFIX): + return get_secret_str(value) + return value + + +def _resolve_env_refs(settings: Mapping[str, object]) -> dict[str, object]: + return {key: _resolve_env_ref(value) for key, value in settings.items()} + + +def _redact_credentials(settings: Mapping[str, object]) -> dict[str, object]: + """Replace credential-bearing values with a fixed marker, keeping the rest intact.""" + return { + key: (_REDACTED_VALUE if key in _SENSITIVE_FIELDS and value is not None else value) + for key, value in settings.items() + } + + +def _redact_all_values(settings: Optional[Mapping[str, object]]) -> dict[str, object]: + """Replace every value with a fixed marker, preserving the key set. + + The audit row shows *which* fields changed without the audit table becoming + a credential-harvest sink. + """ + if not settings: + return {} + return {key: _REDACTED_VALUE for key in settings} + + +def _credential_values(settings: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + str(value) for key, value in settings.items() if key in _SENSITIVE_FIELDS and isinstance(value, (str, int)) + ) + + +def _scrub_credentials(message: str, settings: Mapping[str, object]) -> str: + """Strip any credential value the caller supplied out of an error string. + + Redis client errors routinely echo the connection url (password inline) or + the auth error back to the caller. + """ + scrubbed = message + for secret in _credential_values(settings): + if secret: + scrubbed = scrubbed.replace(secret, _REDACTED_VALUE) + return scrubbed + + +def _merge_over_saved( + incoming: Mapping[str, object], + saved: Mapping[str, object], +) -> dict[str, object]: + """Restore the real credential behind every value the caller echoed back redacted. + + GET returns credentials as ``***REDACTED***``; an admin who edits the + non-secret fields and re-submits would otherwise test (and save) the marker + as the password. + """ + return { + key: (saved[key] if value == _REDACTED_VALUE and key in saved else value) for key, value in incoming.items() + } + + +def _validated_params(settings: Mapping[str, object]) -> CoordinationRedisParams: + """Validate settings the way startup does: resolve env refs, then require a connection target.""" + try: + params = CoordinationRedisParams(**_resolve_env_refs(settings)) + except ValidationError as e: + invalid_fields = sorted({str(error["loc"][0]) for error in e.errors() if error["loc"]}) + raise HTTPException( + status_code=400, + detail={"error": f"Invalid coordination_redis settings for fields: {invalid_fields}"}, + ) + if not params.has_connection_target(): + raise HTTPException( + status_code=400, + detail={ + "error": ( + "coordination_redis needs a connection target: " + "set one of host, url, startup_nodes, or sentinel_nodes" + ) + }, + ) + return params + + +async def _read_general_settings() -> dict[str, object]: + """Read the persisted `general_settings` config row (empty when unset or no DB).""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return {} + config_param = await ConfigRepository(prisma_client).get_param(_GENERAL_SETTINGS_PARAM_NAME) + if config_param is None or config_param.param_value is None: + return {} + return _SETTINGS_ADAPTER.validate_python(config_param.param_value) + + +async def get_persisted_coordination_redis_settings() -> Optional[dict[str, object]]: + """The coordination_redis block saved to the database, if any. + + Read at startup so settings saved from the admin UI take effect on the next + boot, and used here so a read reports what the proxy would boot with. + """ + persisted = (await _read_general_settings()).get(_COORDINATION_REDIS_KEY) + if isinstance(persisted, dict): + return _SETTINGS_ADAPTER.validate_python(persisted) + return None + + +async def _current_coordination_redis_settings() -> Optional[dict[str, object]]: + """The coordination_redis block the proxy would boot with. + + The persisted row wins over the yaml-loaded config state because startup + applies the DB `general_settings` row over the file config. + """ + from litellm.proxy.proxy_server import proxy_config + + persisted = await get_persisted_coordination_redis_settings() + if persisted is not None: + return persisted + + config_state = _SETTINGS_ADAPTER.validate_python(proxy_config.get_config_state()) + general_settings = config_state.get(_GENERAL_SETTINGS_PARAM_NAME) + if not isinstance(general_settings, dict): + return None + from_file = general_settings.get(_COORDINATION_REDIS_KEY) + if isinstance(from_file, dict): + return _SETTINGS_ADAPTER.validate_python(from_file) + return None + + +def _coordination_redis_source(settings: Optional[Mapping[str, object]]) -> Optional[CoordinationRedisSource]: + """Which source the proxy's coordination Redis comes from, in startup precedence order. + + Mirrors `ProxyConfig._init_coordination_redis` -> `ProxyConfig._init_cache`: + an explicit block wins, else a plain-Redis response-cache backend is + borrowed, else the REDIS_* environment fallback applies. + """ + from litellm.proxy.proxy_server import _environment_has_redis_connection_target + + if settings: + return "coordination_redis" + cache_backend = litellm.cache.cache if litellm.cache is not None else None + if isinstance(cache_backend, (RedisCache, RedisClusterCache)): + return "cache_backend" + if _environment_has_redis_connection_target(): + return "environment" + return None + + +def _log_audit_task_exception(task: "asyncio.Task[None]") -> None: + """Surface a fire-and-forget audit-log task failure as a warning.""" + if task.cancelled(): + return + exc = task.exception() + if exc is not None: + verbose_proxy_logger.warning("Failed to write coordination-redis-settings audit log: %s", exc) + + +async def _emit_coordination_redis_audit_log( + *, + action: AUDIT_ACTIONS, + before_settings: Optional[Mapping[str, object]], + after_settings: Optional[Mapping[str, object]], + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str], +) -> None: + """Emit an audit-log row for a /coordination_redis/settings mutation.""" + if litellm.store_audit_logs is not True: + return + + from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + task = asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.CONFIG_TABLE_NAME, + object_id=_COORDINATION_REDIS_KEY, + action=action, + updated_values=json.dumps({"settings": _redact_all_values(after_settings)}, default=str), + before_value=json.dumps({"settings": _redact_all_values(before_settings)}, default=str), + ) + ) + ) + task.add_done_callback(_log_audit_task_exception) + + +class CoordinationRedisSettingsResponse(BaseModel): + values: dict[str, object] = Field(description="Current coordination Redis settings, with credentials redacted") + fields: list[CoordinationRedisSettingsField] = Field( + description="List of all configurable coordination Redis settings with metadata" + ) + source: Optional[CoordinationRedisSource] = Field( + description="Where the proxy's coordination Redis comes from; null when it has none" + ) + + +class CoordinationRedisSettingsRequest(BaseModel): + settings: dict[str, object] = Field(description="Coordination Redis connection params") + + +class CoordinationRedisTestResponse(BaseModel): + status: str = Field(description="Connection status: 'healthy' or 'unhealthy'") + error: Optional[str] = Field(default=None, description="Error message if the connection failed") + + +@router.get( + "/coordination_redis/settings", + tags=["Coordination Redis Settings"], + dependencies=[Depends(user_api_key_auth)], + response_model=CoordinationRedisSettingsResponse, +) +async def get_coordination_redis_settings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> CoordinationRedisSettingsResponse: + """ + Get the coordination Redis configuration and available settings. + + Returns: + - values: current coordination Redis settings, with password/sentinel_password/url redacted + - fields: all configurable settings with their metadata (type, description, default, section) + - source: "coordination_redis" | "cache_backend" | "environment" | null + """ + _enforce_proxy_admin(user_api_key_dict) + + settings = await _current_coordination_redis_settings() + source = _coordination_redis_source(settings) + + values = _redact_credentials(settings or {}) + fields = [field.model_copy(deep=True) for field in COORDINATION_REDIS_SETTINGS_FIELDS] + for field in fields: + if field.field_name in values: + field.field_value = values[field.field_name] + + return CoordinationRedisSettingsResponse(values=values, fields=fields, source=source) + + +@router.post( + "/coordination_redis/settings", + tags=["Coordination Redis Settings"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_coordination_redis_settings( + request: CoordinationRedisSettingsRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +) -> dict[str, object]: + """ + Save coordination Redis settings under `general_settings.coordination_redis`. + + Parameters: + - settings: dict - Redis connection params (host, port, username, password, url, ssl, startup_nodes, sentinel_nodes, sentinel_password, service_name). Values may be `os.environ/VAR` references, which are stored as written and resolved at startup + + The settings are written to the `general_settings` row of LiteLLM_Config, + which startup merges over the yaml config; the proxy picks them up on its + next restart. + """ + from litellm.proxy.proxy_server import prisma_client, store_model_in_db + + _enforce_proxy_admin(user_api_key_dict) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) + + if store_model_in_db is not True: + raise HTTPException( + status_code=500, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, + ) + + saved_settings = await _current_coordination_redis_settings() + settings = _merge_over_saved(request.settings, saved_settings or {}) + _validated_params(settings) + + general_settings = await _read_general_settings() + before_settings = general_settings.get(_COORDINATION_REDIS_KEY) + action: AUDIT_ACTIONS = "updated" if isinstance(before_settings, dict) else "created" + + await ConfigRepository(prisma_client).set_param( + param_name=_GENERAL_SETTINGS_PARAM_NAME, + param_value={**general_settings, _COORDINATION_REDIS_KEY: settings}, + ) + await invalidate_config_param(_GENERAL_SETTINGS_PARAM_NAME) + + # coordination_redis carries Redis credentials and decides where cross-pod + # rate-limit and spend state lives; an admin repointing it is a + # data-routing pivot, so make the change traceable. + await _emit_coordination_redis_audit_log( + action=action, + before_settings=before_settings if isinstance(before_settings, dict) else None, + after_settings=settings, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + return { + "message": "Coordination Redis settings updated successfully. Restart the proxy to apply them.", + "status": "success", + "settings": _redact_credentials(settings), + } + + +@router.post( + "/coordination_redis/settings/test", + tags=["Coordination Redis Settings"], + dependencies=[Depends(user_api_key_auth)], + response_model=CoordinationRedisTestResponse, +) +async def check_coordination_redis_connection( + request: CoordinationRedisSettingsRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> CoordinationRedisTestResponse: + """ + Test a coordination Redis connection with the provided credentials. + + Parameters: + - settings: dict - Redis connection params to test. Credential fields sent back as `***REDACTED***` fall back to the saved value + + Builds a throwaway client (never touching global state) and pings it. + """ + from litellm.proxy.proxy_server import _build_redis_usage_cache + + _enforce_proxy_admin(user_api_key_dict) + + saved_settings = await _current_coordination_redis_settings() + settings = _merge_over_saved(request.settings, saved_settings or {}) + params = _validated_params(settings) + + redis_cache: Optional[RedisCache] = None + try: + redis_cache = _build_redis_usage_cache(params.model_dump(exclude_none=True)) + await asyncio.wait_for(redis_cache.ping(), timeout=_PING_TIMEOUT_SECONDS) + return CoordinationRedisTestResponse(status="healthy") + except asyncio.TimeoutError: + return CoordinationRedisTestResponse( + status="unhealthy", + error=f"Connection timed out after {_PING_TIMEOUT_SECONDS}s", + ) + except Exception as e: # noqa: BLE001 # any client/connection failure is a health verdict, not a 500 + return CoordinationRedisTestResponse(status="unhealthy", error=_scrub_credentials(str(e), settings)) + finally: + if redis_cache is not None: + with suppress(Exception): + await redis_cache.disconnect() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4114bda47c9..f0ca1f6396f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -75,6 +75,7 @@ from litellm.proxy._types import ( ConfigGeneralSettings, ConfigList, ConfigYAML, + CoordinationRedisParams, EnterpriseLicenseData, FieldDetail, InvitationClaim, @@ -212,6 +213,7 @@ from contextlib import asynccontextmanager from functools import lru_cache import litellm +import litellm._redis from litellm import Router from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache @@ -361,6 +363,10 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import ( from litellm.proxy.management_endpoints.callback_management_endpoints import ( router as callback_management_endpoints_router, ) +from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( + get_persisted_coordination_redis_settings, + router as coordination_redis_settings_router, +) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, _user_has_admin_view, @@ -925,6 +931,16 @@ async def proxy_startup_event(app: FastAPI): asyncio.create_task(_run_pw_migration()) + ## A coordination_redis block saved from the admin UI lives in the database, + ## which is only reachable once the prisma client exists. Apply it here, before + ## the coordination Redis is published to its consumers below. + db_coordination_redis_cache = await ProxyStartupEvent._init_coordination_redis_from_db( + litellm_settings=proxy_config.get_config_state().get("litellm_settings") or {}, + llm_router=llm_router, + ) + if db_coordination_redis_cache is not None: + _set_redis_usage_cache(db_coordination_redis_cache) + ## use_redis_transaction_buffer: fall back to a standalone Redis (REDIS_* env) ## when the proxy cache backend is not Redis ## transaction_buffer_redis_cache = redis_usage_cache @@ -3549,6 +3565,101 @@ def _apply_ssrf_general_settings(settings: Mapping[str, object]) -> None: ) +def _set_redis_usage_cache(coordination_redis_cache: RedisCache | None) -> None: + """Publish the resolved coordination Redis to the consumers that read it directly.""" + global redis_usage_cache + redis_usage_cache = coordination_redis_cache + + +def _resolve_coordination_redis_env_refs(raw_params: Mapping[str, object]) -> dict[str, object]: + """Resolve `os.environ/VAR` references in a coordination_redis block.""" + return { + key: (get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value) + for key, value in raw_params.items() + } + + +def _build_redis_usage_cache(redis_params: Mapping[str, object]) -> RedisCache: + """ + Builds the proxy's coordination Redis client from resolved connection + params. Cluster-mode targets (explicit `startup_nodes` or the + REDIS_CLUSTER_NODES env var) get a `RedisClusterCache`, so consumers that + branch on cluster mode (e.g. the v3 rate limiter) take the cluster path; + everything else (host/url/sentinel) gets a plain `RedisCache`. + """ + startup_nodes = redis_params.get("startup_nodes") + if startup_nodes is None: + env_cluster_nodes = get_secret_str("REDIS_CLUSTER_NODES") + if env_cluster_nodes is not None: + startup_nodes = json.loads(env_cluster_nodes) + non_node_params = {key: value for key, value in redis_params.items() if key != "startup_nodes"} + if startup_nodes: + return RedisClusterCache(startup_nodes=startup_nodes, **non_node_params) + return RedisCache(**non_node_params) + + +def _environment_has_redis_connection_target() -> bool: + """ + Whether the REDIS_* environment variables name a Redis to connect to (host, + url, cluster nodes, or sentinel nodes). Read-only: callers that only need to + know whether the env fallback would apply use this instead of building a + client. + """ + redis_env_kwargs = litellm._redis._redis_kwargs_from_environment() + return ( + "host" in redis_env_kwargs + or "url" in redis_env_kwargs + or get_secret_str("REDIS_CLUSTER_NODES") is not None + or get_secret_str("REDIS_SENTINEL_NODES") is not None + ) + + +def _build_redis_usage_cache_from_environment() -> RedisCache | None: + """ + Builds a standalone coordination Redis from REDIS_* environment variables. + + Lets the proxy's coordination Redis (cross-pod tpm/rpm rate limits, spend + tracking, pod lock manager) run when the response-cache backend is not a + plain Redis KV cache (e.g. a semantic cache, disk, or s3). + + Returns None when the environment carries no connection target (host, url, + cluster nodes, or sentinel nodes). + """ + if not _environment_has_redis_connection_target(): + return None + return _build_redis_usage_cache(litellm._redis._redis_kwargs_from_environment()) + + +def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: bool) -> None: + """ + Wires an established coordination Redis into the proxy-level caches that + consume it directly: the spend counter cache, the cluster-wide config + cache, and (only when opted in) the virtual-key auth cache. + """ + spend_counter_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + if enable_redis_auth_cache is True: + user_api_key_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + verbose_proxy_logger.info( + "enable_redis_auth_cache=True: attached Redis to " + "user_api_key_cache — virtual-key lookups are now " + "shared across all proxy workers." + ) + else: + verbose_proxy_logger.info( + "enable_redis_auth_cache is not set: user_api_key_cache " + "remains in-memory only (per-worker). Set " + "litellm_settings.enable_redis_auth_cache: true to share " + "the auth cache across workers and reduce DB load." + ) + litellm_config_cache.redis_cache = redis_cache + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -3747,12 +3858,52 @@ class ProxyConfig: team_config = self._get_team_config(team_id=team_id, all_teams_config=all_teams_config) return team_config + def _init_coordination_redis(self, config: dict) -> RedisCache | None: + """ + Builds the coordination Redis from `general_settings.coordination_redis` + when present, attaching it to the proxy-level caches. Runs before cache + init, so an explicit block takes precedence over borrowing the + response-cache Redis and over the REDIS_* env fallback. Returns the + built client (None when the block is absent) for the caller to publish. + """ + settings = config.get("general_settings") or {} + litellm_settings = config.get("litellm_settings") or {} + raw_params = settings.get("coordination_redis") + if raw_params is None: + return None + if not isinstance(raw_params, dict): + raise ValueError("general_settings.coordination_redis must be a mapping of Redis connection params") + + coordination_params = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(raw_params)) + if not coordination_params.has_connection_target(): + raise ValueError( + "general_settings.coordination_redis needs a connection target: " + "set one of host, url, startup_nodes, or sentinel_nodes" + ) + + coordination_redis_cache = _build_redis_usage_cache(coordination_params.model_dump(exclude_none=True)) + _attach_redis_usage_cache( + coordination_redis_cache, + enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, + ) + verbose_proxy_logger.info( + "coordination_redis: using a standalone Redis from general_settings " + "for usage tracking, rate limiting, and cross-pod coordination." + ) + return coordination_redis_cache + def _init_cache( self, cache_params: dict, enable_redis_auth_cache: bool = False, - ): - global redis_usage_cache, llm_router, general_settings + ) -> RedisCache | None: + """ + Initializes the response cache and resolves the coordination Redis. + + Returns the coordination Redis for the caller to publish: an explicit + coordination_redis block already set wins, else a plain-Redis response + cache backend is borrowed, else the REDIS_* environment fallback applies. + """ from litellm import Cache if "default_in_memory_ttl" in cache_params: @@ -3763,37 +3914,29 @@ class ProxyConfig: litellm.cache = Cache(**cache_params) - if litellm.cache is not None and isinstance(litellm.cache.cache, (RedisCache, RedisClusterCache)): - ## INIT PROXY REDIS USAGE CLIENT ## - redis_usage_cache = litellm.cache.cache - spend_counter_cache.attach_redis_cache( - redis_usage_cache, - default_redis_ttl=litellm.default_redis_ttl, - ) - # Note: PKCE verifier storage uses redis_usage_cache directly (not - # user_api_key_cache) to avoid routing all API-key lookups through Redis. - if enable_redis_auth_cache is True: - user_api_key_cache.attach_redis_cache( - redis_usage_cache, - default_redis_ttl=litellm.default_redis_ttl, - ) - verbose_proxy_logger.info( - "enable_redis_auth_cache=True: attached Redis to " - "user_api_key_cache — virtual-key lookups are now " - "shared across all proxy workers." - ) + resolved_usage_cache = redis_usage_cache + cache_backend = litellm.cache.cache if litellm.cache is not None else None + if resolved_usage_cache is None: + if isinstance(cache_backend, (RedisCache, RedisClusterCache)): + ## INIT PROXY REDIS USAGE CLIENT ## + resolved_usage_cache = cache_backend else: - verbose_proxy_logger.info( - "enable_redis_auth_cache is not set: user_api_key_cache " - "remains in-memory only (per-worker). Set " - "litellm_settings.enable_redis_auth_cache: true to share " - "the auth cache across workers and reduce DB load." - ) - litellm_config_cache.redis_cache = redis_usage_cache + resolved_usage_cache = _build_redis_usage_cache_from_environment() + if resolved_usage_cache is not None: + verbose_proxy_logger.info( + "Cache backend %s is not a Redis KV cache; built a standalone " + "Redis from REDIS_* environment variables for usage tracking, " + "rate limiting, and cross-pod coordination.", + type(cache_backend).__name__, + ) + + if resolved_usage_cache is not None: # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. + _attach_redis_usage_cache(resolved_usage_cache, enable_redis_auth_cache) elif litellm_config_cache.redis_cache is None: verbose_proxy_logger.info("litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled.") + return resolved_usage_cache def switch_on_llm_response_caching(self): """ @@ -4038,6 +4181,11 @@ class ProxyConfig: self._load_environment_variables(config=config) + ## Coordination Redis (before cache init, so the explicit block wins) + coordination_redis_cache = self._init_coordination_redis(config=config) + if coordination_redis_cache is not None: + _set_redis_usage_cache(coordination_redis_cache) + ## Callback settings callback_settings = config.get("callback_settings", {}) if callback_settings: @@ -4118,9 +4266,11 @@ class ProxyConfig: cache_params[key] = get_secret(value) ## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = `, _redis.py checks for REDIS specific environment variables - self._init_cache( - cache_params=cache_params, - enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, + _set_redis_usage_cache( + self._init_cache( + cache_params=cache_params, + enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, + ) ) if litellm.cache is not None: verbose_proxy_logger.debug(f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}") @@ -7265,6 +7415,46 @@ class ProxyStartupEvent: "Redis for the transaction buffer." ) + @staticmethod + async def _init_coordination_redis_from_db( + litellm_settings: Mapping[str, object], + llm_router: Optional[Router], + ) -> RedisCache | None: + """ + Applies a coordination_redis block saved to the database, which the admin + UI writes and the config file therefore never carries. + + Returns None when nothing is persisted or the persisted block names no + connection target, leaving the file/env resolution untouched. + """ + try: + persisted = await get_persisted_coordination_redis_settings() + except Exception as e: # noqa: BLE001 # a config-row read failure must not block proxy startup + verbose_proxy_logger.warning("Could not read coordination_redis from the database: %s", e) + return None + if persisted is None: + return None + + coordination_params = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(persisted)) + if not coordination_params.has_connection_target(): + verbose_proxy_logger.warning( + "coordination_redis saved in the database names no connection target; ignoring it." + ) + return None + + coordination_redis_cache = _build_redis_usage_cache(coordination_params.model_dump(exclude_none=True)) + _attach_redis_usage_cache( + coordination_redis_cache, + enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, + ) + if llm_router is not None and llm_router.cache.redis_cache is None: + llm_router._update_redis_cache(cache=coordination_redis_cache) + verbose_proxy_logger.info( + "coordination_redis: using the standalone Redis saved in the database " + "for usage tracking, rate limiting, and cross-pod coordination." + ) + return coordination_redis_cache + @staticmethod def _get_transaction_buffer_redis_cache( general_settings: dict, @@ -7277,7 +7467,6 @@ class ProxyStartupEvent: Returns None when the buffer is disabled, or when no Redis host or url is set in the environment. """ - from litellm._redis import _redis_kwargs_from_environment from litellm.secret_managers.main import str_to_bool _use_redis_transaction_buffer: bool | str | None = general_settings.get("use_redis_transaction_buffer", False) @@ -7287,11 +7476,7 @@ class ProxyStartupEvent: if not _use_redis_transaction_buffer: return None - redis_env_kwargs = _redis_kwargs_from_environment() - if "host" not in redis_env_kwargs and "url" not in redis_env_kwargs: - return None - - return RedisCache(**redis_env_kwargs) + return _build_redis_usage_cache_from_environment() @classmethod async def _initialize_semantic_tool_filter( @@ -15811,6 +15996,7 @@ app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) +app.include_router(coordination_redis_settings_router) app.include_router(user_agent_analytics_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 1aaa38cea14..6b1ca3564e3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2020,8 +2020,6 @@ class LiteLLMCompletionResponsesConfig: output_details_dict: dict[str, int] = {} if hasattr(completion_details, "reasoning_tokens") and completion_details.reasoning_tokens is not None: output_details_dict["reasoning_tokens"] = completion_details.reasoning_tokens - else: - output_details_dict["reasoning_tokens"] = 0 if hasattr(completion_details, "text_tokens") and completion_details.text_tokens is not None: output_details_dict["text_tokens"] = completion_details.text_tokens diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 8e3be2bc12d..92dd5a513b9 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -127,7 +127,7 @@ def mock_responses_api_response( "input_tokens": 36, "input_tokens_details": {"cached_tokens": 0}, "output_tokens": 87, - "output_tokens_details": {"reasoning_tokens": 0}, + "output_tokens_details": {}, "total_tokens": 123, }, "user": None, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 3618331f0f5..eb78e6f9c8d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -17,6 +17,7 @@ from litellm.constants import ( LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, ) +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -26,7 +27,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ) from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig -from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook @@ -47,6 +48,44 @@ def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) - verbose_logger.error("%s failed: %s", task_name, exception) +_CLIENT_ERROR_CODES: frozenset[str] = frozenset( + ( + "invalid_request_error", + "context_length_exceeded", + "content_policy_violation", + "model_not_found", + ) +) + + +def _error_event_fields(error_obj: object) -> tuple[str, Optional[str], Optional[str]]: + if isinstance(error_obj, dict): + raw_message = error_obj.get("message") + raw_type = error_obj.get("type") + raw_code = error_obj.get("code") + elif error_obj is not None: + raw_message = getattr(error_obj, "message", None) + raw_type = getattr(error_obj, "type", None) + raw_code = getattr(error_obj, "code", None) + else: + raw_message = None + raw_type = None + raw_code = None + message = str(raw_message) if raw_message is not None else "Response API in-stream error" + error_type = raw_type if isinstance(raw_type, str) else None + code = raw_code if isinstance(raw_code, str) else None + return message, error_type, code + + +def _status_code_for_error_fields(error_type: Optional[str], error_code: Optional[str]) -> int: + fields = tuple(field for field in (error_type, error_code) if field is not None) + if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields): + return 429 + if any(field in _CLIENT_ERROR_CODES for field in fields): + return 400 + return 500 + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -73,6 +112,8 @@ class BaseResponsesAPIStreamingIterator: self.completed_response: Optional[Any] = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called + self._yielded_first_chunk = False + self._generated_content = "" self._completed_response_cached = False self._completed_response_logged = False self._completed_response_cache_hit: Optional[bool] = None @@ -160,6 +201,10 @@ class BaseResponsesAPIStreamingIterator: # Encode container_id on streaming events so proxy/UI follow-ups route correctly _event_type = getattr(openai_responses_api_chunk, "type", None) + if _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: + _delta = getattr(openai_responses_api_chunk, "delta", None) + if isinstance(_delta, str): + self._generated_content += _delta _stream_model_id = ( self.litellm_metadata.get("model_info", {}).get("id") if self.litellm_metadata else None ) @@ -327,17 +372,66 @@ class BaseResponsesAPIStreamingIterator: """ response_obj = getattr(self.completed_response, "response", None) if self.completed_response else None error_info = getattr(response_obj, "error", None) if response_obj else None - error_message = "Response failed" - if isinstance(error_info, dict): - error_message = error_info.get("message", str(error_info)) + error_message, error_type, error_code = _error_event_fields(error_info) + self._record_failed_response_usage(response_obj) exception = litellm.APIError( - status_code=500, + status_code=_status_code_for_error_fields(error_type, error_code), message=error_message, llm_provider=self.custom_llm_provider or "", model=self.model or "", ) self._handle_failure(exception) + def _record_failed_response_usage(self, response_obj: Optional[Any]) -> None: + if response_obj is None or self.logging_obj is None: + return + usage_obj = getattr(response_obj, "usage", None) + if usage_obj is None: + return + try: + self.logging_obj.model_call_details["combined_usage_object"] = ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage_obj) + ) + except (TypeError, ValueError) as usage_error: + verbose_logger.debug( + "could not record usage for failed responses stream: %s", + usage_error, + ) + return + self.logging_obj.model_call_details["response_cost"] = ( + self.logging_obj._response_cost_calculator(result=response_obj) or 0.0 + ) + + def _maybe_raise_for_error_event(self, result: object) -> None: + chunk_type = getattr(result, "type", None) + if chunk_type not in ("error", "response.failed"): + return + + error_obj: object = ( + getattr(getattr(result, "response", None), "error", None) + if chunk_type == "response.failed" + else getattr(result, "error", None) + ) + + error_message, error_type, error_code = _error_event_fields(error_obj) + status_code = _status_code_for_error_fields(error_type, error_code) + mapped_exception = litellm.APIError( + status_code=status_code, + message=error_message, + llm_provider=self.custom_llm_provider or "", + model=self.model or "", + ) + if 400 <= status_code < 500 and status_code != 429: + raise mapped_exception + raise MidStreamFallbackError( + message=str(mapped_exception), + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + original_exception=mapped_exception, + generated_content=self._generated_content, + is_pre_first_chunk=not self._yielded_first_chunk, + ) + def _get_completed_response_object(self) -> Optional[Any]: openai_types = _get_openai_response_types() completed_response = self.completed_response @@ -611,11 +705,13 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): if self.finished: raise StopAsyncIteration elif result is not None: + self._maybe_raise_for_error_event(result) # Await hook directly instead of run_async_function # (which spawns a thread + event loop per call) result = await self._call_post_streaming_deployment_hook( chunk=result, ) + self._yielded_first_chunk = True return result # If result is None, continue the loop to get the next chunk @@ -685,11 +781,13 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): if self.finished: raise StopIteration elif result is not None: + self._maybe_raise_for_error_event(result) # Sync path: use run_async_function for the hook result = run_async_function( async_function=self._call_post_streaming_deployment_hook, chunk=result, ) + self._yielded_first_chunk = True return result # If result is None, continue the loop to get the next chunk diff --git a/litellm/router.py b/litellm/router.py index 5ffe60c2da0..6e773a06c7f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2344,6 +2344,8 @@ class Router: self.completed_response = None self.start_time = getattr(source_iterator, "start_time", datetime.now()) self._failure_handled = False + self._yielded_first_chunk = False + self._generated_content = "" self._completed_response_cached = False self._completed_response_logged = False self._completed_response_cache_hit = None diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 4f9135ad9ed..16bf9a87a41 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -4,16 +4,21 @@ Complexity-based Auto Router A rule-based routing strategy that uses weighted scoring across multiple dimensions to classify requests by complexity and route them to appropriate models. -No external API calls - all scoring is local and <1ms. +By default, scoring is local (regex/keyword-based) with no external API calls and <1ms +latency. Optionally, classifier_type="llm" routes classification through a configured +model instead, trading that latency/cost guarantee for potentially better accuracy. Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter """ import re -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union + +from pydantic import BaseModel from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import ModelResponse from .config import ( DEFAULT_CODE_KEYWORDS, @@ -32,6 +37,24 @@ else: PreRoutingHookResponse = Any +class TierClassification(BaseModel): + """Structured response schema for the LLM-based complexity classifier.""" + + tier: Literal["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"] + + +_CLASSIFICATION_PROMPT_TEMPLATE = """Classify the complexity of the following user request into exactly one tier. + +Tiers: +- SIMPLE: factual lookups, greetings, short direct questions with no reasoning or code involved. +- MEDIUM: everyday requests needing some explanation or minor code/technical content. +- COMPLEX: requests involving non-trivial code, architecture, or multi-step technical work. +- REASONING: requests explicitly requiring step-by-step reasoning, analysis, or weighing tradeoffs. + +{system_context}Request: +{prompt}""" + + def _append_custom_keywords(base_keywords: list[str], custom_keywords: Optional[list[str]]) -> list[str]: if not custom_keywords: return base_keywords @@ -40,6 +63,20 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: Optional[ return [*base_keywords, *deduped_custom.values()] +# Metadata keys that carry the parent request's budget reservation. These must not +# reach the classifier's internal acompletion call: the reservation belongs to the +# routed completion that the classifier is deciding on, not to the classifier call +# itself, and forwarding it would let the classifier's cost-tracking reconcile +# against a reservation it isn't responsible for. +_BUDGET_RESERVATION_METADATA_KEYS = frozenset({"user_api_key_budget_reservation", "user_api_key_auth"}) + + +def _classifier_call_metadata(metadata: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]: + if not metadata: + return metadata + return {k: v for k, v in metadata.items() if k not in _BUDGET_RESERVATION_METADATA_KEYS} + + class DimensionScore: """Represents a score for a single dimension with optional signal.""" @@ -53,10 +90,10 @@ class DimensionScore: class ComplexityRouter(CustomLogger): """ - Rule-based complexity router that classifies requests and routes to appropriate models. + Complexity router that classifies requests and routes to appropriate models. - Handles requests in <1ms with zero external API calls by using weighted scoring - across multiple dimensions: + By default, handles requests in <1ms with zero external API calls, using weighted + scoring across multiple dimensions: - Token count (short=simple, long=complex) - Code presence (code keywords → complex) - Reasoning markers ("step by step", "think through" → reasoning tier) @@ -297,6 +334,63 @@ class ComplexityRouter(CustomLogger): return tier, weighted_score, signals + async def aclassify( + self, + prompt: str, + system_prompt: Optional[str] = None, + request_kwargs: Optional[dict[str, Any]] = None, + ) -> tuple[ComplexityTier, float, list[str]]: + """ + Classify a prompt by complexity, using the LLM classifier when configured. + + Falls back to the local heuristic scorer if classifier_type is "heuristic", + or if the LLM call fails, times out, or returns an unparseable response. + """ + if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None: + return self.classify(prompt, system_prompt) + + try: + tier = await self._classify_with_llm(prompt, system_prompt, request_kwargs) + return tier, 1.0, [f"llm-classifier:{tier.value}"] + except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the heuristic scorer + verbose_router_logger.warning( + f"ComplexityRouter: LLM classifier failed ({e}), falling back to heuristic scoring" + ) + return self.classify(prompt, system_prompt) + + async def _classify_with_llm( + self, + prompt: str, + system_prompt: Optional[str] = None, + request_kwargs: Optional[dict[str, Any]] = None, + ) -> ComplexityTier: + """Call the configured classifier model and parse its structured tier response.""" + llm_config = self.config.classifier_llm_config + if llm_config is None: + raise ValueError("classifier_llm_config is not set") + + system_context = f"Context: {system_prompt}\n\n" if system_prompt else "" + classification_prompt = _CLASSIFICATION_PROMPT_TEMPLATE.format(system_context=system_context, prompt=prompt) + + # Forward the original request's metadata so the classifier call's spend is + # attributed to the calling key/team instead of being dropped. Excludes the + # parent request's budget reservation, which the routed completion (not this + # internal classifier call) is responsible for reconciling. + metadata = _classifier_call_metadata((request_kwargs or {}).get("litellm_metadata")) + + response: ModelResponse = await self.litellm_router_instance.acompletion( + model=llm_config.model, + messages=[{"role": "user", "content": classification_prompt}], + response_format=TierClassification, + timeout=llm_config.timeout_ms / 1000, + metadata=metadata, + ) + content = response.choices[0].message.content + if not content: + raise ValueError("LLM classifier returned empty content") + result = TierClassification.model_validate_json(content) + return ComplexityTier[result.tier] + def get_model_for_tier(self, tier: ComplexityTier) -> str: """ Get the model name for a given complexity tier. @@ -445,7 +539,7 @@ class ComplexityRouter(CustomLogger): messages=messages if has_original_messages else None, ) - tier, score, signals = self.classify(user_message, system_prompt) + tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs) routed_model = self.get_model_for_tier(tier) verbose_router_logger.info( diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index ae04d5fc5d4..6407c9a4590 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -6,9 +6,9 @@ All values are configurable via proxy config.yaml. """ from enum import Enum -from typing import Dict, List, Optional +from typing import Dict, List, Literal, Optional -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator class ComplexityTier(str, Enum): @@ -197,6 +197,18 @@ DEFAULT_TIER_MODELS: Dict[str, str] = { } +class ClassifierLLMConfig(BaseModel): + """Configuration for the LLM-based complexity classifier.""" + + model: str = Field( + description="Model name (from the router's model_list) to call for classification", + ) + timeout_ms: int = Field( + default=3000, + description="Timeout budget for the classification call, in milliseconds", + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -257,8 +269,24 @@ class ComplexityRouterConfig(BaseModel): description="Default model to use if tier cannot be determined", ) + # Classifier strategy + classifier_type: Literal["heuristic", "llm"] = Field( + default="heuristic", + description="Classification strategy: local regex/keyword scoring, or an LLM call", + ) + classifier_llm_config: Optional[ClassifierLLMConfig] = Field( + default=None, + description="Configuration for the LLM classifier; required when classifier_type is 'llm'", + ) + model_config = ConfigDict(extra="allow") # Allow additional fields + @model_validator(mode="after") + def _validate_llm_classifier_config(self) -> "ComplexityRouterConfig": + if self.classifier_type == "llm" and self.classifier_llm_config is None: + raise ValueError("classifier_llm_config is required when classifier_type is 'llm'") + return self + # Combined default config DEFAULT_COMPLEXITY_CONFIG = ComplexityRouterConfig() diff --git a/litellm/types/integrations/datadog.py b/litellm/types/integrations/datadog.py index 89faac27830..e0f43519b3d 100644 --- a/litellm/types/integrations/datadog.py +++ b/litellm/types/integrations/datadog.py @@ -6,6 +6,7 @@ from typing_extensions import NotRequired, TypedDict from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams DD_MAX_BATCH_SIZE = 1000 +DD_MAX_PAYLOAD_SIZE_BYTES = 4_000_000 class DataDogStatus(str, Enum): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 3ab5a7b736e..daac1e4506f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1185,7 +1185,7 @@ class ResponsesAPIRequestParams(ResponsesAPIOptionalRequestParams, total=False): class OutputTokensDetails(BaseLiteLLMOpenAIResponseObject): - reasoning_tokens: int = 0 + reasoning_tokens: Optional[int] = None text_tokens: Optional[int] = None @@ -1720,7 +1720,7 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject): type: str # e.g., 'invalid_request_error' code: str # e.g., 'context_length_exceeded' message: str - param: Optional[str] = None + param: Optional[Union[str, Dict[str, Any]]] = None class ErrorEvent(BaseLiteLLMOpenAIResponseObject): diff --git a/litellm/types/management_endpoints/__init__.py b/litellm/types/management_endpoints/__init__.py index 5c5bcb2e754..3b501443edd 100644 --- a/litellm/types/management_endpoints/__init__.py +++ b/litellm/types/management_endpoints/__init__.py @@ -7,6 +7,12 @@ from .cache_settings_endpoints import ( REDIS_TYPE_DESCRIPTIONS, CacheSettingsField, ) +from .coordination_redis_endpoints import ( + COORDINATION_REDIS_SETTINGS_FIELDS, + CoordinationRedisSection, + CoordinationRedisSettingsField, + CoordinationRedisSource, +) from .router_settings_endpoints import ( ROUTER_SETTINGS_FIELDS, ROUTING_STRATEGY_DESCRIPTIONS, @@ -20,4 +26,8 @@ __all__ = [ "CACHE_SETTINGS_FIELDS", "REDIS_TYPE_DESCRIPTIONS", "CacheSettingsField", + "COORDINATION_REDIS_SETTINGS_FIELDS", + "CoordinationRedisSection", + "CoordinationRedisSettingsField", + "CoordinationRedisSource", ] diff --git a/litellm/types/management_endpoints/coordination_redis_endpoints.py b/litellm/types/management_endpoints/coordination_redis_endpoints.py new file mode 100644 index 00000000000..b6889d83323 --- /dev/null +++ b/litellm/types/management_endpoints/coordination_redis_endpoints.py @@ -0,0 +1,105 @@ +""" +Types and field definitions for coordination Redis settings management endpoints +""" + +from typing import Literal, Optional + +from pydantic import BaseModel + +CoordinationRedisSection = Literal["connection", "cluster", "sentinel"] + +CoordinationRedisSource = Literal["coordination_redis", "cache_backend", "environment"] + + +class CoordinationRedisSettingsField(BaseModel): + field_name: str + field_type: str + field_value: Optional[object] = None + field_description: str + field_default: Optional[object] = None + ui_field_name: str + section: CoordinationRedisSection + + +COORDINATION_REDIS_SETTINGS_FIELDS: list[CoordinationRedisSettingsField] = [ + CoordinationRedisSettingsField( + field_name="host", + field_type="String", + field_description="Redis server hostname or IP address", + ui_field_name="Host", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="port", + field_type="Integer", + field_description="Redis server port number", + field_default=6379, + ui_field_name="Port", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="username", + field_type="String", + field_description="Redis server username (if required)", + ui_field_name="Username", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="password", + field_type="String", + field_description="Redis server password", + ui_field_name="Password", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="url", + field_type="String", + field_description=( + "Full Redis connection URL (e.g. redis://:password@host:6379/1). " + "Set this instead of the discrete host/port/username/password fields." + ), + ui_field_name="Redis URL", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="ssl", + field_type="Boolean", + field_description="Connect to Redis over TLS", + field_default=False, + ui_field_name="SSL", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="startup_nodes", + field_type="List", + field_description=( + "Cluster-mode startup nodes (e.g. [{'host': '127.0.0.1', 'port': 7001}]). " + "When set, a Redis Cluster client is used." + ), + ui_field_name="Cluster Startup Nodes", + section="cluster", + ), + CoordinationRedisSettingsField( + field_name="sentinel_nodes", + field_type="List", + field_description=( + "Sentinel [host, port] pairs (e.g. [['localhost', 26379]]). When set, a Sentinel-managed client is used." + ), + ui_field_name="Sentinel Nodes", + section="sentinel", + ), + CoordinationRedisSettingsField( + field_name="sentinel_password", + field_type="String", + field_description="Password for the Redis Sentinel nodes", + ui_field_name="Sentinel Password", + section="sentinel", + ), + CoordinationRedisSettingsField( + field_name="service_name", + field_type="String", + field_description="Master service name for Redis Sentinel", + ui_field_name="Service Name", + section="sentinel", + ), +] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8a5acc24d19..e33e2335525 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -143,6 +143,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_web_search: Optional[bool] supports_reasoning: Optional[bool] supports_adaptive_thinking: Optional[bool] + supports_mid_conversation_system: Optional[bool] supports_url_context: Optional[bool] supports_none_reasoning_effort: Optional[bool] supports_minimal_reasoning_effort: Optional[bool] diff --git a/litellm/utils.py b/litellm/utils.py index 5af7b62b332..731be992af0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -61,7 +61,7 @@ from litellm._lazy_imports import ( ) from litellm._uuid import uuid from litellm.litellm_core_utils.fallback_generalizations import ( - match_fallback_generalization, + match_all_fallback_generalizations, ) from litellm.constants import ( DEFAULT_CHAT_COMPLETION_PARAM_VALUES, @@ -5046,9 +5046,10 @@ def _get_model_info_from_generalization( """Resolve an unmapped model via a declarative fallback-generalization rule. Tries the same name candidates as the exact lookups, in the same order, and - returns ``(matched_name, model_info)`` for the first candidate whose rule also - satisfies the provider constraint. O(number of rules); only call after the - exact lookups have missed. + returns ``(matched_name, model_info)`` for the first matching rule that also + satisfies the provider constraint; a rule scoped to another provider is + skipped in favor of later rules rather than discarding the candidate. + O(number of rules); only call after the exact lookups have missed. """ candidates = [ potential_model_names["combined_model_name"], @@ -5058,11 +5059,9 @@ def _get_model_info_from_generalization( potential_model_names["stripped_model_name"], ] for candidate in candidates: - generalized_info = match_fallback_generalization(candidate) - if generalized_info is not None and _check_provider_match( - model_info=generalized_info, custom_llm_provider=custom_llm_provider - ): - return candidate, generalized_info + for generalized_info in match_all_fallback_generalizations(candidate): + if _check_provider_match(model_info=generalized_info, custom_llm_provider=custom_llm_provider): + return candidate, generalized_info return None @@ -5472,6 +5471,7 @@ def _get_model_info_helper( supports_url_context=_model_info.get("supports_url_context", None), supports_reasoning=_model_info.get("supports_reasoning", None), supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), + supports_mid_conversation_system=_model_info.get("supports_mid_conversation_system", None), supports_none_reasoning_effort=_model_info.get("supports_none_reasoning_effort", None), supports_minimal_reasoning_effort=_model_info.get("supports_minimal_reasoning_effort", None), supports_low_reasoning_effort=_model_info.get("supports_low_reasoning_effort", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e79ddbe35d2..31a6146cbd6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1359,6 +1359,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1393,6 +1394,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1427,6 +1429,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1461,6 +1464,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1481,6 +1485,7 @@ "anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1516,6 +1521,7 @@ "global.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1551,6 +1557,7 @@ "us.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1586,6 +1593,7 @@ "eu.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1621,6 +1629,43 @@ "au.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "input_cost_per_token": 5.5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true + }, + "jp.anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1703,6 +1748,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1737,6 +1783,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1771,6 +1818,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1805,6 +1853,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1839,6 +1888,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -1873,6 +1923,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -45231,6 +45282,17 @@ }, "fallback_generalizations": { "rules": [ + { + "name": "bedrock-anthropic-claude-mid-conversation-system", + "pattern": "anthropic\\.claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d)(?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", + "description": "Bedrock Invoke ids for Claude 4.8 or higher: anthropic.claude- with minor 4.8 through 4.99, any 5.x or later major-minor, or a bare 5+ major, which also admits new families such as fable. These models accept mid-conversation role system messages in place (verified live on Opus 4.8, Sonnet 5 and Fable 5), so unmapped future Bedrock Claudes keep the cache-preserving in-place handling instead of the hoist-all default. Listed first so bare-id provider inference, which takes the first pattern hit, resolves these Bedrock ids to bedrock; model-info resolution skips provider-mismatched rules either way.", + "extends": "anthropic-claude", + "model_info": { + "litellm_provider": "bedrock", + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true + } + }, { "name": "anthropic-claude-adaptive-thinking", "pattern": "(?:opus|sonnet|haiku)[-._](?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|(?:[5-9]|[1-9]\\d{1,})[-._]\\d{1,2}(?!\\d))", diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 6feffe036bd..1a51dd0d0a7 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -189,7 +189,7 @@ litellm_settings: langfuse_host: https://us.cloud.langfuse.com # cache: true # [OPTIONAL] use for caching responses # enable_caching_on_provider_specific_optional_params: True # Include provider-specific params in cache keys - # cache_params: # And for shared health check + # cache_params: # type: redis # host: localhost # port: 6379 @@ -228,8 +228,11 @@ general_settings: proxy_batch_write_at: 1 database_connection_pool_limit: 10 # background_health_checks: true - # use_shared_health_check: true + # use_shared_health_check: true # needs a coordination Redis (below) # health_check_interval: 30 + # coordination_redis: # standalone Redis for cross-pod coordination: rate limits, spend tracking, pod locks, shared health checks + # host: localhost + # port: 6379 # cancel_on_disconnect: true # cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot) # database_url: "postgresql://:@:/" # [OPTIONAL] use for token-based auth to proxy diff --git a/terraform/litellm/README.md b/terraform/litellm/README.md index 8f09cb53407..d4b40741052 100644 --- a/terraform/litellm/README.md +++ b/terraform/litellm/README.md @@ -178,6 +178,7 @@ only where the underlying cloud forces it. | Force destroy of object store | `s3_force_destroy` | `gcs_force_destroy` | | Database deletion protection | `skip_final_snapshot` | `cloudsql_deletion_protection` | | `proxy_config` (typed YAML map) | `proxy_config` | `proxy_config` | +| Coordination Redis | `REDIS_*` from ElastiCache (automatic) | `REDIS_*` from Memorystore (automatic) | | Extra plain env per component | `gateway_extra_env`, `backend_extra_env` | `gateway_extra_env`, `backend_extra_env` | | Extra secret-backed env | `gateway_extra_secrets`, `backend_extra_secrets` (ARNs) | `gateway_extra_secrets`, `backend_extra_secrets` (resource IDs) | | Uvicorn `--workers` on gateway | `gateway_num_workers` | `gateway_num_workers` | @@ -189,6 +190,19 @@ Each module stamps its own stack-identity tag (`litellm:stack` on AWS, merges `var.tags` / `var.labels` on top. Provider `default_tags` on AWS merge on top of all of these. +Coordination Redis needs no input on either cloud. Each module provisions the +managed Redis (ElastiCache on AWS, Memorystore on GCP) and exports `REDIS_HOST`, +`REDIS_PORT` and `REDIS_SSL` (plus `REDIS_SSL_CA_CERTS` on GCP) into the gateway +and backend env. The proxy falls back to those variables to build its +coordination Redis, which backs cross-pod tpm/rpm rate limits, spend tracking +and the pod lock manager. This is independent of LLM response caching, which +stays off unless you enable `litellm_settings.cache` in `proxy_config`. + +To coordinate through a Redis the module does not manage, set +`general_settings.coordination_redis` in `var.proxy_config`. An explicit block +overrides the `REDIS_*` env fallback; see the commented example in each +stack's `examples/default/terraform.tfvars.example` + OTel is opt-in on both clouds: leave `otel_endpoint` empty and nothing OTel-related is added to the container env; set it and both gateway and backend get `LITELLM_OTEL_V2=true` plus the full `OTEL_*` block, with diff --git a/terraform/litellm/aws/examples/default/terraform.tfvars.example b/terraform/litellm/aws/examples/default/terraform.tfvars.example index 4fdfb47e678..061ca2a9b82 100644 --- a/terraform/litellm/aws/examples/default/terraform.tfvars.example +++ b/terraform/litellm/aws/examples/default/terraform.tfvars.example @@ -56,6 +56,19 @@ env = "stage" # general_settings = { # master_key = "os.environ/LITELLM_MASTER_KEY" # database_url = "os.environ/DATABASE_URL" +# +# # Optional. The module already exports REDIS_HOST / REDIS_PORT / REDIS_SSL +# # from the ElastiCache group it provisions, and the proxy falls back to +# # those for cross-pod rate limits, spend tracking and the pod lock manager. +# # Set this block only to coordinate through a Redis the module does not +# # manage; it overrides the REDIS_* env fallback. Cluster mode takes +# # `startup_nodes` and sentinel takes `sentinel_nodes` + `service_name` +# # coordination_redis = { +# # host = "os.environ/COORDINATION_REDIS_HOST" +# # port = "os.environ/COORDINATION_REDIS_PORT" +# # password = "os.environ/COORDINATION_REDIS_PASSWORD" +# # ssl = true +# # } # } # } diff --git a/terraform/litellm/gcp/examples/default/terraform.tfvars.example b/terraform/litellm/gcp/examples/default/terraform.tfvars.example index 6358ec96e6d..4416cf0ee5d 100644 --- a/terraform/litellm/gcp/examples/default/terraform.tfvars.example +++ b/terraform/litellm/gcp/examples/default/terraform.tfvars.example @@ -51,6 +51,20 @@ env = "stage" # general_settings = { # master_key = "os.environ/LITELLM_MASTER_KEY" # database_url = "os.environ/DATABASE_URL" +# +# # Optional. The module already exports REDIS_HOST / REDIS_PORT / REDIS_SSL +# # (plus REDIS_SSL_CA_CERTS) from the Memorystore instance it provisions, and +# # the proxy falls back to those for cross-pod rate limits, spend tracking +# # and the pod lock manager. Set this block only to coordinate through a +# # Redis the module does not manage; it overrides the REDIS_* env fallback. +# # Cluster mode takes `startup_nodes` and sentinel takes `sentinel_nodes` +# # plus `service_name` +# # coordination_redis = { +# # host = "os.environ/COORDINATION_REDIS_HOST" +# # port = "os.environ/COORDINATION_REDIS_PORT" +# # password = "os.environ/COORDINATION_REDIS_PASSWORD" +# # ssl = true +# # } # } # } diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 5a09403c570..4d6e0528ed4 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -172,3 +172,4 @@ langgraph: >=1.0.10 # MIT License langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE pytest-rerunfailures: >=15.1 # MPL 2.0 license pytest-recording: >=0.13.4 # MIT license +expression: >=5.6.0 # MIT License - https://github.com/cognitedata/Expression/blob/main/LICENSE diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index ea8b8fa886c..bd1517dbffb 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1628,23 +1628,27 @@ async def test_openai_responses_api_token_limit_error(): """ Relevant issue: https://github.com/BerriAI/litellm/issues/15785 - - When this fails you'll see: - "pydantic_core._pydantic_core.ValidationError: 3 validation errors for ErrorEvent" - in the console. + Parsing the in-stream ErrorEvent must not raise + "pydantic_core._pydantic_core.ValidationError: 3 validation errors for ErrorEvent". + The iterator now surfaces the event as litellm.APIError with status 400 + (invalid_request_error is a non-retriable client error, so no + MidStreamFallbackError wrapping) carrying the provider's message. """ litellm._turn_on_debug() # Generate text with >400k tokens to trigger token limit error oversized_text = "This is a test sentence. " * 50000 # ~400k tokens - # This will raise ValidationError instead of showing the real error response = await litellm.aresponses( model="gpt-5-mini", input=oversized_text, stream=True ) - async for event in response: - print(event) # Never reaches here - ValidationError is raised + with pytest.raises(litellm.APIError) as exc_info: + async for event in response: + print(event) + + assert exc_info.value.status_code == 400 + assert "exceeds the context window" in str(exc_info.value) async def test_openai_streaming_logging(): diff --git a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py index 25bf79cd575..2fb7bdfceb5 100644 --- a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py +++ b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py @@ -266,3 +266,220 @@ async def test_aresponses_with_streaming_fallbacks_wraps_streaming_iterator(): ) assert out is wrapped mock_wrap.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_aresponses_fallback_on_in_stream_error_event(): + """A retriable in-stream error event (429) must trigger the router's mid-stream + fallback path: the wrapper catches MidStreamFallbackError raised by the source + iterator and yields the fallback stream instead of surfacing the error.""" + import json + from unittest.mock import Mock + + import litellm + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + from litellm.types.llms.openai import ErrorEvent, ErrorEventError + + router = _make_router() + + error_payload = { + "type": "error", + "error": {"type": "tokens", "code": "rate_limit_exceeded", "message": "rate limited"}, + } + sse_bytes = f"data: {json.dumps(error_payload)}\n\n".encode() + + async def mock_aiter_bytes(): + yield sse_bytes + + mock_response = Mock() + mock_response.headers = {} + mock_response.aiter_bytes = mock_aiter_bytes + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.completion_start_time = None + mock_config = Mock(spec=BaseResponsesAPIConfig) + mock_config.transform_streaming_response.return_value = ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=0, + error=ErrorEventError(type="tokens", code="rate_limit_exceeded", message="rate limited"), + ) + + source = ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-5", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + fallback_event = _make_completed_event(1, 1, 2) + + class _FallbackStream: + def __init__(self) -> None: + self._done = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + self._done = True + return fallback_event + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FallbackStream()), + ) as mock_fallback: + wrapped = await router._aresponses_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "input": "original question"}, + ) + collected = [ev async for ev in wrapped] + + assert collected == [fallback_event] + mock_fallback.assert_awaited_once() + raised = mock_fallback.await_args.kwargs["e"] + assert isinstance(raised, MidStreamFallbackError) + assert raised.status_code == 429 + assert isinstance(raised.original_exception, litellm.APIError) + assert raised.original_exception.status_code == 429 + assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question" + + +@pytest.mark.asyncio +async def test_aresponses_fallback_uses_continuation_input_after_partial_content(): + """When output text was already streamed before the error, the fallback re-entry + must carry a continuation input with the partial assistant text instead of + retrying the original input from scratch (which would duplicate streamed content).""" + import json + from unittest.mock import Mock + + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + from litellm.types.llms.openai import ErrorEvent, ErrorEventError + + router = _make_router() + + events = [ + {"type": "response.output_text.delta", "delta": "partial answer"}, + {"type": "error", "error": {"type": "server_error", "code": "internal_error", "message": "boom"}}, + ] + sse_payload = b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + async def mock_aiter_bytes(): + yield sse_payload + + mock_response = Mock() + mock_response.headers = {} + mock_response.aiter_bytes = mock_aiter_bytes + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.completion_start_time = None + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def transform(model, parsed_chunk, logging_obj): + if parsed_chunk.get("type") == "error": + return ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=0, + error=ErrorEventError(**parsed_chunk["error"]), + ) + delta_event = Mock() + delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + delta_event.delta = parsed_chunk["delta"] + return delta_event + + mock_config.transform_streaming_response.side_effect = transform + + source = ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-5", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + fallback_event = _make_completed_event(1, 1, 2) + + class _FallbackStream: + def __init__(self) -> None: + self._done = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + self._done = True + return fallback_event + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FallbackStream()), + ) as mock_fallback: + wrapped = await router._aresponses_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "input": "original question"}, + ) + collected = [ev async for ev in wrapped] + + assert collected[-1] == fallback_event + raised = mock_fallback.await_args.kwargs["e"] + assert isinstance(raised, MidStreamFallbackError) + assert raised.is_pre_first_chunk is False + assert raised.generated_content == "partial answer" + continuation = mock_fallback.await_args.kwargs["kwargs"]["input"] + assert isinstance(continuation, list) + assert continuation[0]["content"][0]["text"] == "original question" + assert continuation[-2]["role"] == "developer" + assert continuation[-1]["role"] == "assistant" + assert continuation[-1]["content"][0]["text"] == "partial answer" + + +@pytest.mark.asyncio +async def test_aresponses_client_error_event_skips_fallback(): + """A 400-mapped in-stream error (raised as APIError, not MidStreamFallbackError) + must surface to the caller without invoking the router's fallback path.""" + import litellm + + router = _make_router() + + class _ClientErrorSource: + completed_response = None + + def __aiter__(self): + return self + + async def __anext__(self): + raise litellm.APIError( + status_code=400, + message="bad request", + llm_provider="openai", + model="gpt-5", + ) + + wrapped = await router._aresponses_streaming_iterator( + response=_ClientErrorSource(), + initial_kwargs={"model": "primary"}, + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + with pytest.raises(litellm.APIError) as exc_info: + async for _ in wrapped: + pass + + assert exc_info.value.status_code == 400 + mock_fallback.assert_not_awaited() diff --git a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py b/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py index d1c7a4032fb..f645379a4f4 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock, Mock, patch import httpx @@ -6,16 +7,20 @@ from httpx import Request, Response from litellm.integrations.datadog.datadog import DataDogLogger from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError -from litellm.types.integrations.datadog import DD_MAX_BATCH_SIZE, DatadogPayload +from litellm.types.integrations.datadog import ( + DD_MAX_BATCH_SIZE, + DD_MAX_PAYLOAD_SIZE_BYTES, + DatadogPayload, +) -def _payloads(n): +def _payloads(n, message=None): return [ DatadogPayload( ddsource="litellm", ddtags="env:test", hostname="host", - message=f'{{"event": {i}}}', + message=f"{message}{i}" if message else f'{{"event": {i}}}', service="svc", status="info", ) @@ -177,6 +182,87 @@ async def test_413_returned_response_also_splits(datadog_env): assert logger.log_queue == [] +def _make_recording_send(sent_batches, delivered): + async def _send(data): + sent_batches.append(list(data)) + delivered.extend(data) + return Response( + 202, request=Request("POST", "https://example.com"), text="Accepted" + ) + + return _send + + +@pytest.mark.asyncio +async def test_oversized_payload_splits_before_any_send(datadog_env): + """Regression for LIT-4325: a batch above Datadog's uncompressed payload limit is + split proactively, so the intake never has to reject it with a 413.""" + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + with patch("asyncio.create_task"): + logger = DataDogLogger() + + events = _payloads(3, message="x" * 3_000_000) + logger.log_queue = list(events) + sent_batches: list = [] + delivered: list = [] + logger.async_send_compressed_data = AsyncMock( + side_effect=_make_recording_send(sent_batches, delivered) + ) + + await logger.async_send_batch() + + assert delivered == events + assert len(sent_batches) == 3 + assert all( + len(safe_dumps(batch).encode("utf-8")) <= DD_MAX_PAYLOAD_SIZE_BYTES + for batch in sent_batches + ) + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_batch_over_max_event_count_splits_before_any_send(datadog_env): + """Datadog caps a payload at 1000 events; a queue that grew past that (e.g. after + re-queues) must be sent in count-compliant chunks.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + events = _payloads(DD_MAX_BATCH_SIZE + 1) + logger.log_queue = list(events) + sent_batches: list = [] + delivered: list = [] + logger.async_send_compressed_data = AsyncMock( + side_effect=_make_recording_send(sent_batches, delivered) + ) + + await logger.async_send_batch() + + assert delivered == events + assert all(len(batch) <= DD_MAX_BATCH_SIZE for batch in sent_batches) + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_single_event_above_payload_cap_is_still_sent(datadog_env): + """A lone event over the byte cap cannot be split further; it must be sent once + (Datadog decides), never looped on.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(1, message="x" * (DD_MAX_PAYLOAD_SIZE_BYTES + 1)) + sent_batches: list = [] + delivered: list = [] + send = AsyncMock(side_effect=_make_recording_send(sent_batches, delivered)) + logger.async_send_compressed_data = send + + await asyncio.wait_for(logger.async_send_batch(), timeout=10) + + assert send.await_count == 1 + assert len(delivered) == 1 + assert logger.log_queue == [] + + @pytest.mark.asyncio async def test_partial_delivery_then_transient_error_requeues_only_undelivered( datadog_env, diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index eb795a64b79..5191414edeb 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -485,6 +485,74 @@ def test_build_span_exporter_variants(): OpenTelemetryV2Config(exporter="otlp_http", endpoint="http://h:4318") ) assert "OTLPSpanExporter" in type(http_exporter).__name__ + + +def test_otlp_logs_endpoint_normalization(): + norm = providers._otlp_logs_endpoint + # A base endpoint gets the signal path appended (the common OTLP env shape). + assert norm("http://collector:4318") == "http://collector:4318/v1/logs" + assert norm("http://collector:4318/") == "http://collector:4318/v1/logs" + # An already-correct path is left intact. + assert norm("http://collector:4318/v1/logs") == "http://collector:4318/v1/logs" + # A sibling signal's path is rewritten to logs, so one OTEL_ENDPOINT works + # for every signal rather than POSTing events at the traces path. + assert norm("http://collector:4318/v1/traces") == "http://collector:4318/v1/logs" + assert norm("http://collector:4318/v1/metrics") == "http://collector:4318/v1/logs" + assert norm(None) is None + + +def test_build_log_exporter_variants(): + from opentelemetry.sdk._logs.export import ConsoleLogExporter, InMemoryLogExporter + + assert isinstance( + providers.build_log_exporter(OpenTelemetryV2Config(exporter="console")), + ConsoleLogExporter, + ) + assert isinstance( + providers.build_log_exporter(OpenTelemetryV2Config(exporter="in_memory")), + InMemoryLogExporter, + ) + # An unrecognized kind falls back to console rather than dropping events. + assert isinstance( + providers.build_log_exporter(OpenTelemetryV2Config(exporter="unknown")), + ConsoleLogExporter, + ) + http_exporter = providers.build_log_exporter( + OpenTelemetryV2Config(exporter="otlp_http", endpoint="http://h:4318") + ) + assert "OTLPLogExporter" in type(http_exporter).__name__ + + +def test_build_logger_provider_picks_processor_by_exporter_kind(): + """Console and in-memory exporters export synchronously (tests depend on it); + every other destination gets the batch processor.""" + from opentelemetry.sdk._logs.export import ( + BatchLogRecordProcessor, + ConsoleLogExporter, + InMemoryLogExporter, + SimpleLogRecordProcessor, + ) + + cfg = OpenTelemetryV2Config(exporter="in_memory") + + def processor_of(provider): + return provider._multi_log_record_processor._log_record_processors[0] + + assert isinstance( + processor_of(providers.build_logger_provider(cfg, log_exporter=InMemoryLogExporter())), + SimpleLogRecordProcessor, + ) + assert isinstance( + processor_of(providers.build_logger_provider(cfg, log_exporter=ConsoleLogExporter())), + SimpleLogRecordProcessor, + ) + http_exporter = providers.build_log_exporter( + OpenTelemetryV2Config(exporter="otlp_http", endpoint="http://h:4318") + ) + assert isinstance( + processor_of(providers.build_logger_provider(cfg, log_exporter=http_exporter)), + BatchLogRecordProcessor, + ) grpc_exporter = providers.build_span_exporter( OpenTelemetryV2Config(exporter="otlp_grpc", endpoint="http://h:4317") ) @@ -719,6 +787,177 @@ def test_success_span_records_no_exception_event(): assert all(e.name != ExceptionEvent.NAME for e in span.events) +def _engine_with_event_recorder(): + from opentelemetry.sdk._logs.export import InMemoryLogExporter + + from litellm.integrations.otel.emitter import SpanEmitter + from litellm.integrations.otel.plumbing.events import GenAIEventRecorder + + cfg = OpenTelemetryV2Config(exporter="in_memory", enable_events=True) + provider, span_exporter = providers.in_memory_provider(cfg) + log_exporter = InMemoryLogExporter() + logger_provider = providers.build_logger_provider(cfg, log_exporter=log_exporter) + recorder = GenAIEventRecorder(providers.get_event_logger(logger_provider)) + engine = SpanEmitter(providers.get_tracer(provider, "t"), cfg, event_recorder=recorder) + return engine, span_exporter, log_exporter + + +def _llm_call_data(error): + return LLMCallSpanData( + operation=GenAIOperation.CHAT, + provider="openai", + request_model="gpt-4o", + response_model=None, + response_id=None, + request_params=LLMRequestParams(), + usage=LLMUsage(), + finish_reasons=(), + error=error, + response_cost=None, + server=None, + identity=RequestIdentity(call_id=None), + ) + + +def test_operation_exception_log_event_emitted_on_failed_llm_call(): + """A failed LLM call records the GenAI semconv ``gen_ai.client.operation.exception`` + event on the logs signal: severity WARN, the full ``exception.*`` trio (including + the stacktrace, which span-side only exists under a vendor key), correlated to + the failed span via trace/span ids. The span-side error surface stays intact.""" + from opentelemetry._logs.severity import SeverityNumber + + from litellm.integrations.otel.model.semconv import ExceptionEvent, GenAIEvent + + engine, span_exporter, log_exporter = _engine_with_event_recorder() + engine.emit( + SpanRole.LLM_CALL, + _llm_call_data( + SpanError( + error_type="RateLimitError", + message="rate limited", + code="429", + stack_trace="Traceback (most recent call last) ...", + llm_provider="openai", + ) + ), + ) + (span,) = span_exporter.get_finished_spans() + (log,) = log_exporter.get_finished_logs() + record = log.log_record + + assert record.attributes["event.name"] == GenAIEvent.OPERATION_EXCEPTION + assert record.severity_number == SeverityNumber.WARN + assert record.attributes[ExceptionEvent.TYPE] == "RateLimitError" + assert record.attributes[ExceptionEvent.MESSAGE] == "rate limited" + assert record.attributes[ExceptionEvent.STACKTRACE] == "Traceback (most recent call last) ..." + assert record.trace_id == span.context.trace_id + assert record.span_id == span.context.span_id + + assert [e.name for e in span.events] == [ExceptionEvent.NAME] + assert span.attributes["error.type"] == "RateLimitError" + + +def test_operation_exception_log_event_omits_absent_stacktrace(): + from litellm.integrations.otel.model.semconv import ExceptionEvent + + engine, _, log_exporter = _engine_with_event_recorder() + engine.emit(SpanRole.LLM_CALL, _llm_call_data(SpanError(error_type="APIError", message="boom"))) + (log,) = log_exporter.get_finished_logs() + + assert ExceptionEvent.STACKTRACE not in log.log_record.attributes + assert log.log_record.attributes[ExceptionEvent.MESSAGE] == "boom" + + +def test_operation_exception_log_event_always_carries_required_pair(): + """``exception.type`` and ``exception.message`` are the semconv-required pair: + they ride the event even when the recorder is handed empty strings, so an + event is never emitted with no required field. Only the stacktrace is + conditional.""" + from opentelemetry.sdk._logs.export import InMemoryLogExporter + from opentelemetry.trace import INVALID_SPAN_CONTEXT + + from litellm.integrations.otel.model.semconv import ExceptionEvent + from litellm.integrations.otel.plumbing.events import GenAIEventRecorder + + cfg = OpenTelemetryV2Config(exporter="in_memory", enable_events=True) + log_exporter = InMemoryLogExporter() + logger_provider = providers.build_logger_provider(cfg, log_exporter=log_exporter) + recorder = GenAIEventRecorder(providers.get_event_logger(logger_provider)) + + recorder.record_operation_exception( + span_context=INVALID_SPAN_CONTEXT, + error_type="", + message="", + stack_trace="", + timestamp_ns=None, + ) + (log,) = log_exporter.get_finished_logs() + attributes = log.log_record.attributes + assert attributes[ExceptionEvent.TYPE] == "" + assert attributes[ExceptionEvent.MESSAGE] == "" + assert ExceptionEvent.STACKTRACE not in attributes + + +def test_operation_exception_log_event_not_emitted_on_success(): + engine, span_exporter, log_exporter = _engine_with_event_recorder() + engine.emit(SpanRole.LLM_CALL, _llm_call_data(None)) + + assert len(span_exporter.get_finished_spans()) == 1 + assert log_exporter.get_finished_logs() == () + + +def test_operation_exception_log_event_only_for_llm_call_role(): + """The event is scoped to GenAI client operations; a failed guardrail span + keeps its span-side error surface but records no GenAI exception event.""" + engine, span_exporter, log_exporter = _engine_with_event_recorder() + engine.emit( + SpanRole.GUARDRAIL, + GuardrailSpanData("presidio", status="failure", error=SpanError(error_type="X", message="denied")), + ) + (span,) = span_exporter.get_finished_spans() + + assert span.attributes["error.type"] == "X" + assert log_exporter.get_finished_logs() == () + + +def test_resolve_logger_provider_honors_explicit_noop_optout(monkeypatch): + """A ``NoOpLoggerProvider`` global is an explicit operator opt-out from the logs + signal: resolve to ``None`` so no recorder (and so no event) is ever built, + rather than emitting into a provider that drops everything.""" + from opentelemetry import _logs + from opentelemetry._logs import NoOpLoggerProvider + + from litellm.integrations.otel.logger import OpenTelemetryV2 + + cfg = OpenTelemetryV2Config(exporter="in_memory", enable_events=True) + tracer_provider, _ = providers.in_memory_provider(cfg) + monkeypatch.setattr(_logs, "get_logger_provider", lambda: NoOpLoggerProvider()) + + assert providers.resolve_logger_provider(cfg) is None + logger = OpenTelemetryV2(config=cfg, tracer_provider=tracer_provider) + assert logger._emitter._event_recorder is None + + +def test_resolve_logger_provider_reuses_operator_sdk_global(monkeypatch): + """Events ride an operator-configured logs pipeline rather than a second one + built by litellm, so they land wherever the operator's other logs land.""" + from opentelemetry import _logs + from opentelemetry.sdk._logs.export import InMemoryLogExporter + + cfg = OpenTelemetryV2Config(exporter="in_memory", enable_events=True) + operator_provider = providers.build_logger_provider(cfg, log_exporter=InMemoryLogExporter()) + monkeypatch.setattr(_logs, "get_logger_provider", lambda: operator_provider) + + assert providers.resolve_logger_provider(cfg) is operator_provider + + +def test_operation_exception_event_keys_are_pinned(): + from litellm.integrations.otel.model.semconv import ExceptionEvent, GenAIEvent + + assert GenAIEvent.OPERATION_EXCEPTION == "gen_ai.client.operation.exception" + assert ExceptionEvent.STACKTRACE == "exception.stacktrace" + + # --- service taxonomy: which calls become spans, and of what kind ----------- # diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 697b9293eea..b5e077e3561 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -217,6 +217,61 @@ def test_async_log_failure_event_marks_error_status(): assert span.attributes["error.type"] == "RateLimitError" +def _logger_with_events(enable_events): + from opentelemetry.sdk._logs.export import InMemoryLogExporter + + cfg = OpenTelemetryV2Config(exporter="in_memory", enable_events=enable_events) + span_exporter = InMemorySpanExporter() + tracer_provider = providers.build_tracer_provider(cfg, exporter=span_exporter) + log_exporter = InMemoryLogExporter() + logger_provider = providers.build_logger_provider(cfg, log_exporter=log_exporter) + logger = OpenTelemetryV2(config=cfg, tracer_provider=tracer_provider, logger_provider=logger_provider) + return logger, span_exporter, log_exporter + + +def test_enable_events_records_operation_exception_through_failure_callback(): + """With ``enable_events`` on, a real failure callback records the GenAI + ``gen_ai.client.operation.exception`` log event, carrying the traceback from + the standard logging payload and correlated to the LLM-call span.""" + from litellm.integrations.otel.model.semconv import ExceptionEvent, GenAIEvent + + logger, span_exporter, log_exporter = _logger_with_events(enable_events=True) + payload = _payload( + status="failure", + error_information={ + "error_class": "RateLimitError", + "error_message": "429 rate limited", + "traceback": "Traceback (most recent call last) ...", + }, + ) + _emit_llm(logger, _kwargs(payload=payload), fail=True) + + (span,) = span_exporter.get_finished_spans() + (log,) = log_exporter.get_finished_logs() + record = log.log_record + assert record.attributes["event.name"] == GenAIEvent.OPERATION_EXCEPTION + assert record.attributes[ExceptionEvent.TYPE] == "RateLimitError" + assert record.attributes[ExceptionEvent.MESSAGE] == "429 rate limited" + assert record.attributes[ExceptionEvent.STACKTRACE] == "Traceback (most recent call last) ..." + assert record.trace_id == span.context.trace_id + assert record.span_id == span.context.span_id + + +def test_events_off_by_default_records_no_log_event_on_failure(): + """``enable_events`` defaults to off: even with a logs pipeline injected, a + failure records only the span-side error surface, no log event.""" + logger, span_exporter, log_exporter = _logger_with_events(enable_events=False) + payload = _payload( + status="failure", + error_information={"error_class": "RateLimitError", "error_message": "429"}, + ) + _emit_llm(logger, _kwargs(payload=payload), fail=True) + + assert len(span_exporter.get_finished_spans()) == 1 + assert log_exporter.get_finished_logs() == () + assert OpenTelemetryV2Config(exporter="in_memory").enable_events is False + + def test_sync_log_event_is_noop(): """V2 closes the span async-only; the sync callback runs out-of-context, so it no-ops (the span stays open on the carrier until the async callback).""" diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py b/tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py new file mode 100644 index 00000000000..ef844c80d4d --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py @@ -0,0 +1,359 @@ +""" +Unit tests for the NoOpMetric guard in _increment_remaining_budget_metrics +and the per-entity guards in _set_*_budget_metrics_after_api_request. + +Regression tests that the specific bug can never happen again: +when budget gauges are excluded from prometheus_metrics_config (and therefore +created as NoOpMetric instances), the DB/cache lookup helpers must not be called. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import NoOpMetric + +_BUDGET_EXCLUDED_CONFIG = [ + { + "group": "core-only", + "metrics": [ + "litellm_requests_metric", + "litellm_total_tokens_metric", + ], + } +] + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + old_config = litellm.prometheus_metrics_config + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + yield + litellm.prometheus_metrics_config = old_config + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def make_logger_with_budget_metrics_disabled() -> PrometheusLogger: + litellm.prometheus_metrics_config = _BUDGET_EXCLUDED_CONFIG + return PrometheusLogger() + + +def make_logger_with_all_metrics_enabled() -> PrometheusLogger: + litellm.prometheus_metrics_config = None + return PrometheusLogger() + + +COMMON_KWARGS = dict( + user_api_team="team-123", + user_api_team_alias="my-team", + user_api_key="hashed-key", + user_api_key_alias="my-key", + litellm_params={"metadata": {}}, + response_cost=0.001, + user_id="user-1", + user_api_key_org_id="org-1", +) + + +class TestBudgetGaugesAreNoopWhenExcluded: + def test_team_gauge_is_noop(self): + logger = make_logger_with_budget_metrics_disabled() + assert isinstance(logger.litellm_remaining_team_budget_metric, NoOpMetric) + + def test_api_key_gauge_is_noop(self): + logger = make_logger_with_budget_metrics_disabled() + assert isinstance(logger.litellm_remaining_api_key_budget_metric, NoOpMetric) + + def test_user_gauge_is_noop(self): + logger = make_logger_with_budget_metrics_disabled() + assert isinstance(logger.litellm_remaining_user_budget_metric, NoOpMetric) + + def test_org_gauge_is_noop(self): + logger = make_logger_with_budget_metrics_disabled() + assert isinstance(logger.litellm_remaining_org_budget_metric, NoOpMetric) + + def test_gauges_are_real_when_all_metrics_enabled(self): + logger = make_logger_with_all_metrics_enabled() + assert not isinstance(logger.litellm_remaining_team_budget_metric, NoOpMetric) + assert not isinstance(logger.litellm_remaining_api_key_budget_metric, NoOpMetric) + assert not isinstance(logger.litellm_remaining_user_budget_metric, NoOpMetric) + assert not isinstance(logger.litellm_remaining_org_budget_metric, NoOpMetric) + + +class TestTopLevelGuard: + @pytest.mark.asyncio + async def test_no_db_lookups_when_all_budget_gauges_are_noop(self): + """Regression: _increment_remaining_budget_metrics must return early + without any I/O when all four budget gauges are NoOpMetric.""" + logger = make_logger_with_budget_metrics_disabled() + + assemble_team = AsyncMock(return_value=MagicMock()) + assemble_key = AsyncMock(return_value=MagicMock()) + assemble_user = AsyncMock(return_value=MagicMock()) + + with ( + patch.object(logger, "_assemble_team_object", assemble_team), + patch.object(logger, "_assemble_key_object", assemble_key), + patch.object(logger, "_assemble_user_object", assemble_user), + ): + await logger._increment_remaining_budget_metrics(**COMMON_KWARGS) + + assemble_team.assert_not_called() + assemble_key.assert_not_called() + assemble_user.assert_not_called() + + @pytest.mark.asyncio + async def test_db_lookups_run_when_budget_gauges_are_real(self): + """When budget gauges are real Prometheus metrics, the assemble helpers + must be called so I/O proceeds normally.""" + logger = make_logger_with_all_metrics_enabled() + + assemble_team = AsyncMock( + return_value=MagicMock( + team_id="team-123", + team_alias="my-team", + spend=0.001, + max_budget=None, + budget_reset_at=None, + ) + ) + assemble_key = AsyncMock( + return_value=MagicMock( + token="hashed-key", + key_alias="my-key", + spend=0.001, + max_budget=None, + budget_reset_at=None, + ) + ) + assemble_user = AsyncMock( + return_value=MagicMock( + user_id="user-1", + spend=0.001, + max_budget=None, + budget_reset_at=None, + user_email=None, + user_alias=None, + ) + ) + + with ( + patch.object(logger, "_assemble_team_object", assemble_team), + patch.object(logger, "_assemble_key_object", assemble_key), + patch.object(logger, "_assemble_user_object", assemble_user), + patch.object(logger, "_set_team_budget_metrics", MagicMock()), + patch.object(logger, "_set_key_budget_metrics", MagicMock()), + patch.object(logger, "_set_user_budget_metrics", MagicMock()), + patch.object(logger, "_set_org_budget_metrics_after_api_request", AsyncMock()), + ): + await logger._increment_remaining_budget_metrics(**COMMON_KWARGS) + + assemble_team.assert_called_once() + assemble_key.assert_called_once() + assemble_user.assert_called_once() + + +class TestPerEntityGuards: + @pytest.mark.asyncio + async def test_team_guard_skips_lookup_when_team_gauge_is_noop(self): + """Per-entity guard: team assemble helper is not called when team gauge is NoOp, + even when key and user gauges are real.""" + logger = make_logger_with_all_metrics_enabled() + logger.litellm_remaining_team_budget_metric = NoOpMetric() + + assemble_team = AsyncMock(return_value=MagicMock()) + assemble_key = AsyncMock( + return_value=MagicMock( + token="hashed-key", + key_alias="my-key", + spend=0.001, + max_budget=None, + budget_reset_at=None, + ) + ) + assemble_user = AsyncMock( + return_value=MagicMock( + user_id="user-1", + spend=0.001, + max_budget=None, + budget_reset_at=None, + user_email=None, + user_alias=None, + ) + ) + + with ( + patch.object(logger, "_assemble_team_object", assemble_team), + patch.object(logger, "_assemble_key_object", assemble_key), + patch.object(logger, "_assemble_user_object", assemble_user), + patch.object(logger, "_set_team_budget_metrics", MagicMock()), + patch.object(logger, "_set_key_budget_metrics", MagicMock()), + patch.object(logger, "_set_user_budget_metrics", MagicMock()), + patch.object(logger, "_set_org_budget_metrics_after_api_request", AsyncMock()), + ): + await logger._increment_remaining_budget_metrics(**COMMON_KWARGS) + + assemble_team.assert_not_called() + assemble_key.assert_called_once() + assemble_user.assert_called_once() + + @pytest.mark.asyncio + async def test_key_guard_skips_lookup_when_key_gauge_is_noop(self): + """Per-entity guard: key assemble helper is not called when key gauge is NoOp, + even when team and user gauges are real.""" + logger = make_logger_with_all_metrics_enabled() + logger.litellm_remaining_api_key_budget_metric = NoOpMetric() + + assemble_team = AsyncMock( + return_value=MagicMock( + team_id="team-123", + team_alias="my-team", + spend=0.001, + max_budget=None, + budget_reset_at=None, + ) + ) + assemble_key = AsyncMock(return_value=MagicMock()) + assemble_user = AsyncMock( + return_value=MagicMock( + user_id="user-1", + spend=0.001, + max_budget=None, + budget_reset_at=None, + user_email=None, + user_alias=None, + ) + ) + + with ( + patch.object(logger, "_assemble_team_object", assemble_team), + patch.object(logger, "_assemble_key_object", assemble_key), + patch.object(logger, "_assemble_user_object", assemble_user), + patch.object(logger, "_set_team_budget_metrics", MagicMock()), + patch.object(logger, "_set_key_budget_metrics", MagicMock()), + patch.object(logger, "_set_user_budget_metrics", MagicMock()), + patch.object(logger, "_set_org_budget_metrics_after_api_request", AsyncMock()), + ): + await logger._increment_remaining_budget_metrics(**COMMON_KWARGS) + + assemble_key.assert_not_called() + assemble_team.assert_called_once() + assemble_user.assert_called_once() + + @pytest.mark.asyncio + async def test_user_guard_skips_lookup_when_user_gauge_is_noop(self): + """Per-entity guard: user assemble helper is not called when user gauge is NoOp, + even when team and key gauges are real.""" + logger = make_logger_with_all_metrics_enabled() + logger.litellm_remaining_user_budget_metric = NoOpMetric() + + assemble_team = AsyncMock( + return_value=MagicMock( + team_id="team-123", + team_alias="my-team", + spend=0.001, + max_budget=None, + budget_reset_at=None, + ) + ) + assemble_key = AsyncMock( + return_value=MagicMock( + token="hashed-key", + key_alias="my-key", + spend=0.001, + max_budget=None, + budget_reset_at=None, + ) + ) + assemble_user = AsyncMock(return_value=MagicMock()) + + with ( + patch.object(logger, "_assemble_team_object", assemble_team), + patch.object(logger, "_assemble_key_object", assemble_key), + patch.object(logger, "_assemble_user_object", assemble_user), + patch.object(logger, "_set_team_budget_metrics", MagicMock()), + patch.object(logger, "_set_key_budget_metrics", MagicMock()), + patch.object(logger, "_set_user_budget_metrics", MagicMock()), + patch.object(logger, "_set_org_budget_metrics_after_api_request", AsyncMock()), + ): + await logger._increment_remaining_budget_metrics(**COMMON_KWARGS) + + assemble_user.assert_not_called() + assemble_team.assert_called_once() + assemble_key.assert_called_once() + + @pytest.mark.asyncio + async def test_set_team_budget_metrics_directly_skips_when_gauge_is_noop(self): + """_set_team_budget_metrics_after_api_request returns early when team gauge is NoOp.""" + logger = make_logger_with_budget_metrics_disabled() + assemble_team = AsyncMock(return_value=MagicMock()) + + with patch.object(logger, "_assemble_team_object", assemble_team): + await logger._set_team_budget_metrics_after_api_request( + user_api_team="team-123", + user_api_team_alias="my-team", + team_spend=0.5, + team_max_budget=10.0, + response_cost=0.001, + ) + + assemble_team.assert_not_called() + + @pytest.mark.asyncio + async def test_set_api_key_budget_metrics_directly_skips_when_gauge_is_noop(self): + """_set_api_key_budget_metrics_after_api_request returns early when key gauge is NoOp.""" + logger = make_logger_with_budget_metrics_disabled() + assemble_key = AsyncMock(return_value=MagicMock()) + + with patch.object(logger, "_assemble_key_object", assemble_key): + await logger._set_api_key_budget_metrics_after_api_request( + user_api_key="hashed-key", + user_api_key_alias="my-key", + response_cost=0.001, + key_max_budget=10.0, + key_spend=0.5, + ) + + assemble_key.assert_not_called() + + @pytest.mark.asyncio + async def test_set_user_budget_metrics_directly_skips_when_gauge_is_noop(self): + """_set_user_budget_metrics_after_api_request returns early when user gauge is NoOp.""" + logger = make_logger_with_budget_metrics_disabled() + assemble_user = AsyncMock(return_value=MagicMock()) + + with patch.object(logger, "_assemble_user_object", assemble_user): + await logger._set_user_budget_metrics_after_api_request( + user_id="user-1", + user_spend=0.5, + user_max_budget=10.0, + response_cost=0.001, + ) + + assemble_user.assert_not_called() + + @pytest.mark.asyncio + async def test_set_org_budget_metrics_directly_skips_when_gauge_is_noop(self): + """_set_org_budget_metrics_after_api_request returns early when org gauge is NoOp. + The guard fires before any import of auth_checks, so prisma_client is never touched.""" + logger = make_logger_with_budget_metrics_disabled() + + set_org_metrics = MagicMock() + with patch.object(logger, "_set_org_budget_metrics", set_org_metrics): + await logger._set_org_budget_metrics_after_api_request( + org_id="org-1", + response_cost=0.001, + ) + + set_org_metrics.assert_not_called() diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 410014958b6..cf74bed9c15 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -16,6 +16,7 @@ sys.path.insert(0, os.path.abspath("../../..")) import litellm from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, + match_all_fallback_generalizations, match_fallback_generalization, set_fallback_generalizations, ) @@ -57,6 +58,48 @@ def test_match_returns_model_info_of_first_matching_rule(restore_generalizations assert matched["tag"] == "first" +def test_match_all_returns_every_matching_rule_in_order(restore_generalizations): + restore_generalizations( + [ + { + "name": "first", + "pattern": r"^acme-", + "model_info": {"litellm_provider": "openai", "tag": "first"}, + }, + { + "name": "second", + "pattern": r"^acme-pro-", + "model_info": {"litellm_provider": "anthropic", "tag": "second"}, + }, + ] + ) + assert [m["tag"] for m in match_all_fallback_generalizations("acme-pro-1")] == ["first", "second"] + assert match_all_fallback_generalizations("gpt-4o") == [] + + +def test_provider_scoped_rule_is_skipped_for_other_providers(restore_generalizations): + """Model-info resolution must fall through a provider-mismatched earlier rule to a + later applicable one, instead of discarding the model name at the first pattern hit.""" + restore_generalizations( + [ + { + "name": "bedrock-scoped", + "pattern": r"^acme-", + "model_info": {"litellm_provider": "bedrock", "supports_vision": False}, + }, + { + "name": "openai-scoped", + "pattern": r"^acme-", + "model_info": {"litellm_provider": "openai", "mode": "chat", "supports_vision": True}, + }, + ] + ) + litellm.get_model_info.cache_clear() + info = litellm.get_model_info("acme-fast-1", custom_llm_provider="openai") + assert info["litellm_provider"] == "openai" + assert info["supports_vision"] is True + + def test_match_is_case_insensitive(restore_generalizations): restore_generalizations( [{"name": "r", "pattern": r"^claude-opus", "model_info": {"ok": True}}] @@ -298,3 +341,47 @@ def test_shipped_adaptive_rule_gates_on_version_not_pricing(shipped_cost_map): assert non_adaptive not in litellm.model_cost assert AnthropicModelInfo._is_adaptive_thinking_model(adaptive) is True assert AnthropicModelInfo._is_adaptive_thinking_model(non_adaptive) is False + + +def test_shipped_bedrock_rule_resolves_unmapped_future_claude_for_bedrock(shipped_cost_map): + """An unmapped Bedrock Claude >= 4.8 resolves via the bedrock-scoped + ``bedrock-anthropic-claude-mid-conversation-system`` rule even when the lookup + carries ``custom_llm_provider="bedrock"``, which the provider check uses to drop + the anthropic-scoped rules. It inherits base capabilities, gains both + version-gated flags, and stays unpriced.""" + model = "us.anthropic.claude-opus-4-9" + assert model not in litellm.model_cost + info = litellm.get_model_info(model, custom_llm_provider="bedrock") + assert info["litellm_provider"] == "bedrock" + assert info["supports_mid_conversation_system"] is True + assert info["supports_adaptive_thinking"] is True + assert info["supports_function_calling"] is True + assert not info.get("input_cost_per_token") + + +def test_shipped_bedrock_mid_conversation_rule_gates_on_version_and_naming(shipped_cost_map): + """The bedrock rule only claims Bedrock-style ids at 4.8+, bare 5+ majors and + new families included; pre-4.8 Bedrock ids and native ids never gain the flag, + and the rule outranks the anthropic-scoped ones for Bedrock ids because it is + listed first.""" + for flagged in ( + "us.anthropic.claude-opus-4-8", + "jp.anthropic.claude-opus-4-8", + "anthropic.claude-sonnet-5", + "us.anthropic.claude-fable-5", + "anthropic.claude-sonnet-5-20260101-v1:0", + ): + matched = match_fallback_generalization(flagged) + assert matched is not None, flagged + assert matched["litellm_provider"] == "bedrock", flagged + assert matched["supports_mid_conversation_system"] is True, flagged + for unflagged in ( + "us.anthropic.claude-opus-4-7", + "us.anthropic.claude-sonnet-4-6", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "claude-opus-4-9", + "claude-sonnet-5", + ): + matched = match_fallback_generalization(unflagged) + assert matched is None or not matched.get("supports_mid_conversation_system"), unflagged diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py new file mode 100644 index 00000000000..06d3effcfbb --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py @@ -0,0 +1,189 @@ +import pytest + +from litellm.constants import ( + DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, +) +from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) + + +def _claude_code_payload(effort="medium", max_tokens=8192, **output_config_extra): + """The exact adaptive-thinking shape Claude Code (claude-cli) sends.""" + output_config = {"effort": effort, **output_config_extra} + return { + "max_tokens": max_tokens, + "thinking": {"type": "adaptive"}, + "output_config": output_config, + } + + +def _transform(model, params, litellm_params=None): + return AnthropicMessagesConfig().transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=dict(params), + litellm_params=litellm_params or {}, + headers={}, + ) + + +def test_effort_translated_to_legacy_thinking_for_haiku_4_5(): + """Core regression: Claude Code sends adaptive thinking + effort to Haiku 4.5 + (thinking-capable, pre-4.6). Effort must be translated to legacy extended + thinking rather than forwarded raw (which Anthropic rejects with "This model + does not support the effort parameter").""" + result = _transform("claude-haiku-4-5", _claude_code_payload(effort="medium")) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert "output_config" not in result + + +def test_effort_high_maps_to_high_budget_for_sonnet_4_5(): + result = _transform("claude-sonnet-4-5", _claude_code_payload(effort="high")) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + } + assert "output_config" not in result + + +def test_adaptive_effort_passes_through_untouched_for_4_6(): + """4.6+ natively supports the adaptive interface, so it must not be rewritten.""" + result = _transform("claude-sonnet-4-6", _claude_code_payload(effort="high")) + + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} + + +def test_thinking_and_effort_dropped_for_non_reasoning_model(): + """A model with no reasoning support cannot take thinking or effort, so both are + silently dropped (no drop_params required) so the request still succeeds.""" + result = _transform("claude-3-5-haiku-latest", _claude_code_payload(effort="medium")) + + assert "thinking" not in result + assert "output_config" not in result + + +def test_residual_output_config_preserved_after_effort_translation(): + """output_config may carry `format` (structured outputs) alongside effort. Only + the consumed effort key is removed; the residual is left for provider subclasses + (bedrock/vertex) to handle, and effort is translated to legacy thinking.""" + result = _transform( + "claude-haiku-4-5", + _claude_code_payload(effort="medium", format={"type": "json_schema"}), + ) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert result["output_config"] == {"format": {"type": "json_schema"}} + + +def test_opus_4_5_keeps_effort_but_drops_adaptive_thinking(): + """Regression: Opus 4.5 advertises supports_output_config (accepts + output_config.effort) but is NOT adaptive, so thinking:{type:adaptive} is + rejected by Anthropic. The effort must be kept and only the adaptive thinking + block dropped, rather than early-returning and forwarding adaptive thinking raw.""" + result = _transform("claude-opus-4-5", _claude_code_payload(effort="medium")) + + assert result["output_config"] == {"effort": "medium"} + assert "thinking" not in result + + +def test_opus_4_5_preserves_native_effort_without_adaptive_thinking(): + """A caller sending output_config.effort alone (no adaptive thinking) to Opus 4.5 + must pass through untouched, since the model supports it natively.""" + result = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-opus-4-5", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={ + "max_tokens": 8192, + "output_config": {"effort": "high"}, + }, + litellm_params={}, + headers={}, + ) + + assert result["output_config"] == {"effort": "high"} + assert "thinking" not in result + + +def test_opus_4_5_unsupported_effort_level_translated_to_legacy_thinking(): + """Opus 4.5 accepts output_config.effort but only levels low/medium/high; + Claude Code defaults to xhigh on newer models, and forwarding that level raw + would be rejected with "effort='xhigh' is not supported by this model". An + unsupported level must fall through to the legacy translation (budget-based + thinking, effort stripped) instead of being preserved.""" + result = _transform("claude-opus-4-5", _claude_code_payload(effort="xhigh", max_tokens=64000)) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, + } + assert "output_config" not in result + + +def test_opus_4_5_effort_only_unsupported_level_left_for_provider_normalization(): + """An effort-only request (no adaptive thinking) must pass through untouched even + when the level exceeds what the model supports: provider subclasses own their + level normalization (bedrock clamps xhigh to the model's ceiling after this base + transform runs), so consuming the effort here breaks that contract.""" + result = _transform( + "claude-opus-4-5", + {"max_tokens": 4096, "output_config": {"effort": "xhigh"}}, + ) + + assert result["output_config"] == {"effort": "xhigh"} + assert "thinking" not in result + + +def test_budget_capped_below_max_tokens(): + """Adaptive thinking carries no budget, so the translated legacy budget must be + capped below max_tokens (Anthropic requires max_tokens > budget_tokens). A + high-effort budget (4096) with max_tokens=3000 must be capped to 2999.""" + result = _transform("claude-haiku-4-5", _claude_code_payload(effort="high", max_tokens=3000)) + + assert result["thinking"] == {"type": "enabled", "budget_tokens": 2999} + + +def test_thinking_dropped_when_max_tokens_too_small_for_min_budget(): + """When max_tokens can't fit even the minimum thinking budget, thinking is + silently dropped so the request still succeeds rather than being rejected.""" + result = _transform("claude-haiku-4-5", _claude_code_payload(effort="medium", max_tokens=512)) + + assert "thinking" not in result + assert "output_config" not in result + + +def test_unrecognized_effort_raises_clean_400(): + """An unrecognized effort value (e.g. a future Anthropic tier) must surface as a + clean AnthropicError 400, matching _translate_reasoning_effort_to_anthropic, + rather than leaking litellm's internal BadRequestError.""" + with pytest.raises(AnthropicError) as exc_info: + _transform("claude-haiku-4-5", _claude_code_payload(effort="turbo")) + + assert exc_info.value.status_code == 400 + + +def test_non_adaptive_request_without_effort_is_untouched(): + """A non-adaptive model receiving a request with no adaptive interface (no + effort, no adaptive thinking) must pass through untouched.""" + result = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params={"max_tokens": 1024}, + litellm_params={}, + headers={}, + ) + + assert "thinking" not in result + assert "output_config" not in result diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a2fbb68bbb9..f7a234c8ae8 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -772,7 +772,6 @@ class TestProxyOAuthHeaderForwarding: self, ): """OAuth Authorization header IS forwarded when x-litellm-api-key was used for proxy auth.""" - from unittest.mock import patch from starlette.datastructures import Headers @@ -1724,3 +1723,48 @@ class TestClaudeOpus48AdaptiveThinking: from litellm.llms.anthropic.common_utils import AnthropicModelInfo assert AnthropicModelInfo._is_adaptive_thinking_model(model) is False + + +class TestDefaultSuffixAdaptiveThinking: + """@default-suffixed Vertex AI model names (e.g. vertex_ai/claude-opus-4-8@default) + must resolve as adaptive thinking. Before the fix, _model_map_lookup_candidates + never stripped the @default suffix, so the lookup fell through to the bare + model name without @default, which may or may not have the flag, and for + provider-prefixed forms the lookup always missed (issue #31760).""" + + @pytest.mark.parametrize( + "model", + [ + "vertex_ai/claude-opus-4-8@default", + "vertex_ai/claude-sonnet-4-6@default", + "vertex_ai/claude-opus-4-7@default", + "vertex_ai/claude-opus-4-6@default", + "vertex_ai/claude-fable-5@default", + ], + ) + def test_default_suffix_models_are_adaptive_thinking( + self, local_model_cost_map, model: str + ) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo._is_adaptive_thinking_model(model) is True, ( + f"{model} not classified as adaptive thinking. " + "Check _model_map_lookup_candidates strips @default suffix." + ) + + @pytest.mark.parametrize( + "model,expected_bare", + [ + ("vertex_ai/claude-opus-4-8@default", "claude-opus-4-8"), + ("vertex_ai/claude-sonnet-4-6@default", "claude-sonnet-4-6"), + ], + ) + def test_lookup_candidates_include_bare_name( + self, model: str, expected_bare: str + ) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + candidates = AnthropicModelInfo._model_map_lookup_candidates(model) + assert expected_bare in candidates, ( + f"Expected '{expected_bare}' in candidates for '{model}', got: {candidates}" + ) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 532c6ff3598..7d104ff1f2b 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1974,6 +1974,217 @@ def test_bedrock_invoke_transform_merges_list_content_system_role_into_system(): ] +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-opus-4-8", + "jp.anthropic.claude-opus-4-8", + "us.anthropic.claude-sonnet-5", + "us.anthropic.claude-fable-5", + ], +) +def test_bedrock_invoke_transform_keeps_mid_conversation_system_role_in_place(local_model_cost_map, model): + """Regression test for the Bedrock prompt-cache collapse: hoisting a + mid-conversation ``role: "system"`` message (e.g. Claude Code's + ``mid-conversation-system-2026-04-07`` reminders) into the top-level + ``system`` field mutates the cache prefix and invalidates the cached message + history, so on models flagged ``supports_mid_conversation_system`` (Claude + 4.8+, which Invoke accepts the role on) such entries must be forwarded + in place. Billing-header blocks must still be stripped from the top-level + ``system`` field even when nothing is hoisted.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [ + {"role": "user", "content": "read the file"}, + {"role": "system", "content": "[Truncated: PARTIAL view of big1.txt]"}, + {"role": "assistant", "content": "reading"}, + {"role": "user", "content": "continue"}, + ] + + result = cfg.transform_anthropic_messages_request( + model=model, + messages=copy.deepcopy(messages), + anthropic_messages_optional_request_params={ + "max_tokens": 256, + "stream": False, + "system": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.205;"}, + {"type": "text", "text": "Base.", "cache_control": {"type": "ephemeral"}}, + ], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["messages"] == messages + assert result["system"] == [ + {"type": "text", "text": "Base.", "cache_control": {"type": "ephemeral"}} + ] + + +def test_bedrock_invoke_transform_hoists_only_leading_system_run(local_model_cost_map): + """On models flagged ``supports_mid_conversation_system``, only the leading + run of ``role: "system"`` messages is hoisted into the top-level ``system`` + field; a later system entry keeps its position in ``messages`` so the + serialized prefix stays stable across turns.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": "Cite sources."}, + {"role": "user", "content": "hi"}, + {"role": "system", "content": "mid-conversation reminder"}, + {"role": "user", "content": "continue"}, + ] + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-8", + messages=copy.deepcopy(messages), + anthropic_messages_optional_request_params={"max_tokens": 256, "stream": False}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["messages"] == [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "mid-conversation reminder"}, + {"role": "user", "content": "continue"}, + ] + assert result["system"] == [ + {"type": "text", "text": "You are terse."}, + {"type": "text", "text": "Cite sources."}, + ] + + +def test_bedrock_invoke_transform_hoists_mid_conversation_system_for_older_claude(local_model_cost_map): + """Regression test for Claude Code 400s on pre-Opus-4.8 Bedrock models: + Invoke rejects ``role: "system"`` in every position on Opus 4.7, Sonnet 4.6, + Haiku 4.5, etc. ("role 'system' is not supported on this model"), so on + models without ``supports_mid_conversation_system`` every system entry must + be hoisted into the top-level ``system`` field, mid-conversation ones + included.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [ + {"role": "user", "content": "read the file"}, + {"role": "system", "content": "[Truncated: PARTIAL view of big1.txt]"}, + {"role": "assistant", "content": "reading"}, + {"role": "user", "content": "continue"}, + ] + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-7", + messages=copy.deepcopy(messages), + anthropic_messages_optional_request_params={ + "max_tokens": 256, + "stream": False, + "system": [{"type": "text", "text": "Base."}], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["messages"] == [ + {"role": "user", "content": "read the file"}, + {"role": "assistant", "content": "reading"}, + {"role": "user", "content": "continue"}, + ] + assert result["system"] == [ + {"type": "text", "text": "Base."}, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ] + + +def test_bedrock_invoke_transform_hoists_all_system_for_unmapped_model(local_model_cost_map): + """A model with no cost-map entry and no fallback-generalization rule gets + the hoist-everything behavior: the safe default is a mutated cache prefix, + never a provider 400 from forwarding a role the model may not accept.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "mid-conversation reminder"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "continue"}, + ] + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-3-9", + messages=copy.deepcopy(messages), + anthropic_messages_optional_request_params={"max_tokens": 256, "stream": False}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["messages"] == [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "continue"}, + ] + assert result["system"] == [{"type": "text", "text": "mid-conversation reminder"}] + + +def test_bedrock_invoke_transform_keeps_system_in_place_for_unmapped_future_claude(local_model_cost_map): + """An unmapped Bedrock Claude at 4.8 or higher resolves through the + ``bedrock-anthropic-claude-mid-conversation-system`` fallback rule, so a + future model that has not landed in the cost map yet keeps the + cache-preserving in-place behavior instead of falling back to hoist-all.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "system", "content": "mid-conversation reminder"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "continue"}, + ] + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-9", + messages=copy.deepcopy(messages), + anthropic_messages_optional_request_params={"max_tokens": 256, "stream": False}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result["messages"] == messages + assert "system" not in result + + +def test_bedrock_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): + """Exact cost-map hits resolve before fallback-generalization rules, so a + mapped Bedrock Claude 4.8+ entry without ``supports_mid_conversation_system`` + silently loses the cache-preserving in-place handling that the + ``bedrock-anthropic-claude-mid-conversation-system`` rule grants unmapped + ids. Every mapped entry the rule's own pattern matches must carry the flag + explicitly.""" + import re + + import litellm + + cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") + with open(cost_map_path) as f: + cost_map = json.load(f) + rules = cost_map["fallback_generalizations"]["rules"] + pattern = re.compile( + next(r["pattern"] for r in rules if r["name"] == "bedrock-anthropic-claude-mid-conversation-system"), + re.IGNORECASE, + ) + missing = [ + key + for key, info in cost_map.items() + if isinstance(info, dict) + and str(info.get("litellm_provider", "")).startswith("bedrock") + and pattern.search(key) + and info.get("supports_mid_conversation_system") is not True + ] + assert missing == [] + + def test_as_system_content_blocks_handles_each_shape(): """``_as_system_content_blocks`` normalizes every system shape: ``None`` -> empty, a string -> a single text block, a list -> a shallow copy, and any other value diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py new file mode 100644 index 00000000000..dca9dc19314 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -0,0 +1,218 @@ +"""Spec tests for the DCR-bridge envelope producer and consumer helpers. + +These pin the contracts the token endpoint and admission edge depend on: the master-key +key derivation is deterministic, domain-separated (keyed HMAC), and always yields a +>= 32-byte signing key; the ``Authorization`` classifier is total over the cases admission +branches on (non-envelope, valid envelope bound to this server, envelope-shaped-but- +unopenable, and envelope minted for a different server); the producer helper round-trips +through the consumer; and no path leaks the upstream token in a repr. +""" + +from datetime import datetime, timedelta, timezone + +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( + BridgeEnvelopeAdmitted, + BridgeEnvelopeInvalid, + NotBridgeEnvelope, + build_bridge_token_response, + envelope_keys_from_master_key, + is_bridge_envelope_shaped, + resolve_bridge_envelope, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + ENVELOPE_PREFIX, + EnvelopeIdentity, + EnvelopeKeys, + EnvelopeTooLarge, + SealedEnvelope, + UpstreamTokenGrant, + mint_envelope, +) + +_NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) +_MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" +_ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" +_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") +_SERVER_ID = _IDENTITY.server_id + + +def _grant() -> UpstreamTokenGrant: + return UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=600) + + +def _sealed_token(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeIdentity = _IDENTITY) -> str: + sealed = mint_envelope(identity, _grant(), keys, now) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def test_key_derivation_is_deterministic(): + assert envelope_keys_from_master_key(_MASTER_KEY) == envelope_keys_from_master_key(_MASTER_KEY) + + +def test_key_derivation_signing_and_encryption_differ(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + assert keys.signing_key.get_secret_value() != keys.encryption_key.get_secret_value() + + +def test_key_derivation_differs_by_master_key(): + a = envelope_keys_from_master_key(_MASTER_KEY) + b = envelope_keys_from_master_key(_MASTER_KEY + "x") + assert a.signing_key.get_secret_value() != b.signing_key.get_secret_value() + assert a.encryption_key.get_secret_value() != b.encryption_key.get_secret_value() + + +def test_key_derivation_signing_key_meets_hs256_floor_for_short_master_key(): + keys = envelope_keys_from_master_key("x") + assert len(keys.signing_key.get_secret_value()) >= 32 + + +def test_key_derivation_is_cached_so_the_memory_hard_kdf_runs_once_per_key(): + """The scrypt KDF is intentionally expensive to resist offline guessing, so it must be cached: + repeated calls for the same master key return the identical object rather than re-deriving, + keeping the per-request admission path free. Returning a distinct object each call would mean + the cache was dropped and every open would pay the memory-hard cost.""" + first = envelope_keys_from_master_key("sk-cache-probe-key-9988776655") + assert envelope_keys_from_master_key("sk-cache-probe-key-9988776655") is first + + +def test_derived_keys_round_trip_mint_and_open(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert result.identity == _IDENTITY + + +def test_resolve_non_envelope_is_not_bridge_envelope(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + assert isinstance(resolve_bridge_envelope("Bearer sk-some-litellm-key", keys, _NOW, _SERVER_ID), NotBridgeEnvelope) + assert isinstance(resolve_bridge_envelope("plain-token", keys, _NOW, _SERVER_ID), NotBridgeEnvelope) + + +def test_resolve_valid_envelope_returns_identity_and_upstream_authorization(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert result.identity == _IDENTITY + assert result.upstream_authorization.get_secret_value() == f"Bearer {_ACCESS_TOKEN}" + + +def test_resolve_strips_optional_bearer_scheme_before_detection(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + token = _sealed_token(keys) + assert token.startswith(ENVELOPE_PREFIX) + bare = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) + prefixed = resolve_bridge_envelope(f"Bearer {token}", keys, _NOW, _SERVER_ID) + lower = resolve_bridge_envelope(f"bearer {token}", keys, _NOW, _SERVER_ID) + assert isinstance(bare, BridgeEnvelopeAdmitted) + assert isinstance(prefixed, BridgeEnvelopeAdmitted) + assert isinstance(lower, BridgeEnvelopeAdmitted) + assert prefixed.upstream_authorization.get_secret_value() == bare.upstream_authorization.get_secret_value() + + +def test_resolve_expired_envelope_is_invalid_not_admitted(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + token = _sealed_token(keys, now=_NOW) + later = _NOW + timedelta(seconds=601) + assert isinstance(resolve_bridge_envelope(token, keys, later, _SERVER_ID), BridgeEnvelopeInvalid) + + +def test_resolve_envelope_minted_under_a_different_master_key_is_invalid(): + minted = envelope_keys_from_master_key(_MASTER_KEY) + other = envelope_keys_from_master_key("a-completely-different-master-key") + assert isinstance(resolve_bridge_envelope(_sealed_token(minted), other, _NOW, _SERVER_ID), BridgeEnvelopeInvalid) + + +def test_resolve_tampered_envelope_is_invalid(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + token = _sealed_token(keys) + tampered = token[:-4] + ("aaaa" if token[-4:] != "aaaa" else "bbbb") + result = resolve_bridge_envelope(tampered, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_resolve_envelope_minted_for_another_server_is_invalid(): + """An envelope sealed for server A must be rejected when presented to server B, so a + captured or misrouted envelope cannot forward one server's upstream credential to + another. The valid access token stays sealed; the mismatch alone fails the resolve.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + other_server_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-OTHER") + token = _sealed_token(keys, identity=other_server_identity) + result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_resolve_matching_server_binding_is_admitted(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, "srv-456") + assert isinstance(result, BridgeEnvelopeAdmitted) + + +def test_resolve_non_ascii_server_id_stays_total_and_does_not_raise(): + """The server-binding check must not raise on a non-ASCII server_id (an admin can register a + unicode server_id); it stays total and returns a typed result. A matching non-ASCII id admits, + a mismatching one is BridgeEnvelopeInvalid, and neither raises.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + unicode_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-café") + token = _sealed_token(keys, identity=unicode_identity) + assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-café"), BridgeEnvelopeAdmitted) + assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-cafe"), BridgeEnvelopeInvalid) + + +def test_resolve_strips_bearer_with_extra_whitespace(): + """A Bearer scheme separated by extra spaces or a tab still yields the envelope, so a client + using non-minimal but legal whitespace is not misclassified as a non-envelope.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + token = _sealed_token(keys) + for header in (f"Bearer {token}", f"Bearer\t{token}", f" Bearer {token}"): + assert isinstance(resolve_bridge_envelope(header, keys, _NOW, _SERVER_ID), BridgeEnvelopeAdmitted) + + +def test_admitted_result_repr_never_leaks_upstream_token(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert _ACCESS_TOKEN not in repr(result) + assert _ACCESS_TOKEN not in str(result) + + +def test_build_bridge_token_response_round_trips_through_the_consumer(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + sealed = build_bridge_token_response(_IDENTITY, _grant(), keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + assert sealed.token.get_secret_value().startswith(ENVELOPE_PREFIX) + result = resolve_bridge_envelope(sealed.token.get_secret_value(), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert result.identity == _IDENTITY + assert result.upstream_authorization.get_secret_value() == f"Bearer {_ACCESS_TOKEN}" + + +def test_build_bridge_token_response_oversized_grant_returns_error_value(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + huge = UpstreamTokenGrant(access_token=SecretStr("x" * 20000), token_type="Bearer") + result = build_bridge_token_response(_IDENTITY, huge, keys, _NOW) + assert isinstance(result, EnvelopeTooLarge) + + +def test_build_bridge_token_response_repr_never_leaks_upstream_token(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + sealed = build_bridge_token_response(_IDENTITY, _grant(), keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + assert _ACCESS_TOKEN not in repr(sealed) + assert _ACCESS_TOKEN not in str(sealed) + + +def test_is_bridge_envelope_shaped_detects_envelope_with_and_without_bearer(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + token = _sealed_token(keys) + assert is_bridge_envelope_shaped(token) is True + assert is_bridge_envelope_shaped(f"Bearer {token}") is True + assert is_bridge_envelope_shaped(f"bearer {token}") is True + + +def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer(): + assert is_bridge_envelope_shaped("Bearer sk-some-litellm-key") is False + assert is_bridge_envelope_shaped("plain-upstream-token") is False + assert is_bridge_envelope_shaped("") is False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py new file mode 100644 index 00000000000..71de206aa1e --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -0,0 +1,487 @@ +"""Spec tests for the sealed-envelope module (oauth_delegate DCR bridge). + +The envelope is the single client-held bearer carrying both a litellm identity and the +encrypted upstream grant, with zero server-side storage. These tests pin the security +contract: an envelope opens only under the exact keys that minted it, tampering with any +signed byte is detected, expiry is enforced against the injected clock (capped by the +module TTL ceiling), oversized envelopes are rejected rather than truncated, and no +error value, model repr, or raised exception ever contains the inner access token. +""" + +import base64 +import hashlib +import hmac +import json +from datetime import datetime, timedelta, timezone + +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import SecretStr, ValidationError + +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + ENVELOPE_ISSUER, + ENVELOPE_PREFIX, + MAX_ENVELOPE_BYTES, + MAX_ENVELOPE_TTL_SECONDS, + BadSignature, + DecryptFailed, + EnvelopeIdentity, + EnvelopeKeys, + EnvelopeTooLarge, + Expired, + MalformedPayload, + NotAnEnvelope, + OpenedEnvelope, + SealedEnvelope, + UpstreamTokenGrant, + is_envelope, + mint_envelope, + open_envelope, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value + +_NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) +_SIGNING_KEY = "unit-test-signing-key-0123456789abcdef0123456789abcdef" +_ENCRYPTION_KEY = "unit-test-encryption-key-fedcba9876543210fedcba9876543210" +_OTHER_SIGNING_KEY = "other-signing-key-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" +_OTHER_ENCRYPTION_KEY = "other-encryption-key-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" +_KEYS = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_ENCRYPTION_KEY)) +_WRONG_SIGNING = EnvelopeKeys(signing_key=SecretStr(_OTHER_SIGNING_KEY), encryption_key=SecretStr(_ENCRYPTION_KEY)) +_WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) +_ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" +_REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" +_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") + + +def _full_grant() -> UpstreamTokenGrant: + return UpstreamTokenGrant( + access_token=SecretStr(_ACCESS_TOKEN), + token_type="Bearer", + refresh_token=SecretStr(_REFRESH_TOKEN), + scope="read:tools write:tools", + expires_in=600, + ) + + +def _minimal_grant() -> UpstreamTokenGrant: + return UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer") + + +def _sealed_token(grant: UpstreamTokenGrant, keys: EnvelopeKeys = _KEYS) -> str: + sealed = mint_envelope(_IDENTITY, grant, keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def _unverified_claims(sealed_token: str) -> dict[str, object]: + return jwt.decode(sealed_token.removeprefix(ENVELOPE_PREFIX), options={"verify_signature": False}) + + +def _forge(claims: dict[str, object], signing_key: str = _SIGNING_KEY) -> str: + return ENVELOPE_PREFIX + jwt.encode(claims, signing_key, algorithm="HS256") + + +def _b64url(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + + +def _hand_crafted_hs256(payload: dict[str, object], signing_key: str = _SIGNING_KEY) -> str: + """Assemble an HS256 envelope from raw bytes, bypassing PyJWT's encode-side claim + guards (it refuses to build a token with a non-string ``iss``). This is the real + attacker path: a client crafts the compact JWT directly, so any registered claim can + carry a hostile JSON type.""" + header = _b64url(json.dumps({"alg": "HS256", "typ": "JWT"}).encode("utf-8")) + body = _b64url(json.dumps(payload).encode("utf-8")) + signing_input = f"{header}.{body}".encode("ascii") + signature = _b64url(hmac.new(signing_key.encode("utf-8"), signing_input, hashlib.sha256).digest()) + return ENVELOPE_PREFIX + f"{header}.{body}.{signature}" + + +def _tampered(sealed_token: str, segment: int, index: int) -> str: + parts = sealed_token.removeprefix(ENVELOPE_PREFIX).split(".") + original = parts[segment][index] + replacement = "A" if original in "QRST" else "Q" + mutated = parts[segment][:index] + replacement + parts[segment][index + 1 :] + rebuilt = ".".join(parts[:segment] + [mutated] + parts[segment + 1 :]) + return ENVELOPE_PREFIX + rebuilt + + +def test_round_trip_recovers_identity_and_grant_exactly(): + grant = _full_grant() + token = _sealed_token(grant) + assert is_envelope(token) + opened = open_envelope(token, _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity == _IDENTITY + assert opened.grant == grant + assert opened.grant.access_token.get_secret_value() == _ACCESS_TOKEN + assert opened.grant.refresh_token is not None + assert opened.grant.refresh_token.get_secret_value() == _REFRESH_TOKEN + + +def test_minimal_grant_round_trips_without_none_leakage_into_claims(): + token = _sealed_token(_minimal_grant()) + claims = _unverified_claims(token) + blob = claims["grant"] + assert isinstance(blob, str) + plaintext = decrypt_value(value=base64.urlsafe_b64decode(blob), signing_key=_ENCRYPTION_KEY) + assert set(json.loads(plaintext)) == {"access_token", "token_type"} + opened = open_envelope(token, _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.grant.refresh_token is None + assert opened.grant.scope is None + assert opened.grant.expires_in is None + + +def test_claim_layout_and_no_plaintext_token_in_envelope(): + token = _sealed_token(_full_grant()) + claims = _unverified_claims(token) + assert set(claims) == {"iss", "iat", "exp", "user_id", "server_id", "grant"} + assert claims["iss"] == ENVELOPE_ISSUER + assert claims["iat"] == int(_NOW.timestamp()) + assert claims["exp"] == int(_NOW.timestamp()) + 600 + assert claims["user_id"] == "user-123" + assert claims["server_id"] == "srv-456" + assert _ACCESS_TOKEN not in token + assert _ACCESS_TOKEN not in json.dumps(claims) + assert _REFRESH_TOKEN not in json.dumps(claims) + + +@pytest.mark.parametrize( + "expires_in, expected_ttl", + [ + (600, 600), + (MAX_ENVELOPE_TTL_SECONDS + 82800, MAX_ENVELOPE_TTL_SECONDS), + (None, MAX_ENVELOPE_TTL_SECONDS), + ], +) +def test_exp_is_min_of_upstream_expires_in_and_cap(expires_in, expected_ttl): + grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=expires_in) + sealed = mint_envelope(_IDENTITY, grant, _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + assert sealed.expires_at == _NOW + timedelta(seconds=expected_ttl) + + +def test_expiry_honored_against_injected_clock(): + token = _sealed_token(_full_grant()) + assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=599)), OpenedEnvelope) + assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=600)), Expired) + assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=601)), Expired) + + +def test_ttl_cap_enforced_on_open_even_when_upstream_token_lives_longer(): + grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=86400) + token = _sealed_token(grant) + just_before_cap = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS - 1) + at_cap = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS) + assert isinstance(open_envelope(token, _KEYS, just_before_cap), OpenedEnvelope) + assert isinstance(open_envelope(token, _KEYS, at_cap), Expired) + + +def test_tampering_any_payload_or_signature_byte_is_bad_signature(): + token = _sealed_token(_full_grant()) + parts = token.removeprefix(ENVELOPE_PREFIX).split(".") + for segment in (1, 2): + for index in range(len(parts[segment])): + result = open_envelope(_tampered(token, segment, index), _KEYS, _NOW) + assert isinstance(result, BadSignature), f"segment {segment} index {index}: {result!r}" + + +def test_tampering_header_bytes_never_opens(): + token = _sealed_token(_full_grant()) + parts = token.removeprefix(ENVELOPE_PREFIX).split(".") + for index in range(len(parts[0])): + result = open_envelope(_tampered(token, 0, index), _KEYS, _NOW) + assert isinstance(result, (BadSignature, MalformedPayload)), f"header index {index}: {result!r}" + + +def test_alg_none_is_rejected(): + claims = _unverified_claims(_sealed_token(_full_grant())) + unsigned = ENVELOPE_PREFIX + jwt.encode(claims, None, algorithm="none") + assert isinstance(open_envelope(unsigned, _KEYS, _NOW), MalformedPayload) + + +def test_wrong_signing_key_is_bad_signature(): + token = _sealed_token(_full_grant()) + assert isinstance(open_envelope(token, _WRONG_SIGNING, _NOW), BadSignature) + + +def test_wrong_encryption_key_is_decrypt_failed(): + token = _sealed_token(_full_grant()) + assert isinstance(open_envelope(token, _WRONG_ENCRYPTION, _NOW), DecryptFailed) + + +def test_ciphertext_swapped_from_another_envelope_is_decrypt_failed(): + claims_a = _unverified_claims(_sealed_token(_full_grant(), keys=_KEYS)) + claims_b = _unverified_claims(_sealed_token(_minimal_grant(), keys=_WRONG_ENCRYPTION)) + swapped = _forge({**claims_a, "grant": claims_b["grant"]}) + assert isinstance(open_envelope(swapped, _KEYS, _NOW), DecryptFailed) + + +def test_wrong_issuer_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + assert isinstance(open_envelope(_forge({**claims, "iss": "evil-issuer"}), _KEYS, _NOW), MalformedPayload) + + +def test_missing_identity_claim_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({key: value for key, value in claims.items() if key != "user_id"}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + + +@pytest.mark.parametrize("identity_claim", ["user_id", "server_id"]) +def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, identity_claim: ""}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + intact = open_envelope(_forge(claims), _KEYS, _NOW) + assert isinstance(intact, OpenedEnvelope) + assert intact.identity == _IDENTITY + + +def test_lone_surrogate_candidate_is_malformed_payload_not_a_raise(): + surrogate_candidate = ENVELOPE_PREFIX + "\ud800abc.def.ghi" + result = open_envelope(surrogate_candidate, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize( + "override", + [ + {"iat": [1]}, + {"iat": {}}, + {"iat": float("inf")}, + {"nbf": None}, + {"nbf": [1]}, + ], +) +def test_hostile_iat_nbf_types_are_malformed_payload_not_a_raise(override): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, **override}) + result = open_envelope(forged, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize("hostile_iss", [["litellm-mcp-bridge"], 5, {"iss": "x"}]) +def test_non_string_issuer_claim_is_malformed_payload_not_a_raise(hostile_iss): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _hand_crafted_hs256({**claims, "iss": hostile_iss}) + result = open_envelope(forged, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize("hostile_exp", ["600", 600.5, [600]]) +def test_non_int_exp_claim_is_malformed_payload_not_a_raise(hostile_exp): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _hand_crafted_hs256({**claims, "exp": hostile_exp}) + result = open_envelope(forged, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +def test_unexpected_extra_claim_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, "role": "admin"}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + + +def test_future_iat_opens_against_injected_now_not_wall_clock(): + future = _NOW + timedelta(seconds=100_000) + sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, future) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, future) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity == _IDENTITY + assert opened.grant == _full_grant() + + +def test_rs256_signed_token_is_rejected_against_the_hs256_pin(): + claims = _unverified_claims(_sealed_token(_full_grant())) + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + rs256_token = ENVELOPE_PREFIX + jwt.encode(claims, private_key, algorithm="RS256") + result = open_envelope(rs256_token, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize("short_key", ["", "too-short", "x" * 31]) +def test_signing_key_below_hs256_minimum_is_rejected_at_construction(short_key): + with pytest.raises(ValidationError): + EnvelopeKeys(signing_key=SecretStr(short_key), encryption_key=SecretStr(_ENCRYPTION_KEY)) + + +def test_signing_key_at_hs256_minimum_is_accepted(): + keys = EnvelopeKeys(signing_key=SecretStr("y" * 32), encryption_key=SecretStr(_ENCRYPTION_KEY)) + assert keys.signing_key.get_secret_value() == "y" * 32 + + +def test_correctly_signed_garbage_grant_blob_is_decrypt_failed(): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, "grant": "not-a-ciphertext"}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), DecryptFailed) + + +def test_decryptable_blob_that_is_not_a_grant_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + wrong_shape = base64.urlsafe_b64encode( + bytes(encrypt_value(value=json.dumps({"nope": 1}), signing_key=_ENCRYPTION_KEY)) + ).decode("ascii") + forged = _forge({**claims, "grant": wrong_shape}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + + +def _mint_with_token_len(n: int) -> SealedEnvelope | EnvelopeTooLarge: + grant = UpstreamTokenGrant(access_token=SecretStr("a" * n), token_type="Bearer") + return mint_envelope(_IDENTITY, grant, _KEYS, _NOW) + + +def _largest_token_len_that_mints(lo: int, hi: int) -> int: + if hi - lo <= 1: + return lo + mid = (lo + hi) // 2 + if isinstance(_mint_with_token_len(mid), SealedEnvelope): + return _largest_token_len_that_mints(mid, hi) + return _largest_token_len_that_mints(lo, mid) + + +def test_oversized_grant_is_a_typed_mint_error_never_truncated(): + result = _mint_with_token_len(30000) + assert isinstance(result, EnvelopeTooLarge) + assert result.tag == "envelope_too_large" + assert result.size_bytes > MAX_ENVELOPE_BYTES + assert result.max_bytes == MAX_ENVELOPE_BYTES + + +def test_size_cap_boundary_just_under_succeeds_and_just_over_fails(): + assert isinstance(_mint_with_token_len(1), SealedEnvelope) + assert isinstance(_mint_with_token_len(30000), EnvelopeTooLarge) + largest = _largest_token_len_that_mints(1, 30000) + assert largest > 6000 + sealed = _mint_with_token_len(largest) + assert isinstance(sealed, SealedEnvelope) + assert len(sealed.token.get_secret_value().encode("utf-8")) <= MAX_ENVELOPE_BYTES + overflowing = _mint_with_token_len(largest + 1) + assert isinstance(overflowing, EnvelopeTooLarge) + assert overflowing.size_bytes > MAX_ENVELOPE_BYTES + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + + +def test_open_size_guard_measures_bytes_not_characters(): + """The open-side size guard must reject on UTF-8 byte length, matching mint's cap, so a + hostile multi-byte candidate whose character count is under the cap but whose byte count is + over it is rejected up front rather than reaching the expensive HMAC/decrypt path. Patching + _decode_claims to fail loudly proves the guard short-circuits before decode.""" + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.outbound_credentials import envelope + + multibyte_body = "é" * 7000 # 7000 chars, 14000 UTF-8 bytes + candidate = ENVELOPE_PREFIX + multibyte_body + assert len(candidate) <= MAX_ENVELOPE_BYTES + assert len(candidate.encode("utf-8")) > MAX_ENVELOPE_BYTES + + with patch.object(envelope, "_decode_claims", side_effect=AssertionError("decode reached")) as decode: + result = open_envelope(candidate, _KEYS, _NOW) + + assert isinstance(result, MalformedPayload) + decode.assert_not_called() + + +def test_open_size_guard_rejects_oversize_character_count_before_decode(): + """A candidate whose character count already exceeds the cap is rejected up front, before the + decode path, so an arbitrarily long hostile string is not run through HMAC/decrypt. The cheap + character precheck makes this O(1) since UTF-8 byte length is never below character length.""" + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.outbound_credentials import envelope + + candidate = ENVELOPE_PREFIX + ("a" * (MAX_ENVELOPE_BYTES + 1)) + assert len(candidate) > MAX_ENVELOPE_BYTES + + with patch.object(envelope, "_decode_claims", side_effect=AssertionError("decode reached")) as decode: + result = open_envelope(candidate, _KEYS, _NOW) + + assert isinstance(result, MalformedPayload) + decode.assert_not_called() + + +def test_is_envelope_detects_only_prefixed_values(): + assert is_envelope(_sealed_token(_full_grant())) + raw_jwt = jwt.encode({"sub": "user-123"}, _SIGNING_KEY, algorithm="HS256") + assert not is_envelope(raw_jwt) + assert not is_envelope("some-random-opaque-token") + assert not is_envelope("") + + +def test_open_on_non_envelope_input_is_not_an_envelope(): + raw_jwt = jwt.encode({"sub": "user-123"}, _SIGNING_KEY, algorithm="HS256") + assert isinstance(open_envelope(raw_jwt, _KEYS, _NOW), NotAnEnvelope) + assert isinstance(open_envelope("", _KEYS, _NOW), NotAnEnvelope) + assert isinstance(open_envelope(_ACCESS_TOKEN, _KEYS, _NOW), NotAnEnvelope) + + +def test_open_on_prefixed_garbage_is_malformed_payload(): + assert isinstance(open_envelope(ENVELOPE_PREFIX + "garbage", _KEYS, _NOW), MalformedPayload) + assert isinstance(open_envelope(ENVELOPE_PREFIX + _ACCESS_TOKEN, _KEYS, _NOW), MalformedPayload) + + +def test_no_result_value_ever_reveals_the_access_token(): + grant = _full_grant() + sealed = mint_envelope(_IDENTITY, grant, _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + token = sealed.token.get_secret_value() + oversized_grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN + "x" * 30000), token_type="Bearer") + values = ( + sealed, + open_envelope(token, _KEYS, _NOW), + mint_envelope(_IDENTITY, oversized_grant, _KEYS, _NOW), + open_envelope(_ACCESS_TOKEN, _KEYS, _NOW), + open_envelope(ENVELOPE_PREFIX + _ACCESS_TOKEN, _KEYS, _NOW), + open_envelope(token, _WRONG_SIGNING, _NOW), + open_envelope(token, _WRONG_ENCRYPTION, _NOW), + open_envelope(token, _KEYS, _NOW + timedelta(seconds=601)), + grant, + ) + for value in values: + assert _ACCESS_TOKEN not in repr(value) + assert _ACCESS_TOKEN not in str(value) + assert _REFRESH_TOKEN not in repr(value) + assert _REFRESH_TOKEN not in str(value) + + +def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): + for bad_expires_in in (0, -5): + with pytest.raises(ValidationError) as excinfo: + UpstreamTokenGrant( + access_token=SecretStr(_ACCESS_TOKEN), + token_type="Bearer", + expires_in=bad_expires_in, + ) + assert _ACCESS_TOKEN not in str(excinfo.value) + assert _ACCESS_TOKEN not in repr(excinfo.value) + + +def test_empty_identity_and_key_fields_are_rejected_at_construction(): + with pytest.raises(ValidationError): + EnvelopeIdentity(user_id="", server_id="srv-456") + with pytest.raises(ValidationError): + EnvelopeIdentity(user_id="user-123", server_id="") + with pytest.raises(ValidationError): + EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) + with pytest.raises(ValidationError): + EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr("")) + with pytest.raises(ValidationError): + UpstreamTokenGrant(access_token=SecretStr(""), token_type="Bearer") + + +def test_public_models_are_frozen(): + sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + with pytest.raises(ValidationError): + sealed.token = SecretStr("overwritten") + with pytest.raises(ValidationError): + opened.grant = _minimal_grant() + with pytest.raises(ValidationError): + _IDENTITY.user_id = "someone-else" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index 00d96f975bb..41d61c8f508 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -11,6 +11,7 @@ import json from unittest.mock import AsyncMock, MagicMock import pytest +from prisma import Json from litellm.proxy._experimental.mcp_server.db import ( create_mcp_server, @@ -19,6 +20,11 @@ from litellm.proxy._experimental.mcp_server.db import ( from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest +def _credentials_cleared(value) -> bool: + """The clear sentinel after the edge translation: prisma Json(None) (SQL null) or a bare None.""" + return value is None or (isinstance(value, Json) and getattr(value, "data", "x") is None) + + def _mock_prisma(): mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable = AsyncMock() @@ -208,7 +214,7 @@ async def test_auth_type_switch_clears_stale_flow_scoped_fields(): "token_exchange_profile", ): assert data_dict[stale_field] is None, f"{stale_field} must be cleared on auth_type switch" - assert data_dict["credentials"] is None + assert _credentials_cleared(data_dict["credentials"]) @pytest.mark.asyncio @@ -558,3 +564,103 @@ async def test_te_update_without_blob_te_keys_leaves_credentials_untouched(): assert data_dict["token_exchange_endpoint"] == "https://new.example.com/token" assert "credentials" not in data_dict + + +# ── client-forwarded credential class: true_passthrough <-> oauth_delegate share one +# stored-app shape, so a switch between them must MERGE (keep the declared app), not REPLACE ── + + +@pytest.mark.asyncio +async def test_cf_pair_switch_without_credentials_keeps_stored_app_and_endpoints(): + """true_passthrough -> oauth_delegate with no credentials in the update must not clear the + stored client or null the endpoint columns: both modes use the same declared app and relay.""" + mock_prisma = _mock_prisma() + existing = _existing_row("true_passthrough", credentials={"client_id": "enc-A", "client_secret": "enc-B"}) + existing.authorization_url = "https://provider.example/authorize" + existing.token_url = "https://provider.example/token" + existing.registration_url = "https://provider.example/register" + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="cf-server", auth_type="oauth_delegate") + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert "credentials" not in data_dict + for scoped_field in ("authorization_url", "token_url", "registration_url", "oauth2_flow"): + assert scoped_field not in data_dict, f"{scoped_field} must not be nulled within the CF class" + + +@pytest.mark.asyncio +async def test_cf_pair_switch_with_partial_credentials_merges_not_replaces(): + """oauth_delegate update carrying only client_id onto a true_passthrough row must MERGE, so the + stored client_secret survives instead of being dropped by a REPLACE.""" + mock_prisma = _mock_prisma() + existing = _existing_row("true_passthrough", credentials={"client_id": "enc-A", "client_secret": "enc-B"}) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="cf-server", auth_type="oauth_delegate", credentials={"client_id": "B"}) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + merged = json.loads(data_dict["credentials"]) + assert merged["client_secret"] == "enc-B" + assert merged["client_id"] != "enc-A" + + +@pytest.mark.asyncio +async def test_null_existing_auth_type_to_cf_counts_as_changed_and_clears_blob(): + """A legacy row with NULL auth_type switched to a client-forwarded mode is a cross-class change, + so the stale blob must be cleared (the two change predicates must agree on this).""" + mock_prisma = _mock_prisma() + existing = _existing_row(None, credentials={"client_id": "enc-old"}) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="cf-server", auth_type="true_passthrough") + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + # The clear must reach prisma as Json(None) (SQL null), never a raw None, which prisma rejects. + assert isinstance(data_dict["credentials"], Json) + assert getattr(data_dict["credentials"], "data", "x") is None + + +@pytest.mark.asyncio +async def test_client_rotation_strips_legacy_minted_token_keys(): + """Rotating the client on a same-class row must drop stale minted token material the update did + not set, so an old access_token/refresh_token never rides forward under the new client.""" + mock_prisma = _mock_prisma() + existing = _existing_row( + "oauth2", credentials={"client_id": "A", "access_token": "T", "refresh_token": "R", "expires_in": 3600} + ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest( + server_id="oauth2-server", auth_type="oauth2", credentials={"client_id": "B", "client_secret": "S"} + ) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + merged = json.loads(data_dict["credentials"]) + assert "access_token" not in merged + assert "refresh_token" not in merged + assert "expires_in" not in merged + assert "client_secret" in merged + + +@pytest.mark.asyncio +async def test_cf_to_non_cf_switch_clears_dcr_bridge(): + """A cross-class switch OUT of a client-forwarded mode (true_passthrough -> api_key) must clear + dcr_bridge: the switch is cross-class so the flow-scoped sweep runs and nulls it, leaving no stale + dcr_bridge=True on a row that no longer supports it.""" + data = UpdateMCPServerRequest(server_id="s", auth_type="api_key") + data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") + assert data_dict["dcr_bridge"] is None + + +@pytest.mark.asyncio +async def test_cf_pair_switch_does_not_clear_dcr_bridge(): + """A within-class switch (true_passthrough <-> oauth_delegate) is not a credential-class change, so + the flow-scoped sweep does not run and dcr_bridge is left intact (both modes use the DCR bridge).""" + data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate") + data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") + assert "dcr_bridge" not in data_dict diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index f5a48b5ec34..5992fd1814f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -399,9 +399,7 @@ class TestMCPServerManagerSigV4: server = next(iter(manager.config_mcp_servers.values())) assert server.auth_type == MCPAuth.aws_sigv4 assert server.aws_access_key_id == "AKIAIOSFODNN7EXAMPLE" - assert ( - server.aws_secret_access_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" - ) + assert server.aws_secret_access_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" assert server.aws_region_name == "us-east-1" assert server.aws_service_name == "bedrock-agentcore" @@ -531,9 +529,7 @@ class TestMCPServerManagerSigV4: "aws_session_name": "my-session", } - result = manager._extract_aws_credentials( - creds, credentials_are_encrypted=False - ) + result = manager._extract_aws_credentials(creds, credentials_are_encrypted=False) assert result["aws_role_name"] == "arn:aws:iam::123456789012:role/TestRole" assert result["aws_session_name"] == "my-session" @@ -561,10 +557,7 @@ class TestSigV4CredentialEncryption: # Secrets should be encrypted assert result["aws_access_key_id"] == "enc:AKIAIOSFODNN7EXAMPLE" - assert ( - result["aws_secret_access_key"] - == "enc:wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" - ) + assert result["aws_secret_access_key"] == "enc:wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" assert result["aws_session_token"] == "enc:FwoGZX..." # Non-secrets should be unchanged assert result["aws_region_name"] == "us-east-1" @@ -606,12 +599,8 @@ class TestCredentialMergeOnUpdate: ) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -650,9 +639,7 @@ class TestCredentialMergeOnUpdate: from litellm.proxy._types import UpdateMCPServerRequest mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -679,12 +666,8 @@ class TestCredentialMergeOnUpdate: existing_record.credentials = None mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -725,12 +708,8 @@ class TestCredentialMergeOnUpdate: ) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -772,12 +751,8 @@ class TestCredentialMergeOnUpdate: ) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -819,9 +794,7 @@ class TestSigV4BuildFromTable: table_record.server_name = "sigv4_server" table_record.alias = None table_record.description = None - table_record.url = ( - "https://bedrock-agentcore.us-east-1.amazonaws.com/invocations" - ) + table_record.url = "https://bedrock-agentcore.us-east-1.amazonaws.com/invocations" table_record.spec_path = None table_record.transport = "http" table_record.auth_type = "aws_sigv4" @@ -867,9 +840,7 @@ class TestSigV4BuildFromTable: with patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.decrypt_value_helper", - side_effect=lambda value, key, exception_type, return_original_value: value.replace( - "enc:", "" - ), + side_effect=lambda value, key, exception_type, return_original_value: value.replace("enc:", ""), ): server = await manager.build_mcp_server_from_table(table_record) @@ -930,9 +901,7 @@ class TestSigV4BuildFromTable: with patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.decrypt_value_helper", - side_effect=lambda value, key, exception_type, return_original_value: value.replace( - "enc:", "" - ), + side_effect=lambda value, key, exception_type, return_original_value: value.replace("enc:", ""), ): server = await manager.build_mcp_server_from_table(table_record) @@ -1018,9 +987,7 @@ class TestRotateCredentials: server.env_vars = None mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( - return_value=[server] - ) + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() with ( @@ -1039,9 +1006,7 @@ class TestRotateCredentials: side_effect=lambda value, new_encryption_key: f"enc_new:{value}", ), ): - await rotate_mcp_server_credentials_master_key( - mock_prisma, "admin", "new-key" - ) + await rotate_mcp_server_credentials_master_key(mock_prisma, "admin", "new-key") update_call = mock_prisma.db.litellm_mcpservertable.update assert update_call.called @@ -1069,9 +1034,7 @@ class TestRotateCredentials: ] mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( - return_value=[server] - ) + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() with ( @@ -1090,9 +1053,7 @@ class TestRotateCredentials: side_effect=lambda value, new_encryption_key: f"enc_new:{value}", ), ): - await rotate_mcp_server_credentials_master_key( - mock_prisma, "admin", "new-key" - ) + await rotate_mcp_server_credentials_master_key(mock_prisma, "admin", "new-key") update_call = mock_prisma.db.litellm_mcpservertable.update assert update_call.called @@ -1116,17 +1077,11 @@ class TestAuthTypeSwitchClearsCredentials: existing_record = MagicMock() existing_record.auth_type = "oauth2" - existing_record.credentials = json.dumps( - {"client_id": "enc:cid", "client_secret": "enc:csec"} - ) + existing_record.credentials = json.dumps({"client_id": "enc:cid", "client_secret": "enc:csec"}) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -1141,8 +1096,12 @@ class TestAuthTypeSwitchClearsCredentials: await update_mcp_server(mock_prisma, data, "test-user") data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] - # Credentials should be cleared (set to None) - assert data_dict.get("credentials") is None + # Credentials should be cleared. The clear reaches prisma as Json(None) (SQL null), which + # prisma-python requires for a Json? field; a bare None is also accepted for older callers. + from prisma import Json + + cleared = data_dict.get("credentials") + assert cleared is None or (isinstance(cleared, Json) and getattr(cleared, "data", "x") is None) class TestInheritCredentials: diff --git a/tests/test_litellm/proxy/client/cli/test_global_options.py b/tests/test_litellm/proxy/client/cli/test_global_options.py index 3a19f735c1b..53b7e4dbc29 100644 --- a/tests/test_litellm/proxy/client/cli/test_global_options.py +++ b/tests/test_litellm/proxy/client/cli/test_global_options.py @@ -1,7 +1,7 @@ # stdlib imports import os import sys -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest from click.testing import CliRunner @@ -36,6 +36,35 @@ def test_cli_version_flag(cli_runner): assert "LiteLLM Proxy Server Version: 1.2.3" in result.output +def test_base_url_trailing_slash_normalized(cli_runner): + """A trailing slash on --base-url must not produce a double slash (e.g. '//sso/cli/start').""" + with ( + patch("webbrowser.open"), + patch( + "requests.post", + return_value=Mock( + status_code=200, + json=Mock( + return_value={ + "login_id": "cli-test-uuid", + "poll_secret": "poll-secret", + "user_code": "ABCD-EFGH", + } + ), + raise_for_status=Mock(), + ), + ) as mock_post, + patch("requests.get", side_effect=ValueError("stop after start request")), + ): + cli_runner.invoke( + cli, ["--base-url", "https://gateway.litellm-sandbox.ai/", "login"] + ) + + mock_post.assert_called_once_with( + "https://gateway.litellm-sandbox.ai/sso/cli/start", timeout=10 + ) + + def test_cli_version_command(cli_runner): """Test that 'version' command prints the correct version, server URL, and server version, and exits successfully""" with ( diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index 2c41b16ba7f..33e45ccb22c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -7,9 +7,7 @@ import sys import pytest -sys.path.insert( - 0, os.path.abspath("../../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path from litellm.proxy.compliance_checks import ComplianceChecker from litellm.types.proxy.compliance_endpoints import ComplianceCheckRequest @@ -385,3 +383,214 @@ class TestGdprCompliant: ) checks = ComplianceChecker(data).check_gdpr() assert all(c.passed for c in checks) + + +class TestModeMatching: + """Direct coverage of ComplianceChecker._mode_matches for every shape. + + LitellmParams.mode is Union[str, List[str], Mode], so a spend-log + guardrail_mode can be None / str / list / tuple / dict. A prior + implementation compared `g_mode == mode`, which silently failed for the + non-str shapes and reported NON-COMPLIANT for every multi-mode guardrail. + A match now reports a mode satisfied only when every configured branch + runs in that mode: fails safe, no false-COMPLIANT. + """ + + @pytest.mark.parametrize( + "g_mode, mode, expected", + [ + (None, "pre_call", True), + (None, "post_call", False), + (None, "during_call", False), + ("pre_call", "pre_call", True), + ("post_call", "pre_call", False), + ("during_call", "during_call", True), + # list/tuple: only guaranteed when every listed mode equals `mode` + (["pre_call"], "pre_call", True), + (["pre_call", "pre_call"], "pre_call", True), + (["pre_call", "post_call"], "pre_call", False), + (["pre_call", "post_call"], "post_call", False), + ([], "pre_call", False), + (("during_call",), "during_call", True), + (("pre_call", "post_call"), "pre_call", False), + # dict: default only + ({"default": "pre_call"}, "pre_call", True), + ({"default": "pre_call"}, "post_call", False), + ({"default": ["pre_call", "post_call"]}, "post_call", False), + ({"default": ["pre_call"]}, "pre_call", True), + # dict with tags: every branch must run in mode + ({"default": "pre_call", "tags": {"a": "pre_call"}}, "pre_call", True), + ({"default": "pre_call", "tags": {"a": ["pre_call"]}}, "pre_call", True), + ({"default": "pre_call", "tags": {"a": ["pre_call", "post_call"]}}, "pre_call", False), + ({"default": "pre_call", "tags": {"eu": "post_call"}}, "pre_call", False), + ({"default": "pre_call", "tags": {"eu": "post_call"}}, "post_call", False), + ({"default": "pre_call", "tags": {"eu": ["during_call"]}}, "during_call", False), + ({"default": ["pre_call", "post_call"], "tags": {"a": "pre_call"}}, "post_call", False), + # Missing default: untagged routing is unknown, nothing guaranteed + ({"tags": {"x": "post_call"}}, "pre_call", False), + ({"tags": {"x": "post_call"}}, "post_call", False), + ({}, "pre_call", False), + ({}, "post_call", False), + ({"default": 123}, "pre_call", False), + # Unknown top-level shapes never match + (5, "pre_call", False), + (object(), "pre_call", False), + ], + ) + def test_mode_matches(self, g_mode, mode, expected): + assert ComplianceChecker._mode_matches(g_mode, mode) is expected + + def test_list_mode_guardrail_not_misclassified(self): + """A guardrail configured with mode ["pre_call", "post_call"] is logged + with the raw list when the writer cannot infer the concrete hook that + ran (e.g. apply_guardrail invocations). The spend log records "this + guardrail could have run at either hook", not "which hook fired this + request". Counting it for both would let a request that only fired + post_call pass a pre_call compliance check. It counts for neither.""" + data = ComplianceCheckRequest( + request_id="req-mode-1", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T00:00:00Z", + guardrail_information=[ + { + "guardrail_name": "pii_masking", + "guardrail_mode": ["pre_call", "post_call"], + "guardrail_status": "success", + } + ], + ) + checker = ComplianceChecker(data) + assert len(checker._get_guardrails_by_mode("pre_call")) == 0 + assert len(checker._get_guardrails_by_mode("post_call")) == 0 + results = {c.check_name: c.passed for c in checker.check_eu_ai_act()} + assert results["Content screened before LLM"] is False + + def test_list_mode_single_value_counts(self): + """A single-entry list ["pre_call"] runs pre_call unconditionally, so it + counts for pre_call and no other mode.""" + data = ComplianceCheckRequest( + request_id="req-mode-1b", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T00:00:00Z", + guardrail_information=[ + { + "guardrail_name": "pii_masking", + "guardrail_mode": ["pre_call"], + "guardrail_status": "success", + } + ], + ) + checker = ComplianceChecker(data) + assert len(checker._get_guardrails_by_mode("pre_call")) == 1 + assert len(checker._get_guardrails_by_mode("post_call")) == 0 + + def test_dict_tag_routed_guardrail_not_misclassified(self): + """A tag-routed guardrail (default=pre_call, a post_call tag) is not + guaranteed to run in either mode, so it counts for neither.""" + data = ComplianceCheckRequest( + request_id="req-mode-2", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T12:00:00Z", + guardrail_information=[ + { + "guardrail_name": "pii_masking", + "guardrail_mode": {"default": "pre_call", "tags": {"eu": "post_call"}}, + "guardrail_status": "success", + } + ], + ) + checker = ComplianceChecker(data) + assert len(checker._get_guardrails_by_mode("pre_call")) == 0 + assert len(checker._get_guardrails_by_mode("post_call")) == 0 + results = {c.check_name: c.passed for c in checker.check_eu_ai_act()} + assert results["Content screened before LLM"] is False + + def test_dict_all_branches_pre_call_counts(self): + """When default and every tag override all run pre_call, the guardrail is + guaranteed pre_call regardless of routing, so it counts for pre_call.""" + data = ComplianceCheckRequest( + request_id="req-mode-4", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T12:00:00Z", + guardrail_information=[ + { + "guardrail_name": "pii_masking", + "guardrail_mode": {"default": "pre_call", "tags": {"eu": "pre_call"}}, + "guardrail_status": "success", + } + ], + ) + checker = ComplianceChecker(data) + assert len(checker._get_guardrails_by_mode("pre_call")) == 1 + + def test_none_mode_defaults_to_pre_call(self): + """A guardrail logged without a mode counts as pre_call only.""" + data = ComplianceCheckRequest( + request_id="req-mode-3", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T12:00:00Z", + guardrail_information=[{"guardrail_name": "pii_masking", "guardrail_status": "success"}], + ) + checker = ComplianceChecker(data) + assert len(checker._get_guardrails_by_mode("pre_call")) == 1 + assert len(checker._get_guardrails_by_mode("post_call")) == 0 + + def test_never_reports_false_compliant(self): + """The core invariant: a match reports `mode` satisfied only when every + configured branch runs in that mode. So True can never claim a hook the + guardrail may not have actually executed. The only allowed error + direction is under-reporting.""" + + def _branch_modes(value): + if isinstance(value, str): + return {value} + if isinstance(value, (list, tuple)): + return {v for v in value if isinstance(v, str)} + return set() + + def _guaranteed_modes(g_mode): + """Modes every branch of ``g_mode`` runs in.""" + if isinstance(g_mode, str): + return {g_mode} + if isinstance(g_mode, (list, tuple)): + sets = [_branch_modes(m) for m in g_mode] + return set.intersection(*sets) if sets else set() + if isinstance(g_mode, dict): + default = g_mode.get("default") + if default is None: + return set() + branches = [default, *(g_mode.get("tags") or {}).values()] + sets = [_branch_modes(b) for b in branches] + return set.intersection(*sets) if sets else set() + return set() + + shapes = [ + None, + "pre_call", + "post_call", + ["pre_call"], + ["pre_call", "post_call"], + [], + {"default": "pre_call"}, + {"default": ["pre_call", "post_call"]}, + {"default": "pre_call", "tags": {"a": "pre_call"}}, + {"default": "pre_call", "tags": {"a": "post_call"}}, + {"default": ["pre_call", "post_call"], "tags": {"a": "pre_call"}}, + {"tags": {"a": "post_call"}}, + {}, + {"default": 123}, + 5, + ] + for g_mode in shapes: + for mode in ("pre_call", "post_call", "during_call"): + matched = ComplianceChecker._mode_matches(g_mode, mode) + if g_mode is None: + assert matched is (mode == "pre_call"), (g_mode, mode) + continue + if matched: + assert mode in _guaranteed_modes(g_mode), (g_mode, mode) diff --git a/tests/test_litellm/proxy/management_endpoints/test_coordination_redis_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_coordination_redis_endpoints.py new file mode 100644 index 00000000000..4e6bfc4c063 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_coordination_redis_endpoints.py @@ -0,0 +1,577 @@ +""" +Unit tests for coordination Redis settings management endpoints +""" + +import asyncio +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path + +import litellm +from litellm.caching.caching import RedisCache +from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.proxy._types import LitellmTableNames, LitellmUserRoles +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( + _REDACTED_VALUE, + CoordinationRedisSettingsRequest, + get_coordination_redis_settings, + check_coordination_redis_connection, + update_coordination_redis_settings, +) +from litellm.types.management_endpoints.coordination_redis_endpoints import ( + COORDINATION_REDIS_SETTINGS_FIELDS, +) + +_SAVED_SETTINGS = { + "host": "coord-redis.example.com", + "port": 6379, + "password": "super-secret-redis-pw", + "url": "redis://:super-secret-redis-pw@coord-redis.example.com:6379", + "sentinel_password": "super-secret-sentinel-pw", +} + + +def _admin_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed", + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + +def _prisma_with_general_settings(general_settings: dict | None) -> MagicMock: + """A prisma client whose LiteLLM_Config `general_settings` row holds ``general_settings``.""" + row = None + if general_settings is not None: + row = MagicMock() + row.param_value = json.dumps(general_settings) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=row) + mock_prisma.db.litellm_config.upsert = AsyncMock() + return mock_prisma + + +def _proxy_config(file_general_settings: dict | None = None) -> MagicMock: + proxy_config = MagicMock() + proxy_config.get_config_state = MagicMock( + return_value={"general_settings": file_general_settings or {}}, + ) + return proxy_config + + +# ── GET /coordination_redis/settings ────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_get_redacts_every_credential_field(): + """password, sentinel_password and the (password-bearing) url never leave the + server in plaintext; non-credential fields come back untouched.""" + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + _prisma_with_general_settings({"coordination_redis": _SAVED_SETTINGS}), + ), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + serialized = json.dumps(response.model_dump()) + assert "super-secret-redis-pw" not in serialized + assert "super-secret-sentinel-pw" not in serialized + + assert response.values["password"] == _REDACTED_VALUE + assert response.values["sentinel_password"] == _REDACTED_VALUE + assert response.values["url"] == _REDACTED_VALUE + assert response.values["host"] == "coord-redis.example.com" + assert response.values["port"] == 6379 + + # field metadata is hydrated with the same redacted values + by_name = {field.field_name: field for field in response.fields} + assert by_name["password"].field_value == _REDACTED_VALUE + assert by_name["host"].field_value == "coord-redis.example.com" + + +@pytest.mark.asyncio +async def test_get_source_is_coordination_redis_when_block_present(monkeypatch): + """An explicit block wins even when a Redis cache backend and REDIS_* env both exist.""" + monkeypatch.setattr(litellm, "cache", MagicMock(cache=MagicMock(spec=RedisCache))) + monkeypatch.setenv("REDIS_HOST", "env-redis") + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + _prisma_with_general_settings({"coordination_redis": _SAVED_SETTINGS}), + ), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + assert response.source == "coordination_redis" + + +@pytest.mark.asyncio +async def test_get_source_reads_block_from_yaml_config_when_db_row_absent(monkeypatch): + """A block set in config.yaml (not the DB) still reports source=coordination_redis.""" + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.delenv("REDIS_HOST", raising=False) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings(None)), + patch( + "litellm.proxy.proxy_server.proxy_config", + _proxy_config({"coordination_redis": {"host": "yaml-redis"}}), + ), + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + assert response.source == "coordination_redis" + assert response.values["host"] == "yaml-redis" + + +@pytest.mark.parametrize("cache_backend_cls", [RedisCache, RedisClusterCache]) +@pytest.mark.asyncio +async def test_get_source_is_cache_backend_when_no_block(monkeypatch, cache_backend_cls): + """With no explicit block, a plain-Redis response-cache backend is borrowed — + which beats the REDIS_* env fallback.""" + monkeypatch.setattr(litellm, "cache", MagicMock(cache=MagicMock(spec=cache_backend_cls))) + monkeypatch.setenv("REDIS_HOST", "env-redis") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + assert response.source == "cache_backend" + assert response.values == {} + + +@pytest.mark.asyncio +async def test_get_source_is_environment_when_no_block_and_non_redis_cache(monkeypatch): + """A non-Redis cache backend falls through to the REDIS_* env fallback.""" + monkeypatch.setattr(litellm, "cache", MagicMock(cache=MagicMock())) + monkeypatch.setenv("REDIS_HOST", "env-redis") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + assert response.source == "environment" + + +@pytest.mark.asyncio +async def test_get_source_is_none_when_nothing_configured(monkeypatch): + monkeypatch.setattr(litellm, "cache", None) + for env_var in ("REDIS_HOST", "REDIS_URL", "REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES"): + monkeypatch.delenv(env_var, raising=False) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + assert response.source is None + + +@pytest.mark.asyncio +async def test_get_source_does_not_build_a_client(monkeypatch): + """The env-fallback probe is read-only: no Redis client is constructed on GET.""" + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.setenv("REDIS_HOST", "env-redis") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server._build_redis_usage_cache") as mock_build, + ): + response = await get_coordination_redis_settings(user_api_key_dict=_admin_auth()) + + assert response.source == "environment" + mock_build.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_rejects_non_admin(): + with pytest.raises(HTTPException) as exc_info: + await get_coordination_redis_settings( + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_role=LitellmUserRoles.INTERNAL_USER) + ) + assert exc_info.value.status_code == 403 + + +def test_fields_cover_every_coordination_redis_param(): + """The declarative field list drives the Admin UI form; it must stay in sync + with the model the backend validates against.""" + from litellm.proxy._types import CoordinationRedisParams + + assert {field.field_name for field in COORDINATION_REDIS_SETTINGS_FIELDS} == set( + CoordinationRedisParams.model_fields.keys() + ) + + by_name = {field.field_name: field for field in COORDINATION_REDIS_SETTINGS_FIELDS} + assert by_name["startup_nodes"].section == "cluster" + assert by_name["sentinel_nodes"].section == "sentinel" + assert by_name["host"].section == "connection" + + +# ── POST /coordination_redis/settings ───────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_update_rejects_settings_without_a_connection_target(monkeypatch): + """A block with no host/url/startup_nodes/sentinel_nodes would blow up at + startup; reject it at write time and persist nothing.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + mock_prisma = _prisma_with_general_settings({}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + with pytest.raises(HTTPException) as exc_info: + await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest(settings={"ssl": True, "service_name": "mymaster"}), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 400 + mock_prisma.db.litellm_config.upsert.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_persists_into_the_general_settings_config_row(monkeypatch): + """Settings land under `general_settings.coordination_redis` in LiteLLM_Config + (the row startup merges over the yaml config), and sibling general_settings + keys survive the write.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + mock_prisma = _prisma_with_general_settings({"master_key": "sk-1234"}) + invalidated: list[str] = [] + + async def _capture_invalidate(param_name: str) -> None: + invalidated.append(param_name) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch( + "litellm.proxy.management_endpoints.coordination_redis_endpoints.invalidate_config_param", + new=_capture_invalidate, + ), + ): + response = await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest( + settings={"host": "coord-redis.example.com", "port": 6379, "password": "pw"} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + upsert_kwargs = mock_prisma.db.litellm_config.upsert.call_args.kwargs + assert upsert_kwargs["where"] == {"param_name": "general_settings"} + persisted = json.loads(upsert_kwargs["data"]["update"]["param_value"]) + assert persisted["coordination_redis"] == { + "host": "coord-redis.example.com", + "port": 6379, + "password": "pw", + } + assert persisted["master_key"] == "sk-1234" + assert invalidated == ["general_settings"] + + # the response echoes the saved settings back redacted + assert response["settings"]["password"] == _REDACTED_VALUE + assert response["settings"]["host"] == "coord-redis.example.com" + + +@pytest.mark.asyncio +async def test_update_persists_os_environ_refs_verbatim(monkeypatch): + """`os.environ/VAR` refs are resolved only to validate; the ref itself is what + gets stored, so the credential never lands in the DB.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setenv("MY_REDIS_HOST", "resolved-host") + mock_prisma = _prisma_with_general_settings({}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch( + "litellm.proxy.management_endpoints.coordination_redis_endpoints.invalidate_config_param", + new=AsyncMock(), + ), + ): + await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest(settings={"host": "os.environ/MY_REDIS_HOST"}), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = json.loads(mock_prisma.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"]) + assert persisted["coordination_redis"] == {"host": "os.environ/MY_REDIS_HOST"} + + +@pytest.mark.asyncio +async def test_update_keeps_saved_credential_when_client_echoes_the_redaction_marker(monkeypatch): + """The UI reads settings back redacted; re-submitting them must not persist + `***REDACTED***` as the password.""" + monkeypatch.setattr(litellm, "store_audit_logs", False) + mock_prisma = _prisma_with_general_settings({"coordination_redis": _SAVED_SETTINGS}) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch( + "litellm.proxy.management_endpoints.coordination_redis_endpoints.invalidate_config_param", + new=AsyncMock(), + ), + ): + await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest( + settings={"host": "new-host", "port": 6380, "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + persisted = json.loads(mock_prisma.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"]) + assert persisted["coordination_redis"]["password"] == "super-secret-redis-pw" + assert persisted["coordination_redis"]["host"] == "new-host" + + +@pytest.mark.asyncio +async def test_update_emits_audit_log_with_values_redacted(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", True) + mock_prisma = _prisma_with_general_settings({}) + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.coordination_redis_endpoints.invalidate_config_param", + new=AsyncMock(), + ), + patch("litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", new=capture), + ): + await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest( + settings={"host": "coord-redis.example.com", "password": "super-secret-redis-pw"} + ), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + for _ in range(3): + await asyncio.sleep(0) + + assert len(audit_calls) == 1 + log = audit_calls[0] + assert log.table_name == LitellmTableNames.CONFIG_TABLE_NAME + assert log.object_id == "coordination_redis" + assert log.action == "created" # no prior block → create + + after = json.loads(log.updated_values) + assert set(after["settings"].keys()) == {"host", "password"} + assert "super-secret-redis-pw" not in log.updated_values + assert "coord-redis.example.com" not in log.updated_values + + +@pytest.mark.asyncio +async def test_update_audit_action_is_updated_when_a_block_already_exists(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", True) + mock_prisma = _prisma_with_general_settings({"coordination_redis": {"host": "old-host"}}) + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.coordination_redis_endpoints.invalidate_config_param", + new=AsyncMock(), + ), + patch("litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", new=capture), + ): + await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest(settings={"host": "new-host"}), + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + for _ in range(3): + await asyncio.sleep(0) + + assert audit_calls[0].action == "updated" + assert json.loads(audit_calls[0].before_value)["settings"] == {"host": _REDACTED_VALUE} + + +@pytest.mark.asyncio +async def test_update_rejects_non_admin(): + with pytest.raises(HTTPException) as exc_info: + await update_coordination_redis_settings( + request=CoordinationRedisSettingsRequest(settings={"host": "coord-redis.example.com"}), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_role=LitellmUserRoles.INTERNAL_USER), + litellm_changed_by=None, + ) + assert exc_info.value.status_code == 403 + + +# ── POST /coordination_redis/settings/test ──────────────────────────────────── + + +@pytest.mark.asyncio +async def test_connection_test_returns_healthy_on_successful_ping(): + mock_client = MagicMock() + mock_client.ping = AsyncMock(return_value=True) + mock_client.disconnect = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server._build_redis_usage_cache", return_value=mock_client) as mock_build, + ): + response = await check_coordination_redis_connection( + request=CoordinationRedisSettingsRequest( + settings={"host": "coord-redis.example.com", "port": 6379, "password": "pw"} + ), + user_api_key_dict=_admin_auth(), + ) + + assert response.status == "healthy" + assert response.error is None + assert mock_build.call_args.args[0] == {"host": "coord-redis.example.com", "port": 6379, "password": "pw"} + mock_client.ping.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_connection_test_reports_unhealthy_without_leaking_the_password(): + """Redis client errors echo the connection url back; the password must be + scrubbed out of the error the admin sees.""" + mock_client = MagicMock() + mock_client.ping = AsyncMock( + side_effect=ConnectionError("Error connecting to redis://:super-secret-redis-pw@coord-redis.example.com:6379") + ) + mock_client.disconnect = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server._build_redis_usage_cache", return_value=mock_client), + ): + response = await check_coordination_redis_connection( + request=CoordinationRedisSettingsRequest( + settings={ + "host": "coord-redis.example.com", + "url": "redis://:super-secret-redis-pw@coord-redis.example.com:6379", + "password": "super-secret-redis-pw", + } + ), + user_api_key_dict=_admin_auth(), + ) + + assert response.status == "unhealthy" + assert response.error is not None + assert "super-secret-redis-pw" not in response.error + assert _REDACTED_VALUE in response.error + mock_client.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_connection_test_uses_the_saved_password_for_a_redacted_field(): + """An admin re-testing settings read back from GET sends `***REDACTED***`; + the saved credential is what actually gets dialed.""" + mock_client = MagicMock() + mock_client.ping = AsyncMock(return_value=True) + mock_client.disconnect = AsyncMock() + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + _prisma_with_general_settings({"coordination_redis": _SAVED_SETTINGS}), + ), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server._build_redis_usage_cache", return_value=mock_client) as mock_build, + ): + response = await check_coordination_redis_connection( + request=CoordinationRedisSettingsRequest( + settings={"host": "coord-redis.example.com", "password": _REDACTED_VALUE} + ), + user_api_key_dict=_admin_auth(), + ) + + assert response.status == "healthy" + assert mock_build.call_args.args[0]["password"] == "super-secret-redis-pw" + + +@pytest.mark.asyncio +async def test_connection_test_times_out_instead_of_hanging(): + async def _never_returns(): + await asyncio.sleep(60) + + mock_client = MagicMock() + mock_client.ping = MagicMock(side_effect=lambda: _never_returns()) + mock_client.disconnect = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + patch("litellm.proxy.proxy_server._build_redis_usage_cache", return_value=mock_client), + patch( + "litellm.proxy.management_endpoints.coordination_redis_endpoints._PING_TIMEOUT_SECONDS", + 0.01, + ), + ): + response = await check_coordination_redis_connection( + request=CoordinationRedisSettingsRequest(settings={"host": "unreachable"}), + user_api_key_dict=_admin_auth(), + ) + + assert response.status == "unhealthy" + assert "timed out" in (response.error or "") + + +@pytest.mark.asyncio +async def test_connection_test_rejects_settings_without_a_connection_target(): + with ( + patch("litellm.proxy.proxy_server.prisma_client", _prisma_with_general_settings({})), + patch("litellm.proxy.proxy_server.proxy_config", _proxy_config()), + ): + with pytest.raises(HTTPException) as exc_info: + await check_coordination_redis_connection( + request=CoordinationRedisSettingsRequest(settings={"ssl": True}), + user_api_key_dict=_admin_auth(), + ) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_connection_test_rejects_non_admin(): + with pytest.raises(HTTPException) as exc_info: + await check_coordination_redis_connection( + request=CoordinationRedisSettingsRequest(settings={"host": "coord-redis.example.com"}), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert exc_info.value.status_code == 403 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d06a1c16ab9..603d5cc15b7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5,6 +5,7 @@ import os import socket import subprocess import sys +import types from datetime import datetime, timedelta, timezone from pathlib import Path from unittest import mock @@ -23,6 +24,10 @@ sys.path.insert( ) # Adds the parent directory to the system-path import litellm +import litellm.proxy.proxy_server as proxy_server_module +from litellm.caching.caching import RedisCache +from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize @@ -9358,3 +9363,326 @@ def test_update_config_redacts_all_environment_variable_values( assert "db.internal" not in data["updated_values"] finally: restore() + + +class _EnvBuiltRedisCache(RedisCache): + """RedisCache stand-in that records its constructor kwargs and never + opens a network connection, so tests can assert which connection params + the proxy used to build its coordination Redis.""" + + def __init__(self, **kwargs): + self.init_kwargs = kwargs + + +def _run_init_cache_with_backend(cache_backend, redis_env_kwargs): + """Run ProxyConfig._init_cache with a stubbed response-cache backend and a + controlled REDIS_* environment, returning (redis_usage_cache, + spend_counter redis, config-cache redis) as observed after the call.""" + mock_litellm_cache = MagicMock() + mock_litellm_cache.cache = cache_backend + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + + with ( + patch.object(proxy_server_module, "redis_usage_cache", None), + patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), + patch.object(proxy_server_module, "user_api_key_cache", DualCache()), + patch.object(proxy_server_module, "llm_router", None), + patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch( + "litellm._redis._redis_kwargs_from_environment", + return_value=redis_env_kwargs, + ), + patch("litellm.Cache", return_value=mock_litellm_cache), + ): + litellm.cache = None + resolved = proxy_server_module.ProxyConfig()._init_cache(cache_params={"type": "qdrant-semantic"}) + return ( + resolved, + fresh_spend_cache.redis_cache, + fresh_config_cache.redis_cache, + ) + + +def test_init_cache_non_redis_backend_builds_usage_redis_from_environment(): + """A semantic (non-Redis-KV) response cache must not disable the proxy's + coordination Redis: when REDIS_* env vars provide a connection, + _init_cache builds a standalone usage cache so cross-pod rate limits, + spend tracking, and the pod lock manager stay Redis-backed.""" + usage_cache, spend_redis, config_redis = _run_init_cache_with_backend( + cache_backend=object(), + redis_env_kwargs={"host": "coordination-redis", "port": "6379"}, + ) + + assert isinstance(usage_cache, _EnvBuiltRedisCache) + assert usage_cache.init_kwargs["host"] == "coordination-redis" + assert spend_redis is usage_cache + assert config_redis is usage_cache + + +def test_init_cache_non_redis_backend_without_redis_env_stays_in_memory(): + """Without any REDIS_* connection info, a non-Redis response cache must + leave the coordination Redis unset instead of building a broken client.""" + usage_cache, spend_redis, config_redis = _run_init_cache_with_backend( + cache_backend=object(), + redis_env_kwargs={}, + ) + + assert usage_cache is None + assert spend_redis is None + assert config_redis is None + + +def test_init_cache_redis_backend_reuses_cache_backend_over_environment(): + """When the response cache itself is a plain Redis KV cache, it must be + reused as the coordination Redis; the REDIS_* environment fallback must + not construct a second client.""" + redis_backend = _EnvBuiltRedisCache(host="cache-params-host") + usage_cache, spend_redis, _ = _run_init_cache_with_backend( + cache_backend=redis_backend, + redis_env_kwargs={"host": "env-host"}, + ) + + assert usage_cache is redis_backend + assert usage_cache.init_kwargs["host"] == "cache-params-host" + assert spend_redis is redis_backend + + +class _EnvBuiltClusterCache(RedisClusterCache): + """RedisClusterCache stand-in that records constructor kwargs and never + opens a network connection.""" + + def __init__(self, **kwargs): + self.init_kwargs = kwargs + + +def _run_init_coordination_redis(config, env=None): + """Run ProxyConfig._init_coordination_redis against a stubbed module state, + returning (redis_usage_cache, spend_counter redis, config-cache redis).""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + + with ( + patch.object(proxy_server_module, "redis_usage_cache", None), + patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), + patch.object(proxy_server_module, "user_api_key_cache", DualCache()), + patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + mock.patch.dict(os.environ, env or {}, clear=False), + ): + built = proxy_server_module.ProxyConfig()._init_coordination_redis(config=config) + return ( + built, + fresh_spend_cache.redis_cache, + fresh_config_cache.redis_cache, + ) + + +def test_init_coordination_redis_explicit_block_builds_standalone_client(): + """general_settings.coordination_redis must build the coordination Redis + even when no response cache is configured at all, and attach it to the + spend counter and config caches.""" + usage_cache, spend_redis, config_redis = _run_init_coordination_redis( + config={"general_settings": {"coordination_redis": {"host": "coord-host", "port": 6380}}}, + ) + + assert isinstance(usage_cache, _EnvBuiltRedisCache) + assert usage_cache.init_kwargs["host"] == "coord-host" + assert usage_cache.init_kwargs["port"] == 6380 + assert spend_redis is usage_cache + assert config_redis is usage_cache + + +def test_init_coordination_redis_resolves_os_environ_references(): + """os.environ/ values inside the coordination_redis block must be resolved + the same way cache_params values are.""" + usage_cache, _, _ = _run_init_coordination_redis( + config={"general_settings": {"coordination_redis": {"host": "os.environ/COORD_REDIS_HOST"}}}, + env={"COORD_REDIS_HOST": "resolved-host"}, + ) + + assert usage_cache.init_kwargs["host"] == "resolved-host" + + +def test_init_coordination_redis_startup_nodes_builds_cluster_client(): + """A coordination_redis block with startup_nodes must construct a cluster + client, so cluster-aware consumers (v3 rate limiter) take the cluster path.""" + usage_cache, _, _ = _run_init_coordination_redis( + config={ + "general_settings": { + "coordination_redis": {"startup_nodes": [{"host": "node-1", "port": 7000}]} + } + }, + ) + + assert isinstance(usage_cache, _EnvBuiltClusterCache) + assert usage_cache.init_kwargs["startup_nodes"] == [{"host": "node-1", "port": 7000}] + + +def test_init_coordination_redis_without_connection_target_raises(): + """A coordination_redis block with no host, url, startup_nodes, or + sentinel_nodes is a config error and must fail startup loudly instead of + silently running without coordination.""" + with pytest.raises(ValueError, match="connection target"): + _run_init_coordination_redis( + config={"general_settings": {"coordination_redis": {"ssl": True}}}, + ) + + +def test_init_coordination_redis_non_mapping_block_raises(): + """A scalar coordination_redis value is a config error.""" + with pytest.raises(ValueError, match="mapping"): + _run_init_coordination_redis( + config={"general_settings": {"coordination_redis": "redis://host:6379"}}, + ) + + +def test_init_coordination_redis_absent_leaves_usage_cache_unset(): + """Without the block, nothing changes: the coordination Redis stays unset + for the downstream borrow / env fallback logic to decide.""" + usage_cache, spend_redis, _ = _run_init_coordination_redis( + config={"general_settings": {}}, + ) + + assert usage_cache is None + assert spend_redis is None + + +def test_explicit_coordination_redis_takes_precedence_over_cache_backend(): + """When both an explicit coordination_redis block and a plain-Redis + response cache are configured, the explicit block must win; the cache + backend must not overwrite it.""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + cache_backend = _EnvBuiltRedisCache(host="cache-backend-host") + mock_litellm_cache = MagicMock() + mock_litellm_cache.cache = cache_backend + + with ( + patch.object(proxy_server_module, "redis_usage_cache", None), + patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), + patch.object(proxy_server_module, "user_api_key_cache", DualCache()), + patch.object(proxy_server_module, "llm_router", None), + patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + patch("litellm.Cache", return_value=mock_litellm_cache), + ): + litellm.cache = None + proxy_config = proxy_server_module.ProxyConfig() + built = proxy_config._init_coordination_redis( + config={"general_settings": {"coordination_redis": {"host": "explicit-coord-host"}}} + ) + assert built is not None + proxy_server_module.redis_usage_cache = built + usage_cache = proxy_config._init_cache(cache_params={"type": "redis"}) + + assert isinstance(usage_cache, _EnvBuiltRedisCache) + assert usage_cache is not cache_backend + assert usage_cache.init_kwargs["host"] == "explicit-coord-host" + assert fresh_spend_cache.redis_cache is usage_cache + + +def test_env_fallback_builds_cluster_client_from_cluster_nodes_env(): + """A deployment whose only Redis env is REDIS_CLUSTER_NODES must still get + a coordination Redis from the env fallback, and it must be a cluster + client so cluster-aware consumers take the cluster path.""" + nodes = '[{"host": "cnode-1", "port": 7000}]' + with ( + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + patch("litellm._redis._redis_kwargs_from_environment", return_value={}), + mock.patch.dict(os.environ, {"REDIS_CLUSTER_NODES": nodes}, clear=False), + ): + result = proxy_server_module._build_redis_usage_cache_from_environment() + + assert isinstance(result, _EnvBuiltClusterCache) + assert result.init_kwargs["startup_nodes"] == [{"host": "cnode-1", "port": 7000}] + + +def test_env_fallback_builds_client_from_sentinel_nodes_env(): + """A sentinel-only environment (REDIS_SENTINEL_NODES, no host or url) must + also produce a coordination Redis from the env fallback.""" + with ( + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + patch("litellm._redis._redis_kwargs_from_environment", return_value={}), + mock.patch.dict(os.environ, {"REDIS_SENTINEL_NODES": '[["s1", 26379]]'}, clear=False), + ): + result = proxy_server_module._build_redis_usage_cache_from_environment() + + assert isinstance(result, _EnvBuiltRedisCache) + + +@pytest.mark.asyncio +async def test_startup_applies_coordination_redis_saved_in_database(): + """A coordination_redis block saved from the admin UI lives only in the + database, so startup must read it and build the coordination Redis from it. + Without this the save endpoint's "restart to apply" promise is false and the + proxy silently coordinates in per-pod memory.""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + + with ( + patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), + patch.object(proxy_server_module, "user_api_key_cache", DualCache()), + patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + patch.object( + proxy_server_module, + "get_persisted_coordination_redis_settings", + AsyncMock(return_value={"host": "db-host", "port": 6381}), + ), + ): + result = await proxy_server_module.ProxyStartupEvent._init_coordination_redis_from_db( + litellm_settings={}, + llm_router=None, + ) + + assert isinstance(result, _EnvBuiltRedisCache) + assert result.init_kwargs["host"] == "db-host" + assert fresh_spend_cache.redis_cache is result + assert fresh_config_cache.redis_cache is result + + +@pytest.mark.asyncio +async def test_startup_ignores_database_coordination_redis_without_connection_target(): + """A persisted block with no host/url/cluster/sentinel must be ignored rather + than crashing startup or building a client that cannot connect.""" + with ( + patch.object(proxy_server_module, "spend_counter_cache", DualCache()), + patch.object(proxy_server_module, "litellm_config_cache", types.SimpleNamespace(redis_cache=None)), + patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + patch.object( + proxy_server_module, + "get_persisted_coordination_redis_settings", + AsyncMock(return_value={"ssl": True}), + ), + ): + result = await proxy_server_module.ProxyStartupEvent._init_coordination_redis_from_db( + litellm_settings={}, + llm_router=None, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_startup_survives_database_read_failure_for_coordination_redis(): + """A config-row read failure must not block proxy startup.""" + with ( + patch.object( + proxy_server_module, + "get_persisted_coordination_redis_settings", + AsyncMock(side_effect=RuntimeError("db unreachable")), + ), + ): + result = await proxy_server_module.ProxyStartupEvent._init_coordination_redis_from_db( + litellm_settings={}, + llm_router=None, + ) + + assert result is None diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 9cf60ea5f91..d8e3f495ced 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1959,6 +1959,110 @@ class TestUsageTransformation: assert response_usage.output_tokens_details.text_tokens == 50 assert response_usage.output_tokens_details.image_tokens == 100 + def test_reasoning_tokens_not_forced_to_zero_when_absent(self): + # Regression: previously the else branch wrote reasoning_tokens=0 even when + # completion_tokens_details had no reasoning (reasoning_tokens=None). That caused + # the proxy to always report reasoning_tokens=0 for non-thinking responses. + usage = Usage( + prompt_tokens=10, + completion_tokens=50, + total_tokens=60, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=50, + # reasoning_tokens intentionally absent -> None + ), + ) + + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="claude-haiku-4-5", + object="chat.completion", + usage=usage, + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Hello!", role="assistant"), + ) + ], + ) + + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=chat_completion_response + ) + + assert response_usage.output_tokens_details is not None + assert response_usage.output_tokens_details.reasoning_tokens is None + + def test_reasoning_tokens_preserved_when_thinking_occurred(self): + # Regression: reasoning_tokens must survive the chat->responses translation + # when the provider actually did thinking. + usage = Usage( + prompt_tokens=100, + completion_tokens=612, + total_tokens=712, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=512, + text_tokens=100, + ), + ) + + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="claude-haiku-4-5", + object="chat.completion", + usage=usage, + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Hello!", role="assistant"), + ) + ], + ) + + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=chat_completion_response + ) + + assert response_usage.output_tokens_details is not None + assert response_usage.output_tokens_details.reasoning_tokens == 512 + + def test_reasoning_tokens_explicit_zero_preserved(self): + usage = Usage( + prompt_tokens=10, + completion_tokens=50, + total_tokens=60, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=0, + text_tokens=50, + ), + ) + + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="gpt-5.6", + object="chat.completion", + usage=usage, + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Hello!", role="assistant"), + ) + ], + ) + + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=chat_completion_response + ) + + assert response_usage.output_tokens_details is not None + assert response_usage.output_tokens_details.reasoning_tokens == 0 + class TestStreamingIDConsistency: """Test cases for consistent IDs across streaming events (issue #14962)""" diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py new file mode 100644 index 00000000000..1a2dcd0fcb7 --- /dev/null +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -0,0 +1,357 @@ +""" +Regression: in-stream error events (type="error", type="response.failed") must +raise instead of being returned as benign chunks, mirroring chat streaming +semantics (_handle_stream_fallback_error): non-retriable 4xx (except 429) +raise litellm.APIError directly; 429 and 5xx are wrapped in +MidStreamFallbackError so the Router's mid-stream fallback machinery fires. + +Status mapping must consider both the OpenAI error `type` (e.g. +"invalid_request_error") and `code` (e.g. "invalid_prompt", +"rate_limit_exceeded") fields — previously only `code` was read, so +type-classified client errors fell through to 500. + +Also covers: ErrorEventError.param must accept dict payloads without raising a +Pydantic ValidationError (previously typed as Optional[str]). +""" + +import json +import os +import sys +from unittest.mock import Mock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ResponsesAPIStreamingIterator, + SyncResponsesAPIStreamingIterator, +) +from litellm.types.llms.openai import ( + ErrorEvent, + ErrorEventError, + ResponseAPIUsage, + ResponsesAPIStreamEvents, +) + + +def _make_iterator() -> BaseResponsesAPIStreamingIterator: + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_config = Mock(spec=BaseResponsesAPIConfig) + mock_response = Mock() + mock_response.headers = {} + return BaseResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-5", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + +def _make_error_chunk(error_type: str, code: str, message: str = "err") -> ErrorEvent: + error_obj = ErrorEventError(type=error_type, code=code, message=message) + return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj) + + +def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback(): + iterator = _make_iterator() + chunk = _make_error_chunk("server_error", "internal_error", "something went wrong") + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 500 + assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert exc_info.value.original_exception.status_code == 500 + + +def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fallback(): + """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped APIError.""" + iterator = _make_iterator() + chunk = _make_error_chunk("tokens", "rate_limit_exceeded", "Too many requests") + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 429 + assert exc_info.value.generated_content == "" + assert exc_info.value.is_pre_first_chunk is True + assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert exc_info.value.original_exception.status_code == 429 + + +def test_maybe_raise_for_error_event_maps_invalid_request_type_to_400(): + """Client errors classified via the `type` field must raise APIError directly (no fallback).""" + iterator = _make_iterator() + chunk = _make_error_chunk("invalid_request_error", "invalid_prompt", "bad request") + with pytest.raises(litellm.APIError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, MidStreamFallbackError) + + +def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): + """Client errors classified via the `code` field alone must still map to 400.""" + iterator = _make_iterator() + chunk = Mock() + chunk.type = "error" + chunk.error = {"code": "context_length_exceeded", "message": "too long"} + with pytest.raises(litellm.APIError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, MidStreamFallbackError) + + +def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): + """OpenAI returns HTTP 429 for insufficient_quota; it must not map to 400 even though its type + is invalid_request_error-adjacent, and it must be wrapped for fallback.""" + iterator = _make_iterator() + chunk = _make_error_chunk("invalid_request_error", "insufficient_quota", "You exceeded your current quota") + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 429 + + +def test_maybe_raise_for_error_event_passes_through_normal_chunk(): + iterator = _make_iterator() + chunk = Mock() + chunk.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + iterator._maybe_raise_for_error_event(chunk) # must not raise + + +def test_error_event_error_param_accepts_dict(): + error_obj = ErrorEventError( + type="invalid_request_error", + code="context_length_exceeded", + message="too long", + param={"field": "messages", "index": 0}, + ) + assert isinstance(error_obj.param, dict) + + +def _make_async_iterator_with_events(events: list) -> ResponsesAPIStreamingIterator: + sse_payload = b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) + + async def mock_aiter_bytes(): + yield sse_payload + + mock_response = Mock() + mock_response.headers = {} + mock_response.aiter_bytes = mock_aiter_bytes + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.completion_start_time = None + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def transform(model, parsed_chunk, logging_obj): + if parsed_chunk.get("type") == "error": + return ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=0, + error=ErrorEventError(**parsed_chunk["error"]), + ) + delta_event = Mock() + delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + delta_event.delta = parsed_chunk.get("delta", "") + return delta_event + + mock_config.transform_streaming_response.side_effect = transform + + return ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-5", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + +@pytest.mark.asyncio +async def test_async_iterator_raises_mid_stream_fallback_on_rate_limit_error_event(): + iterator = _make_async_iterator_with_events( + [ + { + "type": "error", + "error": {"type": "tokens", "code": "rate_limit_exceeded", "message": "rate limited"}, + } + ] + ) + + with pytest.raises(MidStreamFallbackError) as exc_info: + async for _ in iterator: + pass + assert exc_info.value.status_code == 429 + assert exc_info.value.is_pre_first_chunk is True + assert exc_info.value.generated_content == "" + assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert exc_info.value.original_exception.status_code == 429 + + +@pytest.mark.asyncio +async def test_async_iterator_error_after_first_chunk_carries_generated_content(): + """An error after streamed output must expose the accumulated text so the router's + fallback can build a continuation input instead of restarting from scratch.""" + iterator = _make_async_iterator_with_events( + [ + {"type": "response.output_text.delta", "delta": "hello "}, + {"type": "response.output_text.delta", "delta": "world"}, + { + "type": "error", + "error": {"type": "server_error", "code": "internal_error", "message": "boom"}, + }, + ] + ) + + chunks = [] + with pytest.raises(MidStreamFallbackError) as exc_info: + async for chunk in iterator: + chunks.append(chunk) + assert len(chunks) == 2 + assert exc_info.value.status_code == 500 + assert exc_info.value.is_pre_first_chunk is False + assert exc_info.value.generated_content == "hello world" + + +def test_maybe_raise_for_response_failed_event_with_dict_error(): + """response.failed chunks carry a dict error on .response.error; covers dict branch.""" + iterator = _make_iterator() + mock_response_obj = Mock() + mock_response_obj.error = {"type": "tokens", "code": "rate_limit_exceeded", "message": "throttled"} + chunk = Mock() + chunk.type = "response.failed" + chunk.response = mock_response_obj + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 429 + + +def test_maybe_raise_for_error_event_null_error_obj(): + """error chunk with no error field: message and code default; wrapped as 500.""" + iterator = _make_iterator() + chunk = Mock() + chunk.type = "error" + chunk.error = None + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 500 + assert "Response API in-stream error" in str(exc_info.value) + + +def _make_failed_chunk(error: dict, usage: ResponseAPIUsage | None = None) -> Mock: + mock_response_obj = Mock() + mock_response_obj.error = error + mock_response_obj.usage = usage + chunk = Mock() + chunk.type = "response.failed" + chunk.response = mock_response_obj + return chunk + + +def test_handle_logging_failed_response_maps_rate_limit_to_429(): + """The exception logged to failure handlers must carry the mapped status, not a hardcoded 500.""" + iterator = _make_iterator() + iterator.completed_response = _make_failed_chunk( + {"type": "tokens", "code": "rate_limit_exceeded", "message": "throttled"} + ) + with ( + patch("litellm.responses.streaming_iterator.run_async_function") as mock_run_async, + patch("litellm.responses.streaming_iterator.executor"), + ): + iterator._handle_logging_failed_response() + logged_exception = mock_run_async.call_args.kwargs["exception"] + assert isinstance(logged_exception, litellm.APIError) + assert logged_exception.status_code == 429 + assert "throttled" in str(logged_exception) + + +def test_handle_logging_failed_response_maps_type_field_to_400(): + """Status derivation for failed-response logging must also read the error `type` field.""" + iterator = _make_iterator() + iterator.completed_response = _make_failed_chunk( + {"type": "invalid_request_error", "code": "invalid_prompt", "message": "bad prompt"} + ) + with ( + patch("litellm.responses.streaming_iterator.run_async_function") as mock_run_async, + patch("litellm.responses.streaming_iterator.executor"), + ): + iterator._handle_logging_failed_response() + logged_exception = mock_run_async.call_args.kwargs["exception"] + assert isinstance(logged_exception, litellm.APIError) + assert logged_exception.status_code == 400 + + +def test_handle_logging_failed_response_records_usage_and_cost(): + """Usage on a response.failed event must reach failure spend accounting via combined_usage_object.""" + iterator = _make_iterator() + usage = ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15) + chunk = _make_failed_chunk( + {"type": "server_error", "code": "server_error", "message": "boom"}, + usage=usage, + ) + iterator.completed_response = chunk + iterator.logging_obj._response_cost_calculator.return_value = 0.0042 + with ( + patch("litellm.responses.streaming_iterator.run_async_function"), + patch("litellm.responses.streaming_iterator.executor"), + ): + iterator._handle_logging_failed_response() + combined_usage = iterator.logging_obj.model_call_details["combined_usage_object"] + assert isinstance(combined_usage, litellm.Usage) + assert combined_usage.prompt_tokens == 10 + assert combined_usage.completion_tokens == 5 + assert combined_usage.total_tokens == 15 + assert iterator.logging_obj.model_call_details["response_cost"] == 0.0042 + iterator.logging_obj._response_cost_calculator.assert_called_once_with(result=chunk.response) + + +def test_handle_logging_failed_response_without_usage_skips_recording(): + iterator = _make_iterator() + iterator.completed_response = _make_failed_chunk( + {"type": "server_error", "code": "server_error", "message": "boom"} + ) + with ( + patch("litellm.responses.streaming_iterator.run_async_function"), + patch("litellm.responses.streaming_iterator.executor"), + ): + iterator._handle_logging_failed_response() + assert "combined_usage_object" not in iterator.logging_obj.model_call_details + iterator.logging_obj._response_cost_calculator.assert_not_called() + + +def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event(): + """SyncResponsesAPIStreamingIterator must wrap retriable error events for fallback.""" + error_payload = { + "type": "error", + "error": {"type": "tokens", "code": "rate_limit_exceeded", "message": "throttled"}, + } + sse_bytes = f"data: {json.dumps(error_payload)}\n\n".encode() + + mock_response = Mock() + mock_response.headers = {} + mock_response.iter_bytes.return_value = iter([sse_bytes]) + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.completion_start_time = None + mock_config = Mock(spec=BaseResponsesAPIConfig) + + error_obj = ErrorEventError(type="tokens", code="rate_limit_exceeded", message="throttled") + mock_config.transform_streaming_response.return_value = ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj + ) + + iterator = SyncResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-5", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + with pytest.raises(MidStreamFallbackError) as exc_info: + for _ in iterator: + pass + assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.original_exception, litellm.APIError) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index d4a7da86734..c40b362333a 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -6,10 +6,11 @@ Tests the rule-based complexity scoring and tier assignment logic. import os import sys -from typing import Dict, List -from unittest.mock import MagicMock, patch +from typing import Dict +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pydantic import ValidationError sys.path.insert( 0, os.path.abspath("../../..") @@ -743,7 +744,7 @@ class TestKeywordFalsePositives: # Should NOT detect code presence from 'api' in 'capital' assert not any( "code" in s.lower() for s in signals - ), f"False positive: got code signal from 'capital'" + ), "False positive: got code signal from 'capital'" # Should be SIMPLE (definition question) assert tier == ComplexityTier.SIMPLE @@ -754,7 +755,7 @@ class TestKeywordFalsePositives: # Should NOT detect code presence from 'git' in 'digital' assert not any( "code" in s.lower() for s in signals - ), f"False positive: got code signal from 'digital'" + ), "False positive: got code signal from 'digital'" def test_try_not_in_entry(self, complexity_router): """'try' should not match in 'entry'.""" @@ -770,7 +771,7 @@ class TestKeywordFalsePositives: tier, score, signals = complexity_router.classify(prompt) assert not any( "code" in s.lower() for s in signals - ), f"False positive: got code signal from 'terrorism'" + ), "False positive: got code signal from 'terrorism'" def test_class_not_in_classical(self, complexity_router): """'class' should not match in 'classical'.""" @@ -778,7 +779,7 @@ class TestKeywordFalsePositives: tier, score, signals = complexity_router.classify(prompt) assert not any( "code" in s.lower() for s in signals - ), f"False positive: got code signal from 'classical'" + ), "False positive: got code signal from 'classical'" def test_merge_not_in_emerged(self, complexity_router): """'merge' should not match in 'emerged'.""" @@ -786,7 +787,7 @@ class TestKeywordFalsePositives: tier, score, signals = complexity_router.classify(prompt) assert not any( "code" in s.lower() for s in signals - ), f"False positive: got code signal from 'emerged'" + ), "False positive: got code signal from 'emerged'" def test_actual_api_keyword_detected(self, complexity_router): """Actual 'api' usage should be detected.""" @@ -1131,3 +1132,180 @@ class TestExtractUserMessageAndSystemPrompt: ) assert user_msg is None assert sys_prompt is None + + +def _llm_response(content: str): + """Build a fake acompletion response with the given message content.""" + response = MagicMock() + response.choices = [MagicMock()] + response.choices[0].message.content = content + return response + + +@pytest.fixture +def llm_classifier_config() -> Dict: + """Config with an LLM-based classifier wired to a 'haiku-classifier' model.""" + return { + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "claude-sonnet-4-20250514", + "REASONING": "o1-preview", + }, + "classifier_type": "llm", + "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400}, + } + + +@pytest.fixture +def llm_complexity_router(mock_router_instance, llm_classifier_config): + """ComplexityRouter configured to classify via an LLM call.""" + return ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=llm_classifier_config, + ) + + +class TestLLMClassifierConfig: + """Test config validation for the LLM classifier option.""" + + def test_llm_classifier_type_requires_config(self): + """classifier_type='llm' without classifier_llm_config must raise.""" + with pytest.raises(ValidationError): + ComplexityRouterConfig(classifier_type="llm") + + def test_heuristic_classifier_type_needs_no_llm_config(self): + """classifier_type='heuristic' (the default) needs no classifier_llm_config.""" + config = ComplexityRouterConfig() + assert config.classifier_type == "heuristic" + assert config.classifier_llm_config is None + + +class TestLLMClassifier: + """Test the LLM-based classifier path (aclassify) and its fallback behavior.""" + + @pytest.mark.asyncio + async def test_aclassify_heuristic_skips_llm_call(self, complexity_router, mock_router_instance): + """When classifier_type is 'heuristic' (default), aclassify must not call the LLM.""" + mock_router_instance.acompletion = AsyncMock() + tier, score, signals = await complexity_router.aclassify("Hello!") + mock_router_instance.acompletion.assert_not_called() + assert tier == ComplexityTier.SIMPLE + + @pytest.mark.asyncio + async def test_aclassify_llm_success_routes_by_llm_verdict( + self, llm_complexity_router, mock_router_instance + ): + """A well-formed structured LLM response should decide the tier directly. + + Uses a prompt that heuristic scoring alone would classify as SIMPLE, to prove + the LLM verdict -- not the heuristic scorer -- is what decided the tier. + """ + mock_router_instance.acompletion = AsyncMock( + return_value=_llm_response('{"tier": "COMPLEX"}') + ) + tier, score, signals = await llm_complexity_router.aclassify("hi") + assert tier == ComplexityTier.COMPLEX + assert "llm-classifier:COMPLEX" in signals + mock_router_instance.acompletion.assert_awaited_once() + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["model"] == "haiku-classifier" + assert call_kwargs["timeout"] == 0.4 + + @pytest.mark.asyncio + async def test_aclassify_forwards_request_metadata_for_spend_tracking( + self, llm_complexity_router, mock_router_instance + ): + """The classifier call must carry the original request's metadata. + + Without this, the proxy's cost-tracking gate (_should_track_cost_callback) + sees no user_api_key/team_id/user_id and silently drops all spend logging + and budget accounting for the classifier call. + """ + mock_router_instance.acompletion = AsyncMock( + return_value=_llm_response('{"tier": "SIMPLE"}') + ) + request_metadata = {"user_api_key": "sk-abc", "user_api_key_team_id": "team-1"} + await llm_complexity_router.aclassify( + "hi", request_kwargs={"litellm_metadata": request_metadata} + ) + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["metadata"] == request_metadata + + @pytest.mark.asyncio + async def test_aclassify_strips_budget_reservation_from_classifier_metadata( + self, llm_complexity_router, mock_router_instance + ): + """The classifier call must not receive the parent request's budget reservation. + + The reservation belongs to the routed completion the classifier is deciding + on, not to this internal classifier call. Forwarding it would let the + classifier's own cost-tracking reconcile against a reservation it has no + business touching, so it must be stripped while the rest of the attribution + metadata (key/team) is preserved. + """ + mock_router_instance.acompletion = AsyncMock( + return_value=_llm_response('{"tier": "SIMPLE"}') + ) + request_metadata = { + "user_api_key": "sk-abc", + "user_api_key_team_id": "team-1", + "user_api_key_budget_reservation": {"reserved_cost": 1.0}, + "user_api_key_auth": {"budget_reservation": {"reserved_cost": 1.0}}, + } + await llm_complexity_router.aclassify( + "hi", request_kwargs={"litellm_metadata": request_metadata} + ) + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["metadata"] == { + "user_api_key": "sk-abc", + "user_api_key_team_id": "team-1", + } + + @pytest.mark.asyncio + async def test_aclassify_falls_back_to_heuristic_on_llm_exception( + self, llm_complexity_router, mock_router_instance + ): + """A timeout/error from the classifier model must fall back to heuristic scoring.""" + mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out")) + tier, score, signals = await llm_complexity_router.aclassify("Hello!") + assert tier == llm_complexity_router.classify("Hello!")[0] + assert tier == ComplexityTier.SIMPLE + + @pytest.mark.asyncio + async def test_aclassify_falls_back_to_heuristic_on_unparseable_response( + self, llm_complexity_router, mock_router_instance + ): + """Non-JSON or schema-violating output must fall back to heuristic scoring, not raise.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response("not json")) + tier, score, signals = await llm_complexity_router.aclassify("Hello!") + assert tier == ComplexityTier.SIMPLE + + @pytest.mark.asyncio + async def test_aclassify_falls_back_to_heuristic_on_empty_content( + self, llm_complexity_router, mock_router_instance + ): + """Empty/None message content (e.g. provider quirk) must fall back, not raise.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(None)) + tier, score, signals = await llm_complexity_router.aclassify("Hello!") + assert tier == ComplexityTier.SIMPLE + + @pytest.mark.asyncio + async def test_pre_routing_hook_uses_llm_classifier_end_to_end( + self, llm_complexity_router, mock_router_instance + ): + """The full pre-routing hook should route using the LLM classifier's verdict.""" + mock_router_instance.acompletion = AsyncMock( + return_value=_llm_response('{"tier": "REASONING"}') + ) + request_metadata = {"user_api_key": "sk-abc", "user_api_key_team_id": "team-1"} + result = await llm_complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={"litellm_metadata": request_metadata}, + messages=[{"role": "user", "content": "hi"}], + ) + assert result is not None + assert result.model == "o1-preview" # REASONING tier model + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["metadata"] == request_metadata diff --git a/tests/test_litellm/test_gpt_5_6_model_metadata.py b/tests/test_litellm/test_gpt_5_6_model_metadata.py index af9bb117f78..5a7b621d521 100644 --- a/tests/test_litellm/test_gpt_5_6_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_6_model_metadata.py @@ -122,6 +122,8 @@ def test_azure_gpt_5_6_regional_model_info(model): assert info is not None, f"{model} not found in model_prices_and_context_window.json" assert info["litellm_provider"] == "azure" + assert info["mode"] == "chat" + input_cost, output_cost, cache_read_cost, _ = STANDARD_PRICING[_tier_key(model)] assert info["input_cost_per_token"] == pytest.approx(input_cost * 1.1) @@ -132,6 +134,13 @@ def test_azure_gpt_5_6_regional_model_info(model): assert info["input_cost_per_token_priority"] == pytest.approx(input_cost * 2.75) assert info["output_cost_per_token_priority"] == pytest.approx(output_cost * 2.75) + assert info["max_input_tokens"] == 1050000 + assert info["max_output_tokens"] == 128000 + assert info["supports_reasoning"] is True + + _, provider, _, _ = get_llm_provider(model=model) + assert provider == "azure" + def test_gpt_5_6_backup_matches_main(): """Ensure the bundled model cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 0081e1c819f..4b9f13c340b 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -629,3 +629,98 @@ def test_sync_client_url_used_when_no_cluster(mock_from_url, monkeypatch): get_redis_client() mock_from_url.assert_called_once() + + +@patch("litellm._redis.redis.Redis.from_url") +def test_explicit_host_outranks_environment_redis_url(mock_from_url, monkeypatch): + """ + An explicitly configured host must win over REDIS_URL in the environment. + + Otherwise the url branch strips the caller's host/port and the client + silently connects to whatever REDIS_URL names, so an explicit config block + (or a connection test typed into the admin UI) targets the wrong server. + """ + monkeypatch.setenv("REDIS_URL", "redis://env-host:6379") + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + + client = get_redis_client(host="explicit-host", port=6380) + + mock_from_url.assert_not_called() + assert client.connection_pool.connection_kwargs["host"] == "explicit-host" + assert client.connection_pool.connection_kwargs["port"] == 6380 + + +@patch("litellm._redis.redis.Redis.from_url") +def test_explicit_url_still_wins_over_environment_host(mock_from_url, monkeypatch): + """An explicit url argument keeps taking the from_url path.""" + monkeypatch.setenv("REDIS_HOST", "env-host") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + + get_redis_client(url="redis://explicit-host:6380") + + mock_from_url.assert_called_once() + assert mock_from_url.call_args.kwargs["url"] == "redis://explicit-host:6380" + + +@patch("litellm._redis.redis.Redis.from_url") +def test_environment_redis_url_used_when_caller_names_no_target(mock_from_url, monkeypatch): + """With no caller-supplied connection target, REDIS_URL still drives the client.""" + monkeypatch.setenv("REDIS_URL", "redis://env-host:6379") + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + + get_redis_client() + + mock_from_url.assert_called_once() + + +@pytest.mark.parametrize("falsy_ssl", [False, None, 0, ""]) +def test_connection_pool_falsy_ssl_uses_plain_connection(falsy_ssl, monkeypatch): + """ + ssl=False must produce a plain (non-TLS) connection pool. + + The admin UI's coordination Redis form always sends ssl explicitly, so a + presence check here turns ssl=False into an SSLConnection; the TLS + handshake against a plaintext Redis then hangs until the ping timeout and + every connection test from the UI fails. + """ + monkeypatch.delenv("REDIS_URL", raising=False) + monkeypatch.delenv("REDIS_SSL", raising=False) + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + + with patch("litellm._redis.async_redis.BlockingConnectionPool") as mock_pool: + get_redis_connection_pool(host="plain-redis.example.com", port=6379, ssl=falsy_ssl) + + call_kwargs = mock_pool.call_args.kwargs + assert call_kwargs.get("connection_class") is not async_redis.SSLConnection, ( + f"ssl={falsy_ssl!r} must not select SSLConnection" + ) + assert "ssl" not in call_kwargs, "ssl must never leak into BlockingConnectionPool kwargs" + + +def test_connection_pool_ssl_true_uses_ssl_connection(monkeypatch): + """ssl=True must still opt in to a TLS connection pool.""" + monkeypatch.delenv("REDIS_URL", raising=False) + monkeypatch.delenv("REDIS_SSL", raising=False) + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + + with patch("litellm._redis.async_redis.BlockingConnectionPool") as mock_pool: + get_redis_connection_pool(host="tls-redis.example.com", port=6380, ssl=True) + + call_kwargs = mock_pool.call_args.kwargs + assert call_kwargs.get("connection_class") is async_redis.SSLConnection + assert "ssl" not in call_kwargs, "ssl must be consumed, not forwarded to the pool" + + +def test_connection_pool_without_ssl_kwarg_uses_plain_connection(monkeypatch): + """Omitting ssl entirely must keep the historical plain-connection default.""" + monkeypatch.delenv("REDIS_URL", raising=False) + monkeypatch.delenv("REDIS_SSL", raising=False) + monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False) + + with patch("litellm._redis.async_redis.BlockingConnectionPool") as mock_pool: + get_redis_connection_pool(host="plain-redis.example.com", port=6379) + + call_kwargs = mock_pool.call_args.kwargs + assert call_kwargs.get("connection_class") is not async_redis.SSLConnection + assert "ssl" not in call_kwargs diff --git a/tests/test_litellm/test_responses_api_bridge_non_stream.py b/tests/test_litellm/test_responses_api_bridge_non_stream.py index 8905293d6b6..25a3bc2dbba 100644 --- a/tests/test_litellm/test_responses_api_bridge_non_stream.py +++ b/tests/test_litellm/test_responses_api_bridge_non_stream.py @@ -269,29 +269,30 @@ def test_transform_usage_with_zero_values(): """ Test transformation when token details are explicitly set to 0. - This ensures 0 values are preserved and not treated as None. + cached_tokens=0 is preserved (cache was available; nothing was cached). + reasoning_tokens=0 is preserved the same way: an explicit provider-reported + zero passes through, while an absent value (None) is omitted. """ completion_response = create_mock_completion_response( model="gpt-4", prompt_tokens=100, completion_tokens=50, total_tokens=150, - cached_tokens=0, # Explicitly 0 - reasoning_tokens=0, # Explicitly 0 + cached_tokens=0, # Explicitly 0 — preserved + reasoning_tokens=0, # Explicitly 0 — preserved ) responses_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( completion_response ) - # Should preserve 0 values assert responses_usage.input_tokens_details is not None assert responses_usage.input_tokens_details.cached_tokens == 0 assert responses_usage.output_tokens_details is not None assert responses_usage.output_tokens_details.reasoning_tokens == 0 - print("✓ Transformation preserves explicit 0 values") + print("✓ Transformation preserves explicit reasoning_tokens=0 and omits absent values") def test_input_tokens_details_requires_cached_tokens(): @@ -315,25 +316,23 @@ def test_input_tokens_details_requires_cached_tokens(): print("✓ InputTokensDetails correctly defaults cached_tokens to 0") -def test_output_tokens_details_requires_reasoning_tokens(): +def test_output_tokens_details_reasoning_tokens(): """ - Test that OutputTokensDetails has reasoning_tokens as an int with default value 0. + Test OutputTokensDetails.reasoning_tokens field semantics. - This ensures backward compatibility while making the field non-optional. + reasoning_tokens is Optional[int] = None: present only when reasoning actually occurred. """ - # Should work with reasoning_tokens=0 - details1 = OutputTokensDetails(reasoning_tokens=0) - assert details1.reasoning_tokens == 0 + details_explicit_zero = OutputTokensDetails(reasoning_tokens=0) + assert details_explicit_zero.reasoning_tokens == 0 - # Should work with reasoning_tokens=100 - details2 = OutputTokensDetails(reasoning_tokens=100) - assert details2.reasoning_tokens == 100 + details_positive = OutputTokensDetails(reasoning_tokens=100) + assert details_positive.reasoning_tokens == 100 - # Should work without reasoning_tokens (defaults to 0) - details3 = OutputTokensDetails() - assert details3.reasoning_tokens == 0 + # Default is None — absence means reasoning did not occur (or was not tracked) + details_default = OutputTokensDetails() + assert details_default.reasoning_tokens is None - print("✓ OutputTokensDetails correctly defaults reasoning_tokens to 0") + print("✓ OutputTokensDetails.reasoning_tokens defaults to None") def test_all_providers_transformation_scenarios(): @@ -419,7 +418,7 @@ if __name__ == "__main__": test_transform_usage_with_both_token_details() test_transform_usage_with_zero_values() test_input_tokens_details_requires_cached_tokens() - test_output_tokens_details_requires_reasoning_tokens() + test_output_tokens_details_reasoning_tokens() test_all_providers_transformation_scenarios() print("\n" + "=" * 60) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d739f9c116a..26ae6ff7f4c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -855,6 +855,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, "supports_adaptive_thinking": {"type": "boolean"}, + "supports_mid_conversation_system": {"type": "boolean"}, "supports_sampling_params": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, "supports_speed": {"type": "boolean"}, @@ -1091,8 +1092,8 @@ def test_get_model_info_bedrock_regional_inference_profile_pricing(local_model_c def test_get_model_info_bedrock_regional_profile_without_entry_falls_back_to_base(local_model_cost_map): """A regional profile with no dedicated cost-map entry must still resolve to its region-stripped base entry.""" - assert "jp.anthropic.claude-opus-4-8" not in litellm.model_cost - info = litellm.get_model_info(model="bedrock/jp.anthropic.claude-opus-4-8") + assert "apac.anthropic.claude-opus-4-8" not in litellm.model_cost + info = litellm.get_model_info(model="bedrock/apac.anthropic.claude-opus-4-8") assert info["key"] == "anthropic.claude-opus-4-8" diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts index d8644babfe3..1e2f2269dbe 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/auth/logout.spec.ts @@ -6,7 +6,8 @@ test.describe("Logout", () => { test("Clicking Logout clears the session and forces re-login on a protected page", async ({ page }) => { await page.goto("/ui"); - await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 }); + // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". + await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 }); // Open the navbar User dropdown. The trigger button exposes an aria-label // of "Account menu — — signed in as ", and the antd Dropdown diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts index 6358fcf438e..9faf6741333 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/auth/proxyLogoutUrl.spec.ts @@ -42,7 +42,8 @@ test.describe("PROXY_LOGOUT_URL redirect", () => { timeout: 30_000, }); await page.goto("/ui"); - await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); + // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". + await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); await settingsLoaded; // Pre-condition: we start authenticated. The admin storage state carries a diff --git a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts index 548639d6877..92e46d6b27c 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/internal-user/internalUserNoTeam.spec.ts @@ -16,7 +16,8 @@ test.describe("Internal User with no team memberships", () => { await page.getByPlaceholder("Enter your username").fill("noteam@test.local"); await page.getByPlaceholder("Enter your password").fill("test"); await page.getByRole("button", { name: "Login", exact: true }).click(); - await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); + // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". + await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); await dismissFeedbackPopup(page); // Open the Create Key modal. diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts index 6008049a2aa..569908c5f75 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/login/internalUserIdentity.spec.ts @@ -14,7 +14,8 @@ test.describe("Navbar identity scoping", () => { test("Internal user navbar dropdown shows their own role and user id, not the admin's", async ({ page }) => { await page.goto("/ui"); - await expect(page.getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 }); + // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". + await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 }); // The account menu button carries the user's role and email/id in its // aria-label (see UserDropdown.tsx). Match by partial role. diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts index d1b64f37156..11febf0ed48 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts @@ -9,7 +9,8 @@ test("user can log in", async ({ page }) => { const loginButton = page.getByRole("button", { name: "Login", exact: true }); await expect(loginButton).toBeEnabled(); await loginButton.click(); - await expect(page.getByText("Virtual Keys")).toBeVisible(); + // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". + await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible(); // Match the navbar account button by its stable aria-label (UserDropdown.tsx // emits "Account menu — — signed in as "). Earlier this used diff --git a/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts index 0a3be326e42..3ad4b217d08 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/migration/migratedPages.spec.ts @@ -17,7 +17,11 @@ const ROOT = process.env.SERVER_ROOT_PATH ?? ""; const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`); -const virtualKeysLink = (page: Page) => page.getByRole("link", { name: "Virtual Keys", exact: true }); +// Scope nav lookups to the sidebar (a `complementary` landmark). The top bar +// now renders a breadcrumb whose current-page item is also a "Virtual Keys" +// link, so an unscoped locator would match two elements. +const sidebar = (page: Page) => page.getByRole("complementary"); +const virtualKeysLink = (page: Page) => sidebar(page).getByRole("link", { name: "Virtual Keys", exact: true }); /** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */ async function expectRendered(page: Page) { @@ -26,16 +30,21 @@ async function expectRendered(page: Page) { /** * Click a migrated page's sidebar link. Migrated items render as ; - * nested ones live under collapsible submenus, so expand submenus until the link is clickable. + * nested ones live under collapsible groups whose children only render while the + * group is open, so expand collapsed groups until the link is clickable. */ async function clickSidebar(page: Page, segment: string) { - const link = page.locator(`a[href$="/ui/${segment}"]`).first(); + const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first(); for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) { - const collapsedSubmenu = page - .locator(".ant-menu-submenu:not(.ant-menu-submenu-open) > .ant-menu-submenu-title") + // A collapsed group is a menu item with a group-toggle button but no + // rendered submenu yet; clicking the toggle expands it. + const collapsedGroup = sidebar(page) + .locator( + '[data-slot="sidebar-menu-item"]:has(> [data-slot="sidebar-menu-button"]):not(:has(> [data-slot="sidebar-menu-sub"])) > [data-slot="sidebar-menu-button"]', + ) .first(); - if (!(await collapsedSubmenu.isVisible().catch(() => false))) break; - await collapsedSubmenu.click(); + if (!(await collapsedGroup.isVisible().catch(() => false))) break; + await collapsedGroup.click(); await page.waitForTimeout(250); } await link.click(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index 7ac2e7df39d..7e42d07ae7c 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -42,7 +42,9 @@ for (const { role, storage } of roles) { throw new Error(`No page mapping found for menu label: ${buttonLabel}`); } - const tab = page.getByRole("menuitem", { name: buttonLabel }); + // Sidebar items are links inside the `complementary` landmark; scoping + // there avoids the top-bar breadcrumb, which also links the page name. + const tab = page.getByRole("complementary").getByRole("link", { name: buttonLabel }); await expect(tab).toBeVisible(); await tab.click(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts index f61532b05a5..c4a14a891d4 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/settings/adminSettings.spec.ts @@ -6,8 +6,11 @@ test.describe("Add Model", () => { test("admin settings test", async ({ page }) => { await page.goto("/ui"); - await page.getByRole("menuitem", { name: /Settings/ }).click(); - await page.getByRole("menuitem", { name: /Admin Settings/ }).click(); + // "Settings" is a collapsible group (button) in the sidebar; expand it, then + // click the "Admin Settings" child link. Scope to the complementary landmark. + const sidebar = page.getByRole("complementary"); + await sidebar.getByRole("button", { name: /Settings/ }).click(); + await sidebar.getByRole("link", { name: /Admin Settings/ }).click(); await page.getByRole("tab", { name: "UI Settings" }).click(); await expect(page.getByText("Configuration for UI-specific")).toBeVisible(); }); diff --git a/ui/litellm-dashboard/eslint-metrics.json b/ui/litellm-dashboard/eslint-metrics.json index dac15a21620..7abbb186771 100644 --- a/ui/litellm-dashboard/eslint-metrics.json +++ b/ui/litellm-dashboard/eslint-metrics.json @@ -1,8 +1,8 @@ { - "@typescript-eslint/no-explicit-any": 1977, - "complexity": 129, + "@typescript-eslint/no-explicit-any": 1978, + "complexity": 130, "local/no-large-inline-object-arg": 509, - "local/no-long-condition-chain": 233, + "local/no-long-condition-chain": 234, "max-depth": 59, "no-console": 16 } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx new file mode 100644 index 00000000000..17d14cd7fac --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx @@ -0,0 +1,135 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen, waitFor, within } from "@testing-library/react"; +import { renderWithProviders } from "../../../../../tests/test-utils"; +import CacheDashboard from "./cache_dashboard"; + +const { adminGlobalCacheActivity, cachingHealthCheckCall } = vi.hoisted(() => ({ + adminGlobalCacheActivity: vi.fn(), + cachingHealthCheckCall: vi.fn(), +})); + +vi.mock("@/components/networking", () => ({ + adminGlobalCacheActivity, + cachingHealthCheckCall, +})); + +const cacheActivity = [ + { + api_key: "sk-1", + model: "gpt-5.1", + call_type: "acompletion", + total_rows: 1500, + cache_hit_true_rows: 300, + cached_completion_tokens: 12000, + generated_completion_tokens: 48000, + }, + { + api_key: "sk-2", + model: "text-embedding-3-large", + call_type: "aembedding", + total_rows: 700, + cache_hit_true_rows: 100, + cached_completion_tokens: 2000, + generated_completion_tokens: 9000, + }, +]; + +const renderDashboard = () => + renderWithProviders( + , + ); + +const findChartCards = async () => { + await screen.findByText("Cache Hits vs API Requests"); + await waitFor(() => { + expect(document.querySelectorAll("path.recharts-rectangle").length).toBeGreaterThan(0); + }); + const cards = Array.from(document.querySelectorAll('[data-slot="card"]')); + expect(cards).toHaveLength(2); + return { requestsCard: cards[0] as HTMLElement, tokensCard: cards[1] as HTMLElement }; +}; + +const barFills = (card: HTMLElement) => + Array.from(card.querySelectorAll(".recharts-bar")).map((bar) => + bar.querySelector("path.recharts-rectangle")?.getAttribute("fill"), + ); + +const legendFillByCategory = (card: HTMLElement) => + Object.fromEntries( + Array.from(card.querySelectorAll('.recharts-legend-wrapper [style*="background-color"]')).map((swatch) => [ + swatch.parentElement?.textContent, + swatch.getAttribute("style")?.match(/background-color:\s*([^;]+);?/)?.[1], + ]), + ); + +describe("CacheDashboard cache analytics charts", () => { + beforeEach(() => { + vi.clearAllMocks(); + adminGlobalCacheActivity.mockResolvedValue(cacheActivity); + }); + + it("renders both chart card titles", async () => { + renderDashboard(); + + expect(await screen.findByText("Cache Hits vs API Requests")).toBeInTheDocument(); + expect(screen.getByText("Cached Completion Tokens vs Generated Completion Tokens")).toBeInTheDocument(); + }); + + it("renders the requests chart with each category legend-bound to its fill and stacked in order", async () => { + renderDashboard(); + const { requestsCard } = await findChartCards(); + + expect(legendFillByCategory(requestsCard)).toEqual({ + "LLM API requests": "var(--color-sky-500, #0ea5e9)", + "Cache hit": "var(--color-teal-500, #14b8a6)", + }); + expect(barFills(requestsCard)).toEqual(["var(--color-sky-500, #0ea5e9)", "var(--color-teal-500, #14b8a6)"]); + }); + + it("renders the tokens chart with each category legend-bound to its fill and stacked in order", async () => { + renderDashboard(); + const { tokensCard } = await findChartCards(); + + expect(legendFillByCategory(tokensCard)).toEqual({ + "Generated Completion Tokens": "var(--color-sky-500, #0ea5e9)", + "Cached Completion Tokens": "var(--color-teal-500, #14b8a6)", + }); + expect(barFills(tokensCard)).toEqual(["var(--color-sky-500, #0ea5e9)", "var(--color-teal-500, #14b8a6)"]); + }); + + it("indexes bars by call_type name on the x axis", async () => { + renderDashboard(); + const { requestsCard, tokensCard } = await findChartCards(); + + for (const card of [requestsCard, tokensCard]) { + expect(within(card).getAllByText("acompletion").length).toBeGreaterThan(0); + expect(within(card).getAllByText("aembedding").length).toBeGreaterThan(0); + } + }); + + it("stacks the two categories into one column per call_type", async () => { + renderDashboard(); + const { requestsCard, tokensCard } = await findChartCards(); + + for (const card of [requestsCard, tokensCard]) { + const rects = Array.from(card.querySelectorAll("path.recharts-rectangle")); + expect(rects).toHaveLength(4); + const xPositions = rects.map((rect) => rect.getAttribute("d")?.split(",")[0]); + expect(new Set(xPositions).size).toBe(2); + } + }); + + it("formats y-axis ticks with compact notation", async () => { + renderDashboard(); + const { requestsCard, tokensCard } = await findChartCards(); + + const compactTicks = (card: HTMLElement) => + within(card) + .getAllByText(/^\d+(\.\d+)?K$/) + .map((tick) => tick.textContent); + + expect(compactTicks(requestsCard).length).toBeGreaterThan(0); + expect(compactTicks(tokensCard)).toContain("60K"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx index 944983da0e1..b8e8dc8adb1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx @@ -1,5 +1,4 @@ import { - BarChart, Card, Col, DateRangePickerValue, @@ -7,7 +6,6 @@ import { Icon, MultiSelect, MultiSelectItem, - Subtitle, Tab, TabGroup, TabList, @@ -18,6 +16,8 @@ import { import React, { useEffect, useState } from "react"; import NotificationsManager from "@/components/molecules/notifications_manager"; import UsageDatePicker from "@/components/shared/usage_date_picker"; +import { BarChart } from "@/components/shared/charts"; +import { Card as ChartCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { RefreshIcon } from "@heroicons/react/outline"; import { adminGlobalCacheActivity, cachingHealthCheckCall } from "@/components/networking"; @@ -25,6 +25,7 @@ import { adminGlobalCacheActivity, cachingHealthCheckCall } from "@/components/n // Import the new component import { CacheHealthTab } from "./cache_health"; import CacheSettings from "./cache_settings"; +import CoordinationRedisSettings from "./coordination_redis_settings"; const formatDateWithoutTZ = (date: Date | undefined) => { if (!date) return undefined; @@ -61,13 +62,13 @@ interface cacheDataItem { // Add other properties as needed } -interface uiData { +type uiData = { name: string; "LLM API requests": number; "Cache hit": number; "Cached Completion Tokens": number; "Generated Completion Tokens": number; -} +}; interface CacheHealthResponse { status?: string; @@ -264,6 +265,7 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole Cache Analytics Cache Health Cache Settings + Coordination Redis
@@ -348,29 +350,41 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole
- Cache Hits vs API Requests - + + + Cache Hits vs API Requests + + + + + - Cached Completion Tokens vs Generated Completion Tokens - + + + + Cached Completion Tokens vs Generated Completion Tokens + + + + + + @@ -383,6 +397,9 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole + + + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFieldSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFieldSection.tsx new file mode 100644 index 00000000000..af807926ed5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFieldSection.tsx @@ -0,0 +1,46 @@ +import React from "react"; +import CoordinationRedisFormField from "./CoordinationRedisFormField"; +import { fieldsForSection } from "./coordinationRedisUtils"; +import { CoordinationRedisType, CoordinationSection } from "./coordinationRedisFields"; + +interface CoordinationRedisFieldSectionProps { + title: string; + section: CoordinationSection; + redisType: CoordinationRedisType; + configuredSecrets: ReadonlySet; + gridCols?: string; + headingLevel?: "h4" | "h5"; +} + +const CoordinationRedisFieldSection: React.FC = ({ + title, + section, + redisType, + configuredSecrets, + gridCols = "grid-cols-1 gap-6 sm:grid-cols-2", + headingLevel = "h4", +}) => { + const fields = fieldsForSection(section, redisType); + if (fields.length === 0) { + return null; + } + + const Heading = headingLevel; + + return ( +
+ {title} +
+ {fields.map((field) => ( + + ))} +
+
+ ); +}; + +export default CoordinationRedisFieldSection; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFormField.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFormField.tsx new file mode 100644 index 00000000000..50c2a39567a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisFormField.tsx @@ -0,0 +1,39 @@ +import { Form, Input, Switch } from "antd"; +import React from "react"; +import { CoordinationField } from "./coordinationRedisFields"; + +export const SECRET_ALREADY_SET_PLACEHOLDER = "Already set. Enter a new value to replace it."; + +interface CoordinationRedisFormFieldProps { + field: CoordinationField; + isSecretConfigured: boolean; +} + +const renderControl = (field: CoordinationField, placeholder: string): React.ReactNode => { + switch (field.type) { + case "boolean": + return ; + case "password": + return ; + case "integer": + return ; + case "list": + return ; + default: + return ; + } +}; + +const CoordinationRedisFormField: React.FC = ({ field, isSecretConfigured }) => ( + + {renderControl(field, isSecretConfigured ? SECRET_ALREADY_SET_PLACEHOLDER : field.helpText)} + +); + +export default CoordinationRedisFormField; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.tsx new file mode 100644 index 00000000000..daab8505890 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.tsx @@ -0,0 +1,33 @@ +import React from "react"; +import { Select } from "antd"; +import { + COORDINATION_REDIS_TYPES, + COORDINATION_REDIS_TYPE_DESCRIPTIONS, + COORDINATION_REDIS_TYPE_LABELS, + CoordinationRedisType, +} from "./coordinationRedisFields"; + +interface CoordinationRedisTypeSelectorProps { + redisType: CoordinationRedisType; + onTypeChange: (type: CoordinationRedisType) => void; +} + +const OPTIONS = COORDINATION_REDIS_TYPES.map((type) => ({ value: type, label: COORDINATION_REDIS_TYPE_LABELS[type] })); + +const CoordinationRedisTypeSelector: React.FC = ({ redisType, onTypeChange }) => ( +
+ +