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/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/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index daee3369a3c..155fba008d4 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -94,25 +94,30 @@ class AmazonAnthropicClaudeMessagesConfig( 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.""" + """Bedrock Invoke rejects a conversation that opens with ``role: "system"`` + entries inside ``messages`` ("messages.0: use the top-level 'system' + parameter for the initial system prompt"); Anthropic Messages carries that + content in the top-level ``system`` field, so hoist the leading run of + system entries there. Mid-conversation system entries (e.g. Claude Code's + ``mid-conversation-system-2026-04-07`` reminders) are accepted by Invoke in + place and MUST stay in place: hoisting one mutates the ``system`` prefix + and invalidates the prompt cache for the entire message history. + 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") - ] + leading_count = next( + (i for i, m in enumerate(messages) if not (isinstance(m, dict) and m.get("role") == "system")), + len(messages), + ) + if leading_count: + anthropic_messages_request["messages"] = messages[leading_count:] 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 messages[:leading_count]), ) for block in self._as_system_content_blocks(source) ] 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/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/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/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/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..ffe45cd5e79 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,79 @@ def test_bedrock_invoke_transform_merges_list_content_system_role_into_system(): ] +def test_bedrock_invoke_transform_keeps_mid_conversation_system_role_in_place(): + """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 such entries must be forwarded in place. Invoke only rejects a + system entry at ``messages.0``. 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="anthropic.claude-opus-4-8", + 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(): + """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_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/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/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/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 }) => ( +
+ +