diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1a29c0b6691..51d489459d9 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,2 +1,3 @@ /ui/ @yuneng-jiang @ryan-crabbe-berri /litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri +/ui/litellm-dashboard/src/lib/http/schema.d.ts diff --git a/docker/build_from_pip/Dockerfile.build_from_pip b/docker/build_from_pip/Dockerfile.build_from_pip index bda742c71a9..372606a5b0f 100644 --- a/docker/build_from_pip/Dockerfile.build_from_pip +++ b/docker/build_from_pip/Dockerfile.build_from_pip @@ -36,7 +36,7 @@ RUN uv venv --python python && \ "opentelemetry-api==1.28.0" \ "opentelemetry-sdk==1.28.0" \ "opentelemetry-exporter-otlp==1.28.0" \ - "ddtrace==2.19.0" \ + "ddtrace==4.11.0" \ "sentry-sdk==2.21.0" \ "mangum==0.17.0" \ "azure-ai-contentsafety==1.0.0" \ diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index ad8aabf77b6..d10b5a2ab09 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -7,7 +7,8 @@ # Thank you users! We ❤️ you! - Krrish & Ishaan ## This provides an LLM Guard Integration for content moderation on the proxy -from typing import Literal, Optional +import asyncio +from typing import Optional import aiohttp from fastapi import HTTPException @@ -18,7 +19,6 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.secret_managers.main import get_secret_str from litellm.types.utils import CallTypesLiteral -from litellm.utils import get_formatted_prompt class _ENTERPRISE_LLMGuard(CustomLogger): @@ -46,45 +46,44 @@ class _ENTERPRISE_LLMGuard(CustomLogger): except Exception: pass - async def moderation_check(self, text: str): + async def moderation_check(self, text: str) -> str: """ + Runs the LLM Guard moderation check on ``text``. + + Raises an HTTPException when the content violates the safety policy; + otherwise returns the sanitized prompt from LLM Guard, falling back to + the original text when the API does not provide one. + [TODO] make this more performant for high-throughput scenario """ try: - async with aiohttp.ClientSession() as session: - if self.mock_redacted_text is not None: - redacted_text = self.mock_redacted_text - else: - # Make the first request to /analyze - analyze_url = f"{self.llm_guard_api_base}analyze/prompt" - verbose_proxy_logger.debug("Making request to: %s", analyze_url) - analyze_payload = {"prompt": text} - redacted_text = None + if self.mock_redacted_text is not None: + redacted_text = self.mock_redacted_text + else: + analyze_url = f"{self.llm_guard_api_base}analyze/prompt" + verbose_proxy_logger.debug("Making request to: %s", analyze_url) + async with aiohttp.ClientSession() as session: async with session.post( - analyze_url, json=analyze_payload + analyze_url, json={"prompt": text} ) as response: redacted_text = await response.json() - verbose_proxy_logger.debug( - f"LLM Guard: Received response - {redacted_text}" + verbose_proxy_logger.debug( + f"LLM Guard: Received response - {redacted_text}" + ) + if redacted_text is None: + raise HTTPException( + status_code=500, + detail={ + "error": f"Invalid content moderation response: {redacted_text}" + }, ) - if redacted_text is not None: - if ( - redacted_text.get("is_valid", None) is not None - and redacted_text["is_valid"] is False - ): - raise HTTPException( - status_code=400, - detail={"error": "Violated content safety policy"}, - ) - else: - pass - else: - raise HTTPException( - status_code=500, - detail={ - "error": f"Invalid content moderation response: {redacted_text}" - }, - ) + if redacted_text.get("is_valid", None) is False: + raise HTTPException( + status_code=400, + detail={"error": "Violated content safety policy"}, + ) + sanitized_prompt = redacted_text.get("sanitized_prompt") + return sanitized_prompt if isinstance(sanitized_prompt, str) else text except Exception as e: verbose_proxy_logger.exception( "litellm.enterprise.enterprise_hooks.llm_guard::moderation_check - Exception occurred - {}".format( @@ -138,23 +137,75 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return self.print_verbose("Makes LLM Guard Check") - try: - assert call_type in [ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - ] - except Exception: + if call_type not in [ + "completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + ]: self.print_verbose( f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']" ) return data - formatted_prompt = get_formatted_prompt(data=data, call_type=call_type) # type: ignore - self.print_verbose(f"LLM Guard, formatted_prompt: {formatted_prompt}") - return await self.moderation_check(text=formatted_prompt) + return await self._moderate_request(data=data) + + async def _moderate_request(self, data: dict) -> dict: + """ + Sanitizes the request in place using the prompt returned by LLM Guard so + the provider-bound request carries the redacted content, then returns it. + """ + messages = data.get("messages") + if messages is not None: + data["messages"] = list( + await asyncio.gather( + *(self._moderate_message(message) for message in messages) + ) + ) + return data + + input_ = data.get("input") + if input_ is not None: + data["input"] = await self._moderate_input(input_) + return data + + prompt = data.get("prompt") + if isinstance(prompt, str): + data["prompt"] = await self.moderation_check(text=prompt) + return data + + async def _moderate_message(self, message: dict) -> dict: + content = message.get("content") + if isinstance(content, str): + return {**message, "content": await self.moderation_check(text=content)} + if isinstance(content, list): + return { + **message, + "content": list( + await asyncio.gather( + *(self._moderate_content_part(part) for part in content) + ) + ), + } + return message + + async def _moderate_content_part(self, part: dict) -> dict: + if part.get("type") == "text" and isinstance(part.get("text"), str): + return {**part, "text": await self.moderation_check(text=part["text"])} + return part + + async def _moderate_input(self, input_: object) -> object: + if isinstance(input_, str): + return await self.moderation_check(text=input_) + if isinstance(input_, list): + return [ + await self.moderation_check(text=item) + if isinstance(item, str) + else item + for item in input_ + ] + return input_ async def async_post_call_streaming_hook( self, user_api_key_dict: UserAPIKeyAuth, response: str diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index 74e70f4aeb4..0edc4d2504b 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -54,6 +54,12 @@ If `db.useStackgresOperator` is used (not yet implemented): | `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` | | `pdb.labels` | Extra metadata labels to add to the PDB | `{}` | +| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` | +| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` | +| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` | +| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` | +| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` | + #### Example `proxy_config` ConfigMap from values (default): ``` @@ -94,6 +100,21 @@ data: type: Opaque ``` +#### Enterprise billable-request metering + +Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on: + +``` +kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key +``` + +``` +billingMetrics: + enabled: true +``` + +Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time. + ### Database Settings | Name | Description | Value | diff --git a/helm/litellm-helm/templates/_helpers.tpl b/helm/litellm-helm/templates/_helpers.tpl index 469d52c03a7..387bc3d5dc4 100644 --- a/helm/litellm-helm/templates/_helpers.tpl +++ b/helm/litellm-helm/templates/_helpers.tpl @@ -50,6 +50,53 @@ app.kubernetes.io/name: {{ include "litellm.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} {{- end }} +{{/* +Enterprise billable-request metering. The client certificate identifies the +deployment to LiteLLM's collector, so it is mounted read-only from an existing +Secret rather than passed through the environment. +*/}} +{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} +{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} + +{{- define "litellm.billingMetricsEnv" -}} +- name: LITELLM_BILLING_METRICS_ENDPOINT + value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_CERT + value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_KEY + value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} +{{- if .Values.billingMetrics.caSecretName }} +- name: LITELLM_BILLING_METRICS_CA_CERT + value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} +{{- end }} +{{- with .Values.billingMetrics.exportIntervalMs }} +- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS + value: {{ . | quote }} +{{- end }} +{{- end -}} + +{{- define "litellm.billingMetricsVolumes" -}} +- name: billing-metrics-mtls + secret: + secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} +{{- if .Values.billingMetrics.caSecretName }} +- name: billing-metrics-mtls-ca + secret: + secretName: {{ .Values.billingMetrics.caSecretName }} +{{- end }} +{{- end -}} + +{{- define "litellm.billingMetricsVolumeMounts" -}} +- name: billing-metrics-mtls + mountPath: {{ include "litellm.billingMetrics.certDir" . }} + readOnly: true +{{- if .Values.billingMetrics.caSecretName }} +- name: billing-metrics-mtls-ca + mountPath: {{ include "litellm.billingMetrics.caDir" . }} + readOnly: true +{{- end }} +{{- end -}} + {{/* Create the name of the service account to use */}} diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index b9cd1be06ec..32bfa4b2647 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -142,6 +142,9 @@ spec: {{- with .Values.extraEnvVars }} {{- toYaml . | nindent 12 }} {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsEnv" . | nindent 12 }} + {{- end }} {{- if .Values.migrationJob.enabled }} # Schema updates are owned by the dedicated migrations Job; skip # the proxy's startup `prisma db push` so N replicas don't race @@ -220,6 +223,9 @@ spec: - name: npm mountPath: /.npm {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} + {{- end }} {{- with .Values.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -252,6 +258,9 @@ spec: items: - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }} path: "config.yaml" + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} + {{- end }} {{- with .Values.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm-helm/tests/billing_metrics_tests.yaml b/helm/litellm-helm/tests/billing_metrics_tests.yaml new file mode 100644 index 00000000000..71803df378c --- /dev/null +++ b/helm/litellm-helm/tests/billing_metrics_tests.yaml @@ -0,0 +1,297 @@ +suite: test billingMetrics wiring on the proxy deployment +templates: + - deployment.yaml + - configmap-litellm.yaml + - migrations-job.yaml +tests: + - it: is off by default, adding no env, volume, or mount + template: deployment.yaml + asserts: + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: litellm-billing-metrics-mtls + - notContains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls + mountPath: /etc/litellm/billing-mtls + readOnly: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + + - it: renders the endpoint and the mounted cert paths when enabled + template: deployment.yaml + set: + billingMetrics: + enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CLIENT_CERT + value: /etc/litellm/billing-mtls/tls.crt + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CLIENT_KEY + value: /etc/litellm/billing-mtls/tls.key + + # The conventional Secret name is the default, so enabling the block is enough. + - it: mounts the default cert secret read-only alongside the config volume + template: deployment.yaml + set: + billingMetrics: + enabled: true + asserts: + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: litellm-billing-metrics-mtls + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls + mountPath: /etc/litellm/billing-mtls + readOnly: true + + - it: honours a secretName override + template: deployment.yaml + set: + billingMetrics: + enabled: true + secretName: my-billing-mtls + asserts: + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: my-billing-mtls + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: litellm-billing-metrics-mtls + + - it: honours an endpoint override + template: deployment.yaml + set: + billingMetrics: + enabled: true + endpoint: https://collector.internal:4318 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://collector.internal:4318 + + # The production collector presents a public web-PKI certificate, so the CA + # override must stay absent unless a private collector is configured. + - it: omits the CA env, volume, and mount when no caSecretName is set + template: deployment.yaml + set: + billingMetrics: + enabled: true + asserts: + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls-ca + secret: + secretName: billing-ca + - notContains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls-ca + mountPath: /etc/litellm/billing-mtls-ca + readOnly: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CA_CERT + value: /etc/litellm/billing-mtls-ca/ca.crt + + - it: mounts the CA secret when caSecretName is set + template: deployment.yaml + set: + billingMetrics: + enabled: true + caSecretName: billing-ca + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CA_CERT + value: /etc/litellm/billing-mtls-ca/ca.crt + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls-ca + secret: + secretName: billing-ca + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls-ca + mountPath: /etc/litellm/billing-mtls-ca + readOnly: true + + - it: passes the export interval through only when set + template: deployment.yaml + set: + billingMetrics: + enabled: true + exportIntervalMs: 5000 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS + value: "5000" + + - it: omits the export interval when unset + template: deployment.yaml + set: + billingMetrics: + enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS + value: "60000" + + # Kubernetes resolves duplicate env names last-wins, so the chart-owned billing + # entries must render after .Values.envVars or a user could silently redirect + # the metering export. The three billing entries are the last ones emitted here + # (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case). + - it: renders the billing endpoint after envVars so it cannot be shadowed + template: deployment.yaml + set: + migrationJob: + enabled: false + billingMetrics: + enabled: true + envVars: + LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://shadowed.example + - equal: + path: spec.template.spec.containers[0].env[-3] + value: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + - equal: + path: spec.template.spec.containers[0].env[-2].name + value: LITELLM_BILLING_METRICS_CLIENT_CERT + - equal: + path: spec.template.spec.containers[0].env[-1].name + value: LITELLM_BILLING_METRICS_CLIENT_KEY + + - it: keeps user-supplied volumes and mounts alongside the billing secret + template: deployment.yaml + set: + billingMetrics: + enabled: true + volumes: + - name: custom-callbacks + configMap: + name: my-callbacks + volumeMounts: + - name: custom-callbacks + mountPath: /app/callbacks + asserts: + - contains: + path: spec.template.spec.volumes + content: + name: custom-callbacks + configMap: + name: my-callbacks + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: litellm-billing-metrics-mtls + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: custom-callbacks + mountPath: /app/callbacks + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls + mountPath: /etc/litellm/billing-mtls + readOnly: true + + - it: still mounts the proxy config when enabled + template: deployment.yaml + set: + billingMetrics: + enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: litellm-config + mountPath: /etc/litellm/config.yaml + subPath: config.yaml + + # Only the proxy serves billable traffic. The migrations Job must never mount + # the client certificate, and it renders its own env and volumes, so nothing + # stops a future edit from wiring the billing include into it by mistake. + - it: does not touch the migrations job when enabled + template: migrations-job.yaml + set: + billingMetrics: + enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + - notExists: + path: spec.template.spec.containers[0].volumeMounts + - notExists: + path: spec.template.spec.volumes + + - it: fails loudly when enabled with an emptied secretName + template: deployment.yaml + set: + billingMetrics: + enabled: true + secretName: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) + + - it: fails loudly when enabled without an endpoint + template: deployment.yaml + set: + billingMetrics: + enabled: true + endpoint: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index d3821a547e5..0529e74d6e4 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -139,6 +139,20 @@ masterkeySecretName: "" # if set, use this secret key for the master key; otherwise, use the default key masterkeySecretKey: "" +# Optional: enterprise billable-request metering. When enabled, the proxy counts +# successful requests to inference, MCP, and A2A endpoints and pushes them to +# LiteLLM's collector over mutual TLS. Requires an enterprise license. +# The client certificate identifies the deployment, so it is mounted read-only +# from an existing Secret and never passed through the environment. +billingMetrics: + enabled: false + endpoint: https://telemetry.litellm.ai # collector to push the counter to + secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key + # Only for private or test collectors whose server certificate is not on the + # public web PKI. The production collector needs no CA override. + caSecretName: "" # existing Secret holding ca.crt + exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + proxyConfigMap: # when true, creates a new configmap create: true diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 7c281aa158b..043a0afc173 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -34,6 +34,57 @@ app.kubernetes.io/managed-by: {{ .Release.Service }} helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }} {{- end -}} +{{/* +Enterprise billable-request metering. Wired into gateway and backend, not the +migrations job. The gateway serves nearly all billable traffic, but the backend +keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a +SpendLogs row, so metering only the gateway would silently drop that traffic. +The client certificate identifies the deployment to LiteLLM's collector, so it is +mounted read-only from an existing Secret rather than passed through the +environment. +*/}} +{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}} +{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}} + +{{- define "litellm.billingMetricsEnv" -}} +- name: LITELLM_BILLING_METRICS_ENDPOINT + value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_CERT + value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }} +- name: LITELLM_BILLING_METRICS_CLIENT_KEY + value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }} +{{- if .Values.billingMetrics.caSecretName }} +- name: LITELLM_BILLING_METRICS_CA_CERT + value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }} +{{- end }} +{{- with .Values.billingMetrics.exportIntervalMs }} +- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS + value: {{ . | quote }} +{{- end }} +{{- end -}} + +{{- define "litellm.billingMetricsVolumes" -}} +- name: billing-metrics-mtls + secret: + secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }} +{{- if .Values.billingMetrics.caSecretName }} +- name: billing-metrics-mtls-ca + secret: + secretName: {{ .Values.billingMetrics.caSecretName }} +{{- end }} +{{- end -}} + +{{- define "litellm.billingMetricsVolumeMounts" -}} +- name: billing-metrics-mtls + mountPath: {{ include "litellm.billingMetrics.certDir" . }} + readOnly: true +{{- if .Values.billingMetrics.caSecretName }} +- name: billing-metrics-mtls-ca + mountPath: {{ include "litellm.billingMetrics.caDir" . }} + readOnly: true +{{- end }} +{{- end -}} + {{/* Per-component selector labels — used in both Service selectors and Deployment matchLabels. */}} diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 8b4552bf302..9d056167fe1 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -44,14 +44,20 @@ spec: - name: CONFIG_FILE_PATH value: /app/config/config.yaml {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsEnv" . | nindent 12 }} + {{- end }} {{- include "litellm.envFrom" .Values.backend | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.backend.volumeMounts }} + {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} + {{- end }} {{- with .Values.backend.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -66,13 +72,16 @@ spec: {{- end }} resources: {{- toYaml .Values.backend.resources | nindent 12 }} - {{- if or .Values.gateway.config.create .Values.backend.volumes }} + {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: name: {{ include "litellm.gateway.fullname" . }}-config {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} + {{- end }} {{- with .Values.backend.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index bd491b69e0f..4c80d784156 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -46,14 +46,20 @@ spec: - name: NUM_WORKERS value: {{ .Values.gateway.numWorkers | quote }} {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsEnv" . | nindent 12 }} + {{- end }} {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} + {{- end }} {{- with .Values.gateway.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} @@ -68,13 +74,16 @@ spec: {{- end }} resources: {{- toYaml .Values.gateway.resources | nindent 12 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumes }} + {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: name: {{ include "litellm.gateway.fullname" . }}-config {{- end }} + {{- if .Values.billingMetrics.enabled }} + {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} + {{- end }} {{- with .Values.gateway.volumes }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/helm/litellm/tests/billing_metrics_tests.yaml b/helm/litellm/tests/billing_metrics_tests.yaml new file mode 100644 index 00000000000..ceba0bd1430 --- /dev/null +++ b/helm/litellm/tests/billing_metrics_tests.yaml @@ -0,0 +1,249 @@ +suite: test billingMetrics wiring on gateway and backend +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - migrations-job.yaml +values: + - ./values/required.yaml +tests: + - it: is off by default, adding no env, volume, or mount + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: billing-mtls + - equal: + path: spec.template.spec.containers[0].volumeMounts + value: + - name: gateway-config + mountPath: /app/config/config.yaml + subPath: config.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + + - it: renders the endpoint and the mounted cert paths when enabled + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CLIENT_CERT + value: /etc/litellm/billing-mtls/tls.crt + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CLIENT_KEY + value: /etc/litellm/billing-mtls/tls.key + + - it: mounts the cert secret read-only alongside the config volume + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + asserts: + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: billing-mtls + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls + mountPath: /etc/litellm/billing-mtls + readOnly: true + + # The production collector presents a public web-PKI certificate, so the CA + # override must stay absent unless a private collector is configured. + - it: omits the CA env, volume, and mount when no caSecretName is set + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + asserts: + - notContains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls-ca + secret: + secretName: billing-ca + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CA_CERT + value: /etc/litellm/billing-mtls-ca/ca.crt + + - it: mounts the CA secret when caSecretName is set + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + caSecretName: billing-ca + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_CA_CERT + value: /etc/litellm/billing-mtls-ca/ca.crt + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls-ca + secret: + secretName: billing-ca + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls-ca + mountPath: /etc/litellm/billing-mtls-ca + readOnly: true + + - it: passes the export interval through only when set + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + exportIntervalMs: 5000 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS + value: "5000" + + - it: keeps user-supplied gateway volumes alongside the billing secret + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + gateway.volumes: + - name: custom-callbacks + configMap: + name: my-callbacks + gateway.volumeMounts: + - name: custom-callbacks + mountPath: /app/callbacks + asserts: + - contains: + path: spec.template.spec.volumes + content: + name: custom-callbacks + configMap: + name: my-callbacks + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: billing-mtls + + # The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp), + # which writes a SpendLogs row, so it must meter too or that traffic is lost. + - it: meters the backend as well, since it serves the MCP transport + template: backend/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: billing-metrics-mtls + mountPath: /etc/litellm/billing-mtls + readOnly: true + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: billing-mtls + + - it: leaves the backend alone when metering is off + template: backend/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + + # The migrations job runs prisma and serves no traffic; it must never receive + # the client key. + - it: never mounts the billing cert on the migrations job + template: migrations-job.yaml + set: + billingMetrics: + enabled: true + secretName: billing-mtls + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BILLING_METRICS_ENDPOINT + value: https://telemetry.litellm.ai + - isNull: + path: spec.template.spec.volumes + + # The conventional Secret name is the default, so enabling metering needs no + # secretName at all; the guard below only fires on an explicitly blanked one. + - it: uses the conventional secret name by default + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + asserts: + - contains: + path: spec.template.spec.volumes + content: + name: billing-metrics-mtls + secret: + secretName: litellm-billing-metrics-mtls + + - it: fails loudly when the secretName is explicitly blanked + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + secretName: "" + asserts: + - failedTemplate: + errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key) + + - it: fails loudly when enabled without an endpoint + template: gateway/deployment.yaml + set: + billingMetrics: + enabled: true + endpoint: "" + secretName: billing-mtls + asserts: + - failedTemplate: + errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index a8f2d39663e..74d02a25b7a 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -73,6 +73,25 @@ masterKey: secretName: litellm-master-key-secret # name of a Secret containing the master key secretKey: master-key +# Optional: enterprise billable-request metering. When enabled, the gateway and +# backend count successful requests to inference, MCP, and A2A endpoints and push +# them to LiteLLM's collector over mutual TLS. Both components serve billable +# routes: the backend keeps the named-server MCP transport. Requires an +# enterprise license. The client certificate identifies the deployment, so it is +# mounted read-only from an existing Secret and never passed through the env. +billingMetrics: + enabled: false + endpoint: https://telemetry.litellm.ai # collector to push the counter to + # An existing Secret holding the client certificate under tls.crt and its key + # under tls.key, usually created from the onboarding artifact. The default is + # the conventional name, so the common path is to create that Secret and set + # enabled: true. Override only if yours is named differently. + secretName: litellm-billing-metrics-mtls + # Only for private or test collectors whose server certificate is not on the + # public web PKI. The production collector needs no CA override. + caSecretName: "" # existing Secret holding ca.crt + exportIntervalMs: "" # push cadence; the proxy defaults to 60000 + # External Postgres connection. database: writer: diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index c860f8e540d..b17e055c7ea 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -85,6 +85,22 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() +def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]: + """ + The caching handler is stored on the Logging object + (``logging_obj._llm_caching_handler``), so keeping ``litellm_logging_obj`` + inside ``request_kwargs`` closes a reference cycle + (Logging -> LLMCachingHandler -> kwargs -> Logging) that keeps the full + request payload (messages included) alive until a generational GC pass + instead of being freed by refcount when the request ends. Nothing in the + caching layer reads the logging object from these kwargs; cache-key + generation ignores litellm-internal params. + """ + if "litellm_logging_obj" not in request_kwargs: + return request_kwargs + return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"} + + def _is_chat_completion_cached_dict(cached_result: dict) -> bool: cached_id = cached_result.get("id") if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"): @@ -118,7 +134,7 @@ class LLMCachingHandler: self.async_streaming_chunks: List[ModelResponse] = [] self.sync_streaming_chunks: List[ModelResponse] = [] - self.request_kwargs = request_kwargs + self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs) self.preset_cache_key: Optional[str] = None self.original_function = original_function self.start_time = start_time @@ -297,7 +313,7 @@ class LLMCachingHandler: new_kwargs.pop("metadata", None) if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) - self.request_kwargs = new_kwargs + self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) print_verbose("Checking Sync Cache") cached_result = litellm.cache.get_cache(**new_kwargs) if cached_result is not None: @@ -693,7 +709,7 @@ class LLMCachingHandler: new_kwargs.pop("metadata", None) if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) - self.request_kwargs = new_kwargs + self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) cached_result: Optional[Any] = None if call_type == CallTypes.aembedding.value: if isinstance(new_kwargs["input"], str): diff --git a/litellm/constants.py b/litellm/constants.py index 715d57e594d..8e0a5cfe50f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1496,6 +1496,7 @@ MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20)) MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)) DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7)) LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16)) +MINIMUM_CUSTOM_KEY_LENGTH = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16)) SECRET_MANAGER_REFRESH_INTERVAL = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400)) LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "default_internal_user_params", diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index ae4f46c38bd..3a1bd72e1a5 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union from litellm.secret_managers.main import get_secret_bool if TYPE_CHECKING: - from ddtrace.tracer import Tracer as DD_TRACER + from ddtrace.trace import Tracer as DD_TRACER else: DD_TRACER = Any diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 461ab62b815..268dcc3df78 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -925,7 +925,6 @@ class Logging(LiteLLMLoggingBaseClass): def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API - litellm.error_logs["PRE_CALL"] = locals() try: self._pre_call( input=input, @@ -1135,7 +1134,6 @@ class Logging(LiteLLMLoggingBaseClass): def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received - litellm.error_logs["POST_CALL"] = locals() if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: @@ -3074,7 +3072,7 @@ class Logging(LiteLLMLoggingBaseClass): def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) - return list(set(dynamic_success_callbacks + global_callbacks)) + return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) def _remove_internal_litellm_callbacks(self, callbacks: List) -> List: """ diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 6e8429839ad..43181e7f5ff 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -77,6 +77,22 @@ def _redact_streaming_response(streaming_response): streaming_response.reasoning = None +def _redact_tool_calls(tool_calls) -> None: + """Redact tool call arguments (assistant tool calls carry prompt-derived data).""" + if not tool_calls: + return + for tool_call in tool_calls: + function = getattr(tool_call, "function", None) + if function is not None and hasattr(function, "arguments"): + function.arguments = "redacted-by-litellm" + + +def _redact_function_call(function_call) -> None: + """Redact legacy assistant function_call arguments.""" + if function_call is not None and hasattr(function_call, "arguments"): + function_call.arguments = "redacted-by-litellm" + + def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" if isinstance(choice, litellm.Choices): @@ -85,12 +101,16 @@ def _redact_choice_content(choice): choice.message.reasoning_content = "redacted-by-litellm" if hasattr(choice.message, "thinking_blocks"): choice.message.thinking_blocks = None + _redact_tool_calls(getattr(choice.message, "tool_calls", None)) + _redact_function_call(getattr(choice.message, "function_call", None)) elif isinstance(choice, litellm.utils.StreamingChoices): choice.delta.content = "redacted-by-litellm" if hasattr(choice.delta, "reasoning_content"): choice.delta.reasoning_content = "redacted-by-litellm" if hasattr(choice.delta, "thinking_blocks"): choice.delta.thinking_blocks = None + _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) + _redact_function_call(getattr(choice.delta, "function_call", None)) def _redact_responses_api_output(output_items): @@ -111,6 +131,9 @@ def _redact_responses_api_output(output_items): if hasattr(summary_item, "text"): summary_item.text = "redacted-by-litellm" + if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): + output_item.arguments = "redacted-by-litellm" + def _redact_responses_api_output_dict(output_items, redacted_str: str): """Helper to redact ResponsesAPIResponse output items in dict form.""" @@ -131,6 +154,9 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): if isinstance(summary_item, dict) and "text" in summary_item: summary_item["text"] = redacted_str + if output_item.get("type") == "function_call" and "arguments" in output_item: + output_item["arguments"] = redacted_str + def _redact_standard_logging_object(model_call_details: dict): """Redact messages and response inside standard_logging_object if present.""" @@ -162,6 +188,19 @@ def _redact_standard_logging_object(model_call_details: dict): standard_logging_object["response"] = {"text": redacted_str} +def _redact_tool_calls_dict(message: dict, redacted_str: str) -> None: + """Redact tool call / function_call arguments in a dict-form message or delta.""" + tool_calls = message.get("tool_calls") + if isinstance(tool_calls, list): + for tool_call in tool_calls: + if isinstance(tool_call, dict) and isinstance(tool_call.get("function"), dict): + tool_call["function"]["arguments"] = redacted_str + + function_call = message.get("function_call") + if isinstance(function_call, dict) and "arguments" in function_call: + function_call["arguments"] = redacted_str + + def _redact_model_response_dict_choices(choices, redacted_str: str): for choice in choices: if isinstance(choice, dict): @@ -173,6 +212,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): choice["message"]["thinking_blocks"] = None if "audio" in choice["message"]: choice["message"]["audio"] = None + _redact_tool_calls_dict(choice["message"], redacted_str) elif "delta" in choice and isinstance(choice["delta"], dict): choice["delta"]["content"] = redacted_str if "reasoning_content" in choice["delta"]: @@ -181,6 +221,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): choice["delta"]["thinking_blocks"] = None if "audio" in choice["delta"]: choice["delta"]["audio"] = None + _redact_tool_calls_dict(choice["delta"], redacted_str) else: _redact_choice_content(choice) diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index b526068589d..455d0f00c35 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -9,6 +9,8 @@ secrets from strings without depending on the logging-configuration module. import re from typing import List +from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH + _REDACTED = "REDACTED" @@ -30,7 +32,7 @@ def _build_secret_patterns() -> "re.Pattern[str]": # Basic auth headers r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", # OpenAI / Anthropic sk- prefixed keys - r"sk-[A-Za-z0-9\-_]{20,}", + rf"sk-[A-Za-z0-9\-_]{{{MINIMUM_CUSTOM_KEY_LENGTH - len('sk-')},}}", # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr) r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}", # x-api-key / api-key header values (handles 'key': 'value' dict repr) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index cd75eed2e6e..4b6617fbeac 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1403,11 +1403,6 @@ class LiteLLMAnthropicMessagesAdapter: assert isinstance(thinking, str) assert isinstance(signature, str) - if thinking and signature: - raise ValueError( - "Both `thinking` and `signature` in a single streaming chunk isn't supported." - ) - return "thinking", ChatCompletionThinkingBlock( type="thinking", thinking=thinking, signature=signature ) @@ -1463,17 +1458,14 @@ class LiteLLMAnthropicMessagesAdapter: if choice.delta.reasoning_content is not None: reasoning_content += choice.delta.reasoning_content - if reasoning_content and reasoning_signature: - raise ValueError("Both `reasoning` and `signature` in a single streaming chunk isn't supported.") - if partial_json is not None: return "input_json_delta", ContentJsonBlockDelta(type="input_json_delta", partial_json=partial_json) - elif reasoning_content: - return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) elif reasoning_signature: return "signature_delta", ContentThinkingSignatureBlockDelta( type="signature_delta", signature=reasoning_signature ) + elif reasoning_content: + return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) else: return "text_delta", ContentTextBlockDelta(type="text_delta", text=text) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 96d0ad48b79..b47fc50e196 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1879,6 +1879,7 @@ class BaseLLMHTTPHandler: litellm_params: GenericLiteLLMParams, api_key: Optional[str], model: str, + timeout: Optional[Union[float, httpx.Timeout]] = None, ) -> httpx.Response: max_attempts = max(provider_config.max_retry_on_anthropic_messages_http_error, 1) litellm_params_dict = dict(litellm_params) @@ -1891,6 +1892,7 @@ class BaseLLMHTTPHandler: data=signed_json_body or json.dumps(request_body), stream=stream or False, logging_obj=logging_obj, + timeout=timeout, ) response.raise_for_status() return response @@ -1925,6 +1927,32 @@ class BaseLLMHTTPHandler: raise RuntimeError("unreachable: anthropic messages HTTP retry loop exited without return") + @staticmethod + def _resolve_anthropic_messages_timeout( + litellm_params: GenericLiteLLMParams, + stream: bool, + custom_llm_provider: str, + ) -> Optional[Union[float, httpx.Timeout]]: + from litellm.litellm_core_utils.completion_timeout import CompletionTimeout + from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, + ) + from litellm.utils import supports_httpx_timeout + + stream_timeout = litellm_params.get("stream_timeout") if stream else None + model_timeout = stream_timeout if stream_timeout is not None else litellm_params.get("timeout") + request_timeout = litellm_params.get("request_timeout") + global_timeout = get_configured_request_timeout() + if model_timeout is None and request_timeout is None and global_timeout is None: + return None + return CompletionTimeout.resolve( + model_timeout, + {"request_timeout": request_timeout}, + custom_llm_provider, + global_timeout=global_timeout, + supports_httpx_timeout=supports_httpx_timeout, + ) + async def async_anthropic_messages_handler( self, model: str, @@ -2075,6 +2103,11 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, api_key=api_key, model=model, + timeout=self._resolve_anthropic_messages_timeout( + litellm_params=litellm_params, + stream=stream or False, + custom_llm_provider=custom_llm_provider, + ), ) # used for logging + cost tracking diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 093dffccac0..d90703d1544 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -247,6 +247,8 @@ class OpenAIResponsesHandler(BaseTranslation): """ Merge remapped guardrailed tools with original tools that were not sent to the guardrail (e.g. web_search, web_search_preview), preserving order. + Tools a guardrail appended (``remapped`` longer than ``original_tools``) + have no original slot and are kept so an injected tool is not dropped. """ if not original_tools: return remapped @@ -262,6 +264,8 @@ class OpenAIResponsesHandler(BaseTranslation): if j < len(remapped): result.append(remapped[j]) j += 1 + # Keep guardrail-appended tools that matched no original slot above. + result.extend(remapped[j:]) return result def _apply_guardrailed_tools_to_data( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 77ca423b866..dedb9bbf40a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -44261,6 +44261,90 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock_mantle/openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.875e-06, + "cache_read_input_token_cost": 5.5e-07, + "output_cost_per_token": 3.3e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/openai.gpt-5.6-terra": { + "input_cost_per_token": 2.75e-06, + "cache_creation_input_token_cost": 3.4375e-06, + "cache_read_input_token_cost": 2.75e-07, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/openai.gpt-5.6-luna": { + "input_cost_per_token": 1.1e-06, + "cache_creation_input_token_cost": 1.375e-06, + "cache_read_input_token_cost": 1.1e-07, + "output_cost_per_token": 6.6e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, @@ -44373,6 +44457,7 @@ "supports_vision": true }, "bedrock_mantle/xai.grok-4.3": { + "use_openai_responses_path": true, "input_cost_per_token": 1.25e-06, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2e-07, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e6e265abb61..8b1d00c2855 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -186,6 +186,61 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = ( ) +def _blank_to_none(value: str | None) -> str | None: + """Collapse an absent, empty, or whitespace-only string to ``None``. + + OAuth endpoint fields are consumed by truthiness-based merges (``row or discovered``) and by the + corroboration gate. A whitespace-only value is truthy to ``or`` but is not a usable endpoint, so + without this the merge would keep the blank value for redirects while the gate treats it as + unpinned and backfills the other fields, yielding a broken half-discovered config. Normalizing + the pinned fields once, at each build entry point, gives every downstream consumer a single + notion of "blank" so those code paths cannot disagree. + """ + if not isinstance(value, str): + return None + return value.strip() or None + + +def _normalized_authorize_endpoint(url: str) -> str: + """Compare authorize endpoints on scheme, host, and path only. The default port is elided and + the host is lowercased so ``https://IDP.example.com:443/authorize/`` and + ``https://idp.example.com/authorize`` are the same identity; query and trailing slash are not.""" + parsed = urlparse(url) + scheme = parsed.scheme.lower() + host = (parsed.hostname or "").lower() + default_port = {"https": 443, "http": 80}.get(scheme) + try: + port = parsed.port + except ValueError: + port = None + authority = host if port is None or port == default_port else f"{host}:{port}" + return f"{scheme}://{authority}{parsed.path.rstrip('/')}" + + +def _endpoints_corroborate_authorization_url( + source_authorization_url: str | None, + trusted_authorization_url: str | None, +) -> bool: + """Whether a source's ``token_url``/``registration_url`` may be paired with a trusted authorize + endpoint. This is the single trust rule for adopting OAuth endpoints from any non-manual source. + + Discovery is rooted at the MCP resource (RFC 9728), so a compromised upstream can advertise an + attacker-run authorization server. When ``authorization_url`` is admin-pinned, pairing it with a + ``token_url`` from a different source is the RFC 9700 authorization-server mix-up: the user signs + in at the trusted authorize endpoint while the gateway redeems the code, with the stored client + secret and PKCE verifier, at the attacker's token endpoint. Endpoints are trustworthy together + only when they share an authorization server, so a source's endpoints are adopted only when the + same source advertised an ``authorization_endpoint`` matching the pinned value. With no pinned + value (``trusted_authorization_url is None``) there is nothing to protect: the authorize endpoint + comes from the same source as the token endpoint, so they corroborate each other by construction. + """ + if not (trusted_authorization_url and trusted_authorization_url.strip()): + return True + return bool(source_authorization_url) and _normalized_authorize_endpoint( + source_authorization_url + ) == _normalized_authorize_endpoint(trusted_authorization_url) + + def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_server: MCPServer | None) -> None: """Keep the last known good OAuth endpoints when a rebuild's re-discovery comes back empty. @@ -193,26 +248,82 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv during re-discovery downgrades a working server (``authorization_url`` set) to a broken one (``None``, /authorize 400s) with no configuration change. Mirrors the ``short_prefix`` carry-forward. Skipped when the server's ``url`` or ``auth_type`` changed, since the previous - endpoints may then belong to a different upstream. ``registration_url`` IS carried here even - though ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only - restores the same in-memory value the previous build already ran with, while persisting it - would flip ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for - dcr_bridge servers that never had one configured. + endpoints may then belong to a different upstream. ``registration_url`` IS carried even though + ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only restores + the same in-memory value the previous build already ran with, while persisting it would flip + ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for dcr_bridge + servers that never had one configured. + + Carry-forward is a non-manual endpoint source, so the same trust rule as discovery applies: the + previous ``token_url``/``registration_url``/``scopes`` are carried only when the previous + ``authorization_url`` corroborates the authorize endpoint this build will use, i.e. when the + incoming build has no pinned authorize endpoint (``None`` -> we adopt the previous one too, a + consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different + server must not keep serving the old server's token endpoint or granted scopes. """ if previous_server is None: return if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type: return + may_carry = _endpoints_corroborate_authorization_url( + previous_server.authorization_url, new_server.authorization_url + ) if new_server.authorization_url is None and previous_server.authorization_url: new_server.authorization_url = previous_server.authorization_url - if new_server.token_url is None and previous_server.token_url: + if may_carry and new_server.token_url is None and previous_server.token_url: new_server.token_url = previous_server.token_url - if new_server.registration_url is None and previous_server.registration_url: + if may_carry and new_server.registration_url is None and previous_server.registration_url: new_server.registration_url = previous_server.registration_url - if not new_server.scopes and previous_server.scopes: + if may_carry and not new_server.scopes and previous_server.scopes: new_server.scopes = previous_server.scopes +def _restrict_discovery_to_corroborated_authorization_server( + metadata: MCPOAuthMetadata | None, + manual_authorization_url: str | None, + server_identifier: str, + is_dcr_bridge: bool, +) -> MCPOAuthMetadata | None: + """Reject discovered token/registration endpoints a manually pinned authorize endpoint cannot + vouch for (the RFC 9700 authorization-server mix-up). + + Discovery is rooted at the MCP resource, so a compromised upstream can advertise an attacker + ``token_endpoint``: with ``authorization_url`` admin-pinned but ``token_url`` blank, the merge + would pair the trusted authorize endpoint with that attacker token endpoint, and the gateway would + post the authorization code and client secret there. So the discovered ``token_url`` and + ``registration_url`` are kept only if the document corroborates the pin (its + ``authorization_endpoint`` matches). ``scopes`` are deliberately NOT gated here: per the MCP + authorization spec Scope Selection Strategy and RFC 9700 §2.3, the scopes a client requests are + resource-driven (the WWW-Authenticate challenge or the RFC 9728 protected-resource + ``scopes_supported``), and scope inflation by a compromised resource is bounded by the + authorization server and user consent (RFC 6749 §3.3), not by the client second-guessing the + request. With no pin there is no trust anchor to protect, so discovery is returned as-is. + """ + if metadata is None or not (manual_authorization_url and manual_authorization_url.strip()): + return metadata + if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url): + return metadata + if not metadata.token_url and not metadata.registration_url: + return metadata + bridge_note = ( + " The discovered registration_url is rejected with it, so this dcr_bridge server stays on the" + " short-circuit registration arm." + if is_dcr_bridge and metadata.registration_url + else "" + ) + verbose_logger.warning( + "MCP OAuth discovery for server %s advertised authorization_endpoint %s, which does not match the " + "manually configured authorization_url %s; rejecting the discovered token_url/registration_url so " + "authorization codes and client credentials only follow the configured authorization server. " + "Configure Token URL manually if the mismatch is intentional.%s", + server_identifier, + _normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "", + _normalized_authorize_endpoint(manual_authorization_url), + bridge_note, + ) + return metadata.model_copy(update={"token_url": None, "registration_url": None}) + + def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: """Drop a cached entry after the user stores or clears their env var values so the next request reads the fresh value instead of a stale one.""" @@ -1026,12 +1137,15 @@ class MCPServerManager: ) auth_type = server_config.get("auth_type", None) + manual_authorization_url = _blank_to_none(server_config.get("authorization_url")) + manual_token_url = _blank_to_none(server_config.get("token_url")) + manual_registration_url = _blank_to_none(server_config.get("registration_url")) if server_url and ( auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES or self._obo_needs_endpoint_discovery( auth_type, server_config.get("token_exchange_endpoint"), - server_config.get("token_url"), + manual_token_url, ) ): mcp_oauth_metadata = await self._descovery_metadata( @@ -1041,20 +1155,29 @@ class MCPServerManager: else: mcp_oauth_metadata = None + gated_oauth_metadata = ( + _restrict_discovery_to_corroborated_authorization_server( + mcp_oauth_metadata, + manual_authorization_url, + server_name or server_id, + bool(server_config.get("dcr_bridge")), + ) + if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + else mcp_oauth_metadata + ) + # Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so # an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the # entra_obo fail-closed scope precondition and POSTs an empty scope to the IdP. resolved_scopes = self._extract_scopes(server_config.get("scopes")) or ( - mcp_oauth_metadata.scopes if mcp_oauth_metadata else None + gated_oauth_metadata.scopes if gated_oauth_metadata else None ) - resolved_authorization_url = server_config.get("authorization_url") or ( - mcp_oauth_metadata.authorization_url if mcp_oauth_metadata else None + resolved_authorization_url = manual_authorization_url or ( + gated_oauth_metadata.authorization_url if gated_oauth_metadata else None ) - resolved_token_url = server_config.get("token_url") or ( - mcp_oauth_metadata.token_url if mcp_oauth_metadata else None - ) - resolved_registration_url = server_config.get("registration_url") or ( - mcp_oauth_metadata.registration_url if mcp_oauth_metadata else None + resolved_token_url = manual_token_url or (gated_oauth_metadata.token_url if gated_oauth_metadata else None) + resolved_registration_url = manual_registration_url or ( + gated_oauth_metadata.registration_url if gated_oauth_metadata else None ) config_oauth2_flow = server_config.get("oauth2_flow", None) @@ -1447,13 +1570,17 @@ class MCPServerManager: auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url + manual_authorization_url = _blank_to_none(mcp_server.authorization_url) + manual_token_url = _blank_to_none(mcp_server.token_url) + manual_registration_url = _blank_to_none(mcp_server.registration_url) + has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes) needs_discovery = bool(server_url) and ( - (auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not mcp_server.authorization_url) + (auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not has_all_upstream_oauth_fields) or self._obo_needs_endpoint_discovery( auth_type, mcp_server.token_exchange_endpoint or (credentials_dict.get("token_exchange_endpoint") if credentials_dict else None), - mcp_server.token_url, + manual_token_url, ) ) mcp_oauth_metadata = ( @@ -1467,12 +1594,22 @@ class MCPServerManager: if needs_discovery and mcp_oauth_metadata is None: verbose_logger.warning( "MCP OAuth discovery yielded no metadata for server %s (%s); " - "OAuth endpoints stay unresolved until a rebuild succeeds", + "OAuth endpoints/scopes stay unresolved until a rebuild succeeds", mcp_server.server_id, server_url, ) + gated_oauth_metadata = ( + _restrict_discovery_to_corroborated_authorization_server( + mcp_oauth_metadata, + manual_authorization_url, + mcp_server.server_id, + bool(getattr(mcp_server, "dcr_bridge", None)), + ) + if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + else mcp_oauth_metadata + ) - resolved_scopes = scopes or (mcp_oauth_metadata.scopes if mcp_oauth_metadata else None) + resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None) new_server = MCPServer( server_id=mcp_server.server_id, @@ -1492,9 +1629,9 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)), scopes=resolved_scopes, - authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), - token_url=mcp_server.token_url or getattr(mcp_oauth_metadata, "token_url", None), - registration_url=mcp_server.registration_url or getattr(mcp_oauth_metadata, "registration_url", None), + authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None), + token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None), + registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None), token_endpoint_auth_method=( credentials_dict.get("token_endpoint_auth_method") if credentials_dict else None ), @@ -1545,16 +1682,16 @@ class MCPServerManager: await self._persist_discovered_obo_token_url( server_id=mcp_server.server_id, auth_type=auth_type, - existing_token_url=mcp_server.token_url, + existing_token_url=manual_token_url, discovered_token_url=new_server.token_url, ) await self._persist_discovered_oauth_endpoints( server_id=mcp_server.server_id, auth_type=auth_type, - existing_authorization_url=mcp_server.authorization_url, - existing_token_url=mcp_server.token_url, + existing_authorization_url=manual_authorization_url, + existing_token_url=manual_token_url, existing_scopes=scopes, - metadata=mcp_oauth_metadata, + metadata=gated_oauth_metadata, ) return new_server diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 33f0641b732..43fe3999291 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -175,16 +175,18 @@ mcp_oauth2_token_cache = MCPOAuth2TokenCache() def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int: """Compute Redis TTL for a per-user token. - Uses server.token_storage_ttl_seconds when configured; otherwise derives - TTL from expires_in minus the expiry buffer; falls back to the default TTL. + Uses server.token_storage_ttl_seconds when configured, capped at the token's + remaining lifetime (expires_in minus the expiry buffer) so a cached entry never + outlives the token itself; otherwise derives TTL from expires_in minus the + expiry buffer; falls back to the default TTL. """ + lifetime_bound = expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS if expires_in is not None else None if server.token_storage_ttl_seconds is not None: - return max(server.token_storage_ttl_seconds, 1) - if expires_in is not None: - return max( - expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, - 1, - ) + if lifetime_bound is None: + return max(server.token_storage_ttl_seconds, 1) + return max(min(server.token_storage_ttl_seconds, lifetime_bound), 1) + if lifetime_bound is not None: + return max(lifetime_bound, 1) return MCP_PER_USER_TOKEN_DEFAULT_TTL diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 27cdc483d4a..67d935e34e3 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -7613,6 +7613,18 @@ ], "title": "Messages" }, + "metadata": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Metadata" + }, "text": { "title": "Text", "type": "string" diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 893e09ece6e..a610e44e69c 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -10,7 +10,7 @@ from fastapi import HTTPException, Request, status import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger -from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS +from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * @@ -1533,4 +1533,6 @@ def get_model_from_request( def abbreviate_api_key(api_key: str) -> str: + if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH: + return "sk-..." return f"sk-...{api_key[-4:]}" diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index ce13a906a36..17041751d15 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -471,6 +471,109 @@ The token minted by `lite login` is a short-lived, per-session agent credential, The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. +### Route Every Claude Code Session Through the Proxy + +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. + +Two things need to already be true: you've run `lite login`, since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. + +```bash +lite login +litellm --config litellm/proxy/dev_config.yaml & +lite up +``` + +`lite up` runs in the foreground and blocks. Press Ctrl-C to stop it, which restores the original settings file and exits. If the process is ever killed uncleanly instead -- `kill -9`, a crash -- the settings file is left patched, and `lite down` is the manual recovery path: run it at any later point to restore from the same backup. + +This is a one-time file patch and restore, not a live traffic interceptor. A Claude Code session already running before `lite up` started keeps whatever `ANTHROPIC_BASE_URL` and token it loaded at its own startup, and a session still running when `lite up` stops keeps routing through the proxy until it exits; only sessions *started* while the patch is in effect are affected, and only *new* sessions after a restore go back to Anthropic directly. + +Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI. + +### QA Complexity-Based Auto-Routing Against Your Real Proxy + +`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session. + +#### Install the CLI + +`lite autoroute up` builds and runs a throwaway litellm proxy locally, so unlike the rest of this CLI it needs the proxy server runtime, not just the thin `litellm[cli]` client. Install `litellm[proxy]` (which ships the `lite` command too) with a single curl command -- no existing Python tooling required, `uv` is bootstrapped automatically if missing: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh +``` + +To QA an unreleased branch or commit instead of the latest PyPI release, set `LITELLM_CLI_REF`: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install.sh | \ + LITELLM_CLI_REF= sh +``` + +The thin `scripts/install-cli.sh` installs only `litellm[cli]`, which is enough for `lite login`, `lite claude`, and `lite up`, but not for `lite autoroute up`; running it against a `litellm[cli]` install fails fast with a message telling you to install the proxy runtime. + +Point the CLI at your real proxy and key before running any `lite model-groups` or `lite autoroute` command -- like every other command in this CLI, they read `LITELLM_PROXY_URL`/`LITELLM_PROXY_API_KEY` (or `--base-url`/`--api-key`), no `lite login` required: + +```bash +export LITELLM_PROXY_URL=http://localhost:4000 +export LITELLM_PROXY_API_KEY=sk-... +``` + +#### List Your Accessible Model Groups + +```bash +lite model-groups list [--format table|json] +``` + +Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. This is also what `lite autoroute configure` uses internally to discover what it can offer you. + +#### Configure the Auto-Router + +```bash +lite autoroute configure +``` + +An interactive wizard. It runs the same model-group discovery as above, splits the results into chat-capable and embedding-capable pools, and asks you to assign one or more models from the chat pool to each of the four complexity tiers -- SIMPLE, MEDIUM, COMPLEX, REASONING. Each tier's picker is a type-to-filter fuzzy search (fzf-style) rather than a scrollable numbered list, so it stays usable even with hundreds of model groups: type a substring to narrow the list, tab to toggle a model into the selection, enter to confirm (assigning more than one model to a tier is exactly when this matters -- complexity_router picks randomly among a tier's pool per request, and adaptive mode specifically depends on having more than one candidate to choose from). From there it optionally offers: classifying prompt complexity with an LLM (again picked from your discovered pool) instead of the free built-in heuristic scorer, semantic keyword matching for tier assignment (needs an embedding model from the pool), and adaptive (bandit-based) selection layered on top of tiering. + +The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key. + +You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.) + +You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first. + +#### Launch the Ephemeral Auto-Router Proxy + +```bash +lite autoroute up +``` + +Starts a local, throwaway litellm proxy on a random free port, running the config `configure` generated, with a freshly-minted random API key baked in for this session only (your real proxy key never leaves the generated config -- it only appears there, forwarding to your real proxy). It waits for the ephemeral proxy to report healthy, then patches `~/.claude/settings.json` the same way `lite up` does, except with a static `ANTHROPIC_AUTH_TOKEN` env var instead of an `apiKeyHelper`, since this key is short-lived and self-issued rather than something needing SSO refresh. Any `claude` session started afterward, from any terminal, routes through the ephemeral proxy. + +`lite autoroute up` runs in the foreground and streams the ephemeral proxy's own log file into your terminal, so you can watch its routing decisions -- which tier and model got picked for each request -- as you use Claude Code normally. Press Ctrl-C (or send SIGTERM) to stop it; this kills the child proxy process and restores your original Claude Code settings, in that order. + +#### Recover From an Unclean Shutdown + +```bash +lite autoroute down +``` + +If the `lite autoroute up` process dies uncleanly -- `kill -9`, a crash -- rather than being stopped with Ctrl-C, `down` is the manual recovery path: it kills any leftover ephemeral proxy process found via a recorded pid file and restores Claude Code's settings from whatever backup is on disk. + +#### Example + +```bash +lite autoroute configure +lite autoroute up +# use Claude Code as normal in another terminal; routing decisions stream live +lite autoroute down # only needed if `up` was killed uncleanly instead of Ctrl-C'd +``` + +#### Caveats + +Adaptive mode's learned state does not persist across `lite autoroute up` sessions -- there is no local database, so every session starts adaptive selection cold. A Claude Code session already running before `up` started, or still running when it stops, keeps whatever settings it loaded at its own startup; like `lite up`, this is a one-time file patch and restore, not a live traffic interceptor. Only Claude Code is supported, for the same reason as `lite up`: no other supported agent (for example Cursor) has an equivalent hot-patchable config file. + +A session that outlives `up` (or is still running the moment you stop it) keeps sending requests, master key included, to that now-freed loopback port until you restart it. Once the ephemeral proxy process exits, nothing stops another local account on the same machine from binding that same port and receiving those requests instead -- unlike `lite up`'s `apiKeyHelper`, which is re-resolved per request, `autoroute`'s master key is a static value, so whoever receives them gets a live-looking token along with the prompt content. Restart any Claude Code session before you consider the machine clean, run `lite autoroute down` promptly rather than leaving a stopped session's settings patched, and do not run `lite autoroute up` on a shared or multi-tenant host. + +Do not run `lite up` and `lite autoroute up` at the same time. Each patches `~/.claude/settings.json` and keeps its own separate backup, with no coordination between them: whichever one you stop or crash out of last is the one whose backup gets restored, which can silently leave the *other* mode's settings (a static master key and a now-dead loopback URL, or a stale `apiKeyHelper`) active. Run `lite down` or `lite autoroute down` (whichever applies) before switching to the other mode. + ## Environment Variables The CLI respects the following environment variables: diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index dfcedb70686..6b6252d8ecb 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -212,7 +212,7 @@ def _is_interactive() -> bool: return sys.stdin.isatty() -def _resolve_api_key(ctx: click.Context) -> str: +def resolve_api_key(ctx: click.Context) -> str: base_url = ctx.obj["base_url"] api_key = ctx.obj.get("api_key") if api_key: @@ -238,7 +238,7 @@ _SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy." def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: base_url = ctx.obj["base_url"] started_interactive = _is_interactive() - api_key = _resolve_api_key(ctx) + api_key = resolve_api_key(ctx) display_name, _ = agent_profile(binary) click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") @@ -288,5 +288,6 @@ __all__ = [ "agent_launch_args", "verify_proxy_key", "agent_profile", + "resolve_api_key", "AgentRunError", ] diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 4eda6817252..785d3b1e37b 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -277,6 +277,28 @@ def prompt_team_selection_fallback( return None +def _response_error_detail(response: requests.Response) -> str | None: + try: + body = response.json() + except ValueError: + return None + detail = body.get("detail") if isinstance(body, dict) else None + if isinstance(detail, str) and detail: + return detail + return None + + +def _polling_error_message(response: requests.Response) -> str: + detail = _response_error_detail(response) + if detail: + return f"Polling error: HTTP {response.status_code}: {detail}" + return f"Polling error: HTTP {response.status_code}" + + +def _is_permanent_polling_error(status_code: int) -> bool: + return 400 <= status_code < 500 and status_code != 429 + + # Polling-based authentication - no local server needed def _poll_for_ready_data( url: str, @@ -308,8 +330,14 @@ def _poll_for_ready_data( click.echo(pending_message) elif other_status_message and other_status_log_every > 0 and attempt % other_status_log_every == 0: click.echo(other_status_message) + elif _is_permanent_polling_error(response.status_code): + detail = _response_error_detail(response) + raise ValueError( + f"The proxy rejected the login session with HTTP {response.status_code}" + + (f": {detail}" if detail else f" and no error detail (from {url})") + ) elif http_error_log_every > 0 and attempt % http_error_log_every == 0: - click.echo(f"Polling error: HTTP {response.status_code}") + click.echo(_polling_error_message(response)) except requests.RequestException as e: if connection_error_log_every > 0 and attempt % connection_error_log_every == 0: click.echo(f"Connection error (will retry): {e}") @@ -342,12 +370,45 @@ def _normalize_teams(teams, team_details): def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]: - response = requests.post(f"{base_url}/sso/cli/start", timeout=10) - response.raise_for_status() - data = response.json() + start_url = f"{base_url}/sso/cli/start" + try: + response = requests.post(start_url, timeout=10) + except requests.RequestException as e: + raise ValueError( + f"Could not reach the proxy at {start_url}: {e}. " + "Check that the proxy is running and that --base-url points at it." + ) from e + + if response.status_code in (404, 405): + raise ValueError( + f"POST {start_url} returned HTTP {response.status_code}. " + "Either --base-url is wrong, or the proxy is older than this CLI and does not support " + "the CLI SSO login flow; upgrade the proxy or use a CLI version that matches it." + ) + if response.status_code != 200: + detail = _response_error_detail(response) + raise ValueError( + f"Starting CLI login failed: HTTP {response.status_code} from {start_url}" + + (f": {detail}" if detail else "") + ) + + try: + data = response.json() + except ValueError: + content_type = response.headers.get("content-type", "unknown") + raise ValueError( + f"The proxy returned a non-JSON response from {start_url} (content-type: {content_type}). " + "A proxy, load balancer, or auth gateway in front of LiteLLM may be intercepting the request. " + f"Response starts with: {response.text[:200]!r}" + ) + required_fields = ("login_id", "poll_secret", "user_code") - if not all(isinstance(data.get(field), str) for field in required_fields): - raise ValueError("Invalid CLI SSO start response") + missing_fields = tuple(field for field in required_fields if not isinstance(data.get(field), str)) + if missing_fields: + raise ValueError( + f"The response from {start_url} is missing required field(s): {', '.join(missing_fields)}. " + "The proxy version may not match this CLI; upgrade whichever is older." + ) return data @@ -577,6 +638,10 @@ def login(ctx: click.Context): return else: click.echo("❌ Authentication timed out. Please try again.") + click.echo( + "The proxy never reported the browser sign-in as finished. If you did complete it, " + "check the proxy logs for /sso/callback errors and confirm SSO is configured on the proxy." + ) return except KeyboardInterrupt: diff --git a/litellm/proxy/client/cli/commands/autoroute/__init__.py b/litellm/proxy/client/cli/commands/autoroute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/client/cli/commands/autoroute/commands.py b/litellm/proxy/client/cli/commands/autoroute/commands.py new file mode 100644 index 00000000000..161907f5b27 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/commands.py @@ -0,0 +1,207 @@ +import atexit +import json +import secrets +import signal +import threading +from types import FrameType + +import click +import yaml +from pydantic import JsonValue, TypeAdapter, ValidationError + +from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup +from ..up import BackupRecord as ClaudeBackupRecord +from .process import ( + AUTOROUTE_DIR, + CONFIG_PATH, + LOG_PATH, + PidRecord, + ProcessLaunchError, + allocate_free_port, + clear_pid_record, + is_running, + launch_proxy, + missing_proxy_runtime_modules, + poll_liveliness, + read_pid_record, + secure_create, + stream_log, + terminate, + write_pid_record, +) +from .settings import merge_claude_settings_static_token +from .wizard import run_configure_wizard + +AUTOROUTE_BACKUP_PATH = AUTOROUTE_DIR / "claude_settings_backup.json" + +_GENERATED_CONFIG_ADAPTER = TypeAdapter(dict[str, JsonValue]) + + +def _mint_and_embed_master_key() -> str: + """Generate a fresh key for this session and write it into the generated config.yaml. + + Must go under general_settings, not litellm_settings -- the proxy server only ever + reads general_settings.master_key (proxy_server.py:4530) to authenticate requests. A + key placed under litellm_settings is silently ignored, leaving the ephemeral proxy with + no real auth: any request reaches it regardless of the token Claude Code sends. + """ + master_key = secrets.token_urlsafe(32) + with open(CONFIG_PATH, "r") as f: + try: + generated = _GENERATED_CONFIG_ADAPTER.validate_python(yaml.safe_load(f)) + except (yaml.YAMLError, ValidationError): + raise click.ClickException( + f"{CONFIG_PATH} is empty or corrupt. Run `lite autoroute configure` again to regenerate it." + ) + general_settings = generated.get("general_settings") + updated_settings: dict[str, JsonValue] = { + **(general_settings if isinstance(general_settings, dict) else {}), + "master_key": master_key, + } + updated: dict[str, JsonValue] = {**generated, "general_settings": updated_settings} + with secure_create(CONFIG_PATH) as f: + yaml.safe_dump(updated, f, sort_keys=False) + return master_key + + +@click.group(name="autoroute") +def autoroute_group() -> None: + """QA complexity-based auto-routing against models your key can already use""" + + +@autoroute_group.command("configure") +@click.pass_context +def configure(ctx: click.Context) -> None: + """Discover accessible models and generate an ephemeral auto-router config""" + run_configure_wizard(ctx) + + +@autoroute_group.command("up") +def up() -> None: + """Launch the ephemeral auto-router proxy and route Claude Code through it""" + if not CONFIG_PATH.exists(): + raise click.ClickException("No config found. Run `lite autoroute configure` first.") + + missing = missing_proxy_runtime_modules() + if missing: + raise click.ClickException( + "lite autoroute up launches a local litellm proxy, which needs the proxy runtime that the " + f"thin `litellm[cli]` install does not include (missing: {', '.join(missing)}). Install the " + "proxy runtime with `uv tool install --force 'litellm[proxy]'`, or to QA a branch, " + "`curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install.sh | " + "LITELLM_CLI_REF= sh`." + ) + + try: + existing_pid = read_pid_record() + except UpError as e: + raise click.ClickException(str(e)) + if existing_pid is not None and is_running(existing_pid.pid): + raise click.ClickException( + "An ephemeral proxy is already running (lite autoroute up looks already active). " + "Run `lite autoroute down` first." + ) + + if AUTOROUTE_BACKUP_PATH.exists(): + raise click.ClickException( + f"{AUTOROUTE_BACKUP_PATH} already exists -- `lite autoroute up` looks like it's already " + "running (or crashed without cleanup). Run `lite autoroute down` first." + ) + + master_key = _mint_and_embed_master_key() + port = allocate_free_port() + base_url = f"http://127.0.0.1:{port}" + process = launch_proxy(CONFIG_PATH, port, LOG_PATH) + write_pid_record(PidRecord(pid=process.pid, port=port, config_path=str(CONFIG_PATH), log_path=str(LOG_PATH))) + + try: + poll_liveliness(base_url, LOG_PATH, process) + except ProcessLaunchError as e: + terminate(process.pid) + clear_pid_record() + raise click.ClickException(str(e)) + + try: + original_existed = CLAUDE_SETTINGS_PATH.exists() + original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH) + write_backup( + ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None), + AUTOROUTE_BACKUP_PATH, + ) + merged = merge_claude_settings_static_token(original_settings, base_url, master_key) + CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True) + with secure_create(CLAUDE_SETTINGS_PATH) as f: + json.dump(merged, f, indent=2) + except UpError as e: + terminate(process.pid) + clear_pid_record() + raise click.ClickException(str(e)) + + click.echo(f"litellm: ephemeral auto-router proxy up at {base_url} (pid {process.pid})") + click.echo("Claude Code sessions started now will route through it. Press Ctrl-C to stop and restore.") + + stop_event = threading.Event() + restored = threading.Lock() + + def _teardown() -> None: + if not restored.acquire(blocking=False): + return + terminate(process.pid) + clear_pid_record() + try: + restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) + except UpError as e: + # Runs from atexit/a signal handler too, outside Click's own exception + # handling -- raising here would only produce an unhandled-exception + # warning on stderr, not a clean message. + click.echo(str(e), err=True) + return + click.echo("\nStopped ephemeral proxy and restored Claude Code settings.") + click.echo( + f"Restart any Claude Code session still open from this session, or another local account could " + f"bind the now-free port {port} and receive its requests. Do not use `lite autoroute up` on a " + f"shared or multi-tenant host." + ) + + def _handle_signal(_signum: int, _frame: FrameType | None) -> None: + stop_event.set() + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + atexit.register(_teardown) + + log_thread = threading.Thread(target=stream_log, args=(LOG_PATH, stop_event), daemon=True) + log_thread.start() + + stop_event.wait() + _teardown() + + +@autoroute_group.command("down") +def down() -> None: + """Restore Claude Code settings and stop a leftover ephemeral proxy, if any""" + try: + record: PidRecord | None = read_pid_record() + except UpError as e: + # down is the crash-recovery path -- a corrupt pid record must not block it; clear the + # unusable record and keep going rather than leaving the user with no way to clean up. + click.echo(f"{e} Clearing it and continuing cleanup.", err=True) + record = None + if record is not None and is_running(record.pid): + terminate(record.pid) + click.echo(f"Stopped leftover ephemeral proxy (pid {record.pid}).") + clear_pid_record() + + try: + restored = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH) + except UpError as e: + raise click.ClickException(str(e)) + if restored is None: + click.echo("Nothing to restore.") + elif restored.existed: + click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.") + else: + click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite autoroute up`).") + + +__all__ = ["autoroute_group"] diff --git a/litellm/proxy/client/cli/commands/autoroute/config.py b/litellm/proxy/client/cli/commands/autoroute/config.py new file mode 100644 index 00000000000..2d760ef0f8a --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/config.py @@ -0,0 +1,249 @@ +from typing import Literal, Union + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter + +TIER_NAMES: tuple[str, ...] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") +AUTOROUTER_MODEL_NAME = "autorouter" + + +class ConfigGenerationError(Exception): + """Raised when an AutorouteConfig references a model the discovery step didn't find.""" + + +class DiscoveredModel(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + mode: str = "chat" + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +class _RawModelGroup(BaseModel): + model_config = ConfigDict(extra="ignore") + + model_group: str + # Optional: some real deployments return an explicit `"mode": null` for models that + # were registered without a mode (seen for embedding models like voyage-4-large). + # ModelGroupInfo's own "chat" default (litellm/types/router.py) only applies when the + # key is missing entirely, not when it's present as null, so this must tolerate None. + mode: str | None = "chat" + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + + +_RAW_MODEL_GROUPS_ADAPTER = TypeAdapter(list[_RawModelGroup]) + + +def parse_discovered_models(raw: list[JsonValue]) -> tuple[DiscoveredModel, ...]: + """Validate a raw `/model_group/info` response into typed models.""" + parsed = _RAW_MODEL_GROUPS_ADAPTER.validate_python(raw) + return tuple( + DiscoveredModel( + name=group.model_group, + # A null mode means the server genuinely doesn't know what this model does; + # "unknown" (rather than guessing "chat") keeps it out of both chat_models() + # and embedding_models() instead of risking a wrong-mode deployment. + mode=group.mode or "unknown", + input_cost_per_token=group.input_cost_per_token, + output_cost_per_token=group.output_cost_per_token, + ) + for group in parsed + ) + + +def chat_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: + return tuple(m for m in models if m.mode == "chat") + + +def embedding_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: + return tuple(m for m in models if m.mode == "embedding") + + +class HeuristicClassifier(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["heuristic"] = "heuristic" + + +class LLMClassifier(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["llm"] = "llm" + model: str + timeout_ms: int = 3000 + + +ClassifierChoice = Union[HeuristicClassifier, LLMClassifier] + + +class NoSemanticMatching(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["none"] = "none" + + +class KeywordTierRule(BaseModel): + model_config = ConfigDict(frozen=True) + keywords: tuple[str, ...] + tier: str + + +# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules" +# invariant with a sane starting point; the wizard lets the user override these per tier. +DEFAULT_KEYWORD_TIER_RULES: tuple[KeywordTierRule, ...] = ( + KeywordTierRule(keywords=("hi", "hello", "thanks"), tier="SIMPLE"), + KeywordTierRule(keywords=("explain", "how does"), tier="MEDIUM"), + KeywordTierRule(keywords=("refactor", "implement", "debug"), tier="COMPLEX"), + KeywordTierRule(keywords=("step by step", "think through", "prove"), tier="REASONING"), +) + + +class SemanticMatching(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["semantic"] = "semantic" + embedding_model: str + match_threshold: float = 0.5 + keyword_tier_rules: tuple[KeywordTierRule, ...] = DEFAULT_KEYWORD_TIER_RULES + + +SemanticMatchingChoice = Union[NoSemanticMatching, SemanticMatching] + + +class AutorouteConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + base_url: str + api_key: str + # Each tier maps to a pool of one or more models; complexity_router picks randomly among + # them per request (or, in adaptive mode, learns which to prefer within the pool). + tiers: dict[str, tuple[str, ...]] + default_model: str + classifier: ClassifierChoice = Field(default_factory=HeuristicClassifier) + semantic_matching: SemanticMatchingChoice = Field(default_factory=NoSemanticMatching) + adaptive: bool = False + + +def validate_config(config: AutorouteConfig, discovered: tuple[DiscoveredModel, ...]) -> None: + """Raise ConfigGenerationError if config references a model discovery didn't return.""" + chat_names: frozenset[str] = frozenset(m.name for m in chat_models(discovered)) + embedding_names: frozenset[str] = frozenset(m.name for m in embedding_models(discovered)) + + for tier, models in config.tiers.items(): + for model in models: + if model not in chat_names: + raise ConfigGenerationError(f"Tier {tier} references unknown chat model '{model}'") + + if config.default_model not in chat_names: + raise ConfigGenerationError(f"default_model '{config.default_model}' is not a known chat model") + + if isinstance(config.classifier, LLMClassifier) and config.classifier.model not in chat_names: + raise ConfigGenerationError(f"classifier model '{config.classifier.model}' is not a known chat model") + + if ( + isinstance(config.semantic_matching, SemanticMatching) + and config.semantic_matching.embedding_model not in embedding_names + ): + raise ConfigGenerationError( + f"embedding model '{config.semantic_matching.embedding_model}' is not a known embedding model" + ) + + +def _litellm_proxy_deployment(name: str, base_url: str, api_key: str) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"litellm_proxy/{name}", + "api_base": base_url, + "api_key": api_key, + }, + } + + +def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]: + """Build the model_list for the ephemeral proxy's config.yaml. + + Every real model referenced anywhere (tier targets, classifier, embedding) is deduplicated + to exactly one `litellm_proxy/` deployment forwarding to the customer's real proxy, + plus one `auto_router/complexity_router` deployment tying the tiers together. + """ + referenced_names = {model for models in config.tiers.values() for model in models} + referenced_names.add(config.default_model) + if isinstance(config.classifier, LLMClassifier): + referenced_names.add(config.classifier.model) + if isinstance(config.semantic_matching, SemanticMatching): + referenced_names.add(config.semantic_matching.embedding_model) + + proxy_deployments = [ + _litellm_proxy_deployment(name, config.base_url, config.api_key) for name in sorted(referenced_names) + ] + + complexity_router_config: dict[str, JsonValue] = { + "tiers": {tier: list(models) for tier, models in config.tiers.items()}, + "default_model": config.default_model, + } + if isinstance(config.classifier, LLMClassifier): + complexity_router_config["classifier_type"] = "llm" + complexity_router_config["classifier_llm_config"] = { + "model": config.classifier.model, + "timeout_ms": config.classifier.timeout_ms, + } + if isinstance(config.semantic_matching, SemanticMatching): + complexity_router_config["semantic_keyword_matching"] = True + complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model + complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold + complexity_router_config["keyword_tier_rules"] = [ + {"keywords": list(rule.keywords), "tier": rule.tier} for rule in config.semantic_matching.keyword_tier_rules + ] + if config.adaptive: + complexity_router_config["adaptive"] = True + + auto_router_litellm_params: dict[str, JsonValue] = { + "model": "auto_router/complexity_router", + "complexity_router_config": complexity_router_config, + } + # A bare "*" model_name looks like the obvious way to catch every request Claude Code + # might send regardless of which model it thinks it's using, but Router's auto-router + # registry is keyed by the literal requested model string (router.py:10711-10717), not + # resolved through pattern/wildcard matching first -- so a "*" entry here would only ever + # match a client that literally sends model="*", never an actual wildcard catch-all. Callers + # instead need to make Claude Code request this "autorouter" name directly (see + # ANTHROPIC_DEFAULT_*_MODEL in settings.py's merge_claude_settings_static_token). + return [ + *proxy_deployments, + {"model_name": AUTOROUTER_MODEL_NAME, "litellm_params": auto_router_litellm_params}, + ] + + +def build_generated_proxy_config(config: AutorouteConfig, master_key: str) -> dict[str, JsonValue]: + """Full config.yaml content for the ephemeral proxy, including its own auth key. + + master_key must live under general_settings, not litellm_settings -- the proxy server + only ever reads general_settings.master_key (proxy_server.py:4530) to authenticate + requests; a key placed under litellm_settings is silently ignored, leaving the proxy + with no real auth at all. + """ + return { + "model_list": build_generated_model_list(config), + "general_settings": {"master_key": master_key}, + } + + +__all__ = [ + "AUTOROUTER_MODEL_NAME", + "TIER_NAMES", + "AutorouteConfig", + "ClassifierChoice", + "ConfigGenerationError", + "DEFAULT_KEYWORD_TIER_RULES", + "DiscoveredModel", + "HeuristicClassifier", + "KeywordTierRule", + "LLMClassifier", + "NoSemanticMatching", + "SemanticMatching", + "SemanticMatchingChoice", + "build_generated_model_list", + "build_generated_proxy_config", + "chat_models", + "embedding_models", + "parse_discovered_models", + "validate_config", +] diff --git a/litellm/proxy/client/cli/commands/autoroute/process.py b/litellm/proxy/client/cli/commands/autoroute/process.py new file mode 100644 index 00000000000..712f2eed2da --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/process.py @@ -0,0 +1,190 @@ +import contextlib +import importlib.util +import json +import os +import signal +import socket +import subprocess +import sys +import threading +import time +from dataclasses import dataclass +from pathlib import Path + +import click +import requests +from pydantic import TypeAdapter, ValidationError + +from ..up import UpError, secure_create + +AUTOROUTE_DIR = Path.home() / ".litellm" / "autorouter" +CONFIG_PATH = AUTOROUTE_DIR / "config.yaml" +LOG_PATH = AUTOROUTE_DIR / "proxy.log" +PID_RECORD_PATH = AUTOROUTE_DIR / "proxy.pid.json" + + +class ProcessLaunchError(Exception): + """Raised when the ephemeral proxy subprocess fails to come up healthy.""" + + +@dataclass(frozen=True, slots=True) +class PidRecord: + pid: int + port: int + config_path: str + log_path: str + + +_PID_RECORD_ADAPTER = TypeAdapter(PidRecord) + + +_PROXY_RUNTIME_MODULES: tuple[str, ...] = ("fastapi", "uvicorn", "backoff", "orjson", "websockets", "apscheduler") + + +def missing_proxy_runtime_modules() -> tuple[str, ...]: + """Proxy-server modules that ``lite autoroute up`` needs but the thin CLI install lacks. + + ``launch_proxy`` runs the full ``litellm.proxy.proxy_cli`` server, whose dependencies live in + the ``proxy`` extra, not the ``cli`` extra that installs the ``lite`` command. On a thin + ``litellm[cli]`` install the subprocess dies with a bare ``ModuleNotFoundError``; detecting the + gap here lets ``up`` fail with an actionable message instead. + """ + return tuple(name for name in _PROXY_RUNTIME_MODULES if importlib.util.find_spec(name) is None) + + +def allocate_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def launch_proxy(config_path: Path, port: int, log_path: Path) -> "subprocess.Popen[bytes]": + log_path.parent.mkdir(parents=True, exist_ok=True) + with open(log_path, "w") as log_file: + return subprocess.Popen( + [ + sys.executable, + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config_path), + "--port", + str(port), + "--host", + "127.0.0.1", + ], + stdout=log_file, + stderr=subprocess.STDOUT, + ) + + +def _tail(log_path: Path, lines: int = 40) -> str: + if not log_path.exists(): + return "(no log output captured)" + return "\n".join(log_path.read_text(errors="replace").splitlines()[-lines:]) + + +def poll_liveliness(base_url: str, log_path: Path, process: "subprocess.Popen[bytes]", timeout: float = 30.0) -> None: + """Poll /health/liveliness until it responds, the process dies, or timeout elapses.""" + deadline = time.monotonic() + timeout + url = base_url.rstrip("/") + "/health/liveliness" + while time.monotonic() < deadline: + if process.poll() is not None: + raise ProcessLaunchError( + f"Ephemeral proxy exited early (code {process.returncode}). Last log lines:\n{_tail(log_path)}" + ) + with contextlib.suppress(requests.RequestException): + if requests.get(url, timeout=2).status_code == 200: + return + time.sleep(0.5) + raise ProcessLaunchError( + f"Ephemeral proxy never became healthy within {timeout}s. Last log lines:\n{_tail(log_path)}" + ) + + +def write_pid_record(record: PidRecord, path: Path | None = None) -> None: + resolved_path = path if path is not None else PID_RECORD_PATH + resolved_path.parent.mkdir(parents=True, exist_ok=True) + with open(resolved_path, "w") as f: + json.dump( + {"pid": record.pid, "port": record.port, "config_path": record.config_path, "log_path": record.log_path}, + f, + indent=2, + ) + + +def read_pid_record(path: Path | None = None) -> PidRecord | None: + resolved_path = path if path is not None else PID_RECORD_PATH + if not resolved_path.exists(): + return None + with open(resolved_path, "r") as f: + content = f.read() + try: + return _PID_RECORD_ADAPTER.validate_json(content) + except ValidationError: + raise UpError(f"{resolved_path} contains invalid or unexpected JSON; cannot proceed safely.") + + +def clear_pid_record(path: Path | None = None) -> None: + resolved_path = path if path is not None else PID_RECORD_PATH + resolved_path.unlink(missing_ok=True) + + +def is_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def terminate(pid: int, grace_period: float = 5.0) -> None: + """Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed.""" + if not is_running(pid): + return + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGTERM) + deadline = time.monotonic() + grace_period + while time.monotonic() < deadline and is_running(pid): + time.sleep(0.2) + if is_running(pid): + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + + +def stream_log(log_path: Path, stop_event: threading.Event) -> None: + """Print new lines appended to log_path until stop_event is set. Blocks the calling thread.""" + while not log_path.exists() and not stop_event.is_set(): + time.sleep(0.1) + if stop_event.is_set() or not log_path.exists(): + return + with open(log_path, "r") as f: + while not stop_event.is_set(): + line = f.readline() + if line: + click.echo(line, nl=False) + else: + time.sleep(0.2) + + +__all__ = [ + "AUTOROUTE_DIR", + "CONFIG_PATH", + "LOG_PATH", + "PID_RECORD_PATH", + "PidRecord", + "ProcessLaunchError", + "allocate_free_port", + "clear_pid_record", + "is_running", + "launch_proxy", + "missing_proxy_runtime_modules", + "poll_liveliness", + "read_pid_record", + "secure_create", + "stream_log", + "terminate", + "write_pid_record", +] diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py new file mode 100644 index 00000000000..4bed184eb34 --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -0,0 +1,46 @@ +from pydantic import JsonValue + +from .config import AUTOROUTER_MODEL_NAME + +ENV_KEY = "env" +API_KEY_HELPER_KEY = "apiKeyHelper" +ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY" +ANTHROPIC_AUTH_TOKEN_KEY = "ANTHROPIC_AUTH_TOKEN" +ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL" +# Force every one of Claude Code's own model tiers to request the auto-router by name. +# Router's auto-router registry is keyed by the literal requested model string +# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" +# model_name can never work as a catch-all -- these overrides are what actually makes +# Claude Code send "autorouter" regardless of /model or its own version-specific defaults. +ANTHROPIC_DEFAULT_MODEL_ENV_KEYS = ( + "ANTHROPIC_DEFAULT_SONNET_MODEL", + "ANTHROPIC_DEFAULT_HAIKU_MODEL", + "ANTHROPIC_DEFAULT_OPUS_MODEL", +) + + +def merge_claude_settings_static_token( + settings: dict[str, JsonValue], base_url: str, auth_token: str +) -> dict[str, JsonValue]: + """Return a new settings dict wired to a local ephemeral proxy with a static token. + + Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real + remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key was just + minted for this session, so a plain env var is simpler and correct. Any existing + apiKeyHelper is cleared so it can't fight with the static token. + """ + raw_env = settings.get(ENV_KEY, {}) + base_env = raw_env if isinstance(raw_env, dict) else {} + env: dict[str, JsonValue] = { + **base_env, + ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), + ANTHROPIC_AUTH_TOKEN_KEY: auth_token, + **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, + } + env.pop(ANTHROPIC_API_KEY_KEY, None) + merged: dict[str, JsonValue] = {**settings, ENV_KEY: env} + merged.pop(API_KEY_HELPER_KEY, None) + return merged + + +__all__ = ["merge_claude_settings_static_token"] diff --git a/litellm/proxy/client/cli/commands/autoroute/wizard.py b/litellm/proxy/client/cli/commands/autoroute/wizard.py new file mode 100644 index 00000000000..60696fb2e7e --- /dev/null +++ b/litellm/proxy/client/cli/commands/autoroute/wizard.py @@ -0,0 +1,150 @@ +import sys +from pathlib import Path + +import click +import yaml +from InquirerPy import inquirer +from InquirerPy.base.control import Choice + +from .... import Client +from .config import ( + DEFAULT_KEYWORD_TIER_RULES, + TIER_NAMES, + AutorouteConfig, + ConfigGenerationError, + DiscoveredModel, + HeuristicClassifier, + KeywordTierRule, + LLMClassifier, + NoSemanticMatching, + SemanticMatching, + build_generated_model_list, + chat_models, + embedding_models, + parse_discovered_models, + validate_config, +) +from .process import CONFIG_PATH, secure_create + + +def _is_interactive() -> bool: + return sys.stdin.isatty() + + +def _fuzzy_pick(models: tuple[DiscoveredModel, ...], prompt_label: str, multiselect: bool) -> list[str]: + """Type-to-filter picker over a (possibly huge) model pool, using InquirerPy's fzf-style fuzzy prompt. + + A plain numbered table + typed index does not scale past a handful of models -- proxies with + hundreds of model groups made that interaction unusable. This lets the user narrow the pool by + typing a substring instead of scrolling/counting. + + Assumes the caller already checked interactivity (run_configure_wizard does, once, up front) -- + checking here too would check the wrong thing under test, where InquirerPy is driven through its + own injected input/output rather than the real process stdin. + """ + choices = [Choice(value=model.name, name=model.name) for model in models] + toggle_hint = "tab to toggle, " if multiselect else "" + while True: + result = inquirer.fuzzy( + message=f"{prompt_label}: type to filter, {toggle_hint}enter to confirm", + choices=choices, + multiselect=multiselect, + max_height="70%", + ).execute() + selected = result if multiselect else [result] + if selected: + return selected + click.echo("Select at least one model.") + + +def _render_and_prompt_for_model(models: tuple[DiscoveredModel, ...], prompt_label: str) -> str: + return _fuzzy_pick(models, prompt_label, multiselect=False)[0] + + +def _render_and_prompt_for_models(models: tuple[DiscoveredModel, ...], prompt_label: str) -> tuple[str, ...]: + return tuple(_fuzzy_pick(models, prompt_label, multiselect=True)) + + +def _parse_keywords(raw: str) -> tuple[str, ...]: + return tuple(keyword.strip() for keyword in raw.split(",") if keyword.strip()) + + +def _prompt_for_keyword_tier_rules() -> tuple[KeywordTierRule, ...]: + """Let the user supply the semantic-matching keywords per tier, since matching those + keywords against the request is the whole point of enabling it. Each prompt is prefilled + with the built-in default, so pressing enter keeps it.""" + click.echo("\nEnter example keywords/phrases per tier (comma-separated); press enter to keep the default:") + defaults = {rule.tier: rule.keywords for rule in DEFAULT_KEYWORD_TIER_RULES} + + def _rule_for(tier: str) -> KeywordTierRule: + default_keywords = defaults.get(tier, ()) + raw = click.prompt(f" {tier} keywords", default=", ".join(default_keywords), show_default=True) + return KeywordTierRule(keywords=_parse_keywords(raw) or default_keywords, tier=tier) + + return tuple(_rule_for(tier) for tier in TIER_NAMES) + + +def run_configure_wizard(ctx: click.Context) -> Path: + """Discover the caller's accessible models, walk them through tier assignment, write config.""" + base_url = ctx.obj["base_url"] + api_key = ctx.obj["api_key"] + client = Client(base_url=base_url, api_key=api_key) + + raw_groups = client.model_groups.info() + if not isinstance(raw_groups, list): + raise click.ClickException( + f"Unexpected response from /model_group/info: expected a list, got {type(raw_groups).__name__}" + ) + discovered = parse_discovered_models(raw_groups) + chat_pool = chat_models(discovered) + embedding_pool = embedding_models(discovered) + + if not chat_pool: + raise click.ClickException("Your key has no chat-capable models available on this proxy.") + + if not _is_interactive(): + raise click.ClickException("`lite autoroute configure` requires an interactive terminal.") + + click.echo("Assign model(s) to each complexity tier (from what your key can access):") + tiers = {tier: _render_and_prompt_for_models(chat_pool, tier) for tier in TIER_NAMES} + default_model = tiers["MEDIUM"][0] + + classifier = HeuristicClassifier() + if click.confirm("\nUse an LLM classifier instead of the free heuristic scorer?", default=False): + classifier_model = _render_and_prompt_for_model(chat_pool, "LLM classifier") + classifier = LLMClassifier(model=classifier_model) + + semantic_matching = NoSemanticMatching() + if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False): + embedding_model = _render_and_prompt_for_model(embedding_pool, "semantic embeddings") + keyword_tier_rules = _prompt_for_keyword_tier_rules() + semantic_matching = SemanticMatching(embedding_model=embedding_model, keyword_tier_rules=keyword_tier_rules) + + adaptive = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False) + + config = AutorouteConfig( + base_url=base_url, + api_key=api_key, + tiers=tiers, + default_model=default_model, + classifier=classifier, + semantic_matching=semantic_matching, + adaptive=adaptive, + ) + try: + validate_config(config, discovered) + except ConfigGenerationError as e: + raise click.ClickException(str(e)) + + model_list = build_generated_model_list(config) + CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) + with secure_create(CONFIG_PATH) as f: + yaml.safe_dump({"model_list": model_list}, f, sort_keys=False) + + click.echo(f"\nWrote {CONFIG_PATH}") + for tier, models in tiers.items(): + click.echo(f" {tier}: {', '.join(models)}") + return CONFIG_PATH + + +__all__ = ["run_configure_wizard"] diff --git a/litellm/proxy/client/cli/commands/model_groups.py b/litellm/proxy/client/cli/commands/model_groups.py new file mode 100644 index 00000000000..7de959a9b78 --- /dev/null +++ b/litellm/proxy/client/cli/commands/model_groups.py @@ -0,0 +1,57 @@ +from typing import Literal + +import click +import rich +import rich.table + +from ... import Client + + +def create_client(ctx: click.Context) -> Client: + return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + + +@click.group(name="model-groups") +def model_groups() -> None: + """Inspect model groups your key can access on the proxy""" + + +@model_groups.command("list") +@click.option( + "--format", + "output_format", + type=click.Choice(["table", "json"]), + default="table", + help="Output format (table or json)", +) +@click.pass_context +def list_model_groups(ctx: click.Context, output_format: Literal["table", "json"]) -> None: + """List model groups accessible to your key, with mode and pricing""" + client = create_client(ctx) + groups = client.model_groups.info() + if not isinstance(groups, list): + raise click.ClickException( + f"Unexpected response from /model_group/info: expected a list, got {type(groups).__name__}" + ) + + if output_format == "json": + rich.print_json(data=groups) + return + + table = rich.table.Table(title="Accessible Model Groups") + table.add_column("Model", style="cyan") + table.add_column("Mode", style="green") + table.add_column("Input $/token", style="yellow") + table.add_column("Output $/token", style="yellow") + + for group in groups: + table.add_row( + str(group.get("model_group", "")), + str(group.get("mode", "chat")), + str(group.get("input_cost_per_token", "")), + str(group.get("output_cost_per_token", "")), + ) + rich.print(table) + + +__all__ = ["model_groups"] diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py new file mode 100644 index 00000000000..dc9157d7ca4 --- /dev/null +++ b/litellm/proxy/client/cli/commands/up.py @@ -0,0 +1,283 @@ +import atexit +import contextlib +import json +import os +import shlex +import shutil +import signal +import sys +import threading +from dataclasses import dataclass +from pathlib import Path +from types import FrameType +from typing import IO, Iterator, Mapping + +import click +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh + +from .agents import AgentRunError, resolve_api_key, verify_proxy_key +from .auth import load_token, login + +ENV_KEY = "env" +API_KEY_HELPER_KEY = "apiKeyHelper" +ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL" +ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY" + +CLAUDE_SETTINGS_PATH = Path.home() / ".claude" / "settings.json" +BACKUP_PATH = Path.home() / ".litellm" / "claude_settings_backup.json" + + +class UpError(Exception): + """Raised for any user-actionable failure while starting/stopping interception.""" + + +@dataclass(frozen=True, slots=True) +class BackupRecord: + """Snapshot of ~/.claude/settings.json taken right before `lite up` patches it.""" + + existed: bool + content: dict[str, JsonValue] | None + + +_SETTINGS_ADAPTER = TypeAdapter(dict[str, JsonValue]) +_BACKUP_RECORD_ADAPTER = TypeAdapter(BackupRecord) + + +def load_json_or_empty(path: Path) -> dict[str, JsonValue]: + if not path.exists(): + return {} + with open(path, "r") as f: + content = f.read() + if not content.strip(): + return {} + try: + return _SETTINGS_ADAPTER.validate_json(content) + except ValidationError: + raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.") + + +def merge_claude_settings( + settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str +) -> dict[str, JsonValue]: + """Return a new settings dict wired to route Claude Code through the proxy. + + Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a + stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued + token (same reasoning as build_agent_env in agents.py). Every other key is + preserved untouched. + """ + raw_env = settings.get(ENV_KEY, {}) + base_env = raw_env if isinstance(raw_env, dict) else {} + env = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")} + env.pop(ANTHROPIC_API_KEY_KEY, None) + return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} + + +@contextlib.contextmanager +def secure_create(path: Path) -> Iterator[IO[str]]: + """Open path for writing with mode 0600 fixed up before any content is written. + + A plain `open(path, "w")` creates a *new* file at the umask-derived default (commonly 0644) + and leaves it world- or group-readable until a later `chmod` call catches up -- a real window + in which a file holding a credential is readable by another local account. Passing the mode to + `os.open` closes that window for a brand-new file, but `O_CREAT`'s mode argument is only + applied on creation: if the file already exists its old, broader permissions carry over + untouched. `os.fchmod` right after opening -- before a single byte of the new content is + written -- covers both cases. + """ + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + os.fchmod(fd, 0o600) + f: IO[str] = os.fdopen(fd, "w") + try: + yield f + finally: + f.close() + + +def write_backup(record: BackupRecord, backup_path: Path | None = None) -> None: + path = backup_path if backup_path is not None else BACKUP_PATH + path.parent.mkdir(exist_ok=True) + with secure_create(path) as f: + json.dump({"existed": record.existed, "content": record.content}, f, indent=2) + + +def read_backup(backup_path: Path | None = None) -> BackupRecord | None: + path = backup_path if backup_path is not None else BACKUP_PATH + if not path.exists(): + return None + with open(path, "r") as f: + content = f.read() + try: + return _BACKUP_RECORD_ADAPTER.validate_json(content) + except ValidationError: + raise UpError(f"{path} contains invalid or unexpected JSON; cannot restore from it safely.") + + +def restore_claude_settings(settings_path: Path | None = None, backup_path: Path | None = None) -> BackupRecord | None: + """Restore settings_path from the backup at backup_path, then delete the backup. + + Returns the restored record, or None if there was nothing to restore. + """ + resolved_settings_path = settings_path if settings_path is not None else CLAUDE_SETTINGS_PATH + resolved_backup_path = backup_path if backup_path is not None else BACKUP_PATH + record = read_backup(resolved_backup_path) + if record is None: + return None + if record.existed and record.content is not None: + resolved_settings_path.parent.mkdir(parents=True, exist_ok=True) + with open(resolved_settings_path, "w") as f: + json.dump(record.content, f, indent=2) + elif resolved_settings_path.exists(): + resolved_settings_path.unlink() + resolved_backup_path.unlink() + return record + + +def resolve_api_key_helper(base_url: str) -> str: + """Build the shell command Claude Code should run for its apiKeyHelper. + + Resolves `lite` to an absolute path so the helper works regardless of the + PATH visible to whatever subprocess Claude Code spawns it from. Passing + --base-url explicitly (rather than relying on the bare invocation Claude + Code would otherwise use) makes `print-token` enforce that the cached + token was actually issued for this proxy -- without it, a token minted + for a different, previously-logged-into proxy would be handed to + whichever server `up` currently points at. + """ + lite_path = shutil.which("lite") + if lite_path is None: + raise UpError( + "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs " + "an absolute path to it, so `lite up` cannot continue." + ) + return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}" + + +def _ensure_fresh_login(ctx: click.Context) -> None: + base_url = ctx.obj["base_url"].rstrip("/") + token_data = load_token() + if token_data and token_data.get("base_url") == base_url and is_cli_token_fresh(token_data): + return + + if not sys.stdin.isatty(): + raise UpError( + "No fresh LiteLLM login found for this proxy. Run `lite login` first (apiKeyHelper " + "reads this token on every Claude Code request)." + ) + + click.echo("No fresh LiteLLM login found for this proxy; starting login...") + ctx.invoke(login) + token_data = load_token() + if not token_data or token_data.get("base_url") != base_url or not is_cli_token_fresh(token_data): + raise UpError("Login did not produce a usable token; cannot start `lite up`.") + + +def _restore_and_report() -> None: + record = restore_claude_settings() + if record is None: + click.echo("Nothing to restore.") + return + if record.existed: + click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.") + else: + click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite up`).") + + +@click.command(name="up") +@click.pass_context +def up(ctx: click.Context) -> None: + """Route every Claude Code session through your LiteLLM proxy until stopped. + + Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own + next startup, from any terminal -- no need to launch it through `lite`. + Press Ctrl-C to stop and restore your original settings. Assumes the proxy + is already running (this does not start one for you). Cursor is not + supported: it has no equivalent file-based config to patch. + """ + base_url = ctx.obj["base_url"] + + try: + _ensure_fresh_login(ctx) + api_key = resolve_api_key(ctx) + verify_proxy_key(base_url, api_key) + + if BACKUP_PATH.exists(): + raise UpError( + f"{BACKUP_PATH} already exists -- `lite up` looks like it's already " + "running (or crashed without cleanup). Run `lite down` first." + ) + + api_key_helper = resolve_api_key_helper(base_url) + original_existed = CLAUDE_SETTINGS_PATH.exists() + original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH) + write_backup( + BackupRecord( + existed=original_existed, + content=original_settings if original_existed else None, + ) + ) + + CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True) + merged = merge_claude_settings(original_settings, base_url, api_key_helper) + with open(CLAUDE_SETTINGS_PATH, "w") as f: + json.dump(merged, f, indent=2) + except (AgentRunError, UpError) as e: + raise click.ClickException(str(e)) + + click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}") + click.echo("Press Ctrl-C to stop and restore your original settings.") + + stop_event = threading.Event() + restored = threading.Lock() + + def _handle_signal(_signum: int, _frame: FrameType | None) -> None: + stop_event.set() + + def _restore_once() -> None: + if not restored.acquire(blocking=False): + return + try: + _restore_and_report() + except UpError as e: + # Runs from atexit/a signal handler, outside Click's own exception + # handling -- raising here would only produce an unhandled-exception + # warning on stderr, not a clean message. + click.echo(str(e), err=True) + + signal.signal(signal.SIGINT, _handle_signal) + signal.signal(signal.SIGTERM, _handle_signal) + atexit.register(_restore_once) + + stop_event.wait() + _restore_once() + + +@click.command(name="down") +def down() -> None: + """Restore ~/.claude/settings.json if a `lite up` session left it patched. + + Use this after a `lite up` process was killed uncleanly (e.g. `kill -9`) + instead of stopped with Ctrl-C. + """ + try: + _restore_and_report() + except UpError as e: + raise click.ClickException(str(e)) + + +__all__ = [ + "BACKUP_PATH", + "CLAUDE_SETTINGS_PATH", + "BackupRecord", + "UpError", + "down", + "load_json_or_empty", + "merge_claude_settings", + "read_backup", + "resolve_api_key_helper", + "restore_claude_settings", + "up", + "write_backup", +] diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 4de3ff5fc87..e641956b2c5 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -9,15 +9,18 @@ from litellm.proxy.client.health import HealthManagementClient from .commands.agents import agent_commands from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami +from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.credentials import credentials from .commands.encryption import encryption from .commands.http import http from .commands.keys import keys +from .commands.model_groups import model_groups # local imports from .commands.models import models from .commands.teams import teams +from .commands.up import down, up from .commands.users import users from .interface import interactive_shell @@ -131,6 +134,13 @@ cli.add_command(users) # Add a top-level command per coding agent (claude, codex, opencode, ...) for agent_command in agent_commands(): cli.add_command(agent_command) +# Add the up/down commands (route Claude Code through the local LiteLLM proxy) +cli.add_command(up) +cli.add_command(down) +# Add the model-groups command group (discover models your key can access) +cli.add_command(model_groups) +# Add the autoroute command group (QA auto-routing against your real proxy) +cli.add_command(autoroute_group, name="autoroute") if __name__ == "__main__": diff --git a/litellm/proxy/enterprise_billing/__init__.py b/litellm/proxy/enterprise_billing/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/enterprise_billing/billing_metrics.py b/litellm/proxy/enterprise_billing/billing_metrics.py new file mode 100644 index 00000000000..f9f8ceaf721 --- /dev/null +++ b/litellm/proxy/enterprise_billing/billing_metrics.py @@ -0,0 +1,323 @@ +""" +Push-based OTLP metering for enterprise litellm deployments. + +Owns a dedicated OpenTelemetry meter provider and an OTLP/HTTP exporter +authenticated to our global collector with a TLS client certificate. The +collector front end terminates mutual TLS: the client certificate presented +here is validated against our CA at the edge, and the verified subject is +what identifies the deployment. It is intentionally isolated from the global +meter provider so the customer's own OTEL metrics are untouched and ours +never leak into their backend. + +The deployment's identity rides on the TLS client certificate, not on the +payload; the secret license key is never sent as an attribute or header. +""" + +import os +import tempfile +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter +from opentelemetry.metrics import Counter +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk.resources import Resource + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.middleware.billable_request_metrics_middleware import ( + BillableCategory, +) + +if TYPE_CHECKING: + from litellm.proxy._types import EnterpriseLicenseData + +ENDPOINT_ENV = "LITELLM_BILLING_METRICS_ENDPOINT" +CLIENT_CERT_ENV = "LITELLM_BILLING_METRICS_CLIENT_CERT" +CLIENT_KEY_ENV = "LITELLM_BILLING_METRICS_CLIENT_KEY" +CA_CERT_ENV = "LITELLM_BILLING_METRICS_CA_CERT" +EXPORT_INTERVAL_ENV = "LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS" +DEFAULT_EXPORT_INTERVAL_MS = 60_000 +SHUTDOWN_FLUSH_TIMEOUT_MS = 5_000 +_METRICS_PATH = "/v1/metrics" + +# The cert env vars take a path or the PEM itself. Secret stores that inject +# values as env content cannot mount them as files, so inline PEM is written out. +_PEM_PREFIX = "-----BEGIN" +_PEM_DIR_PREFIX = "litellm-billing-mtls-" +_PEM_FILE_MODE = 0o600 +_CLIENT_CERT_FILENAME = "client.crt" +_CLIENT_KEY_FILENAME = "client.key" +_CA_CERT_FILENAME = "ca.crt" + +METRIC_NAME = "litellm.enterprise.billable_requests" +METER_NAME = "litellm.enterprise.billing" + +AttributeValue = Union[str, int] + + +@dataclass(frozen=True, slots=True) +class BillingMetricsConfig: + endpoint: str + client_cert_path: str + client_key_path: str + ca_cert_path: Optional[str] + export_interval_ms: int + litellm_version: str + license_id: Optional[str] + + +def _metrics_endpoint(endpoint: str) -> str: + """The OTLP/HTTP metric exporter wants the full URL including the signal path.""" + trimmed = endpoint.rstrip("/") + return trimmed if trimmed.endswith(_METRICS_PATH) else f"{trimmed}{_METRICS_PATH}" + + +def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]: + base: dict[str, AttributeValue] = { + "service.name": "litellm-proxy", + "litellm.version": config.litellm_version, + } + license_attr: dict[str, AttributeValue] = {"litellm.license.id": config.license_id} if config.license_id else {} + return {**base, **license_attr} + + +def _billable_attributes( + category: BillableCategory, route: str, status_code: int, model_id: Optional[str] +) -> dict[str, AttributeValue]: + base: dict[str, AttributeValue] = { + "litellm.endpoint.category": category.value, + "http.route": route, + "http.response.status_code": status_code, + } + model_attr: dict[str, AttributeValue] = {"litellm.model_id": model_id} if model_id else {} + return {**base, **model_attr} + + +def build_mtls_meter_provider(config: BillingMetricsConfig) -> MeterProvider: + """OTLP/HTTP exporter presenting a TLS client certificate. + + The collector's load balancer terminates mutual TLS and validates the client + certificate against our CA. Server verification uses the system trust store + (the collector presents a public web-PKI certificate); ca_cert_path overrides + it only for private/test collectors. + """ + exporter = OTLPMetricExporter( + endpoint=_metrics_endpoint(config.endpoint), + # None -> exporter falls back to the system trust store. + certificate_file=config.ca_cert_path, + client_certificate_file=config.client_cert_path, + client_key_file=config.client_key_path, + ) + reader = PeriodicExportingMetricReader(exporter, export_interval_millis=config.export_interval_ms) + return MeterProvider(metric_readers=[reader], resource=Resource.create(_resource_attributes(config))) + + +class BillingMetricsRecorder: + """Increments one OTLP counter per billable request. The meter provider is injected (see the factory).""" + + def __init__(self, provider: MeterProvider) -> None: + self._provider = provider + self._counter: Counter = provider.get_meter(METER_NAME).create_counter( + name=METRIC_NAME, + unit="{request}", + description="Count of 2xx HTTP requests to billable LLM/MCP/A2A endpoints", + ) + + def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None: + self._counter.add(1, _billable_attributes(category, route, status_code, model_id)) + + def shutdown(self) -> None: + """Final flush + exporter-thread stop. Without this, up to one export + interval of billable counts is dropped on every proxy restart.""" + self._provider.shutdown(timeout_millis=SHUTDOWN_FLUSH_TIMEOUT_MS) + + +def _export_interval_ms() -> int: + raw = os.getenv(EXPORT_INTERVAL_ENV) + if raw is None: + return DEFAULT_EXPORT_INTERVAL_MS + try: + return int(raw) + except ValueError: + verbose_proxy_logger.warning( + "Invalid %s=%r, falling back to %d ms", EXPORT_INTERVAL_ENV, raw, DEFAULT_EXPORT_INTERVAL_MS + ) + return DEFAULT_EXPORT_INTERVAL_MS + + +@dataclass(frozen=True, slots=True) +class _CredentialPaths: + client_cert_path: str + client_key_path: str + ca_cert_path: Optional[str] + + +def _is_pem_content(value: str) -> bool: + return value.lstrip().startswith(_PEM_PREFIX) + + +def _write_pem(directory: str, filename: str, pem: str) -> str: + path = os.path.join(directory, filename) + with open(path, "w", encoding="utf-8") as handle: + handle.write(pem if pem.endswith("\n") else f"{pem}\n") + os.chmod(path, _PEM_FILE_MODE) + return path + + +def _resolve_credential_paths(*, client_cert: str, client_key: str, ca_cert: Optional[str]) -> _CredentialPaths: + """ + Accept either a filesystem path or inline PEM content for each credential. + + Secret stores that inject values as environment content rather than mounted + files (ECS tasks reading AWS Secrets Manager, Cloud Run reading Secret + Manager) can only deliver the certificate as a string. The OTLP exporter + takes paths, so inline PEM is written to a private directory once, when the + recorder is built. Raises OSError if that write fails; the caller disables + metering rather than propagating. + """ + inline = tuple(value for value in (client_cert, client_key, ca_cert) if value and _is_pem_content(value)) + if not inline: + return _CredentialPaths(client_cert, client_key, ca_cert) + + # mkdtemp is 0o700, so the 0o600 key file it holds is unreachable by other users. + directory = tempfile.mkdtemp(prefix=_PEM_DIR_PREFIX) + return _CredentialPaths( + client_cert_path=( + _write_pem(directory, _CLIENT_CERT_FILENAME, client_cert) if _is_pem_content(client_cert) else client_cert + ), + client_key_path=( + _write_pem(directory, _CLIENT_KEY_FILENAME, client_key) if _is_pem_content(client_key) else client_key + ), + ca_cert_path=( + _write_pem(directory, _CA_CERT_FILENAME, ca_cert) if ca_cert and _is_pem_content(ca_cert) else ca_cert + ), + ) + + +def load_billing_metrics_config( + *, license_data: Optional["EnterpriseLicenseData"], litellm_version: str +) -> Optional[BillingMetricsConfig]: + endpoint = os.getenv(ENDPOINT_ENV) + client_cert = os.getenv(CLIENT_CERT_ENV) + client_key = os.getenv(CLIENT_KEY_ENV) + # Optional: only for private/test collectors whose server cert is not on the + # public web PKI. The production collector needs no CA override. + ca_cert = os.getenv(CA_CERT_ENV) + + missing = [ + name + for name, value in ( + (ENDPOINT_ENV, endpoint), + (CLIENT_CERT_ENV, client_cert), + (CLIENT_KEY_ENV, client_key), + ) + if not value + ] + if not endpoint or not client_cert or not client_key: + verbose_proxy_logger.warning( + "Enterprise billing metrics disabled: licensed deployment missing config (%s)", + ", ".join(missing), + ) + return None + + try: + paths = _resolve_credential_paths(client_cert=client_cert, client_key=client_key, ca_cert=ca_cert) + except OSError as exc: + verbose_proxy_logger.warning( + "Enterprise billing metrics disabled: could not write inline certificate content to disk: %s", exc + ) + return None + + # Report the variable names, never their values. A value that is neither a + # readable path nor recognizable PEM is still secret material, and this + # warning would otherwise copy a client key straight into the proxy logs. + unreadable = [ + env_name + for env_name, path in ( + (CLIENT_CERT_ENV, paths.client_cert_path), + (CLIENT_KEY_ENV, paths.client_key_path), + (CA_CERT_ENV, paths.ca_cert_path), + ) + if path and not os.path.isfile(path) + ] + if unreadable: + verbose_proxy_logger.warning( + "Enterprise billing metrics disabled: %s did not resolve to a readable certificate file. " + "Set each to a file path, or to inline PEM content beginning with '%s'.", + ", ".join(unreadable), + _PEM_PREFIX, + ) + return None + + return BillingMetricsConfig( + endpoint=endpoint, + client_cert_path=paths.client_cert_path, + client_key_path=paths.client_key_path, + ca_cert_path=paths.ca_cert_path, + export_interval_ms=_export_interval_ms(), + litellm_version=litellm_version, + license_id=(license_data or {}).get("user_id"), + ) + + +class _ActiveRecorderRegistry: + """One-slot registry linking the factory-built recorder to the shutdown + hook; the middleware instance holding the recorder is not reachable from + proxy_shutdown_event.""" + + def __init__(self) -> None: + self._recorder: Optional[BillingMetricsRecorder] = None + + def set(self, recorder: BillingMetricsRecorder) -> None: + self._recorder = recorder + + def pop(self) -> Optional[BillingMetricsRecorder]: + recorder = self._recorder + self._recorder = None + return recorder + + +_ACTIVE_RECORDER = _ActiveRecorderRegistry() + + +def build_billing_metrics_recorder( + *, premium: bool, license_data: Optional["EnterpriseLicenseData"], litellm_version: str +) -> Optional[BillingMetricsRecorder]: + """Build the recorder, or None when the deployment is not licensed or metering is unconfigured.""" + if not premium: + # Debug, not warning: unlicensed is the common case and a warning here + # would be noise on every OSS proxy. Every other disable path warns. + verbose_proxy_logger.debug("Enterprise billing metrics disabled: deployment is not licensed") + return None + + config = load_billing_metrics_config(license_data=license_data, litellm_version=litellm_version) + if config is None: + return None + + try: + recorder = BillingMetricsRecorder(build_mtls_meter_provider(config)) + except Exception as exc: # noqa: BLE001 -- metering must never break proxy startup + verbose_proxy_logger.warning("Enterprise billing metrics disabled: failed to initialize exporter: %s", exc) + return None + _ACTIVE_RECORDER.set(recorder) + # The only positive signal that this component meters. Without it, a silent + # return above is indistinguishable from a working exporter in the logs, and + # a component that carries the cert but no license would look healthy. + verbose_proxy_logger.info( + "Enterprise billing metrics enabled: exporting to %s every %d ms", + config.endpoint, + config.export_interval_ms, + ) + return recorder + + +def shutdown_billing_metrics_recorder() -> None: + """Flush and stop the active recorder, if any. Idempotent; never raises.""" + recorder = _ACTIVE_RECORDER.pop() + if recorder is None: + return + try: + recorder.shutdown() + except Exception as exc: # noqa: BLE001 -- shutdown must never block or fail proxy exit + verbose_proxy_logger.warning("Enterprise billing metrics: final flush failed: %s", exc) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index b6a2d8d9069..1ed67a93d94 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2238,7 +2238,10 @@ async def apply_guardrail( if litellm_logging_obj is not None: _patch_logging_obj_for_guardrail(litellm_logging_obj, request) - request_data: dict = {"messages": request.messages} if request.messages else {} + request_data: dict = { + **({"messages": request.messages} if request.messages is not None else {}), + **({"metadata": request.metadata} if request.metadata is not None else {}), + } _input_type = _resolve_guardrail_input_type(active_guardrail, request.input_type) guardrailed_inputs = await active_guardrail.apply_guardrail( inputs={"texts": [request.text]}, diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py new file mode 100644 index 00000000000..73d31f7aec0 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/__init__.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import ( + GuardrailEventHooks, + Mode, + SupportedGuardrailIntegrations, +) + +from .compresr import CompresrGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def _coerce_event_hook( + mode: str | list[str] | Mode, +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: + if isinstance(mode, Mode): + return mode + if isinstance(mode, list): + return [GuardrailEventHooks(item) for item in mode] + return GuardrailEventHooks(mode) + + +def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object: + if optional_params is not None: + value = getattr(optional_params, attribute_name, None) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> CompresrGuardrail: + import litellm + + optional_params = getattr(litellm_params, "optional_params", None) + + _callback = CompresrGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + model=litellm_params.model, + target_compression_ratio=_get_optional_value(litellm_params, optional_params, "target_compression_ratio"), + coarse=_get_optional_value(litellm_params, optional_params, "coarse"), + min_chars_to_compress=_get_optional_value(litellm_params, optional_params, "min_chars_to_compress"), + compress_tool_outputs=_get_optional_value(litellm_params, optional_params, "compress_tool_outputs"), + compress_system=_get_optional_value(litellm_params, optional_params, "compress_system"), + compress_history=_get_optional_value(litellm_params, optional_params, "compress_history"), + compress_last_user=_get_optional_value(litellm_params, optional_params, "compress_last_user"), + enable_retrieval=_get_optional_value(litellm_params, optional_params, "enable_retrieval"), + max_bytes_per_call=_get_optional_value(litellm_params, optional_params, "max_bytes_per_call"), + allow_bypass_header=_get_optional_value(litellm_params, optional_params, "allow_bypass_header"), + dynamic=_get_optional_value(litellm_params, optional_params, "dynamic"), + dynamic_min_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_min_ratio"), + dynamic_max_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_max_ratio"), + compression_params=_get_optional_value(litellm_params, optional_params, "compression_params"), + guardrail_name=guardrail["guardrail_name"], + event_hook=_coerce_event_hook(litellm_params.mode), + default_on=litellm_params.default_on or False, + unreachable_fallback=litellm_params.unreachable_fallback, + ) + litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped + _callback + ) + return _callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.COMPRESR.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.COMPRESR.value: CompresrGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py new file mode 100644 index 00000000000..a95bdb670c3 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -0,0 +1,1214 @@ +"""Compresr guardrail — query-aware, recoverable context compression. + +Compresses bulky message content (tool outputs by default) through the +Compresr API before the request reaches the LLM. Each compressed message +carries a hash marker; a ``compresr_retrieve`` tool is injected so the model +can fetch the original content back through the agentic loop when the +compressed version is not enough — making compression recoverable instead +of lossy. + +Unlike gateway-side compressors that operate on whole message lists, each +target is compressed *query-aware*: the query sent to Compresr is the intent +of the tool call that produced the message (``name + arguments``, resolved +via ``tool_call_id``), falling back to the last user message. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import ipaddress +import json +import time +from collections import Counter, OrderedDict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal +from urllib.parse import urlparse + +import httpx +from fastapi import HTTPException +from httpx import Response as HttpxResponse +from typing_extensions import TypeGuard + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.litellm_core_utils.prompt_templates.factory import ( + get_attribute_or_key, + get_tool_calls_from_response, + has_tool_with_name, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.base import ( + GuardrailConfigModel, + ) + +BYPASS_HEADER = "x-compresr-bypass" +COMPRESR_RETRIEVE_TOOL_NAME = "compresr_retrieve" +DEFAULT_API_BASE = "https://api.compresr.ai" +DEFAULT_COMPRESSION_MODEL = "latte_v2" +DEFAULT_TARGET_COMPRESSION_RATIO = 0.5 +DEFAULT_MIN_CHARS_TO_COMPRESS = 500 +_ORIGINALS_TTL_SECONDS = 15 * 60 +_NO_SCOPE_WARNING_INTERVAL_SECONDS = 15 * 60 +_MAX_TRACKED_CALLS = 256 +_DEFAULT_MAX_BYTES_PER_CALL = 10 * 1024 * 1024 +# Aggregate ceiling across all recovery-store entries. max_bytes_per_call only +# bounds a single call; this caps the whole store so many calls cannot exhaust it. +_MAX_TOTAL_STORE_BYTES = 256 * 1024 * 1024 +# Max compresr_retrieve calls expanded into a single follow-up (repeats deduped). +_MAX_RETRIEVALS_PER_LOOP = 8 +# The shared client's 600s read timeout is far too long for an on-request +# guardrail; bound the compress call so a stall hits the fail policy quickly. +_COMPRESS_TIMEOUT_SECONDS = 60.0 +_SOURCE_TAG = "integration:litellm" +# Request-content fields the compression_params passthrough must never +# override — they carry the actual message content/queries being compressed. +_RESERVED_COMPRESSION_PARAM_KEYS = frozenset({"context", "query", "inputs"}) +_BLOCKED_METADATA_HOSTS = frozenset( + { + "metadata.google.internal", + "metadata.goog", + "metadata.azure.com", + "metadata.azure.internal", + } +) +_BLOCKED_METADATA_IPS = frozenset( + ipaddress.ip_address(ip) for ip in ("169.254.169.254", "fd00:ec2::254", "100.100.100.200", "168.63.129.16") +) + + +def _parse_ip_literal(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None: + """Parse ``host`` as an IP literal, covering the alternate spellings the + socket layer accepts (decimal/hex single-integer IPv4, IPv4-mapped IPv6) + so a blocked address cannot be smuggled past a string comparison.""" + try: + addr = ipaddress.ip_address(host) + except ValueError: + try: + addr = ipaddress.ip_address(int(host, 0)) + except (TypeError, ValueError): + return None + if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: + return addr.ipv4_mapped + return addr + + +def _validate_api_base(url: str) -> str: + """Return ``url`` if it passes basic outbound-target checks, else raise. + + Best-effort defense in depth for a mis/maliciously-configured ``api_base``: + rejects non-http(s) schemes and cloud-metadata IPs/hosts (incl. alternate IP + encodings); private ranges are allowed for on-prem deployments. NOT a complete + SSRF control — no DNS resolution, and the shared client follows redirects and + re-resolves DNS (TOCTOU / rebinding); ``api_base`` is trusted operator config, + so this is an accepted limitation. + """ + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"Compresr guardrail api_base must be http or https, got scheme={parsed.scheme!r}") + host = (parsed.hostname or "").lower() + if not host: + raise ValueError("Compresr guardrail api_base has no host") + ip_literal = _parse_ip_literal(host) + if host in _BLOCKED_METADATA_HOSTS or (ip_literal is not None and ip_literal in _BLOCKED_METADATA_IPS): + raise ValueError(f"Compresr guardrail api_base {host!r} is a blocked cloud-metadata host") + return url + + +def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, dict) + + +def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, list) + + +def _content_to_text(content: object) -> str: + """Collapse a message ``content`` (str or list-of-parts) to plain text. + + For the multimodal list shape, joins ``{type: "text", text: ...}`` parts + with blank-line separators; non-text parts are ignored. + """ + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + text = part.get("text") + if isinstance(text, str): + parts.append(text) + return "\n\n".join(parts) + return "" + + +def _replace_text_in_content(content: object, new_text: str) -> object: + """Write ``new_text`` back into a ``content`` value, preserving shape. + + ``str`` content is replaced directly. For list-of-parts content the first + text part carries ``new_text``, later text parts are dropped, and + non-text parts (images, audio, files) pass through untouched. + """ + if isinstance(content, str): + return new_text + if isinstance(content, list): + out: list[object] = [] + replaced = False + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + if not replaced: + out.append({**part, "text": new_text}) + replaced = True + continue + out.append(part) + if not replaced: + out.insert(0, {"type": "text", "text": new_text}) + return out + return new_text + + +def _render_tool_intent(fn: dict[str, object]) -> str: + name = str(fn.get("name") or "").strip() + args = fn.get("arguments") + if isinstance(args, dict): + try: + args_str = json.dumps(args, separators=(",", ":")) + except (TypeError, ValueError): + args_str = str(args) + else: + args_str = str(args).strip() if args is not None else "" + if name and args_str: + return f"{name}: {args_str}" + return name or args_str + + +def _query_for_target(messages: list[dict[str, object]], target_idx: int, fallback: str) -> str: + """Query used to compress ``messages[target_idx]``. + + Tool/function outputs are compressed against the intent of the tool call + that produced them (found via ``tool_call_id`` on a prior assistant + message); everything else uses the last user message. + """ + msg = messages[target_idx] + if msg.get("role") not in ("tool", "function"): + return fallback + + tool_call_id = msg.get("tool_call_id") + fn_name = msg.get("name") + for j in range(target_idx - 1, -1, -1): + prev = messages[j] + if prev.get("role") != "assistant": + continue + tool_calls = prev.get("tool_calls") + if isinstance(tool_calls, list): + for tc in tool_calls: + if not isinstance(tc, dict): + continue + if tool_call_id and tc.get("id") == tool_call_id: + fn = tc.get("function") + intent = _render_tool_intent(fn if isinstance(fn, dict) else {}) + if intent: + return intent + # Legacy function_call fallback: require a name match, else an earlier + # function_call turn would attribute the wrong intent. + fc = prev.get("function_call") + if isinstance(fc, dict) and fn_name and fc.get("name") == fn_name: + intent = _render_tool_intent(fc) + if intent: + return intent + return fallback + + +def _safe_int(value: object) -> int: + """Parse a token-stat field defensively. + + A malformed-but-200 response must not raise here: ``_call_compress`` has + already returned successfully, so the fail_open/fail_closed decision is + behind us. A bare ``int()`` on a non-numeric field would surface as an + unhandled 500 even when ``fail_open`` is configured. + """ + try: + return int(value) if value is not None else 0 + except (TypeError, ValueError): + return 0 + + +def _safe_response_text(response: object, limit: int = 500) -> str: + """Read a response body for error logging without letting the read itself + raise. A corrupt ``Content-Encoding`` makes ``httpx``'s ``.text`` raise a + ``DecodingError``; if that happened while building a failure detail it would + turn an already-handled error into an unhandled 500.""" + try: + text = getattr(response, "text", "") + except httpx.DecodingError: + return "" + return (text or "")[:limit] + + +def _content_hash(text: str) -> str: + # surrogatepass so a lone surrogate in untrusted content (valid via a JSON + # \uXXXX escape) hashes instead of raising past the fail policy. + return hashlib.sha256(text.encode("utf-8", "surrogatepass")).hexdigest()[:24] + + +def _entry_bytes(originals: dict[str, str]) -> int: + """UTF-8 byte size of one recovery-store entry (surrogatepass, like _content_hash).""" + return sum(len(value.encode("utf-8", "surrogatepass")) for value in originals.values()) + + +def _display_hash(hash_value: str) -> str: + """Bound a model-supplied hash for logs/fallback text. A real marker hash is + 24 hex chars; a prompt-injected ``compresr_retrieve`` call could pass a huge + or control-character-laden string, so strip non-printables (no forged log + lines / ANSI escapes) and cap length before echoing into logs and the + conversation.""" + printable = "".join(ch for ch in hash_value if ch.isprintable()) + return printable if len(printable) <= 32 else f"{printable[:32]}…" + + +def _recovery_marker(hash_value: str) -> str: + return ( + f"\n\n[compresr hash={hash_value}: parts of this content were compressed " + f"away. If you need the full original, call the " + f"{COMPRESR_RETRIEVE_TOOL_NAME} tool with this hash.]" + ) + + +def _build_compresr_retrieve_tool() -> dict[str, object]: + return { + "type": "function", + "function": { + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "description": ( + "Retrieve the original, uncompressed content behind a Compresr " + "compression marker. Call this when a compression marker's hash " + "points at content you need in full." + ), + "parameters": { + "type": "object", + "properties": { + "hash": { + "type": "string", + "description": "The 24-character hex hash from the compression marker.", + }, + }, + "required": ["hash"], + }, + }, + } + + +def has_compresr_retrieve_tool(tools: object) -> bool: + return has_tool_with_name(tools, COMPRESR_RETRIEVE_TOOL_NAME) + + +def _merge_retrieve_tool(existing_tools: object) -> list[object] | None: + """The request's tools plus the retrieve tool, or None when the incoming + shape is not a list (leave the caller's tools untouched; markers stay + inert text).""" + if existing_tools is not None and not isinstance(existing_tools, list): + return None + retrieve_tool = _build_compresr_retrieve_tool() + if existing_tools is None: + return [retrieve_tool] + if has_compresr_retrieve_tool(existing_tools): + return list(existing_tools) + return list(existing_tools) + [retrieve_tool] + + +def _extract_compresr_tool_calls(response: object) -> list[dict[str, object]]: + return [ + {"id": tc.get("id"), "type": "function", "name": tc.get("name"), "arguments": tc.get("arguments", {})} + for tc in get_tool_calls_from_response(response) + if tc.get("name") == COMPRESR_RETRIEVE_TOOL_NAME + ] + + +def _resolve_call_id(logging_obj: object) -> str | None: + """The call id from the framework logging object. + + This value ultimately derives from the client-settable ``x-litellm-call-id`` + header and is echoed back in responses, so it is NOT a trust boundary on its + own — ``_scoped_store_key`` prefixes it with the caller's virtual-key hash to + partition the recovery store per tenant. Request-body/kwargs call ids are + deliberately not consulted here. + """ + logging_call_id = getattr(logging_obj, "litellm_call_id", None) + if isinstance(logging_call_id, str) and logging_call_id: + return logging_call_id + return None + + +def _caller_scope(logging_obj: object) -> str: + """The caller's virtual-key hash, used to partition the recovery store. + + Trust is anchored on the ``UserAPIKeyAuth`` object the proxy sets + server-side (``metadata.user_api_key_auth``, litellm_pre_call_utils). Its + ``api_key`` is the hash of the authenticated key. Both metadata spellings + are scanned (``/v1/messages`` and ``/v1/responses`` carry it under + ``litellm_metadata``), but the bare ``user_api_key`` *string* is never + trusted on its own: a JSON request body can place one in the client-supplied + ``metadata`` field, which is only sanitized on the route's canonical + container. Returns "" when the proxy runs without per-key auth, in which case + all traffic is a single trust domain and the call id alone suffices. + """ + details = getattr(logging_obj, "model_call_details", None) + if not _is_str_object_dict(details): + return "" + litellm_params = details.get("litellm_params") + for container in (litellm_params, details): + if not _is_str_object_dict(container): + continue + for meta_key in ("metadata", "litellm_metadata"): + metadata = container.get(meta_key) + if not _is_str_object_dict(metadata): + continue + auth = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and isinstance(auth.api_key, str) and auth.api_key: + return auth.api_key + return "" + + +def _scoped_store_key(logging_obj: object) -> str | None: + """Key for the recovery store: caller identity plus framework call id. + + Keying on the call id alone is unsafe: it comes from the client-settable + ``x-litellm-call-id`` header and is echoed back in responses, so one caller + could read or evict another's originals by reusing the id. Prefixing the + unforgeable virtual-key hash binds each entry to the tenant that created it. + Returns None when there is no call id, which disables recovery for the call. + """ + call_id = _resolve_call_id(logging_obj) + if call_id is None: + return None + scope = _caller_scope(logging_obj) + return f"{scope}\x00{call_id}" if scope else call_id + + +def _is_responses_api_response(response: object) -> bool: + return isinstance(get_attribute_or_key(response, "output", None), list) + + +def _is_anthropic_messages_response(response: object) -> bool: + return isinstance(get_attribute_or_key(response, "content", None), list) + + +def _assistant_text_from_response(response: object) -> str | None: + """The assistant's natural-language text from a model response, across chat, + Anthropic, and Responses shapes. Preserved when the turn is rebuilt for the + retrieval follow-up so the model's reasoning is not lost.""" + choices = get_attribute_or_key(response, "choices", None) + if isinstance(choices, list) and choices: + message = get_attribute_or_key(choices[0], "message", None) + if message is not None: + text = _content_to_text(get_attribute_or_key(message, "content", None)) + if text: + return text + content = get_attribute_or_key(response, "content", None) + if isinstance(content, list): + parts = [ + text + for block in content + if get_attribute_or_key(block, "type", None) == "text" + for text in (get_attribute_or_key(block, "text", None),) + if isinstance(text, str) and text + ] + if parts: + return "".join(parts) + output = get_attribute_or_key(response, "output", None) + if isinstance(output, list): + parts = [] + for item in output: + if get_attribute_or_key(item, "type", None) != "message": + continue + item_content = get_attribute_or_key(item, "content", None) + if not isinstance(item_content, list): + continue + for chunk in item_content: + if get_attribute_or_key(chunk, "type", None) == "output_text": + text = get_attribute_or_key(chunk, "text", None) + if isinstance(text, str) and text: + parts.append(text) + if parts: + return "".join(parts) + return None + + +def _build_assistant_message_from_response( + response: object, + retrieved: list[tuple[dict[str, object], str]], +) -> dict[str, object]: + """Rebuild the chat-completions assistant turn for the retrieval follow-up. + + Only the ``compresr_retrieve`` calls are echoed, each answered by a tool + result below. Other tool calls made in the same turn are omitted on purpose: + the follow-up re-runs the model with the recovered content so it re-plans + them. Echoing them would leave tool_calls with no matching tool result and + the provider would reject the request. + """ + return { + "role": "assistant", + "content": _assistant_text_from_response(response), + "tool_calls": [ + { + "id": tool_call.get("id"), + "type": "function", + "function": { + "name": tool_call.get("name"), + "arguments": json.dumps(tool_call.get("arguments", {})), + }, + } + for tool_call, _ in retrieved + ], + } + + +def _build_anthropic_followup_messages( + response: object, + retrieved: list[tuple[dict[str, object], str]], +) -> list[dict[str, object]]: + """Anthropic requires the tool_use block echoed back in an assistant + message paired with a tool_result block keyed by the same tool_use_id. The + assistant text is preserved; non-retrieve tool calls are re-planned by the + follow-up (see _build_assistant_message_from_response).""" + assistant_content: list[dict[str, object]] = [] + text = _assistant_text_from_response(response) + if text: + assistant_content.append({"type": "text", "text": text}) + assistant_content.extend( + { + "type": "tool_use", + "id": tool_call.get("id"), + "name": tool_call.get("name"), + "input": tool_call.get("arguments", {}), + } + for tool_call, _ in retrieved + ) + assistant_message: dict[str, object] = {"role": "assistant", "content": assistant_content} + user_message: dict[str, object] = { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": tool_call.get("id"), "content": content} + for tool_call, content in retrieved + ], + } + return [assistant_message, user_message] + + +def _build_responses_followup_items( + response: object, + retrieved: list[tuple[dict[str, object], str]], +) -> list[dict[str, object]]: + """The Responses API requires the model's function_call echoed back paired + with a function_call_output keyed by the same call_id. The assistant text is + preserved; non-retrieve tool calls are re-planned by the follow-up.""" + items: list[dict[str, object]] = [] + text = _assistant_text_from_response(response) + if text: + items.append({"role": "assistant", "content": text}) + for tool_call, content in retrieved: + call_id = tool_call.get("id") + items.append( + { + "type": "function_call", + "call_id": call_id, + "name": tool_call.get("name"), + "arguments": json.dumps(tool_call.get("arguments", {})), + } + ) + items.append({"type": "function_call_output", "call_id": call_id, "output": content}) + return items + + +@dataclass +class _CompressionResult: + """Outcome of applying compression results to a message list.""" + + compressed_messages: list[dict[str, object]] + originals: dict[str, str] = field(default_factory=dict) + # original text -> compressed text, plus the machinery the Responses `texts` + # mirror needs to replace only where it is unambiguous. + text_replacements: dict[str, str] = field(default_factory=dict) + replaced_text_counts: dict[str, int] = field(default_factory=dict) + ambiguous_texts: set[str] = field(default_factory=set) + messages_compressed: int = 0 + tokens_before: int = 0 + tokens_after: int = 0 + + +class CompresrGuardrail(CustomGuardrail): + def __init__( + self, + api_base: str | None = None, + api_key: str | None = None, + model: str | None = None, + target_compression_ratio: float | None = None, + coarse: bool | None = None, + min_chars_to_compress: int | None = None, + compress_tool_outputs: bool | None = None, + compress_system: bool | None = None, + compress_history: bool | None = None, + compress_last_user: bool | None = None, + enable_retrieval: bool | None = None, + guardrail_name: str | None = None, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, + default_on: bool = False, + unreachable_fallback: str | None = None, + max_bytes_per_call: int | None = None, + allow_bypass_header: bool | None = None, + dynamic: bool | None = None, + dynamic_min_ratio: float | None = None, + dynamic_max_ratio: float | None = None, + compression_params: dict[str, object] | None = None, + ): + raw_api_base = (api_base or get_secret_str("COMPRESR_API_BASE") or DEFAULT_API_BASE).rstrip("/") + self.compresr_api_base = _validate_api_base(raw_api_base) + self.compresr_api_key = api_key or get_secret_str("COMPRESR_API_KEY") + if not self.compresr_api_key: + raise ValueError( + "Compresr guardrail requires an API key. Set `api_key` in the " + "guardrail config or the COMPRESR_API_KEY env var." + ) + self.compression_model = model or DEFAULT_COMPRESSION_MODEL + self.target_compression_ratio = ( + DEFAULT_TARGET_COMPRESSION_RATIO if target_compression_ratio is None else target_compression_ratio + ) + self.coarse = True if coarse is None else coarse + self.min_chars_to_compress = ( + DEFAULT_MIN_CHARS_TO_COMPRESS if min_chars_to_compress is None else min_chars_to_compress + ) + self.compress_tool_outputs = True if compress_tool_outputs is None else compress_tool_outputs + self.compress_system = False if compress_system is None else compress_system + self.compress_history = False if compress_history is None else compress_history + self.compress_last_user = False if compress_last_user is None else compress_last_user + self.enable_retrieval = True if enable_retrieval is None else enable_retrieval + self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( + "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" + ) + self.max_bytes_per_call = _DEFAULT_MAX_BYTES_PER_CALL if max_bytes_per_call is None else max_bytes_per_call + if self.max_bytes_per_call < 0: + raise ValueError("max_bytes_per_call must be >= 0 (0 disables the cap; positive values enforce it)") + self.allow_bypass_header = False if allow_bypass_header is None else allow_bypass_header + # Dynamic (adaptive) compression — latte_v2 only, on by default: the server + # picks the ratio per input instead of honoring target_compression_ratio. + self.dynamic = True if dynamic is None else dynamic + self.dynamic_min_ratio = dynamic_min_ratio + self.dynamic_max_ratio = dynamic_max_ratio + # Passthrough of extra compression params forwarded verbatim, so a new + # Compresr feature works without changing this guardrail. Named fields win; + # request-content fields are stripped. + reserved_keys = _RESERVED_COMPRESSION_PARAM_KEYS.intersection(compression_params or {}) + if reserved_keys: + verbose_proxy_logger.warning( + "Compresr: ignoring reserved compression_params keys %s", sorted(reserved_keys) + ) + self.compression_params: dict[str, object] = { + k: v for k, v in (compression_params or {}).items() if k not in _RESERVED_COMPRESSION_PARAM_KEYS + } + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + self._originals_by_call_id: OrderedDict[str, tuple[dict[str, str], float]] = OrderedDict() + # Running byte size of the store, kept in sync to enforce the global cap cheaply. + self._store_total_bytes = 0 + # Rate-limits the "recovery skipped, no auth scope" warning so an ongoing + # misconfiguration stays visible without flooding hot-path logs. + self._no_scope_warning_expiry = 0.0 + if self.enable_retrieval: + verbose_proxy_logger.warning( + "Compresr: enable_retrieval is on; the recovery store is per-process. " + "For multi-worker deployments, set enable_retrieval=false or run with --workers 1." + ) + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + ) + + def _should_bypass(self, request_data: dict) -> bool: + if not self.allow_bypass_header: + return False + psr = request_data.get("proxy_server_request") + if not _is_str_object_dict(psr): + return False + headers = psr.get("headers") + if not _is_str_object_dict(headers): + return False + return str(headers.get(BYPASS_HEADER)).lower() == "true" + + def _request_headers(self) -> dict[str, str]: + return { + "Content-Type": "application/json", + "X-API-Key": self.compresr_api_key or "", + } + + def _handle_compress_failure(self, error: str, log_detail: dict[str, object]) -> None: + """fail_open logs and returns (caller forwards uncompressed); + fail_closed raises. ``log_detail`` may include upstream response bodies + and is written only to server logs; the raised ``HTTPException`` carries + a generic message so a malicious ``api_base`` cannot exfiltrate response + bytes through the client-visible error.""" + if self.unreachable_fallback == "fail_open": + verbose_proxy_logger.warning( + "Compresr: %s; fail_open configured, forwarding request uncompressed. detail=%s", + error, + log_detail, + ) + return + verbose_proxy_logger.error("Compresr: %s. detail=%s", error, log_detail) + raise HTTPException(status_code=502, detail={"error": error}) + + def _evict_oldest(self) -> None: + """Drop the front (oldest) entry and decrement the running byte total.""" + _key, (evicted, _expiry) = self._originals_by_call_id.popitem(last=False) + self._store_total_bytes -= _entry_bytes(evicted) + + def _prune_originals(self) -> None: + # Insertion order == expiry order (shared TTL); prune from the front. + now = time.monotonic() + store = self._originals_by_call_id + while store and store[next(iter(store))][1] <= now: + self._evict_oldest() + while len(store) > _MAX_TRACKED_CALLS: + self._evict_oldest() + # Global byte budget; keep the most-recent entry so the current call's + # originals survive (a single call is already bounded by max_bytes_per_call). + while len(store) > 1 and self._store_total_bytes > _MAX_TOTAL_STORE_BYTES: + self._evict_oldest() + + def _existing_originals(self, store_key: str | None) -> dict[str, str]: + """Originals already stored under this key, so the per-call byte budget + can account for an earlier turn that reused the store key.""" + if store_key is None: + return {} + return self._originals_by_call_id.get(store_key, ({}, 0.0))[0] + + def _store_originals(self, store_key: str, originals: dict[str, str]) -> None: + existing, _ = self._originals_by_call_id.get(store_key, ({}, 0.0)) + merged = self._bound_call_bytes({**existing, **originals}) + # Keep the running total in sync: drop the overwritten entry, add the new one. + self._store_total_bytes += _entry_bytes(merged) - _entry_bytes(existing) + self._originals_by_call_id[store_key] = ( + merged, + time.monotonic() + _ORIGINALS_TTL_SECONDS, + ) + self._originals_by_call_id.move_to_end(store_key) + self._prune_originals() + + def _bound_call_bytes(self, merged: dict[str, str]) -> dict[str, str]: + """Drop oldest entries (dict insertion order) until the aggregate byte + size fits ``self.max_bytes_per_call``. Prevents one call with many + large tool outputs from growing proxy memory without bound.""" + if self.max_bytes_per_call <= 0: + return merged + total = _entry_bytes(merged) + if total <= self.max_bytes_per_call: + return merged + bounded = dict(merged) + for key in list(bounded.keys()): + if total <= self.max_bytes_per_call: + break + total -= len(bounded[key].encode("utf-8", "surrogatepass")) + del bounded[key] + verbose_proxy_logger.warning("Compresr: originals-store byte cap hit, evicted hash=%s", key) + return bounded + + def _retrieve_original(self, store_key: str | None, hash_value: str) -> str | None: + """Stored original for a marker hash, or None if not issued for this + request (unknown, expired, or from another caller's scope).""" + if store_key: + originals, expiry = self._originals_by_call_id.get(store_key, ({}, 0.0)) + if expiry > time.monotonic() and hash_value in originals: + return originals[hash_value] + verbose_proxy_logger.warning( + "Compresr retrieve: rejecting hash=%s (not issued for this request, or expired)", + _display_hash(hash_value), + ) + return None + + def _resolve_retrievals( + self, store_key: str | None, tool_calls: list[dict[str, object]] + ) -> tuple[list[tuple[dict[str, object], str]], bool]: + """Resolve compresr_retrieve calls to (call, result_text) pairs, deduping + repeated hashes and capping the count so the follow-up cannot be amplified. + The bool is True iff at least one call resolved to real stored content.""" + retrieved: list[tuple[dict[str, object], str]] = [] + seen: set[str] = set() + resolved_any = False + for idx, tc in enumerate(tool_calls): + arguments = tc.get("arguments", {}) + hash_value = str(arguments.get("hash", "")) if isinstance(arguments, dict) else "" + if idx >= _MAX_RETRIEVALS_PER_LOOP: + result = "[compresr: retrieval limit reached for this turn]" + elif hash_value in seen: + result = "[compresr: already retrieved above for this hash]" + else: + content = self._retrieve_original(store_key, hash_value) + if content is None: + result = f"[compresr: hash={_display_hash(hash_value)} not found, expired, or not issued for this request]" + else: + seen.add(hash_value) + resolved_any = True + result = content + verbose_proxy_logger.debug("Compresr retrieve: hash=%s -> %d chars", _display_hash(hash_value), len(result)) + retrieved.append((tc, result)) + return retrieved, resolved_any + + async def _call_compress( + self, + contexts: list[str], + queries: list[str], + ) -> list[dict[str, object]] | None: + """Compress ``contexts`` (query-aware). Returns one result dict per + context, or None when the service failed and fail_open applies.""" + common: dict[str, object] = { + # Passthrough first so the named fields below always win on collision. + **self.compression_params, + "compression_model_name": self.compression_model, + "target_compression_ratio": self.target_compression_ratio, + "coarse": self.coarse, + "dynamic": self.dynamic, + "source": _SOURCE_TAG, + } + # Only send the bounds the operator actually set; otherwise let the + # server apply its own floor/ceiling. + if self.dynamic_min_ratio is not None: + common["dynamic_min_ratio"] = self.dynamic_min_ratio + if self.dynamic_max_ratio is not None: + common["dynamic_max_ratio"] = self.dynamic_max_ratio + if len(contexts) == 1: + url = f"{self.compresr_api_base}/api/compress/question-specific/" + payload: dict[str, object] = { + "context": contexts[0], + "query": queries[0], + **common, + } + else: + url = f"{self.compresr_api_base}/api/compress/question-specific/batch" + payload = { + "inputs": [{"context": ctx, "query": q} for ctx, q in zip(contexts, queries)], + **common, + } + + try: + raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped + url=url, + json=payload, + headers=self._request_headers(), + timeout=_COMPRESS_TIMEOUT_SECONDS, + ) + except asyncio.CancelledError: + raise + except httpx.HTTPStatusError as e: + # The shared handler calls raise_for_status(), so a non-2xx reply arrives + # here as an error carrying the upstream body + our API key header; route + # it through the fail policy so none of that reaches the client. + resp = getattr(e, "response", None) + self._handle_compress_failure( + "Compresr compression service returned an error", + { + "status_code": getattr(resp, "status_code", None), + "body": _safe_response_text(resp), + }, + ) + return None + except (httpx.RequestError, litellm.Timeout) as e: + # Every request-side httpx failure is a RequestError; route the whole + # class through the fail policy so none escapes as a 500 under fail_open. + # (HTTPStatusError is handled above and is not a RequestError.) + self._handle_compress_failure( + "Compresr compression service request failed", + {"detail": str(e)}, + ) + return None + if raw_response is None or not 200 <= raw_response.status_code < 300: + self._handle_compress_failure( + "Compresr compression service returned an error", + { + "status_code": getattr(raw_response, "status_code", None), + "body": _safe_response_text(raw_response), + }, + ) + return None + + try: + body: object = raw_response.json() + except (ValueError, httpx.DecodingError, RecursionError): + # RecursionError: a deeply nested JSON body overflows the parser; + # route it through the fail policy rather than let it escape as a 500. + self._handle_compress_failure( + "Compresr compression service returned an unreadable response", + {"body": _safe_response_text(raw_response)}, + ) + return None + if not _is_str_object_dict(body) or not _is_str_object_dict(body.get("data")): + self._handle_compress_failure( + "Compresr compression service returned unexpected response shape", + {"body": _safe_response_text(raw_response)}, + ) + return None + data: dict[str, object] = body["data"] # pyright: ignore[reportAssignmentType] # dict-guarded above; subscript does not narrow + + if len(contexts) == 1: + return [data] + results = data.get("results") + if ( + not _is_object_list(results) + or len(results) != len(contexts) + or not all(_is_str_object_dict(r) for r in results) + ): + # Anything but a 1:1 dict-per-context mapping would misalign + # results with their target messages. + self._handle_compress_failure( + "Compresr batch response missing or mismatched 'results'", + {"expected": len(contexts), "got": len(results) if _is_object_list(results) else None}, + ) + return None + return results # pyright: ignore[reportReturnType] # every element dict-checked above; list[object] does not narrow + + def _select_targets(self, messages: list[dict[str, object]], query_idx: int | None) -> list[int]: + """Indices of messages whose text content should be compressed.""" + targets: list[int] = [] + for idx, msg in enumerate(messages): + if idx == query_idx and not self.compress_last_user: + continue + role = msg.get("role") + if role in ("tool", "function"): + if not self.compress_tool_outputs: + continue + elif role == "system": + if not self.compress_system: + continue + elif role == "user": + if idx != query_idx and not self.compress_history: + continue + else: + continue + if len(_content_to_text(msg.get("content"))) < self.min_chars_to_compress: + continue + targets.append(idx) + return targets + + @staticmethod + def _extract_fallback_query( + messages: list[dict[str, object]], + ) -> tuple[str, int | None]: + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") == "user": + return _content_to_text(messages[idx].get("content")), idx + return "", None + + def _apply_compression_results( + self, + messages: list[dict[str, object]], + targets: list[int], + contexts: list[str], + results: list[dict[str, object]], + recovery_enabled: bool, + existing_originals: dict[str, str] | None = None, + ) -> _CompressionResult: + """Write each compression result into a copy of ``messages``. + + A result is a real compression only when it is a non-empty string that + differs from the original; identical text is treated as a no-op so an + untouched request is not needlessly rewritten downstream. + """ + out = _CompressionResult(compressed_messages=list(messages)) + existing = existing_originals or {} + cap = self.max_bytes_per_call + # Seed with what is already stored under this store key: markers are + # attached only while the store (existing + this call's originals) stays + # within the cap, so _store_originals never has to evict a hash this call + # just shipped a marker for -- including on a later turn that reuses the + # store key. A hash already stored (or repeated here) costs no new bytes. + recovery_bytes = _entry_bytes(existing) + for target_idx, original_text, result in zip(targets, contexts, results): + compressed_text = result.get("compressed_context") + if not isinstance(compressed_text, str) or not compressed_text or compressed_text == original_text: + continue + out.messages_compressed += 1 + if recovery_enabled: + hash_value = _content_hash(original_text) + already_stored = hash_value in existing or hash_value in out.originals + new_bytes = 0 if already_stored else len(original_text.encode("utf-8", "surrogatepass")) + if cap <= 0 or recovery_bytes + new_bytes <= cap: + recovery_bytes += new_bytes + out.originals[hash_value] = original_text + compressed_text += _recovery_marker(hash_value) + previous = out.text_replacements.get(original_text) + if previous is not None and previous != compressed_text: + # Two targets with identical text but different query-specific + # compressions; a value-keyed replacement cannot tell them apart. + out.ambiguous_texts.add(original_text) + else: + out.text_replacements[original_text] = compressed_text + out.replaced_text_counts[original_text] = out.replaced_text_counts.get(original_text, 0) + 1 + original_msg = out.compressed_messages[target_idx] + out.compressed_messages[target_idx] = { + **original_msg, + "content": _replace_text_in_content(original_msg.get("content"), compressed_text), + } + out.tokens_before += _safe_int(result.get("original_tokens")) + out.tokens_after += _safe_int(result.get("compressed_tokens")) + return out + + @staticmethod + def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None: + """Compressed content mirrored into the Responses `texts` channel. + + The chat/Anthropic handlers round-trip ``structured_messages``; the + Responses translation cannot rebuild its input from chat messages and + instead writes back through ``texts``. This matches by value, so a + replacement is applied only when it is unambiguous: one compression per + text, and every occurrence in ``texts`` accounted for by a compressed + target. Anything else is left uncompressed rather than risk a wrong or + out-of-policy replacement. Returns None when nothing safe applies. + """ + if not applied.text_replacements or not isinstance(input_texts, list): + return None + counts = Counter(text for text in input_texts if isinstance(text, str)) + safe = { + text: replacement + for text, replacement in applied.text_replacements.items() + if text not in applied.ambiguous_texts and counts.get(text) == applied.replaced_text_counts.get(text) + } + if not safe: + return None + return [safe.get(text, text) if isinstance(text, str) else text for text in input_texts] + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + if self._should_bypass(request_data): + verbose_proxy_logger.debug("Compresr: %s header set; skipping compression", BYPASS_HEADER) + return inputs + + structured_messages = inputs.get("structured_messages") + if not _is_object_list(structured_messages) or not structured_messages: + return inputs + messages = [m for m in structured_messages if _is_str_object_dict(m)] + if len(messages) != len(structured_messages): + return inputs + + fallback_query, query_idx = self._extract_fallback_query(messages) + targets: list[int] = [] + queries: list[str] = [] + for idx in self._select_targets(messages, query_idx): + query = _query_for_target(messages, idx, fallback_query) + # latte models require a non-empty query; leave targets we cannot + # derive one for uncompressed rather than erroring. + if not query.strip(): + continue + targets.append(idx) + queries.append(query) + if not targets: + verbose_proxy_logger.debug("Compresr: no messages eligible for compression") + return inputs + + contexts = [_content_to_text(messages[idx].get("content")) for idx in targets] + + start_time = time.monotonic() + results = await self._call_compress(contexts=contexts, queries=queries) + end_time = time.monotonic() + if results is None: # service failed, fail_open configured + return inputs + + # Recovery needs a per-tenant scope; without per-key auth the key would fall + # back to the client-settable call id (cross-tenant reads), so skip it. + store_key = _scoped_store_key(logging_obj) + scope = _caller_scope(logging_obj) + recovery_enabled = self.enable_retrieval and store_key is not None and bool(scope) + if self.enable_retrieval and not scope and time.monotonic() >= self._no_scope_warning_expiry: + # Surface the silent no-recovery case (compressed, but no auth scope + # to inject the retrieve tool), re-warning once per interval. + self._no_scope_warning_expiry = time.monotonic() + _NO_SCOPE_WARNING_INTERVAL_SECONDS + verbose_proxy_logger.warning( + "Compresr: enable_retrieval is on but this request has no per-key auth scope; " + "compressing without recovery (compresr_retrieve tool not injected). " + "Configure virtual-key auth to enable recovery." + ) + + existing_originals = self._existing_originals(store_key) + applied = self._apply_compression_results( + messages, targets, contexts, results, recovery_enabled, existing_originals + ) + if applied.messages_compressed == 0: + # Nothing replaced: return the original inputs object (handlers detect + # edits by identity; a fresh list forces write-back that strips Anthropic + # cache_control from thinking blocks). + verbose_proxy_logger.debug("Compresr: service returned no compressed content; request unchanged") + return inputs + + stats: dict[str, object] = { + "messages_compressed": applied.messages_compressed, + "tokens_before": applied.tokens_before, + "tokens_after": applied.tokens_after, + "tokens_saved": applied.tokens_before - applied.tokens_after, + "compression_model": self.compression_model, + } + verbose_proxy_logger.debug( + "Compresr: compressed %s message(s), %s -> %s tokens", + applied.messages_compressed, + applied.tokens_before, + applied.tokens_after, + ) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=stats, + request_data=request_data, + guardrail_status="success", + guardrail_provider="compresr", + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + + compressed_inputs: dict[str, object] = {**inputs, "structured_messages": applied.compressed_messages} + mirrored_texts = self._mirror_texts_channel(inputs.get("texts"), applied) + if mirrored_texts is not None: + compressed_inputs["texts"] = mirrored_texts + + originals = applied.originals + if not recovery_enabled or not originals or store_key is None: + return compressed_inputs # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + + self._store_originals(store_key, originals) + + merged_tools = _merge_retrieve_tool(inputs.get("tools")) + if merged_tools is not None: + compressed_inputs["tools"] = merged_tools + return compressed_inputs # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + if not has_compresr_retrieve_tool(tools): + return False, {} + tool_calls = _extract_compresr_tool_calls(response) + if not tool_calls: + return False, {} + return True, {"tool_calls": tool_calls} + + async def async_build_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: dict, + logging_obj: Any, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + tool_calls: list[dict[str, object]] = tools.get("tool_calls", []) # pyright: ignore[reportAssignmentType] # gate hook builds this dict with list values only + + self._prune_originals() + store_key = _scoped_store_key(logging_obj) + retrieved, resolved_any = self._resolve_retrievals(store_key, tool_calls) + if not resolved_any: + # Nothing this guardrail stored resolved; skip the extra provider round-trip. + return AgenticLoopPlan(run_agentic_loop=False) + + if _is_responses_api_response(response): + follow_up_messages = list(messages) + _build_responses_followup_items(response, retrieved) + elif _is_anthropic_messages_response(response): + follow_up_messages = list(messages) + _build_anthropic_followup_messages(response, retrieved) + else: + assistant_message = _build_assistant_message_from_response(response, retrieved) + tool_results = [ + {"role": "tool", "tool_call_id": tc.get("id"), "content": content} for tc, content in retrieved + ] + follow_up_messages = list(messages) + [assistant_message] + tool_results + + anthropic_max = anthropic_messages_optional_request_params.get("max_tokens") + max_tokens: int | None = anthropic_max if anthropic_max is not None else kwargs.get("max_tokens") + optional_params_without_max_tokens = { + k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" + } + + full_model_name = model + if logging_obj is not None: + agentic_params = getattr(logging_obj, "model_call_details", {}).get("agentic_loop_params", {}) + candidate = agentic_params.get("model", model) + if isinstance(candidate, str) and candidate: + full_model_name = candidate + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + model=full_model_name, + messages=follow_up_messages, + max_tokens=max_tokens, + optional_params=optional_params_without_max_tokens, + kwargs=self._sanitized_follow_up_kwargs(kwargs), + ), + metadata={"tool_type": "compresr_retrieve"}, + ) + + def _sanitized_follow_up_kwargs(self, kwargs: dict) -> dict[str, object]: + """Copy of the request kwargs for the retrieval follow-up with other + guardrails' pre-call-executed markers stripped, so input guardrails + re-inspect the restored originals; only this guardrail's own marker is + kept, to avoid recompressing what it just retrieved.""" + out: dict[str, object] = { + k: v for k, v in kwargs.items() if not k.startswith("_compresr") and k != "litellm_logging_obj" + } + own_marker = self._pre_call_marker() + for meta_key in ("metadata", "litellm_metadata"): + meta = out.get(meta_key) + if not isinstance(meta, dict): + continue + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if not isinstance(executed, list): + continue + kept = [m for m in executed if own_marker is not None and m == own_marker] + out[meta_key] = ( + {**meta, PRE_CALL_EXECUTED_GUARDRAILS_KEY: kept} + if kept + else {k: v for k, v in meta.items() if k != PRE_CALL_EXECUTED_GUARDRAILS_KEY} + ) + return out + + @staticmethod + def get_config_model() -> type[GuardrailConfigModel[object]] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, + ) + + return CompresrGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index e9eee3a1a8a..1f2c9e0c182 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -494,6 +494,7 @@ class InMemoryGuardrailHandler: guardrail_id=guardrail.get("guardrail_id"), guardrail_name=guardrail["guardrail_name"], litellm_params=litellm_params, + guardrail_info=guardrail.get("guardrail_info"), ) # store references to the guardrail in memory @@ -612,6 +613,27 @@ class InMemoryGuardrailHandler: """ return self._sources.get(guardrail_id) + def list_config_guardrails(self) -> List[Guardrail]: + """ + List in-memory guardrails owned by config.yaml. + + DB-sourced entries are excluded: a read surface that also queries the DB + would double-count live ones, and a DB-sourced entry that's missing from + the DB is stale (deleted on another pod, awaiting reconciliation here). + """ + return [g for gid, g in self.IN_MEMORY_GUARDRAILS.items() if self._sources.get(gid) == "config"] + + def get_config_guardrail_by_id(self, guardrail_id: str) -> Optional[Guardrail]: + """ + Get a config-owned in-memory guardrail by its ID, or None. + + Mirrors the fallback in get_guardrail_info: a DB-sourced in-memory entry + that missed the DB lookup is stale and must not be surfaced. + """ + if self._sources.get(guardrail_id) != "config": + return None + return self.IN_MEMORY_GUARDRAILS.get(guardrail_id) + def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]: """ Drop in-memory entries that originated from the DB but are no longer diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index e03bdbb95d2..f56b22ddd49 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -137,10 +137,26 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: return [{"date": d, "passed": v["passed"], "blocked": v["blocked"]} for d, v in sorted(chart_by_date.items())] +def _get_guardrail_field(g: Any, field: str) -> Any: + """Read `field` off a guardrail whether it's a Prisma row (attr) or a dict/TypedDict (key).""" + if isinstance(g, dict): + return g.get(field) + return getattr(g, field, None) + + +def _to_dict(value: Any) -> Dict[str, Any]: + """Coerce a pydantic model (e.g. LitellmParams) / dict value into a plain dict.""" + if isinstance(value, BaseModel): + return value.model_dump(exclude_none=True) + if isinstance(value, dict): + return value + return {} + + def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" - gid = getattr(g, "guardrail_id", None) or (g.get("guardrail_id") if isinstance(g, dict) else None) - name = getattr(g, "guardrail_name", None) or (g.get("guardrail_name") if isinstance(g, dict) else None) + gid = _get_guardrail_field(g, "guardrail_id") + name = _get_guardrail_field(g, "guardrail_name") return gid, (name or gid or "") @@ -163,9 +179,9 @@ def _guardrail_overview_rows( break req, blocked = a["requests"], a["blocked"] fail_rate = (100.0 * blocked / req) if req else 0.0 - litellm_params = (g.litellm_params or {}) if isinstance(g.litellm_params, dict) else {} + litellm_params = _to_dict(_get_guardrail_field(g, "litellm_params")) provider = str(litellm_params.get("guardrail", "Unknown")) - guardrail_info = (g.guardrail_info or {}) if isinstance(g.guardrail_info, dict) else {} + guardrail_info = _to_dict(_get_guardrail_field(g, "guardrail_info")) gtype = str(guardrail_info.get("type", "Guardrail")) prev_fail = 0.0 for k in lookup_keys: @@ -262,9 +278,15 @@ async def guardrails_usage_overview( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + try: - # Guardrails from DB - guardrails = await GuardrailsRepository(prisma_client).table.find_many() + db_guardrails = await GuardrailsRepository(prisma_client).table.find_many() + seen_ids = {gid for g in db_guardrails if (gid := _get_guardrail_field(g, "guardrail_id")) is not None} + config_guardrails = [ + g for g in IN_MEMORY_GUARDRAIL_HANDLER.list_config_guardrails() if g.get("guardrail_id") not in seen_ids + ] + guardrails: List[Any] = [*db_guardrails, *config_guardrails] # Daily metrics in range metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( @@ -321,16 +343,18 @@ async def guardrails_usage_detail( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") - guardrail = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) - if not guardrail: + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) + if guardrail is None: + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id) + if guardrail is None: from fastapi import HTTPException raise HTTPException(status_code=404, detail="Guardrail not found") # Metrics are keyed by logical name (from spend log metadata), not UUID - logical_id = getattr(guardrail, "guardrail_name", None) or ( - guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None - ) + logical_id = _get_guardrail_field(guardrail, "guardrail_name") metric_ids = [i for i in (logical_id, guardrail_id) if i] metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( @@ -367,17 +391,9 @@ async def guardrails_usage_detail( {"date": d, "passed": v["passed"], "blocked": v["blocked"], "score": None} for d, v in sorted(ts_by_date.items()) ] - _litellm_params = getattr(guardrail, "litellm_params", None) or ( - guardrail.get("litellm_params") if isinstance(guardrail, dict) else None - ) - litellm_params = _litellm_params if isinstance(_litellm_params, dict) else {} - _guardrail_info = getattr(guardrail, "guardrail_info", None) or ( - guardrail.get("guardrail_info") if isinstance(guardrail, dict) else None - ) - guardrail_info = _guardrail_info if isinstance(_guardrail_info, dict) else {} - _guardrail_name = getattr(guardrail, "guardrail_name", None) or ( - guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None - ) + litellm_params = _to_dict(_get_guardrail_field(guardrail, "litellm_params")) + guardrail_info = _to_dict(_get_guardrail_field(guardrail, "guardrail_info")) + _guardrail_name = _get_guardrail_field(guardrail, "guardrail_name") return UsageDetailResponse( guardrail_id=guardrail_id, @@ -548,11 +564,15 @@ async def guardrails_usage_logs( # Query by both so we match regardless of which was written. effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else [] if guardrail_id: - guardrail = await GuardrailsRepository(prisma_client).table.find_unique( + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) + if guardrail is None: + guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id) if guardrail: - logical_name = getattr(guardrail, "guardrail_name", None) + logical_name = _get_guardrail_field(guardrail, "guardrail_name") if logical_name and logical_name not in effective_guardrail_ids: effective_guardrail_ids.append(logical_name) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index df311bed7b2..01f4e040e58 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -32,6 +32,7 @@ from litellm._uuid import uuid from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, + MINIMUM_CUSTOM_KEY_LENGTH, UI_SESSION_TOKEN_TEAM_ID, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds @@ -1022,6 +1023,14 @@ async def _common_key_generation_helper( detail={"error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {_masked}"}, ) + if data.key is not None and len(data.key) < MINIMUM_CUSTOM_KEY_LENGTH: + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid key format. LiteLLM Virtual Key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long." + }, + ) + # check org key limits - done here to handle inheriting org id from team if data.organization_id is not None: from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -1474,7 +1483,7 @@ async def generate_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - key_alias: Optional[str] - User defined key alias - - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + - key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you. - team_id: Optional[str] - The team id of the key - user_id: Optional[str] - The user id of the key - agent_id: Optional[str] - The agent id associated with the key. @@ -1688,7 +1697,7 @@ async def generate_service_account_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - key_alias: Optional[str] - User defined key alias - - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. + - key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you. - team_id: Optional[str] - The team id of the key - user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. @@ -4356,7 +4365,6 @@ async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: if data and data.new_key is not None: # Reject custom key values if disabled by admin await _check_custom_key_allowed(data.new_key) - new_token = data.new_key if not data.new_key.startswith("sk-"): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -4364,6 +4372,12 @@ async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: "error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key." }, ) + if len(data.new_key) < MINIMUM_CUSTOM_KEY_LENGTH: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": f"New key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long."}, + ) + new_token = data.new_key else: new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}" return new_token @@ -4470,7 +4484,7 @@ async def _execute_virtual_key_regeneration( new_token = await get_new_token(data=data) new_token_hash = hash_token(new_token) - new_token_key_name = f"sk-...{new_token[-4:]}" + new_token_key_name = abbreviate_api_key(api_key=new_token) update_data = {"token": new_token_hash, "key_name": new_token_key_name} non_default_values = {} @@ -4550,7 +4564,7 @@ async def regenerate_key_fn( - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update - key: Optional[str] - The key to regenerate. - new_master_key: Optional[str] - The new master key to use, if key is the master key. - - new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used. + - new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used. - key_alias: Optional[str] - User-friendly key alias - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 797f4600857..d267a3cac69 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1621,6 +1621,7 @@ async def update_team( - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit per model for this team. Example: {"gpt-4": 100, "gpt-3.5-turbo": 200} - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit per model for this team. Example: {"gpt-4": 10000, "gpt-3.5-turbo": 20000} + - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. Example - update team TPM Limit - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d8015bb8031..0475566192e 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -241,13 +241,34 @@ def _check_cli_sso_start_rate_limit( def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict: + if isinstance(login_id, str) and login_id.startswith("sk-"): + raise HTTPException( + status_code=400, + detail=( + "Your litellm CLI is out of date and uses a login flow this proxy no longer supports. " + "Upgrade it with `pip install -U 'litellm[proxy]'` and run `litellm-proxy login` again." + ), + ) if not _is_valid_cli_sso_login_id(login_id): - raise HTTPException(status_code=400, detail="Invalid CLI login session") + raise HTTPException(status_code=400, detail="Invalid CLI login session id") cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) flow = cache.get_cache(key=cache_key) if not isinstance(flow, dict) or "poll_secret_hash" not in flow: - raise HTTPException(status_code=400, detail="Invalid CLI login session") + verbose_proxy_logger.warning( + "CLI SSO login session not found in cache for login_id=%s. If the proxy runs multiple replicas, " + "a shared Redis cache (enable_redis_auth_cache: true) is required for CLI login to work.", + login_id, + ) + raise HTTPException( + status_code=400, + detail=( + "CLI login session not found or expired. Run `litellm-proxy login` again. " + "If this happens immediately after starting a login, the proxy is likely running multiple " + "replicas without a shared cache; configure Redis with `enable_redis_auth_cache: true` " + "so every replica can see the login session." + ), + ) return flow diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py new file mode 100644 index 00000000000..72bac1bcd6f --- /dev/null +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -0,0 +1,234 @@ +""" +Counts billable HTTP requests on enterprise deployments. + +A billable request is an inbound request to an LLM inference, MCP, or A2A +endpoint that returns a 2xx status. The actual export happens in an injected +recorder (see litellm.proxy.enterprise_billing.billing_metrics); when no +recorder is injected (non-enterprise, or metering misconfigured) this +middleware is a transparent pass-through. +""" + +import re +import threading +from enum import Enum +from typing import Callable, Optional, Protocol, Sequence, runtime_checkable + +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLMRoutes + + +class BillableCategory(str, Enum): + LLM = "llm" + MCP = "mcp" + A2A = "a2a" + + +@runtime_checkable +class BillingRecorder(Protocol): + def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None: ... + + +_MODEL_ID_HEADER = b"x-litellm-model-id" + +# Ordered: a longer suffix that shares an ending with a shorter one must come +# first, e.g. "/chat/completions" before "/completions". This is the POST +# inference surface that writes a SpendLogs row on success, so the exported +# count lines up with the admin UI usage page for inference traffic. Billing +# is a deliberate lower bound on SpendLogs rows: management writes that also +# log (batch/file/fine-tuning creation, interaction cancel) and non-POST calls +# that log (passthrough reads) never bill, so drift only ever undercounts. +_LLM_ROUTE_SUFFIXES: tuple[str, ...] = ( + "/chat/completions", + "/completions", + "/embeddings", + "/responses", + "/rerank", + "/moderations", + "/images/generations", + "/images/edits", + "/images/variations", + "/audio/transcriptions", + "/audio/translations", + "/audio/speech", + "/videos", # create; GET list is excluded by the POST gate + "/remix", # /v1/videos/{id}/remix + "/ocr", + "/search", # /v1/search and /v1/vector_stores/{id}/search + "/rag/query", + "/rag/ingest", + ":generateContent", # Gemini-native /v1beta/models/{model}:generateContent + ":streamGenerateContent", +) + +# Exact paths only: a suffix match would also catch non-inference resources that +# share the ending, e.g. the OpenAI Assistants route /v1/threads/{id}/messages +# writes no SpendLogs row and must not bill, unlike Anthropic /v1/messages. +_LLM_ROUTE_EXACT: tuple[str, ...] = ( + "/v1/messages", + "/interactions", # Google Interactions create; /{id} reads and /cancel do not match + "/v1beta/interactions", +) + +# Provider passthrough prefixes (e.g. /bedrock/..., /vertex-ai/...) carry real +# inference calls that write SpendLogs rows, so they bill. Anchored to the +# routes enum so new providers are picked up without touching this module. +# /langfuse forwards observability traffic, not inference: it writes no +# SpendLogs row and must not bill. +_NON_BILLABLE_PASSTHROUGH_PREFIXES = frozenset({"/langfuse"}) +_PASSTHROUGH_PREFIXES: tuple[str, ...] = tuple( + prefix + for prefix in LiteLLMRoutes.mapped_pass_through_routes.value + if prefix not in _NON_BILLABLE_PASSTHROUGH_PREFIXES +) + + +def _classify_llm_route(path: str) -> Optional[str]: + exact_match = next((route for route in _LLM_ROUTE_EXACT if path == route), None) + if exact_match is not None: + return exact_match + suffix_match = next((suffix for suffix in _LLM_ROUTE_SUFFIXES if path == suffix or path.endswith(suffix)), None) + if suffix_match is not None: + return suffix_match + # Deep passthrough paths only: the bare prefix itself is not an inference call. + return next((prefix for prefix in _PASSTHROUGH_PREFIXES if path.startswith(f"{prefix}/")), None) + + +_MCP_MANAGEMENT_PREFIX = "/v1/mcp" +_MCP_DYNAMIC_TRANSPORT = re.compile(r"/(?:toolset/)?[^/]+/mcp") +# The REST wrapper's tool-call endpoint executes a tool and fires the same MCP +# spend logging as the /mcp transport; its list/test siblings do not bill. +_MCP_REST_TOOL_CALL = "/mcp-rest/tools/call" + +_A2A_INVOKE_SUFFIX = "/message/send" +_A2A_TRANSPORT_PREFIXES: tuple[str, ...] = ("/v1/a2a/", "/a2a/") +# Bare POST /a2a/{agent_id} carries the JSON-RPC method in the body, not the +# path. Only message/send and message/stream write a SpendLogs row there; the +# task RPCs (tasks/get, tasks/cancel, tasks/pushNotificationConfig/*, ...) are +# forwarded upstream and write none. A path-only classifier cannot separate +# them, so the bare route does not bill: counting a task RPC would overcount, +# while missing a bare-path message/send only undercounts, and undercounting is +# the sole direction this metric is allowed to drift. The /mcp transport is +# method-agnostic by contrast because its list path logs a SpendLogs row too. + + +def _classify_mcp_route(path: str) -> Optional[str]: + if path == _MCP_MANAGEMENT_PREFIX or path.startswith(f"{_MCP_MANAGEMENT_PREFIX}/"): + return None + if path == "/mcp" or path.startswith("/mcp/"): + return "/mcp" + if path == _MCP_REST_TOOL_CALL: + return "/mcp" + if _MCP_DYNAMIC_TRANSPORT.fullmatch(path) is not None: + return "/mcp" + return None + + +def _classify_a2a_route(path: str) -> Optional[str]: + if path.endswith(_A2A_INVOKE_SUFFIX) and any(path.startswith(prefix) for prefix in _A2A_TRANSPORT_PREFIXES): + return "/a2a" + return None + + +def classify_billable_request(path: str, method: str = "POST") -> Optional[tuple[BillableCategory, str]]: + """Map a request path to its (category, normalized route), or None if not billable.""" + normalized = path.rstrip("/") or "/" + + mcp_route = _classify_mcp_route(normalized) + if mcp_route is not None: + return (BillableCategory.MCP, mcp_route) + + a2a_route = _classify_a2a_route(normalized) + if a2a_route is not None: + return (BillableCategory.A2A, a2a_route) + + # POST-only is a conservative gate: non-POST calls can still write a + # SpendLogs row (passthrough reads, resource GETs) but must not bill, so + # any classifier-vs-dashboard mismatch is an undercount, never an overcount. + if method.upper() != "POST": + return None + + llm_route = _classify_llm_route(normalized) + if llm_route is not None: + return (BillableCategory.LLM, llm_route) + return None + + +def _extract_model_id(headers: Sequence[tuple[bytes, bytes]]) -> Optional[str]: + return next( + (value.decode("latin-1") for name, value in headers if name.lower() == _MODEL_ID_HEADER and value), + None, + ) + + +class BillableRequestMetricsMiddleware: + """ + Pure ASGI middleware that records one billable request per 2xx response to a + billable endpoint. Modeled on InFlightRequestsMiddleware: it wraps `send`, + reads the final status and the x-litellm-model-id header off the + `http.response.start` message, and never blocks or fails the request path. + """ + + def __init__( + self, + app: ASGIApp, + recorder: Optional[BillingRecorder] = None, + recorder_factory: Optional[Callable[[], Optional[BillingRecorder]]] = None, + ) -> None: + self.app = app + self.recorder = recorder + # The factory defers recorder construction to the first request, AFTER the + # startup event has loaded the YAML config's environment_variables (license + # and cert env vars). Building at import time captured recorder=None for + # deployments configured that way. Resolved exactly once; the result + # (including None) is cached. + self._recorder_factory = recorder_factory + self._resolved = recorder_factory is None + self._resolve_lock = threading.Lock() + + def _resolve_recorder(self) -> Optional[BillingRecorder]: + if self._resolved: + return self.recorder + # The lock keeps concurrent first requests from each building their own + # MeterProvider (and leaking its background exporter thread). + with self._resolve_lock: + if not self._resolved: + factory = self._recorder_factory + self.recorder = factory() if factory is not None else self.recorder + self._resolved = True + return self.recorder + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + recorder = self._resolve_recorder() + if recorder is None: + await self.app(scope, receive, send) + return + + classification = classify_billable_request(scope.get("path", ""), scope.get("method", "POST")) + if classification is None: + await self.app(scope, receive, send) + return + + category, route = classification + status_code = 0 + model_id: Optional[str] = None + + async def send_wrapper(message: Message) -> None: + nonlocal status_code, model_id + if message["type"] == "http.response.start": + status_code = message["status"] + model_id = _extract_model_id(message.get("headers", [])) + await send(message) + + await self.app(scope, receive, send_wrapper) + + if 200 <= status_code < 300: + try: + recorder.record(category=category, route=route, status_code=status_code, model_id=model_id) + except Exception: # noqa: BLE001 -- metering must never fail a request that was already served + verbose_proxy_logger.warning("billable request metering failed for %s", route, exc_info=True) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a0b07ece4dd..6661474d215 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -438,10 +438,30 @@ from litellm.proxy.management_helpers.audit_logs import ( create_object_audit_log, ) from litellm.proxy.memory.memory_endpoints import router as memory_router +from litellm.proxy.middleware.billable_request_metrics_middleware import ( + BillableRequestMetricsMiddleware, + BillingRecorder, +) from litellm.proxy.plugin_routes import ( - router as plugin_router, register_plugins_from_config, ) +from litellm.proxy.plugin_routes import ( + router as plugin_router, +) + +try: + from litellm.proxy.enterprise_billing.billing_metrics import ( + build_billing_metrics_recorder as _build_billing_metrics_recorder, + ) + from litellm.proxy.enterprise_billing.billing_metrics import ( + shutdown_billing_metrics_recorder as _shutdown_billing_metrics_recorder, + ) + + build_billing_metrics_recorder: Optional[Callable[..., Optional[BillingRecorder]]] = _build_billing_metrics_recorder + shutdown_billing_metrics_recorder: Optional[Callable[[], None]] = _shutdown_billing_metrics_recorder +except ImportError: + build_billing_metrics_recorder = None + shutdown_billing_metrics_recorder = None from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -461,13 +481,11 @@ from litellm.proxy.openai_files_endpoints.files_endpoints import ( ) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, + vertex_ai_live_websocket_passthrough, ) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( router as llm_passthrough_router, ) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - vertex_ai_live_websocket_passthrough, -) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( initialize_pass_through_endpoints, ) @@ -552,21 +570,19 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( from litellm.types.realtime import RealtimeQueryParams from litellm.types.router import ( DeploymentTypedDict, -) -from litellm.types.router import ModelInfo as RouterModelInfo -from litellm.types.router import ( RouterGeneralSettings, + RoutingPlugin, SearchToolTypedDict, updateDeployment, ) +from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.scheduler import DefaultPriorities from litellm.types.secret_managers.main import ( KeyManagementSettings, KeyManagementSystem, ) -from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer +from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import RawRequestTypedDict, StandardLoggingPayload from litellm.utils import _add_custom_logger_callback_to_specific_event try: @@ -767,6 +783,11 @@ async def proxy_shutdown_event(): if db_writer_client is not None: await db_writer_client.close() # type: ignore[reportGeneralTypeIssues] + # final flush of billable-request counts: without it, up to one export + # interval of enterprise billing data is dropped on every restart + if shutdown_billing_metrics_recorder is not None: + shutdown_billing_metrics_recorder() + # flush remaining langfuse logs if "langfuse" in litellm.success_callback: try: @@ -972,11 +993,11 @@ async def proxy_startup_event(app: FastAPI): if is_otel_v2_enabled(): from opentelemetry import trace as _otel_trace - from litellm.litellm_core_utils.litellm_logging import _in_memory_loggers from litellm.integrations.otel.logger import ( OpenTelemetryV2, publish_global_otel_v2_provider, ) + from litellm.litellm_core_utils.litellm_logging import _in_memory_loggers registered = open_telemetry_logger if isinstance(open_telemetry_logger, OpenTelemetryV2) else None publish_global_otel_v2_provider( @@ -1780,6 +1801,31 @@ app.add_middleware( ) app.add_middleware(PrometheusAuthMiddleware) +# Added before InFlightRequestsMiddleware so it nests *inside* it: Starlette +# makes the last-added middleware outermost. The billable count is recorded +# after the inner app returns, so if this sat outside the in-flight tracker a +# request could be counted as drained while its record() had not yet run, and +# proxy_shutdown_event could flush and stop the exporter underneath it. +app.add_middleware( + BillableRequestMetricsMiddleware, + # Factory, not an instance: the recorder is resolved on the first request so + # it sees premium_user and the billing env vars AFTER proxy_startup_event has + # loaded the YAML config's environment_variables. Building it here at import + # time would permanently capture recorder=None for YAML-configured + # deployments. The lambda reads the module globals at call time. + recorder_factory=lambda: ( + build_billing_metrics_recorder( + premium=premium_user, + # Read from the license check, not the premium_user_data module + # global: that global is bound once at import and goes stale when + # the license arrives via the YAML config's environment_variables. + license_data=_license_check.airgapped_license_data, + litellm_version=version, + ) + if build_billing_metrics_recorder is not None + else None + ), +) app.add_middleware(InFlightRequestsMiddleware) app.add_middleware(SecurityHeadersMiddleware) @@ -3660,6 +3706,45 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: litellm_config_cache.redis_cache = redis_cache +def resolve_complexity_router_plugins( + model_name: str, + complexity_router_config: dict, + config_file_path: str | None, +) -> None: + """ + Resolves `complexity_router_config["plugins"]` dotted-path strings to live + instances via `get_instance_fn` (the same convention `litellm_settings.callbacks` + uses), in place. Raises at config-load time if a path resolves to something that + doesn't implement `RoutingPlugin`, rather than deferring to a confusing + `AttributeError` on the first request that reaches the plugin pipeline. + """ + plugin_paths = complexity_router_config.get("plugins") + if not isinstance(plugin_paths, list): + return + + resolved_plugins = [ + get_instance_fn(value=plugin_path, config_file_path=config_file_path) + if isinstance(plugin_path, str) + else plugin_path + for plugin_path in plugin_paths + ] + for plugin_path, resolved_plugin in zip(plugin_paths, resolved_plugins): + # `@runtime_checkable` only checks that `run` exists as an attribute, not that + # it's a coroutine function -- a synchronous `def run(self, context)` would pass + # isinstance() here and only fail at request time with a confusing `TypeError: + # object RoutingContext can't be used in 'await' expression`. + if not isinstance(resolved_plugin, RoutingPlugin) or not inspect.iscoroutinefunction( + getattr(resolved_plugin, "run", None) + ): + raise ValueError( + f"complexity_router_config.plugins entry {plugin_path!r} on model {model_name!r} " + f"resolved to {resolved_plugin!r}, which does not implement the RoutingPlugin " + "interface (an async `run(context)` method). Fix the referenced module before " + "starting the proxy." + ) + complexity_router_config["plugins"] = resolved_plugins + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -4720,6 +4805,13 @@ class ProxyConfig: for k, v in model["litellm_params"].items(): if isinstance(v, str) and v.startswith("os.environ/"): model["litellm_params"][k] = get_secret(v) + complexity_router_config = model["litellm_params"].get("complexity_router_config") + if isinstance(complexity_router_config, dict): + resolve_complexity_router_plugins( + model_name=model.get("model_name", ""), + complexity_router_config=complexity_router_config, + config_file_path=config_file_path, + ) print(f"\033[32m {model.get('model_name', '')}\033[0m") # noqa: T201 litellm_model_name = model["litellm_params"]["model"] litellm_model_api_base = model["litellm_params"].get("api_base", None) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 8c980f33b01..94871ff072d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -6,6 +6,7 @@ from fastapi import HTTPException, status import litellm from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_utils.common_utils import _is_proxy_admin_request # Router-internal mock_testing_* flag names — kept in sync with # ``litellm.types.router.MockRouterTestingParams`` by the test @@ -363,6 +364,7 @@ async def route_request( team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] + is_proxy_admin_without_team = team_id is None and _is_proxy_admin_request(data) # Preprocess Google GenAI generate content requests if route_type in ["agenerate_content", "agenerate_content_stream"]: @@ -517,6 +519,13 @@ async def route_request( data["model"] = team_model_name return getattr(llm_router, f"{route_type}")(**data) + elif ( + is_proxy_admin_without_team + and data["model"] not in router_model_names + and data["model"] in llm_router.team_public_model_names + ): + return getattr(llm_router, f"{route_type}")(**data) + elif data["model"] in router_model_names or llm_router.has_model_id(data["model"]): return getattr(llm_router, f"{route_type}")(**data) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 62fc28256cd..9f36e729330 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -833,7 +833,7 @@ class ProxyLogging: def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) - return list(set(dynamic_success_callbacks + global_callbacks)) + return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) def _parse_pre_mcp_call_hook_response( self, diff --git a/litellm/router.py b/litellm/router.py index 6e8127110cc..f408d030b8b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -26,6 +26,7 @@ from typing import ( AsyncGenerator, Callable, Dict, + FrozenSet, Generator, List, Literal, @@ -108,6 +109,7 @@ from litellm.router_utils.clientside_credential_handler import ( is_clientside_credential, ) from litellm.router_utils.common_utils import ( + _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, ) @@ -494,6 +496,7 @@ class Router: self.model_name_to_deployment_indices: Dict[str, List[int]] = {} # Maps (team_id, team_public_model_name) -> list of indices in model_list self.team_model_to_deployment_indices: Dict[Tuple[str, str], List[int]] = {} + self.team_public_model_names: FrozenSet[str] = frozenset() # Initialize cache attributes that ``_invalidate_model_group_info_cache`` # touches *before* the first ``set_model_list`` below (which calls @@ -2983,7 +2986,7 @@ class Router: # here before it's wiped below, instead of relying on that attempt's # (possibly still-pending) failure event to do it. refund_stale_reservation_before_retry(self.cache, kwargs) - set_io_token_rate_limit_request_kwargs(kwargs) + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment)) ## DEPLOYMENT-LEVEL TAGS deployment_tags = deployment.get("litellm_params", {}).get("tags") @@ -7284,6 +7287,14 @@ class Router: raise e return returned_healthy_deployments + @staticmethod + def _json_default_stable_id(value: object) -> str: + """json.dumps default= for _generate_model_id: plain str() on an arbitrary + object (e.g. a RoutingPlugin instance) falls back to object.__repr__'s + ``, so the hash -- and deployment id -- would + change every restart. Use the class name instead, stable across restarts.""" + return f"{type(value).__module__}.{type(value).__qualname__}" + def _generate_model_id(self, model_group: str, litellm_params: dict): """ Helper function to consistently generate the same id for a deployment @@ -7299,14 +7310,14 @@ class Router: if isinstance(k, str): parts.append(k) elif isinstance(k, dict): - parts.append(json.dumps(k)) + parts.append(json.dumps(k, default=self._json_default_stable_id)) else: parts.append(str(k)) if isinstance(v, str): parts.append(v) elif isinstance(v, dict): - parts.append(json.dumps(v)) + parts.append(json.dumps(v, default=self._json_default_stable_id)) else: parts.append(str(v)) @@ -7792,6 +7803,7 @@ class Router: self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index self.team_model_to_deployment_indices = {} # Reset the team_model index + self.team_public_model_names = frozenset() # Reset per-strategy router registries so hot-reload doesn't leave # stale routers pointing at the old model_list. self.quality_routers = {} @@ -8143,6 +8155,9 @@ class Router: self.team_model_to_deployment_indices[key] = updated_indices else: del self.team_model_to_deployment_indices[key] + self.team_public_model_names = frozenset( + public_model_name for _, public_model_name in self.team_model_to_deployment_indices + ) def _update_team_model_index(self, model: dict, idx: int) -> None: """ @@ -8156,6 +8171,7 @@ class Router: team_public_model_name = (model.get("model_info") or {}).get("team_public_model_name") if team_id and team_public_model_name: key = (team_id, team_public_model_name) + self.team_public_model_names = self.team_public_model_names | frozenset({team_public_model_name}) if key not in self.team_model_to_deployment_indices: self.team_model_to_deployment_indices[key] = [] if idx not in self.team_model_to_deployment_indices[key]: @@ -9110,6 +9126,7 @@ class Router: """ self.model_name_to_deployment_indices.clear() self.team_model_to_deployment_indices.clear() + self.team_public_model_names = frozenset() for idx, model in enumerate(model_list): model_name = model.get("model_name") @@ -10018,7 +10035,10 @@ class Router: return [m for m in self.model_list if m["litellm_params"]["model"] == model] def _try_early_resolve_deployments_for_model_not_in_names( - self, model: str, request_team_id: Optional[str] + self, + model: str, + request_team_id: Optional[str], + include_team_models: bool = False, ) -> Optional[Tuple[str, Union[List, Dict]]]: """ When ``model`` is not in ``self.model_names``, try team routes, pattern routes, @@ -10033,6 +10053,30 @@ class Router: team_deployments = self._get_all_deployments(model_name=model, team_id=request_team_id) if team_deployments: return model, team_deployments + elif include_team_models: + team_deployments = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] + team_ids = { + team_id + for deployment in team_deployments + for team_id in [(deployment.get("model_info") or {}).get("team_id")] + if team_id is not None + } + if len(team_ids) > 1: + raise litellm.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + if team_deployments: + return model, team_deployments pattern_deployments = self.pattern_router.get_deployments_by_pattern( model=model, @@ -10097,7 +10141,11 @@ class Router: if _model_from_alias is not None: model = _model_from_alias - early = self._try_early_resolve_deployments_for_model_not_in_names(model=model, request_team_id=request_team_id) + early = self._try_early_resolve_deployments_for_model_not_in_names( + model=model, + request_team_id=request_team_id, + include_team_models=_is_proxy_admin_request(request_kwargs), + ) if early is not None: return early diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index bd3b300b558..e85987870e1 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -98,9 +98,9 @@ def _sanitize_user_api_key_auth(auth: Any) -> Any: return auth -def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None: +def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: if not metadata: - return metadata + return {} return { k: _sanitize_user_api_key_auth(v) if k == "user_api_key_auth" else v for k, v in metadata.items() @@ -468,6 +468,38 @@ class ComplexityRouter(CustomLogger): def _tier_pools(self) -> dict[str, list[str]]: return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()} + async def _pick_model_for_tier( + self, + tier: ComplexityTier, + raw_messages: list[dict[str, Any]] | None, + resolved_messages: list[dict[str, Any]] | None, + request_kwargs: dict, + ) -> str: + if not self.config.plugins: + return self.get_model_for_tier(tier) + + from litellm.types.router import RoutingContext + + tier_key = tier.value + metadata_key = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + context = RoutingContext( + raw_messages=raw_messages or [], + structured_messages=resolved_messages or [], + candidate_models=list(self._tier_pools().get(tier_key, [])), + metadata=request_kwargs.get(metadata_key) or {}, + ) + for plugin in self.config.plugins: + context = await plugin.run(context) + + if not context.candidate_models: + # A plugin narrowing a tier to zero candidates is a policy decision (e.g. no + # model this tenant's budget allows) -- falling back to default_model here + # (which was never checked against the plugins) would let that policy be + # silently bypassed. Raise instead, matching the Router-level plugin + # pipeline's own fail-closed behavior for the same situation. + raise ValueError(f"No candidate models left for tier {tier_key} after routing-plugin filtering") + return self._pick_from_tier_value(context.candidate_models, tier_key) + def _ensure_adaptive_router(self) -> Any | None: if not self.config.adaptive: return None @@ -731,8 +763,8 @@ class ComplexityRouter(CustomLogger): # embedding call. Forwarding it would let the embedding's cost callback finalize the # reservation, so the routed completion's own callback then skips incrementing the # key/team budget. Key/team attribution fields are preserved for spend logging. - metadata = _classifier_call_metadata(request_kwargs.get("metadata")) or {} - litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata")) or {} + metadata = _classifier_call_metadata(request_kwargs.get("metadata")) + litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata")) query_vector = ( await encoder.aencode_queries([user_message], metadata=metadata, litellm_metadata=litellm_metadata) )[0] @@ -861,10 +893,16 @@ class ComplexityRouter(CustomLogger): When `session_affinity` is enabled and a session_id is resolvable on the request, pins the model chosen on the session's first turn and reuses it for every later turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`. + + Skipped entirely when `plugins` are configured: reusing a stale pin would bypass + the plugin pipeline on every turn after the first, since a pinned model was never + re-checked against a policy plugin whose decision can change between turns (e.g. a + budget plugin, once the session's spend crosses its cap). """ from litellm.types.router import PreRoutingHookResponse - session_id = self._get_session_id_from_request_kwargs(request_kwargs) if self.config.session_affinity else None + use_session_affinity = self.config.session_affinity and not self.config.plugins + session_id = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None if cache_key is not None: @@ -947,14 +985,26 @@ class ComplexityRouter(CustomLogger): if user_message is None: verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model") + if not self.config.plugins and self.config.default_model: + # No plugins configured: preserve the pre-existing default_model-first + # priority exactly (changing it would be a silent behavior change for + # every non-plugin user, not just a security fix). + routed_model = self.config.default_model + else: + # Plugins configured: default_model must never bypass them, so it's not + # checked here at all -- _pick_model_for_tier -> get_model_for_tier still + # falls back to it (after the MEDIUM tier) once the plugin pipeline runs. + routed_model = await self._pick_model_for_tier( + ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs + ) return PreRoutingHookResponse( - model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM), + model=routed_model, messages=messages if has_original_messages else None, ) override_tier = await self._resolve_keyword_tier_override(user_message, request_kwargs) if override_tier is not None: - routed_model = self.get_model_for_tier(override_tier) + routed_model = await self._pick_model_for_tier(override_tier, messages, resolved_messages, request_kwargs) cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match" verbose_router_logger.info( f"ComplexityRouter: routing decision cause={cause}, " @@ -980,7 +1030,7 @@ class ComplexityRouter(CustomLogger): f"signals={signals}, routed_model={routed_model}" ) else: - routed_model = self.get_model_for_tier(tier) + routed_model = await self._pick_model_for_tier(tier, messages, resolved_messages, request_kwargs) verbose_router_logger.info( f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, " f"score={score:.3f}, signals={signals}, routed_model={routed_model}" diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index e4bd36505e6..1f984798970 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -10,7 +10,7 @@ from typing import Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from litellm.types.router import AdaptiveRouterWeights +from litellm.types.router import AdaptiveRouterWeights, RoutingPlugin class ComplexityTier(str, Enum): @@ -363,10 +363,13 @@ class ComplexityRouterConfig(BaseModel): # Session affinity: pin the first turn's routed model for the rest of the session session_affinity: bool = Field( - default=False, + default=True, description=( "When True and a session_id is resolvable on the request, pin the model chosen on the " - "session's first turn and reuse it for every later turn, skipping re-classification." + "session's first turn and reuse it for every later turn, skipping re-classification. " + "On by default so multi-turn sessions stay on one model, preserving provider prompt " + "caches and avoiding cross-model conversation-history errors. Set False to reclassify " + "every turn." ), ) session_affinity_ttl_seconds: int = Field( @@ -375,7 +378,12 @@ class ComplexityRouterConfig(BaseModel): description="TTL for the session affinity pin; refreshed on every cache hit", ) - model_config = ConfigDict(extra="allow") # Allow additional fields + plugins: list[RoutingPlugin] | None = Field( + default=None, + description="RoutingPlugin instances that narrow the classified tier's candidate models before selection", + ) + + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) # Allow additional fields @field_validator("tiers", mode="before") @classmethod @@ -421,6 +429,15 @@ class ComplexityRouterConfig(BaseModel): raise ValueError("keyword_tier_rules must be non-empty when semantic_keyword_matching is enabled") return self + @model_validator(mode="after") + def _validate_plugins_adaptive_combo(self) -> "ComplexityRouterConfig": + if self.plugins and self.adaptive: + raise ValueError( + "plugins and adaptive=True cannot both be set: adaptive's bandit selection doesn't yet " + "consume plugin-narrowed candidate pools. Disable adaptive or remove plugins." + ) + return self + # Combined default config DEFAULT_COMPLEXITY_CONFIG = ComplexityRouterConfig() diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 18bce2348f6..5cfea5e3bf2 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -1,14 +1,27 @@ import hashlib import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Dict, List, Optional, Union if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject +from litellm.exceptions import BadRequestError from litellm.types.router import CredentialLiteLLMParams from litellm._logging import verbose_logger +def _is_proxy_admin_request(request_kwargs: Optional[Mapping[str, object]]) -> bool: + if request_kwargs is None: + return False + metadata_value = request_kwargs.get("metadata") + litellm_metadata_value = request_kwargs.get("litellm_metadata") + metadata = metadata_value if isinstance(metadata_value, Mapping) else {} + litellm_metadata = litellm_metadata_value if isinstance(litellm_metadata_value, Mapping) else {} + user_api_key_auth = metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth") + return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model @@ -59,6 +72,40 @@ def filter_team_based_models( metadata = request_kwargs.get("metadata") or {} litellm_metadata = request_kwargs.get("litellm_metadata") or {} request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list): + requested_model = ( + request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") + ) + candidate_deployments = tuple( + (deployment.get("model_name"), deployment.get("model_info") or {}) for deployment in healthy_deployments + ) + team_ids = frozenset( + team_id + for _, model_info in candidate_deployments + for team_id in [model_info.get("team_id")] + if team_id is not None + ) + matches_requested_model = ( + isinstance(requested_model, str) + and bool(candidate_deployments) + and all( + model_info.get("team_id") is not None + and (model_name == requested_model or model_info.get("team_public_model_name") == requested_model) + for model_name, model_info in candidate_deployments + ) + ) + if matches_requested_model and len(team_ids) > 1: + raise BadRequestError( + message=( + f"Model name '{requested_model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=requested_model, + llm_provider="", + ) + if matches_requested_model: + return healthy_deployments + ids_to_remove = set() if isinstance(healthy_deployments, dict): return healthy_deployments diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 62c99a1c9f6..803fdc4b353 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -43,14 +43,21 @@ ITPM_CACHE_KEY = "_litellm_itpm_cache_key" OTPM_CACHE_KEY = "_litellm_otpm_cache_key" -def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]]) -> None: +def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]], store_in_context: bool = True) -> None: # The reservation sentinels are server-only, but `metadata` is caller # controlled on proxy requests. Strip any client-supplied copies here (this # runs before the router stashes its own reservation) so a forged # reservation can't drive the post-call reconcile/refund against an # arbitrary counter and bypass the configured limits. _clear_reservation_from_kwargs(kwargs) - _io_token_rate_limit_request_kwargs.set(kwargs) + # The context slot pins the entire request kwargs (messages included) for + # the lifetime of the surrounding context, which outlives the request when + # the context is captured by pooled resources (e.g. a redis connection + # created mid-request). Only ITPM/OTPM-limited deployments read it, so for + # every other deployment overwrite the slot with None instead of the + # kwargs; overwriting (rather than skipping) also releases a previous + # request's kwargs when a context is reused. + _io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None) def get_io_token_rate_limit_request_kwargs() -> Optional[dict[str, Any]]: diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 03b6b36ec2e..f3410935ec7 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -56,6 +56,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( HeadroomGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, +) """ Pydantic object defining how to set guardrails on litellm proxy @@ -123,6 +126,7 @@ class SupportedGuardrailIntegrations(Enum): VIGIL_GUARD = "vigil_guard" REPELLOAI = "repelloai" HEADROOM = "headroom" + COMPRESR = "compresr" class Role(Enum): @@ -806,7 +810,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up default="fail_closed", description=( "Behavior when a guardrail endpoint is unreachable due to network errors. " - "Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', and 'headroom'. " + "Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. " "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed." ), ) @@ -899,6 +903,7 @@ class LitellmParams( BedrockGuardrailConfigModel, LakeraV2GuardrailConfigModel, HeadroomGuardrailConfigModel, + CompresrGuardrailConfigModel, RepelloAIGuardrailConfigModel, LassoGuardrailConfigModel, PillarGuardrailConfigModel, @@ -1034,6 +1039,7 @@ class ApplyGuardrailRequest(BaseModel): entities: Optional[List[PiiEntityType]] = None input_type: str = "request" messages: Optional[List[Dict[str, Any]]] = None + metadata: Dict[str, Any] | None = None class ApplyGuardrailResponse(BaseModel): diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 82ff15303f3..801436c774a 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -17,6 +17,12 @@ MCPInfo = Dict[str, Any] class MCPOAuthMetadata(BaseModel): scopes: Optional[List[str]] = None + """Resource-driven scopes for the authorization request: the RFC 9728 protected-resource + ``scopes_supported``, or the ``scope`` from the WWW-Authenticate 401 challenge when the resource + supplied one, else the authorization server's ``scopes_supported``. This is the scope value a + client requests per the MCP authorization spec Scope Selection Strategy; scope minimization and + inflation control are the authorization server's and user's job at consent (RFC 6749 §3.3), not + the client's.""" authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None @@ -122,9 +128,10 @@ class MCPServer(BaseModel): # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). # Tokens that fail validation are rejected before storage. token_validation: Optional[Dict[str, Any]] = None - # Optional TTL override (seconds) for the Redis per-user token cache. - # Defaults to the token's expires_in minus the expiry buffer, or - # MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. + # Optional TTL override (seconds) for the Redis per-user token cache, capped + # at the token's expires_in minus the expiry buffer so a cached entry never + # outlives the token. Defaults to the token's expires_in minus the expiry + # buffer, or MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. token_storage_ttl_seconds: Optional[int] = None timeout: Optional[float] = None # Max concurrent outbound tool calls to this server; excess calls queue. diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py new file mode 100644 index 00000000000..dad61f83b7d --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py @@ -0,0 +1,135 @@ +from typing import Any, Dict, Literal + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class CompresrGuardrailOptionalParams(BaseModel): + """Optional tuning knobs for the Compresr guardrail.""" + + target_compression_ratio: float | None = Field( + default=None, + description=( + "Compression strength. 0-1 is the fraction of tokens to remove " + "(0.5 = remove ~50%, the default); a value >1 is an Nx reduction " + "factor (e.g. 4 = ~4x smaller)." + ), + ) + coarse: bool | None = Field( + default=None, + description=("Paragraph-level compression (default, faster) instead of token-level (finer-grained)."), + ) + min_chars_to_compress: int | None = Field( + default=None, + description=("Skip messages whose text is shorter than this many characters. Defaults to 500."), + ) + compress_tool_outputs: bool | None = Field( + default=None, + description=("Compress tool/function result messages (search hits, RAG chunks, API dumps). Defaults to True."), + ) + compress_system: bool | None = Field( + default=None, + description="Also compress system messages. Defaults to False.", + ) + compress_history: bool | None = Field( + default=None, + description="Also compress prior (non-last) user messages. Defaults to False.", + ) + compress_last_user: bool | None = Field( + default=None, + description=( + "Also compress the last user message. The query sent to Compresr " + "is always the original verbatim text. Defaults to False." + ), + ) + enable_retrieval: bool | None = Field( + default=None, + description=( + "Make compression recoverable: inject a `compresr_retrieve` tool " + "so the model can fetch the original content behind a compression " + "marker via the agentic loop. Defaults to True. Set to False (or " + "run the proxy with --workers 1) for multi-worker deployments: " + "the recovery store is per-process, so pre-call and retrieval hooks " + "on different workers cannot see each other's originals." + ), + ) + max_bytes_per_call: int | None = Field( + default=None, + description=( + "Cap on aggregate bytes of stored originals per litellm_call_id. " + "When a call exceeds this, oldest entries are evicted so the " + "in-process store cannot grow without bound. Defaults to 10 MiB." + ), + ) + allow_bypass_header: bool | None = Field( + default=None, + description=( + "Honor the `x-compresr-bypass: true` request header to skip " + "compression for a single call. Off by default because the " + "header is caller-settable; enable only on trusted deployments." + ), + ) + dynamic: bool | None = Field( + default=None, + description=( + "latte_v2 only. Let the server choose the compression amount per input " + "(Kneedle elbow) instead of using target_compression_ratio. Defaults to True." + ), + ) + dynamic_min_ratio: float | None = Field( + default=None, + description=( + "latte_v2 only. Floor on the adaptive ratio when `dynamic` is on. " + "Unset lets the server default apply (~1.5)." + ), + ) + dynamic_max_ratio: float | None = Field( + default=None, + description=( + "latte_v2 only. Ceiling on the adaptive ratio when `dynamic` is on. " + "Unset lets the server default apply (~10.0)." + ), + ) + compression_params: Dict[str, Any] | None = Field( + default=None, + description=( + "Passthrough of extra parameters forwarded verbatim in the Compresr " + "compress payload (e.g. `heuristic_chunking`, or any newer knob), so " + "a new Compresr feature works without a guardrail update. The named " + "fields above take precedence on collision." + ), + ) + + +class CompresrGuardrailConfigModel(GuardrailConfigModel[CompresrGuardrailOptionalParams]): + api_key: str | None = Field( + default=None, + description=("Compresr API key. Falls back to the COMPRESR_API_KEY env var."), + ) + api_base: str | None = Field( + default=None, + description=( + "Base URL of the Compresr API. Falls back to the COMPRESR_API_BASE " + "env var, then https://api.compresr.ai. Point at your internal " + "service URL for on-prem deployments." + ), + ) + model: str | None = Field( + default=None, + description=( + "Compresr compression model (not the LLM). Defaults to 'latte_v2', the query-aware compression model." + ), + ) + unreachable_fallback: Literal["fail_closed", "fail_open"] = Field( + default="fail_closed", + description=( + "Behavior when the Compresr compression service is unreachable or errors. " + "'fail_closed' raises an error (default). 'fail_open' logs a critical error and " + "forwards the request uncompressed instead of blocking it." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Compresr (context compression)" diff --git a/litellm/types/router.py b/litellm/types/router.py index 3bedd97c20c..d62c613bf57 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -9,7 +9,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hi import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from typing_extensions import Protocol, Required, TypedDict +from typing_extensions import Protocol, Required, TypedDict, runtime_checkable from litellm._uuid import uuid @@ -852,6 +852,7 @@ class RoutingContext(BaseModel): signals: dict[str, Any] = Field(default_factory=dict) +@runtime_checkable class RoutingPlugin(Protocol): """Interface a custom routing plugin must implement to run in `Router(plugins=[...])`.""" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index df6e8c992ef..e10dde793d1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -44382,6 +44382,90 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock_mantle/openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.875e-06, + "cache_read_input_token_cost": 5.5e-07, + "output_cost_per_token": 3.3e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/openai.gpt-5.6-terra": { + "input_cost_per_token": 2.75e-06, + "cache_creation_input_token_cost": 3.4375e-06, + "cache_read_input_token_cost": 2.75e-07, + "output_cost_per_token": 1.65e-05, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/openai.gpt-5.6-luna": { + "input_cost_per_token": 1.1e-06, + "cache_creation_input_token_cost": 1.375e-06, + "cache_read_input_token_cost": 1.1e-07, + "output_cost_per_token": 6.6e-06, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, @@ -44494,6 +44578,7 @@ "supports_vision": true }, "bedrock_mantle/xai.grok-4.3": { + "use_openai_responses_path": true, "input_cost_per_token": 1.25e-06, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2e-07, diff --git a/pyproject.toml b/pyproject.toml index f890cc976f3..2c19bf64b4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "litellm" version = "1.94.0" description = "Library to easily interface with LLM API providers" readme = "README.md" -requires-python = ">=3.10, <3.14" +requires-python = ">=3.10, <3.15" license = "MIT" license-files = ["LICENSE"] authors = [ @@ -66,6 +66,7 @@ proxy = [ "litellm-enterprise==0.1.50", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", + "InquirerPy>=0.3.4,<1.0", "polars>=1.38.1,<2.0", "soundfile>=0.12.1,<1.0", "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", @@ -74,11 +75,12 @@ proxy = [ ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports (fastapi, cryptography, ...) are all guarded, so it runs on the base -# SDK plus just these three; none of the server runtime in `proxy` is pulled in. +# SDK plus just these four; none of the server runtime in `proxy` is pulled in. cli = [ "rich>=13.9.4,<14.0", "pyyaml>=6.0.3,<7.0", "requests>=2.32.0,<3.0", + "InquirerPy>=0.3.4,<1.0", ] extra_proxy = [ "prisma>=0.11.0,<1.0", @@ -129,7 +131,7 @@ proxy-runtime = [ "opentelemetry-sdk==1.28.0", "opentelemetry-exporter-otlp==1.28.0", "opentelemetry-instrumentation-fastapi==0.49b0", - "ddtrace>=2.19.0,<3.0", + "ddtrace>=4.8.2,<5.0", "sentry-sdk>=2.21.0,<3.0", "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", diff --git a/scripts/install-cli.sh b/scripts/install-cli.sh index d147286fcac..a39b73c2e5a 100755 --- a/scripts/install-cli.sh +++ b/scripts/install-cli.sh @@ -11,12 +11,21 @@ # Python itself (honouring litellm's requires-python), downloading a managed one # when the host has no suitable interpreter. # +# To try an unreleased branch instead of the latest PyPI release (for example, to +# QA a CLI feature before it ships), set LITELLM_CLI_REF to a branch, tag, or commit: +# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install-cli.sh | \ +# LITELLM_CLI_REF= sh +# # NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian # ignores the shebang when invoked as `sh` and does not support `pipefail`). set -eu -# NOTE: before merging, this must stay as "litellm[cli]" to install from PyPI. -LITELLM_PACKAGE="litellm[cli]" +# Defaults to the PyPI release; LITELLM_CLI_REF opts into installing from source instead. +if [ -n "${LITELLM_CLI_REF:-}" ]; then + LITELLM_PACKAGE="litellm[cli] @ git+https://github.com/BerriAI/litellm.git@${LITELLM_CLI_REF}" +else + LITELLM_PACKAGE="litellm[cli]" +fi UV_VERSION="0.10.9" # ── colours ──────────────────────────────────────────────────────────────── @@ -90,7 +99,11 @@ fi # otherwise download a managed one. Either way uv honours litellm's requires-python, # so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced. echo "" -header "Installing litellm[cli]…" +if [ -n "${LITELLM_CLI_REF:-}" ]; then + header "Installing litellm[cli] from ${LITELLM_CLI_REF}…" +else + header "Installing litellm[cli]…" +fi echo "" "$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \ diff --git a/scripts/install.sh b/scripts/install.sh index 06e6249c9ba..213f8a7b440 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -5,12 +5,24 @@ # Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible # Python itself (reusing a suitable system one, else downloading a managed build). # +# To install from an unreleased branch, tag, or commit instead of the latest PyPI +# release, set LITELLM_CLI_REF: +# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm//scripts/install.sh | \ +# LITELLM_CLI_REF= sh +# # NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian # ignores the shebang when invoked as `sh` and does not support `pipefail`). set -eu # NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI. -LITELLM_PACKAGE="litellm[proxy]" +# LITELLM_CLI_REF opts into installing from a branch, tag, or commit instead (for +# example, to QA lite autoroute against an unreleased branch, which needs this proxy +# runtime, not the thin litellm[cli] install). +if [ -n "${LITELLM_CLI_REF:-}" ]; then + LITELLM_PACKAGE="litellm[proxy] @ git+https://github.com/BerriAI/litellm.git@${LITELLM_CLI_REF}" +else + LITELLM_PACKAGE="litellm[proxy]" +fi UV_VERSION="0.10.9" # ── colours ──────────────────────────────────────────────────────────────── @@ -81,7 +93,11 @@ fi # ── install ──────────────────────────────────────────────────────────────── echo "" -header "Installing litellm[proxy]…" +if [ -n "${LITELLM_CLI_REF:-}" ]; then + header "Installing litellm[proxy] from ${LITELLM_CLI_REF}…" +else + header "Installing litellm[proxy]…" +fi echo "" # --python-preference system: reuse a compatible system Python when present, diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 7d4ef0a14fb..40a9da66c70 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -158,6 +158,38 @@ AgentOps) live under `proxy_config.litellm_settings.callbacks` and are orthogonal to the OTLP variables above; their credentials still go in `*_extra_secrets`. +### Enterprise billing metrics + +License-gated request metering is opt-in and gated entirely on +`billing_metrics_endpoint`. Empty (default) and no billing env is added to +the container, so existing deployments are unchanged. Set it and both +gateway and backend export billable-request counts over OTLP/HTTP, +authenticating to the collector with the mTLS client certificate issued for +your deployment. + +The proxy accepts the certificate, key, and CA bundle as either a file path +or literal PEM content. This stack takes the PEM, writes each one to its own +Secrets Manager entry, grants the task-execution role +`secretsmanager:GetSecretValue` on them, and injects them as +`LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and `_CA_CERT` when +set), so no volume mount is needed on Fargate. + +```hcl +billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics" +``` + +```bash +export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)" +export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)" +``` + +`billing_metrics_ca_cert_pem` is only for private or test collectors whose +CA is not in the system trust store; leave it empty against +`telemetry.litellm.ai`. Metering requires an enterprise license, so pair +this with `litellm_license`. To tune the export cadence, set +`LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / +`backend_extra_env` + ## Tenant deployment Every resource the stack creates is named `${tenant}-litellm-${env}` (or diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 54ab80de9f4..4df41c278e8 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -76,6 +76,33 @@ locals { { name = "OTEL_HEADERS", valueFrom = var.otel_headers_secret_arn }, ] : [] + # Enterprise request metering, gated on billing_metrics_endpoint. The + # endpoint rides in as a plain env var; the mTLS material is stored in + # Secrets Manager (secrets.tf) and injected as PEM-valued env vars, which + # the proxy accepts in place of file paths. Each PEM is wired only when the + # operator supplied it, so an empty ca_cert_pem falls back to the system + # trust store. + billing_metrics_enabled = var.billing_metrics_endpoint != "" + billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != "" + billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != "" + billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != "" + + billing_metrics_env = local.billing_metrics_enabled ? [ + { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint }, + ] : [] + + billing_metrics_secrets = concat( + local.billing_metrics_client_cert_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_client_cert[0].arn }, + ] : [], + local.billing_metrics_client_key_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", valueFrom = aws_secretsmanager_secret.billing_metrics_client_key[0].arn }, + ] : [], + local.billing_metrics_ca_cert_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CA_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_ca_cert[0].arn }, + ] : [], + ) + shared_env = [ { name = "IAM_TOKEN_DB_AUTH", value = "true" }, { name = "DATABASE_HOST", value = aws_rds_cluster.this.endpoint }, @@ -108,6 +135,7 @@ locals { { name = "LITELLM_LICENSE", valueFrom = aws_secretsmanager_secret.license[0].arn }, ], local.otel_secrets, + local.billing_metrics_secrets, ) # Backend-only managed secrets. UI_PASSWORD is consumed by the management @@ -179,6 +207,30 @@ locals { # ---------- Gateway ---------- resource "aws_ecs_task_definition" "gateway" { + # Metering needs a client certificate AND its key. Each secret is created only + # when its own PEM is supplied, so an endpoint set with a missing key would + # otherwise apply cleanly and leave the proxy logging "missing config" and + # never exporting. ca_cert_pem stays optional: empty means fall back to the + # system trust store. + # + # The guard lives here, on an unconditional resource, rather than on the cert + # secret: that secret is count-gated on the cert itself, so it has zero + # instances in exactly the case this must catch. Adding count or for_each to + # this resource would silently stop the guard from evaluating. + # + # endpoint cert key -> result + # "" any any -> metering off, no secrets created + # set set set -> metering on + # set any-missing -> plan fails here + lifecycle { + precondition { + condition = var.billing_metrics_endpoint == "" || ( + var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + ) + error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + } + } + family = "${local.name}-gateway" network_mode = "awsvpc" requires_compatibilities = ["FARGATE"] @@ -198,6 +250,7 @@ resource "aws_ecs_task_definition" "gateway" { environment = concat( local.shared_env, local.gateway_otel_env, + local.billing_metrics_env, local.gateway_extra_env_list, local.proxy_config_env, ) @@ -264,6 +317,18 @@ resource "aws_ecs_service" "gateway" { # ---------- Backend ---------- resource "aws_ecs_task_definition" "backend" { + # Same guard as the gateway: the backend meters too (it serves the named-server + # MCP transport), and a targeted apply of just this resource must not slip a + # billing endpoint through without the credentials to use it. + lifecycle { + precondition { + condition = var.billing_metrics_endpoint == "" || ( + var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + ) + error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + } + } + family = "${local.name}-backend" network_mode = "awsvpc" requires_compatibilities = ["FARGATE"] @@ -284,6 +349,7 @@ resource "aws_ecs_task_definition" "backend" { local.shared_env, local.backend_default_env, local.backend_otel_env, + local.billing_metrics_env, local.backend_extra_env_list, local.proxy_config_env, ) diff --git a/terraform/litellm/aws/iam.tf b/terraform/litellm/aws/iam.tf index 64e1b1ad5f9..63c6c26f184 100644 --- a/terraform/litellm/aws/iam.tf +++ b/terraform/litellm/aws/iam.tf @@ -53,6 +53,9 @@ data "aws_iam_policy_document" "secrets_access" { [aws_secretsmanager_secret.master_key.arn], aws_secretsmanager_secret.license[*].arn, aws_secretsmanager_secret.ui_password[*].arn, + aws_secretsmanager_secret.billing_metrics_client_cert[*].arn, + aws_secretsmanager_secret.billing_metrics_client_key[*].arn, + aws_secretsmanager_secret.billing_metrics_ca_cert[*].arn, local.extra_secret_arns, var.otel_headers_secret_arn == "" ? [] : [var.otel_headers_secret_arn], ) diff --git a/terraform/litellm/aws/secrets.tf b/terraform/litellm/aws/secrets.tf index 300d38e4053..85d3eb4502c 100644 --- a/terraform/litellm/aws/secrets.tf +++ b/terraform/litellm/aws/secrets.tf @@ -74,6 +74,61 @@ resource "aws_secretsmanager_secret_version" "ui_password" { secret_string = var.ui_password } +# Billing-metrics mTLS material — only created when metering is enabled +# (billing_metrics_endpoint non-empty) and the operator supplied the PEM. +# The task-execution role gets GetSecretValue via iam.tf, and gateway + +# backend pick the env vars up through shared_secrets in ecs.tf. +resource "aws_secretsmanager_secret" "billing_metrics_client_cert" { + count = local.billing_metrics_client_cert_enabled ? 1 : 0 + + name = "${local.name}-billing-metrics-client-cert" + description = "LITELLM_BILLING_METRICS_CLIENT_CERT for gateway + backend." + recovery_window_in_days = 0 + + tags = local.tags +} + +resource "aws_secretsmanager_secret_version" "billing_metrics_client_cert" { + count = local.billing_metrics_client_cert_enabled ? 1 : 0 + + secret_id = aws_secretsmanager_secret.billing_metrics_client_cert[0].id + secret_string = var.billing_metrics_client_cert_pem +} + +resource "aws_secretsmanager_secret" "billing_metrics_client_key" { + count = local.billing_metrics_client_key_enabled ? 1 : 0 + + name = "${local.name}-billing-metrics-client-key" + description = "LITELLM_BILLING_METRICS_CLIENT_KEY for gateway + backend." + recovery_window_in_days = 0 + + tags = local.tags +} + +resource "aws_secretsmanager_secret_version" "billing_metrics_client_key" { + count = local.billing_metrics_client_key_enabled ? 1 : 0 + + secret_id = aws_secretsmanager_secret.billing_metrics_client_key[0].id + secret_string = var.billing_metrics_client_key_pem +} + +resource "aws_secretsmanager_secret" "billing_metrics_ca_cert" { + count = local.billing_metrics_ca_cert_enabled ? 1 : 0 + + name = "${local.name}-billing-metrics-ca-cert" + description = "LITELLM_BILLING_METRICS_CA_CERT for gateway + backend." + recovery_window_in_days = 0 + + tags = local.tags +} + +resource "aws_secretsmanager_secret_version" "billing_metrics_ca_cert" { + count = local.billing_metrics_ca_cert_enabled ? 1 : 0 + + secret_id = aws_secretsmanager_secret.billing_metrics_ca_cert[0].id + secret_string = var.billing_metrics_ca_cert_pem +} + resource "aws_secretsmanager_secret" "db_master_password" { name = "${local.name}-db-master-password" description = "Aurora master-user password - bootstrap only. Runtime auth is IAM-token." diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 8db4935664b..c2ed1db14b1 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -533,3 +533,65 @@ variable "otel_headers_secret_arn" { type = string default = "" } + +# ---------- Enterprise billing metrics ---------- +# +# License-gated request metering. Opt-in and gated entirely on +# billing_metrics_endpoint: leave it empty (the default) and nothing +# metering-related lands in the container env. Set it and gateway + backend +# export billable-request counts over OTLP/HTTP, authenticating to the +# collector with an mTLS client cert. The proxy accepts the cert, key, and CA +# as either a file path or literal PEM content, so on Fargate they are +# injected straight from Secrets Manager as env vars and no volume is needed. + +variable "billing_metrics_endpoint" { + description = <<-EOT + OTLP/HTTP endpoint for enterprise billing metrics (sets + LITELLM_BILLING_METRICS_ENDPOINT). Non-empty enables request metering; + empty (default) disables it and adds no billing env to the container. + Requires an enterprise license. Example: + "https://telemetry.litellm.ai/v1/metrics" + EOT + type = string + default = "" +} + +variable "billing_metrics_client_cert_pem" { + description = <<-EOT + PEM content of the mTLS client certificate issued for this deployment. + When billing_metrics_endpoint is set, the stack stores this in a + `-litellm--billing-metrics-client-cert` Secrets Manager + entry, grants the task-execution role GetSecretValue on it, and exposes + it to gateway + backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required + whenever metering is enabled. + EOT + type = string + default = "" + sensitive = true +} + +variable "billing_metrics_client_key_pem" { + description = <<-EOT + PEM content of the private key matching + billing_metrics_client_cert_pem. Stored in a + `-litellm--billing-metrics-client-key` Secrets Manager + entry and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required + whenever metering is enabled. + EOT + type = string + default = "" + sensitive = true +} + +variable "billing_metrics_ca_cert_pem" { + description = <<-EOT + PEM content of the CA bundle used to verify the metering collector. + Only needed for private or test collectors whose CA is not in the + system trust store; telemetry.litellm.ai is publicly trusted, so leave + this empty for production. When set, it is exposed as + LITELLM_BILLING_METRICS_CA_CERT. + EOT + type = string + default = "" + sensitive = true +} diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 1e0bf4319df..88e9979148f 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -204,6 +204,40 @@ Behavior matches the AWS stack 1:1; the only naming differences are `otel_headers_secret` (a Secret Manager resource ID) vs AWS's `otel_headers_secret_arn` (a Secrets Manager ARN). +### Enterprise billing metrics + +License-gated request metering is opt-in and gated entirely on +`billing_metrics_endpoint`. Empty (default) and no billing env is added to +the container, so existing deployments are unchanged. Set it and both +gateway and backend export billable-request counts over OTLP/HTTP, +authenticating to the collector with the mTLS client certificate issued for +your deployment. + +The proxy accepts the certificate, key, and CA bundle as either a file path +or literal PEM content. This stack takes the PEM, writes each one to its own +Secret Manager entry, grants the runtime service account +`roles/secretmanager.secretAccessor` on them, and injects them as Cloud Run +secret env vars `LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and +`_CA_CERT` when set), so no volume mount is needed. + +```hcl +billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics" +``` + +```bash +export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)" +export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)" +``` + +`billing_metrics_ca_cert_pem` is only for private or test collectors whose +CA is not in the system trust store; leave it empty against +`telemetry.litellm.ai`. Metering requires an enterprise license, so pair +this with `litellm_license`. To tune the export cadence, set +`LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / +`backend_extra_env` + +Behavior matches the AWS stack 1:1; the variable names are identical + ## Tenant deployment Every resource the stack creates is named `${tenant}-litellm-${env}` (or diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 7b1bb901e20..57533b71731 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -59,6 +59,33 @@ locals { { name = "OTEL_HEADERS", secret = var.otel_headers_secret, version = "latest" }, ] : [] + # Enterprise request metering, gated on billing_metrics_endpoint. The + # endpoint rides in as a plain env var; the mTLS material lives in Secret + # Manager (secrets.tf) and is injected as PEM-valued env vars, which the + # proxy accepts in place of file paths. Each PEM is wired only when the + # operator supplied it, so an empty ca_cert_pem falls back to the system + # trust store. + billing_metrics_enabled = var.billing_metrics_endpoint != "" + billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != "" + billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != "" + billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != "" + + billing_metrics_env_kv = local.billing_metrics_enabled ? [ + { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint }, + ] : [] + + billing_metrics_env_secrets = concat( + local.billing_metrics_client_cert_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", secret = google_secret_manager_secret.billing_metrics_client_cert[0].id, version = "latest" }, + ] : [], + local.billing_metrics_client_key_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", secret = google_secret_manager_secret.billing_metrics_client_key[0].id, version = "latest" }, + ] : [], + local.billing_metrics_ca_cert_enabled ? [ + { name = "LITELLM_BILLING_METRICS_CA_CERT", secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id, version = "latest" }, + ] : [], + ) + # Cloud Run v2 secret env vars use value_source.secret_key_ref pointing at a # secret resource ID. Shared between gateway and backend (the migrations # job has its own narrower env list — see migrations_env_secrets below). @@ -138,6 +165,30 @@ locals { # ---------- Gateway ---------- resource "google_cloud_run_v2_service" "gateway" { + # Metering needs a client certificate AND its key. Each secret is created only + # when its own PEM is supplied, so an endpoint set with a missing key would + # otherwise apply cleanly and leave the proxy logging "missing config" and + # never exporting. ca_cert_pem stays optional: empty means fall back to the + # system trust store. + # + # The guard lives here, on an unconditional resource, rather than on the cert + # secret: that secret is count-gated on the cert itself, so it has zero + # instances in exactly the case this must catch. Adding count or for_each to + # this resource would silently stop the guard from evaluating. + # + # endpoint cert key -> result + # "" any any -> metering off, no secrets created + # set set set -> metering on + # set any-missing -> plan fails here + lifecycle { + precondition { + condition = var.billing_metrics_endpoint == "" || ( + var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + ) + error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + } + } + name = "${local.name}-gateway" location = var.region ingress = "INGRESS_TRAFFIC_INTERNAL_LOAD_BALANCER" @@ -175,7 +226,7 @@ resource "google_cloud_run_v2_service" "gateway" { } dynamic "env" { - for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.gateway_extra_env_kv, local.proxy_config_env) + for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env) content { name = env.value.name value = env.value.value @@ -183,7 +234,7 @@ resource "google_cloud_run_v2_service" "gateway" { } dynamic "env" { - for_each = concat(local.shared_env_secrets, local.otel_env_secrets, local.gateway_extra_secret_kv) + for_each = concat(local.shared_env_secrets, local.otel_env_secrets, local.billing_metrics_env_secrets, local.gateway_extra_secret_kv) content { name = env.value.name value_source { @@ -242,6 +293,9 @@ resource "google_cloud_run_v2_service" "gateway" { google_secret_manager_secret_iam_member.license, google_secret_manager_secret_iam_member.extras, google_secret_manager_secret_iam_member.otel_headers, + google_secret_manager_secret_iam_member.billing_metrics_client_cert, + google_secret_manager_secret_iam_member.billing_metrics_client_key, + google_secret_manager_secret_iam_member.billing_metrics_ca_cert, google_storage_bucket_iam_member.proxy_config_runtime, google_sql_user.app, # Don't go live until the schema is migrated; otherwise the proxy boots, @@ -252,6 +306,18 @@ resource "google_cloud_run_v2_service" "gateway" { # ---------- Backend ---------- resource "google_cloud_run_v2_service" "backend" { + # Same guard as the gateway: the backend meters too (it serves the named-server + # MCP transport), and a targeted apply of just this resource must not slip a + # billing endpoint through without the credentials to use it. + lifecycle { + precondition { + condition = var.billing_metrics_endpoint == "" || ( + var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + ) + error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." + } + } + name = "${local.name}-backend" location = var.region ingress = "INGRESS_TRAFFIC_INTERNAL_LOAD_BALANCER" @@ -289,7 +355,7 @@ resource "google_cloud_run_v2_service" "backend" { } dynamic "env" { - for_each = concat(local.shared_env_kv, local.backend_default_env_kv, local.backend_otel_env_kv, local.backend_extra_env_kv, local.proxy_config_env) + for_each = concat(local.shared_env_kv, local.backend_default_env_kv, local.backend_otel_env_kv, local.billing_metrics_env_kv, local.backend_extra_env_kv, local.proxy_config_env) content { name = env.value.name value = env.value.value @@ -297,7 +363,7 @@ resource "google_cloud_run_v2_service" "backend" { } dynamic "env" { - for_each = concat(local.shared_env_secrets, local.backend_managed_env_secrets, local.otel_env_secrets, local.backend_extra_secret_kv) + for_each = concat(local.shared_env_secrets, local.backend_managed_env_secrets, local.otel_env_secrets, local.billing_metrics_env_secrets, local.backend_extra_secret_kv) content { name = env.value.name value_source { @@ -357,6 +423,9 @@ resource "google_cloud_run_v2_service" "backend" { google_secret_manager_secret_iam_member.ui_password, google_secret_manager_secret_iam_member.extras, google_secret_manager_secret_iam_member.otel_headers, + google_secret_manager_secret_iam_member.billing_metrics_client_cert, + google_secret_manager_secret_iam_member.billing_metrics_client_key, + google_secret_manager_secret_iam_member.billing_metrics_ca_cert, google_storage_bucket_iam_member.proxy_config_runtime, google_sql_user.app, terraform_data.migration, diff --git a/terraform/litellm/gcp/iam.tf b/terraform/litellm/gcp/iam.tf index dc3ae5e0912..09df5e7dff0 100644 --- a/terraform/litellm/gcp/iam.tf +++ b/terraform/litellm/gcp/iam.tf @@ -79,3 +79,29 @@ resource "google_secret_manager_secret_iam_member" "otel_headers" { role = "roles/secretmanager.secretAccessor" member = "serviceAccount:${google_service_account.runtime.email}" } + +# Billing-metrics mTLS accessors — only created when request metering is +# enabled and the matching PEM was supplied. +resource "google_secret_manager_secret_iam_member" "billing_metrics_client_cert" { + count = local.billing_metrics_client_cert_enabled ? 1 : 0 + + secret_id = google_secret_manager_secret.billing_metrics_client_cert[0].id + role = "roles/secretmanager.secretAccessor" + member = "serviceAccount:${google_service_account.runtime.email}" +} + +resource "google_secret_manager_secret_iam_member" "billing_metrics_client_key" { + count = local.billing_metrics_client_key_enabled ? 1 : 0 + + secret_id = google_secret_manager_secret.billing_metrics_client_key[0].id + role = "roles/secretmanager.secretAccessor" + member = "serviceAccount:${google_service_account.runtime.email}" +} + +resource "google_secret_manager_secret_iam_member" "billing_metrics_ca_cert" { + count = local.billing_metrics_ca_cert_enabled ? 1 : 0 + + secret_id = google_secret_manager_secret.billing_metrics_ca_cert[0].id + role = "roles/secretmanager.secretAccessor" + member = "serviceAccount:${google_service_account.runtime.email}" +} diff --git a/terraform/litellm/gcp/secrets.tf b/terraform/litellm/gcp/secrets.tf index f93514bb70b..6ec77139996 100644 --- a/terraform/litellm/gcp/secrets.tf +++ b/terraform/litellm/gcp/secrets.tf @@ -63,3 +63,58 @@ resource "google_secret_manager_secret_version" "ui_password" { secret = google_secret_manager_secret.ui_password[0].id secret_data = var.ui_password } + +# Billing-metrics mTLS material — only created when metering is enabled +# (billing_metrics_endpoint non-empty) and the operator supplied the PEM. +# The runtime SA gets accessor permission via iam.tf, and gateway + backend +# pick the env vars up through billing_metrics_env_secrets in cloudrun.tf. +resource "google_secret_manager_secret" "billing_metrics_client_cert" { + count = local.billing_metrics_client_cert_enabled ? 1 : 0 + + secret_id = "${local.name}-billing-metrics-client-cert" + labels = local.labels + replication { + auto {} + } +} + +resource "google_secret_manager_secret_version" "billing_metrics_client_cert" { + count = local.billing_metrics_client_cert_enabled ? 1 : 0 + + secret = google_secret_manager_secret.billing_metrics_client_cert[0].id + secret_data = var.billing_metrics_client_cert_pem +} + +resource "google_secret_manager_secret" "billing_metrics_client_key" { + count = local.billing_metrics_client_key_enabled ? 1 : 0 + + secret_id = "${local.name}-billing-metrics-client-key" + labels = local.labels + replication { + auto {} + } +} + +resource "google_secret_manager_secret_version" "billing_metrics_client_key" { + count = local.billing_metrics_client_key_enabled ? 1 : 0 + + secret = google_secret_manager_secret.billing_metrics_client_key[0].id + secret_data = var.billing_metrics_client_key_pem +} + +resource "google_secret_manager_secret" "billing_metrics_ca_cert" { + count = local.billing_metrics_ca_cert_enabled ? 1 : 0 + + secret_id = "${local.name}-billing-metrics-ca-cert" + labels = local.labels + replication { + auto {} + } +} + +resource "google_secret_manager_secret_version" "billing_metrics_ca_cert" { + count = local.billing_metrics_ca_cert_enabled ? 1 : 0 + + secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id + secret_data = var.billing_metrics_ca_cert_pem +} diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 4355192e9f1..1162e100bb2 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -490,3 +490,66 @@ variable "otel_capture_message_content" { error_message = "otel_capture_message_content must be one of: no_content, prompt_and_completion." } } + +# ---------- Enterprise billing metrics ---------- +# +# License-gated request metering. Opt-in and gated entirely on +# billing_metrics_endpoint: leave it empty (the default) and nothing +# metering-related is added to the container env. Set it and gateway + +# backend export billable-request counts over OTLP/HTTP, authenticating to +# the collector with an mTLS client cert. The proxy accepts the cert, key, +# and CA as either a file path or literal PEM content, so on Cloud Run they +# are injected straight from Secret Manager as env vars and no volume is +# needed. + +variable "billing_metrics_endpoint" { + description = <<-EOT + OTLP/HTTP endpoint for enterprise billing metrics (sets + LITELLM_BILLING_METRICS_ENDPOINT). Non-empty enables request metering; + empty (default) disables it and adds no billing env to the container. + Requires an enterprise license. Example: + "https://telemetry.litellm.ai/v1/metrics" + EOT + type = string + default = "" +} + +variable "billing_metrics_client_cert_pem" { + description = <<-EOT + PEM content of the mTLS client certificate issued for this deployment. + When billing_metrics_endpoint is set, the stack stores this in a + `-litellm--billing-metrics-client-cert` Secret Manager + entry, grants the runtime SA accessor on it, and exposes it to gateway + + backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required whenever + metering is enabled. + EOT + type = string + default = "" + sensitive = true +} + +variable "billing_metrics_client_key_pem" { + description = <<-EOT + PEM content of the private key matching + billing_metrics_client_cert_pem. Stored in a + `-litellm--billing-metrics-client-key` Secret Manager entry + and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required whenever + metering is enabled. + EOT + type = string + default = "" + sensitive = true +} + +variable "billing_metrics_ca_cert_pem" { + description = <<-EOT + PEM content of the CA bundle used to verify the metering collector. + Only needed for private or test collectors whose CA is not in the + system trust store; telemetry.litellm.ai is publicly trusted, so leave + this empty for production. When set, it is exposed as + LITELLM_BILLING_METRICS_CA_CERT. + EOT + type = string + default = "" + sensitive = true +} diff --git a/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py b/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py index b1745008fac..a3569ebdb49 100644 --- a/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py +++ b/tests/e2e/claude_code/_builder_unit_tests/test_v0_layout.py @@ -16,8 +16,8 @@ from pathlib import Path import pytest import yaml -REPO_ROOT = Path(__file__).resolve().parents[1] -MANIFEST_PATH = REPO_ROOT / "manifest.yaml" +SUITE_ROOT = Path(__file__).resolve().parents[1] +MANIFEST_PATH = SUITE_ROOT / "manifest.yaml" # The PRD's "Features in v0" section, in row order. EXPECTED_FEATURE_IDS = [ @@ -90,14 +90,14 @@ def test_manifest_every_feature_has_human_readable_name(manifest): @pytest.mark.parametrize("feature_id", EXPECTED_FEATURE_IDS) def test_feature_directory_exists(feature_id): - feature_dir = REPO_ROOT / feature_id + feature_dir = SUITE_ROOT / feature_id assert feature_dir.is_dir(), f"missing feature directory: {feature_dir}" @pytest.mark.parametrize("feature_id", EXPECTED_FEATURE_IDS) @pytest.mark.parametrize("provider", EXPECTED_PROVIDERS) def test_per_provider_test_file_exists(feature_id, provider): - test_file = REPO_ROOT / feature_id / f"test_{provider}.py" + test_file = SUITE_ROOT / feature_id / f"test_{provider}.py" assert test_file.is_file(), f"missing per-provider test file: {test_file}" @@ -106,7 +106,7 @@ def test_feature_directory_has_init_file(feature_id): """Each feature directory needs an __init__.py so pytest collects the per-provider test files as a package — matches the layout established by `basic_messaging_non_streaming/`.""" - init_file = REPO_ROOT / feature_id / "__init__.py" + init_file = SUITE_ROOT / feature_id / "__init__.py" assert init_file.is_file(), f"missing __init__.py: {init_file}" @@ -117,7 +117,7 @@ def test_feature_directory_has_init_file(feature_id): # a broken post-v0 directory still fails CI. @pytest.mark.parametrize("feature_id", ALL_FEATURE_IDS) def test_every_manifest_feature_has_directory(feature_id): - feature_dir = REPO_ROOT / feature_id + feature_dir = SUITE_ROOT / feature_id assert feature_dir.is_dir(), ( f"manifest declares {feature_id!r} but {feature_dir} is missing — " "feature_id MUST match its on-disk directory (see manifest.yaml header)." @@ -126,7 +126,7 @@ def test_every_manifest_feature_has_directory(feature_id): @pytest.mark.parametrize("feature_id", ALL_FEATURE_IDS) def test_every_manifest_feature_has_init_file(feature_id): - init_file = REPO_ROOT / feature_id / "__init__.py" + init_file = SUITE_ROOT / feature_id / "__init__.py" assert init_file.is_file(), f"missing __init__.py: {init_file}" @@ -137,7 +137,7 @@ def test_every_manifest_feature_has_per_provider_test_file(feature_id, provider) backed by a per-provider test file. Without this check, a missing file silently becomes a `not_tested` cell in the published matrix rather than a CI failure surfacing the layout drift.""" - test_file = REPO_ROOT / feature_id / f"test_{provider}.py" + test_file = SUITE_ROOT / feature_id / f"test_{provider}.py" assert test_file.is_file(), f"missing per-provider test file: {test_file}" @@ -151,7 +151,7 @@ def test_per_provider_test_file_imports_and_parametrizes_three_models( use plain aliases or per-provider-suffixed aliases (e.g. `claude-opus-4-7-bedrock-invoke`), so we check for the tier substrings rather than exact alias names.""" - text = (REPO_ROOT / feature_id / f"test_{provider}.py").read_text() + text = (SUITE_ROOT / feature_id / f"test_{provider}.py").read_text() for tier in ("haiku-4-5", "sonnet-4-6", "opus-4-7"): assert ( tier in text @@ -171,7 +171,7 @@ def test_azure_test_file_drives_the_proxy(feature_id): that wraps them — both shapes drive the proxy, and we don't want this layout pin to block legitimate de-duplication of test bodies. """ - text = (REPO_ROOT / feature_id / "test_azure.py").read_text() + text = (SUITE_ROOT / feature_id / "test_azure.py").read_text() assert "run_claude" in text or "run_basic_messaging_cell" in text, ( f"{feature_id}/test_azure.py must drive the claude CLI via run_claude() " "or a shared helper that wraps it; the not_applicable stub was removed " diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py b/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py index 786f6029993..f1a0534906e 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_cli_driver.py @@ -19,6 +19,7 @@ from claude_code.cli_driver import ( ClaudeCLIError, DriverResult, failure_diagnostic, + is_rate_limit_shaped, run_claude, run_claude_models_parallel, ) @@ -793,3 +794,199 @@ def test_failure_diagnostic_uses_last_result_event_status(): diag = failure_diagnostic(result) assert "api_status=429" in diag assert "500" not in diag + + +_RATE_LIMITED_STDOUT = ( + json.dumps( + { + "type": "assistant", + "message": { + "content": [ + {"type": "text", "text": "API Error: 429 Too Many Requests"} + ] + }, + } + ) + + "\n" + + json.dumps({"type": "result", "api_error_status": 429}) + + "\n" +) + +_OK_STDOUT = ( + json.dumps( + { + "type": "assistant", + "message": {"content": [{"type": "text", "text": "pong"}]}, + } + ) + + "\n" +) + + +class _FlakyRunner: + """Fake runner that rate-limits each model N times before succeeding. + + Keeps a per-model call count so tests can assert exactly how many + attempts the retry loop made — the load-bearing detail a canned + single-response runner can't express. + """ + + def __init__(self, failures_before_success: dict): + self.failures_before_success = dict(failures_before_success) + self.calls: dict = {} + + def __call__(self, cmd, env, capture_output, text, timeout, check, input=None): + model = cmd[cmd.index("--model") + 1] + self.calls[model] = self.calls.get(model, 0) + 1 + if self.calls[model] <= self.failures_before_success.get(model, 0): + return _Completed(returncode=1, stdout=_RATE_LIMITED_STDOUT) + return _Completed(returncode=0, stdout=_OK_STDOUT) + + +@pytest.mark.parametrize( + "outcome,expected", + [ + (ClaudeCLIError("claude CLI timed out after 120.0s"), True), + (ClaudeCLIError("claude CLI not found at 'claude'"), False), + ( + DriverResult( + text="", + events=[{"type": "result", "api_error_status": 429}], + exit_code=1, + ), + True, + ), + (DriverResult(text="Too Many Requests", exit_code=1), True), + (DriverResult(text="", stderr="throttled by upstream", exit_code=1), True), + (DriverResult(text="rate limit exceeded", exit_code=0), False), + (DriverResult(text="", stderr="auth failed", exit_code=2), False), + ], +) +def test_is_rate_limit_shaped_classification(outcome, expected): + """The retry trigger must match 429/throttle/timeout markers on + failures only — a passing result mentioning '429' in its reply text + must never be classified as retryable.""" + assert is_rate_limit_shaped(outcome) is expected + + +def test_run_claude_models_parallel_retries_rate_limited_model_until_success(): + """A model that 429s once must be retried after the backoff sleep and + end up green, while an untroubled sibling model runs exactly once.""" + runner = _FlakyRunner({"flaky": 1}) + sleeps: List[float] = [] + + outcomes = run_claude_models_parallel( + models=["flaky", "steady"], + prompt="hi", + base_url="http://x", + api_key="k", + runner=runner, + rate_limit_retries=2, + rate_limit_backoff_seconds=0.5, + sleep=sleeps.append, + ) + + assert isinstance(outcomes["flaky"], DriverResult) + assert outcomes["flaky"].exit_code == 0 + assert outcomes["flaky"].text == "pong" + assert runner.calls == {"flaky": 2, "steady": 1} + assert sleeps == [0.5] + + +def test_run_claude_models_parallel_does_not_retry_non_rate_limit_failures(): + """A deterministic failure (bad auth) must fail fast: no sleeps, one + attempt — retrying it would just triple the matrix wall time.""" + + def runner(cmd, env, capture_output, text, timeout, check, input=None): + return _Completed(returncode=2, stdout="", stderr="auth failed") + + sleeps: List[float] = [] + outcomes = run_claude_models_parallel( + models=["a"], + prompt="hi", + base_url="http://x", + api_key="k", + runner=runner, + rate_limit_retries=2, + rate_limit_backoff_seconds=0.5, + sleep=sleeps.append, + ) + + assert outcomes["a"].exit_code == 2 + assert sleeps == [] + + +def test_run_claude_models_parallel_returns_last_failure_when_retries_exhausted(): + """A persistently rate-limited model exhausts its budget (initial + attempt + N retries, each preceded by one backoff sleep) and still + surfaces the 429 diagnostic instead of masking it.""" + runner = _FlakyRunner({"stuck": 99}) + sleeps: List[float] = [] + + outcomes = run_claude_models_parallel( + models=["stuck"], + prompt="hi", + base_url="http://x", + api_key="k", + runner=runner, + rate_limit_retries=2, + rate_limit_backoff_seconds=0.25, + sleep=sleeps.append, + ) + + assert runner.calls == {"stuck": 3} + assert sleeps == [0.25, 0.25] + assert outcomes["stuck"].exit_code == 1 + assert "429" in failure_diagnostic(outcomes["stuck"]) + + +def test_run_claude_models_parallel_retries_timeout_shaped_cli_errors(): + """CLI timeouts are how saturated upstreams usually present (the CLI + retries 429s internally until the harness kills it), so a timeout + must be retried like an explicit 429.""" + calls: List[int] = [] + + def runner(cmd, env, capture_output, text, timeout, check, input=None): + calls.append(1) + if len(calls) == 1: + raise subprocess.TimeoutExpired(cmd="claude", timeout=1) + return _Completed(returncode=0, stdout=_OK_STDOUT) + + sleeps: List[float] = [] + outcomes = run_claude_models_parallel( + models=["a"], + prompt="hi", + base_url="http://x", + api_key="k", + runner=runner, + rate_limit_retries=1, + rate_limit_backoff_seconds=0.5, + sleep=sleeps.append, + ) + + assert isinstance(outcomes["a"], DriverResult) + assert outcomes["a"].text == "pong" + assert len(calls) == 2 + assert sleeps == [0.5] + + +def test_run_claude_models_parallel_zero_retries_disables_backoff(): + """`rate_limit_retries=0` must restore the old single-attempt + behavior exactly: one call, no sleeps, failure returned as-is.""" + runner = _FlakyRunner({"stuck": 99}) + sleeps: List[float] = [] + + outcomes = run_claude_models_parallel( + models=["stuck"], + prompt="hi", + base_url="http://x", + api_key="k", + runner=runner, + rate_limit_retries=0, + rate_limit_backoff_seconds=0.5, + sleep=sleeps.append, + ) + + assert runner.calls == {"stuck": 1} + assert sleeps == [] + assert outcomes["stuck"].exit_code == 1 diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py b/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py new file mode 100644 index 00000000000..2d17a84d418 --- /dev/null +++ b/tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py @@ -0,0 +1,195 @@ +"""Unit tests for the shared `run_passthrough_cell` helper. + +These tests inject a fake `run_models` callable and an explicit `env` +mapping (both are first-class parameters, no monkeypatching), so they +exercise the helper's branching -- env-missing guard, base-URL +assembly, extra-env forwarding, per-model pass/fail -- without +spawning the real CLI. + +The env-builder tests pin the provider-mode contract itself: the +CLAUDE_CODE_USE_* / CLAUDE_CODE_SKIP_*_AUTH flags and the passthrough +route each mode must target. Those values are the feature -- e.g. +dropping the `/v1` from the vertex base URL produces a request Google +404s on -- so a mutation to any of them must fail here before it burns +a live matrix run. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Mapping, Optional + +import pytest + +from claude_code._passthrough import ( + ANTHROPIC_PASSTHROUGH_BASE_PATH, + CLIENT_SIDE_AWS_REGION, + VERTEX_PLACEHOLDER_PROJECT, + VERTEX_PLACEHOLDER_REGION, + bedrock_extra_env, + foundry_extra_env, + run_passthrough_cell, + vertex_extra_env, +) +from claude_code.cli_driver import ClaudeCLIError, DriverResult + +PROXY_ENV = { + "LITELLM_PROXY_BASE_URL": "http://localhost:4000", + "LITELLM_PROXY_API_KEY": "sk-test", +} + + +class _FakeResult: + def __init__(self) -> None: + self.rows: List[Dict[str, Any]] = [] + self.single: Optional[Dict[str, Any]] = None + + def set(self, payload: Mapping[str, Any]) -> None: + self.single = dict(payload) + + def add(self, payload: Mapping[str, Any]) -> None: + self.rows.append(dict(payload)) + + +def _fake_run_models(outcomes_by_model, captured: Dict[str, Any]): + def fake(*, models, prompt, base_url, api_key, extra_env=None, **_kwargs): + captured["models"] = list(models) + captured["prompt"] = prompt + captured["base_url"] = base_url + captured["api_key"] = api_key + captured["extra_env"] = dict(extra_env) if extra_env is not None else None + return {model: outcomes_by_model[model] for model in models} + + return fake + + +def test_env_missing_guard_reports_fail_and_aborts(): + fake_result = _FakeResult() + with pytest.raises(pytest.fail.Exception): + run_passthrough_cell( + compat_result=fake_result, + models=["claude-haiku-4-5"], + prompt="ping", + env={}, + ) + assert fake_result.single is not None + assert fake_result.single["status"] == "fail" + assert "LITELLM_PROXY_BASE_URL" in fake_result.single["error"] + + +def test_anthropic_base_path_appended_to_normalized_proxy_url(): + fake_result = _FakeResult() + captured: Dict[str, Any] = {} + outcome = DriverResult(text="pong") + + run_passthrough_cell( + compat_result=fake_result, + models=["claude-haiku-4-5"], + prompt="ping", + passthrough_base_path=ANTHROPIC_PASSTHROUGH_BASE_PATH, + run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured), + env={**PROXY_ENV, "LITELLM_PROXY_BASE_URL": "http://localhost:4000/"}, + ) + + assert captured["base_url"] == "http://localhost:4000/anthropic" + assert captured["extra_env"] is None + assert fake_result.rows == [{"status": "pass"}] + + +def test_extra_env_builder_receives_normalized_base_and_is_forwarded(): + fake_result = _FakeResult() + captured: Dict[str, Any] = {} + outcome = DriverResult(text="pong") + seen_bases: List[str] = [] + + def build(proxy_base: str) -> Dict[str, str]: + seen_bases.append(proxy_base) + return {"SOME_FLAG": "1"} + + run_passthrough_cell( + compat_result=fake_result, + models=["claude-haiku-4-5"], + prompt="ping", + build_extra_env=build, + run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured), + env={**PROXY_ENV, "LITELLM_PROXY_BASE_URL": "http://localhost:4000/"}, + ) + + assert seen_bases == ["http://localhost:4000"] + assert captured["extra_env"] == {"SOME_FLAG": "1"} + assert captured["base_url"] == "http://localhost:4000" + + +def test_per_model_failures_reported_individually(): + fake_result = _FakeResult() + captured: Dict[str, Any] = {} + outcomes = { + "claude-haiku-4-5": DriverResult(text="pong"), + "claude-sonnet-4-6": ClaudeCLIError("claude CLI timed out after 120s"), + "claude-opus-4-7": DriverResult(text="", exit_code=1), + } + + with pytest.raises(pytest.fail.Exception): + run_passthrough_cell( + compat_result=fake_result, + models=list(outcomes.keys()), + prompt="ping", + run_models=_fake_run_models(outcomes, captured), + env=PROXY_ENV, + ) + + statuses = [row["status"] for row in fake_result.rows] + assert statuses == ["pass", "fail", "fail"] + assert "timed out" in fake_result.rows[1]["error"] + assert "claude CLI failed" in fake_result.rows[2]["error"] + + +def test_empty_assistant_text_is_a_fail(): + fake_result = _FakeResult() + captured: Dict[str, Any] = {} + outcomes = {"claude-haiku-4-5": DriverResult(text=" ")} + + with pytest.raises(pytest.fail.Exception): + run_passthrough_cell( + compat_result=fake_result, + models=["claude-haiku-4-5"], + prompt="ping", + run_models=_fake_run_models(outcomes, captured), + env=PROXY_ENV, + ) + + assert fake_result.rows == [ + { + "status": "fail", + "error": "[claude-haiku-4-5] claude returned empty assistant text", + } + ] + + +def test_bedrock_extra_env_targets_proxy_bedrock_route(): + env = bedrock_extra_env("http://localhost:4000") + assert env == { + "CLAUDE_CODE_USE_BEDROCK": "1", + "CLAUDE_CODE_SKIP_BEDROCK_AUTH": "1", + "ANTHROPIC_BEDROCK_BASE_URL": "http://localhost:4000/bedrock", + "AWS_REGION": CLIENT_SIDE_AWS_REGION, + } + + +def test_vertex_extra_env_keeps_the_api_version_in_the_base_url(): + env = vertex_extra_env("http://localhost:4000") + assert env == { + "CLAUDE_CODE_USE_VERTEX": "1", + "CLAUDE_CODE_SKIP_VERTEX_AUTH": "1", + "ANTHROPIC_VERTEX_BASE_URL": "http://localhost:4000/vertex_ai/v1", + "ANTHROPIC_VERTEX_PROJECT_ID": VERTEX_PLACEHOLDER_PROJECT, + "CLOUD_ML_REGION": VERTEX_PLACEHOLDER_REGION, + } + + +def test_foundry_extra_env_targets_proxy_azure_route(): + env = foundry_extra_env("http://localhost:4000") + assert env == { + "CLAUDE_CODE_USE_FOUNDRY": "1", + "CLAUDE_CODE_SKIP_FOUNDRY_AUTH": "1", + "ANTHROPIC_FOUNDRY_BASE_URL": "http://localhost:4000/azure", + } diff --git a/tests/e2e/claude_code/_passthrough.py b/tests/e2e/claude_code/_passthrough.py new file mode 100644 index 00000000000..24b6d694d3f --- /dev/null +++ b/tests/e2e/claude_code/_passthrough.py @@ -0,0 +1,196 @@ +"""Shared body for the `passthrough` × compat cells. + +Every other matrix row drives the proxy's `/v1/messages` translation +layer: Claude Code speaks the first-party Anthropic wire and LiteLLM +transforms the request per provider. This row instead exercises +LiteLLM's *native passthrough* routes -- the "LLM gateway" +configuration documented at https://code.claude.com/docs/en/gateway -- +where Claude Code speaks each cloud's own wire format and the proxy +forwards it, attaching provider credentials on the way out: + + anthropic ANTHROPIC_BASE_URL={proxy}/anthropic. The CLI's + first-party wire, forwarded verbatim to + api.anthropic.com, so the model ids are real + Anthropic ids rather than proxy aliases. + bedrock_invoke CLAUDE_CODE_USE_BEDROCK=1 + + ANTHROPIC_BEDROCK_BASE_URL={proxy}/bedrock. The + CLI POSTs /model/{model}/invoke-with-response-stream; + the proxy recognizes a router alias in the model + segment, rewrites it to the deployment's upstream + model id, and SigV4-signs with its own AWS creds. + vertex_ai CLAUDE_CODE_USE_VERTEX=1 + + ANTHROPIC_VERTEX_BASE_URL={proxy}/vertex_ai/v1. + The CLI POSTs + .../models/{model}:streamRawPredict; the proxy + resolves a router alias in the model segment and + takes project, location, and credentials from the + deployment (which is why the deployment must set + `use_in_pass_through: true` -- see + test_config.yaml). + azure CLAUDE_CODE_USE_FOUNDRY=1 + + ANTHROPIC_FOUNDRY_BASE_URL={proxy}/azure. Foundry + mode sends the model in the JSON body, not the + URL, so the proxy's /azure route cannot resolve a + router alias and falls back to the env-configured + AZURE_API_BASE / AZURE_API_KEY target. + bedrock_converse not applicable -- Claude Code's bedrock mode is + InvokeModel-only; no Converse-wire client exists. + +Auth is the same in every mode: the CLI's provider-native signing is +disabled via CLAUDE_CODE_SKIP__AUTH, and the LiteLLM virtual +key travels as `Authorization: Bearer` (ANTHROPIC_AUTH_TOKEN), exactly +like the translation rows. The proxy holds the real provider +credentials. + +The per-mode env vars and URL shapes above were captured from a real +`claude` CLI (2.1.210) run against a request-logging sink, not from +docs; if a CLI release changes them, the cells fail with the CLI's own +diagnostic rather than silently testing the wrong wire. + +`run_models` and `env` are injection seams for +`_driver_unit_tests/test_passthrough.py`; production callers leave +them unset. +""" + +from __future__ import annotations + +import os +from typing import Any, Callable, Dict, Mapping, Optional, Sequence + +import pytest + +from claude_code.cli_driver import ( + ClaudeCLIError, + failure_diagnostic, + run_claude_models_parallel, +) + +PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" +PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" + +ANTHROPIC_PASSTHROUGH_BASE_PATH = "/anthropic" + +CLIENT_SIDE_AWS_REGION = "us-east-1" +"""Satisfies the CLI's embedded AWS SDK, which refuses to construct a +client without a region. The value never influences routing: the proxy +signs the upstream request with its own credentials and region.""" + +VERTEX_PLACEHOLDER_PROJECT = "proxy-resolved-project" +VERTEX_PLACEHOLDER_REGION = "us-east5" +"""The CLI refuses to build a Vertex URL without a project id and +region, but the proxy replaces both path segments with the resolved +deployment's `vertex_project` / `vertex_location` before forwarding, +so deliberately-fake values prove the resolution actually happened.""" + + +def bedrock_extra_env(proxy_base_url: str) -> Dict[str, str]: + return { + "CLAUDE_CODE_USE_BEDROCK": "1", + "CLAUDE_CODE_SKIP_BEDROCK_AUTH": "1", + "ANTHROPIC_BEDROCK_BASE_URL": f"{proxy_base_url}/bedrock", + "AWS_REGION": CLIENT_SIDE_AWS_REGION, + } + + +def vertex_extra_env(proxy_base_url: str) -> Dict[str, str]: + """Vertex-mode CLI env pointed at the proxy's /vertex_ai route. + + The `/v1` suffix on ANTHROPIC_VERTEX_BASE_URL is load-bearing: the + CLI's Vertex SDK ships its API version inside its *default* base + URL (`https://{region}-aiplatform.googleapis.com/v1`), so + overriding the base drops the version from the request path unless + the override carries it. LiteLLM's /vertex_ai route reuses the + incoming path verbatim when it contains `/projects/.../locations/...`, + so a version-less path would reach Google as + `aiplatform.googleapis.com/projects/...` and 404. + """ + return { + "CLAUDE_CODE_USE_VERTEX": "1", + "CLAUDE_CODE_SKIP_VERTEX_AUTH": "1", + "ANTHROPIC_VERTEX_BASE_URL": f"{proxy_base_url}/vertex_ai/v1", + "ANTHROPIC_VERTEX_PROJECT_ID": VERTEX_PLACEHOLDER_PROJECT, + "CLOUD_ML_REGION": VERTEX_PLACEHOLDER_REGION, + } + + +def foundry_extra_env(proxy_base_url: str) -> Dict[str, str]: + return { + "CLAUDE_CODE_USE_FOUNDRY": "1", + "CLAUDE_CODE_SKIP_FOUNDRY_AUTH": "1", + "ANTHROPIC_FOUNDRY_BASE_URL": f"{proxy_base_url}/azure", + } + + +def run_passthrough_cell( + *, + compat_result, + models: Sequence[str], + prompt: str, + passthrough_base_path: str = "", + build_extra_env: Optional[Callable[[str], Mapping[str, str]]] = None, + run_models: Callable[..., Mapping[str, Any]] = run_claude_models_parallel, + env: Optional[Mapping[str, str]] = None, +) -> None: + """Run the shared `passthrough` × cell body. + + `passthrough_base_path` is appended to the proxy base URL and + becomes the CLI's ANTHROPIC_BASE_URL (only the anthropic column + uses it; the cloud columns ignore ANTHROPIC_BASE_URL entirely once + their CLAUDE_CODE_USE_* flag is set). `build_extra_env` receives + the trailing-slash-normalized proxy base URL and returns the + provider-mode env for the CLI subprocess. + """ + environ = env if env is not None else os.environ + base_url = environ.get(PROXY_BASE_URL_ENV) + api_key = environ.get(PROXY_API_KEY_ENV) + if not base_url or not api_key: + compat_result.set( + { + "status": "fail", + "error": ( + f"missing required env: set {PROXY_BASE_URL_ENV} and " + f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy" + ), + } + ) + pytest.fail( + f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured", + pytrace=False, + ) + + proxy_base = base_url.rstrip("/") + extra_env = dict(build_extra_env(proxy_base)) if build_extra_env else None + + outcomes = run_models( + models=models, + prompt=prompt, + base_url=proxy_base + passthrough_base_path, + api_key=api_key, + extra_env=extra_env, + ) + + failures = [] + for model in models: + outcome = outcomes[model] + if isinstance(outcome, ClaudeCLIError): + error = f"[{model}] {outcome}" + compat_result.add({"status": "fail", "error": error}) + failures.append(error) + continue + + if outcome.exit_code != 0: + error = f"[{model}] claude CLI failed: {failure_diagnostic(outcome)}" + compat_result.add({"status": "fail", "error": error}) + failures.append(error) + continue + + if not outcome.text.strip(): + error = f"[{model}] claude returned empty assistant text" + compat_result.add({"status": "fail", "error": error}) + failures.append(error) + continue + + compat_result.add({"status": "pass"}) + + if failures: + pytest.fail("; ".join(failures), pytrace=False) diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 97eaa0e6847..5b18c1c291a 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -15,6 +15,7 @@ from __future__ import annotations import json import os +import re import shutil import subprocess import sys @@ -40,6 +41,29 @@ DEFAULT_TIMEOUT_SECONDS = float( os.environ.get("LITELLM_COMPAT_CLI_TIMEOUT_SECONDS") or 120 ) +RATE_LIMIT_SHAPED_RE = re.compile( + r"(?:\b429\b|rate[\s_-]?limit|too\s+many\s+requests|throttl(?:ed|ing)|" + r"claude\s+CLI\s+timed\s+out)", + re.IGNORECASE, +) +"""Heuristic shared with the conftest rate-limit summary: 429s and +throttle markers anywhere in the failure text, plus CLI timeouts -- +the CLI retries 429s internally until the harness timeout kills it, +so a saturated upstream usually surfaces as a timeout rather than a +clean 429.""" + +DEFAULT_RATE_LIMIT_RETRIES = int( + os.environ.get("LITELLM_COMPAT_RATE_LIMIT_RETRIES") or 2 +) +DEFAULT_RATE_LIMIT_BACKOFF_SECONDS = float( + os.environ.get("LITELLM_COMPAT_RATE_LIMIT_BACKOFF_SECONDS") or 65 +) +"""Rate-limit-shaped failures are retried after a backoff long enough +for a per-minute quota window (the dominant 429 source across +Anthropic / Bedrock / Vertex) to reset. Both knobs are env-tunable so +a matrix run can trade wall time for resilience without code edits; +retries=0 disables the behavior entirely.""" + # Env vars the `claude` Node CLI legitimately needs to function: # locating its own binary + node, basic locale/terminal plumbing. # Deliberately excludes every credential-bearing var that the @@ -265,6 +289,22 @@ def run_claude( ModelResult = Union[DriverResult, ClaudeCLIError] +def is_rate_limit_shaped(outcome: ModelResult) -> bool: + """Classify an outcome as a retryable rate-limit-shaped failure. + + A `ClaudeCLIError` matches on its message (which is where the + driver's own timeout diagnostic lands); a failing `DriverResult` + matches on its full `failure_diagnostic` so 429s buried in the + CLI's stdout text or `api_error_status` are both caught. Passing + results are never rate-limit-shaped. + """ + if isinstance(outcome, ClaudeCLIError): + return bool(RATE_LIMIT_SHAPED_RE.search(str(outcome))) + if outcome.exit_code == 0: + return False + return bool(RATE_LIMIT_SHAPED_RE.search(failure_diagnostic(outcome))) + + def run_claude_models_parallel( *, models: Sequence[str], @@ -277,6 +317,9 @@ def run_claude_models_parallel( cli_path: str = CLAUDE_CLI_DEFAULT, timeout: float = DEFAULT_TIMEOUT_SECONDS, runner: Optional[Callable[..., Any]] = None, + rate_limit_retries: Optional[int] = None, + rate_limit_backoff_seconds: Optional[float] = None, + sleep: Callable[[float], None] = time.sleep, ) -> Dict[str, ModelResult]: """Invoke `run_claude` for every `models[i]` concurrently and collect outcomes. @@ -290,6 +333,14 @@ def run_claude_models_parallel( keep the synchronous CLI driver unchanged so unit tests can keep injecting a fake `runner`. + Rate-limit-shaped failures (see `is_rate_limit_shaped`) are retried + per model up to `rate_limit_retries` times, sleeping + `rate_limit_backoff_seconds` before each retry so per-minute quota + windows can reset; both default to the `LITELLM_COMPAT_RATE_LIMIT_*` + env knobs. Each retry goes back through `run_claude`, so it + re-acquires a token from the provider rate limiter like any other + invocation. `sleep` is an injection seam for unit tests. + Returns a dict keyed by model id. Each value is either the `DriverResult` produced by `run_claude` or the `ClaudeCLIError` that aborted that model's run — callers decide how to map either @@ -300,14 +351,20 @@ def run_claude_models_parallel( if not models: raise ValueError("models must be a non-empty sequence") - def _one(model: str) -> Tuple[str, ModelResult, float]: - # Per-model wall clock: this is what the matrix run actually pays for. - # We record it whether the run succeeded or raised so the breakdown - # log below covers both code paths and surfaces "which model is the - # long pole?" without requiring per-test instrumentation. - started = time.monotonic() + retries = ( + DEFAULT_RATE_LIMIT_RETRIES + if rate_limit_retries is None + else max(0, rate_limit_retries) + ) + backoff = ( + DEFAULT_RATE_LIMIT_BACKOFF_SECONDS + if rate_limit_backoff_seconds is None + else max(0.0, rate_limit_backoff_seconds) + ) + + def _run_once(model: str) -> ModelResult: try: - result = run_claude( + return run_claude( prompt=prompt, model=model, base_url=base_url, @@ -319,14 +376,8 @@ def run_claude_models_parallel( timeout=timeout, runner=runner, ) - elapsed = time.monotonic() - started - # Stamp the duration onto the DriverResult so callers (tests, - # diagnostics) can attribute slow cells without re-timing. - result.duration_ms = int(elapsed * 1000) - return model, result, elapsed except ClaudeCLIError as exc: - elapsed = time.monotonic() - started - return model, exc, elapsed + return exc except Exception as exc: # Honor the documented "errors as values" contract for any # exception type — not just ClaudeCLIError. The rate @@ -334,13 +385,38 @@ def run_claude_models_parallel( # raise ValueError on edge-case model strings, and a future # bug elsewhere in the call stack must not abort the entire # parallel batch and lose the other models' outcomes. - elapsed = time.monotonic() - started wrapped = ClaudeCLIError( f"unexpected error running model {model!r}: " f"{type(exc).__name__}: {exc}" ) wrapped.__cause__ = exc - return model, wrapped, elapsed + return wrapped + + def _one(model: str) -> Tuple[str, ModelResult, float]: + # Per-model wall clock: this is what the matrix run actually pays + # for, retries and backoff sleeps included. We record it whether + # the run succeeded or raised so the breakdown log below covers + # both code paths and surfaces "which model is the long pole?" + # without requiring per-test instrumentation. + started = time.monotonic() + outcome = _run_once(model) + for attempt in range(retries): + if not is_rate_limit_shaped(outcome): + break + print( + f"[retry] {model}: rate-limit-shaped failure; sleeping " + f"{backoff:.0f}s before attempt {attempt + 2}/{retries + 1}", + file=sys.stderr, + flush=True, + ) + sleep(backoff) + outcome = _run_once(model) + elapsed = time.monotonic() - started + if isinstance(outcome, DriverResult): + # Stamp the duration onto the DriverResult so callers (tests, + # diagnostics) can attribute slow cells without re-timing. + outcome.duration_ms = int(elapsed * 1000) + return model, outcome, elapsed outcomes: Dict[str, ModelResult] = {} durations: Dict[str, float] = {} diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index d2bfa1a54bf..6ee8b940648 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -33,7 +33,6 @@ from __future__ import annotations import functools import json import os -import re import sys from collections import Counter, defaultdict from dataclasses import dataclass, field @@ -43,6 +42,8 @@ from typing import Any, Dict, FrozenSet, List, Optional, Tuple import pytest import yaml +from claude_code.cli_driver import RATE_LIMIT_SHAPED_RE + VALID_STATUSES = {"pass", "fail", "not_applicable", "not_tested"} RESULTS_ARTIFACT_ENV = "COMPAT_RESULTS_PATH" DEFAULT_ARTIFACT_PATH = "compat-results.json" @@ -62,11 +63,10 @@ DEFAULT_RATE_LIMIT_SUMMARY_PATH = "compat-rate-limit-summary.json" # the rate limiter is supposed to back off from. False positives on a # genuinely slow upstream are tolerable here because the worst case is # the binary search runs at a slightly lower rate than necessary. -_RATE_LIMIT_RE = re.compile( - r"(?:\b429\b|rate[\s_-]?limit|too\s+many\s+requests|throttl(?:ed|ing)|" - r"claude\s+CLI\s+timed\s+out)", - re.IGNORECASE, -) +# +# The pattern lives in `cli_driver` so the driver's retry-on-rate-limit +# logic and this summary classify failures identically. +_RATE_LIMIT_RE = RATE_LIMIT_SHAPED_RE @dataclass diff --git a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example b/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example index 11633810533..5ca7937a426 100644 --- a/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example +++ b/tests/e2e/claude_code/cron_vm/litellm-compat-matrix.env.example @@ -27,6 +27,15 @@ VERTEXAI_LOCATION=global AZURE_FOUNDRY_API_KEY= AZURE_FOUNDRY_API_BASE= +# Azure cell of the `passthrough` row. Foundry-mode Claude Code sends +# the model in the request body, so the proxy's /azure passthrough +# cannot resolve a router alias and falls back to these env vars. +# AZURE_API_BASE is the Foundry resource's Anthropic surface, i.e. +# https://.services.ai.azure.com/anthropic ; AZURE_API_KEY +# is the same key as AZURE_FOUNDRY_API_KEY. +AZURE_API_BASE= +AZURE_API_KEY= + # REQUIRED for publishing: PAT for the `agent-shin` user, used to push # the daily compat-matrix branch to its fork (agent-shin/litellm-docs) # and open the cross-repo PR against BerriAI/litellm-docs. Scopes: diff --git a/tests/e2e/claude_code/manifest.yaml b/tests/e2e/claude_code/manifest.yaml index f7cccf0cef2..e5a956991cb 100644 --- a/tests/e2e/claude_code/manifest.yaml +++ b/tests/e2e/claude_code/manifest.yaml @@ -91,6 +91,23 @@ features: # Code releases. The HTTP probe hits the bug surface LiteLLM # has actually shipped fixes for (2.1.117, 2.1.72, 2.1.70 per # the Claude Code release notes). + - id: passthrough + name: Native API passthrough + # Drives the CLI in each cloud's native mode against LiteLLM's + # passthrough routes instead of the /v1/messages translation + # layer -- the "LLM gateway" setup from + # https://code.claude.com/docs/en/gateway. anthropic uses + # ANTHROPIC_BASE_URL={proxy}/anthropic; bedrock_invoke uses + # CLAUDE_CODE_USE_BEDROCK=1 against {proxy}/bedrock (InvokeModel + # wire, alias resolved from the URL by the router); vertex_ai + # uses CLAUDE_CODE_USE_VERTEX=1 against {proxy}/vertex_ai/v1 + # (rawPredict wire, alias + project + location resolved from the + # deployment, which therefore needs `use_in_pass_through: true`); + # azure uses CLAUDE_CODE_USE_FOUNDRY=1 against {proxy}/azure and + # needs AZURE_API_BASE/AZURE_API_KEY on the proxy (see + # passthrough/test_azure.py and the cron env example). + # bedrock_converse is structurally not_applicable: Claude Code + # has no Converse-wire client. - id: long_context_1m name: Long context (1M) # Sends a ~210k-token padded prompt with the diff --git a/tests/e2e/claude_code/passthrough/__init__.py b/tests/e2e/claude_code/passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/claude_code/passthrough/test_anthropic.py b/tests/e2e/claude_code/passthrough/test_anthropic.py new file mode 100644 index 00000000000..aa0443e0625 --- /dev/null +++ b/tests/e2e/claude_code/passthrough/test_anthropic.py @@ -0,0 +1,44 @@ +"""passthrough x Anthropic. + +Drive the real `claude` CLI in its default first-party mode, but with +ANTHROPIC_BASE_URL aimed at the proxy's `/anthropic` passthrough route +instead of the `/v1/messages` translation endpoint. The proxy forwards +the request verbatim to api.anthropic.com, swapping the virtual-key +bearer for its own ANTHROPIC_API_KEY. + +The (feature, provider) for this cell is inferred from the file path by +`tests/e2e/claude_code/conftest.py`: + + tests/e2e/claude_code/passthrough/test_anthropic.py + ^^^^^^^^^^^ ^^^^^^^^^ + feature_id provider + +Because nothing is translated, the model ids are the real Anthropic API +ids (which happen to equal the proxy aliases for this column). A red +cell here means the passthrough route broke forwarding itself -- auth +header swap, streaming SSE relay, or beta-header propagation -- since +no per-provider transformation is involved. +""" + +from __future__ import annotations + +from claude_code._passthrough import ( + ANTHROPIC_PASSTHROUGH_BASE_PATH, + run_passthrough_cell, +) + +ANTHROPIC_MODELS = [ + "claude-haiku-4-5", + "claude-sonnet-4-6", + "claude-opus-4-7", +] + + +def test_passthrough_anthropic(compat_result): + """Drive the `claude` CLI through `{proxy}/anthropic` and assert a reply.""" + run_passthrough_cell( + compat_result=compat_result, + models=ANTHROPIC_MODELS, + prompt="Reply with the single word 'pong' and nothing else.", + passthrough_base_path=ANTHROPIC_PASSTHROUGH_BASE_PATH, + ) diff --git a/tests/e2e/claude_code/passthrough/test_azure.py b/tests/e2e/claude_code/passthrough/test_azure.py new file mode 100644 index 00000000000..09b0824047a --- /dev/null +++ b/tests/e2e/claude_code/passthrough/test_azure.py @@ -0,0 +1,59 @@ +"""passthrough x Azure (Microsoft Foundry). + +Drive the real `claude` CLI in foundry mode (CLAUDE_CODE_USE_FOUNDRY=1) +with ANTHROPIC_FOUNDRY_BASE_URL aimed at the proxy's `/azure` +passthrough route. The CLI POSTs `/v1/messages` with the model in the +JSON body -- unlike the bedrock/vertex modes there is no model segment +in the URL, so the proxy's router-alias resolution cannot engage and +the `/azure` route falls back to its env-configured target: the proxy +must set AZURE_API_BASE to the Foundry resource's Anthropic surface +(`https://.services.ai.azure.com/anthropic`) and +AZURE_API_KEY to the Foundry key (see +cron_vm/litellm-compat-matrix.env.example). The model ids are the +Foundry deployment names, which this matrix provisions to match the +Anthropic ids. + +The (feature, provider) for this cell is inferred from the file path by +`tests/e2e/claude_code/conftest.py`: + + tests/e2e/claude_code/passthrough/test_azure.py + ^^^^^^^^^^^ ^^^^^ + feature_id provider + +Wiring verified live at authoring time: through `{proxy}/azure` the +Foundry Anthropic surface accepted the `api-key` / `Authorization: +Bearer` headers the fallback sends (a bogus key 401s, the real key +proceeds to deployment lookup), so a red cell here means missing +AZURE_API_BASE/AZURE_API_KEY on the proxy, missing Foundry deployments +for the three tiers, or a genuine forwarding gap -- not an auth-scheme +mismatch. + +Known-red at authoring time against a healthy Foundry resource: the +`/azure` fallback assembles only its own auth headers and drops the +rest of the client's headers, including the `anthropic-version` header +the CLI sends, and Foundry's Anthropic surface rejects the request +with 400 "anthropic-version: header is required" (the same request +sent directly to Foundry with that header succeeds). This cell stays +red until that forwarding gap is fixed, which is precisely the class +of bug the row exists to surface. +""" + +from __future__ import annotations + +from claude_code._passthrough import foundry_extra_env, run_passthrough_cell + +AZURE_MODELS = [ + "claude-haiku-4-5", + "claude-sonnet-4-6", + "claude-opus-4-7", +] + + +def test_passthrough_azure(compat_result): + """Drive the `claude` CLI through `{proxy}/azure` and assert a reply.""" + run_passthrough_cell( + compat_result=compat_result, + models=AZURE_MODELS, + prompt="Reply with the single word 'pong' and nothing else.", + build_extra_env=foundry_extra_env, + ) diff --git a/tests/e2e/claude_code/passthrough/test_bedrock_converse.py b/tests/e2e/claude_code/passthrough/test_bedrock_converse.py new file mode 100644 index 00000000000..d1093a7a958 --- /dev/null +++ b/tests/e2e/claude_code/passthrough/test_bedrock_converse.py @@ -0,0 +1,34 @@ +"""passthrough x Bedrock (Converse). + +Structurally not applicable. In bedrock mode the `claude` CLI speaks +only the InvokeModel wire (`/model/{id}/invoke-with-response-stream`); +it has no Converse-wire client, so there is no Claude Code traffic a +Converse passthrough could serve. LiteLLM's `/bedrock` route does +accept `/model/{id}/converse-stream`, but exercising it would test a +wire no Claude Code user can produce, which is out of scope for this +matrix. + +The (feature, provider) for this cell is inferred from the file path by +`tests/e2e/claude_code/conftest.py`: + + tests/e2e/claude_code/passthrough/test_bedrock_converse.py + ^^^^^^^^^^^ ^^^^^^^^^^^^^^^^ + feature_id provider +""" + +from __future__ import annotations + + +def test_passthrough_bedrock_converse(compat_result): + """Report not_applicable: Claude Code has no Converse-wire mode.""" + compat_result.set( + { + "status": "not_applicable", + "reason": ( + "Claude Code's bedrock mode speaks only the InvokeModel wire " + "(/model/{id}/invoke-with-response-stream); it has no " + "Converse-wire client, so there is no Claude Code surface " + "for Converse passthrough." + ), + } + ) diff --git a/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py b/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py new file mode 100644 index 00000000000..6e84dea6779 --- /dev/null +++ b/tests/e2e/claude_code/passthrough/test_bedrock_invoke.py @@ -0,0 +1,42 @@ +"""passthrough x Bedrock (Invoke). + +Drive the real `claude` CLI in bedrock mode (CLAUDE_CODE_USE_BEDROCK=1) +with ANTHROPIC_BEDROCK_BASE_URL aimed at the proxy's `/bedrock` +passthrough route. The CLI speaks the native InvokeModel wire -- +`POST /model/{model}/invoke-with-response-stream` -- with the proxy +alias in the model segment; the proxy resolves the alias through its +router, rewrites the path to the deployment's upstream model id, and +SigV4-signs the forwarded request with its own AWS credentials +(CLAUDE_CODE_SKIP_BEDROCK_AUTH=1 keeps the CLI from signing). + +The (feature, provider) for this cell is inferred from the file path by +`tests/e2e/claude_code/conftest.py`: + + tests/e2e/claude_code/passthrough/test_bedrock_invoke.py + ^^^^^^^^^^^ ^^^^^^^^^^^^^^ + feature_id provider + +The CLI also fires a best-effort `GET /bedrock/inference-profiles` +listing at startup; its failure is non-fatal and does not gate this +cell. +""" + +from __future__ import annotations + +from claude_code._passthrough import bedrock_extra_env, run_passthrough_cell + +BEDROCK_INVOKE_MODELS = [ + "claude-haiku-4-5-bedrock-invoke", + "claude-sonnet-4-6-bedrock-invoke", + "claude-opus-4-7-bedrock-invoke", +] + + +def test_passthrough_bedrock_invoke(compat_result): + """Drive the `claude` CLI through `{proxy}/bedrock` and assert a reply.""" + run_passthrough_cell( + compat_result=compat_result, + models=BEDROCK_INVOKE_MODELS, + prompt="Reply with the single word 'pong' and nothing else.", + build_extra_env=bedrock_extra_env, + ) diff --git a/tests/e2e/claude_code/passthrough/test_vertex_ai.py b/tests/e2e/claude_code/passthrough/test_vertex_ai.py new file mode 100644 index 00000000000..5e3c6bce419 --- /dev/null +++ b/tests/e2e/claude_code/passthrough/test_vertex_ai.py @@ -0,0 +1,45 @@ +"""passthrough x Vertex AI. + +Drive the real `claude` CLI in vertex mode (CLAUDE_CODE_USE_VERTEX=1) +with ANTHROPIC_VERTEX_BASE_URL aimed at the proxy's `/vertex_ai` +passthrough route. The CLI speaks the native rawPredict wire -- +`POST .../projects/{p}/locations/{l}/publishers/anthropic/models/{model}:streamRawPredict` +-- with the proxy alias in the model segment; the proxy resolves the +alias through its router, replaces the placeholder project/location +path segments with the deployment's `vertex_project` / +`vertex_location`, and attaches its own Google credentials +(CLAUDE_CODE_SKIP_VERTEX_AUTH=1 keeps the CLI from minting a token). + +The (feature, provider) for this cell is inferred from the file path by +`tests/e2e/claude_code/conftest.py`: + + tests/e2e/claude_code/passthrough/test_vertex_ai.py + ^^^^^^^^^^^ ^^^^^^^^^ + feature_id provider + +This cell requires the vertex deployments in the proxy config to carry +`use_in_pass_through: true` (see test_config.yaml) -- that is what +registers their credentials with the passthrough router. Without it +the proxy forwards the CLI's own headers (the virtual-key bearer) to +Google and every tier fails with a 401. +""" + +from __future__ import annotations + +from claude_code._passthrough import run_passthrough_cell, vertex_extra_env + +VERTEX_MODELS = [ + "claude-haiku-4-5-vertex", + "claude-sonnet-4-6-vertex", + "claude-opus-4-7-vertex", +] + + +def test_passthrough_vertex_ai(compat_result): + """Drive the `claude` CLI through `{proxy}/vertex_ai` and assert a reply.""" + run_passthrough_cell( + compat_result=compat_result, + models=VERTEX_MODELS, + prompt="Reply with the single word 'pong' and nothing else.", + build_extra_env=vertex_extra_env, + ) diff --git a/tests/e2e/claude_code/test_config.yaml b/tests/e2e/claude_code/test_config.yaml index eec68d11dcf..e9253da2b3c 100644 --- a/tests/e2e/claude_code/test_config.yaml +++ b/tests/e2e/claude_code/test_config.yaml @@ -59,21 +59,31 @@ model_list: aws_region_name: us-east-1 # ---- Vertex AI ---- + # `use_in_pass_through: true` registers each deployment's + # project/location/credentials with the /vertex_ai passthrough + # router, which the `passthrough` row needs to resolve + # .../models/{alias}:streamRawPredict URLs. That registration only + # reads the canonical `vertex_project`/`vertex_location` param names + # (not the `vertex_ai_*` aliases); the chat translation path accepts + # both. - model_name: claude-haiku-4-5-vertex litellm_params: model: vertex_ai/claude-haiku-4-5 - vertex_ai_project: os.environ/VERTEXAI_PROJECT - vertex_ai_location: os.environ/VERTEXAI_LOCATION + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: os.environ/VERTEXAI_LOCATION + use_in_pass_through: true - model_name: claude-sonnet-4-6-vertex litellm_params: model: vertex_ai/claude-sonnet-4-6 - vertex_ai_project: os.environ/VERTEXAI_PROJECT - vertex_ai_location: os.environ/VERTEXAI_LOCATION + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: os.environ/VERTEXAI_LOCATION + use_in_pass_through: true - model_name: claude-opus-4-7-vertex litellm_params: model: vertex_ai/claude-opus-4-7 - vertex_ai_project: os.environ/VERTEXAI_PROJECT - vertex_ai_location: os.environ/VERTEXAI_LOCATION + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: os.environ/VERTEXAI_LOCATION + use_in_pass_through: true # ---- Microsoft Foundry (Anthropic deployments on Azure) ---- - model_name: claude-haiku-4-5-azure diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index c4b26387c13..afb6dbc964e 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -9,6 +9,7 @@ - {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"} - {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"} - {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"} +- {id: logging.otel.stream.exports_metric, module: logging, tier: P0, event: stream, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses], source: "integrations/otel/logger.py", rationale: "Streaming closes the LLM span from the stream path; historically prone to duplicate/orphaned spans"} - {id: logging.otel.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions, messages], source: "integrations/otel/logger.py", rationale: "Error spans for observability continuity"} - {id: logging.braintrust.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/braintrust_logging.py", rationale: "Evals platform spend"} - {id: logging.langsmith.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/langsmith.py", rationale: "LangChain ecosystem"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 3746b029331..ba192c2912e 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -17,6 +17,7 @@ - {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"} - {id: reliability.routing.usage_based.picks_under_tpm, module: reliability, tier: P0, behavior: routing, variant: usage_based, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_tpm_rpm_v2.py", rationale: "Routes to lowest-TPM deployment; prevents over-allocation"} - {id: reliability.routing.least_busy.picks_lowest_traffic, module: reliability, tier: P1, behavior: routing, variant: least_busy, assertions: [picks_lowest_traffic], exercised_on: [chat_completions, messages], source: "router_strategy/least_busy.py", rationale: "Fewest in-flight requests"} +- {id: reliability.routing.complexity_llm_classifier.routes_by_llm_tier, module: reliability, tier: P1, behavior: routing, variant: complexity_llm_classifier, assertions: [routes_by_llm_tier], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py", fail_before_fix: proven, rationale: "v2 auto-router LLM complexity classifier runs over the proxy and routes by semantic tier instead of silently crashing on absent litellm_metadata and falling back to heuristic scoring"} - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 12f47331dbb..b64f3d8dbfd 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -66,6 +66,23 @@ configs: model: openai/text-embedding-3-small api_key: os.environ/OPENAI_API_KEY + # v2 auto-router with the LLM complexity classifier. SIMPLE stays on the + # openai backend; every higher tier routes to the anthropic backend, so the + # served deployment (read back from the spend log's model) reveals whether + # the LLM classifier actually ran or silently fell back to heuristic scoring. + - model_name: complexity-smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: llm + classifier_llm_config: + model: gpt-5.5 + tiers: + SIMPLE: gpt-5.5 + MEDIUM: claude-haiku-4-5 + COMPLEX: claude-haiku-4-5 + REASONING: claude-haiku-4-5 + services: litellm: image: ghcr.io/berriai/litellm:main-latest @@ -79,6 +96,7 @@ services: env_file: .env environment: LITELLM_MASTER_KEY: sk-1234 + STORE_MODEL_IN_DB: "True" LITELLM_OTEL_V2: "true" PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces PHOENIX_API_KEY: local-jaeger-noauth diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index ad53b2b4aa2..ff296969079 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -121,6 +121,11 @@ class StreamingResponse(BaseModel): headers: dict[str, str] = {} body: str chunks: int = 0 # streamed events (0 for non-streaming) + # First in-stream error event, if any. A streamed call commits its HTTP 200 + # before the upstream completes, so upstream failures (e.g. insufficient + # quota) arrive as SSE error events inside an otherwise-successful response; + # the consumed body is elided, so this is the only place they surface. + stream_error: str | None = None @property def ok(self) -> bool: @@ -289,7 +294,19 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon body=resp.text, ) lines = cast("Iterator[bytes]", resp.iter_lines()) - chunks = sum(1 for line in lines if line) + chunks = 0 + stream_error: str | None = None + for line in lines: + if not line: + continue + chunks += 1 + if stream_error is None and ( + line.startswith(b"event: error") + or b'"type":"error"' in line + or b'"type": "error"' in line + or line.startswith(b'data: {"error"') + ): + stream_error = line.decode(errors="replace")[:300] return StreamingResponse( status_code=resp.status_code, call_id=call_id, @@ -298,6 +315,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon headers=headers, body="", chunks=chunks, + stream_error=stream_error, ) diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index bffdf71ed80..e37d6175705 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -77,11 +77,12 @@ WEATHER_TOOL = ChatTool( class ResponsesRequestBody(BaseModel): - """OpenAI Responses API /v1/responses request (non-streaming).""" + """OpenAI Responses API /v1/responses request.""" model: str input: str max_output_tokens: int + stream: bool | None = None class TeamCallbackBody(BaseModel): @@ -464,30 +465,43 @@ class LoggingClient: json=body, ) - def messages_raw(self, key: str, model: str, text: str, *, max_tokens: int = 16) -> StreamingResponse: - """Non-streaming POST /v1/messages (Anthropic-native body): raw outcome - judged by status/body/headers, for tests that need x-litellm-call-id.""" + def messages_raw( + self, key: str, model: str, text: str, *, max_tokens: int = 16, stream: bool = False + ) -> StreamingResponse: + """POST /v1/messages (Anthropic-native body): raw outcome judged by + status/body/headers, for tests that need x-litellm-call-id. With + ``stream=True`` the SSE body is consumed and its events counted.""" + body = AnthropicMessagesBody( + model=model, + max_tokens=max_tokens, + messages=[ChatMessage(role="user", content=text)], + stream=True if stream else None, + ) + if stream: + return self.gateway.transport.stream( + "/v1/messages", headers=self.gateway.transport.bearer(key), json=body + ) return self.gateway.transport.send( - "/v1/messages", - headers=self.gateway.transport.bearer(key), - json=AnthropicMessagesBody( - model=model, - max_tokens=max_tokens, - messages=[ChatMessage(role="user", content=text)], - ), + "/v1/messages", headers=self.gateway.transport.bearer(key), json=body ) def responses_raw( - self, key: str, model: str, text: str, *, max_output_tokens: int = 64 + self, key: str, model: str, text: str, *, max_output_tokens: int = 64, stream: bool = False ) -> StreamingResponse: - """Non-streaming POST /v1/responses (OpenAI Responses API): raw outcome - judged by status/body/headers, for tests that need x-litellm-call-id. + """POST /v1/responses (OpenAI Responses API): raw outcome judged by + status/body/headers, for tests that need x-litellm-call-id. max_output_tokens caps reasoning-model output cost; a capped response is - still a 200 and still exports the trace.""" + still a 200 and still exports the trace. With ``stream=True`` the SSE + body is consumed and its events counted.""" + body = ResponsesRequestBody( + model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None + ) + if stream: + return self.gateway.transport.stream( + "/v1/responses", headers=self.gateway.transport.bearer(key), json=body + ) return self.gateway.transport.send( - "/v1/responses", - headers=self.gateway.transport.bearer(key), - json=ResponsesRequestBody(model=model, input=text, max_output_tokens=max_output_tokens), + "/v1/responses", headers=self.gateway.transport.bearer(key), json=body ) def scrape_metrics(self) -> str: diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py index 90887ea5510..b00fd91be3c 100644 --- a/tests/e2e/logging/test_otel_trace_e2e.py +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -21,13 +21,14 @@ import time from collections.abc import Callable import pytest -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, ValidationError from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker from e2e_http import NoBody, StreamingResponse, require_successful_call from lifecycle import ResourceManager -from logging_client import LoggingClient -from otel_client import JaegerTrace, OtelReader +from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient +from models import LiteLLMParamsBody +from otel_client import JaegerSpan, JaegerTrace, OtelReader pytestmark = pytest.mark.e2e @@ -96,7 +97,9 @@ def _chain_reaches(span_id: str, root_id: str, trace: JaegerTrace) -> bool: return False -def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: str) -> None: +def _assert_complete_trace( + hits: list[JaegerTrace], *, route: str, genai_span: str, require_cost_span: bool = True +) -> None: """The enforced behavior: the destination holds exactly one trace for the call, rooted at the SERVER span, with auth/db/cost children and the gen-AI span all connected into that one tree - no dangling parent references.""" @@ -135,7 +138,8 @@ def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: s assert any(name.startswith(DB_SPAN_PREFIX) for name in names), ( f"no db ('{DB_SPAN_PREFIX}*') span in the trace; spans: {names}" ) - assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}" + if require_cost_span: + assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}" genai = next((span for span in trace.spans if span.operation_name == genai_span), None) assert genai is not None, f"gen-AI span {genai_span!r} missing; spans: {names}" @@ -146,8 +150,79 @@ def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: s ) -def _settled_names(*, route: str, genai_span: str) -> set[str]: - return {f"POST {route}", f"auth {route}", COST_SPAN, genai_span} +def _settled_names(*, route: str, genai_span: str, require_cost_span: bool = True) -> set[str]: + names = {f"POST {route}", f"auth {route}", genai_span} + return (names | {COST_SPAN}) if require_cost_span else names + + +def _tag(span: JaegerSpan, key: str) -> str | int | float | bool | None: + for tag in span.tags: + if tag.key == key: + return tag.value + return None + + +#: The attribute contract a failed call's gen-AI span must carry (LIT-4179), as +#: one reviewable payload. Exact-match values; error.message is additionally +#: proven untruncated by _assert_error_span_contract, which parses the provider +#: error JSON embedded in it - a truncated message stops parsing. +EXPECTED_ERROR_SPAN_ATTRIBUTES: dict[str, str] = { + "error": "True", + "error.type": "AuthenticationError", + "otel.status_code": "ERROR", + "litellm.provider.error.code": "401", + "litellm.provider.error.llm_provider": "anthropic", +} + + +class _ProviderErrorDetail(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: str + message: str + + +class _ProviderError(BaseModel): + """The provider error JSON embedded in error.message; validating it proves + the attribute survived untruncated (a cut-off message stops parsing).""" + + model_config = ConfigDict(extra="ignore") + + error: _ProviderErrorDetail + + +def _assert_error_span_contract(span: JaegerSpan) -> None: + """The failed call's gen-AI span carries the LIT-4179 error contract: the + exact attributes in EXPECTED_ERROR_SPAN_ATTRIBUTES, plus an untruncated + error.message whose embedded provider error JSON still parses and whose + text also rides the span status description.""" + for key, expected in EXPECTED_ERROR_SPAN_ATTRIBUTES.items(): + actual = _tag(span, key) + assert str(actual) == expected, f"error span attribute {key!r} must be {expected!r}, got {actual!r}" + + message = _tag(span, "error.message") + assert isinstance(message, str) and message, "error span must carry a non-empty error.message" + assert "AnthropicException" in message, ( + f"error.message must carry the upstream provider exception, got: {message[:200]}" + ) + start, end = message.find("{"), message.rfind("}") + assert start != -1 and end > start, ( + f"error.message carries no parseable provider error JSON (truncated?): {message[:200]}" + ) + try: + provider_error = _ProviderError.model_validate_json(message[start : end + 1]) + except ValidationError: + pytest.fail(f"the embedded provider error JSON does not parse (truncated?): {message[:300]}") + assert provider_error.error.message == "invalid x-api-key", ( + f"the embedded provider error must survive untruncated; parsed: {provider_error}" + ) + assert _tag(span, "otel.status_description") == message, ( + "the span status description must carry the same untruncated message as error.message" + ) + stack = _tag(span, "litellm.provider.error.stack_trace") + assert isinstance(stack, str) and stack, ( + "the error span must carry a non-empty litellm.provider.error.stack_trace" + ) class TestOtelTraceCompleteness: @@ -257,3 +332,238 @@ class TestOtelTraceCompleteness: settled_prefixes={DB_SPAN_PREFIX}, ) _assert_complete_trace(hits, route=route, genai_span=genai_span) + + @pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["chat_completions"]) + def test_chat_completions_stream_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """A successful streamed `/chat/completions` request should export one + complete OTEL trace. The trace must contain a single root `SERVER` + span, with the auth, database, cost, and gen-AI `CLIENT` spans all + connected back to that root. + + Streaming has an additional lifecycle risk because the gen-AI span is + closed by the stream-consumption path after the final chunk has + arrived and usage has been aggregated. Historically, this has caused + duplicate or orphaned spans. + + The test therefore confirms that: + + * The response actually streams. + * Exactly one gen-AI span is created for the request. + * The gen-AI span contains `litellm.request.streaming=true`. + """ + route = "/chat/completions" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-stream-chat-{unique_marker()}", models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, + lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", stream=True, max_tokens=16), + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}" + assert outcome.chunks > 0, "the stream must deliver at least one event" + assert outcome.stream_error is None, ( + f"the stream carried an upstream error event despite the 200: {outcome.stream_error}" + ) + + genai_span = f"chat {MODEL}" + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=genai_span), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=genai_span) + + genai_spans = [span for span in hits[0].spans if span.operation_name == genai_span] + assert len(genai_spans) == 1, ( + f"a streamed call must produce exactly ONE gen-AI span, got {len(genai_spans)}; " + f"spans: {hits[0].span_names()}" + ) + assert _tag(genai_spans[0], "litellm.request.streaming") is True, ( + "the gen-AI span must record litellm.request.streaming=true; its absence means " + "the stream flag was dropped before the model call" + ) + + @pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["messages"]) + def test_messages_stream_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """A successful streamed `/v1/messages` request should export one + complete OTEL trace. The trace must contain a single root `SERVER` + span, with the auth, database, cost, and gen-AI `CLIENT` spans all + connected back to that root. + + This endpoint has the same streaming lifecycle risk as + `/chat/completions`: the gen-AI span is closed by the + stream-consumption path after the final chunk has arrived and usage + has been aggregated. + + The test therefore confirms that: + + * The response actually streams. + * Exactly one gen-AI span is created for the request. + * The gen-AI span contains `litellm.request.streaming=true`. + """ + route = "/v1/messages" + _assert_otel_destination_configured(client) + + key = client.key_with_alias(f"otel-stream-messages-{unique_marker()}", models=[MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, + lambda: client.messages_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16, stream=True), + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}" + assert outcome.chunks > 0, "the stream must deliver at least one event" + assert outcome.stream_error is None, ( + f"the stream carried an upstream error event despite the 200: {outcome.stream_error}" + ) + + genai_span = f"chat {MODEL}" + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=genai_span), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=genai_span) + + genai_spans = [span for span in hits[0].spans if span.operation_name == genai_span] + assert len(genai_spans) == 1, ( + f"a streamed call must produce exactly ONE gen-AI span, got {len(genai_spans)}; " + f"spans: {hits[0].span_names()}" + ) + assert _tag(genai_spans[0], "litellm.request.streaming") is True, ( + "the gen-AI span must record litellm.request.streaming=true; its absence means " + "the stream flag was dropped before the model call" + ) + + @pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["responses"]) + def test_responses_stream_exports_complete_trace( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """A successful streamed /v1/responses request should export one complete + OTEL trace. The trace must contain a single root SERVER span, with the + auth, database, and gen-AI CLIENT spans all connected back to that + root. + + This endpoint has the same streaming lifecycle risk as the other + streaming surfaces: the gen-AI span is closed by the + stream-consumption path after the final event has arrived and usage + has been aggregated. + + The test therefore confirms that: + + * The response actually streams. + * Exactly one gen-AI span is created for the request. + * Spend is recorded correctly. + """ + route = "/v1/responses" + _assert_otel_destination_configured(client) + + key = client.key_with_alias( + f"otel-stream-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL] + ) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = _first_ok( + client, + lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}", stream=True), + ) + assert outcome.call_id is not None, "success response must carry x-litellm-call-id" + assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}" + assert outcome.chunks > 0, "the stream must deliver at least one event" + assert outcome.stream_error is None, ( + f"the stream carried an upstream error event despite the 200: {outcome.stream_error}" + ) + + genai_span = f"chat {CHEAP_OPENAI_MODEL}" + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False) + + genai_spans = [span for span in hits[0].spans if span.operation_name == genai_span] + assert len(genai_spans) == 1, ( + f"a streamed call must produce exactly ONE gen-AI span, got {len(genai_spans)}; " + f"spans: {hits[0].span_names()}" + ) + + spend_row = client.poll_proxy_spend_for_key(key) + assert spend_row is not None and spend_row.spend is not None and spend_row.spend > 0, ( + "a successful streamed responses call must record a positive-spend row in /spend/logs " + "(the cost-write SPAN is knowingly absent on this surface, LIT-4428, but the spend " + f"itself must land); got {spend_row!r}" + ) + assert spend_row.call_type == "aresponses", ( + f"the spend row must be attributed to the responses call type, got {spend_row.call_type!r}" + ) + + @pytest.mark.covers("logging.otel.failure.exports_metric", exercised_on=["chat_completions"]) + def test_failed_chat_completions_error_span_attributes( + self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager + ) -> None: + """This test checks that a failed `/chat/completions` request produces a + single, complete OTEL trace. The model-call span should include all + expected error attributes, a non-empty stack trace, and the full, + untruncated `error.message`. The provider error embedded in that + message must also remain valid JSON. The root server span should + record the same 401 response returned to the client. + + The test uses a deployment with an invalid upstream API key. This + allows the request to pass LiteLLM’s proxy authentication and fail at + the provider, which is necessary to generate a model-call error span. + There should be no cost-write span because failed requests are not + billed.""" + route = "/chat/completions" + _assert_otel_destination_configured(client) + + model_name = f"otel-err-{unique_marker()}" + model_id = client.create_model( + model_name, + LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = client.key_with_alias(f"otel-err-{unique_marker()}", models=[model_name]) + resources.defer(lambda: client.delete_key(key)) + + deadline = time.monotonic() + client.gateway.poll_timeout + while True: + outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16) + assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid" + if "AnthropicException" in outcome.body or time.monotonic() >= deadline: + break + time.sleep(client.gateway.poll_interval) + assert "AnthropicException" in outcome.body, ( + "never saw the upstream provider failure before the deadline; the key may still be " + f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" + ) + assert outcome.status_code == 401, ( + f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}" + ) + assert outcome.call_id is not None, "failed responses must still carry x-litellm-call-id" + + genai_span = f"chat {model_name}" + hits = otel_reader.poll_traces_for_call( + call_id=outcome.call_id, + settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False), + settled_prefixes={DB_SPAN_PREFIX}, + ) + _assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False) + + root = next(span for span in hits[0].spans if not span.references) + assert str(_tag(root, "http.status_code")) == "401", ( + f"the SERVER span must record the 401 the client received, got {_tag(root, 'http.status_code')!r}" + ) + genai = next(span for span in hits[0].spans if span.operation_name == genai_span) + _assert_error_span_contract(genai) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bf90426188f..4140967f3e0 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -154,6 +154,7 @@ class AnthropicMessagesBody(BaseModel): model: str messages: list[ChatMessage] max_tokens: int + stream: bool | None = None class AnthropicMessagesResponse(BaseModel): diff --git a/tests/e2e/router/complexity_router_client.py b/tests/e2e/router/complexity_router_client.py new file mode 100644 index 00000000000..929acbb3461 --- /dev/null +++ b/tests/e2e/router/complexity_router_client.py @@ -0,0 +1,20 @@ +"""Client for the complexity auto-router e2e tests. + +The suite drives the shared /chat/completions and spend-log reads on the Gateway, +so this client only carries the Gateway the shared lifecycle needs for cleanup. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from e2e_gateway import Gateway, build_gateway + + +@dataclass(frozen=True, slots=True) +class ComplexityRouterClient: + gateway: Gateway + + +def build_client() -> ComplexityRouterClient: + return ComplexityRouterClient(gateway=build_gateway()) diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py new file mode 100644 index 00000000000..e8c05520b10 --- /dev/null +++ b/tests/e2e/router/conftest.py @@ -0,0 +1,15 @@ +"""Router suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker +live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared +Gateway, so the `resources` fixture cleans up keys this suite creates. +""" + +import pytest + +from complexity_router_client import ComplexityRouterClient, build_client + + +@pytest.fixture(scope="session") +def client() -> ComplexityRouterClient: + return build_client() diff --git a/tests/e2e/router/test_complexity_router_e2e.py b/tests/e2e/router/test_complexity_router_e2e.py new file mode 100644 index 00000000000..88d79a9cac0 --- /dev/null +++ b/tests/e2e/router/test_complexity_router_e2e.py @@ -0,0 +1,62 @@ +"""Live e2e: the v2 auto-router's LLM complexity classifier actually runs over the +proxy and drives routing, instead of silently crashing and falling back to the +local heuristic scorer. + +The regression this guards (complexity_router.py `_classifier_call_metadata` +returning None when the request carries no `litellm_metadata`, which the classifier +sub-call then fed into a `.update`, raising `'NoneType' object has no attribute +'update'`) was invisible from the outside: the router caught the error and answered +from heuristic scoring, so every request still returned 200. The only tell is which +tier, and therefore which backend, served the request. + +`complexity-smart-router` (see the inline config in docker-compose.yml) pins SIMPLE +to the openai backend and every higher tier to the anthropic backend. "Is P equal +to NP?" is lexically trivial, so the heuristic scorer lands it in SIMPLE (openai), +but any competent LLM classifier reads it as a hard reasoning question and lands it +above SIMPLE (anthropic). The served deployment is read back from the spend log's +`model`, so anthropic proves the classifier ran and openai proves it silently fell +back - the exact failure before the fix. +""" + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_http import unwrap +from models import ChatBody, ChatMessage + +pytestmark = pytest.mark.e2e + +ROUTER_MODEL = "complexity-smart-router" +# Lexically simple (heuristic -> SIMPLE) but a hard reasoning question (LLM -> above SIMPLE). +LEXICALLY_SIMPLE_HARD_PROMPT = "Is P equal to NP?" +# SIMPLE tier backend; served only when the classifier silently falls back to heuristic. +HEURISTIC_TIER_MODEL = "openai/gpt-5.5" +# MEDIUM/COMPLEX/REASONING tier backend; served only when the LLM classifier runs. +LLM_TIER_MODEL = "anthropic/claude-haiku-4-5" + + +class TestComplexityRouterLlmClassifier: + @pytest.mark.covers("reliability.routing.complexity_llm_classifier.routes_by_llm_tier") + def test_llm_classifier_runs_and_routes_by_semantic_tier( + self, client: ComplexityRouterClient, scoped_key: str + ) -> None: + chat = unwrap( + client.gateway.chat( + scoped_key, + ChatBody( + model=ROUTER_MODEL, + messages=[ChatMessage(role="user", content=LEXICALLY_SIMPLE_HARD_PROMPT)], + max_tokens=16, + ), + ) + ) + assert chat.choices, f"router returned no choices: {chat}" + + rows = client.gateway.poll_logs_for_key(scoped_key, min_rows=1) + served = [row.model for row in rows] + assert served == [LLM_TIER_MODEL], ( + f"expected the request to be served by {LLM_TIER_MODEL!r} (the higher-tier " + f"backend the LLM classifier picks for a hard prompt), but the spend log shows " + f"{served!r}. {HEURISTIC_TIER_MODEL!r} means the LLM classifier silently failed " + f"and the router fell back to heuristic scoring (SIMPLE) - the pre-fix regression" + ) diff --git a/tests/local_testing/test_llm_guard.py b/tests/local_testing/test_llm_guard.py index da881e19d3c..78bbd1c0af8 100644 --- a/tests/local_testing/test_llm_guard.py +++ b/tests/local_testing/test_llm_guard.py @@ -28,8 +28,11 @@ from litellm.caching.caching import DualCache @pytest.mark.asyncio async def test_llm_guard_valid_response(): """ - Tests to see llm guard raises an error for a flagged response + A valid (is_valid=True) LLM Guard response must apply the returned + sanitized_prompt back onto the request data so the provider receives the + redacted content. """ + litellm.llm_guard_mode = "all" input_a_anonymizer_results = { "sanitized_prompt": "hello world", "is_valid": True, @@ -44,21 +47,65 @@ async def test_llm_guard_valid_response(): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) local_cache = DualCache() - try: - await llm_guard.async_moderation_hook( - data={ - "messages": [ - { - "role": "user", - "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl", - } - ] - }, - user_api_key_dict=user_api_key_dict, - call_type="completion", - ) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") + data = { + "messages": [ + { + "role": "user", + "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl", + } + ] + } + + result = await llm_guard.async_moderation_hook( + data=data, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + assert result is data + assert data["messages"][0]["content"] == "hello world" + + +@pytest.mark.asyncio +async def test_llm_guard_sanitizes_multimodal_and_input(): + """ + Sanitization must reach text parts of multimodal message content and the + ``input`` field (embeddings/moderation) while leaving non-text parts intact. + """ + litellm.llm_guard_mode = "all" + llm_guard = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={ + "sanitized_prompt": "email: [REDACTED]", + "is_valid": True, + "scanners": {"Regex": 0.0}, + }, + ) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-12345")) + + image_part = {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}} + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "email: person@example.com"}, + image_part, + ], + } + ] + } + result = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type="completion" + ) + assert result["messages"][0]["content"][0]["text"] == "email: [REDACTED]" + assert result["messages"][0]["content"][1] == image_part + + input_data = {"input": ["email: person@example.com", "another prompt"]} + input_result = await llm_guard.async_moderation_hook( + data=input_data, user_api_key_dict=user_api_key_dict, call_type="embeddings" + ) + assert input_result["input"] == ["email: [REDACTED]", "email: [REDACTED]"] @pytest.mark.asyncio diff --git a/tests/ocr_tests/test_ocr_azure_ai.py b/tests/ocr_tests/test_ocr_azure_ai.py index 172682175c5..acb44958fd9 100644 --- a/tests/ocr_tests/test_ocr_azure_ai.py +++ b/tests/ocr_tests/test_ocr_azure_ai.py @@ -23,7 +23,7 @@ class TestAzureAIOCR(BaseOCRTest): Return the base OCR call args for Azure AI. """ return { - "model": "azure_ai/mistral-document-ai-2505", + "model": "azure_ai/mistral-document-ai-2512", "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), } diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 01327529410..1136a0b7e7b 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -556,3 +556,31 @@ async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries(): assert cache_hit # token_counter over "hello world" yields a nonzero count — fallback path still runs assert response.usage.prompt_tokens > 0 + + +def test_request_kwargs_does_not_retain_logging_obj(): + """ + The caching handler lives on logging_obj._llm_caching_handler, so keeping + litellm_logging_obj inside request_kwargs closes a reference cycle + (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the + full request payload alive until a generational GC pass instead of being + freed by refcount when the request finishes; under bursts of large-token + requests this presents as stepwise RSS growth that never returns to + baseline. Other kwargs (messages included) must be preserved. + """ + logging_obj = MagicMock() + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "litellm_logging_obj": logging_obj, + } + + handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs=kwargs, + start_time=datetime.now(), + ) + + assert "litellm_logging_obj" not in handler.request_kwargs + assert handler.request_kwargs["messages"] == kwargs["messages"] + assert handler.request_kwargs["model"] == "gpt-4o" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index ade2c677745..f99cfb953ba 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -32,6 +32,13 @@ def logging_obj(): ) +def test_get_combined_callback_list_preserves_insertion_order(logging_obj): + assert logging_obj.get_combined_callback_list( + dynamic_success_callbacks=["prometheus", "langfuse", "datadog", "otel", "s3"], + global_callbacks=["langfuse", "gcs_bucket", "arize", "logfire"], + ) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"] + + def test_get_masked_api_base(logging_obj): api_base = "https://api.openai.com/v1" masked_api_base = logging_obj._get_masked_api_base(api_base) @@ -3773,3 +3780,19 @@ def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0 assert payload["total_tokens"] == 0 assert payload["completion_tokens"] == 0 + + +def test_pre_call_does_not_pin_request_in_module_state(logging_obj): + """ + pre_call/post_call must not stash their locals (full messages, the Logging + object, complete_input_dict) into module-level state. That pinned the most + recent request's entire payload in memory for the life of the worker, + which with multi-hundred-KB requests is a permanent per-worker leak. + """ + litellm.error_logs.clear() + big_input = [{"role": "user", "content": "x" * 10_000}] + + logging_obj.pre_call(input=big_input, api_key="sk-test") + logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test") + + assert litellm.error_logs == {} diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 0f7f492ddb6..e1ffabb3515 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -320,6 +320,176 @@ class TestPerformRedaction: assert choice.message.content == "redacted-by-litellm" assert choice.message.reasoning_content == "redacted-by-litellm" + def test_redacts_tool_call_arguments_in_model_response_dict(self): + """Assistant tool call arguments must not leak when redaction is on.""" + result = { + "choices": [ + { + "message": { + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + }, + } + ], + "function_call": { + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + }, + } + } + ] + } + + redacted = perform_redaction({}, result) + + message = redacted["choices"][0]["message"] + assert message["content"] == "redacted-by-litellm" + tool_call = message["tool_calls"][0] + assert tool_call["function"]["arguments"] == "redacted-by-litellm" + assert tool_call["function"]["name"] == "get_weather" + assert message["function_call"]["arguments"] == "redacted-by-litellm" + + def test_redacts_tool_call_arguments_in_streaming_delta_dict(self): + result = { + "choices": [ + { + "delta": { + "content": None, + "tool_calls": [ + { + "index": 0, + "function": { + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + }, + } + ], + } + } + ] + } + + redacted = perform_redaction({}, result) + + delta = redacted["choices"][0]["delta"] + assert delta["tool_calls"][0]["function"]["arguments"] == "redacted-by-litellm" + + def test_redacts_tool_call_arguments_on_model_response_object(self): + result = litellm.ModelResponse( + id="resp-1", + choices=[ + litellm.Choices( + message=litellm.Message( + content=None, + role="assistant", + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + }, + } + ], + ) + ) + ], + model="gpt-4o", + ) + + redacted = perform_redaction({}, result) + + tool_call = redacted.choices[0].message.tool_calls[0] + assert tool_call.function.arguments == "redacted-by-litellm" + assert tool_call.function.name == "get_weather" + assert result.choices[0].message.tool_calls[0].function.arguments == ( + '{"city": "sensitive-city"}' + ) + + def test_redacts_tool_call_arguments_on_streaming_response_object(self): + """Reproduces the Stream=True path where tool calls arrive as deltas.""" + streaming_choice = litellm.utils.StreamingChoices( + delta=litellm.utils.Delta( + content=None, + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + }, + } + ], + ) + ) + streaming_response = SimpleNamespace(choices=[streaming_choice]) + details = { + "stream": True, + "complete_streaming_response": streaming_response, + } + + perform_redaction(details, None) + + tool_call = streaming_response.choices[0].delta.tool_calls[0] + assert tool_call.function.arguments == "redacted-by-litellm" + + def test_redacts_tool_call_arguments_in_standard_logging_object(self): + details = { + "standard_logging_object": { + "response": { + "choices": [ + { + "message": { + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + }, + } + ], + } + } + ] + } + } + } + + perform_redaction(details, None) + + message = details["standard_logging_object"]["response"]["choices"][0]["message"] + assert message["tool_calls"][0]["function"]["arguments"] == "redacted-by-litellm" + + def test_redacts_responses_api_function_call_arguments_dict(self): + result = { + "output": [ + { + "type": "function_call", + "name": "get_weather", + "arguments": '{"city": "sensitive-city"}', + "call_id": "call_1", + } + ] + } + + redacted = perform_redaction({}, result) + + assert redacted["output"][0]["arguments"] == "redacted-by-litellm" + assert redacted["output"][0]["name"] == "get_weather" + def test_redacts_response_output_objects_with_top_level_text(self): output_items = [ SimpleNamespace(text="top-level output"), diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 807f1fe95f5..c5422e0d70f 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations, specifically testing edge cases with empty choic import os import sys -from typing import Any, List, Literal, Optional +from typing import Any, Literal, Optional from unittest.mock import MagicMock, patch import pytest @@ -295,6 +295,81 @@ class TestAnthropicMessagesHandlerInputProcessing: assert "input_schema" in tools[1] +class ToolAppendingGuardrail(CustomGuardrail): + """Guardrail that appends a new OpenAI-format function tool, mimicking a + guardrail that injects a retrieval/recovery tool the model can later call.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tools = list(inputs.get("tools") or []) + tools.append( + { + "type": "function", + "function": { + "name": "injected_tool", + "description": "injected by guardrail", + "parameters": {"type": "object", "properties": {}}, + }, + } + ) + inputs["tools"] = tools + return inputs + + +class TestAnthropicMessagesHandlerToolInjection: + """A tool a guardrail injects in OpenAI format must survive the write-back + to Anthropic format alongside the request's original tools.""" + + @pytest.mark.asyncio + async def test_injected_tool_survives_when_request_already_has_tools(self): + handler = AnthropicMessagesHandler() + guardrail = ToolAppendingGuardrail(guardrail_name="test") + + data = { + "model": "claude-opus-4-6", + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather at a specific location", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string"}}, + }, + } + ], + } + + result = await handler.process_input_messages( + data=data, guardrail_to_apply=guardrail, litellm_logging_obj=MagicMock() + ) + + names = [t.get("name") for t in result["tools"]] + assert "get_weather" in names + assert "injected_tool" in names + + @pytest.mark.asyncio + async def test_injected_tool_survives_when_request_has_no_tools(self): + handler = AnthropicMessagesHandler() + guardrail = ToolAppendingGuardrail(guardrail_name="test") + + data = { + "model": "claude-opus-4-6", + "messages": [{"role": "user", "content": "hi"}], + } + + result = await handler.process_input_messages( + data=data, guardrail_to_apply=guardrail, litellm_logging_obj=MagicMock() + ) + + assert [t.get("name") for t in result["tools"]] == ["injected_tool"] + + if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index f9c55db72b5..dfe7e0c3a51 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -256,7 +256,14 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_signature_block( } -def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_signature_content_block(): +def test_translate_streaming_openai_chunk_to_anthropic_content_block_thinking_and_signature(): + """The content-block classifier must treat a chunk carrying both ``thinking`` + and ``signature`` as a ``thinking`` block instead of raising. + + Such a chunk is the terminal signature event of an already-open thinking block, + so classifying it as ``thinking`` keeps the stream on the same block rather than + 500'ing. Before the fix this raised ``ValueError``. + """ choices = [ StreamingChoices( finish_reason=None, @@ -289,10 +296,14 @@ def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_ ) ] - with pytest.raises(ValueError): - LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( - choices=choices - ) + ( + block_type, + content_block_start, + ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "thinking" def test_translate_anthropic_messages_to_openai_thinking_blocks(): @@ -738,7 +749,17 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_thinking(): assert content_block_delta["signature"] == "sigsig" -def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_signature(): +def test_translate_streaming_openai_chunk_to_anthropic_emits_signature_when_thinking_and_signature(): + """A single streaming chunk carrying both ``thinking`` and ``signature`` must + translate to a ``signature_delta``, not crash. + + litellm's Anthropic streaming handler emits the ``signature_delta`` event as an + OpenAI chunk whose ``thinking_blocks`` entry re-states the full accumulated + thinking text alongside the signature (see anthropic/chat/handler.py). That text + was already streamed as ``thinking_delta`` chunks, so the signature must win and + the duplicate thinking must not be re-emitted. Before the fix this raised + ``ValueError`` and 500'd the whole stream, breaking Claude Code through the proxy. + """ choices = [ StreamingChoices( finish_reason=None, @@ -771,10 +792,25 @@ def test_translate_streaming_openai_chunk_to_anthropic_raises_when_thinking_and_ ) ] - with pytest.raises(ValueError): - LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic( - choices=choices - ) + adapter = LiteLLMAnthropicMessagesAdapter() + + ( + type_of_content, + content_block_delta, + ) = adapter._translate_streaming_openai_chunk_to_anthropic(choices=choices) + + assert type_of_content == "signature_delta" + assert content_block_delta["type"] == "signature_delta" + assert content_block_delta["signature"] == "sigsig" + + ( + block_type, + content_block_start, + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "thinking" def test_translate_anthropic_messages_to_openai_user_message_with_base64_image(): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index d67de0dcaf8..6973340101e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -12,6 +12,7 @@ content survives. import asyncio import json from types import SimpleNamespace +from typing import AsyncIterator from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, @@ -202,3 +203,89 @@ def test_split_clears_reasoning_and_thinking_on_finish_chunk(): assert content_chunk.choices[0].delta.thinking_blocks == [{"type": "thinking"}] assert finish_chunk.choices[0].delta.reasoning_content is None assert finish_chunk.choices[0].delta.thinking_blocks is None + + +def _thinking_delta_chunk(thinking: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=thinking, + thinking_blocks=[{"type": "thinking", "thinking": thinking, "signature": None}], + provider_specific_fields={ + "thinking_blocks": [{"type": "thinking", "thinking": thinking, "signature": None}] + }, + ), + finish_reason=None, + ) + ], + ) + + +def _signature_chunk(recap: str, signature: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=recap, + thinking_blocks=[{"type": "thinking", "thinking": recap, "signature": signature}], + provider_specific_fields={ + "thinking_blocks": [{"type": "thinking", "thinking": recap, "signature": signature}] + }, + ), + finish_reason=None, + ) + ], + ) + + +def test_thinking_then_signature_chunk_does_not_crash_stream(): + """Regression for the /v1/messages streaming crash reported on autoroute. + + Anthropic streams extended thinking as incremental ``thinking_delta`` chunks, then a + closing chunk that recaps the full accumulated thinking AND carries the signature. The + adapter used to raise ``ValueError`` on that closing chunk, killing the whole stream. It + must instead emit a single ``signature_delta`` for the recap chunk and never re-emit the + recap thinking, so the incremental thinking text is not duplicated. + """ + chunks = [ + _thinking_delta_chunk("First, "), + _thinking_delta_chunk("reason."), + _signature_chunk("First, reason.", "sig-abc"), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="Done"), finish_reason=None)], + ), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8), + ), + ] + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in chunks: + yield chunk + + wrapper = AnthropicStreamWrapper(completion_stream=_aiter(), model="claude-haiku-4-5") + sse = _collect_async(wrapper) + + signature_deltas = [ + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"signature_delta"' in line + ] + assert len(signature_deltas) == 1 + assert signature_deltas[0]["delta"]["signature"] == "sig-abc" + + thinking_text = "".join( + json.loads(line[len("data: ") :])["delta"]["thinking"] + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"thinking_delta"' in line + ) + assert thinking_text == "First, reason." + + assert "message_stop" in sse + assert "Done" in sse diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 6ed79763753..19ec1a04b45 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -438,6 +438,58 @@ async def test_empty_reasoning_delta_mid_thinking_block_is_suppressed_async(): _assert_empty_reasoning_delta_suppressed(await _drain_async(wrapper)) +def _full_snapshot_signature_chunks() -> List[MagicMock]: + """Mirror litellm's real Anthropic streaming: incremental ``thinking_delta`` + chunks (empty signature), then a terminal chunk whose ``thinking_blocks`` entry + re-states the *full accumulated thinking text* together with the signature + (anthropic/chat/handler.py builds the signature_delta event this way), then the + answer text. + """ + return [ + _thinking_chunk("Let me "), + _thinking_chunk("think about it."), + _thinking_chunk("Let me think about it.", signature="sig-abc"), + _make_chunk(Delta(content="42")), + _make_chunk(Delta(content=None), finish_reason="stop"), + ] + + +def _assert_full_snapshot_signature_handled(events: List[dict]) -> None: + _assert_deltas_match_their_block_type(events) + # The full-text snapshot on the signature chunk must NOT be re-emitted as an + # extra thinking_delta (it was already streamed incrementally) - otherwise the + # client renders the reasoning twice. + assert _thinking_deltas(events) == ["Let me ", "think about it."] + assert "".join(_thinking_deltas(events)) == "Let me think about it." + assert _signature_deltas(events) == ["sig-abc"] + assert _text_deltas(events) == ["42"] + + +def test_full_thinking_snapshot_with_signature_emits_signature_only_sync(): + """Regression: a terminal thinking chunk carrying both the full thinking text + and the signature used to raise ``ValueError`` (500) mid-stream, breaking every + Claude Code request routed through the proxy with an extended-thinking model. It + must instead emit a single ``signature_delta`` without duplicating the thinking. + """ + wrapper = AnthropicStreamWrapper( + completion_stream=iter(_full_snapshot_signature_chunks()), + model="claude-x", + ) + _assert_full_snapshot_signature_handled(_drain_sync(wrapper)) + + +@pytest.mark.asyncio +async def test_full_thinking_snapshot_with_signature_emits_signature_only_async(): + """Async twin - the proxy serves the async iterator, so the crash must be gone + on that path too. + """ + wrapper = AnthropicStreamWrapper( + completion_stream=_AsyncStream(_full_snapshot_signature_chunks()), + model="claude-x", + ) + _assert_full_snapshot_signature_handled(await _drain_async(wrapper)) + + def test_empty_content_chunk_mid_text_block_is_suppressed_sync(): """An empty-content chunk arriving mid-text-block (no transition) used to emit a pointless ``text_delta {"text": ""}``; it must be dropped without diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py index 5254808e315..e9d4d625421 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py @@ -32,6 +32,37 @@ def _transform(model, params, litellm_params=None): ) +def test_adaptive_thinking_only_translated_to_legacy_for_haiku_4_5(): + """The minimal autoroute repro: Claude Code sends bare ``thinking={type: adaptive}`` + (no ``output_config``) and the complexity router picks Haiku 4.5, which does not + support adaptive thinking. Anthropic 400s with "adaptive thinking is not supported on + this model" unless the flag is dropped, so it must be translated to the legacy extended + thinking the model does support rather than forwarded raw.""" + result = _transform("claude-haiku-4-5", {"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + + assert result["thinking"] == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert "output_config" not in result + + +def test_adaptive_thinking_only_dropped_for_non_reasoning_model(): + """Bare adaptive thinking on a model with no reasoning support at all is silently + dropped so the request still succeeds instead of being rejected.""" + result = _transform("claude-3-5-haiku-latest", {"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + + assert "thinking" not in result + + +def test_adaptive_thinking_only_preserved_for_4_6(): + """A 4.6+ model natively supports adaptive thinking, so a bare adaptive flag must not + be rewritten even without output_config.""" + result = _transform("claude-sonnet-4-6", {"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + + assert result["thinking"] == {"type": "adaptive"} + + def test_effort_translated_to_legacy_thinking_for_haiku_4_5(): """Core regression: Claude Code sends adaptive thinking + effort to Haiku 4.5 (thinking-capable, pre-4.6). Effort must be translated to legacy extended diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index aafaf401700..816a025e11a 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -397,6 +397,20 @@ class TestBedrockMantleResponsesRegistry: assert isinstance(cfg, BedrockMantleResponsesAPIConfig) assert cfg.use_openai_path is True + @pytest.mark.parametrize( + "model", + ["openai.gpt-5.6-sol", "openai.gpt-5.6-terra", "openai.gpt-5.6-luna"], + ) + def test_registry_returns_config_for_gpt_5_6_family(self, local_cost_map, model): + from litellm.utils import ProviderConfigManager + + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model=model, + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True + def test_registry_returns_native_config_for_gpt_oss(self, local_cost_map): # Core regression: gpt-oss-120b supports the native Responses API (AWS # model card), so it must get a BedrockMantleResponsesAPIConfig on the @@ -460,7 +474,12 @@ class TestBedrockMantleResponsesRegistry: model="xai.grok-4.3", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) - assert cfg.use_openai_path is False + # grok-4.3 is a third-party frontier model on Bedrock Mantle, served on the + # /openai/v1 base (like gpt-5.x / gemma-4), not the standard /v1 path used by + # open-weights models such as gpt-oss. The standard /v1 base returns + # "Berm is not enabled for this account", so the price-map entry carries + # use_openai_responses_path=true. + assert cfg.use_openai_path is True def test_unmapped_frontier_model_falls_through_to_none(self, restore_model_cost): # The gate is data-driven, not name-based: an unseen model not yet in the @@ -1237,6 +1256,25 @@ class TestBedrockMantleResponsesPricing: assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07) assert info["max_input_tokens"] == 272000 + @pytest.mark.parametrize( + "model, input_cost, cache_creation_cost, cache_read_cost, output_cost", + [ + ("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05), + ("openai.gpt-5.6-terra", 2.75e-06, 3.4375e-06, 2.75e-07, 1.65e-05), + ("openai.gpt-5.6-luna", 1.1e-06, 1.375e-06, 1.1e-07, 6.6e-06), + ], + ) + def test_gpt_5_6_pricing_and_mode( + self, local_cost_map, model, input_cost, cache_creation_cost, cache_read_cost, output_cost + ): + info = litellm.get_model_info(f"bedrock_mantle/{model}") + assert info["mode"] == "responses" + assert info["input_cost_per_token"] == pytest.approx(input_cost) + assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) + assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) + assert info["output_cost_per_token"] == pytest.approx(output_cost) + assert info["max_input_tokens"] == 272000 + def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 926f40a6c67..fddd8d09dfc 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1084,6 +1084,147 @@ def test_sync_delete_responses_sets_json_content_type(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "litellm_params_kwargs, stream, global_timeout, expected", + [ + ({"timeout": 12.0}, False, None, 12.0), + ({"request_timeout": 30.0}, False, None, 30.0), + ({}, False, 1500.0, 1500.0), + ({"timeout": 5.0, "stream_timeout": 50.0}, True, None, 50.0), + ({"timeout": 5.0, "stream_timeout": 50.0}, False, None, 5.0), + ({"timeout": 5.0, "request_timeout": 30.0}, False, None, 5.0), + ({}, False, None, None), + ({}, True, None, None), + ], +) +def test_resolve_anthropic_messages_timeout( + monkeypatch, litellm_params_kwargs, stream, global_timeout, expected +): + from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS + + if global_timeout is None: + monkeypatch.setattr( + "litellm.request_timeout", + float(DEFAULT_REQUEST_TIMEOUT_SECONDS), + raising=False, + ) + monkeypatch.setattr( + "litellm.request_timeout_explicitly_set", + False, + raising=False, + ) + else: + monkeypatch.setattr("litellm.request_timeout", global_timeout, raising=False) + monkeypatch.setattr( + "litellm.request_timeout_explicitly_set", True, raising=False + ) + + resolved = BaseLLMHTTPHandler._resolve_anthropic_messages_timeout( + litellm_params=GenericLiteLLMParams(**litellm_params_kwargs), + stream=stream, + custom_llm_provider="anthropic", + ) + + assert resolved == expected + + +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_forwards_request_timeout(monkeypatch): + from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "request_timeout", float(DEFAULT_REQUEST_TIMEOUT_SECONDS)) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) + handler = BaseLLMHTTPHandler() + + mock_config = Mock() + mock_config.validate_anthropic_messages_environment = Mock( + return_value=({"x-api-key": "k"}, "https://api.anthropic.com") + ) + mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False) + mock_config.transform_anthropic_messages_request = Mock( + return_value={"model": "claude", "messages": []} + ) + mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages") + mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None)) + mock_config.max_retry_on_anthropic_messages_http_error = 1 + expected_response = {"id": "msg_1", "content": []} + mock_config.transform_anthropic_messages_response = Mock(return_value=expected_response) + + ok_response = Mock() + ok_response.raise_for_status = Mock(return_value=None) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=ok_response) + + logging_obj = Mock() + logging_obj.model_call_details = {} + logging_obj.dynamic_success_callbacks = [] + + result = await handler.async_anthropic_messages_handler( + model="claude", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=mock_config, + anthropic_messages_optional_request_params={}, + custom_llm_provider="anthropic", + litellm_params=GenericLiteLLMParams(request_timeout=0.3), + logging_obj=logging_obj, + client=mock_client, + kwargs={}, + ) + + assert result is expected_response + assert mock_client.post.await_args.kwargs["timeout"] == 0.3 + + +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_forwards_stream_timeout(monkeypatch): + from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "request_timeout", float(DEFAULT_REQUEST_TIMEOUT_SECONDS)) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) + handler = BaseLLMHTTPHandler() + + mock_config = Mock() + mock_config.validate_anthropic_messages_environment = Mock( + return_value=({"x-api-key": "k"}, "https://api.anthropic.com") + ) + mock_config.should_filter_anthropic_beta_headers = Mock(return_value=False) + mock_config.transform_anthropic_messages_request = Mock( + return_value={"model": "claude", "messages": []} + ) + mock_config.get_complete_url = Mock(return_value="https://api.anthropic.com/v1/messages") + mock_config.sign_request = Mock(return_value=({"x-api-key": "k"}, None)) + mock_config.max_retry_on_anthropic_messages_http_error = 1 + mock_config.get_async_streaming_response_iterator = Mock(return_value=Mock()) + + ok_response = Mock() + ok_response.raise_for_status = Mock(return_value=None) + ok_response.headers = httpx.Headers({}) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=ok_response) + + logging_obj = Mock() + logging_obj.model_call_details = {} + logging_obj.dynamic_success_callbacks = [] + + await handler.async_anthropic_messages_handler( + model="claude", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=mock_config, + anthropic_messages_optional_request_params={}, + custom_llm_provider="anthropic", + litellm_params=GenericLiteLLMParams(timeout=9.0, stream_timeout=0.7), + logging_obj=logging_obj, + client=mock_client, + stream=True, + kwargs={}, + ) + + assert mock_client.post.await_args.kwargs["stream"] is True + assert mock_client.post.await_args.kwargs["timeout"] == 0.7 + + @pytest.mark.asyncio async def test_anthropic_post_uses_prebuilt_body_without_redumping(): """When the caller passes a pre-serialized (unsigned) body, attempt 0 must @@ -1894,7 +2035,9 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques ok_response = httpx.Response(200, json={"id": "msg_1"}, request=httpx.Request("POST", request_url)) class FakeAsyncClient: - async def post(self, url, headers, data, stream=False, logging_obj=None): + async def post( + self, url, headers, data, stream=False, logging_obj=None, timeout=None + ): posts.append({"headers": dict(headers), "data": data}) return invalid_signature_response if len(posts) == 1 else ok_response diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 49cd1b71ef2..4c45eaac7b9 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1168,3 +1168,69 @@ class TestGetStructuredMessages: data = {"input": None} result = handler.get_structured_messages(data) assert result is None + + +class ToolAppendingGuardrail(CustomGuardrail): + """Guardrail that appends a new function tool, mimicking a guardrail that + injects a retrieval/recovery tool the model can later call.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tools = list(inputs.get("tools") or []) + tools.append( + { + "type": "function", + "function": { + "name": "injected_tool", + "description": "injected by guardrail", + "parameters": {"type": "object", "properties": {}}, + }, + } + ) + inputs["tools"] = tools + return inputs + + +class TestOpenAIResponsesHandlerToolInjection: + """A tool a guardrail injects must survive the write-back to Responses format.""" + + def test_merge_keeps_guardrail_appended_tool(self): + """_merge_tools_after_guardrail must not drop the extra appended tool.""" + handler = OpenAIResponsesHandler() + original = [{"type": "function", "name": "a"}] + remapped = [ + {"type": "function", "name": "a"}, + {"type": "function", "name": "b"}, + ] + merged = handler._merge_tools_after_guardrail(original, remapped) + assert [t["name"] for t in merged] == ["a", "b"] + + @pytest.mark.asyncio + async def test_injected_tool_survives_when_request_already_has_tools(self): + """Regression: the merge dropped the injected tool whenever the request + already carried tools, so the model never saw it.""" + handler = OpenAIResponsesHandler() + guardrail = ToolAppendingGuardrail(guardrail_name="test") + + data = { + "input": [{"role": "user", "content": "hi", "type": "message"}], + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + } + ], + "model": "gpt-4", + } + + result = await handler.process_input_messages(data, guardrail) + + names = [t.get("name") for t in result["tools"]] + assert "get_weather" in names + assert "injected_tool" in names diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index adcfff6fe9d..bba00ed1819 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -356,6 +356,86 @@ class TestMCPServerManager: assert server.oauth2_flow == "authorization_code" assert server.needs_user_oauth_token is True + @pytest.mark.asyncio + async def test_load_servers_from_config_rejects_uncorroborated_endpoints_but_keeps_resource_scopes(self): + """A yaml server with a manual authorization_url has the same config-time mix-up exposure as a + DB row: a document advertising a different authorize endpoint has its token_url rejected. The + resource-driven scopes are kept, because scope selection is resource-driven (MCP Scope + Selection Strategy) and scope inflation is bounded by the authorization server at consent, not + by dropping scopes when an endpoint mismatches.""" + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://attacker.example.com/authorize", + token_url="https://attacker.example.com/token", + scopes=["read", "admin"], + ) + config = self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize", + token_url=None, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.authorization_url == "https://idp.example.com/authorize" + assert server.token_url is None + assert server.scopes == ["read", "admin"] + + @pytest.mark.asyncio + async def test_load_servers_from_config_fills_token_url_when_metadata_corroborates_manual_authorization_url(self): + """Corroborated metadata keeps the self-heal on the config path: when the discovered document + advertises the same authorize endpoint the admin pinned, its token_url fills the blank field + and scopes come through resource-driven (the discovered document's resource-preferred scopes), + not the authorization server's own capability list.""" + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read", "admin"], + ) + config = self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize/", + token_url=None, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.token_url == "https://idp.example.com/token" + assert server.scopes == ["read", "admin"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("blank_authorization_url", ["", " "]) + async def test_load_servers_from_config_blank_authorization_url_is_not_a_pin(self, blank_authorization_url): + """A blank authorization_url — empty or whitespace-only — is not a trust anchor, so discovery + backfills the whole set (authorize endpoint, token_url, and its resource-preferred scopes) + from the same chain, exactly as if the field had been omitted. The merge and the corroboration + gate must agree that blank means unpinned; a whitespace value that the merge kept for redirects + while the gate treated as unpinned would strand a broken half-discovered config.""" + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + ) + config = self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url=blank_authorization_url, + token_url=None, + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.authorization_url == "https://idp.example.com/authorize" + assert server.token_url == "https://idp.example.com/token" + assert server.scopes == ["read"] + @pytest.mark.asyncio async def test_load_servers_from_config_non_oauth2_needs_no_flow(self): manager = MCPServerManager() @@ -1026,7 +1106,6 @@ class TestMCPServerManager: """The gateway's relayed authorize flow (used by the browser-only Authorize) needs the upstream's authorization_url on the registry entry, and these rows never persist one, so the DB build must discover it the same way oauth2 rows do.""" - from types import SimpleNamespace manager = MCPServerManager() row = LiteLLM_MCPServerTable( @@ -1053,6 +1132,181 @@ class TestMCPServerManager: assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url == "https://idp.example.com/token" + @pytest.mark.asyncio + async def test_build_from_table_backfills_resource_driven_scopes_for_pinned_authorization_url(self): + """When authorization_url is admin-pinned and corroborated, scopes backfill as the + resource-driven value (the WWW-Authenticate challenge scope, else the RFC 9728 + protected-resource scopes_supported), per the MCP authorization spec Scope Selection Strategy. + The client does not restrict scopes to the authorization server's own scopes_supported; scope + minimization and inflation control are the authorization server's and user's job at consent + (RFC 6749 §3.3).""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="manual-auth-url-1", + alias="manual_auth_url", + description="manual authorization_url, blank scopes", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read", "admin"], + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)) as mock_discovery: + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + mock_discovery.assert_awaited_once() + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.scopes == ["read", "admin"] + + @pytest.mark.asyncio + async def test_build_from_table_whitespace_authorization_url_is_not_a_pin(self): + """A whitespace-only authorization_url on the row must not be kept for redirects while the + gate treats it as unpinned. It is normalized to unpinned everywhere, so the built server + takes the discovered authorize endpoint, token_url, and scopes as one consistent group + rather than serving the whitespace value with half-discovered fields.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="whitespace-auth-url", + alias="whitespace_auth_url", + description="whitespace authorization_url is not a pin", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url=" ", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["read"], + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.scopes == ["read"] + + @pytest.mark.asyncio + async def test_build_from_table_fills_endpoints_when_metadata_corroborates_manual_authorization_url(self): + """A discovered token_url is only trusted next to a manual authorization_url when the same + metadata document advertises that authorize endpoint, and the comparison must tolerate + formatting-only differences (host case, trailing slash, query params like ?prompt=consent) + so hand-copied URLs still self-heal.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="manual-auth-url-2", + alias="manual_auth_url_match", + description="manual authorization_url matching discovery, blank token_url", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://IDP.example.com/authorize/?prompt=consent", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["read"], + ) + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.authorization_url == "https://IDP.example.com/authorize/?prompt=consent" + assert built.token_url == "https://idp.example.com/token" + assert built.registration_url == "https://idp.example.com/register" + assert built.scopes == ["read"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "advertised_authorization_url", + ["https://attacker.example.com/authorize", None], + ) + async def test_build_from_table_rejects_uncorroborated_endpoints_but_keeps_resource_scopes( + self, advertised_authorization_url + ): + """Resource-rooted discovery lets a compromised upstream advertise its own authorization + server. With a manual authorization_url pinned, a document that does not corroborate it has + its token_url and registration_url dropped: accepting them would send the code, client secret, + and PKCE verifier to the attacker (config-time RFC 9700 mix-up). The resource-driven scopes + are kept, because scope selection is resource-driven (MCP Scope Selection Strategy) and scope + inflation is bounded by the authorization server at consent (RFC 6749 §3.3), not by dropping + scopes on an endpoint mismatch. Both the in-memory merge and the persisted metadata drop only + the uncorroborated endpoints.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="manual-auth-url-3", + alias="manual_auth_url_mismatch", + description="manual authorization_url, hostile discovery document", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + metadata = MCPOAuthMetadata( + authorization_url=advertised_authorization_url, + token_url="https://attacker.example.com/token", + registration_url="https://attacker.example.com/register", + scopes=["read", "admin"], + ) + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)), + patch.object(manager, "_persist_discovered_oauth_endpoints", new=AsyncMock()) as mock_persist, + ): + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url is None + assert built.registration_url is None + assert built.scopes == ["read", "admin"] + persisted_metadata = mock_persist.await_args.kwargs["metadata"] + assert persisted_metadata.token_url is None + assert persisted_metadata.registration_url is None + assert persisted_metadata.scopes == ["read", "admin"] + + @pytest.mark.asyncio + async def test_build_from_table_skips_discovery_when_all_upstream_oauth_fields_present(self): + """A fully hand-configured server (authorization_url, token_url, and scopes all set) has + nothing left for discovery to fill, so the build must not fetch upstream metadata.""" + manager = MCPServerManager() + row = LiteLLM_MCPServerTable( + server_id="fully-manual-1", + alias="fully_manual", + description="all upstream oauth fields set by the admin", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/manual-authorize", + token_url="https://idp.example.com/manual-token", + credentials={"scopes": ["calendar.read"]}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)) as mock_discovery: + built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + mock_discovery.assert_not_awaited() + assert built.authorization_url == "https://idp.example.com/manual-authorize" + assert built.token_url == "https://idp.example.com/manual-token" + assert built.scopes == ["calendar.read"] + async def _capture_subject_token(self, call) -> Optional[str]: """Run a manager method (via ``call(manager)``) and return the subject_token it threaded into ``_create_mcp_client``.""" @@ -2127,6 +2381,46 @@ class TestMCPServerManager: assert result.scopes == ["api://some-scope/.default"] assert result.from_origin_fallback is False + @pytest.mark.asyncio + async def test_descovery_metadata_scopes_are_resource_driven(self): + """The effective `scopes` are resource-driven: the RFC 9728 protected-resource advertisement + (or WWW-Authenticate challenge) overrides the authorization server's own scopes_supported. This + is the MCP Scope Selection Strategy: the client requests what the resource needs, not the AS's + full capability list.""" + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + authorization_server_metadata = MCPOAuthMetadata( + scopes=["as.read", "as.write"], + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ), + patch.object( + manager, + "_attempt_well_known_discovery", + AsyncMock(return_value=(["https://idp.example.com"], ["resource.only"])), + ), + patch.object( + manager, + "_fetch_authorization_server_metadata", + AsyncMock(return_value=authorization_server_metadata), + ), + ): + result = await manager._descovery_metadata("https://up.example.com/mcp") + + assert result is not None + assert result.scopes == ["resource.only"] + @pytest.mark.asyncio async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path( self, @@ -2252,6 +2546,10 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_load_servers_from_config_overrides_discovery_metadata(self): + """Config values win per field. The discovered token_url/registration_url do NOT fill the + blanks here: the document advertises a different authorization_endpoint than the manually + configured one, so combining its endpoints with the pinned authorize URL would be the + config-time mix-up the discovery gate exists to prevent.""" manager = MCPServerManager() discovered_metadata = MCPOAuthMetadata( @@ -2285,8 +2583,8 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.scopes == ["config"] # config overrides discovery assert server.authorization_url == "https://config.example.com/auth" - assert server.token_url == "https://discovered.example.com/token" - assert server.registration_url == "https://discovered.example.com/register" + assert server.token_url is None + assert server.registration_url is None @pytest.mark.asyncio async def test_load_servers_from_config_filters_blank_scopes(self): @@ -5093,6 +5391,102 @@ class TestMCPServerTimestamps: _carry_forward_resolved_oauth_endpoints(new_server=explicit, previous_server=previous) assert explicit.authorization_url == "https://configured.example.com/auth" + def test_carry_forward_does_not_revive_token_url_across_authorization_url_change(self): + """Carry-forward is a non-manual endpoint source, so it obeys the same trust rule as + discovery: a previous token_url/registration_url belongs to the previous authorization + server, so it must not be pinned to a NEW authorization_url the admin re-pointed to. Without + this, re-pointing authorize to server B while the same MCP url keeps serving A's token + endpoint recreates the RFC 9700 mix-up, durably, and the discovery gate alone cannot catch + it because the stale endpoint comes from the registry, not from discovery.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _carry_forward_resolved_oauth_endpoints, + ) + + previous = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp-a.example.com/authorize", + token_url="https://idp-a.example.com/token", + registration_url="https://idp-a.example.com/register", + ) + repointed = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp-b.example.com/authorize", + ) + + _carry_forward_resolved_oauth_endpoints(new_server=repointed, previous_server=previous) + + assert repointed.authorization_url == "https://idp-b.example.com/authorize" + assert repointed.token_url is None + assert repointed.registration_url is None + + def test_carry_forward_restores_endpoints_when_authorization_url_unchanged(self): + """The last-known-good path still works: a rebuild whose discovery blipped (no authorize + endpoint) adopts the previous authorize endpoint AND its token endpoint together as a + consistent group, and a rebuild that re-pins the same authorize endpoint (formatting aside) + keeps carrying the corroborated token endpoint.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _carry_forward_resolved_oauth_endpoints, + ) + + def previous() -> MCPServer: + return MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + ) + + blipped = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url=None, + ) + _carry_forward_resolved_oauth_endpoints(new_server=blipped, previous_server=previous()) + assert blipped.authorization_url == "https://idp.example.com/authorize" + assert blipped.token_url == "https://idp.example.com/token" + assert blipped.registration_url == "https://idp.example.com/register" + + same_authorize = MCPServer( + server_id="s1", + name="s1", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://IDP.example.com:443/authorize/", + ) + _carry_forward_resolved_oauth_endpoints(new_server=same_authorize, previous_server=previous()) + assert same_authorize.token_url == "https://idp.example.com/token" + assert same_authorize.registration_url == "https://idp.example.com/register" + + def test_normalized_authorize_endpoint_treats_default_port_and_slash_as_identity(self): + """The corroboration check must not fail on formatting-only differences an IdP legitimately + emits: default port, trailing slash, host case, and query string are not identity, but a + non-default port is.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _normalized_authorize_endpoint, + ) + + canonical = _normalized_authorize_endpoint("https://idp.example.com/authorize") + assert _normalized_authorize_endpoint("https://idp.example.com:443/authorize") == canonical + assert _normalized_authorize_endpoint("https://IDP.example.com/authorize/") == canonical + assert _normalized_authorize_endpoint("https://idp.example.com/authorize?prompt=consent") == canonical + assert _normalized_authorize_endpoint("https://idp.example.com:8443/authorize") != canonical + def test_build_mcp_server_table_preserves_timestamps(self): """_build_mcp_server_table must use the MCPServer's stored timestamps, not datetime.now().""" manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py index 7d2cf4442a5..47112fc9900 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py @@ -227,3 +227,66 @@ async def test_client_credentials_uses_client_secret_basic_when_configured(): assert "client_secret" not in kwargs["data"] assert "client_id" not in kwargs["data"] assert kwargs["data"]["grant_type"] == "client_credentials" + + +def test_storage_ttl_capped_at_token_lifetime(): + """A token_storage_ttl_seconds longer than the token's own lifetime must be capped at + expires_in minus the expiry buffer. Before the cap, the configured TTL won outright and the + Redis fast path (which never re-checks expires_at) kept serving the dead token until eviction, + while the stored refresh_token sat unused because refresh only runs on the DB read-through.""" + from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + _compute_per_user_token_ttl, + ) + + server = _server(oauth2_flow=None, token_storage_ttl_seconds=604800) + assert _compute_per_user_token_ttl(server, expires_in=86400) == 86400 - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS + + +def test_storage_ttl_shorter_than_token_lifetime_wins(): + """A configured TTL below the token lifetime is the operative value: the knob's purpose is to + force earlier DB re-checks (staleness backstop), so the shorter side must win the min().""" + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + _compute_per_user_token_ttl, + ) + + server = _server(oauth2_flow=None, token_storage_ttl_seconds=3600) + assert _compute_per_user_token_ttl(server, expires_in=86400) == 3600 + + +def test_storage_ttl_verbatim_when_token_lifetime_unknown(): + """With no expires_in from the upstream there is nothing to cap against, so the configured + TTL applies as-is (matching the pre-cap behavior for lifetime-less tokens).""" + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + _compute_per_user_token_ttl, + ) + + server = _server(oauth2_flow=None, token_storage_ttl_seconds=604800) + assert _compute_per_user_token_ttl(server, expires_in=None) == 604800 + + +def test_storage_ttl_floors_at_one_second_for_nearly_dead_token(): + """A token already inside the expiry buffer yields the 1-second floor, not zero or a negative + TTL, mirroring the floor the default (unconfigured) path has always had.""" + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + _compute_per_user_token_ttl, + ) + + server = _server(oauth2_flow=None, token_storage_ttl_seconds=3600) + assert _compute_per_user_token_ttl(server, expires_in=30) == 1 + + +def test_default_ttl_paths_unchanged_without_storage_ttl(): + """With token_storage_ttl_seconds unset the TTL still derives from expires_in minus the + buffer, and falls back to MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.""" + from litellm.constants import ( + MCP_PER_USER_TOKEN_DEFAULT_TTL, + MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + ) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + _compute_per_user_token_ttl, + ) + + server = _server(oauth2_flow=None) + assert _compute_per_user_token_ttl(server, expires_in=86400) == 86400 - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS + assert _compute_per_user_token_ttl(server, expires_in=None) == MCP_PER_USER_TOKEN_DEFAULT_TTL diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 042fc107f40..17ff700791f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -659,7 +659,16 @@ def test_get_model_from_request_ignores_session_model_on_non_realtime_routes(): def test_abbreviate_api_key(): - assert abbreviate_api_key("sk-test-1234") == "sk-...1234" + assert abbreviate_api_key("sk-test-1234-abcdefgh") == "sk-...efgh" + assert abbreviate_api_key("sk-abcdefghijklm") == "sk-...jklm" + + +def test_abbreviate_api_key_short_key_is_fully_masked(): + """Regression test for LIT-4355: for keys shorter than the enforced minimum, + showing the last 4 characters can reveal the entire key (sk-1234 -> sk-...1234).""" + assert abbreviate_api_key("sk-1234") == "sk-..." + assert abbreviate_api_key("sk-test-1234") == "sk-..." + assert abbreviate_api_key("") == "sk-..." def test_get_customer_user_header_returns_none_when_no_customer_role(): diff --git a/tests/test_litellm/proxy/auth/test_cli_auth.py b/tests/test_litellm/proxy/auth/test_cli_auth.py index a4f72ef90ef..c9b31a1d776 100644 --- a/tests/test_litellm/proxy/auth/test_cli_auth.py +++ b/tests/test_litellm/proxy/auth/test_cli_auth.py @@ -60,13 +60,13 @@ async def test_normalize_teams_with_details_with_aliases(): @patch("litellm.proxy.client.cli.commands.auth.requests.post") def test_start_cli_sso_flow_rejects_invalid_response(request_mock): - """Test CLI SSO start rejects malformed server responses""" + """Test CLI SSO start rejects malformed server responses and names the missing fields""" response = Mock() - response.raise_for_status = Mock() + response.status_code = 200 response.json.return_value = {"login_id": "cli-session", "user_code": "ABCD-EFGH"} request_mock.return_value = response - with pytest.raises(ValueError, match="Invalid CLI SSO start response"): + with pytest.raises(ValueError, match="missing required field\\(s\\): poll_secret"): _start_cli_sso_flow("https://litellm.com") @@ -75,15 +75,13 @@ def test_start_cli_sso_flow_rejects_invalid_response(request_mock): "litellm.proxy.client.cli.commands.auth.requests.get", side_effect=[Mock(status_code=404)], ) -@patch("litellm.proxy.client.cli.commands.auth.click.echo") @patch("litellm.proxy.client.cli.commands.auth.time.sleep") -async def test_poll_for_ready_404(sleep_mock, click_mock, request_mock): - """Test poll_for_ready function""" - actual = _poll_for_ready_data( - "https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42 - ) - assert actual is None - click_mock.assert_called_once_with("Polling error: HTTP 404") +async def test_poll_for_ready_404(sleep_mock, request_mock): + """Test polling treats HTTP 404 as a permanent error and raises instead of retrying""" + with pytest.raises(ValueError, match="rejected the login session with HTTP 404"): + _poll_for_ready_data( + "https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42 + ) request_mock.assert_called_once_with("https://litellm.com", timeout=42) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/__init__.py b/tests/test_litellm/proxy/client/cli/autoroute/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py new file mode 100644 index 00000000000..9efde03e04c --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -0,0 +1,319 @@ +import json +import stat +from typing import Optional + +import yaml +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands.autoroute import commands as commands_module +from litellm.proxy.client.cli.commands.autoroute import process as process_module +from litellm.proxy.client.cli.commands.autoroute.commands import down, up +from litellm.proxy.client.cli.commands.autoroute.process import PidRecord, ProcessLaunchError, write_pid_record +from litellm.proxy.client.cli.commands.up import BackupRecord as ClaudeBackupRecord +from litellm.proxy.client.cli.commands.up import write_backup + + +class FakeProcess: + def __init__(self, pid: int): + self.pid = pid + self.returncode: Optional[int] = None + + def poll(self) -> Optional[int]: + return self.returncode + + +def _patch_paths(monkeypatch, tmp_path): + config_path = tmp_path / "config.yaml" + log_path = tmp_path / "proxy.log" + claude_settings_path = tmp_path / "claude_settings.json" + backup_path = tmp_path / "backup.json" + pid_record_path = tmp_path / "pid.json" + + monkeypatch.setattr(commands_module, "CONFIG_PATH", config_path) + monkeypatch.setattr(commands_module, "LOG_PATH", log_path) + monkeypatch.setattr(commands_module, "CLAUDE_SETTINGS_PATH", claude_settings_path) + monkeypatch.setattr(commands_module, "AUTOROUTE_BACKUP_PATH", backup_path) + monkeypatch.setattr(process_module, "PID_RECORD_PATH", pid_record_path) + + return config_path, log_path, claude_settings_path, backup_path, pid_record_path + + +def _silence_signal_handling(monkeypatch): + monkeypatch.setattr(commands_module.signal, "signal", lambda *a, **k: None) + monkeypatch.setattr(commands_module.atexit, "register", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "stream_log", lambda *a, **k: None) + + +class TestUpCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_refuses_when_never_configured(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "lite autoroute configure" in result.output + + def test_surfaces_clean_error_on_empty_config_file(self, monkeypatch, tmp_path): + """A `configure` killed between secure_create's O_TRUNC and the write completing leaves an + empty config.yaml on disk -- yaml.safe_load(empty) returns None, and validating None as the + generated-config model raises a raw pydantic.ValidationError if uncaught.""" + config_path, _log_path, _settings_path, _backup_path, _pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text("") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert result.exception is None or isinstance(result.exception, SystemExit) + assert "lite autoroute configure" in result.output + + def test_refuses_with_actionable_error_when_proxy_runtime_missing(self, monkeypatch, tmp_path): + """`up` launches a real litellm proxy, which the thin `litellm[cli]` install cannot run. + It must fail fast with an actionable message pointing at the proxy install, before it ever + tries to launch the doomed subprocess (which would otherwise die with a bare ImportError).""" + config_path, _log_path, _settings_path, _backup_path, _pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + monkeypatch.setattr(commands_module, "missing_proxy_runtime_modules", lambda: ("fastapi", "websockets")) + + def _fail_if_launched(*args, **kwargs): + raise AssertionError("launch_proxy must not run when the proxy runtime is missing") + + monkeypatch.setattr(commands_module, "launch_proxy", _fail_if_launched) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "fastapi, websockets" in result.output + assert "litellm[proxy]" in result.output + + def test_refuses_when_pid_record_exists_and_process_still_running(self, monkeypatch, tmp_path): + config_path, _log_path, _settings_path, _backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + write_pid_record( + PidRecord(pid=123, port=4000, config_path=str(config_path), log_path="/tmp/proxy.log"), pid_record_path + ) + monkeypatch.setattr(commands_module, "is_running", lambda pid: True) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "already running" in result.output + assert "lite autoroute down" in result.output + assert config_path.read_text() == yaml.safe_dump({"model_list": []}) + + def test_refuses_when_backup_exists_after_an_unclean_crash(self, monkeypatch, tmp_path): + """A prior `up` that was SIGKILL'd leaves no live pid but does leave a stale backup file. + + Without this guard, a fresh `up` would overwrite that backup with the currently-patched + (not original) Claude settings, so `down`/Ctrl-C would restore the wrong content forever. + """ + config_path, _log_path, claude_settings_path, backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "stale-patched-token"}})) + write_backup(ClaudeBackupRecord(existed=True, content={"theme": "dark"}), backup_path) + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "already exists" in result.output + assert "lite autoroute down" in result.output + assert json.loads(backup_path.read_text())["content"] == {"theme": "dark"} + + def test_happy_path_patches_settings_then_restores_everything_on_stop(self, monkeypatch, tmp_path): + config_path, log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + original_settings = {"theme": "dark"} + claude_settings_path.write_text(json.dumps(original_settings)) + _silence_signal_handling(monkeypatch) + + fake_process = FakeProcess(pid=99999) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 54321) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + captured = {} + + def fake_wait(self, timeout=None): + captured["settings"] = json.loads(claude_settings_path.read_text()) + captured["backup_existed"] = backup_path.exists() + captured["settings_mode"] = stat.S_IMODE(claude_settings_path.stat().st_mode) + return True + + monkeypatch.setattr("threading.Event.wait", fake_wait) + + result = self.runner.invoke(up) + + assert result.exit_code == 0, result.output + assert captured["backup_existed"] is True + assert captured["settings"]["theme"] == "dark" + assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:54321" + assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" + assert "apiKeyHelper" not in captured["settings"] + assert captured["settings_mode"] == 0o600 + + assert terminate_calls == [99999] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + written_config = yaml.safe_load(config_path.read_text()) + assert written_config["general_settings"]["master_key"] == "fixed-master-key" + assert stat.S_IMODE(config_path.stat().st_mode) == 0o600 + + def test_teardown_reports_clean_error_when_backup_is_corrupt(self, monkeypatch, tmp_path): + """A corrupt backup at teardown time (e.g. a concurrent process wrote garbage to it) must + not crash the whole command -- _restore_once in up.py handles the identical case in + lite up the same way, echoing the error instead of propagating it.""" + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text(json.dumps({"theme": "dark"})) + _silence_signal_handling(monkeypatch) + + fake_process = FakeProcess(pid=11111) + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 65432) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: None) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + def fake_wait(self, timeout=None): + backup_path.write_text("not json at all {{{") + return True + + monkeypatch.setattr("threading.Event.wait", fake_wait) + + result = self.runner.invoke(up) + + assert result.exit_code == 0, result.output + assert "invalid or unexpected JSON" in result.output + assert not pid_record_path.exists() + + def test_surfaces_clean_error_and_cleans_up_when_health_check_fails(self, monkeypatch, tmp_path): + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + original_settings = {"theme": "dark"} + claude_settings_path.write_text(json.dumps(original_settings)) + + fake_process = FakeProcess(pid=555) + terminate_calls = [] + + def _raise_launch_error(*args, **kwargs): + raise ProcessLaunchError("boom: proxy never became healthy") + + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", _raise_launch_error) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 12345) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "boom" in result.output + assert terminate_calls == [555] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_terminates_ephemeral_proxy_when_claude_settings_is_corrupt(self, monkeypatch, tmp_path): + """The health check can pass and the proxy can come up fine, but if + ~/.claude/settings.json turns out to be corrupt, the just-started proxy must not be left + running with no pid record -- exactly the leak `lite autoroute down` exists to clean up.""" + config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths(monkeypatch, tmp_path) + config_path.write_text(yaml.safe_dump({"model_list": []})) + claude_settings_path.write_text("not json at all {{{") + + fake_process = FakeProcess(pid=777) + terminate_calls = [] + monkeypatch.setattr(commands_module, "launch_proxy", lambda *a, **k: fake_process) + monkeypatch.setattr(commands_module, "poll_liveliness", lambda *a, **k: None) + monkeypatch.setattr(commands_module, "allocate_free_port", lambda: 23456) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + monkeypatch.setattr(commands_module.secrets, "token_urlsafe", lambda n: "fixed-master-key") + + result = self.runner.invoke(up) + + assert result.exit_code != 0 + assert "invalid JSON" in result.output + assert terminate_calls == [777] + assert not pid_record_path.exists() + assert not backup_path.exists() + + +class TestDownCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_restores_settings_and_terminates_when_process_still_running(self, monkeypatch, tmp_path): + _config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + original_settings = {"theme": "dark"} + write_backup(ClaudeBackupRecord(existed=True, content=original_settings), backup_path) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "fixed-master-key"}})) + write_pid_record(PidRecord(pid=777, port=1234, config_path="c", log_path="l"), pid_record_path) + + terminate_calls = [] + monkeypatch.setattr(commands_module, "is_running", lambda pid: True) + monkeypatch.setattr(commands_module, "terminate", lambda pid, **k: terminate_calls.append(pid)) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Stopped leftover ephemeral proxy" in result.output + assert "Restored" in result.output + assert terminate_calls == [777] + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_is_a_clean_no_op_when_nothing_is_running_and_no_backup_exists(self, monkeypatch, tmp_path): + _config_path, _log_path, claude_settings_path, _backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Nothing to restore." in result.output + assert not claude_settings_path.exists() + + def test_clears_a_corrupt_pid_record_and_still_restores_settings(self, monkeypatch, tmp_path): + """down is specifically the crash-recovery path -- a pid file truncated by a mid-write + crash must not block it from clearing the record and restoring Claude settings anyway.""" + _config_path, _log_path, claude_settings_path, backup_path, pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + pid_record_path.parent.mkdir(parents=True, exist_ok=True) + pid_record_path.write_text("not json at all {{{") + original_settings = {"theme": "dark"} + write_backup(ClaudeBackupRecord(existed=True, content=original_settings), backup_path) + claude_settings_path.write_text(json.dumps({"env": {"ANTHROPIC_AUTH_TOKEN": "fixed-master-key"}})) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "invalid or unexpected JSON" in result.output + assert "Restored" in result.output + assert not pid_record_path.exists() + assert not backup_path.exists() + assert json.loads(claude_settings_path.read_text()) == original_settings + + def test_surfaces_clean_error_when_backup_is_corrupt(self, monkeypatch, tmp_path): + _config_path, _log_path, _claude_settings_path, backup_path, _pid_record_path = _patch_paths( + monkeypatch, tmp_path + ) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("not json at all {{{") + + result = self.runner.invoke(down) + + assert result.exit_code != 0 + assert "invalid or unexpected JSON" in result.output diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_config.py b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py new file mode 100644 index 00000000000..f8d82476ef0 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_config.py @@ -0,0 +1,208 @@ +from typing import Any, Dict, Tuple + +import pytest + +from litellm.proxy.client.cli.commands.autoroute.config import ( + DEFAULT_KEYWORD_TIER_RULES, + AutorouteConfig, + ConfigGenerationError, + DiscoveredModel, + HeuristicClassifier, + KeywordTierRule, + LLMClassifier, + NoSemanticMatching, + SemanticMatching, + build_generated_model_list, + build_generated_proxy_config, + chat_models, + embedding_models, + parse_discovered_models, + validate_config, +) + +DISCOVERED: Tuple[DiscoveredModel, ...] = ( + DiscoveredModel(name="gpt-4o-mini", mode="chat"), + DiscoveredModel(name="gpt-4o", mode="chat"), + DiscoveredModel(name="o1", mode="chat"), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), +) + + +def _base_config(**overrides: Any) -> AutorouteConfig: + defaults: Dict[str, Any] = { + "base_url": "http://real-proxy.internal:4000", + "api_key": "sk-real-key", + "tiers": { + "SIMPLE": ("gpt-4o-mini",), + "MEDIUM": ("gpt-4o",), + "COMPLEX": ("gpt-4o",), + "REASONING": ("o1",), + }, + "default_model": "gpt-4o", + } + defaults.update(overrides) + return AutorouteConfig(**defaults) + + +class TestParseDiscoveredModels: + def test_parses_valid_raw_list_into_typed_tuple(self): + raw = [ + { + "model_group": "gpt-4o", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + {"model_group": "text-embedding-3-small", "mode": "embedding"}, + ] + result = parse_discovered_models(raw) + assert result == ( + DiscoveredModel(name="gpt-4o", mode="chat", input_cost_per_token=0.01, output_cost_per_token=0.02), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), + ) + + def test_ignores_unknown_extra_fields(self): + raw = [{"model_group": "gpt-4o", "mode": "chat", "totally_unknown_field": "whatever"}] + result = parse_discovered_models(raw) + assert result == (DiscoveredModel(name="gpt-4o", mode="chat"),) + + def test_missing_mode_defaults_to_chat(self): + raw = [{"model_group": "gpt-4o"}] + result = parse_discovered_models(raw) + assert result[0].mode == "chat" + + +class TestChatAndEmbeddingFiltering: + def test_filters_by_mode(self): + models = ( + DiscoveredModel(name="gpt-4o", mode="chat"), + DiscoveredModel(name="text-embedding-3-small", mode="embedding"), + DiscoveredModel(name="claude", mode="chat"), + ) + assert chat_models(models) == (models[0], models[2]) + assert embedding_models(models) == (models[1],) + + +class TestBuildGeneratedModelList: + def test_dedups_model_used_in_multiple_roles(self): + config = _base_config(classifier=LLMClassifier(model="gpt-4o")) + model_list = build_generated_model_list(config) + gpt4o_entries = [m for m in model_list if m["model_name"] == "gpt-4o"] + assert len(gpt4o_entries) == 1 + + def test_every_proxy_deployment_points_back_at_customer_proxy(self): + config = _base_config() + model_list = build_generated_model_list(config) + proxy_entries = [m for m in model_list if m["model_name"] not in ("autorouter", "*")] + names = {m["model_name"] for m in proxy_entries} + assert names == {"gpt-4o-mini", "gpt-4o", "o1"} + for entry in proxy_entries: + assert entry["litellm_params"]["model"] == f"litellm_proxy/{entry['model_name']}" + assert entry["litellm_params"]["api_base"] == config.base_url + assert entry["litellm_params"]["api_key"] == config.api_key + + def test_no_wildcard_deployment_is_generated(self): + # A bare "*" model_name looks like the obvious catch-all, but Router's auto-router + # registry is keyed by the literal requested model string with no wildcard resolution + # (litellm/router.py:10711-10717), so a "*" entry here would silently never match real + # traffic. Regression guard: don't reintroduce it. + config = _base_config() + model_list = build_generated_model_list(config) + assert not any(m["model_name"] == "*" for m in model_list) + + def test_complexity_router_config_reflects_llm_classifier(self): + config = _base_config(classifier=LLMClassifier(model="gpt-4o", timeout_ms=1234)) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["classifier_type"] == "llm" + assert router_config["classifier_llm_config"] == {"model": "gpt-4o", "timeout_ms": 1234} + assert "semantic_keyword_matching" not in router_config + assert "adaptive" not in router_config + + def test_complexity_router_config_reflects_semantic_matching(self): + config = _base_config( + semantic_matching=SemanticMatching(embedding_model="text-embedding-3-small", match_threshold=0.7) + ) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["semantic_keyword_matching"] is True + assert router_config["embedding_model"] == "text-embedding-3-small" + assert router_config["match_threshold"] == 0.7 + assert router_config["keyword_tier_rules"] + assert "classifier_type" not in router_config + + def test_semantic_matching_defaults_emit_builtin_keyword_rules(self): + config = _base_config(semantic_matching=SemanticMatching(embedding_model="text-embedding-3-small")) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["keyword_tier_rules"] == [ + {"keywords": list(rule.keywords), "tier": rule.tier} for rule in DEFAULT_KEYWORD_TIER_RULES + ] + + def test_semantic_matching_serializes_custom_keyword_rules(self): + config = _base_config( + semantic_matching=SemanticMatching( + embedding_model="text-embedding-3-small", + keyword_tier_rules=( + KeywordTierRule(keywords=("yo", "sup"), tier="SIMPLE"), + KeywordTierRule(keywords=("architect", "design a system"), tier="COMPLEX"), + ), + ) + ) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert router_config["keyword_tier_rules"] == [ + {"keywords": ["yo", "sup"], "tier": "SIMPLE"}, + {"keywords": ["architect", "design a system"], "tier": "COMPLEX"}, + ] + + def test_complexity_router_config_reflects_adaptive(self): + config = _base_config(adaptive=True) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + assert autorouter["litellm_params"]["complexity_router_config"]["adaptive"] is True + + def test_default_classifier_and_semantic_matching_add_no_extra_keys(self): + config = _base_config(classifier=HeuristicClassifier(), semantic_matching=NoSemanticMatching()) + autorouter = next(m for m in build_generated_model_list(config) if m["model_name"] == "autorouter") + router_config = autorouter["litellm_params"]["complexity_router_config"] + assert set(router_config.keys()) == {"tiers", "default_model"} + + +class TestBuildGeneratedProxyConfig: + def test_embeds_master_key_under_general_settings(self): + config = _base_config() + proxy_config = build_generated_proxy_config(config, "sk-master-123") + assert proxy_config["general_settings"] == {"master_key": "sk-master-123"} + assert proxy_config["model_list"] == build_generated_model_list(config) + + +class TestValidateConfig: + def test_passes_for_fully_valid_config(self): + validate_config(_base_config(), DISCOVERED) + + def test_raises_for_tier_referencing_unknown_model(self): + config = _base_config( + tiers={ + "SIMPLE": ("unknown-model",), + "MEDIUM": ("gpt-4o",), + "COMPLEX": ("gpt-4o",), + "REASONING": ("o1",), + } + ) + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_default_model(self): + config = _base_config(default_model="unknown-model") + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_llm_classifier_model(self): + config = _base_config(classifier=LLMClassifier(model="unknown-model")) + with pytest.raises(ConfigGenerationError, match="unknown-model"): + validate_config(config, DISCOVERED) + + def test_raises_for_unknown_semantic_embedding_model(self): + config = _base_config(semantic_matching=SemanticMatching(embedding_model="unknown-embedding")) + with pytest.raises(ConfigGenerationError, match="unknown-embedding"): + validate_config(config, DISCOVERED) diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_process.py b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py new file mode 100644 index 00000000000..a4f85ea44ff --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_process.py @@ -0,0 +1,158 @@ +import os +import socket +from typing import Optional +from unittest.mock import patch + +import pytest + +from litellm.proxy.client.cli.commands.autoroute import process as process_module +from litellm.proxy.client.cli.commands.autoroute.process import ( + PidRecord, + ProcessLaunchError, + UpError, + allocate_free_port, + clear_pid_record, + is_running, + launch_proxy, + missing_proxy_runtime_modules, + poll_liveliness, + read_pid_record, + write_pid_record, +) + + +class FakeProcess: + def __init__(self, returncode: Optional[int] = None): + self.returncode = returncode + + def poll(self) -> Optional[int]: + return self.returncode + + +class FakeResponse: + def __init__(self, status_code: int): + self.status_code = status_code + + +def test_allocate_free_port_returns_a_bindable_port(): + port = allocate_free_port() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", port)) + + +class TestLaunchProxy: + def test_binds_loopback_only_not_all_interfaces(self, tmp_path): + """proxy_cli.py's own --host default is 0.0.0.0 -- without an explicit override here, the + ephemeral proxy would be reachable from other hosts on the network despite base_url always + being built from 127.0.0.1, exposing its unauthenticated-until-master-key-lands routes.""" + config_path = tmp_path / "config.yaml" + log_path = tmp_path / "proxy.log" + + with patch.object(process_module.subprocess, "Popen") as mock_popen: + launch_proxy(config_path, 12345, log_path) + + args = mock_popen.call_args[0][0] + assert "--host" in args + assert args[args.index("--host") + 1] == "127.0.0.1" + + +class TestPidRecordRoundTrip: + def test_write_then_read_round_trips(self, tmp_path): + path = tmp_path / "pid.json" + record = PidRecord(pid=123, port=4000, config_path="/tmp/config.yaml", log_path="/tmp/proxy.log") + + write_pid_record(record, path) + + assert read_pid_record(path) == record + + def test_read_missing_file_returns_none(self, tmp_path): + assert read_pid_record(tmp_path / "missing.json") is None + + def test_read_raises_clean_error_on_corrupt_content(self, tmp_path): + path = tmp_path / "pid.json" + path.write_text("not json at all {{{") + + with pytest.raises(UpError, match="invalid or unexpected JSON"): + read_pid_record(path) + + def test_clear_removes_an_existing_record(self, tmp_path): + path = tmp_path / "pid.json" + write_pid_record(PidRecord(pid=1, port=1, config_path="a", log_path="b"), path) + assert path.exists() + + clear_pid_record(path) + + assert not path.exists() + + def test_clear_missing_file_is_a_no_op(self, tmp_path): + clear_pid_record(tmp_path / "missing.json") + + def test_write_creates_parent_directories(self, tmp_path): + path = tmp_path / "nested" / "dir" / "pid.json" + + write_pid_record(PidRecord(pid=1, port=1, config_path="a", log_path="b"), path) + + assert path.exists() + + +class TestIsRunning: + def test_current_process_is_running(self): + assert is_running(os.getpid()) is True + + def test_huge_unlikely_pid_is_not_running(self): + assert is_running(2**30) is False + + def test_permission_error_from_kill_is_treated_as_running(self, monkeypatch): + def fake_kill(pid: int, sig: int) -> None: + raise PermissionError("not permitted to signal this pid") + + monkeypatch.setattr(process_module.os, "kill", fake_kill) + + assert is_running(999) is True + + +class TestPollLiveliness: + def test_succeeds_when_health_check_returns_200_quickly(self, monkeypatch, tmp_path): + monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(200)) + + poll_liveliness("http://127.0.0.1:4000", tmp_path / "proxy.log", FakeProcess(), timeout=5.0) + + def test_raises_with_log_tail_when_timeout_elapses(self, monkeypatch, tmp_path): + log_path = tmp_path / "proxy.log" + log_path.write_text("line one\nline two\nline three\n") + monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(500)) + monkeypatch.setattr(process_module.time, "sleep", lambda seconds: None) + + with pytest.raises(ProcessLaunchError) as exc_info: + poll_liveliness("http://127.0.0.1:4000", log_path, FakeProcess(), timeout=0.05) + + assert "never became healthy" in str(exc_info.value) + assert "line three" in str(exc_info.value) + + def test_raises_immediately_when_process_already_exited(self, tmp_path): + log_path = tmp_path / "proxy.log" + log_path.write_text("crash log line") + + with pytest.raises(ProcessLaunchError) as exc_info: + poll_liveliness("http://127.0.0.1:4000", log_path, FakeProcess(returncode=1), timeout=5.0) + + assert "exited early" in str(exc_info.value) + assert "crash log line" in str(exc_info.value) + + +class TestMissingProxyRuntimeModules: + def test_flags_absent_modules_only(self, monkeypatch): + """A thin litellm[cli] install lacks the proxy runtime; the missing ones must be reported + (by name, for an actionable error) while modules that are importable are not.""" + monkeypatch.setattr( + process_module, + "_PROXY_RUNTIME_MODULES", + ("os", "litellm_autoroute_definitely_absent_pkg", "socket"), + ) + + assert missing_proxy_runtime_modules() == ("litellm_autoroute_definitely_absent_pkg",) + + def test_empty_when_all_present(self, monkeypatch): + monkeypatch.setattr(process_module, "_PROXY_RUNTIME_MODULES", ("os", "socket")) + + assert missing_proxy_runtime_modules() == () diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py new file mode 100644 index 00000000000..40d3e7f2aee --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -0,0 +1,56 @@ +from litellm.proxy.client.cli.commands.autoroute.settings import ( + ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, + merge_claude_settings_static_token, +) + + +def test_preserves_unrelated_top_level_keys(): + merged = merge_claude_settings_static_token({"theme": "dark"}, "http://127.0.0.1:4000", "token-abc") + assert merged["theme"] == "dark" + + +def test_preserves_unrelated_env_keys(): + settings = {"env": {"SOME_OTHER_VAR": "value"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["SOME_OTHER_VAR"] == "value" + + +def test_sets_base_url_and_auth_token(): + merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" + assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + + +def test_drops_stray_api_key(): + settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert "ANTHROPIC_API_KEY" not in merged["env"] + + +def test_removes_existing_api_key_helper(): + settings = {"apiKeyHelper": "/usr/local/bin/lite auth print-token"} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert "apiKeyHelper" not in merged + + +def test_does_not_mutate_input(): + settings = {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} + merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert settings == {"env": {"FOO": "bar"}, "apiKeyHelper": "old-helper"} + + +def test_forces_all_claude_code_default_model_tiers_to_the_autorouter(): + # A bare "*" model_name deployment looks like the obvious way to catch every request + # regardless of which model Claude Code thinks it's using, but Router's auto-router + # registry is keyed by the literal requested model string with no wildcard resolution + # (litellm/router.py:10711-10717) -- so the only reliable way to make every one of Claude + # Code's own tiers hit the auto-router is to override the env vars it reads per tier. + merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000", "token-abc") + for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: + assert merged["env"][key] == "autorouter" + + +def test_overrides_a_preexisting_default_model_env_var(): + settings = {"env": {"ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-opus-4-8"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "autorouter" diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py new file mode 100644 index 00000000000..2b9240aafc7 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py @@ -0,0 +1,329 @@ +import asyncio +from typing import Any, Dict, List, Tuple +from unittest.mock import patch + +import click +import pytest +import yaml +from click.testing import CliRunner +from InquirerPy.base.control import Choice +from prompt_toolkit.application import create_app_session +from prompt_toolkit.input import create_pipe_input +from prompt_toolkit.output import DummyOutput + +from litellm.proxy.client.cli.commands.autoroute import wizard as wizard_module +from litellm.proxy.client.cli.commands.autoroute.config import DiscoveredModel +from litellm.proxy.client.cli.commands.autoroute.wizard import run_configure_wizard + +CHAT_AND_EMBEDDING_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "gpt-4o-mini", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + {"model_group": "gpt-4o", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + {"model_group": "claude-opus", "mode": "chat"}, + {"model_group": "o1", "mode": "chat"}, + {"model_group": "text-embedding-3-small", "mode": "embedding"}, +] + +CHAT_ONLY_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "gpt-4o-mini", "mode": "chat"}, + {"model_group": "gpt-4o", "mode": "chat"}, + {"model_group": "claude-opus", "mode": "chat"}, + {"model_group": "o1", "mode": "chat"}, +] + +EMBEDDING_ONLY_GROUPS: List[Dict[str, Any]] = [ + {"model_group": "text-embedding-3-small", "mode": "embedding"}, +] + + +@click.command() +@click.pass_context +def _invoke_wizard(ctx: click.Context) -> None: + run_configure_wizard(ctx) + + +def _run( + tmp_path, + raw_groups: List[Dict[str, Any]], + tier_picks: Dict[str, Tuple[str, ...]], + input_str: str, + classifier_pick: str = "", + embedding_pick: str = "", +): + """Drives run_configure_wizard's orchestration logic (discovery, validation, config writing, + classifier/semantic/adaptive branching) by mocking the fuzzy picker itself, since that widget + is a real prompt_toolkit application tested separately in TestFuzzyPickWidget. CliRunner's + injected input still drives the plain click.confirm() y/n prompts.""" + config_path = tmp_path / "config.yaml" + runner = CliRunner() + + def _fake_prompt_for_models(models, prompt_label): + return tier_picks[prompt_label] + + def _fake_prompt_for_model(models, prompt_label): + if prompt_label == "LLM classifier": + return classifier_pick + if prompt_label == "semantic embeddings": + return embedding_pick + raise AssertionError(f"unexpected single-pick prompt_label {prompt_label!r}") + + with ( + patch.object(wizard_module, "Client") as mock_client_cls, + patch.object(wizard_module, "CONFIG_PATH", config_path), + patch.object(wizard_module, "_is_interactive", return_value=True), + patch.object(wizard_module, "_render_and_prompt_for_models", side_effect=_fake_prompt_for_models), + patch.object(wizard_module, "_render_and_prompt_for_model", side_effect=_fake_prompt_for_model), + ): + mock_client_cls.return_value.model_groups.info.return_value = raw_groups + result = runner.invoke( + _invoke_wizard, + obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}, + input=input_str, + ) + return result, config_path + + +def _router_config(config_path) -> Dict[str, Any]: + written = yaml.safe_load(config_path.read_text()) + autorouter = next(m for m in written["model_list"] if m["model_name"] == "autorouter") + return autorouter["litellm_params"]["complexity_router_config"] + + +_SIMPLE_TIER_PICKS: Dict[str, Tuple[str, ...]] = { + "SIMPLE": ("gpt-4o-mini",), + "MEDIUM": ("gpt-4o",), + "COMPLEX": ("claude-opus",), + "REASONING": ("o1",), +} + + +class TestRunConfigureWizardHappyPath: + def test_assigns_tiers_and_declines_everything(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\nn\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["tiers"] == { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": ["gpt-4o"], + "COMPLEX": ["claude-opus"], + "REASONING": ["o1"], + } + assert router_config["default_model"] == "gpt-4o" + assert "classifier_type" not in router_config + assert "classifier_llm_config" not in router_config + assert "semantic_keyword_matching" not in router_config + assert "adaptive" not in router_config + + def test_assigns_multiple_models_to_a_single_tier(self, tmp_path): + tier_picks = {**_SIMPLE_TIER_PICKS, "SIMPLE": ("gpt-4o-mini", "gpt-4o")} + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, tier_picks, input_str="n\nn\nn\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["tiers"]["SIMPLE"] == ["gpt-4o-mini", "gpt-4o"] + assert router_config["default_model"] == "gpt-4o" + + def test_writes_config_file_with_restricted_permissions(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\nn\n") + + assert result.exit_code == 0, result.output + assert config_path.exists() + assert oct(config_path.stat().st_mode)[-3:] == "600" + + def test_no_embedding_pool_skips_semantic_prompt_entirely(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_ONLY_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert "semantic_keyword_matching" not in router_config + + +class TestRunConfigureWizardLLMClassifier: + def test_accepting_llm_classifier_records_chosen_model(self, tmp_path): + result, config_path = _run( + tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="y\nn\nn\n", classifier_pick="gpt-4o" + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["classifier_type"] == "llm" + assert router_config["classifier_llm_config"]["model"] == "gpt-4o" + + +class TestRunConfigureWizardSemanticMatching: + def test_accepting_semantic_matching_records_embedding_model(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + _SIMPLE_TIER_PICKS, + input_str="n\ny\n\n\n\n\nn\n", + embedding_pick="text-embedding-3-small", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["semantic_keyword_matching"] is True + assert router_config["embedding_model"] == "text-embedding-3-small" + + def test_blank_keyword_answers_keep_the_builtin_defaults(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + _SIMPLE_TIER_PICKS, + input_str="n\ny\n\n\n\n\nn\n", + embedding_pick="text-embedding-3-small", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["keyword_tier_rules"] == [ + {"keywords": ["hi", "hello", "thanks"], "tier": "SIMPLE"}, + {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, + {"keywords": ["refactor", "implement", "debug"], "tier": "COMPLEX"}, + {"keywords": ["step by step", "think through", "prove"], "tier": "REASONING"}, + ] + + def test_custom_keyword_answers_are_recorded_per_tier(self, tmp_path): + result, config_path = _run( + tmp_path, + CHAT_AND_EMBEDDING_GROUPS, + _SIMPLE_TIER_PICKS, + input_str="n\ny\nyo, sup\n\nbuild a service, migrate\nderive, prove rigorously\nn\n", + embedding_pick="text-embedding-3-small", + ) + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["keyword_tier_rules"] == [ + {"keywords": ["yo", "sup"], "tier": "SIMPLE"}, + {"keywords": ["explain", "how does"], "tier": "MEDIUM"}, + {"keywords": ["build a service", "migrate"], "tier": "COMPLEX"}, + {"keywords": ["derive", "prove rigorously"], "tier": "REASONING"}, + ] + + +class TestRunConfigureWizardAdaptive: + def test_accepting_adaptive_sets_adaptive_flag(self, tmp_path): + result, config_path = _run(tmp_path, CHAT_AND_EMBEDDING_GROUPS, _SIMPLE_TIER_PICKS, input_str="n\nn\ny\n") + + assert result.exit_code == 0, result.output + router_config = _router_config(config_path) + assert router_config["adaptive"] is True + + +class TestRunConfigureWizardNoChatModels: + def test_fails_cleanly_without_prompting_when_no_chat_models(self, tmp_path): + result, config_path = _run(tmp_path, EMBEDDING_ONLY_GROUPS, {}, input_str="") + + assert result.exit_code != 0 + assert "no chat-capable models" in result.output.lower() + assert not config_path.exists() + + def test_surfaces_clean_error_when_response_is_not_a_list(self, tmp_path): + result, config_path = _run(tmp_path, {"data": CHAT_AND_EMBEDDING_GROUPS}, {}, input_str="") + + assert result.exit_code != 0 + assert result.exception is None or not isinstance(result.exception, AssertionError) + assert "Unexpected response from /model_group/info" in result.output + assert not config_path.exists() + + +class TestRunConfigureWizardNotInteractive: + def test_fails_cleanly_when_not_a_tty(self, tmp_path): + config_path = tmp_path / "config.yaml" + runner = CliRunner() + with ( + patch.object(wizard_module, "Client") as mock_client_cls, + patch.object(wizard_module, "CONFIG_PATH", config_path), + patch.object(wizard_module, "_is_interactive", return_value=False), + ): + mock_client_cls.return_value.model_groups.info.return_value = CHAT_AND_EMBEDDING_GROUPS + result = runner.invoke(_invoke_wizard, obj={"base_url": "http://localhost:4000", "api_key": "sk-test"}) + + assert result.exit_code != 0 + assert "interactive terminal" in result.output + assert not config_path.exists() + + +def _drive_fuzzy_pick( + models: Tuple[DiscoveredModel, ...], + prompt_label: str, + multiselect: bool, + key_events: List[Tuple[str, float]], +) -> List[str]: + """Drives the real InquirerPy fuzzy prompt through prompt_toolkit's own test input/output, + exercising the actual widget (filtering, tab-to-toggle, enter-to-confirm) rather than mocking + it away. asyncio.to_thread propagates the create_app_session context into the worker thread + running _fuzzy_pick's synchronous .execute() call.""" + + async def _run() -> List[str]: + with create_pipe_input() as pipe_input: + with create_app_session(input=pipe_input, output=DummyOutput()): + task = asyncio.ensure_future( + asyncio.to_thread(wizard_module._fuzzy_pick, models, prompt_label, multiselect) + ) + await asyncio.sleep(0.05) + for text, delay in key_events: + pipe_input.send_text(text) + await asyncio.sleep(delay) + return await task + + return asyncio.run(_run()) + + +class TestFuzzyPickWidget: + def _models(self) -> Tuple[DiscoveredModel, ...]: + return tuple(DiscoveredModel(name=f"model-{i}") for i in range(20)) + + def test_single_select_filters_and_returns_highlighted_match(self): + result = _drive_fuzzy_pick( + self._models(), "test", multiselect=False, key_events=[("model-13", 0.3), ("\r", 0.1)] + ) + assert result == ["model-13"] + + def test_multiselect_requires_tab_to_toggle_before_enter(self): + result = _drive_fuzzy_pick( + self._models(), "test", multiselect=True, key_events=[("model-7", 0.3), ("\t", 0.1), ("\r", 0.1)] + ) + assert result == ["model-7"] + + def test_multiselect_can_pick_more_than_one_across_filters(self): + result = _drive_fuzzy_pick( + self._models(), + "test", + multiselect=True, + key_events=[ + ("model-3", 0.3), + ("\t", 0.1), + *[("\x7f", 0.02) for _ in range("model-3".__len__())], + ("model-15", 0.3), + ("\t", 0.1), + ("\r", 0.1), + ], + ) + assert set(result) == {"model-3", "model-15"} + + def test_choice_wraps_name_and_value_to_the_same_model_name(self): + model = DiscoveredModel(name="only-model") + choice = Choice(value=model.name, name=model.name) + assert choice.value == choice.name == "only-model" + + +class TestRenderAndPromptForModelWrappers: + def test_single_pick_wrapper_returns_bare_string(self): + with patch.object(wizard_module, "_fuzzy_pick", return_value=["model-a"]) as mock_pick: + result = wizard_module._render_and_prompt_for_model((), "tier") + assert result == "model-a" + mock_pick.assert_called_once_with((), "tier", multiselect=False) + + def test_multi_pick_wrapper_returns_tuple(self): + with patch.object(wizard_module, "_fuzzy_pick", return_value=["model-a", "model-b"]) as mock_pick: + result = wizard_module._render_and_prompt_for_models((), "tier") + assert result == ("model-a", "model-b") + mock_pick.assert_called_once_with((), "tier", multiselect=True) + + +@pytest.mark.parametrize("isatty_value", [True, False]) +def test_is_interactive_reflects_stdin_isatty(isatty_value): + with patch.object(wizard_module.sys.stdin, "isatty", return_value=isatty_value): + assert wizard_module._is_interactive() is isatty_value diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 6be43c9da44..e101515d4b1 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -8,6 +8,7 @@ from unittest.mock import Mock, mock_open, patch sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path +import pytest from click.testing import CliRunner from litellm.constants import CLI_JWT_EXPIRATION_HOURS @@ -40,6 +41,138 @@ def _mock_cli_sso_start_response( return mock_response +class TestPollingErrorSurfacing: + def test_client_error_raises_with_server_detail_and_stops_polling(self): + from litellm.proxy.client.cli.commands.auth import _poll_for_ready_data + + mock_response = Mock() + mock_response.status_code = 400 + mock_response.json.return_value = { + "detail": "Your litellm CLI is out of date and uses a login flow this proxy no longer supports." + } + + with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): + with pytest.raises(ValueError) as exc_info: + _poll_for_ready_data("http://test/sso/cli/poll/sk-legacy") + + assert mock_get.call_count == 1 + assert ( + "The proxy rejected the login session with HTTP 400: Your litellm CLI is out of date " + "and uses a login flow this proxy no longer supports." in str(exc_info.value) + ) + + def test_login_command_shows_server_rejection_to_user(self): + mock_context = Mock() + mock_context.obj = {"base_url": "https://test.example.com"} + + mock_poll_response = Mock() + mock_poll_response.status_code = 400 + mock_poll_response.json.return_value = {"detail": "CLI login session not found or expired."} + + with ( + patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), + patch("requests.get", return_value=mock_poll_response), + patch("time.sleep"), + ): + result = CliRunner().invoke(login, obj=mock_context.obj) + + assert result.exit_code == 0 + assert "❌ Authentication failed:" in result.output + assert "CLI login session not found or expired." in result.output + assert "Authentication timed out" not in result.output + + def test_server_error_without_json_body_retries_until_timeout(self, capsys): + from litellm.proxy.client.cli.commands.auth import _poll_for_ready_data + + mock_response = Mock() + mock_response.status_code = 500 + mock_response.json.side_effect = ValueError("no json") + + with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): + result = _poll_for_ready_data("http://test/sso/cli/poll/cli-abc", total_timeout=6, poll_interval=2) + + assert result is None + assert mock_get.call_count == 3 + assert "Polling error: HTTP 500" in capsys.readouterr().out + + def test_rate_limit_is_retried_not_aborted(self, capsys): + from litellm.proxy.client.cli.commands.auth import _poll_for_ready_data + + mock_response = Mock() + mock_response.status_code = 429 + mock_response.json.return_value = {"detail": "Too many CLI login attempts. Try again later."} + + with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): + result = _poll_for_ready_data("http://test/sso/cli/poll/cli-abc", total_timeout=4, poll_interval=2) + + assert result is None + assert mock_get.call_count == 2 + assert "Polling error: HTTP 429: Too many CLI login attempts. Try again later." in capsys.readouterr().out + + +class TestStartCliSsoFlowErrors: + def test_endpoint_not_found_explains_version_or_base_url(self): + from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow + + mock_response = Mock() + mock_response.status_code = 404 + + with patch("requests.post", return_value=mock_response): + with pytest.raises(ValueError) as exc_info: + _start_cli_sso_flow("https://old-proxy.example.com") + + message = str(exc_info.value) + assert "HTTP 404" in message + assert "--base-url" in message + assert "older than this CLI" in message + + def test_http_error_includes_server_detail(self): + from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow + + mock_response = Mock() + mock_response.status_code = 429 + mock_response.json.return_value = {"detail": "Too many CLI login attempts. Try again later."} + + with patch("requests.post", return_value=mock_response): + with pytest.raises(ValueError) as exc_info: + _start_cli_sso_flow("https://test.example.com") + + assert "HTTP 429" in str(exc_info.value) + assert "Too many CLI login attempts. Try again later." in str(exc_info.value) + + def test_non_json_response_names_interception(self): + from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.side_effect = ValueError("no json") + mock_response.headers = {"content-type": "text/html"} + mock_response.text = "Sign in to corporate VPN" + + with patch("requests.post", return_value=mock_response): + with pytest.raises(ValueError) as exc_info: + _start_cli_sso_flow("https://test.example.com") + + message = str(exc_info.value) + assert "non-JSON response" in message + assert "text/html" in message + assert "Sign in to corporate VPN" in message + + def test_connection_error_points_at_base_url(self): + import requests + + from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow + + with patch("requests.post", side_effect=requests.ConnectionError("Connection refused")): + with pytest.raises(ValueError) as exc_info: + _start_cli_sso_flow("https://unreachable.example.com") + + message = str(exc_info.value) + assert "Could not reach the proxy" in message + assert "https://unreachable.example.com/sso/cli/start" in message + + class TestTokenUtilities: """Test token file utility functions""" @@ -664,14 +797,18 @@ class TestPrintTokenCommand: verbatim as the bearer token, so any diagnostic text on stdout would corrupt authentication. - apiKeyHelper is configured as a bare command (managed-settings.json sets - just `"apiKeyHelper": "lite auth print-token"`, no --base-url flag) -- - so in the common case ctx.obj has no explicit base_url at all, and the - command must resolve the server from whatever `lite login` stored in - token.json, not from a CLI default. `--base-url`/`LITELLM_PROXY_URL` - only matters when a caller explicitly overrides it (tracked via - ctx.obj["base_url_explicit"], set by the `cli` group from - click's ParameterSource). + `lite up` now writes `apiKeyHelper` with an explicit `--base-url` bound + to whatever proxy it was pointed at (resolve_api_key_helper), so + print-token enforces that the cached token was actually issued for that + server -- a token minted for a different, previously-logged-into proxy + must never be handed to whichever server the helper is invoked for. + Settings patched by an older `lite up`, or a manually-configured + apiKeyHelper, can still invoke this bare (no --base-url at all); that + case falls back to trusting whatever `lite login` stored in token.json, + since there is no explicit target to check it against. `--base-url`/ + `LITELLM_PROXY_URL` only enforces the match when a caller explicitly + passes it (tracked via ctx.obj["base_url_explicit"], set by the `cli` + group from click's ParameterSource). """ def setup_method(self): @@ -685,8 +822,9 @@ class TestPrintTokenCommand: assert "Not authenticated" in result.output def test_bare_invocation_resolves_server_from_stored_token(self): - """The apiKeyHelper's real invocation shape: no --base-url given at - all. Must use token.json's own base_url, not a hardcoded default.""" + """The legacy/manual invocation shape: no --base-url given at all + (e.g. settings patched before resolve_api_key_helper started binding + one). Must use token.json's own base_url, not a hardcoded default.""" with ( patch( "litellm.proxy.client.cli.commands.auth.load_token", @@ -706,7 +844,10 @@ class TestPrintTokenCommand: def test_explicit_base_url_mismatch_fails_cleanly(self): """When the caller *does* explicitly pass --base-url, a token issued - for a different server must never be printed.""" + for a different server must never be printed. This is the exact + scenario `lite up`'s own bound --base-url now guards against: a + token minted for proxy A must not reach a helper invocation aimed + at proxy B, even though the token itself is otherwise fresh.""" with patch( "litellm.proxy.client.cli.commands.auth.load_token", return_value={ @@ -723,6 +864,25 @@ class TestPrintTokenCommand: assert result.exit_code != 0 assert "sk-should-not-print" not in result.output + def test_explicit_base_url_match_prints_token(self): + """`lite up`'s own bound invocation shape: --base-url matching the token's origin + must succeed exactly like the bare/legacy invocation does.""" + with patch( + "litellm.proxy.client.cli.commands.auth.load_token", + return_value={ + "base_url": "http://localhost:4000", + "key": "sk-matches", + "timestamp": time.time(), + }, + ): + result = self.runner.invoke( + print_token, + obj={"base_url": "http://localhost:4000", "base_url_explicit": True}, + ) + + assert result.exit_code == 0 + assert result.output.strip() == "sk-matches" + def test_fresh_cached_key_printed_without_network_call(self): """A recently-issued key should be printed straight from cache -- no refresh call on every single invocation (apiKeyHelper gets called diff --git a/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py b/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py new file mode 100644 index 00000000000..c2809a90ba2 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_model_groups_commands.py @@ -0,0 +1,114 @@ +import json +import os +from typing import Any, Dict, List +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli import cli + +SAMPLE_MODEL_GROUPS: List[Dict[str, Any]] = [ + { + "model_group": "gpt-4o", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + { + "model_group": "text-embedding-3-small", + "mode": "embedding", + "input_cost_per_token": 0.0001, + "output_cost_per_token": None, + }, +] + + +@pytest.fixture +def mock_client(): + with patch("litellm.proxy.client.cli.commands.model_groups.Client") as MockClient: + yield MockClient + + +@pytest.fixture +def cli_runner(): + return CliRunner() + + +@pytest.fixture(autouse=True) +def mock_env(): + with patch.dict( + os.environ, + { + "LITELLM_PROXY_URL": "http://localhost:4000", + "LITELLM_PROXY_API_KEY": "sk-test", + }, + ): + yield + + +def test_list_table_format_shows_model_names_and_modes(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = SAMPLE_MODEL_GROUPS + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code == 0, result.output + assert "gpt-4o" in result.output + assert "chat" in result.output + assert "text-embedding-3-small" in result.output + assert "embedding" in result.output + assert "0.01" in result.output + assert "0.02" in result.output + + mock_client.assert_called_once_with(base_url="http://localhost:4000", api_key="sk-test") + mock_client.return_value.model_groups.info.assert_called_once() + + +def test_list_table_format_defaults_missing_mode_to_chat(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = [{"model_group": "some-model"}] + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code == 0, result.output + assert "some-model" in result.output + assert "chat" in result.output + + +def test_list_json_format_round_trips_raw_data(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = SAMPLE_MODEL_GROUPS + + result = cli_runner.invoke(cli, ["model-groups", "list", "--format", "json"]) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == SAMPLE_MODEL_GROUPS + + +def test_list_with_custom_base_url_and_api_key(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = [] + + result = cli_runner.invoke( + cli, + ["--base-url", "http://custom.server:8000", "--api-key", "custom-key", "model-groups", "list"], + ) + + assert result.exit_code == 0, result.output + mock_client.assert_called_once_with(base_url="http://custom.server:8000", api_key="custom-key") + + +def test_list_error_handling(mock_client, cli_runner): + mock_client.return_value.model_groups.info.side_effect = Exception("API Error") + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code != 0 + assert "API Error" in str(result.exception) + + +def test_list_surfaces_clean_error_when_response_is_not_a_list(mock_client, cli_runner): + mock_client.return_value.model_groups.info.return_value = {"data": SAMPLE_MODEL_GROUPS} + + result = cli_runner.invoke(cli, ["model-groups", "list"]) + + assert result.exit_code != 0 + assert result.exception is None or not isinstance(result.exception, AssertionError) + assert "Unexpected response from /model_group/info" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py new file mode 100644 index 00000000000..1b182553644 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -0,0 +1,398 @@ +import json +import shutil +import stat +import sys +from unittest.mock import patch + +import click +import pytest +from click.testing import CliRunner + +from litellm.proxy.client.cli.commands import up as up_module +from litellm.proxy.client.cli.commands.agents import AgentRunError +from litellm.proxy.client.cli.commands.up import ( + BackupRecord, + UpError, + _ensure_fresh_login, + down, + load_json_or_empty, + merge_claude_settings, + read_backup, + resolve_api_key_helper, + restore_claude_settings, + up, + write_backup, +) + +UP_MODULE = "litellm.proxy.client.cli.commands.up" + + +def _patch_paths(monkeypatch, tmp_path): + settings_path = tmp_path / "claude_settings.json" + backup_path = tmp_path / "backup.json" + monkeypatch.setattr(up_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(up_module, "BACKUP_PATH", backup_path) + return settings_path, backup_path + + +class TestMergeClaudeSettings: + def test_preserves_unrelated_top_level_keys(self): + merged = merge_claude_settings({"theme": "dark"}, "http://localhost:4000", "helper") + assert merged["theme"] == "dark" + + def test_preserves_unrelated_env_keys(self): + settings = {"env": {"SOME_OTHER_VAR": "value"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert merged["env"]["SOME_OTHER_VAR"] == "value" + + def test_overrides_base_url_and_helper(self): + settings = { + "env": {"ANTHROPIC_BASE_URL": "https://old.example.com"}, + "apiKeyHelper": "old-helper", + } + merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") + assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert merged["apiKeyHelper"] == "new-helper" + + def test_drops_stray_api_key(self): + settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert "ANTHROPIC_API_KEY" not in merged["env"] + + def test_works_from_empty_settings(self): + merged = merge_claude_settings({}, "http://localhost:4000", "helper") + assert merged["env"] == {"ANTHROPIC_BASE_URL": "http://localhost:4000"} + assert merged["apiKeyHelper"] == "helper" + + def test_does_not_mutate_input(self): + settings = {"env": {"FOO": "bar"}} + merge_claude_settings(settings, "http://localhost:4000", "helper") + assert settings == {"env": {"FOO": "bar"}} + + +class TestLoadJsonOrEmpty: + def test_returns_empty_dict_when_file_does_not_exist(self, tmp_path): + assert load_json_or_empty(tmp_path / "missing.json") == {} + + def test_returns_empty_dict_when_file_is_empty(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text("") + assert load_json_or_empty(path) == {} + + def test_returns_empty_dict_when_file_is_whitespace_only(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text(" \n") + assert load_json_or_empty(path) == {} + + def test_parses_real_content(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text(json.dumps({"theme": "dark"})) + assert load_json_or_empty(path) == {"theme": "dark"} + + def test_raises_clean_error_on_invalid_json(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text("not json at all {{{") + with pytest.raises(UpError, match="invalid JSON"): + load_json_or_empty(path) + + def test_raises_clean_error_on_non_object_root(self, tmp_path): + path = tmp_path / "settings.json" + path.write_text(json.dumps([1, 2, 3])) + with pytest.raises(UpError, match="invalid JSON"): + load_json_or_empty(path) + + +class TestBackupRoundTrip: + def test_restores_original_content_when_file_existed(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + original = {"apiKeyHelper": "old-helper", "theme": "dark"} + settings_path.write_text(json.dumps(original)) + + write_backup(BackupRecord(existed=True, content=original)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + restored = restore_claude_settings() + + assert restored is not None + assert restored.existed is True + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + def test_deletes_settings_file_when_it_did_not_exist_before(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=False, content=None)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + restored = restore_claude_settings() + + assert restored is not None + assert restored.existed is False + assert not settings_path.exists() + assert not backup_path.exists() + + def test_no_backup_is_a_no_op_returning_none(self, monkeypatch, tmp_path): + settings_path, _backup_path = _patch_paths(monkeypatch, tmp_path) + assert restore_claude_settings() is None + assert not settings_path.exists() + + def test_recreates_claude_dir_if_it_was_deleted_while_up_was_running(self, monkeypatch, tmp_path): + """If ~/.claude/ is removed while `lite up` holds it open, restoring must recreate the + directory rather than crash with FileNotFoundError and strand the backup file, which + would otherwise permanently break every future `lite down`.""" + claude_dir = tmp_path / "claude_dir" + settings_path = claude_dir / "settings.json" + backup_path = tmp_path / "backup.json" + monkeypatch.setattr(up_module, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(up_module, "BACKUP_PATH", backup_path) + original = {"theme": "dark"} + claude_dir.mkdir(parents=True) + write_backup(BackupRecord(existed=True, content=original)) + shutil.rmtree(claude_dir) + + restored = restore_claude_settings() + + assert restored is not None + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + def test_read_backup_round_trips_write_backup(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=True, content={"a": 1})) + assert read_backup() == BackupRecord(existed=True, content={"a": 1}) + + def test_read_backup_missing_file_returns_none(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + assert read_backup() is None + + def test_read_backup_raises_clean_error_on_corrupt_content(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("not json at all {{{") + + with pytest.raises(UpError, match="invalid or unexpected JSON"): + read_backup() + + def test_write_backup_restricts_permissions_for_a_new_file(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=True, content={"a": 1})) + assert stat.S_IMODE(backup_path.stat().st_mode) == 0o600 + + def test_write_backup_restricts_permissions_of_a_preexisting_permissive_file(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("{}") + backup_path.chmod(0o644) + + write_backup(BackupRecord(existed=True, content={"a": 1})) + + assert stat.S_IMODE(backup_path.stat().st_mode) == 0o600 + + def test_backup_file_always_removed_after_restore(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=False, content=None)) + assert backup_path.exists() + + restore_claude_settings() + + assert not backup_path.exists() + + +class TestResolveApiKeyHelper: + def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") + helper = resolve_api_key_helper("http://localhost:4000") + assert helper == "/usr/local/bin/lite auth print-token --base-url http://localhost:4000" + + def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") + helper = resolve_api_key_helper("http://example.com/path; rm -rf /") + assert helper == "/usr/local/bin/lite auth print-token --base-url 'http://example.com/path; rm -rf /'" + + def test_raises_when_lite_not_on_path(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: None) + with pytest.raises(UpError, match="Could not find `lite`"): + resolve_api_key_helper("http://localhost:4000") + + +def _make_ctx(base_url): + return click.Context(click.Command("test"), obj={"base_url": base_url}) + + +class TestEnsureFreshLogin: + """A token that is fresh but was issued for a *different* proxy must not be trusted: without + this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get an + apiKeyHelper wired up around proxy A's real token, which print-token would then hand to proxy B.""" + + def test_reuses_a_fresh_token_issued_for_the_same_proxy(self, monkeypatch): + monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"}) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + login_calls = [] + monkeypatch.setattr(up_module, "login", lambda ctx: login_calls.append(ctx)) + + _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + + assert login_calls == [] + + def test_forces_a_fresh_login_when_the_cached_token_is_for_a_different_proxy(self, monkeypatch): + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: True) + tokens = iter( + [ + {"key": "sk-a", "base_url": "http://proxy-a:4000"}, + {"key": "sk-b", "base_url": "http://proxy-b:4000"}, + ] + ) + monkeypatch.setattr(up_module, "load_token", lambda: next(tokens)) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + login_calls = [] + + @click.pass_context + def fake_login(ctx): + login_calls.append(ctx.obj["base_url"]) + + monkeypatch.setattr(up_module, "login", fake_login) + + _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + + assert login_calls == ["http://proxy-b:4000"] + + def test_fails_cleanly_non_interactively_when_only_a_different_proxys_token_is_cached(self, monkeypatch): + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: False) + monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"}) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + + with pytest.raises(UpError, match="lite login"): + _ensure_fresh_login(_make_ctx("http://proxy-b:4000")) + + +class TestUpCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_refuses_double_start_without_touching_settings_file(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + existing_backup = {"existed": False, "content": None} + backup_path.write_text(json.dumps(existing_backup)) + + with ( + patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), + patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), + patch(f"{UP_MODULE}.verify_proxy_key"), + ): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code != 0 + assert "already" in result.output + assert "lite down" in result.output + assert not settings_path.exists() + assert json.loads(backup_path.read_text()) == existing_backup + + def test_no_fresh_login_non_interactive_fails_cleanly(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + monkeypatch.setattr(sys.stdin, "isatty", lambda: False) + + with patch(f"{UP_MODULE}.load_token", return_value=None): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code != 0 + assert "lite login" in result.output + + def test_unreachable_proxy_fails_cleanly(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + with ( + patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), + patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), + patch( + f"{UP_MODULE}.verify_proxy_key", + side_effect=AgentRunError("Could not reach the LiteLLM proxy at http://localhost:4000"), + ), + ): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code != 0 + assert "Could not reach the LiteLLM proxy" in result.output + + def test_happy_path_writes_settings_and_backup_then_restores_on_stop(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + original = {"theme": "dark"} + settings_path.write_text(json.dumps(original)) + + captured = {} + + def fake_wait(self, timeout=None): + captured["settings"] = json.loads(settings_path.read_text()) + captured["backup_existed"] = backup_path.exists() + return True + + with ( + patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), + patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), + patch(f"{UP_MODULE}.verify_proxy_key"), + patch( + f"{UP_MODULE}.resolve_api_key_helper", + return_value="/usr/local/bin/lite auth print-token", + ), + patch(f"{UP_MODULE}.signal.signal"), + patch(f"{UP_MODULE}.atexit.register"), + patch("threading.Event.wait", new=fake_wait), + ): + result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) + + assert result.exit_code == 0, result.output + assert captured["backup_existed"] is True + assert captured["settings"]["theme"] == "dark" + assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert captured["settings"]["apiKeyHelper"] == "/usr/local/bin/lite auth print-token" + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + +class TestDownCommand: + def setup_method(self): + self.runner = CliRunner() + + def test_restores_when_backup_exists(self, monkeypatch, tmp_path): + settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + original = {"apiKeyHelper": "old-helper"} + write_backup(BackupRecord(existed=True, content=original)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Restored" in result.output + assert json.loads(settings_path.read_text()) == original + assert not backup_path.exists() + + def test_removes_settings_file_when_it_did_not_exist_before(self, monkeypatch, tmp_path): + settings_path, _backup_path = _patch_paths(monkeypatch, tmp_path) + write_backup(BackupRecord(existed=False, content=None)) + settings_path.write_text(json.dumps({"apiKeyHelper": "lite-helper"})) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Removed" in result.output + assert not settings_path.exists() + + def test_prints_nothing_to_restore_when_no_backup(self, monkeypatch, tmp_path): + _patch_paths(monkeypatch, tmp_path) + + result = self.runner.invoke(down) + + assert result.exit_code == 0, result.output + assert "Nothing to restore." in result.output + + def test_surfaces_clean_error_on_a_corrupt_backup_file(self, monkeypatch, tmp_path): + _settings_path, backup_path = _patch_paths(monkeypatch, tmp_path) + backup_path.parent.mkdir(parents=True, exist_ok=True) + backup_path.write_text("not json at all {{{") + + result = self.runner.invoke(down) + + assert result.exit_code != 0 + assert result.exception is None or isinstance(result.exception, SystemExit) + assert "invalid or unexpected JSON" in result.output diff --git a/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py b/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py new file mode 100644 index 00000000000..0446cfeeab0 --- /dev/null +++ b/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py @@ -0,0 +1,432 @@ +""" +Tests for the enterprise billing-metrics recorder and its factory. + +These verify the license gate, the missing-config and missing-cert disable +paths, the OTLP/HTTP exporter wiring (client cert+key authenticate us to the +collector's mTLS-terminating front end; CA override optional for private +collectors), the metric attribute mapping, and that recording produces the +expected OTLP counter via an in-memory reader. +""" + +import os +import socket +import stat +from pathlib import Path +from typing import Dict, List, Optional + +import pytest +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader + +from litellm.proxy.enterprise_billing import billing_metrics as bm +from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory + +_ENV_VARS = ( + bm.ENDPOINT_ENV, + bm.CLIENT_CERT_ENV, + bm.CLIENT_KEY_ENV, + bm.CA_CERT_ENV, + bm.EXPORT_INTERVAL_ENV, +) + + +@pytest.fixture(autouse=True) +def clear_env(monkeypatch): + for name in _ENV_VARS: + monkeypatch.delenv(name, raising=False) + yield + bm.shutdown_billing_metrics_recorder() + + +def _write_certs(tmp_path: Path) -> Dict[str, str]: + files = { + bm.CA_CERT_ENV: ("ca.pem", b"ca-bytes"), + bm.CLIENT_CERT_ENV: ("client.pem", b"client-cert-bytes"), + bm.CLIENT_KEY_ENV: ("client.key", b"client-key-bytes"), + } + paths = {} + for env_name, (filename, content) in files.items(): + path = tmp_path / filename + path.write_bytes(content) + paths[env_name] = str(path) + return paths + + +def _set_full_env(monkeypatch, tmp_path: Path) -> Dict[str, str]: + paths = _write_certs(tmp_path) + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV]) + monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) + return paths + + +def _config(tmp_path: Path, license_id: Optional[str] = "org-1") -> bm.BillingMetricsConfig: + paths = _write_certs(tmp_path) + return bm.BillingMetricsConfig( + endpoint="https://collector.example:4317", + client_cert_path=paths[bm.CLIENT_CERT_ENV], + client_key_path=paths[bm.CLIENT_KEY_ENV], + ca_cert_path=paths[bm.CA_CERT_ENV], + export_interval_ms=60_000, + litellm_version="1.2.3", + license_id=license_id, + ) + + +# ── Factory gating ──────────────────────────────────────────────────────────── + + +def test_not_premium_returns_none(tmp_path, monkeypatch): + _set_full_env(monkeypatch, tmp_path) + assert bm.build_billing_metrics_recorder(premium=False, license_data=None, litellm_version="1.0") is None + + +def test_premium_without_config_returns_none(monkeypatch): + assert bm.build_billing_metrics_recorder(premium=True, license_data={"user_id": "x"}, litellm_version="1.0") is None + + +def test_premium_with_missing_cert_files_returns_none(monkeypatch, tmp_path): + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CA_CERT_ENV, str(tmp_path / "missing-ca.pem")) + monkeypatch.setenv(bm.CLIENT_CERT_ENV, str(tmp_path / "missing-cert.pem")) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, str(tmp_path / "missing-key.pem")) + assert bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") is None + + +def test_premium_with_full_config_builds_recorder(monkeypatch, tmp_path): + """Builds a real MeterProvider, so the exporter is stubbed: the live one + resolves the collector and opens a TLS connection during the shutdown flush. + The getaddrinfo spy keeps that stub from being quietly dropped later.""" + _set_full_env(monkeypatch, tmp_path) + monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class({})) + + resolved: List[str] = [] + real_getaddrinfo = socket.getaddrinfo + + def _spy_getaddrinfo(host, port, *args, **kwargs): + resolved.append(str(host)) + return real_getaddrinfo(host, port, *args, **kwargs) + + monkeypatch.setattr(socket, "getaddrinfo", _spy_getaddrinfo) + + recorder = bm.build_billing_metrics_recorder( + premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0" + ) + assert isinstance(recorder, bm.BillingMetricsRecorder) + recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id=None) + bm.shutdown_billing_metrics_recorder() + + assert [host for host in resolved if "collector.example" in host] == [] + + +def test_building_the_recorder_logs_an_affirmative_line(monkeypatch, tmp_path): + """ + Every disable path logs; a successful build must log too. Otherwise an + operator cannot tell a metering component from one that silently returned + None, which is how an unlicensed component looks healthy while exporting + nothing. + """ + _set_full_env(monkeypatch, tmp_path) + monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000") + monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class({})) + + infos: List[str] = [] + monkeypatch.setattr(bm.verbose_proxy_logger, "info", lambda msg, *args: infos.append(msg % args if args else msg)) + + recorder = bm.build_billing_metrics_recorder(premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0") + + assert recorder is not None + joined = "\n".join(infos) + assert "https://collector.example:4317" in joined + assert "5000" in joined + + +def test_unlicensed_build_does_not_warn(monkeypatch, tmp_path): + """Unlicensed is the common OSS case; warning there would be pure noise.""" + _set_full_env(monkeypatch, tmp_path) + + warnings: List[str] = [] + monkeypatch.setattr(bm.verbose_proxy_logger, "warning", lambda msg, *args: warnings.append(str(msg))) + + assert bm.build_billing_metrics_recorder(premium=False, license_data=None, litellm_version="1.0") is None + assert warnings == [] + + +def test_shutdown_flushes_active_recorder_once(monkeypatch, tmp_path): + """The shutdown hook must flush the recorder the factory built (buffered + counts are lost on restart otherwise) and be idempotent for repeat calls.""" + _set_full_env(monkeypatch, tmp_path) + shutdowns = [] + + class _SpyProvider: + def get_meter(self, name): + return MeterProvider().get_meter(name) + + def shutdown(self, timeout_millis=None): + shutdowns.append(timeout_millis) + + monkeypatch.setattr(bm, "build_mtls_meter_provider", lambda config: _SpyProvider()) + recorder = bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") + assert recorder is not None + + bm.shutdown_billing_metrics_recorder() + bm.shutdown_billing_metrics_recorder() + assert shutdowns == [bm.SHUTDOWN_FLUSH_TIMEOUT_MS] + + +def test_shutdown_without_active_recorder_is_noop(): + bm.shutdown_billing_metrics_recorder() + + +# ── Config loading ──────────────────────────────────────────────────────────── + + +def test_load_config_carries_license_id(monkeypatch, tmp_path): + _set_full_env(monkeypatch, tmp_path) + config = bm.load_billing_metrics_config(license_data={"user_id": "org-42"}, litellm_version="9.9") + assert config is not None and config.license_id == "org-42" and config.litellm_version == "9.9" + + +def test_load_config_with_empty_string_env_is_disabled(monkeypatch, tmp_path): + """An env var set to the empty string is as unusable as an unset one and + must disable metering rather than produce a config with a blank endpoint.""" + paths = _write_certs(tmp_path) + monkeypatch.setenv(bm.ENDPOINT_ENV, "") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + + +_CLIENT_CERT_PEM = "-----BEGIN CERTIFICATE-----\nclient-cert-body\n-----END CERTIFICATE-----" +_CLIENT_KEY_PEM = "-----BEGIN PRIVATE KEY-----\nclient-key-body\n-----END PRIVATE KEY-----" +_CA_CERT_PEM = "-----BEGIN CERTIFICATE-----\nca-body\n-----END CERTIFICATE-----" + + +def test_load_config_materializes_inline_pem_content(monkeypatch): + """ + ECS and Cloud Run inject secrets as env content, not as mounted files, so the + cert env vars must accept PEM directly. The exporter takes paths, so the PEM + is written to disk and the config points at those files. + """ + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) + monkeypatch.setenv(bm.CA_CERT_ENV, _CA_CERT_PEM) + + config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + + assert config is not None + assert config.ca_cert_path is not None + written = { + config.client_cert_path: _CLIENT_CERT_PEM, + config.client_key_path: _CLIENT_KEY_PEM, + config.ca_cert_path: _CA_CERT_PEM, + } + for path, pem in written.items(): + assert path != pem, "config must carry a file path, not the PEM itself" + assert os.path.isfile(path) + assert Path(path).read_text(encoding="utf-8") == f"{pem}\n" + + # The private key must not be world- or group-readable. + assert stat.S_IMODE(os.stat(config.client_key_path).st_mode) == 0o600 + + +def test_load_config_accepts_a_mix_of_pem_content_and_file_paths(monkeypatch, tmp_path): + """A deployment may mount the CA but inject the client credentials inline.""" + paths = _write_certs(tmp_path) + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) + monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV]) + + config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + + assert config is not None + assert config.ca_cert_path == paths[bm.CA_CERT_ENV] + assert Path(config.client_cert_path).read_text(encoding="utf-8") == f"{_CLIENT_CERT_PEM}\n" + + +def test_load_config_leaves_file_paths_untouched(monkeypatch, tmp_path): + """Path-valued env vars keep working; nothing is copied or rewritten.""" + paths = _set_full_env(monkeypatch, tmp_path) + + config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + + assert config is not None + assert config.client_cert_path == paths[bm.CLIENT_CERT_ENV] + assert config.client_key_path == paths[bm.CLIENT_KEY_ENV] + assert config.ca_cert_path == paths[bm.CA_CERT_ENV] + + +def test_load_config_with_inline_pem_disabled_when_unwritable(monkeypatch): + """A failure to materialize the PEM disables metering instead of raising.""" + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM) + + def _explode(prefix=None): + raise OSError("read-only filesystem") + + monkeypatch.setattr(bm.tempfile, "mkdtemp", _explode) + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + + +def test_load_config_never_logs_credential_values(monkeypatch): + """ + A value that is neither a readable path nor `-----BEGIN`-prefixed PEM is + still secret material. The disable warning must name the env vars, never + echo their contents, or a malformed key lands in the proxy logs. + """ + secret_material = "MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQ-not-pem-prefixed" + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, secret_material) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, secret_material) + + logged: List[str] = [] + + def _capture(msg, *args): + logged.append(msg % args if args else msg) + + monkeypatch.setattr(bm.verbose_proxy_logger, "warning", _capture) + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + + joined = "\n".join(logged) + assert secret_material not in joined + assert bm.CLIENT_CERT_ENV in joined and bm.CLIENT_KEY_ENV in joined + + +def test_load_config_with_empty_pem_env_is_disabled(monkeypatch): + """Empty stays empty: an unset secret must not be mistaken for inline PEM.""" + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, "") + monkeypatch.setenv(bm.CLIENT_KEY_ENV, "") + + assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None + + +def test_export_interval_default_and_override(monkeypatch): + assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS + monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000") + assert bm._export_interval_ms() == 5000 + monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "not-a-number") + assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS + + +# ── OTLP/HTTP exporter wiring ───────────────────────────────────────────────── + + +def test_metrics_endpoint_appends_signal_path(): + assert bm._metrics_endpoint("https://telemetry.example.com") == "https://telemetry.example.com/v1/metrics" + assert bm._metrics_endpoint("https://telemetry.example.com/") == "https://telemetry.example.com/v1/metrics" + assert bm._metrics_endpoint("https://telemetry.example.com/v1/metrics") == "https://telemetry.example.com/v1/metrics" + + +def _fake_exporter_class(captured: Dict[str, object]) -> type: + """A no-network stand-in for OTLPMetricExporter. Tests that build a real + MeterProvider must install this: the real exporter resolves the collector + host and opens a TLS connection on the reader's first export and on the + shutdown flush.""" + + class _FakeExporter: + # PeriodicExportingMetricReader probes these on the exporter it wraps. + _preferred_temporality: dict = {} + _preferred_aggregation: dict = {} + + def __init__(self, **kwargs): + captured.update(kwargs) + + def export(self, *args, **kwargs): + return None + + def shutdown(self, *args, **kwargs): + return None + + def force_flush(self, *args, **kwargs): + return True + + return _FakeExporter + + +def test_meter_provider_wires_client_cert_into_http_exporter(tmp_path, monkeypatch): + """Client cert+key authenticate us at the collector's mTLS front end; CA override rides certificate_file.""" + captured: Dict[str, object] = {} + + monkeypatch.setattr(bm, "OTLPMetricExporter", _fake_exporter_class(captured)) + config = _config(tmp_path) + provider = bm.build_mtls_meter_provider(config) + provider.shutdown() + + assert captured["endpoint"] == "https://collector.example:4317/v1/metrics" + assert captured["client_certificate_file"] == config.client_cert_path + assert captured["client_key_file"] == config.client_key_path + assert captured["certificate_file"] == config.ca_cert_path + + +def test_load_config_without_ca_is_valid(monkeypatch, tmp_path): + """The production collector presents a public web-PKI cert: no CA override required.""" + paths = _write_certs(tmp_path) + monkeypatch.setenv(bm.ENDPOINT_ENV, "https://telemetry.example.com") + monkeypatch.setenv(bm.CLIENT_CERT_ENV, paths[bm.CLIENT_CERT_ENV]) + monkeypatch.setenv(bm.CLIENT_KEY_ENV, paths[bm.CLIENT_KEY_ENV]) + config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") + assert config is not None and config.ca_cert_path is None + + +# ── Resource and metric attributes ──────────────────────────────────────────── + + +def test_resource_attributes_include_license_id(tmp_path): + attrs = bm._resource_attributes(_config(tmp_path, license_id="org-7")) + assert attrs["service.name"] == "litellm-proxy" + assert attrs["litellm.version"] == "1.2.3" + assert attrs["litellm.license.id"] == "org-7" + + +def test_resource_attributes_omit_license_id_when_absent(tmp_path): + attrs = bm._resource_attributes(_config(tmp_path, license_id=None)) + assert "litellm.license.id" not in attrs + + +def test_billable_attributes_with_model_id(): + attrs = bm._billable_attributes(BillableCategory.LLM, "/chat/completions", 200, "deploy-3") + assert attrs == { + "litellm.endpoint.category": "llm", + "http.route": "/chat/completions", + "http.response.status_code": 200, + "litellm.model_id": "deploy-3", + } + + +def test_billable_attributes_omit_model_id_when_none(): + attrs = bm._billable_attributes(BillableCategory.MCP, "/mcp", 200, None) + assert "litellm.model_id" not in attrs + + +# ── End-to-end recording via in-memory reader ───────────────────────────────── + + +def _counter_points(reader: InMemoryMetricReader): + data = reader.get_metrics_data() + for resource_metric in data.resource_metrics: + for scope_metric in resource_metric.scope_metrics: + for metric in scope_metric.metrics: + if metric.name == bm.METRIC_NAME: + return list(metric.data.data_points) + return [] + + +def test_record_increments_counter_with_attributes(): + reader = InMemoryMetricReader() + recorder = bm.BillingMetricsRecorder(MeterProvider(metric_readers=[reader])) + + recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id="m1") + recorder.record(category=BillableCategory.LLM, route="/chat/completions", status_code=200, model_id="m1") + recorder.record(category=BillableCategory.MCP, route="/mcp", status_code=200, model_id=None) + + points = _counter_points(reader) + by_category = {point.attributes["litellm.endpoint.category"]: point.value for point in points} + assert by_category["llm"] == 2 + assert by_category["mcp"] == 1 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py new file mode 100644 index 00000000000..f6f29eee5bc --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_compresr.py @@ -0,0 +1,2091 @@ +""" +Unit tests for the Compresr guardrail. + +Tests cover: +- apply_guardrail compresses eligible messages query-aware (tool-call intent + resolved via tool_call_id, falling back to the last user message) +- target selection: tool outputs by default, system/history opt-in, min-chars + threshold, targets without a derivable query are left uncompressed +- multimodal content: text parts replaced, non-text parts preserved +- recovery: hash marker appended, compresr_retrieve tool injected, originals + stored per litellm_call_id, agentic loop returns the original content and + rejects hashes not issued for the current request +- x-compresr-bypass header, response-type passthrough +- fail_closed raises HTTPException; fail_open forwards uncompressed +""" + +import hashlib +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, create_autospec, patch + +import httpx +import pytest +from fastapi import HTTPException + +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import ( + COMPRESR_RETRIEVE_TOOL_NAME, + CompresrGuardrail, + _content_hash, + _extract_compresr_tool_calls, + _scoped_store_key, + has_compresr_retrieve_tool, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +FAKE_API_BASE = "https://compresr.example.com" +FAKE_API_KEY = "cmp_test-key" + +TOOL_OUTPUT = "Result 1: EV range comparison. " * 40 # > 500 chars +USER_QUESTION = "Which 2026 EV has the longest range?" + +AGENT_MESSAGES = [ + {"role": "system", "content": "You are a research assistant."}, + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "2026 EV range"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TOOL_OUTPUT}, +] + + +def _make_guardrail(**kwargs) -> CompresrGuardrail: + defaults = dict( + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + guardrail_name="compresr", + default_on=True, + ) + defaults.update(kwargs) + return CompresrGuardrail(**defaults) + + +def _make_single_compress_response( + compressed_context: str = "compressed summary", + original_tokens: int = 1000, + compressed_tokens: int = 400, + status: int = 200, +) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = { + "success": True, + "data": { + "compressed_context": compressed_context, + "original_tokens": original_tokens, + "compressed_tokens": compressed_tokens, + "actual_compression_ratio": 0.6, + "tokens_saved": original_tokens - compressed_tokens, + "duration_ms": 42, + }, + } + mock.text = "" + return mock + + +def _make_batch_compress_response(compressed_contexts: list, status: int = 200) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = { + "success": True, + "data": { + "results": [ + { + "compressed_context": ctx, + "original_tokens": 1000, + "compressed_tokens": 400, + "actual_compression_ratio": 0.6, + "tokens_saved": 600, + "duration_ms": 42, + } + for ctx in compressed_contexts + ], + "count": len(compressed_contexts), + }, + } + mock.text = "" + return mock + + +def _make_openai_response_with_tool_call(tool_name: str, arguments: dict, tool_id: str = "call_abc123") -> MagicMock: + fn = MagicMock() + fn.name = tool_name + fn.arguments = json.dumps(arguments) + + tc = MagicMock() + tc.id = tool_id + tc.type = "function" + tc.function = fn + + message = MagicMock() + message.content = None + message.tool_calls = [tc] + + choice = MagicMock() + choice.message = message + + response = MagicMock() + response.choices = [choice] + # Plain chat-completion shape: no responses-API `output` list, no + # anthropic `content` list. + response.output = None + response.content = None + return response + + +def _make_openai_response_with_tool_calls(tool_calls: list, content: object = None) -> MagicMock: + """Chat-completion response carrying several tool calls in one turn + (parallel tool calling). ``tool_calls`` items are (name, arguments, id).""" + tcs = [] + for name, arguments, tool_id in tool_calls: + fn = MagicMock() + fn.name = name + fn.arguments = json.dumps(arguments) + tc = MagicMock() + tc.id = tool_id + tc.type = "function" + tc.function = fn + tcs.append(tc) + + message = MagicMock() + message.content = content + message.tool_calls = tcs + + choice = MagicMock() + choice.message = message + + response = MagicMock() + response.choices = [choice] + response.output = None + response.content = None + return response + + +def _apply_inputs(messages: list) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs(structured_messages=[dict(m) for m in messages]) + + +def _logging_obj(call_id: str) -> SimpleNamespace: + # Default fixture models a proxy with per-key auth enabled (the production + # shape). Recovery requires a caller scope; tests that need the no-auth + # path should build the object explicitly. + from litellm.proxy._types import UserAPIKeyAuth + + return SimpleNamespace( + litellm_call_id=call_id, + model_call_details={ + "litellm_params": {"metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="hash-default")}} + }, + ) + + +def _logging_obj_with_key(call_id: str, user_api_key: str, meta_key: str = "metadata") -> SimpleNamespace: + """Logging object carrying the server-set UserAPIKeyAuth object, the way the + proxy populates it for an authenticated request (the bare user_api_key + string alone is never trusted — a client could forge that).""" + from litellm.proxy._types import UserAPIKeyAuth + + return SimpleNamespace( + litellm_call_id=call_id, + model_call_details={"litellm_params": {meta_key: {"user_api_key_auth": UserAPIKeyAuth(api_key=user_api_key)}}}, + ) + + +def _retrieve_tool_call(hash_value: str, tool_id: str) -> dict: + return { + "id": tool_id, + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + + +def _retrieve_tool_stub() -> dict: + return { + "type": "function", + "function": {"name": COMPRESR_RETRIEVE_TOOL_NAME, "parameters": {}}, + } + + +@pytest.fixture +def guardrail() -> CompresrGuardrail: + return _make_guardrail() + + +# ── init ────────────────────────────────────────────────────────────── + + +def test_init_raises_without_api_key(monkeypatch): + monkeypatch.delenv("COMPRESR_API_KEY", raising=False) + with pytest.raises(ValueError, match="API key"): + CompresrGuardrail(guardrail_name="compresr") + + +def test_init_defaults(): + g = _make_guardrail() + assert g.compresr_api_base == FAKE_API_BASE + assert g.compression_model == "latte_v2" + assert g.target_compression_ratio == 0.5 + assert g.coarse is True + assert g.min_chars_to_compress == 500 + assert g.compress_tool_outputs is True + assert g.compress_system is False + assert g.compress_history is False + assert g.compress_last_user is False + assert g.enable_retrieval is True + assert g.unreachable_fallback == "fail_closed" + + +def test_init_coerces_unknown_unreachable_fallback_to_fail_closed(): + g = _make_guardrail(unreachable_fallback="banana") + assert g.unreachable_fallback == "fail_closed" + + +# ── compression core ───────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_apply_guardrail_compresses_tool_output_with_intent_query( + guardrail: CompresrGuardrail, +): + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + _, call_kwargs = mock_post.call_args + assert call_kwargs["url"] == f"{FAKE_API_BASE}/api/compress/question-specific/" + assert call_kwargs["headers"]["X-API-Key"] == FAKE_API_KEY + payload = call_kwargs["json"] + assert payload["context"] == TOOL_OUTPUT + # Query is the tool call's intent, not the user question. + assert payload["query"] == 'web_search: {"query": "2026 EV range"}' + assert payload["compression_model_name"] == "latte_v2" + assert payload["target_compression_ratio"] == 0.5 + + out = result["structured_messages"] + assert out[3]["content"].startswith("compressed summary") + # Untouched messages pass through byte-identical. + assert out[0] == AGENT_MESSAGES[0] + assert out[1] == AGENT_MESSAGES[1] + assert out[2] == AGENT_MESSAGES[2] + + +@pytest.mark.asyncio +async def test_apply_guardrail_mirrors_compression_into_texts_channel( + guardrail: CompresrGuardrail, +): + """The /v1/responses translation writes compressed output back through the + `texts` channel, not structured_messages. Compression must be mirrored there + or that surface silently forwards the original content uncompressed.""" + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_unknown", "content": TOOL_OUTPUT}, + ] + inputs = GenericGuardrailAPIInputs( + texts=[USER_QUESTION, TOOL_OUTPUT], + structured_messages=[dict(m) for m in messages], + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + texts = result["texts"] + # The compressed tool output replaces the original in the texts channel... + assert texts[1].startswith("compressed summary") + assert texts[1] != TOOL_OUTPUT + # ...while untouched text passes through byte-identical. + assert texts[0] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_apply_guardrail_returns_inputs_unchanged_when_nothing_compressed( + guardrail: CompresrGuardrail, +): + """A 200 response whose compressed_context is empty is a functional no-op. + The exact inputs object must come back: handlers detect guardrail edits by + identity, and a fresh structured_messages list would force a full write-back + of an untouched request (on Anthropic, reconversion strips cache_control + from thinking blocks).""" + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response(compressed_context="")) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + assert result is inputs + + +@pytest.mark.asyncio +async def test_texts_mirror_skips_duplicate_content_with_diverging_compressions( + guardrail: CompresrGuardrail, +): + """Two targets with identical text but different query-specific compressions: + the value-keyed texts mirror cannot tell the occurrences apart, so it must + leave them uncompressed rather than apply an arbitrary variant to both.""" + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "search_docs", "arguments": '{"q": "a"}'}}, + {"id": "call_2", "type": "function", "function": {"name": "search_web", "arguments": '{"q": "b"}'}}, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TOOL_OUTPUT}, + {"role": "tool", "tool_call_id": "call_2", "content": TOOL_OUTPUT}, + ] + inputs = GenericGuardrailAPIInputs( + texts=[USER_QUESTION, TOOL_OUTPUT, TOOL_OUTPUT], + structured_messages=[dict(m) for m in messages], + ) + mock_post = AsyncMock(return_value=_make_batch_compress_response(["compressed for docs", "compressed for web"])) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + # Each message position still gets its own query-specific compression... + out = result["structured_messages"] + assert out[2]["content"].startswith("compressed for docs") + assert out[3]["content"].startswith("compressed for web") + # ...but the texts mirror leaves the ambiguous occurrences untouched. + assert result["texts"] == [USER_QUESTION, TOOL_OUTPUT, TOOL_OUTPUT] + + +@pytest.mark.asyncio +async def test_texts_mirror_skips_text_that_also_appears_outside_targets( + guardrail: CompresrGuardrail, +): + """compress_system is off, so a system message whose text happens to equal + a compressed tool output must not be rewritten through the texts mirror.""" + messages = [ + {"role": "system", "content": TOOL_OUTPUT}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "content": TOOL_OUTPUT}, + ] + inputs = GenericGuardrailAPIInputs( + texts=[TOOL_OUTPUT, USER_QUESTION, TOOL_OUTPUT], + structured_messages=[dict(m) for m in messages], + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + out = result["structured_messages"] + assert out[0]["content"] == TOOL_OUTPUT # system message untouched + assert out[2]["content"].startswith("compressed summary") + # One compressed target cannot account for two occurrences in texts. + assert result["texts"] == [TOOL_OUTPUT, USER_QUESTION, TOOL_OUTPUT] + + +@pytest.mark.asyncio +async def test_tool_output_without_matching_call_uses_user_question( + guardrail: CompresrGuardrail, +): + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_unknown", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["query"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_function_result_without_name_does_not_bind_unrelated_call( + guardrail: CompresrGuardrail, +): + """A legacy function-role result missing its name must not adopt the intent + of an arbitrary earlier assistant function_call; it falls back to the last + user message.""" + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "function_call": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + {"role": "function", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["query"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_target_without_derivable_query_left_uncompressed( + guardrail: CompresrGuardrail, +): + # No user message and no tool-call intent anywhere -> nothing to compress. + messages = [{"role": "tool", "tool_call_id": "call_x", "content": TOOL_OUTPUT}] + mock_post = AsyncMock() + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + mock_post.assert_not_called() + assert result["structured_messages"][0]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_system_and_history_not_compressed_by_default( + guardrail: CompresrGuardrail, +): + long_system = "Rules. " * 200 + messages = [ + {"role": "system", "content": long_system}, + {"role": "user", "content": "Old question? " * 100}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + # Only one (single, non-batch) call: the tool output. + assert mock_post.call_count == 1 + assert mock_post.call_args.kwargs["json"]["context"] == TOOL_OUTPUT + out = result["structured_messages"] + assert out[0]["content"] == long_system + assert out[2]["content"] == USER_QUESTION + + +@pytest.mark.asyncio +async def test_opt_in_system_uses_batch_endpoint(): + guardrail = _make_guardrail(compress_system=True) + long_system = "Rules. " * 200 + messages = [ + {"role": "system", "content": long_system}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_batch_compress_response(["short system", "short tool"])) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"].endswith("/api/compress/question-specific/batch") + batch_inputs = call_kwargs["json"]["inputs"] + assert [i["context"] for i in batch_inputs] == [long_system, TOOL_OUTPUT] + out = result["structured_messages"] + assert out[0]["content"].startswith("short system") + assert out[2]["content"].startswith("short tool") + + +@pytest.mark.asyncio +async def test_short_messages_skipped(guardrail: CompresrGuardrail): + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": "tiny result"}, + ] + mock_post = AsyncMock() + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + mock_post.assert_not_called() + assert result["structured_messages"][1]["content"] == "tiny result" + + +@pytest.mark.asyncio +async def test_multimodal_text_replaced_non_text_preserved( + guardrail: CompresrGuardrail, +): + image_part = {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}} + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "tool", + "tool_call_id": "c1", + "content": [{"type": "text", "text": TOOL_OUTPUT}, image_part], + }, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + content = result["structured_messages"][1]["content"] + assert isinstance(content, list) + assert content[0]["type"] == "text" + assert content[0]["text"].startswith("compressed summary") + assert content[1] == image_part + + +# ── passthrough / bypass ───────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_bypass_header_skips_compression_when_allowed(): + guardrail = _make_guardrail(allow_bypass_header=True) + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock() + request_data = { + "model": "gpt-4o", + "proxy_server_request": {"headers": {"x-compresr-bypass": "true"}}, + } + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") + + mock_post.assert_not_called() + assert result is inputs + + +@pytest.mark.asyncio +async def test_bypass_header_ignored_by_default(guardrail: CompresrGuardrail): + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + request_data = { + "model": "gpt-4o", + "proxy_server_request": {"headers": {"x-compresr-bypass": "true"}}, + } + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") + + mock_post.assert_called_once() + + +@pytest.mark.asyncio +async def test_response_input_type_passthrough(guardrail: CompresrGuardrail): + inputs = _apply_inputs(AGENT_MESSAGES) + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response") + assert result is inputs + + +@pytest.mark.asyncio +async def test_missing_structured_messages_passthrough(guardrail: CompresrGuardrail): + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + assert result is inputs + + +# ── failure policy ──────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_transport_error_raises_when_fail_closed(guardrail: CompresrGuardrail): + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(side_effect=httpx.ConnectError("boom")), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_transport_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = _apply_inputs(AGENT_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(side_effect=httpx.ConnectError("boom")), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_non_json_response_raises_when_fail_closed(guardrail: CompresrGuardrail): + mock = MagicMock() + mock.status_code = 200 + mock.json.side_effect = ValueError("not json") + mock.text = "gateway error" + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock)): + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_http_exception_does_not_reflect_upstream_body(guardrail: CompresrGuardrail): + mock = MagicMock() + mock.status_code = 500 + mock.json.side_effect = ValueError("not json") + mock.text = "SECRET_INSTANCE_METADATA_TOKEN=aws-imds-response" + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert "SECRET_INSTANCE_METADATA_TOKEN" not in json.dumps(exc_info.value.detail) + + +def test_init_rejects_non_http_api_base(): + with pytest.raises(ValueError, match="scheme"): + CompresrGuardrail( + api_base="file:///etc/passwd", + api_key=FAKE_API_KEY, + guardrail_name="compresr", + ) + + +def test_init_rejects_cloud_metadata_api_base(): + with pytest.raises(ValueError, match="metadata"): + CompresrGuardrail( + api_base="http://169.254.169.254", + api_key=FAKE_API_KEY, + guardrail_name="compresr", + ) + + +@pytest.mark.parametrize( + "api_base", + [ + "http://2852039166", # decimal encoding of 169.254.169.254 + "http://0xa9fea9fe", # hex encoding + "http://[::ffff:169.254.169.254]", # IPv4-mapped IPv6 + "http://metadata.azure.com", + "http://metadata.azure.internal", + "http://168.63.129.16", # Azure WireServer + ], +) +def test_init_rejects_encoded_cloud_metadata_api_base(api_base): + with pytest.raises(ValueError, match="metadata"): + CompresrGuardrail( + api_base=api_base, + api_key=FAKE_API_KEY, + guardrail_name="compresr", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_ignores_user_supplied_call_id(guardrail: CompresrGuardrail): + mock_post = AsyncMock(return_value=_make_single_compress_response()) + attacker_call_id = "victim-tenant-call-id" + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o", "litellm_call_id": attacker_call_id}, + input_type="request", + logging_obj=_logging_obj("real-framework-call-id"), + ) + + assert not any(attacker_call_id in k for k in guardrail._originals_by_call_id) + assert any(k.endswith("real-framework-call-id") for k in guardrail._originals_by_call_id) + + +@pytest.mark.asyncio +async def test_agentic_plan_ignores_user_supplied_call_id(guardrail: CompresrGuardrail): + hash_value = "d" * 24 + guardrail._store_originals("victim-tenant-call-id", {hash_value: "victim-original"}) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_abc" + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("attacker-call-id"), + stream=False, + kwargs={"litellm_call_id": "victim-tenant-call-id"}, + ) + + # Attacker's scope resolves nothing, so the loop is vetoed and the victim + # original never surfaces. + assert plan.run_agentic_loop is False + assert plan.request_patch is None + + +@pytest.mark.asyncio +async def test_recovery_store_partitioned_by_caller_identity(guardrail: CompresrGuardrail): + """Two tenants that set the SAME client-forgeable x-litellm-call-id must not + read each other's stored originals, and each still reads its own.""" + shared_call_id = "shared-call-id" + expected_hash = hashlib.sha256(TOOL_OUTPUT.encode()).hexdigest()[:24] + + async def _plan_for(user_api_key: str, tool_id: str): + return await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(expected_hash, tool_id)]}, + model="gpt-4o", + messages=[], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": expected_hash}, tool_id=tool_id + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj_with_key(shared_call_id, user_api_key), + stream=False, + kwargs={}, + ) + + # Tenant A compresses and stores its original. + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=_make_single_compress_response())): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj_with_key(shared_call_id, "hash-tenant-A"), + ) + + # Tenant B, same call id, different virtual-key hash → different bucket, so + # nothing resolves and the loop is vetoed (Tenant A's original never leaks). + plan_b = await _plan_for("hash-tenant-B", "call_b") + assert plan_b.run_agentic_loop is False + assert plan_b.request_patch is None + + # Tenant A retrieves its own content successfully. + plan_a = await _plan_for("hash-tenant-A", "call_a") + assert TOOL_OUTPUT in plan_a.request_patch.messages[-1]["content"] + + +@pytest.mark.asyncio +async def test_caller_scope_read_from_litellm_metadata(guardrail: CompresrGuardrail): + """/v1/messages and /v1/responses carry the auth object under + litellm_metadata rather than metadata; the store key must be scoped by it + there too, without relying on upstream's metadata backfill.""" + logging_obj = _logging_obj_with_key("call-lm", "hash-tenant-lm", meta_key="litellm_metadata") + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + + assert "hash-tenant-lm\x00call-lm" in guardrail._originals_by_call_id + assert "call-lm" not in guardrail._originals_by_call_id + + +@pytest.mark.asyncio +async def test_caller_scope_rejects_forged_user_api_key_string(guardrail: CompresrGuardrail): + """A client-supplied metadata.user_api_key STRING (no server-set + UserAPIKeyAuth object) must not be trusted as a tenant scope — otherwise a + caller could forge another tenant's recovery bucket on /v1/messages.""" + logging_obj = SimpleNamespace( + litellm_call_id="call-forge", + model_call_details={"litellm_params": {"metadata": {"user_api_key": "victim-tenant-hash"}}}, + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + + # Forged string is ignored: scope resolves to empty, so recovery is + # disabled entirely (no bucket keyed on victim-tenant-hash, no unscoped + # bucket that another caller could reuse). + assert not guardrail._originals_by_call_id + + +@pytest.mark.asyncio +async def test_compress_post_called_with_real_handler_signature(): + """AsyncHTTPHandler.post has a fixed signature; an autospec mock enforces it + (unlike AsyncMock(spec=...), which silently accepts any kwarg) so a kwarg the + real handler rejects — which would raise TypeError past the fail policy — + fails the test instead.""" + guardrail = _make_guardrail() + autospec_post = create_autospec(guardrail.async_handler.post, return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", autospec_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + assert result["structured_messages"][3]["content"].startswith("compressed summary") + + +@pytest.mark.asyncio +async def test_batch_result_count_mismatch_raises_when_fail_closed(): + guardrail = _make_guardrail(compress_system=True) + messages = [ + {"role": "system", "content": "Rules. " * 200}, + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "c1", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_batch_compress_response(["only one"])) + + with patch.object(guardrail.async_handler, "post", mock_post): + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +# ── recovery (compresr_retrieve) ────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_recovery_marker_tool_injection_and_original_stored( + guardrail: CompresrGuardrail, +): + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + compressed_content = result["structured_messages"][3]["content"] + expected_hash = hashlib.sha256(TOOL_OUTPUT.encode()).hexdigest()[:24] + assert f"compresr hash={expected_hash}" in compressed_content + + tools = result.get("tools") + assert tools is not None and has_compresr_retrieve_tool(tools) + + scoped_key = next(k for k in guardrail._originals_by_call_id if k.endswith("call-id-1")) + originals, _expiry = guardrail._originals_by_call_id[scoped_key] + assert originals[expected_hash] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_enable_retrieval_false_no_marker_no_tool(): + guardrail = _make_guardrail(enable_retrieval=False) + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result["structured_messages"][3]["content"] == "compressed summary" + assert not has_compresr_retrieve_tool(result.get("tools") or []) + assert guardrail._originals_by_call_id == {} + + +@pytest.mark.asyncio +async def test_existing_tools_preserved_when_injecting(guardrail: CompresrGuardrail): + existing_tool = {"type": "function", "function": {"name": "my_tool", "parameters": {}}} + inputs = GenericGuardrailAPIInputs( + structured_messages=[dict(m) for m in AGENT_MESSAGES], + tools=[existing_tool], + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + tools = result["tools"] + assert existing_tool in tools + assert has_compresr_retrieve_tool(tools) + assert len(tools) == 2 + + +@pytest.mark.asyncio +async def test_non_list_tools_left_untouched_when_injecting(guardrail: CompresrGuardrail): + # An unexpected non-list tools value must survive unchanged rather than be + # clobbered by the injected retrieve tool. + odd_tools = {"type": "function", "function": {"name": "my_tool"}} + inputs = GenericGuardrailAPIInputs( + structured_messages=[dict(m) for m in AGENT_MESSAGES], + tools=odd_tools, # type: ignore[typeddict-item] + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + + assert result["tools"] is odd_tools + + +def test_extract_compresr_tool_calls_tolerates_missing_keys(): + # A retrieve call missing id/arguments must not KeyError in the post-call + # hook; it extracts with safe defaults and resolves to a rejection later. + with patch( + "litellm.proxy.guardrails.guardrail_hooks.compresr.compresr.get_tool_calls_from_response", + return_value=[{"name": COMPRESR_RETRIEVE_TOOL_NAME}, {"id": "x"}], + ): + extracted = _extract_compresr_tool_calls(object()) + + assert extracted == [{"id": None, "type": "function", "name": COMPRESR_RETRIEVE_TOOL_NAME, "arguments": {}}] + + +# ── agentic loop ────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_true_for_retrieve_call( + guardrail: CompresrGuardrail, +): + response = _make_openai_response_with_tool_call(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": "a" * 24}) + tools = [dict(t) for t in [_retrieve_tool_stub()]] + + should_run, gate_tools = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=tools, + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + assert should_run is True + assert gate_tools["tool_calls"][0]["name"] == COMPRESR_RETRIEVE_TOOL_NAME + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_false_without_retrieve_tool( + guardrail: CompresrGuardrail, +): + response = _make_openai_response_with_tool_call("other_tool", {"x": 1}) + should_run, _ = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + assert should_run is False + + +@pytest.mark.asyncio +async def test_agentic_plan_returns_stored_original(guardrail: CompresrGuardrail): + hash_value = hashlib.sha256(TOOL_OUTPUT.encode()).hexdigest()[:24] + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_abc" + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + assert plan.run_agentic_loop is True + follow_up = plan.request_patch.messages + tool_result = follow_up[-1] + assert tool_result["role"] == "tool" + assert tool_result["tool_call_id"] == "call_abc" + assert tool_result["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_agentic_plan_preserves_list_shaped_assistant_text(guardrail: CompresrGuardrail): + """Some providers return chat assistant content as list-of-parts; the + retrieval follow-up must keep that text, not drop it to None.""" + hash_value = "a" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: "original"}) + response = _make_openai_response_with_tool_calls( + [(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, "call_r")], + content=[{"type": "text", "text": "Let me fetch the original."}], + ) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_r")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + assistant_message = plan.request_patch.messages[-2] + assert assistant_message["role"] == "assistant" + assert assistant_message["content"] == "Let me fetch the original." + + +@pytest.mark.asyncio +async def test_agentic_plan_strips_other_guardrails_executed_markers(guardrail: CompresrGuardrail): + # The retrieval follow-up restores content other pre-call guardrails may + # never have inspected, so their executed markers must not be replayed; + # only this guardrail's own marker survives (no recompression loop). + hash_value = "a" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + response = _make_openai_response_with_tool_calls([(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, "call_r")]) + own_marker = guardrail._pre_call_marker() + assert own_marker is not None + kwargs = { + "metadata": { + "user_api_key": "key-hash", + PRE_CALL_EXECUTED_GUARDRAILS_KEY: [own_marker, "token:pii_guardrail"], + }, + "litellm_metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["token:other_guardrail"]}, + } + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_r")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs=kwargs, + ) + + out = plan.request_patch.kwargs + assert out["metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY] == [own_marker] + assert out["metadata"]["user_api_key"] == "key-hash" + assert PRE_CALL_EXECUTED_GUARDRAILS_KEY not in out["litellm_metadata"] + assert kwargs["metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY] == [own_marker, "token:pii_guardrail"] + assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY] == ["token:other_guardrail"] + + +@pytest.mark.asyncio +async def test_agentic_plan_rejects_hash_from_other_request( + guardrail: CompresrGuardrail, +): + hash_value = "b" * 24 + guardrail._store_originals("someone-elses-call", {hash_value: "secret"}) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[], + response=_make_openai_response_with_tool_call( + COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_abc" + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("my-call"), + stream=False, + kwargs={}, + ) + + # Hash belongs to another caller's scope; the loop is vetoed and the secret + # never surfaces. + assert plan.run_agentic_loop is False + assert plan.request_patch is None + + +@pytest.mark.asyncio +async def test_agentic_loop_vetoed_when_no_recovery_state(guardrail: CompresrGuardrail): + # A caller-defined compresr_retrieve tool with no stored original must not + # trigger an extra provider round-trip. + hash_value = "f" * 24 + response = _make_openai_response_with_tool_call(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, tool_id="call_x") + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_x")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + assert plan.run_agentic_loop is False + assert plan.request_patch is None + + +@pytest.mark.asyncio +async def test_agentic_loop_dedupes_repeated_retrievals(guardrail: CompresrGuardrail): + # Retrieving the same marker many times expands the original once; repeats + # get a short marker (no follow-up amplification). + hash_value = "a" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + calls = [(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, f"call_{i}") for i in range(5)] + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, f"call_{i}") for i in range(5)]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=_make_openai_response_with_tool_calls(calls), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + tool_results = [m for m in plan.request_patch.messages if m.get("role") == "tool"] + assert len(tool_results) == 5 + assert sum(1 for m in tool_results if m["content"] == TOOL_OUTPUT) == 1 + assert all("already retrieved" in m["content"] for m in tool_results if m["content"] != TOOL_OUTPUT) + + +@pytest.mark.asyncio +async def test_agentic_loop_caps_retrieval_count(guardrail: CompresrGuardrail): + # Beyond _MAX_RETRIEVALS_PER_LOOP retrievals, extra calls get a bounded marker. + n = 10 + hashes = [f"{i:024x}" for i in range(n)] + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {h: f"original-{h}" for h in hashes}) + calls = [(COMPRESR_RETRIEVE_TOOL_NAME, {"hash": h}, f"call_{i}") for i, h in enumerate(hashes)] + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(h, f"call_{i}") for i, h in enumerate(hashes)]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=_make_openai_response_with_tool_calls(calls), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + tool_results = [m for m in plan.request_patch.messages if m.get("role") == "tool"] + assert len(tool_results) == n + over_limit = [m for m in tool_results if "retrieval limit reached" in m["content"]] + assert len(over_limit) == n - 8 # only the first 8 expand + + +def test_display_hash_strips_control_characters(): + """The compresr_retrieve `hash` argument is model/tool-output-influenced, so + control characters (newlines, ANSI escapes) must be stripped — not just + length-capped — before it is echoed into logs or the fallback message.""" + from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import _display_hash + + assert _display_hash("a" * 24) == "a" * 24 # a real marker hash passes through + assert "\n" not in _display_hash("abc\ndef\rFORGED LOG LINE") + assert "\x1b" not in _display_hash("hash\x1b[31mred") + capped = _display_hash("z" * 100) + assert capped.endswith("…") and len(capped) <= 33 + + +@pytest.mark.asyncio +async def test_agentic_plan_builds_anthropic_followup_shape( + guardrail: CompresrGuardrail, +): + hash_value = "c" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = MagicMock() + response.output = None + response.content = [{"type": "tool_use", "id": "toolu_1"}] + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="claude-sonnet-5", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + follow_up = plan.request_patch.messages + assistant_msg, user_msg = follow_up[-2], follow_up[-1] + assert assistant_msg["role"] == "assistant" + assert assistant_msg["content"][0]["type"] == "tool_use" + assert user_msg["content"][0]["type"] == "tool_result" + assert user_msg["content"][0]["tool_use_id"] == "toolu_1" + assert user_msg["content"][0]["content"] == TOOL_OUTPUT + assert plan.request_patch.max_tokens == 1024 + + +@pytest.mark.asyncio +async def test_agentic_plan_builds_responses_followup_shape( + guardrail: CompresrGuardrail, +): + """The /v1/responses path echoes the function_call and pairs it with a + function_call_output keyed by the same call_id.""" + hash_value = "e" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = MagicMock() + response.output = [{"type": "function_call", "call_id": "fc_1"}] # responses-API shape + response.content = None + + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + call_item, output_item = plan.request_patch.messages[-2], plan.request_patch.messages[-1] + assert call_item["type"] == "function_call" + assert call_item["call_id"] == "fc_1" + assert output_item["type"] == "function_call_output" + assert output_item["call_id"] == "fc_1" + assert output_item["output"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_agentic_plan_chat_parallel_tool_calls_echoes_only_retrieve( + guardrail: CompresrGuardrail, +): + """When the model calls a real tool alongside compresr_retrieve in one turn, + only the retrieve call may be echoed in the reconstructed assistant message: + every echoed tool_call must have a matching tool result or the provider 400s. + The real call is re-planned by the follow-up; the assistant text is kept.""" + hash_value = "f" * 24 + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = _make_openai_response_with_tool_calls( + [ + ("get_weather", {"city": "Paris"}, "call_weather"), + (COMPRESR_RETRIEVE_TOOL_NAME, {"hash": hash_value}, "call_retrieve"), + ], + content="Let me expand that note and check the weather.", + ) + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "call_retrieve")]}, + model="gpt-4o", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + follow_up = plan.request_patch.messages + assistant_msg = follow_up[-2] + echoed_ids = {tc["id"] for tc in assistant_msg["tool_calls"]} + result_ids = {m["tool_call_id"] for m in follow_up if m.get("role") == "tool"} + # get_weather is not echoed; every echoed tool_call is answered. + assert echoed_ids == {"call_retrieve"} + assert echoed_ids == result_ids + assert assistant_msg["content"] == "Let me expand that note and check the weather." + assert follow_up[-1]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_agentic_plan_anthropic_parallel_preserves_text_and_balances( + guardrail: CompresrGuardrail, +): + """Anthropic parallel-tool-call turn: the assistant text is preserved, the + real tool_use is dropped (re-planned), and the reconstructed turn stays + balanced — one tool_result per echoed tool_use.""" + hash_value = "a" * 23 + "9" + guardrail._store_originals(_scoped_store_key(_logging_obj("call-id-1")), {hash_value: TOOL_OUTPUT}) + + response = MagicMock() + response.output = None + response.content = [ + {"type": "text", "text": "Checking the weather and expanding the note."}, + {"type": "tool_use", "id": "toolu_weather", "name": "get_weather", "input": {"city": "Paris"}}, + { + "type": "tool_use", + "id": "toolu_retrieve", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "input": {"hash": hash_value}, + }, + ] + + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [_retrieve_tool_call(hash_value, "toolu_retrieve")]}, + model="claude-sonnet-5", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"max_tokens": 1024}, + logging_obj=_logging_obj("call-id-1"), + stream=False, + kwargs={}, + ) + + assistant_msg, user_msg = plan.request_patch.messages[-2], plan.request_patch.messages[-1] + assert assistant_msg["content"][0] == { + "type": "text", + "text": "Checking the weather and expanding the note.", + } + echoed_ids = [b["id"] for b in assistant_msg["content"] if b["type"] == "tool_use"] + answered_ids = [b["tool_use_id"] for b in user_msg["content"]] + # get_weather dropped; balanced tool_use/tool_result pairing. + assert echoed_ids == ["toolu_retrieve"] + assert answered_ids == echoed_ids + + +# ── store hygiene ───────────────────────────────────────────────────── + + +def test_originals_store_prunes_expired(guardrail: CompresrGuardrail): + guardrail._originals_by_call_id["old"] = ({"a" * 24: "x"}, 0.0) # already expired + guardrail._store_originals("new", {"b" * 24: "y"}) + assert "old" not in guardrail._originals_by_call_id + assert "new" in guardrail._originals_by_call_id + + +def test_originals_store_caps_tracked_calls(guardrail: CompresrGuardrail): + for i in range(300): + guardrail._store_originals(f"call-{i}", {("%024x" % i): "x"}) + assert len(guardrail._originals_by_call_id) <= 256 + # Most recent entries survive. + assert "call-299" in guardrail._originals_by_call_id + + +def test_originals_store_caps_bytes_per_call(): + guardrail = _make_guardrail(max_bytes_per_call=1000) + hashes = tuple(f"{i:024x}" for i in range(5)) + values = tuple("x" * 400 for _ in range(5)) + guardrail._store_originals("c", dict(zip(hashes, values))) + + stored, _expiry = guardrail._originals_by_call_id["c"] + assert sum(len(v.encode("utf-8")) for v in stored.values()) <= 1000 + # Oldest entries are evicted first; newest survives. + assert hashes[-1] in stored + assert hashes[0] not in stored + + +def test_originals_store_byte_cap_survives_lone_surrogates(): + # Regression: eviction path must use surrogatepass to match the hash + # function; a bare encode("utf-8") crashed on lone surrogates. + guardrail = _make_guardrail(max_bytes_per_call=500) + surrogate_value = "\ud800" * 60 + hashes = tuple(f"{i:024x}" for i in range(3)) + guardrail._store_originals("c", dict(zip(hashes, (surrogate_value, surrogate_value, surrogate_value)))) + + stored, _expiry = guardrail._originals_by_call_id["c"] + assert hashes[-1] in stored + assert hashes[0] not in stored + + +def test_originals_store_caps_total_bytes_across_calls(monkeypatch: pytest.MonkeyPatch): + # Global byte budget: many distinct call ids must not retain unbounded memory. + monkeypatch.setattr( + "litellm.proxy.guardrails.guardrail_hooks.compresr.compresr._MAX_TOTAL_STORE_BYTES", + 10_000, + ) + guardrail = _make_guardrail(max_bytes_per_call=4_000) + for i in range(20): + guardrail._store_originals(f"call-{i}", {f"{i:024x}": "x" * 3_000}) + + total = sum( + len(v.encode("utf-8")) + for originals, _expiry in guardrail._originals_by_call_id.values() + for v in originals.values() + ) + assert total <= 10_000 + assert guardrail._store_total_bytes == total # running counter stays exact + # Oldest calls evicted; the most-recent call's originals survive. + assert "call-0" not in guardrail._originals_by_call_id + assert "call-19" in guardrail._originals_by_call_id + + +def test_originals_store_global_cap_keeps_current_when_single_call_is_large( + monkeypatch: pytest.MonkeyPatch, +): + # One call over the global cap is still kept (only max_bytes_per_call trims it); + # global eviction never empties the store. + monkeypatch.setattr( + "litellm.proxy.guardrails.guardrail_hooks.compresr.compresr._MAX_TOTAL_STORE_BYTES", + 1_000, + ) + guardrail = _make_guardrail(max_bytes_per_call=5_000) + guardrail._store_originals("solo", {f"{0:024x}": "x" * 4_000}) + assert "solo" in guardrail._originals_by_call_id + + +def test_recovery_markers_respect_per_call_byte_cap(): + # Regression: markers were built from every original before _store_originals + # applied the byte cap, so an evicted original left a dangling marker the + # model could never retrieve. Recovery must be attached only for originals + # that fit the cap, so every shipped marker stays retrievable. + guardrail = _make_guardrail(max_bytes_per_call=1000) + contexts = ["a" * 400, "b" * 400, "c" * 400] + messages = [{"role": "tool", "content": text} for text in contexts] + results = [{"compressed_context": f"small-{i}"} for i in range(3)] + + applied = guardrail._apply_compression_results( + messages, [0, 1, 2], contexts, results, recovery_enabled=True + ) + + # 400 + 400 fit under 1000; the third (which would reach 1200) is skipped. + assert applied.messages_compressed == 3 + assert len(applied.originals) == 2 + third_hash = _content_hash("c" * 400) + assert third_hash not in applied.originals + assert f"compresr hash={third_hash}" not in applied.compressed_messages[2]["content"] + + # Every marker still shipped must resolve to a stored original. + guardrail._store_originals("c", applied.originals) + for hash_value in applied.originals: + assert guardrail._retrieve_original("c", hash_value) is not None + assert f"compresr hash={hash_value}" in "".join( + str(m["content"]) for m in applied.compressed_messages + ) + + +def test_recovery_markers_respect_byte_cap_across_reused_store_key(): + # Regression: the per-call budget must also count bytes already stored under + # the same store key (a later turn reusing the call id). Otherwise merging + # this call's originals with the existing entry overflows the cap and + # _store_originals evicts an original this call just shipped a marker for. + guardrail = _make_guardrail(max_bytes_per_call=100) + old_hash = _content_hash("A" * 50) + guardrail._store_originals("k", {old_hash: "A" * 50}) + guardrail._store_originals("k", {_content_hash("C" * 40): "C" * 40}) + existing = guardrail._originals_by_call_id["k"][0] + + # This turn recompresses the same "A" (already stored) plus a new "D". + contexts = ["D" * 40, "A" * 50] + messages = [{"role": "tool", "content": text} for text in contexts] + results = [{"compressed_context": "dd"}, {"compressed_context": "aa"}] + applied = guardrail._apply_compression_results( + messages, [0, 1], contexts, results, recovery_enabled=True, existing_originals=existing + ) + + guardrail._store_originals("k", applied.originals) + # No marker shipped this turn may dangle after the store enforces the cap. + for hash_value in applied.originals: + assert guardrail._retrieve_original("k", hash_value) is not None + assert f"compresr hash={hash_value}" in "".join( + str(m["content"]) for m in applied.compressed_messages + ) + # The zero-cost repeat of an already-stored original stays retrievable. + assert old_hash in applied.originals + assert guardrail._retrieve_original("k", old_hash) is not None + + +# ── dynamic (adaptive) compression — latte_v2 Kneedle ───────────────── + + +@pytest.mark.asyncio +async def test_dynamic_flag_in_payload(): + """dynamic=True must appear in the compress payload; unset bounds omitted.""" + guardrail = _make_guardrail(dynamic=True) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["dynamic"] is True + assert "dynamic_min_ratio" not in payload + assert "dynamic_max_ratio" not in payload + + +@pytest.mark.asyncio +async def test_dynamic_bounds_in_payload_when_set(): + guardrail = _make_guardrail(dynamic=True, dynamic_min_ratio=2.0, dynamic_max_ratio=8.0) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["dynamic"] is True + assert payload["dynamic_min_ratio"] == 2.0 + assert payload["dynamic_max_ratio"] == 8.0 + + +@pytest.mark.asyncio +async def test_dynamic_on_by_default(): + guardrail = _make_guardrail() # dynamic defaults on (latte_v2) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert mock_post.call_args.kwargs["json"]["dynamic"] is True + + +# ── generic passthrough compression params ──────────────────────────── + + +@pytest.mark.asyncio +async def test_compression_params_passthrough_in_payload(): + """Extra params in compression_params are forwarded verbatim; named fields + still win on collision.""" + guardrail = _make_guardrail(compression_params={"heuristic_chunking": True, "coarse": False}) + messages = [ + {"role": "user", "content": USER_QUESTION}, + {"role": "tool", "tool_call_id": "call_x", "name": "search", "content": TOOL_OUTPUT}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + payload = mock_post.call_args.kwargs["json"] + assert payload["heuristic_chunking"] is True + # named `coarse` (default True) wins over the passthrough's coarse=False + assert payload["coarse"] is True + + +@pytest.mark.asyncio +async def test_compression_params_cannot_override_request_content_fields(): + """context/query/inputs carry the actual content being compressed; a + passthrough collision on them must be dropped, not silently win.""" + guardrail = _make_guardrail( + compression_params={ + "context": "injected", + "query": "injected", + "inputs": [], + "heuristic_chunking": True, + } + ) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["context"] == TOOL_OUTPUT + assert payload["query"] == 'web_search: {"query": "2026 EV range"}' + assert "inputs" not in payload + assert payload["heuristic_chunking"] is True + + +# ── compress_last_user ──────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_compress_last_user_compresses_with_verbatim_query(): + """compress_last_user=True compresses the last user message, but the query + sent to Compresr is still the original verbatim user text.""" + guardrail = _make_guardrail(compress_last_user=True) + long_question = "Which 2026 EV has the longest range? " * 20 # > 500 chars + messages = [{"role": "user", "content": long_question}] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + payload = mock_post.call_args.kwargs["json"] + assert payload["context"] == long_question + assert payload["query"] == long_question # verbatim, not the compressed text + assert result["structured_messages"][0]["content"] == "compressed summary" + + +# ── malformed-but-200 token stats (must not defeat fail policy) ─────── + + +@pytest.mark.asyncio +async def test_non_numeric_token_stats_do_not_raise(guardrail: CompresrGuardrail): + """A 200 response with non-numeric token counts must not raise: _call_compress + already succeeded, so a bare int() here would 500 even under fail policy.""" + resp = _make_single_compress_response() + resp.json.return_value["data"]["original_tokens"] = "not-a-number" + resp.json.return_value["data"]["compressed_tokens"] = None + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=resp)): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result["structured_messages"][3]["content"].startswith("compressed summary") + + +# ── HTTP status errors (non-2xx from the shared handler) ────────────── + + +def _http_status_error(status: int = 500, text: str = "upstream error body") -> httpx.HTTPStatusError: + # The shared AsyncHTTPHandler.post() raises HTTPStatusError on any non-2xx, + # carrying the upstream body and request headers; this simulates that. + request = httpx.Request("POST", f"{FAKE_API_BASE}/api/compress/question-specific/") + response = httpx.Response(status, text=text, request=request) + return httpx.HTTPStatusError(str(status), request=request, response=response) + + +@pytest.mark.asyncio +async def test_http_status_error_raises_when_fail_closed(guardrail: CompresrGuardrail): + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_http_status_error(500))): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_http_status_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_http_status_error(429))): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_http_status_error_does_not_leak_upstream_body(guardrail: CompresrGuardrail): + secret = "SECRET_INSTANCE_METADATA_TOKEN=aws-imds-response" + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_http_status_error(500, text=secret))): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert secret not in json.dumps(exc_info.value.detail) + + +# ── non-transport httpx errors must still honor the fail policy ──────── +# TooManyRedirects and DecodingError are httpx.RequestError but NOT +# httpx.TransportError, so a narrow except would let them escape as a 500 +# even under fail_open. These lock in that they are routed through the policy. + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + httpx.TooManyRedirects("redirect loop"), + httpx.DecodingError("bad content-encoding"), + ], +) +async def test_request_errors_raise_when_fail_closed(guardrail: CompresrGuardrail, error): + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=error)): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", + [ + httpx.TooManyRedirects("redirect loop"), + httpx.DecodingError("bad content-encoding"), + ], +) +async def test_request_errors_fail_open_forwards_uncompressed(error): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=error)): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_undecodable_body_on_200_forwards_uncompressed_when_fail_open(): + """A 200 whose body raises DecodingError on .json()/.text must not 500.""" + guardrail = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock() + resp.status_code = 200 + resp.json.side_effect = httpx.DecodingError("bad content-encoding") + type(resp).text = property(lambda self: (_ for _ in ()).throw(httpx.DecodingError("bad content-encoding"))) + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=resp)): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + assert result["structured_messages"][3]["content"] == TOOL_OUTPUT + + +@pytest.mark.asyncio +async def test_recursion_error_on_json_forwards_when_fail_open(): + """A deeply nested JSON body can raise RecursionError while parsing; it must + route through the fail policy, not escape as a 500.""" + guardrail = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock() + resp.status_code = 200 + resp.json.side_effect = RecursionError("maximum recursion depth exceeded") + resp.text = "" + inputs = _apply_inputs(AGENT_MESSAGES) + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=resp)): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_lone_surrogate_in_content_does_not_crash(guardrail: CompresrGuardrail): + """A lone Unicode surrogate (reachable via a JSON \\uXXXX escape) in content + must not crash hashing/byte-accounting after the fail-policy decision.""" + surrogate_output = ("x" * 600) + "\ud800" + messages = [ + {"role": "user", "content": USER_QUESTION}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "web_search", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": surrogate_output}, + ] + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(messages), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + assert result["structured_messages"][2]["content"].startswith("compressed summary") + # The original (surrogate included) is recoverable by its hash. + stored = next(iter(guardrail._originals_by_call_id.values()))[0] + assert surrogate_output in stored.values() + + +@pytest.mark.asyncio +async def test_identical_compressed_text_treated_as_noop(guardrail: CompresrGuardrail): + """If the service returns text byte-identical to the original, nothing + changed: the exact inputs object is returned so no write-back is forced.""" + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response(compressed_context=TOOL_OUTPUT)) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=_logging_obj("call-id-1"), + ) + assert result is inputs + + +# ── recovery requires a framework-issued call id ────────────────────── + + +@pytest.mark.asyncio +async def test_recovery_disabled_without_call_id(): + # enable_retrieval defaults True, but with no framework litellm_call_id we + # cannot scope stored originals to the request, so compression proceeds + # without markers, the retrieve tool, or any stored originals. + guardrail = _make_guardrail() + inputs = _apply_inputs(AGENT_MESSAGES) + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + assert result["structured_messages"][3]["content"] == "compressed summary" + assert not has_compresr_retrieve_tool(result.get("tools") or []) + assert guardrail._originals_by_call_id == {} + + +def test_config_model_exposes_unreachable_fallback(): + from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( + CompresrGuardrailConfigModel, + ) + + field = CompresrGuardrailConfigModel.model_fields.get("unreachable_fallback") + assert field is not None + assert field.default == "fail_closed" + + +# ── audit fixes ─────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_cancelled_error_propagates_not_swallowed(): + # Regression: CancelledError is a BaseException, not caught by + # (RequestError, Timeout). It must re-raise so cooperative cancellation + # (asyncio.wait_for, client disconnect) still fires. + import asyncio as _asyncio + + guardrail = _make_guardrail(unreachable_fallback="fail_open") + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=_asyncio.CancelledError())): + with pytest.raises(_asyncio.CancelledError): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + +def test_max_bytes_per_call_negative_rejected(): + # Regression: a negative value silently disabled the byte cap (< 0 behaves + # like 0 in _bound_call_bytes). Validate at construction so the footgun + # surfaces as a ValueError at startup, not silent unbounded storage. + with pytest.raises(ValueError, match="max_bytes_per_call"): + _make_guardrail(max_bytes_per_call=-1) + + +@pytest.mark.asyncio +async def test_max_tokens_zero_from_optional_params_wins_over_kwargs(): + # Regression: `or` short-circuits on falsy values, so an explicit + # max_tokens=0 from optional_params fell through to kwargs["max_tokens"]. + # Must use `is not None`. + guardrail = _make_guardrail() + hash_value = "deadbeef" + guardrail._store_originals(_scoped_store_key(_logging_obj("call-1")), {hash_value: TOOL_OUTPUT}) + response = MagicMock() + response.content = [{"type": "tool_use", "id": "toolu_1"}] + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "toolu_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="claude-sonnet-5", + messages=[{"role": "user", "content": "q"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"max_tokens": 0}, + logging_obj=_logging_obj("call-1"), + stream=False, + kwargs={"max_tokens": 999}, + ) + assert plan.request_patch.max_tokens == 0 + + +@pytest.mark.asyncio +async def test_recovery_disabled_when_no_caller_scope(): + # Regression: on a no-auth deployment (no UserAPIKeyAuth in metadata) the + # store key would fall back to the client-settable call id alone, letting + # any caller retrieve any other caller's originals. Recovery must be off. + guardrail = _make_guardrail() + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-abc" + logging_obj.model_call_details = {"litellm_params": {"metadata": {}}} + mock_post = AsyncMock(return_value=_make_single_compress_response()) + with patch.object(guardrail.async_handler, "post", mock_post): + result = await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + assert not has_compresr_retrieve_tool(result.get("tools") or []) + assert guardrail._originals_by_call_id == {} + + +@pytest.mark.asyncio +async def test_warns_once_per_interval_when_recovery_skipped_without_scope(): + # enable_retrieval is on but the request has no per-key auth scope: recovery + # is silently skipped, so a call-time warning must surface it, rate-limited + # within the interval but re-arming after it so an ongoing misconfiguration + # stays visible. + guardrail = _make_guardrail() + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-abc" + logging_obj.model_call_details = {"litellm_params": {"metadata": {}}} + mock_post = AsyncMock(return_value=_make_single_compress_response()) + + def _no_scope_warnings(mock_log): + return [c for c in mock_log.warning.call_args_list if "no per-key auth scope" in str(c)] + + with patch.object(guardrail.async_handler, "post", mock_post): + with patch("litellm.proxy.guardrails.guardrail_hooks.compresr.compresr.verbose_proxy_logger") as mock_log: + for _ in range(3): + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + assert len(_no_scope_warnings(mock_log)) == 1 + + guardrail._no_scope_warning_expiry = 0.0 + await guardrail.apply_guardrail( + inputs=_apply_inputs(AGENT_MESSAGES), + request_data={"model": "gpt-4o"}, + input_type="request", + logging_obj=logging_obj, + ) + assert len(_no_scope_warnings(mock_log)) == 2 diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 21e7186fca3..359b1807344 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1291,6 +1291,136 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker): } +def _patch_apply_guardrail_env(mocker, guardrail_result): + mock_guardrail = mocker.Mock() + mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result) + + mock_registry = mocker.Mock() + mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry + ) + + mock_logging_obj = mocker.Mock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_processor = mocker.Mock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj) + ) + mocker.patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", + return_value=mock_processor, + ) + + mock_proxy_logging = mocker.Mock() + mock_proxy_logging.post_call_success_hook = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + mocker.patch("litellm.proxy.proxy_server.general_settings", {}) + mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock()) + mocker.patch("litellm.proxy.proxy_server.version", "test") + mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor") + + return mock_guardrail + + +@pytest.mark.asyncio +async def test_apply_guardrail_forwards_metadata_to_guardrail(mocker): + """Client-supplied metadata must reach apply_guardrail via request_data so + parameterized custom guardrails can read per-request configuration.""" + mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]}) + + request = ApplyGuardrailRequest( + guardrail_name="test-guardrail", + text="What are tax loopholes?", + metadata={"forbidden_topics": ["tax"]}, + ) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + mock_guardrail.apply_guardrail.assert_awaited_once_with( + inputs={"texts": ["What are tax loopholes?"]}, + request_data={"metadata": {"forbidden_topics": ["tax"]}}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker): + """metadata and messages must coexist in request_data; the dict merge must + not clobber messages when both fields are sent.""" + mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]}) + + messages = [{"role": "user", "content": "What are tax loopholes?"}] + request = ApplyGuardrailRequest( + guardrail_name="test-guardrail", + text="What are tax loopholes?", + messages=messages, + metadata={"forbidden_topics": ["tax"]}, + ) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + mock_guardrail.apply_guardrail.assert_awaited_once_with( + inputs={"texts": ["What are tax loopholes?"]}, + request_data={ + "messages": messages, + "metadata": {"forbidden_topics": ["tax"]}, + }, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_omits_metadata_when_not_sent(mocker): + """Without metadata, request_data stays empty (backward-compatible).""" + mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]}) + + request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello") + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + mock_guardrail.apply_guardrail.assert_awaited_once_with( + inputs={"texts": ["hello"]}, + request_data={}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_forwards_explicit_empty_messages_and_metadata(mocker): + """Explicitly-sent empty messages/metadata must be forwarded, not dropped; + only omitted fields stay out of request_data.""" + mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]}) + + request = ApplyGuardrailRequest( + guardrail_name="test-guardrail", + text="hello", + messages=[], + metadata={}, + ) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + mock_guardrail.apply_guardrail.assert_awaited_once_with( + inputs={"texts": ["hello"]}, + request_data={"messages": [], "metadata": {}}, + input_type="request", + ) + + @pytest.mark.asyncio async def test_get_guardrail_info_endpoint_config_guardrail(mocker): """ diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 0ef9ad857f9..26feddadf79 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -44,9 +44,7 @@ def test_update_in_memory_guardrail(): "123", Guardrail( guardrail_name="test-guardrail", - litellm_params=LitellmParams( - guardrail="test-guardrail", mode="pre_call", default_on=True - ), + litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True), ), ) @@ -56,10 +54,7 @@ def test_update_in_memory_guardrail(): ) is True ) - assert ( - handler.guardrail_id_to_custom_guardrail["123"].event_hook - is GuardrailEventHooks.pre_call - ) + assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: @@ -135,6 +130,34 @@ def test_delete_in_memory_guardrail_clears_source_marker(): assert handler.get_source("a") is None +def test_list_config_guardrails_excludes_db_sourced(): + """LIT-2529: read surfaces union DB rows with config guardrails; db-sourced + in-memory entries would double-count (or resurrect stale ones), so exclude them.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one") + handler._sources["db"] = "db" + + config_guardrails = handler.list_config_guardrails() + + assert [g["guardrail_id"] for g in config_guardrails] == ["cfg"] + + +def test_get_config_guardrail_by_id_returns_config_only(): + """LIT-2529: the detail/logs fallback must return config-owned guardrails and + treat a db-sourced (stale) or missing id as a miss.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one") + handler._sources["db"] = "db" + + assert handler.get_config_guardrail_by_id("cfg")["guardrail_name"] == "config-one" + assert handler.get_config_guardrail_by_id("db") is None + assert handler.get_config_guardrail_by_id("missing") is None + + def test_initialize_guardrail_early_return_updates_source_marker(): """ When initialize_guardrail is called for a guardrail that already exists @@ -152,9 +175,7 @@ def test_initialize_guardrail_early_return_updates_source_marker(): g = Guardrail( guardrail_id="collide", guardrail_name="bedrock", - litellm_params=LitellmParams( - guardrail="bedrock", mode="pre_call", default_on=False - ), + litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call", default_on=False), ) handler.initialize_guardrail(guardrail=g, source="config") @@ -331,10 +352,7 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances(): def distinct_runner_instances() -> int: seen = set() for callback in litellm.logging_callback_manager._get_all_callbacks(): - if ( - isinstance(callback, CustomGuardrail) - and getattr(callback, "guardrail_name", None) == name - ): + if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == name: seen.add(id(callback)) return len(seen) diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index a511229942a..83593c20110 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -5,9 +5,7 @@ from unittest.mock import MagicMock, patch 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.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -36,8 +34,31 @@ def test_initialize_presidio_guardrail(): ) assert result["guardrail_name"] == "test_presidio_guardrail" - assert ( - result["litellm_params"].guardrail - == SupportedGuardrailIntegrations.PRESIDIO.value - ) + assert result["litellm_params"].guardrail == SupportedGuardrailIntegrations.PRESIDIO.value assert result["litellm_params"].mode == "pre_call" + + +def test_initialize_guardrail_preserves_guardrail_info(): + """ + Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the + stored in-memory Guardrail. Dropping it left the Guardrail Monitor's usage + endpoints unable to render type/description for YAML-defined guardrails. + """ + test_guardrail = { + "guardrail_name": "test_presidio_with_info", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": "pre_call", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + }, + "guardrail_info": {"type": "PII", "description": "masks PII"}, + } + + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + assert result is not None + assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"} + stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]] + assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"} diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py new file mode 100644 index 00000000000..bf7b1b3b238 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -0,0 +1,239 @@ +""" +Tests for the /guardrails/usage/* endpoints backing the dashboard Guardrail Monitor. + +Regression (LIT-2529): guardrails defined in config.yaml live only in +IN_MEMORY_GUARDRAIL_HANDLER, so the monitor's overview/detail/logs endpoints — +which read the litellm_guardrailstable Prisma table — could not see them: +detail 404'd, overview omitted them (or rendered them as Custom/Guardrail +orphans), and logs missed their logical-name alias. +""" + +import os +import sys +from datetime import datetime +from typing import Any, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.proxy.guardrails.usage_endpoints import ( + guardrails_usage_detail, + guardrails_usage_logs, + guardrails_usage_overview, +) +from litellm.types.guardrails import Guardrail, LitellmParams + +ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) +# Query() defaults don't resolve to None when the handler is called directly. +START, END = "2026-04-20", "2026-04-27" + + +def _config_handler(*guardrails: Guardrail) -> InMemoryGuardrailHandler: + """A real handler seeded with config-sourced YAML guardrails (no callbacks).""" + handler = InMemoryGuardrailHandler() + for g in guardrails: + gid = g["guardrail_id"] + handler.IN_MEMORY_GUARDRAILS[gid] = g + handler._sources[gid] = "config" + return handler + + +def _yaml_guardrail( + guardrail_id: str = "yaml-1", + name: str = "yaml-pii", + provider: str = "presidio", + info: Optional[dict] = None, +) -> Guardrail: + return Guardrail( + guardrail_id=guardrail_id, + guardrail_name=name, + litellm_params=LitellmParams(guardrail=provider, mode="pre_call"), + guardrail_info=info if info is not None else {"type": "PII", "description": "yaml-defined"}, + ) + + +def _db_row(guardrail_id: str = "db-1", name: str = "db-guard", provider: str = "aim") -> Any: + """A Prisma-style row: attribute access, litellm_params/guardrail_info as plain dicts.""" + row = MagicMock(spec=["guardrail_id", "guardrail_name", "litellm_params", "guardrail_info"]) + row.guardrail_id = guardrail_id + row.guardrail_name = name + row.litellm_params = {"guardrail": provider, "mode": "pre_call"} + row.guardrail_info = {"type": "ContentSafety", "description": "db-defined"} + return row + + +def _metric(guardrail_id: str, date: str = "2026-04-25", requests: int = 10, passed: int = 8, blocked: int = 2) -> Any: + m = MagicMock() + m.guardrail_id = guardrail_id + m.date = date + m.requests_evaluated = requests + m.passed_count = passed + m.blocked_count = blocked + m.flagged_count = 0 + return m + + +def _prisma( + *, + find_many=None, + find_unique=None, + metrics=None, + index_find_many=None, +) -> MagicMock: + client = MagicMock() + db = client.db + db.litellm_guardrailstable.find_many = AsyncMock(return_value=find_many or []) + db.litellm_guardrailstable.find_unique = AsyncMock(return_value=find_unique) + db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=metrics or []) + db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=index_find_many or []) + db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0) + db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + return client + + +def _patches(prisma: MagicMock, handler: InMemoryGuardrailHandler): + return ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler), + ) + + +# ---- detail ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_detail_returns_yaml_guardrail_when_db_misses(): + prisma = _prisma(find_unique=None) + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.guardrail_id == "yaml-1" + assert resp.guardrail_name == "yaml-pii" + assert resp.provider == "presidio" # coerced from the LitellmParams pydantic model + assert resp.type == "PII" # from guardrail_info + assert resp.description == "yaml-defined" + + +@pytest.mark.asyncio +async def test_detail_404_when_neither_db_nor_config(): + prisma = _prisma(find_unique=None) + handler = _config_handler() # empty + p1, p2 = _patches(prisma, handler) + with p1, p2, pytest.raises(HTTPException) as exc: + await guardrails_usage_detail(guardrail_id="ghost", start_date=START, end_date=END, user_api_key_dict=ADMIN) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_detail_does_not_surface_db_sourced_in_memory_entry(): + """A stale in-memory entry (source=db, gone from DB) must 404, not resurface.""" + prisma = _prisma(find_unique=None) + handler = InMemoryGuardrailHandler() + stale = _yaml_guardrail(guardrail_id="stale-1", name="stale") + handler.IN_MEMORY_GUARDRAILS["stale-1"] = stale + handler._sources["stale-1"] = "db" + p1, p2 = _patches(prisma, handler) + with p1, p2, pytest.raises(HTTPException) as exc: + await guardrails_usage_detail(guardrail_id="stale-1", start_date=START, end_date=END, user_api_key_dict=ADMIN) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_detail_db_row_still_resolves(): + prisma = _prisma(find_unique=_db_row(guardrail_id="db-1", provider="aim")) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.provider == "aim" + assert resp.type == "ContentSafety" + + +# ---- overview --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_overview_includes_yaml_guardrail_with_no_metrics(): + """The core bug: a YAML guardrail with zero metrics must still appear as a row.""" + prisma = _prisma(find_many=[]) # no DB guardrails + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + rows = [r for r in resp.rows if r.id == "yaml-1"] + assert len(rows) == 1 + assert rows[0].name == "yaml-pii" + assert rows[0].provider == "presidio" + assert rows[0].type == "PII" + assert rows[0].requestsEvaluated == 0 + + +@pytest.mark.asyncio +async def test_overview_yaml_metrics_matched_by_logical_name(): + """Daily metrics are keyed by logical name; the YAML row must pick them up.""" + prisma = _prisma( + find_many=[], + metrics=[_metric("yaml-pii", requests=10, blocked=2)], # keyed by name, not uuid + ) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + rows = [r for r in resp.rows if r.id == "yaml-uuid"] + assert len(rows) == 1 + assert rows[0].requestsEvaluated == 10 + assert rows[0].failRate == 20.0 + # must not also emit an orphan row keyed by the logical name + assert [r for r in resp.rows if r.id == "yaml-pii"] == [] + + +@pytest.mark.asyncio +async def test_overview_excludes_db_sourced_in_memory_entry(): + """union must not resurrect a stale db-sourced in-memory guardrail.""" + prisma = _prisma(find_many=[]) + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg"] = _yaml_guardrail(guardrail_id="cfg", name="cfg-guard") + handler._sources["cfg"] = "config" + handler.IN_MEMORY_GUARDRAILS["stale"] = _yaml_guardrail(guardrail_id="stale", name="stale-guard") + handler._sources["stale"] = "db" + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + ids = {r.id for r in resp.rows} + assert "cfg" in ids + assert "stale" not in ids + + +# ---- logs ------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_logs_resolves_config_guardrail_logical_name(): + """The index query must include the YAML guardrail's logical name alias.""" + prisma = _prisma(find_unique=None) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + await guardrails_usage_logs( + guardrail_id="yaml-uuid", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + where = prisma.db.litellm_spendlogguardrailindex.find_many.call_args.kwargs["where"] + assert where["guardrail_id"] == {"in": ["yaml-uuid", "yaml-pii"]} diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index aa24b0199ab..dffca3093fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -563,7 +563,7 @@ async def test_generate_key_debug_log_never_contains_raw_token(monkeypatch, capl generate_key_fn, ) - raw_key = "sk-short-secret" + raw_key = "sk-short-secret-a1b2" with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): await generate_key_fn( data=GenerateKeyRequest(key=raw_key), @@ -1336,10 +1336,10 @@ async def test_get_new_token_with_valid_key(monkeypatch): ) # Test with valid new_key - data = RegenerateKeyRequest(new_key="sk-test123456789") + data = RegenerateKeyRequest(new_key="sk-test1234567890abc") result = await get_new_token(data) - assert result == "sk-test123456789" + assert result == "sk-test1234567890abc" @pytest.mark.asyncio @@ -1370,6 +1370,110 @@ async def test_get_new_token_with_invalid_key(monkeypatch): assert "New key must start with 'sk-'" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_get_new_token_rejects_short_new_key(monkeypatch): + """Regression test for LIT-4355: a short custom key like sk-99 must be rejected, + otherwise the stored key_name (sk-...{last 4 chars}) reveals the entire key.""" + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + get_new_token, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + data = RegenerateKeyRequest(new_key="sk-99") + + with pytest.raises(HTTPException) as exc_info: + await get_new_token(data) + + assert exc_info.value.status_code == 400 + assert "at least 16 characters" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("short_key", ["sk-1234", "sk-abcdefghijkl"]) +async def test_generate_key_fn_rejects_short_custom_key(monkeypatch, short_key): + """Regression test for LIT-4355: /key/generate must reject custom keys shorter + than the minimum length (including the 15-char boundary); sk-1234 used to be + accepted and fully exposed via key_name.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles, ProxyException + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + assert len(short_key) < 16 + + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=GenerateKeyRequest(key=short_key), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + ) + + assert exc_info.value.code == "400" + assert "at least 16 characters" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_generate_key_fn_accepts_custom_key_at_minimum_length(monkeypatch): + """Custom keys at exactly the minimum length (16 chars) are still accepted.""" + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None, object_permission=None) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + custom_key = "sk-abcdefghijklm" + assert len(custom_key) == 16 + + response = await generate_key_fn( + data=GenerateKeyRequest(key=custom_key), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + ) + + assert response.key == custom_key + + @pytest.mark.asyncio async def test_check_custom_key_allowed_when_disabled(monkeypatch): """_check_custom_key_allowed raises 403 when disable_custom_api_keys is true.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 2a4e2ed6b25..92d1b870d75 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2186,6 +2186,36 @@ class TestCLIKeyRegenerationFlow: assert not _is_valid_cli_sso_login_id("cli-test\x001234567890") assert not _is_valid_cli_sso_login_id("sk-test1234567890") + def test_cli_sso_flow_lookup_tells_legacy_clients_to_upgrade(self): + """Legacy CLIs send self-generated sk- login ids; the 400 must say the CLI is outdated""" + from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_or_raise, + ) + + mock_cache = MagicMock() + mock_cache.get_cache.return_value = None + + with pytest.raises(HTTPException) as legacy_exc: + _get_cli_sso_flow_or_raise( + login_id="sk-85c789af-fc21-474c-9dc9-b5d794fe07ec", + cache=mock_cache, + ) + assert legacy_exc.value.status_code == 400 + assert "out of date" in legacy_exc.value.detail + assert "pip install" in legacy_exc.value.detail + mock_cache.get_cache.assert_not_called() + + with pytest.raises(HTTPException) as generic_exc: + _get_cli_sso_flow_or_raise(login_id="not-a-valid-id", cache=mock_cache) + assert generic_exc.value.status_code == 400 + assert generic_exc.value.detail == "Invalid CLI login session id" + + with pytest.raises(HTTPException) as expired_exc: + _get_cli_sso_flow_or_raise(login_id="cli-test_1234567890", cache=mock_cache) + assert expired_exc.value.status_code == 400 + assert "session not found or expired" in expired_exc.value.detail + assert "enable_redis_auth_cache" in expired_exc.value.detail + @pytest.mark.asyncio async def test_cli_sso_start_creates_bound_flow(self): """Test CLI SSO start creates a polling secret bound flow""" diff --git a/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py b/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py new file mode 100644 index 00000000000..9363c50407d --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py @@ -0,0 +1,458 @@ +""" +Tests for BillableRequestMetricsMiddleware and route classification. + +These verify the metering gate (records only on 2xx to a billable endpoint), +correct category/route classification, model-id extraction, and that the +middleware is a transparent pass-through when no recorder is injected. +""" + +import asyncio +import threading +from typing import List, Optional, Tuple + +import pytest +from starlette.applications import Starlette +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route +from starlette.testclient import TestClient + +from litellm.proxy.middleware.billable_request_metrics_middleware import ( + BillableCategory, + BillableRequestMetricsMiddleware, + _extract_model_id, + classify_billable_request, +) +from litellm.proxy.middleware.in_flight_requests_middleware import ( + InFlightRequestsMiddleware, +) + + +class FakeRecorder: + def __init__(self) -> None: + self.calls: List[dict] = [] + + def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None: + self.calls.append( + {"category": category, "route": route, "status_code": status_code, "model_id": model_id} + ) + + +def _make_app(recorder: Optional[FakeRecorder], status_code: int = 200, model_id: Optional[str] = None) -> Starlette: + async def handler(request: Request) -> Response: + headers = {"x-litellm-model-id": model_id} if model_id else {} + return JSONResponse({}, status_code=status_code, headers=headers) + + paths = [ + "/v1/chat/completions", + "/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/mcp", + "/github/mcp", + "/toolset/my-tools/mcp", + "/v1/mcp/tools", + "/v1/mcp/server", + "/a2a/agent-1/message/send", + "/v1/a2a/discover", + "/health", + "/ui", + ] + app = Starlette(routes=[Route(p, handler, methods=["GET", "POST"]) for p in paths]) + app.add_middleware(BillableRequestMetricsMiddleware, recorder=recorder) + return app + + +# ── Structure ─────────────────────────────────────────────────────────────── + + +def test_is_pure_asgi_not_base_http_middleware(): + assert not issubclass(BillableRequestMetricsMiddleware, BaseHTTPMiddleware) + + +# ── classify_billable_request ───────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "path,expected", + [ + ("/v1/chat/completions", (BillableCategory.LLM, "/chat/completions")), + ("/chat/completions", (BillableCategory.LLM, "/chat/completions")), + ("/openai/deployments/gpt-4o/chat/completions", (BillableCategory.LLM, "/chat/completions")), + ("/engines/gpt-4o/chat/completions", (BillableCategory.LLM, "/chat/completions")), + ("/v1/completions", (BillableCategory.LLM, "/completions")), + ("/completions", (BillableCategory.LLM, "/completions")), + ("/v1/embeddings", (BillableCategory.LLM, "/embeddings")), + ("/v1/responses", (BillableCategory.LLM, "/responses")), + ("/v1/rerank", (BillableCategory.LLM, "/rerank")), + ("/v1/audio/transcriptions", (BillableCategory.LLM, "/audio/transcriptions")), + # Routes from the metering-bypass finding: authenticated inference + # endpoints that must bill and previously classified as None. + ("/v1/images/edits", (BillableCategory.LLM, "/images/edits")), + ("/images/edits", (BillableCategory.LLM, "/images/edits")), + ("/openai/deployments/dall-e/images/edits", (BillableCategory.LLM, "/images/edits")), + ("/v1/images/variations", (BillableCategory.LLM, "/images/variations")), + ("/v1/messages", (BillableCategory.LLM, "/v1/messages")), + ("/interactions", (BillableCategory.LLM, "/interactions")), + ("/v1beta/interactions", (BillableCategory.LLM, "/v1beta/interactions")), + ("/v1/videos", (BillableCategory.LLM, "/videos")), + ("/v1/videos/video_123/remix", (BillableCategory.LLM, "/remix")), + ("/v1/ocr", (BillableCategory.LLM, "/ocr")), + ("/v1beta/models/gemini-2.5-pro:generateContent", (BillableCategory.LLM, ":generateContent")), + ("/v1beta/models/gemini-2.5-pro:streamGenerateContent", (BillableCategory.LLM, ":streamGenerateContent")), + # SpendLogs-producing routes surfaced by the route-inventory audit + ("/v1/search", (BillableCategory.LLM, "/search")), + ("/v1/vector_stores/vs_1/search", (BillableCategory.LLM, "/search")), + ("/v1/rag/query", (BillableCategory.LLM, "/rag/query")), + ("/rag/ingest", (BillableCategory.LLM, "/rag/ingest")), + # Provider passthrough carries real inference and writes SpendLogs + ("/bedrock/model/anthropic.claude-v2/invoke", (BillableCategory.LLM, "/bedrock")), + ("/vertex-ai/publishers/google/models/gemini:predict", (BillableCategory.LLM, "/vertex-ai")), + ("/cohere/v2/chat", (BillableCategory.LLM, "/cohere")), + # Passthrough inference bills under its provider prefix + ("/anthropic/v1/messages", (BillableCategory.LLM, "/anthropic")), + ("/mcp", (BillableCategory.MCP, "/mcp")), + ("/mcp/", (BillableCategory.MCP, "/mcp")), + ("/mcp/tools/list", (BillableCategory.MCP, "/mcp")), + ("/github/mcp", (BillableCategory.MCP, "/mcp")), + ("/github/mcp/", (BillableCategory.MCP, "/mcp")), + ("/toolset/my-tools/mcp", (BillableCategory.MCP, "/mcp")), + ("/github,slack/mcp", (BillableCategory.MCP, "/mcp")), + # REST wrapper tool execution fires the same MCP spend logging as /mcp + ("/mcp-rest/tools/call", (BillableCategory.MCP, "/mcp")), + ("/a2a/agent-1/message/send", (BillableCategory.A2A, "/a2a")), + ("/v1/a2a/agent-9/message/send", (BillableCategory.A2A, "/a2a")), + ], +) +def test_classify_billable(path: str, expected: Tuple[BillableCategory, str]): + assert classify_billable_request(path) == expected + + +@pytest.mark.parametrize( + "path", + [ + "/health", + "/health/readiness", + "/metrics", + "/ui", + "/", + "/v1/models", + "/key/generate", + "/v1/files", + # tokenization helper, not an inference call + "/v1/messages/count_tokens", + # OpenAI Assistants thread messages write no SpendLogs row + "/v1/threads/thread_abc123/messages", + "/threads/thread_abc123/messages", + # Google Interactions reads and cancel are not inference calls + "/interactions/int_123", + "/v1beta/interactions/int_123", + "/interactions/int_123/cancel", + "/v1beta/interactions/int_123/cancel", + # observability passthrough writes no SpendLogs row + "/langfuse/api/public/ingestion", + # a bare provider prefix is not an inference call + "/bedrock", + "/v1/mcp", + "/v1/mcp/tools", + "/v1/mcp/server", + "/v1/mcp/server/health", + "/v1/mcp/server/some-id", + "/v1/mcp/server/register", + "/v1/mcp/oauth/some-id/authorize", + "/a2a/agent-1/.well-known/agent-card.json", + "/v1/a2a/discover", + "/.well-known/oauth-protected-resource/github/mcp", + "/mcp-rest/tools/list", + "/mcp-rest/test/connection", + "/mcp-rest/test/tools/list", + ], +) +def test_classify_non_billable_returns_none(path: str): + assert classify_billable_request(path) is None + + +@pytest.mark.parametrize( + "path", + [ + "/v1/mcp/tools", + "/v1/mcp/server", + "/v1/mcp/server/register", + "/v1/a2a/discover", + ], +) +def test_classify_management_writes_are_not_billable(path: str): + assert classify_billable_request(path, "POST") is None + + +@pytest.mark.parametrize( + "path", + [ + "/a2a/agent-1", + "/a2a/agent-1/", + "/v1/a2a/agent-1", + ], +) +def test_classify_bare_a2a_route_is_not_billable(path: str): + """ + The bare A2A route multiplexes JSON-RPC methods off the request body. Only + message/send and message/stream write a SpendLogs row; tasks/get, + tasks/cancel and the pushNotificationConfig RPCs are forwarded upstream and + write none. Billing the path would count those task RPCs and push the metric + above the dashboard's successful-request count, so it must stay unbilled. + """ + assert classify_billable_request(path, "POST") is None + + +@pytest.mark.parametrize( + "path", + ["/v1/videos", "/v1/responses", "/v1/chat/completions", "/v1/messages"], +) +def test_classify_get_reads_are_not_billable(path: str): + """GETs on inference resources (list videos, fetch a response) write no + SpendLogs row and must not bill; only POST inference calls count.""" + assert classify_billable_request(path, "GET") is None + + +def test_classify_mcp_not_method_gated(): + assert classify_billable_request("/mcp/tools/list", "GET") == (BillableCategory.MCP, "/mcp") + + +def test_chat_completions_not_misclassified_as_plain_completions(): + """The /chat/completions suffix must win over /completions so the route label is correct.""" + category, route = classify_billable_request("/v1/chat/completions") + assert route == "/chat/completions" + + +# ── _extract_model_id ───────────────────────────────────────────────────────── + + +def test_extract_model_id_present(): + headers = [(b"content-type", b"application/json"), (b"x-litellm-model-id", b"deploy-123")] + assert _extract_model_id(headers) == "deploy-123" + + +def test_extract_model_id_case_insensitive(): + assert _extract_model_id([(b"X-LiteLLM-Model-Id", b"deploy-9")]) == "deploy-9" + + +def test_extract_model_id_absent(): + assert _extract_model_id([(b"content-type", b"application/json")]) is None + + +# ── Middleware recording behaviour ──────────────────────────────────────────── + + +def test_records_once_on_2xx_llm_with_model_id(): + recorder = FakeRecorder() + TestClient(_make_app(recorder, status_code=200, model_id="deploy-7")).post("/v1/chat/completions") + assert recorder.calls == [ + {"category": BillableCategory.LLM, "route": "/chat/completions", "status_code": 200, "model_id": "deploy-7"} + ] + + +def test_records_mcp_category(): + recorder = FakeRecorder() + TestClient(_make_app(recorder)).post("/github/mcp") + assert len(recorder.calls) == 1 and recorder.calls[0]["category"] == BillableCategory.MCP + + +def test_records_a2a_category(): + recorder = FakeRecorder() + TestClient(_make_app(recorder)).post("/a2a/agent-1/message/send") + assert len(recorder.calls) == 1 and recorder.calls[0]["category"] == BillableCategory.A2A + + +def test_does_not_record_mcp_management_read(): + recorder = FakeRecorder() + TestClient(_make_app(recorder)).get("/v1/mcp/tools") + assert recorder.calls == [] + + +def test_does_not_record_mcp_management_write(): + recorder = FakeRecorder() + TestClient(_make_app(recorder)).post("/v1/mcp/server") + assert recorder.calls == [] + + +def test_does_not_record_a2a_discovery(): + recorder = FakeRecorder() + TestClient(_make_app(recorder)).post("/v1/a2a/discover") + assert recorder.calls == [] + + +def test_does_not_record_on_4xx(): + recorder = FakeRecorder() + TestClient(_make_app(recorder, status_code=404)).post("/v1/chat/completions") + assert recorder.calls == [] + + +def test_does_not_record_on_5xx(): + recorder = FakeRecorder() + TestClient(_make_app(recorder, status_code=503)).post("/v1/chat/completions") + assert recorder.calls == [] + + +def test_does_not_record_non_billable_path(): + recorder = FakeRecorder() + TestClient(_make_app(recorder, status_code=200)).get("/health") + assert recorder.calls == [] + + +def test_no_model_id_when_header_absent(): + recorder = FakeRecorder() + TestClient(_make_app(recorder, status_code=200, model_id=None)).post("/github/mcp") + assert recorder.calls[0]["model_id"] is None + + +def test_passthrough_when_recorder_is_none(): + """Non-enterprise: middleware records nothing and does not break the response.""" + response = TestClient(_make_app(None, status_code=200)).post("/v1/chat/completions") + assert response.status_code == 200 + + +def test_record_raising_does_not_fail_the_request(): + """A broken exporter must never surface to the client: the response was + already served when record() runs, so exceptions are swallowed and logged.""" + + class ExplodingRecorder: + def record(self, *, category, route, status_code, model_id): + raise RuntimeError("exporter down") + + app = _make_app(None, status_code=200) + app.user_middleware.clear() + app.add_middleware(BillableRequestMetricsMiddleware, recorder=ExplodingRecorder()) + response = TestClient(app).post("/v1/chat/completions") + assert response.status_code == 200 + + +def test_non_http_scope_is_ignored(): + recorder = FakeRecorder() + + class _Inner: + async def __call__(self, scope, receive, send): + return None + + mw = BillableRequestMetricsMiddleware(_Inner(), recorder=recorder) + asyncio.run(mw({"type": "lifespan"}, None, None)) # type: ignore[arg-type] + assert recorder.calls == [] + + +# ── lazy recorder factory ───────────────────────────────────────────────────── + + +def test_recorder_factory_not_called_at_init(): + """The factory must run on the first request, not at middleware construction: + building at import time captured recorder=None before the YAML config's + environment_variables loaded the license and cert env vars.""" + calls = [] + + def factory(): + calls.append(1) + return FakeRecorder() + + class _Inner: + async def __call__(self, scope, receive, send): + return None + + BillableRequestMetricsMiddleware(_Inner(), recorder_factory=factory) + assert calls == [] + + +def test_recorder_factory_resolved_once_on_first_request(): + recorder = FakeRecorder() + calls = [] + + def factory(): + calls.append(1) + return recorder + + client = TestClient(_make_app_with_factory(factory, status_code=200)) + client.post("/v1/chat/completions") + client.post("/v1/chat/completions") + assert calls == [1] + assert len(recorder.calls) == 2 + + +def test_recorder_factory_returning_none_is_cached(): + calls = [] + + def factory(): + calls.append(1) + return None + + client = TestClient(_make_app_with_factory(factory, status_code=200)) + assert client.post("/v1/chat/completions").status_code == 200 + assert client.post("/v1/chat/completions").status_code == 200 + assert calls == [1] + + +def test_recorder_factory_resolved_once_under_concurrency(): + """Concurrent first requests must not each build a recorder: every extra + build leaks a MeterProvider and its background exporter thread.""" + calls = [] + release = threading.Event() + + def slow_factory(): + calls.append(1) + release.wait(timeout=2) + return FakeRecorder() + + class _Inner: + async def __call__(self, scope, receive, send): + return None + + mw = BillableRequestMetricsMiddleware(_Inner(), recorder_factory=slow_factory) + threads = [threading.Thread(target=mw._resolve_recorder) for _ in range(8)] + for t in threads: + t.start() + release.set() + for t in threads: + t.join(timeout=5) + assert calls == [1] + + +def _make_app_with_factory(factory, status_code: int) -> Starlette: + async def handler(request: Request) -> Response: + return JSONResponse({}, status_code=status_code) + + app = Starlette(routes=[Route("/v1/chat/completions", handler, methods=["POST"])]) + app.add_middleware(BillableRequestMetricsMiddleware, recorder_factory=factory) + return app + + +# ── Shutdown ordering ─────────────────────────────────────────────────────── + + +def test_record_runs_before_request_leaves_the_in_flight_tracker(): + """ + The count is recorded after the inner app returns. If this middleware sat + outside InFlightRequestsMiddleware, a request could be seen as drained while + its record() had not run, letting proxy_shutdown_event flush and stop the + exporter underneath it. Nested inside, the in-flight count still covers it. + """ + observed: List[int] = [] + + class _CountingRecorder: + def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None: + observed.append(InFlightRequestsMiddleware.get_count()) + + async def inner(scope, receive, send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + stack = InFlightRequestsMiddleware(BillableRequestMetricsMiddleware(inner, recorder=_CountingRecorder())) + assert TestClient(stack).post("/v1/chat/completions").status_code == 200 + + assert observed == [1] + assert InFlightRequestsMiddleware.get_count() == 0 + + +def test_billable_middleware_is_registered_inside_the_in_flight_tracker(): + """Starlette makes the last-added middleware outermost, so the in-flight + tracker must be registered after the billing middleware to wrap it.""" + from litellm.proxy.proxy_server import app as proxy_app + + classes = [middleware.cls for middleware in proxy_app.user_middleware] + assert classes.index(InFlightRequestsMiddleware) < classes.index(BillableRequestMetricsMiddleware) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 98b270788dc..45d35419680 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -21,6 +21,7 @@ from litellm.proxy.proxy_server import ( _is_remote_module_url, _scrub_db_overlay_remote_module_loads, _scrub_guardrail_inner, + resolve_complexity_router_plugins, ) from .conftest import normalize @@ -112,6 +113,78 @@ def test__scrub_db_overlay_remote_module_loads_invalid_non_dict_returns_input(): assert _scrub_db_overlay_remote_module_loads("litellm_settings", "raw") == "raw" +# --------------------------------------------------------------------------- +# resolve_complexity_router_plugins +# --------------------------------------------------------------------------- + + +def test_resolve_complexity_router_plugins_no_plugins_key_is_a_noop(): + config: Dict[str, Any] = {"tiers": {"SIMPLE": "gpt-4o-mini"}} + resolve_complexity_router_plugins( + model_name="smart-router", complexity_router_config=config, config_file_path=None + ) + assert config == {"tiers": {"SIMPLE": "gpt-4o-mini"}} + + +def test_resolve_complexity_router_plugins_resolves_dotted_path_to_live_instance(tmp_path): + plugin_file = tmp_path / "my_plugin.py" + plugin_file.write_text( + "class _Plugin:\n" + " async def run(self, context):\n" + " return context\n" + "\n" + "my_plugin_instance = _Plugin()\n" + ) + config: Dict[str, Any] = {"plugins": ["my_plugin.my_plugin_instance"]} + + resolve_complexity_router_plugins( + model_name="smart-router", + complexity_router_config=config, + config_file_path=str(tmp_path / "config.yaml"), + ) + + assert len(config["plugins"]) == 1 + assert hasattr(config["plugins"][0], "run") + assert type(config["plugins"][0]).__name__ == "_Plugin" + + +def test_resolve_complexity_router_plugins_rejects_non_routing_plugin_object(tmp_path): + plugin_file = tmp_path / "bad_plugin.py" + plugin_file.write_text("not_a_plugin = object()\n") + config: Dict[str, Any] = {"plugins": ["bad_plugin.not_a_plugin"]} + + with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"): + resolve_complexity_router_plugins( + model_name="smart-router", + complexity_router_config=config, + config_file_path=str(tmp_path / "config.yaml"), + ) + + +def test_resolve_complexity_router_plugins_rejects_synchronous_run_method(tmp_path): + """Regression: @runtime_checkable only checks that `run` exists as an attribute, + not that it's a coroutine function. A plugin with a synchronous `run` passes a bare + isinstance() check and would only fail at request time with a confusing + `TypeError: object RoutingContext can't be used in 'await' expression`. Reported + by Greptile on PR #33251.""" + plugin_file = tmp_path / "sync_plugin.py" + plugin_file.write_text( + "class _SyncPlugin:\n" + " def run(self, context):\n" + " return context\n" + "\n" + "sync_plugin_instance = _SyncPlugin()\n" + ) + config: Dict[str, Any] = {"plugins": ["sync_plugin.sync_plugin_instance"]} + + with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"): + resolve_complexity_router_plugins( + model_name="smart-router", + complexity_router_config=config, + config_file_path=str(tmp_path / "config.yaml"), + ) + + # --------------------------------------------------------------------------- # ProxyConfig.__init__ # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 74a0efba43d..303871e3981 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -8,7 +8,7 @@ sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to from unittest.mock import MagicMock -from litellm.proxy.route_llm_request import route_request +from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_request @pytest.mark.parametrize( @@ -42,6 +42,200 @@ async def test_route_request_dynamic_credentials(route_type): getattr(llm_router, route_type).assert_called_once_with(**data) +@pytest.mark.asyncio +async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_without_team_id(): + import litellm + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "internal-team-azure-east", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://east.example.openai.azure.com", + "api_version": "2024-02-15-preview", + "mock_response": "east", + }, + "model_info": { + "id": "team-azure-east", + "team_id": "team-a", + "team_public_model_name": "team-azure", + }, + }, + { + "model_name": "internal-team-azure-west", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://west.example.openai.azure.com", + "api_version": "2024-02-15-preview", + "mock_response": "west", + }, + "model_info": { + "id": "team-azure-west", + "team_id": "team-a", + "team_public_model_name": "team-azure", + }, + }, + ] + ) + admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + data = { + "model": "team-azure", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"user_api_key_auth": admin_auth}, + } + + llm_call = await route_request( + data=data, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=admin_auth, + ) + response = await llm_call + deployments = await router.async_get_healthy_deployments( + model="team-azure", + request_kwargs=data, + ) + + assert response.choices[0].message.content in {"east", "west"} + assert {deployment["model_info"]["id"] for deployment in deployments} == { + "team-azure-east", + "team-azure-west", + } + + non_admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + with pytest.raises(ProxyModelNotFoundError): + await route_request( + data={ + **data, + "metadata": {"user_api_key_auth": non_admin_auth}, + }, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=non_admin_auth, + ) + + from litellm.types.router import Deployment + + router.add_deployment( + Deployment( + model_name="internal-team-only", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://internal.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={ + "id": "internal-team-only-id", + "team_id": "team-a", + }, + ) + ) + internal_deployments = await router.async_get_healthy_deployments( + model="internal-team-only", + request_kwargs={ + **data, + "model": "internal-team-only", + }, + ) + + assert {deployment["model_info"]["id"] for deployment in internal_deployments} == {"internal-team-only-id"} + + router.add_deployment( + Deployment( + model_name="internal-other-team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://other.example.openai.azure.com", + "api_version": "2024-02-15-preview", + "mock_response": "other", + }, + model_info={ + "id": "other-team-azure", + "team_id": "team-b", + "team_public_model_name": "team-azure", + }, + ) + ) + + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + ambiguous_call = await route_request( + data=data, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=admin_auth, + ) + await ambiguous_call + + router.add_deployment( + Deployment( + model_name="team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://legacy.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={ + "id": "legacy-team-azure", + "team_id": "team-a", + "team_public_model_name": "team-azure", + }, + ) + ) + router.add_deployment( + Deployment( + model_name="team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://other-legacy.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={ + "id": "other-legacy-team-azure", + "team_id": "team-b", + "team_public_model_name": "team-azure", + }, + ) + ) + + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + await router.async_get_healthy_deployments( + model="team-azure", + request_kwargs=data, + ) + + router.add_deployment( + Deployment( + model_name="team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://global.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={"id": "global-team-azure"}, + ) + ) + + collision_deployments = await router.async_get_healthy_deployments( + model="team-azure", + request_kwargs=data, + ) + + assert {deployment["model_info"]["id"] for deployment in collision_deployments} == {"global-team-azure"} + + @pytest.mark.asyncio async def test_route_request_no_model_required(): """Test route types that don't require model parameter""" diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py b/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py index 45b81acbce1..9452e8042bd 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py @@ -330,6 +330,13 @@ def test_get_combined_callback_list_matrix(proxy_logging): } +def test_get_combined_callback_list_preserves_insertion_order(proxy_logging): + assert proxy_logging.get_combined_callback_list( + dynamic_success_callbacks=["prometheus", "langfuse", "datadog", "otel", "s3"], + global_callbacks=["langfuse", "gcs_bucket", "arize", "logfire"], + ) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"] + + def test_get_combined_callback_list_unhashable_dynamic_raises(proxy_logging): with pytest.raises(TypeError): proxy_logging.get_combined_callback_list( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index f7c9f343f80..12b2c9abefb 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2387,6 +2387,16 @@ class TestSubCallMetadataSanitization: assert sanitized["user_api_key_auth"] is not None assert _get_budget_reservation_from_metadata(sanitized) is None + def test_returns_empty_dict_for_missing_metadata(self): + from litellm.router_strategy.complexity_router.complexity_router import ( + _classifier_call_metadata, + ) + + for absent in (None, {}): + result = _classifier_call_metadata(absent) + assert result == {} + assert isinstance(result, dict) + def test_sanitized_auth_keeps_access_group_fields_and_leaves_original_untouched(self): from litellm.proxy._types import UserAPIKeyAuth from litellm.router_strategy.complexity_router.complexity_router import ( @@ -2489,7 +2499,7 @@ class TestRoutingDecisionCauseLogging: class TestSessionAffinity: - """Test the opt-in session_affinity sticky-routing behavior.""" + """Test the session_affinity sticky-routing behavior (on by default).""" REASONING_MESSAGE = [ { @@ -2503,14 +2513,19 @@ class TestSessionAffinity: def session_affinity_config(self, basic_config) -> Dict: return {**basic_config, "session_affinity": True} + @pytest.fixture + def session_affinity_disabled_config(self, basic_config) -> Dict: + return {**basic_config, "session_affinity": False} + @staticmethod def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} @pytest.mark.asyncio - async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): - """Regression: session_affinity defaults to False, so a shared session_id must - not pin the model -- each turn is still classified independently.""" + async def test_enabled_by_default_pins_model(self, mock_router_instance, basic_config): + """Regression: session_affinity defaults to True, so a shared session_id pins the + first turn's model and later turns reuse it instead of reclassifying.""" + assert "session_affinity" not in basic_config mock_router_instance.cache = DualCache() router = ComplexityRouter( model_name="test-router", @@ -2525,6 +2540,28 @@ class TestSessionAffinity: model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE ) assert first.model == "o1-preview" + assert second.model == "o1-preview" + + @pytest.mark.asyncio + async def test_can_be_disabled_reclassifies_every_turn( + self, mock_router_instance, session_affinity_disabled_config + ): + """Regression: session_affinity=False must still reclassify every turn even when a + shared session_id is present, so the opt-out keeps working.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_disabled_config, + ) + request_kwargs = self._request_kwargs("session-1") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" assert second.model == "gpt-4o-mini" @pytest.mark.asyncio @@ -2700,3 +2737,250 @@ class TestSessionAffinity: spy_aclassify.assert_not_called() assert second.model == "cheap" assert request_kwargs_2["metadata"]["adaptive_router_chosen_model"] == "cheap" + + +class _DummyPlugin: + async def run(self, context): + return context + + +class TestRoutingPlugins: + """Test the `complexity_router_config.plugins` field: narrows the classified + tier's candidate pool before a model is picked. Discussion: + https://github.com/BerriAI/litellm/discussions/32168""" + + @pytest.mark.asyncio + async def test_plugin_narrows_tier_candidates(self, mock_router_instance): + class ExcludeGpt4oMini: + async def run(self, context): + context.candidate_models = [m for m in context.candidate_models if m != "gpt-4o-mini"] + return context + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"]}, + "plugins": [ExcludeGpt4oMini()], + }, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "hi"}], + ) + assert result is not None + assert result.model == "gpt-4o-nano" + + @pytest.mark.asyncio + async def test_plugin_narrowing_to_zero_raises_even_with_default_model_configured(self, mock_router_instance): + """Regression: default_model must never be used as an escape hatch around a + plugin's narrowing decision -- it was never checked against the plugins, so + falling back to it would let a tenant/budget policy be silently bypassed. + Reported by Veria AI on PR #33251.""" + + class BlockEverything: + async def run(self, context): + context.candidate_models = [] + return context + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "default_model": "gpt-4o-fallback", + "plugins": [BlockEverything()], + }, + ) + with pytest.raises(ValueError, match="No candidate models left for tier"): + await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "hi"}], + ) + + @pytest.mark.asyncio + async def test_plugin_narrowing_to_zero_without_default_model_raises(self, mock_router_instance): + class BlockEverything: + async def run(self, context): + context.candidate_models = [] + return context + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "plugins": [BlockEverything()], + }, + ) + with pytest.raises(ValueError, match="No candidate models left for tier"): + await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "hi"}], + ) + + @pytest.mark.asyncio + async def test_plugin_receives_metadata_from_request_kwargs(self, mock_router_instance): + captured = {} + + class CaptureMetadata: + async def run(self, context): + captured.update(context.metadata) + return context + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "plugins": [CaptureMetadata()], + }, + ) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs={"metadata": {"tenant": "acme-corp"}}, + messages=[{"role": "user", "content": "hi"}], + ) + assert captured.get("tenant") == "acme-corp" + + @pytest.mark.asyncio + async def test_plugin_applies_to_keyword_tier_override(self, mock_router_instance): + """A policy plugin must not be bypassable via the keyword_tier_rules override path.""" + + class ExcludeGpt4oMini: + async def run(self, context): + context.candidate_models = [m for m in context.candidate_models if m != "gpt-4o-mini"] + return context + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini", "gpt-4o-nano"]}, + "keyword_tier_rules": [{"keywords": ["hello"], "tier": "SIMPLE"}], + "plugins": [ExcludeGpt4oMini()], + }, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "hello there"}], + ) + assert result is not None + assert result.model == "gpt-4o-nano" + + @pytest.mark.asyncio + async def test_plugin_applies_to_no_user_message_default_tier_path(self, mock_router_instance): + """Regression: `self.config.default_model or await self._pick_model_for_tier(...)` + short-circuited on a truthy default_model, so the no-user-message path never ran + the plugin pipeline at all when default_model was configured. A policy plugin + must not be bypassable via this path either. Reported by Veria AI on PR #33251.""" + + class ExcludeDefaultModel: + async def run(self, context): + context.candidate_models = [m for m in context.candidate_models if m != "gpt-4o-default"] + return context + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"MEDIUM": ["gpt-4o-default", "gpt-4o-nano"]}, + "default_model": "gpt-4o-default", + "plugins": [ExcludeDefaultModel()], + }, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[ + {"role": "system", "content": "You are helpful."}, + {"role": "assistant", "content": "Hello!"}, + ], + ) + assert result is not None + assert result.model == "gpt-4o-nano" + + @pytest.mark.asyncio + async def test_no_user_message_prefers_default_model_over_medium_tier_without_plugins( + self, mock_router_instance + ): + """Regression: without plugins configured, the no-user-message path must keep its + pre-existing default_model-first priority over the MEDIUM tier exactly as before -- + closing the plugin-bypass gap must not silently flip model selection for the (much + larger) population of users who don't use plugins at all. Flagged by Greptile on + PR #33251 after the plugin-bypass fix changed this priority unconditionally.""" + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"MEDIUM": ["gpt-4o-medium-tier"]}, + "default_model": "gpt-4o-configured-default", + }, + ) + result = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[ + {"role": "system", "content": "You are helpful."}, + {"role": "assistant", "content": "Hello!"}, + ], + ) + assert result is not None + assert result.model == "gpt-4o-configured-default" + + def test_plugins_and_adaptive_together_raises(self): + with pytest.raises(ValidationError, match="plugins and adaptive=True cannot both be set"): + ComplexityRouterConfig( + tiers={"SIMPLE": ["gpt-4o-mini"]}, + adaptive=True, + plugins=[_DummyPlugin()], + ) + + @pytest.mark.asyncio + async def test_no_plugins_configured_is_unaffected(self, complexity_router): + """Regression guard: a ComplexityRouter with no `plugins` configured behaves exactly as before.""" + result = await complexity_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "Hello!"}], + ) + assert result is not None + + @pytest.mark.asyncio + async def test_session_affinity_pin_shortcut_disabled_when_plugins_configured(self, mock_router_instance): + """Regression: the session_affinity cache-pin shortcut returned a stale pinned + model without ever re-running it through plugins, so a policy plugin's decision + (e.g. a budget cap crossed mid-session) was only ever enforced on a session's + first turn. With plugins configured, every turn must go through + _classify_and_route (and therefore the plugin pipeline) again.""" + mock_router_instance.cache = DualCache() + + class AllowAll: + async def run(self, context): + return context + + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "session_affinity": True, + "plugins": [AllowAll()], + }, + ) + request_kwargs = {"metadata": {"session_id": "session-1"}} + + with patch.object(router, "_classify_and_route", wraps=router._classify_and_route) as spy: + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=[{"role": "user", "content": "hi"}] + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=[{"role": "user", "content": "hi again"}] + ) + assert first.model == "gpt-4o-mini" + assert second.model == "gpt-4o-mini" + assert spy.call_count == 2 diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index e9c12d009e2..78ed71f5ffd 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -218,3 +218,36 @@ def test_filter_by_routing_plugin_candidates_narrows_and_raises_when_empty(): healthy_deployments=healthy_deployments, request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["nonexistent/model"]}}, ) + + +def test_json_default_stable_id_is_stable_across_instances(): + """_generate_model_id's json.dumps `default=` fallback must not embed an object's + memory address (e.g. plain str() on an object with no custom __repr__ falls back + to object.__repr__'s ``) -- that would make the + deployment id churn on every process restart for any deployment whose + litellm_params contain a live plugin instance.""" + router = Router(model_list=_smart_router_model_list()) + + assert router._json_default_stable_id(LanguageDetector()) == router._json_default_stable_id(LanguageDetector()) + assert router._json_default_stable_id(LanguageDetector()) != router._json_default_stable_id(TenantPolicy()) + + +def test_generate_model_id_is_stable_when_litellm_params_contain_a_plugin_instance(): + """End-to-end: a deployment id built from litellm_params containing a routing + plugin instance (e.g. complexity_router_config.plugins) must be identical across + separate calls, not just non-crashing.""" + router = Router(model_list=_smart_router_model_list()) + litellm_params = { + "model": "auto_router/complexity_router", + "complexity_router_config": {"plugins": [LanguageDetector()]}, + } + + id1 = router._generate_model_id("smart-router", litellm_params) + id2 = router._generate_model_id( + "smart-router", + { + "model": "auto_router/complexity_router", + "complexity_router_config": {"plugins": [LanguageDetector()]}, + }, + ) + assert id1 == id2 diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py index 939e3189596..a5a68271111 100644 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -977,3 +977,65 @@ class TestRouterIOTokenIntegration: assert info is not None assert info.itpm == 100 assert info.otpm == 20 + + +class TestContextSlotRetention: + def test_setter_stores_kwargs_only_for_io_limited_deployments(self): + """ + The context slot pins the entire request kwargs (messages included) + for the lifetime of the surrounding asyncio context, and pooled + resources created mid-request (e.g. redis connections) capture that + context, extending the pin far past the request. Only ITPM/OTPM + pre-call checks read the slot, so the setter must store None for + deployments without io token limits and still clear reservation + sentinels from kwargs either way. + """ + kwargs = { + "messages": [{"role": "user", "content": "x" * 1000}], + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + } + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + assert ITPM_RESERVED_KEY not in kwargs["metadata"] + assert ITPM_CACHE_KEY not in kwargs["metadata"] + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True) + assert get_io_token_rate_limit_request_kwargs() is kwargs + + set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_does_not_pin_kwargs_without_io_limits(self): + router = Router( + model_list=[ + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + } + ] + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("plain") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is None + + @pytest.mark.asyncio + async def test_router_pins_kwargs_for_io_limited_deployment(self): + router = Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + set_io_token_rate_limit_request_kwargs(None) + kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + deployment = router.get_deployment_by_model_group_name("limited") + assert deployment is not None + router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) + assert get_io_token_rate_limit_request_kwargs() is kwargs diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 85430ba752b..7152eea7c9f 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -65,6 +65,16 @@ def test_redact_string_catches_secret_patterns(): assert redact_string(normal) == normal +def test_redact_string_catches_minimum_length_virtual_key(): + """Regression test for LIT-4355: keys at the enforced 16-char minimum + (MINIMUM_CUSTOM_KEY_LENGTH) must be treated as key-shaped by the scrubber.""" + minimum_length_key = "sk-abcdefghijklm" + assert len(minimum_length_key) == 16 + result = redact_string("msg: " + minimum_length_key) + assert minimum_length_key not in result + assert "REDACTED" in result + + def test_filter_redacts_secrets_in_logger_output(): def log_messages(): verbose_logger.debug("Key: " + SECRET) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 44d4a111abd..a2c91205b15 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1145,11 +1145,6 @@ "count": 1 } }, - "src/app/(dashboard)/tag-management/_components/TagTable.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx": { "no-restricted-imports": { "count": 1 @@ -1639,17 +1634,6 @@ "count": 1 } }, - "src/components/UsageIndicator.tsx": { - "no-nested-ternary": { - "count": 4 - }, - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/static-components": { - "count": 1 - } - }, "src/components/activity_metrics.tsx": { "no-nested-ternary": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.test.tsx index 94ed707f31d..d6f5f22616b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.test.tsx @@ -53,7 +53,72 @@ describe("GuardrailTestPanel", () => { // Verify onSubmit was called with the correct text await waitFor(() => { - expect(mockOnSubmit).toHaveBeenCalledWith("Test input text"); + expect(mockOnSubmit).toHaveBeenCalledWith("Test input text", null); }); }); + + it("should submit parsed metadata when a JSON object is provided", async () => { + /** + * Tests that a JSON object typed into the Metadata field is parsed and + * passed to onSubmit so it reaches the apply_guardrail request body. + */ + const user = userEvent.setup(); + + render( + , + ); + + const textarea = screen.getByPlaceholderText("Enter text to test with guardrails..."); + await user.type(textarea, "Test input text"); + + const metadataField = screen.getByPlaceholderText('{"forbidden_topics": ["tax", "finance"]}'); + await user.click(metadataField); + await user.paste('{"forbidden_topics": ["tax"]}'); + + await user.click(screen.getByRole("button", { name: /Test 2 guardrails/ })); + + await waitFor(() => { + expect(mockOnSubmit).toHaveBeenCalledWith("Test input text", { forbidden_topics: ["tax"] }); + }); + }); + + it("should block submission and show an error for invalid metadata JSON", async () => { + /** + * Tests that invalid JSON in the Metadata field prevents submission + * instead of silently sending a request without metadata. + */ + const user = userEvent.setup(); + + render( + , + ); + + const textarea = screen.getByPlaceholderText("Enter text to test with guardrails..."); + await user.type(textarea, "Test input text"); + + const metadataField = screen.getByPlaceholderText('{"forbidden_topics": ["tax", "finance"]}'); + await user.click(metadataField); + await user.paste("{not json"); + + await user.click(screen.getByRole("button", { name: /Test 2 guardrails/ })); + + await waitFor(() => { + expect(screen.getByText("Invalid JSON")).toBeInTheDocument(); + }); + expect(mockOnSubmit).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx index a65b980d8b8..e310da55063 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx @@ -10,7 +10,7 @@ const { Text } = Typography; interface GuardrailTestPanelProps { guardrailNames: string[]; - onSubmit: (text: string) => void; + onSubmit: (text: string, metadata?: Record | null) => void; isLoading: boolean; results: Array<{ guardrailName: string; response_text: string; latency: number }> | null; errors: Array<{ guardrailName: string; error: Error; latency: number }> | null; @@ -26,6 +26,23 @@ export function GuardrailTestPanel({ onClose, }: GuardrailTestPanelProps) { const [inputText, setInputText] = useState(""); + const [metadataText, setMetadataText] = useState(""); + const [metadataError, setMetadataError] = useState(null); + + const parseMetadata = (raw: string): { metadata: Record | null; error: string | null } => { + if (!raw.trim()) { + return { metadata: null, error: null }; + } + try { + const parsed = JSON.parse(raw); + if (parsed === null || typeof parsed !== "object" || Array.isArray(parsed)) { + return { metadata: null, error: "Metadata must be a JSON object" }; + } + return { metadata: parsed, error: null }; + } catch { + return { metadata: null, error: "Invalid JSON" }; + } + }; const handleSubmit = () => { if (!inputText.trim()) { @@ -33,7 +50,15 @@ export function GuardrailTestPanel({ return; } - onSubmit(inputText); + const { metadata, error } = parseMetadata(metadataText); + if (error) { + setMetadataError(error); + NotificationsManager.fromBackend(`Metadata: ${error}`); + return; + } + setMetadataError(null); + + onSubmit(inputText, metadata); }; const handleKeyDown = (e: React.KeyboardEvent) => { @@ -142,6 +167,33 @@ export function GuardrailTestPanel({ +
+
+ + + + +
+