diff --git a/helm/litellm/README.md b/helm/litellm/README.md index 959657b1b8f..53dc300b886 100644 --- a/helm/litellm/README.md +++ b/helm/litellm/README.md @@ -5,3 +5,21 @@ This chart deploys separate gateway, backend and UI services. Configure Lens thr Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for the complete values and verification flow. It also links the GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published + +For an ingress with separate component hostnames, set all three `ingress.componentHosts` values, leave `ingress.host` and `ingress.extraPaths` empty, configure `ui.backendUrl` as the backend's public origin and `backend.gatewayUrl` as the gateway's public origin, and set `ui.publicUrl` to the browser-facing UI origin. The UI Deployment receives `LITELLM_BACKEND_URL` from `ui.backendUrl` only when `ingress.componentHosts` is configured. + +```yaml +ingress: + enabled: true + host: "" + componentHosts: + ui: ui.example.com + backend: backend.example.com + gateway: gateway.example.com + extraPaths: [] +ui: + publicUrl: https://ui.example.com + backendUrl: https://backend.example.com +backend: + gatewayUrl: https://gateway.example.com +``` diff --git a/helm/litellm/templates/NOTES.txt b/helm/litellm/templates/NOTES.txt index 468cf621b32..135dcb876e0 100644 --- a/helm/litellm/templates/NOTES.txt +++ b/helm/litellm/templates/NOTES.txt @@ -22,12 +22,14 @@ Reminders: - database.writer.{host,port,dbname} (writer connection pieces) - database.writer.passwordSecret.{name,usernameKey,passwordKey} (Secret holding the writer DB username + password) + - database.writer.passwordSecret.{hostKey,urlKey} (optional, read writer host or full URL from that Secret) - database.writer.useIAMAuth: true (optional — chart sets IAM_TOKEN_DB_AUTH=true and omits DATABASE_PASSWORD / DATABASE_URL so the proxy mints the URL from an IAM token at startup) - - database.reader.host (optional — enables read-replica routing; reader - .passwordSecret.name is required when set, unless - .useIAMAuth is true) + - database.reader.host / database.reader.passwordSecret.{hostKey,urlKey} + (optional, enables read-replica routing. The Secret + name is required for hostKey, urlKey, and for host + with password auth) - database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true — chart emits DATABASE_*_READ_REPLICA env vars and omits DATABASE_PASSWORD_READ_REPLICA / @@ -35,6 +37,8 @@ Reminders: reader URL from an IAM token at startup) - redis.passwordSecret.name (optional — set when redis.host is provided and the cache requires auth) + - redis.passwordSecret.{hostKey,urlKey} (optional, read Redis host or full URL from that Secret) + (the Secret name is required when either key is set) - redis.cluster: true (optional — chart sets REDIS_CLUSTER_NODES from redis.host / redis.port so the proxy's Cache() constructs a RedisClusterCache; the cluster client diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 294ce90c907..647e435bdf6 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -27,6 +27,30 @@ Common naming + label helpers shared by gateway, backend, and ui templates. {{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}} {{- end -}} +{{- define "litellm.ui.publicOrigin" -}} +{{- $origin := .Values.ui.publicUrl | trim | trimSuffix "/" -}} +{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}} +{{- fail "ui.publicUrl must be an origin like https://ui.example.com, without a path" -}} +{{- end -}} +{{- $origin -}} +{{- end -}} + +{{- define "litellm.ui.backendOrigin" -}} +{{- $origin := .Values.ui.backendUrl | trim | trimSuffix "/" -}} +{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}} +{{- fail "ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set" -}} +{{- end -}} +{{- $origin -}} +{{- end -}} + +{{- define "litellm.backend.gatewayOrigin" -}} +{{- $origin := .Values.backend.gatewayUrl | trim | trimSuffix "/" -}} +{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}} +{{- fail "backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set" -}} +{{- end -}} +{{- $origin -}} +{{- end -}} + {{- define "litellm.commonLabels" -}} app.kubernetes.io/name: {{ include "litellm.name" . }} app.kubernetes.io/instance: {{ .Release.Name }} @@ -202,7 +226,9 @@ the chart context (needed for .Values), `component` selects which component's Sensitive values (master key, DB username + password, Redis password) come only from referenced Secrets; the chart never accepts inline values for them. -The chart never assembles DATABASE_URL itself. It emits only the discrete +The chart never assembles DATABASE_URL itself. When +database.writer.passwordSecret.urlKey is set, the chart passes the operator's +URL through from the Secret. Otherwise it emits only the discrete DATABASE_HOST/PORT/USER/NAME/SCHEMA (+ DATABASE_PASSWORD for password auth) vars; the proxy's entrypoint (DatabaseURLSettings in litellm/proxy/db/db_url_settings.py) builds the URL from them and @@ -218,11 +244,11 @@ token instead of a static password. `database.writer.useAzureEntraAuth: true` does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token, for Azure Database for PostgreSQL. The two are mutually exclusive. -The read replica is opt-in via `database.reader.host`. The chart emits -DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA / -DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth -modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for -password auth. When `database.reader.useIAMAuth: true` (or +The read replica is opt-in via `database.reader.host`, `hostKey`, or `urlKey`. +The chart emits DATABASE_URL_READ_REPLICA alone for urlKey, or discrete +DATABASE_*_READ_REPLICA vars for host and hostKey configurations, plus +DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for password auth. +When `database.reader.useIAMAuth: true` (or `database.reader.useAzureEntraAuth: true`) it omits DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the same way. Reader token auth only takes effect when the writer uses the same @@ -242,8 +268,36 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. value: {{ $component.logLevel | quote }} {{- end }} {{- with $root.Values.database.writer }} +{{- $hostKey := .passwordSecret.hostKey }} +{{- $urlKey := .passwordSecret.urlKey }} +{{- if and .host $hostKey }} +{{- fail "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" }} +{{- end }} +{{- if and $urlKey (or .host $hostKey .dbname .schema .sslMode .sslRootCert .useIAMAuth .useAzureEntraAuth) }} +{{- fail "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" }} +{{- end }} +{{- if and (not .host) (not $hostKey) (not $urlKey) }} +{{- fail "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" }} +{{- end }} +{{- if and (or $hostKey $urlKey) (not .passwordSecret.name) }} +{{- fail "database.writer.passwordSecret.name is required when hostKey or urlKey is set" }} +{{- end }} +{{- if $urlKey }} +- name: DATABASE_URL + valueFrom: + secretKeyRef: + name: {{ .passwordSecret.name | quote }} + key: {{ $urlKey | quote }} +{{- else }} - name: DATABASE_HOST - value: {{ required "database.writer.host is required" .host | quote }} + {{- if $hostKey }} + valueFrom: + secretKeyRef: + name: {{ .passwordSecret.name | quote }} + key: {{ $hostKey | quote }} + {{- else }} + value: {{ .host | quote }} + {{- end }} - name: DATABASE_PORT value: {{ .port | default 5432 | quote }} - name: DATABASE_USER @@ -282,8 +336,29 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. key: {{ .passwordSecret.passwordKey | default "password" }} {{- end }} {{- end }} +{{- end }} {{- with $root.Values.database.reader }} -{{- if .host }} +{{- $readerSecret := .passwordSecret | default dict }} +{{- $hostKey := $readerSecret.hostKey }} +{{- $urlKey := $readerSecret.urlKey }} +{{- $readerEnabled := or .host $hostKey $urlKey }} +{{- if and .host $hostKey }} +{{- fail "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" }} +{{- end }} +{{- if and $urlKey (or .host $hostKey .dbname .schema .useIAMAuth .useAzureEntraAuth) }} +{{- fail "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" }} +{{- end }} +{{- if and (or $hostKey $urlKey) (not $readerSecret.name) }} +{{- fail "database.reader.passwordSecret.name is required when hostKey or urlKey is set" }} +{{- end }} +{{- if $readerEnabled }} +{{- if $urlKey }} +- name: DATABASE_URL_READ_REPLICA + valueFrom: + secretKeyRef: + name: {{ $readerSecret.name | quote }} + key: {{ $urlKey | quote }} +{{- else }} {{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }} {{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }} {{- end }} @@ -291,11 +366,18 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. {{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }} {{- end }} - name: DATABASE_HOST_READ_REPLICA + {{- if $hostKey }} + valueFrom: + secretKeyRef: + name: {{ $readerSecret.name | quote }} + key: {{ $hostKey | quote }} + {{- else }} value: {{ .host | quote }} + {{- end }} - name: DATABASE_PORT_READ_REPLICA value: {{ .port | default 5432 | quote }} - name: DATABASE_NAME_READ_REPLICA - value: {{ required "database.reader.dbname is required when database.reader.host is set" .dbname | quote }} + value: {{ required "database.reader.dbname is required when database.reader.host, hostKey, or urlKey is set" .dbname | quote }} {{- if .schema }} - name: DATABASE_SCHEMA_READ_REPLICA value: {{ .schema | quote }} @@ -325,6 +407,7 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets. {{- end }} {{- end }} {{- end }} +{{- end }} {{/* The migrations Job (helm.sh/hook: pre-upgrade) is the single owner of `prisma migrate deploy`. Without this, every gateway/backend pod also runs @@ -341,9 +424,34 @@ harmless no-op for the Job and authoritative for the app pods. tracking, pod lock manager) via its REDIS_* env fallback. An explicit `general_settings.coordination_redis` block in proxy_config takes precedence over anything emitted here. */}} -{{- if $root.Values.redis.host }} +{{- $redisSecret := $root.Values.redis.passwordSecret | default dict }} +{{- $hostKey := $redisSecret.hostKey }} +{{- $urlKey := $redisSecret.urlKey }} +{{- if and $root.Values.redis.host $hostKey }} +{{- fail "redis.host and redis.passwordSecret.hostKey cannot both be set" }} +{{- end }} +{{- if and $urlKey (or $root.Values.redis.host $hostKey $root.Values.redis.cluster) }} +{{- fail "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" }} +{{- end }} +{{- if and (or $hostKey $urlKey) (not $redisSecret.name) }} +{{- fail "redis.passwordSecret.name is required when hostKey or urlKey is set" }} +{{- end }} +{{- if $urlKey }} +- name: REDIS_URL + valueFrom: + secretKeyRef: + name: {{ $redisSecret.name | quote }} + key: {{ $urlKey | quote }} +{{- else if or $root.Values.redis.host $hostKey }} - name: REDIS_HOST + {{- if $hostKey }} + valueFrom: + secretKeyRef: + name: {{ $redisSecret.name | quote }} + key: {{ $hostKey | quote }} + {{- else }} value: {{ $root.Values.redis.host | quote }} + {{- end }} - name: REDIS_PORT value: {{ $root.Values.redis.port | quote }} {{- if $root.Values.redis.passwordSecret.name }} @@ -360,7 +468,11 @@ harmless no-op for the Job and authoritative for the app pods. configured endpoint; the cluster client discovers the remaining nodes from CLUSTER SLOTS at startup. */}} - name: REDIS_CLUSTER_NODES + {{- if $hostKey }} + value: {{ printf "[{\"host\":\"$(REDIS_HOST)\",\"port\":%v}]" (int $root.Values.redis.port) | quote }} + {{- else }} value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }} + {{- end }} {{- end }} {{- end }} {{- with $component.extraEnv }} diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 40aa70ec58e..7360261dab1 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -1,4 +1,6 @@ {{- if .Values.backend.enabled }} +{{- $componentHosts := .Values.ingress.componentHosts | default dict -}} +{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}} apiVersion: apps/v1 kind: Deployment metadata: @@ -58,7 +60,17 @@ spec: protocol: TCP env: {{- include "litellm.lensConnectionEnv" . | nindent 12 }} + {{- if .Values.ui.publicUrl }} + - name: LITELLM_CORS_ORIGINS + value: {{ include "litellm.ui.publicOrigin" . | quote }} + - name: LITELLM_CORS_ALLOW_CREDENTIALS + value: "true" + {{- end }} {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }} + {{- if .Values.backend.gatewayUrl }} + - name: LITELLM_GATEWAY_URL + value: {{ if $hasComponentHosts }}{{ include "litellm.backend.gatewayOrigin" . | quote }}{{ else }}{{ .Values.backend.gatewayUrl | quote }}{{ end }} + {{- end }} {{- if .Values.gateway.config.create }} - name: CONFIG_FILE_PATH value: /app/config/config.yaml diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 9b3b58c97ed..b0e7f0dcebb 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -56,6 +56,10 @@ spec: protocol: TCP env: {{- include "litellm.lensConnectionEnv" . | nindent 12 }} + {{- if .Values.ui.publicUrl }} + - name: LITELLM_CORS_ORIGINS + value: {{ include "litellm.ui.publicOrigin" . | quote }} + {{- end }} {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }} {{- if .Values.gateway.config.create }} - name: CONFIG_FILE_PATH diff --git a/helm/litellm/templates/ingress.yaml b/helm/litellm/templates/ingress.yaml index a933b5e85c6..22ff8a95566 100644 --- a/helm/litellm/templates/ingress.yaml +++ b/helm/litellm/templates/ingress.yaml @@ -1,3 +1,36 @@ +{{- $componentHosts := .Values.ingress.componentHosts | default dict -}} +{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}} +{{- if $hasComponentHosts }} +{{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }} +{{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }} +{{- end }} +{{- if not .Values.gateway.enabled }} +{{- fail "ingress.componentHosts requires gateway.enabled=true" }} +{{- end }} +{{- if not .Values.backend.enabled }} +{{- fail "ingress.componentHosts requires backend.enabled=true" }} +{{- end }} +{{- if not .Values.ui.enabled }} +{{- fail "ingress.componentHosts requires ui.enabled=true" }} +{{- end }} +{{- if .Values.ingress.host }} +{{- fail "ingress.host must be empty when ingress.componentHosts is configured" }} +{{- end }} +{{- if .Values.ingress.extraPaths }} +{{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }} +{{- end }} +{{- if not (trim .Values.ui.backendUrl) }} +{{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }} +{{- end }} +{{- $_ := include "litellm.ui.backendOrigin" . }} +{{- if not (trim .Values.backend.gatewayUrl) }} +{{- fail "backend.gatewayUrl is required when ingress.componentHosts is configured" }} +{{- end }} +{{- $_ := include "litellm.backend.gatewayOrigin" . }} +{{- if not (trim .Values.ui.publicUrl) }} +{{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }} +{{- end }} +{{- end }} {{- if .Values.ingress.enabled -}} {{- $gatewayName := include "litellm.gateway.fullname" . -}} {{- $backendName := include "litellm.backend.fullname" . -}} @@ -90,6 +123,66 @@ adding to it. */}} {{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}} +{{- if $hasComponentHosts }} +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: {{ include "litellm.fullname" . }} + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + {{- with .Values.ingress.annotations }} + annotations: + {{- toYaml . | nindent 4 }} + {{- end }} +spec: + {{- with .Values.ingress.className }} + ingressClassName: {{ . | quote }} + {{- end }} + {{- with .Values.ingress.tls }} + tls: + {{- toYaml . | nindent 4 }} + {{- end }} + rules: + - host: {{ $componentHosts.ui | quote }} + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: {{ $uiName }} + port: + number: {{ $uiPort }} + - host: {{ $componentHosts.backend | quote }} + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: {{ $backendName }} + port: + number: {{ $backendPort }} + - host: {{ $componentHosts.gateway | quote }} + http: + paths: + {{- if or (eq (include "litellm.lens.mode" .) "bundled") (and (eq (include "litellm.lens.mode" .) "external") .Values.lensWorker.externalServiceName) }} + - path: /lens-ingest + pathType: Prefix + backend: + service: + name: {{ .Values.lensWorker.externalServiceName | default (printf "%s-lens-worker" (include "litellm.fullname" .)) }} + port: + number: {{ .Values.lensWorker.service.port }} + {{- end }} + - path: / + pathType: Prefix + backend: + service: + name: {{ $gatewayName }} + port: + number: {{ $gatewayPort }} +{{- else }} apiVersion: networking.k8s.io/v1 kind: Ingress metadata: @@ -217,3 +310,4 @@ spec: port: number: {{ $backendPort }} {{- end }} +{{- end }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index efee2d5fc34..394eea17451 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -1,3 +1,5 @@ +{{- $componentHosts := .Values.ingress.componentHosts | default dict -}} +{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}} {{- if .Values.ui.enabled }} apiVersion: apps/v1 kind: Deployment @@ -56,9 +58,9 @@ spec: - name: LITELLM_LOG value: {{ .Values.ui.logLevel | quote }} {{- end }} - {{- if .Values.ui.backendUrl }} + {{- if and .Values.ui.backendUrl $hasComponentHosts }} - name: LITELLM_BACKEND_URL - value: {{ .Values.ui.backendUrl | quote }} + value: {{ include "litellm.ui.backendOrigin" . | quote }} {{- end }} {{- with .Values.ui.extraEnv }} {{- toYaml . | nindent 12 }} diff --git a/helm/litellm/tests/component_hosts_tests.yaml b/helm/litellm/tests/component_hosts_tests.yaml new file mode 100644 index 00000000000..f9c9ce6fad2 --- /dev/null +++ b/helm/litellm/tests/component_hosts_tests.yaml @@ -0,0 +1,441 @@ +suite: ingress component hosts +templates: + - ingress.yaml + - backend/deployment.yaml + - gateway/deployment.yaml + - gateway/configmap.yaml + - ui/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: routes each hostname to its component and injects the external URLs + template: ingress.yaml + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + lensWorker.mode: disabled + ingress.tls: + - hosts: + - ui.example.com + - backend.example.com + - gateway.example.com + secretName: component-hosts + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + asserts: + - equal: + path: spec.rules[0].host + value: ui.example.com + - equal: + path: spec.tls[0] + value: + hosts: + - ui.example.com + - backend.example.com + - gateway.example.com + secretName: component-hosts + - equal: + path: spec.rules[0].http.paths + value: + - path: / + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-ui + port: + number: 3000 + - equal: + path: spec.rules[1].host + value: backend.example.com + - equal: + path: spec.rules[1].http.paths + value: + - path: / + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-backend + port: + number: 4001 + - equal: + path: spec.rules[2].host + value: gateway.example.com + - equal: + path: spec.rules[2].http.paths + value: + - path: / + pathType: Prefix + backend: + service: + name: RELEASE-NAME-litellm-gateway + port: + number: 4000 + + - it: keeps Lens ingestion ahead of the gateway catch-all + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + lensWorker.mode: external + lensWorker.externalServiceName: existing-lens + lensWorker.externalUrl: https://lens.example.com + lensWorker.publicUrl: https://lens.example.com + lensWorker.serviceTokenSecret.name: existing-lens-signing + lensWorker.gateway.secretName: existing-lens-signing + asserts: + - equal: + path: spec.rules[2].http.paths[0].path + value: /lens-ingest + - equal: + path: spec.rules[2].http.paths[0].backend.service.name + value: existing-lens + - equal: + path: spec.rules[2].http.paths[1].path + value: / + template: ingress.yaml + + - it: rejects a partial component host configuration + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + asserts: + - failedTemplate: + errorMessage: ingress.componentHosts.ui, backend, and gateway must all be set together + template: ingress.yaml + + - it: rejects component hosts when the gateway is disabled + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + gateway.enabled: false + asserts: + - failedTemplate: + errorMessage: ingress.componentHosts requires gateway.enabled=true + template: ingress.yaml + + - it: rejects component hosts when the backend is disabled + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + backend.enabled: false + asserts: + - failedTemplate: + errorMessage: ingress.componentHosts requires backend.enabled=true + template: ingress.yaml + + - it: rejects component hosts when the UI is disabled + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + ui.enabled: false + asserts: + - failedTemplate: + errorMessage: ingress.componentHosts requires ui.enabled=true + template: ingress.yaml + + - it: rejects a legacy host alongside component hosts + set: + ingress.enabled: true + ingress.host: legacy.example.com + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + asserts: + - failedTemplate: + errorMessage: ingress.host must be empty when ingress.componentHosts is configured + template: ingress.yaml + + - it: rejects extra paths alongside component hosts + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ingress.extraPaths: + - path: /custom + asserts: + - failedTemplate: + errorMessage: ingress.extraPaths must be empty when ingress.componentHosts is configured + template: ingress.yaml + + - it: requires the UI backend URL in component host mode + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + asserts: + - failedTemplate: + errorMessage: ui.backendUrl is required when ingress.componentHosts is configured + template: ingress.yaml + + - it: requires the backend gateway URL in component host mode + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + asserts: + - failedTemplate: + errorMessage: backend.gatewayUrl is required when ingress.componentHosts is configured + template: ingress.yaml + + - it: requires the public UI origin in component host mode + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + backend.gatewayUrl: https://gateway.example.com + asserts: + - failedTemplate: + errorMessage: ui.publicUrl is required when ingress.componentHosts is configured + template: ingress.yaml + + - it: exposes the gateway URL in the backend Deployment + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: " https://gateway.example.com/ " + template: backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_GATEWAY_URL + value: https://gateway.example.com + + - it: preserves the backend gateway URL verbatim in single-host mode + set: + ingress.enabled: true + ingress.host: litellm.example.com + backend.gatewayUrl: " https://gateway.example.com/api/ " + template: backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_GATEWAY_URL + value: " https://gateway.example.com/api/ " + + - it: exposes the backend URL in the UI Deployment + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: " https://backend.example.com/ " + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + template: ui/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BACKEND_URL + value: https://backend.example.com + + - it: omits the backend URL from the UI Deployment in single-host mode + set: + ingress.enabled: true + ingress.host: litellm.example.com + ui.backendUrl: https://litellm.example.com/api + template: ui/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_BACKEND_URL + any: true + + - it: enables credentialed CORS from the public UI origin on the backend + set: + ui.publicUrl: https://ui.example.com + template: backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ORIGINS + value: https://ui.example.com + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ALLOW_CREDENTIALS + value: "true" + + - it: lets backend extraEnv override generated CORS values + set: + ui.publicUrl: https://ui.example.com + backend.extraEnv: + - name: LITELLM_CORS_ORIGINS + value: https://backend-override.example.com + - name: LITELLM_CORS_ALLOW_CREDENTIALS + value: "false" + template: backend/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].env[0].name + value: LITELLM_CORS_ORIGINS + - equal: + path: spec.template.spec.containers[0].env[0].value + value: https://ui.example.com + - equal: + path: spec.template.spec.containers[0].env[1].name + value: LITELLM_CORS_ALLOW_CREDENTIALS + - equal: + path: spec.template.spec.containers[0].env[1].value + value: "true" + - equal: + path: spec.template.spec.containers[0].env[10].name + value: LITELLM_CORS_ORIGINS + - equal: + path: spec.template.spec.containers[0].env[10].value + value: https://backend-override.example.com + - equal: + path: spec.template.spec.containers[0].env[11].name + value: LITELLM_CORS_ALLOW_CREDENTIALS + - equal: + path: spec.template.spec.containers[0].env[11].value + value: "false" + + - it: exposes the public UI origin to the gateway without enabling credentials + set: + ui.publicUrl: https://ui.example.com + template: gateway/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ORIGINS + value: https://ui.example.com + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ALLOW_CREDENTIALS + value: "true" + + - it: trims the public UI origin before setting backend CORS + set: + ui.publicUrl: " https://ui.example.com/ " + template: backend/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ORIGINS + value: https://ui.example.com + + - it: trims the public UI origin before setting gateway CORS + set: + ui.publicUrl: " https://ui.example.com/ " + template: gateway/deployment.yaml + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_CORS_ORIGINS + value: https://ui.example.com + + - it: rejects a UI public URL with a path + set: + ui.publicUrl: https://ui.example.com/app + template: gateway/deployment.yaml + asserts: + - failedTemplate: + errorMessage: ui.publicUrl must be an origin like https://ui.example.com, without a path + + - it: lets gateway extraEnv override the generated CORS origin + set: + ui.publicUrl: https://ui.example.com + gateway.extraEnv: + - name: LITELLM_CORS_ORIGINS + value: https://gateway-override.example.com + template: gateway/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].env[0].name + value: LITELLM_CORS_ORIGINS + - equal: + path: spec.template.spec.containers[0].env[0].value + value: https://ui.example.com + - equal: + path: spec.template.spec.containers[0].env[9].name + value: LITELLM_CORS_ORIGINS + - equal: + path: spec.template.spec.containers[0].env[9].value + value: https://gateway-override.example.com + +--- +suite: component host UI backend origin validation +templates: + - ingress.yaml +values: + - ./values/required.yaml +tests: + - it: rejects a pathful UI backend URL in component host mode + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com/api + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com + asserts: + - failedTemplate: + errorMessage: ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set + - it: rejects a pathful backend gateway URL in component host mode + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: https://gateway.example.com/v1 + asserts: + - failedTemplate: + errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set + - it: rejects a backend gateway URL without a scheme in component host mode + set: + ingress.enabled: true + ingress.componentHosts.ui: ui.example.com + ingress.componentHosts.backend: backend.example.com + ingress.componentHosts.gateway: gateway.example.com + ui.backendUrl: https://backend.example.com + ui.publicUrl: https://ui.example.com + backend.gatewayUrl: gateway.example.com + asserts: + - failedTemplate: + errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set diff --git a/helm/litellm/tests/database_auth_tests.yaml b/helm/litellm/tests/database_auth_tests.yaml index add531f68bb..bc6653de700 100644 --- a/helm/litellm/tests/database_auth_tests.yaml +++ b/helm/litellm/tests/database_auth_tests.yaml @@ -191,3 +191,400 @@ tests: content: name: DATABASE_PASSWORD_READ_REPLICA any: true + + - it: writer hostKey reads the host from the Secret and supports IAM auth + template: gateway/deployment.yaml + set: + database.writer.host: "" + database.writer.passwordSecret.hostKey: hostname + database.writer.useIAMAuth: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: hostname + - contains: + path: spec.template.spec.containers[0].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PASSWORD + any: true + + - it: writer urlKey supplies only DATABASE_URL to each database client container and keeps pool env + set: + database.writer.host: "" + database.writer.dbname: "" + database.writer.schema: "" + database.writer.sslMode: "" + database.writer.sslRootCert: "" + database.writer.passwordSecret.urlKey: connection-url + database.connectionPool.enabled: true + gateway.collector.enabled: true + gateway.metricsServer.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_URL + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: connection-url + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PORT + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_USER + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_NAME + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SCHEMA + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLMODE + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SSLROOTCERT + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PASSWORD + any: true + template: gateway/deployment.yaml + - notExists: + path: >- + spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" || + @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" || + @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" || + @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" || + @.name == "DATABASE_PASSWORD")] + template: gateway/deployment.yaml + - notExists: + path: >- + spec.template.spec.containers[2].env[?(@.name == "DATABASE_HOST" || + @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" || + @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" || + @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" || + @.name == "DATABASE_PASSWORD")] + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[2].env + content: + name: DATABASE_URL + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: connection-url + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[2].env + content: + name: DATABASE_USER + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[2].env + content: + name: DATABASE_PASSWORD + any: true + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[2].env + content: + name: DATABASE_URL + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: connection-url + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_URL + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: connection-url + template: backend/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_USER + any: true + template: backend/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PASSWORD + any: true + template: backend/deployment.yaml + - notExists: + path: >- + spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" || + @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" || + @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" || + @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" || + @.name == "DATABASE_PASSWORD")] + template: backend/deployment.yaml + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_URL + valueFrom: + secretKeyRef: + name: litellm-writer-secret + key: connection-url + template: migrations-job.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_USER + any: true + template: migrations-job.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PASSWORD + any: true + template: migrations-job.yaml + - notExists: + path: >- + spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" || + @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" || + @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" || + @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" || + @.name == "DATABASE_PASSWORD")] + template: migrations-job.yaml + + - it: writer rejects host together with hostKey + template: gateway/deployment.yaml + set: + database.writer.passwordSecret.hostKey: hostname + asserts: + - failedTemplate: + errorMessage: "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" + + - it: writer rejects urlKey with discrete connection settings + template: gateway/deployment.yaml + set: + database.writer.host: "" + database.writer.passwordSecret.urlKey: connection-url + database.writer.schema: public + database.writer.sslMode: require + database.writer.sslRootCert: /etc/ssl/ca.crt + asserts: + - failedTemplate: + errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" + + - it: writer rejects urlKey with IAM auth + template: gateway/deployment.yaml + set: + database.writer.host: "" + database.writer.dbname: "" + database.writer.passwordSecret.urlKey: connection-url + database.writer.useIAMAuth: true + asserts: + - failedTemplate: + errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" + + - it: writer rejects urlKey with Entra auth + template: gateway/deployment.yaml + set: + database.writer.host: "" + database.writer.dbname: "" + database.writer.passwordSecret.urlKey: connection-url + database.writer.useAzureEntraAuth: true + asserts: + - failedTemplate: + errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" + + - it: writer requires an endpoint + template: gateway/deployment.yaml + set: + database.writer.host: "" + asserts: + - failedTemplate: + errorMessage: "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" + + - it: writer hostKey requires a Secret name + template: gateway/deployment.yaml + set: + database.writer.host: "" + database.writer.passwordSecret.hostKey: hostname + database.writer.passwordSecret.name: "" + asserts: + - failedTemplate: + errorMessage: "database.writer.passwordSecret.name is required when hostKey or urlKey is set" + + - it: reader renders without replica env vars when passwordSecret is null + template: gateway/deployment.yaml + set: + database.reader.passwordSecret: null + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST_READ_REPLICA + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_URL_READ_REPLICA + any: true + + - it: reader hostKey enables replica routing and reads its host from the Secret + template: gateway/deployment.yaml + set: + database.writer.useIAMAuth: true + database.reader.passwordSecret.hostKey: hostname + database.reader.dbname: litellm + database.reader.useIAMAuth: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST_READ_REPLICA + valueFrom: + secretKeyRef: + name: litellm-reader-secret + key: hostname + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_NAME_READ_REPLICA + value: litellm + - contains: + path: spec.template.spec.containers[0].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" + + - it: reader urlKey emits only the replica URL + template: gateway/deployment.yaml + set: + database.reader.passwordSecret.urlKey: connection-url + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_URL_READ_REPLICA + valueFrom: + secretKeyRef: + name: litellm-reader-secret + key: connection-url + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_HOST_READ_REPLICA + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PORT_READ_REPLICA + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_USER_READ_REPLICA + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_PASSWORD_READ_REPLICA + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_NAME_READ_REPLICA + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: DATABASE_SCHEMA_READ_REPLICA + any: true + + - it: reader rejects host together with hostKey + template: gateway/deployment.yaml + set: + database.reader.host: reader.example.com + database.reader.passwordSecret.hostKey: hostname + asserts: + - failedTemplate: + errorMessage: "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" + + - it: reader rejects urlKey with discrete connection settings + template: gateway/deployment.yaml + set: + database.reader.passwordSecret.urlKey: connection-url + database.reader.dbname: litellm + database.reader.schema: public + asserts: + - failedTemplate: + errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" + + - it: reader rejects urlKey with token auth + template: gateway/deployment.yaml + set: + database.writer.useIAMAuth: true + database.reader.passwordSecret.urlKey: connection-url + database.reader.useIAMAuth: true + asserts: + - failedTemplate: + errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" + + - it: reader urlKey requires a Secret name + template: gateway/deployment.yaml + set: + database.reader.passwordSecret.urlKey: connection-url + database.reader.passwordSecret.name: "" + asserts: + - failedTemplate: + errorMessage: "database.reader.passwordSecret.name is required when hostKey or urlKey is set" diff --git a/helm/litellm/tests/redis_env_tests.yaml b/helm/litellm/tests/redis_env_tests.yaml index 684d7071b35..c9304cedd9c 100644 --- a/helm/litellm/tests/redis_env_tests.yaml +++ b/helm/litellm/tests/redis_env_tests.yaml @@ -21,6 +21,27 @@ tests: name: REDIS_CLUSTER_NODES any: true + - it: gateway renders without redis env vars when passwordSecret is null + template: gateway/deployment.yaml + set: + redis.passwordSecret: null + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_URL + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + any: true + - it: gateway emits host, port and password when redis is configured template: gateway/deployment.yaml set: @@ -68,6 +89,27 @@ tests: name: redis-secret key: redis-password + - it: backend renders without redis env vars when passwordSecret is null + template: backend/deployment.yaml + set: + redis.passwordSecret: null + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_URL + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + any: true + - it: gateway omits REDIS_PASSWORD for an auth-less redis template: gateway/deployment.yaml set: @@ -107,3 +149,174 @@ tests: content: name: REDIS_CLUSTER_NODES any: true + + - it: gateway gets redis host from hostKey and keeps port and password + template: gateway/deployment.yaml + set: + redis.host: "" + redis.port: 6380 + redis.passwordSecret.name: redis-secret + redis.passwordSecret.hostKey: hostname + redis.passwordSecret.passwordKey: redis-password + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + valueFrom: + secretKeyRef: + name: redis-secret + key: hostname + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PORT + value: "6380" + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: redis-secret + key: redis-password + + - it: backend gets redis host from hostKey with the same coordination env + template: backend/deployment.yaml + set: + redis.host: "" + redis.passwordSecret.name: redis-secret + redis.passwordSecret.hostKey: hostname + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + valueFrom: + secretKeyRef: + name: redis-secret + key: hostname + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PORT + value: "6379" + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: redis-secret + key: password + + - it: gateway expands the redis cluster seed from hostKey at runtime + template: gateway/deployment.yaml + set: + redis.host: "" + redis.passwordSecret.name: redis-secret + redis.passwordSecret.hostKey: hostname + redis.cluster: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_CLUSTER_NODES + value: '[{"host":"$(REDIS_HOST)","port":6379}]' + - equal: + path: spec.template.spec.containers[0].env[8].name + value: REDIS_HOST + - equal: + path: spec.template.spec.containers[0].env[11].name + value: REDIS_CLUSTER_NODES + + - it: gateway gets REDIS_URL from urlKey without discrete Redis env + template: gateway/deployment.yaml + set: + redis.passwordSecret.name: redis-secret + redis.passwordSecret.urlKey: connection-url + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_URL + valueFrom: + secretKeyRef: + name: redis-secret + key: connection-url + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_HOST + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PORT + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_PASSWORD + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: REDIS_CLUSTER_NODES + any: true + + - it: redis rejects host together with hostKey + template: gateway/deployment.yaml + set: + redis.host: redis.example.com + redis.passwordSecret.hostKey: hostname + redis.passwordSecret.name: redis-secret + asserts: + - failedTemplate: + errorMessage: "redis.host and redis.passwordSecret.hostKey cannot both be set" + + - it: redis urlKey rejects host + template: gateway/deployment.yaml + set: + redis.host: redis.example.com + redis.passwordSecret.urlKey: connection-url + redis.passwordSecret.name: redis-secret + asserts: + - failedTemplate: + errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" + + - it: redis urlKey rejects hostKey + template: gateway/deployment.yaml + set: + redis.passwordSecret.hostKey: hostname + redis.passwordSecret.urlKey: connection-url + redis.passwordSecret.name: redis-secret + asserts: + - failedTemplate: + errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" + + - it: redis urlKey rejects cluster mode + template: gateway/deployment.yaml + set: + redis.cluster: true + redis.passwordSecret.urlKey: connection-url + redis.passwordSecret.name: redis-secret + asserts: + - failedTemplate: + errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" + + - it: redis hostKey requires a Secret name + template: gateway/deployment.yaml + set: + redis.passwordSecret.hostKey: hostname + asserts: + - failedTemplate: + errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set" + + - it: redis urlKey requires a Secret name + template: gateway/deployment.yaml + set: + redis.passwordSecret.urlKey: connection-url + asserts: + - failedTemplate: + errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set" diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index a3c620b3434..d8e77085010 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -23,7 +23,11 @@ ingress: # wildcard pathType, so that rule could never match there. controller: alb annotations: {} - host: "" # optional; if set, becomes the rule's host + host: "" # optional; if set, becomes the single ingress rule's host + componentHosts: + ui: "" + backend: "" + gateway: "" tls: [] # Extra HTTP paths appended to the ingress rule. Additive: every built-in # UI / gateway / backend path is still rendered, these entries are placed @@ -208,6 +212,8 @@ database: name: litellm-writer-secret usernameKey: username passwordKey: password + hostKey: "" # Secret key containing the writer database host + urlKey: "" # Secret key containing the complete writer database URL # libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the # in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS. # sslRootCert on its own implies sslMode verify-full @@ -229,6 +235,8 @@ database: name: litellm-reader-secret usernameKey: username passwordKey: password + hostKey: "" # Secret key containing the reader database host + urlKey: "" # Secret key containing the complete reader database URL # In-container connection pool (PgBouncer, transaction mode) shared by every # gateway worker in the pod. Without it each of the `gateway.numWorkers` @@ -274,6 +282,8 @@ redis: passwordSecret: name: "" # Leave empty for auth-less Redis passwordKey: password + hostKey: "" # Secret key containing the Redis host + urlKey: "" # Secret key containing the complete Redis URL # ---------- gateway (LLM data plane) ---------- gateway: @@ -493,6 +503,7 @@ gateway: backend: enabled: true logLevel: INFO + gatewayUrl: "" extraEnv: [] envConfigMaps: [] envSecrets: [] @@ -575,9 +586,10 @@ ui: service: type: ClusterIP port: 3000 - # The dashboard expects to know where to reach the backend API. Set this to - # the externally-routable URL (typically the ingress host + /api or similar). + # The backend's public origin, passed to the UI only when ingress.componentHosts is configured. backendUrl: "" + # Public origin of the UI, used for backend and gateway CORS. + publicUrl: "" resources: requests: cpu: 500m diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py index 671ce0f5c3e..89857859134 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py @@ -13,6 +13,7 @@ from litellm_proxy_extras.prisma_toolchain import ( ) MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big") +PRISMA_ADVISORY_LOCK_KEY: Final = 72707369 if TYPE_CHECKING: import psycopg @@ -32,7 +33,7 @@ def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]: if not database_url: return environment migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL")) - return {**environment, "DATABASE_URL": migration_url} + return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"} @dataclass(frozen=True, slots=True) @@ -55,10 +56,15 @@ class MigrationCoordinator: def check_connection(self) -> None: self.connection.execute("SELECT 1") + def release_statement_snapshot(self) -> None: + # a parameterized statement leaves its portal snapshot registered on the backend until the + # next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away + self.connection.execute("SELECT 1") + def acquire_prisma_lock(self) -> None: deadline: Final = time.monotonic() + migration_lock_timeout() while time.monotonic() < deadline: - if _try_lock(self.connection, 72707369): + if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY): return time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic()))) raise RuntimeError( @@ -95,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N ) +@contextmanager +def prisma_migration_lock(database_url: str) -> Generator[None, None, None]: + """Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one + pinned connection while a Prisma command runs with its session advisory lock disabled, so a + transaction pooler can never strand that lock on an idle backend after the command exits.""" + with migration_lock(database_url) as coordinator: + coordinator.acquire_prisma_lock() + coordinator.release_statement_snapshot() + yield + + @contextmanager def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]: """A session-level, non-blocking hold of the migration coordinator lock on an autocommit diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql new file mode 100644 index 00000000000..35197a9e63c --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql @@ -0,0 +1,4 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_add_access_group_skill_names/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_add_access_group_skill_names/migration.sql new file mode 100644 index 00000000000..eafb75d83ef --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_add_access_group_skill_names/migration.sql @@ -0,0 +1,4 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_skill_names" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index ca58247a2c9..12438614755 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable { // Resource memberships - explicit arrays per type access_model_names String[] @default([]) access_mcp_server_ids String[] @default([]) + mcp_tool_permissions Json? access_agent_ids String[] @default([]) + access_skill_names String[] @default([]) assigned_team_ids String[] @default([]) assigned_key_ids String[] @default([]) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2a109130ff8..4a311331fc0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -9,6 +9,7 @@ import tempfile import threading import time from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, replace from pathlib import Path from typing import TYPE_CHECKING, Final, Optional, Protocol, Union @@ -1021,7 +1022,11 @@ class ProxyExtrasDBManager: def _setup_database_v2(use_migrate: bool) -> bool: if not use_migrate: return ProxyExtrasDBManager._run_database_v2(False) - from litellm_proxy_extras.migration_lock import migration_environment, migration_lock + from litellm_proxy_extras.migration_lock import ( + migration_environment, + migration_lock, + prisma_migration_lock, + ) from litellm_proxy_extras.migration_recovery import ( baseline_current_schema, recover_completed_migration, @@ -1055,7 +1060,9 @@ class ProxyExtrasDBManager: migration_environment(_get_prisma_env()), ) - while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing): + while not ProxyExtrasDBManager._run_database_v2( + True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url) + ): continue return True @@ -1064,6 +1071,7 @@ class ProxyExtrasDBManager: use_migrate: bool, recover_completed: Callable[[str], bool] = lambda name: False, baseline_existing: "Callable[[str], None] | None" = None, + deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext, ) -> bool: """ v2 migration resolver (opt-in via --use_v2_migration_resolver). @@ -1126,11 +1134,12 @@ class ProxyExtrasDBManager: while not budget.exhausted: attempt_started = time.monotonic() try: - result = prisma_toolchain.run_prisma( - [_get_prisma_command(), "migrate", "deploy"], - timeout=deploy_timeout, - env=migration_env, - ) + with deploy_lock(): + result = prisma_toolchain.run_prisma( + [_get_prisma_command(), "migrate", "deploy"], + timeout=deploy_timeout, + env=migration_env, + ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") return True @@ -1155,6 +1164,7 @@ class ProxyExtrasDBManager: e, budget, time.monotonic() - attempt_started, + deploy_lock=deploy_lock, ) if next_budget.attempts_left < budget.attempts_left: @@ -1187,15 +1197,19 @@ class ProxyExtrasDBManager: return match.group(1) if match else None @staticmethod - def _v2_roll_back_migration_best_effort(migration_name: str) -> None: + def _v2_roll_back_migration_best_effort( + migration_name: str, + deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext, + ) -> None: from litellm_proxy_extras.migration_lock import migration_environment try: - prisma_toolchain.run_prisma( - [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], - timeout=prisma_command_timeout(), - env=migration_environment(_get_prisma_env()), - ) + with deploy_lock(): + prisma_toolchain.run_prisma( + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + timeout=prisma_command_timeout(), + env=migration_environment(_get_prisma_env()), + ) except (subprocess.CalledProcessError, subprocess.TimeoutExpired): pass @@ -1204,6 +1218,7 @@ class ProxyExtrasDBManager: error: subprocess.CalledProcessError, budget: "_MigrateAttemptBudget", attempt_seconds: float = 0.0, + deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext, ) -> "_MigrateAttemptBudget": """Recover from one failed `prisma migrate deploy`, and price the pass. @@ -1225,14 +1240,14 @@ class ProxyExtrasDBManager: "deadlock race, rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs: logger.info( "Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at): logger.info( @@ -1276,7 +1291,7 @@ class ProxyExtrasDBManager: "rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr: @@ -1286,7 +1301,7 @@ class ProxyExtrasDBManager: "its ledger row back and retrying. Check for long-running transactions on the migrated tables.", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() raise RuntimeError( diff --git a/litellm/_redis.py b/litellm/_redis.py index 499d713e5e4..186f8523519 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -675,7 +675,10 @@ def _get_redis_client_logic(**env_overrides): if redis_kwargs.get("startup_nodes") else redis.Redis ) - return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client) + coerced_kwargs: Final = _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client) + if coerced_kwargs.get("socket_timeout") is None: + coerced_kwargs["socket_timeout"] = REDIS_SOCKET_TIMEOUT + return coerced_kwargs def init_redis_cluster(redis_kwargs) -> redis.RedisCluster: @@ -722,7 +725,6 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis: sentinel_password: Final = redis_kwargs.get("sentinel_password") service_name: Final = redis_kwargs.get("service_name") connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs) - connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT) sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password) if not sentinel_nodes or not service_name: @@ -755,7 +757,6 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: sentinel_password: Final = redis_kwargs.get("sentinel_password") service_name: Final = redis_kwargs.get("service_name") connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs) - connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT) sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password) if not sentinel_nodes or not service_name: diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 6385aa96272..060b4fb878e 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -152,6 +152,11 @@ class ServiceLogging(CustomLogger): except Exception as e: verbose_logger.exception("Error dispatching service event - %s", e) + def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None: + prometheus_logger: Final[PrometheusServicesLogger | None] = getattr(self, "prometheusServicesLogger", None) + if prometheus_logger is not None: + prometheus_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class) + def service_success_hook( self, service: ServiceTypes, diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 0e34709d128..6027f3774fd 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,7 +18,9 @@ import json import logging import threading import time -from collections.abc import Awaitable, Callable, Iterator, Sequence +import weakref +from collections import Counter +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -51,12 +53,12 @@ from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from prometheus_client import Counter as _PromCounter - from prometheus_client import Gauge as _PromGauge from redis.asyncio import Redis, RedisCluster from redis.asyncio.client import Pipeline from redis.asyncio.cluster import ClusterPipeline + from litellm._service_logger import ServiceLogging + pipeline = Pipeline cluster_pipeline = ClusterPipeline async_redis_client = Redis @@ -243,7 +245,9 @@ class RedisCircuitBreaker: recovery_timeout: int, enabled: bool = True, timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, + service_logger: "ServiceLogging | None" = None, ) -> None: + self._service_logger = service_logger self.failure_threshold = failure_threshold self.recovery_timeout = recovery_timeout self.enabled = enabled @@ -255,7 +259,12 @@ class RedisCircuitBreaker: self._opened_at: float | None = None self._state = self.CLOSED self._generation = 0 - _breaker_metrics().record_state_change(None, self._state) + _LIVE_BREAKERS.add(self) + self._report(transitions_to=None, failure_class=None) + + @property + def state(self) -> str: + return self._state @property def generation(self) -> int: @@ -301,7 +310,7 @@ class RedisCircuitBreaker: self._timeout_count = 0 self._timeout_streak_started_at = None self._opened_at = now - _breaker_metrics().record_failure("timeout" if is_timeout else "connectivity") + self._report(transitions_to=None, failure_class="timeout" if is_timeout else "connectivity") if self._should_open(now): if self._state != self.OPEN: verbose_logger.warning( @@ -327,10 +336,13 @@ class RedisCircuitBreaker: def _set_state(self, state: str) -> None: if state == self._state: return - _breaker_metrics().record_transition(state) - _breaker_metrics().record_state_change(self._state, state) self._state = state self._generation += 1 + self._report(transitions_to=state, failure_class=None) + + def _report(self, transitions_to: str | None, failure_class: str | None) -> None: + if self._service_logger is not None: + self._service_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class) _RedisCallResult = TypeVar("_RedisCallResult") @@ -400,12 +412,19 @@ def _redis_timeout_error_types() -> tuple[type, ...]: ``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither. + ``MaxConnectionsError`` is the async cluster client's pool-exhausted signal (its nodes have + no blocking pool, so it is raised instantly instead of after a wait): a busy pool, not a + dead Redis, like the standalone pool's chained wait timeout. """ try: - from redis.exceptions import TimeoutError as RedisTimeoutError + from redis import exceptions as redis_exceptions except ImportError: return (TimeoutError, asyncio.TimeoutError) - return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError) + max_connections_error: Final = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError + "type[Exception]", + redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # types-redis 4.6 stubs predate it + ) + return (redis_exceptions.TimeoutError, TimeoutError, asyncio.TimeoutError, max_connections_error) _MAX_EXCEPTION_CAUSE_DEPTH: Final = 20 @@ -430,60 +449,6 @@ def is_redis_timeout_failure(exc: BaseException) -> bool: return any(isinstance(link, timeout_types) for link in _explicit_causes(exc)) -class _BreakerMetrics: - """Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent. - - Registered lazily on the default registry (which /metrics serves) via the module-level - ``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers. - """ - - def __init__(self) -> None: - self._state_gauge: _PromGauge | None = None - self._transitions: _PromCounter | None = None - self._failures: _PromCounter | None = None - try: - from prometheus_client import Counter as PromCounter - from prometheus_client import Gauge - except ImportError: - return - self._state_gauge = Gauge( - "litellm_redis_circuit_breaker_state", - "Number of Redis circuit breakers currently in each state", - labelnames=("state",), - multiprocess_mode="livesum", - ) - self._transitions = PromCounter( - "litellm_redis_circuit_breaker_transitions", - "Redis circuit breaker state transitions", - labelnames=("state",), - ) - self._failures = PromCounter( - "litellm_redis_circuit_breaker_failures", - "Redis health failures counted by the circuit breaker", - labelnames=("failure_class",), - ) - - def record_state_change(self, old_state: str | None, new_state: str) -> None: - if self._state_gauge is None: - return - if old_state is not None: - self._state_gauge.labels(old_state).dec() - self._state_gauge.labels(new_state).inc() - - def record_transition(self, state: str) -> None: - if self._transitions is not None: - self._transitions.labels(state).inc() - - def record_failure(self, failure_class: str) -> None: - if self._failures is not None: - self._failures.labels(failure_class).inc() - - -@functools.lru_cache(maxsize=1) -def _breaker_metrics() -> _BreakerMetrics: - return _BreakerMetrics() - - def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None: """Record a Redis failure that the calling method is about to swallow. @@ -499,6 +464,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet() + + +def redis_breaker_state_counts() -> Mapping[str, int]: + """How many live Redis circuit breakers sit in each state, for the Prometheus services logger.""" + return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS))) + + class RedisCircuitBreakerOpenError(Exception): pass @@ -638,9 +611,14 @@ def _redis_circuit_breaker_guard(method): """ @functools.wraps(method) - async def wrapper(self, *args, **kwargs): + async def wrapper(self: "RedisCache", *args, **kwargs): + async def guarded() -> object: + return await method(self, *args, **kwargs) + return await _run_under_circuit_breaker( - self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs) + self._circuit_breaker, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator + wrapper.__name__, + lambda: self._admit_async_call(guarded), # pyright: ignore[reportPrivateUsage] # same decorator ) return wrapper @@ -685,12 +663,9 @@ class RedisCache(BaseCache): redis_kwargs["socket_timeout"] = socket_timeout ### HEALTH MONITORING OBJECT ### - if kwargs.get("service_logger_obj", None) is not None and isinstance( - kwargs["service_logger_obj"], ServiceLogging - ): + self.service_logger_obj: ServiceLogging = ServiceLogging() + if isinstance(kwargs.get("service_logger_obj"), ServiceLogging): self.service_logger_obj = kwargs.pop("service_logger_obj") - else: - self.service_logger_obj = ServiceLogging() redis_kwargs.update(kwargs) self.redis_client = get_redis_client(**redis_kwargs) @@ -717,6 +692,7 @@ class RedisCache(BaseCache): failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, enabled=REDIS_CIRCUIT_BREAKER_ENABLED, + service_logger=self.service_logger_obj, ) self._setup_health_pings() @@ -726,6 +702,9 @@ class RedisCache(BaseCache): else: super().__init__() # defaults to 60s + def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: + return call() + def _setup_health_pings(self): """Setup async and sync health pings for Redis.""" # ASYNC HEALTH PING @@ -1106,7 +1085,9 @@ class RedisCache(BaseCache): litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor) return await executor(keys=keys, args=args, client=client) - return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute) + return await _run_under_circuit_breaker( + self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute) + ) return run_script diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index d0f3c11bf10..5e926fd2ed8 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -5,9 +5,17 @@ Key differences: - RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created """ -from typing import TYPE_CHECKING, Any, Final +import asyncio +import contextvars +import os +import weakref +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar +from litellm._logging import verbose_logger from litellm.caching.redis_cache import RedisCache +from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -23,9 +31,100 @@ else: Span = Any +_T: Final = TypeVar("_T") + + +def _connection_limit(max_connections: object) -> int | None: + try: + limit: Final = int(max_connections) # pyright: ignore[reportArgumentType] # config value, str or int + except (TypeError, ValueError): + return None + return limit if limit > 0 else None + + +_ClientSemaphores: TypeAlias = weakref.WeakValueDictionary[str, asyncio.Semaphore] + + +@dataclass(slots=True) +class _SemaphoreRegistry: + """Weak on both layers: a semaphore binds to its loop once a waiter queues on it, so a strong value would keep a + closed loop alive. In-flight ``ClusterConnectionGate.run`` calls hold the semaphore they acquired.""" + + _per_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _ClientSemaphores]" = field( + default_factory=weakref.WeakKeyDictionary, repr=False + ) + + def get(self, client_key: str, limit: int) -> asyncio.Semaphore: + per_loop: Final = self._per_loop.setdefault(asyncio.get_running_loop(), weakref.WeakValueDictionary()) + existing: Final = per_loop.get(client_key) + if existing is not None: + return existing + created: Final = asyncio.Semaphore(limit) + per_loop[client_key] = created + return created + + +_Admission: TypeAlias = tuple[str, "asyncio.Task[object] | None"] +_SEMAPHORES: Final = _SemaphoreRegistry() +_ADMITTED: Final[contextvars.ContextVar[frozenset[_Admission]]] = contextvars.ContextVar( + "redis_cluster_gate_admitted", default=frozenset() +) + + +@dataclass(frozen=True, slots=True) +class ClusterConnectionGate: + """The async ``RedisCluster`` has no blocking pool, so past ``max_connections`` a node raises + ``MaxConnectionsError`` instead of waiting and the circuit breaker opens on a healthy Redis. This gate makes + guarded commands wait for a slot the way the standalone ``BlockingConnectionPool`` does. A command that finds no + slot within ``wait_timeout`` runs anyway, so a genuinely saturated pool still reaches the client's own error + handling and the breaker.""" + + limit: int + wait_timeout: float = REDIS_CONNECTION_POOL_TIMEOUT + + @classmethod + def from_max_connections(cls, max_connections: object) -> "ClusterConnectionGate | None": + limit: Final = _connection_limit(max_connections) + return None if limit is None else cls(limit=limit) + + async def run(self, call: Callable[[], Awaitable[_T]], client_key: str) -> _T: + admission: Final[_Admission] = (client_key, asyncio.current_task()) + if admission in _ADMITTED.get(): + return await call() + semaphore: Final = _SEMAPHORES.get(client_key, self.limit) + try: + await asyncio.wait_for(semaphore.acquire(), timeout=self.wait_timeout) + except asyncio.TimeoutError: + verbose_logger.debug( + "Redis Cluster: no free connection slot within %ss (max_connections=%s), sending anyway", + self.wait_timeout, + self.limit, + ) + return await call() + token: Final = _ADMITTED.set(_ADMITTED.get() | {admission}) + try: + return await call() + finally: + _ADMITTED.reset(token) + semaphore.release() + + class RedisClusterCache(RedisCache): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, *args, max_connections: object = None, connection_gate: object = None, **kwargs): + pool_kwargs: Final = {} if max_connections is None else {"max_connections": max_connections} + super().__init__(*args, **pool_kwargs, **kwargs) + self._connection_gate: Final = ( + connection_gate + if isinstance(connection_gate, ClusterConnectionGate) + else ClusterConnectionGate.from_max_connections( + os.environ.get("REDIS_MAX_CONNECTIONS") if max_connections is None else max_connections + ) + ) + + def _admit_async_call(self, call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]: + if self._connection_gate is None: + return call() + return self._connection_gate.run(call, self._get_async_client_cache_key()) def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]: """ diff --git a/litellm/constants.py b/litellm/constants.py index 4ac6adaba73..e78e3e36364 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -60,6 +60,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( "enable_tag_filtering", "tag_routing_prefix", "optional_pre_call_checks", + "fallback_on_unsupported_params", } ) ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( @@ -75,6 +76,9 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( } ) DEFAULT_BATCH_SIZE: Final = get_env_int("DEFAULT_BATCH_SIZE", 512) +SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE: Final = get_env_int_in_range( + "SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE", 1000, 1, 10_000 +) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_FLUSH_INTERVAL_SECONDS", 5) CLICKHOUSE_BATCH_SIZE: Final = get_env_int("CLICKHOUSE_BATCH_SIZE", 10_000) CLICKHOUSE_FLUSH_INTERVAL_SECONDS: Final = get_env_float("CLICKHOUSE_FLUSH_INTERVAL_SECONDS", 1.0) @@ -547,7 +551,7 @@ DEFAULT_POLLING_INTERVAL: Final = get_env_float( AZURE_OPERATION_POLLING_TIMEOUT: Final = get_env_int("AZURE_OPERATION_POLLING_TIMEOUT", 120) AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30")) AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = get_env_int("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96) -REDIS_SOCKET_TIMEOUT: Final = get_env_float("REDIS_SOCKET_TIMEOUT", 0.1) +REDIS_SOCKET_TIMEOUT: Final = get_env_float("REDIS_SOCKET_TIMEOUT", 5.0) CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0 REDIS_CONNECTION_POOL_TIMEOUT: Final = get_env_int("REDIS_CONNECTION_POOL_TIMEOUT", 5) REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = get_env_int("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index ed7593c4961..126d84a0e10 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -799,12 +799,16 @@ class AnthropicCacheControlHook(CustomPromptManagement): The hook sees neither the tools nor the request kwargs, so the target dialect and the client's breakpoint count outside the message list ride on the points. - Builds copies because config-owned point dicts are shared across requests. + A zero count replaces an earlier layer's stamp (the chat bridge runs behind the + Responses layer, whose outside marks sit in the bridge's own message list) and is + otherwise left off. Builds copies because config-owned point dicts are shared + across requests. """ with_dialect: Final = AnthropicCacheControlHook._stamped_with_dialect( points, model, custom_llm_provider, api_base, prompt_cache_options ) - if external_breakpoints == 0: + stamped_earlier: Final = any(EXTERNAL_BREAKPOINTS_STAMP in point for point in with_dialect) + if external_breakpoints == 0 and not stamped_earlier: return with_dialect return AnthropicCacheControlHook._stamped(with_dialect, EXTERNAL_BREAKPOINTS_STAMP, external_breakpoints) diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index d0db4888c50..42a3ffac855 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -3,7 +3,9 @@ # On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers) -from typing import Final +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Final import litellm from litellm._logging import print_verbose, verbose_logger @@ -17,6 +19,42 @@ from litellm.types.services import ( FAILED_REQUESTS_LABELS: Final = ["error_class", "function_name"] +if TYPE_CHECKING: + from prometheus_client import Counter as PrometheusCounter + from prometheus_client import Gauge + + +@dataclass(frozen=True, slots=True) +class _RedisBreakerCollectors: + state: "Gauge" + transitions: "PrometheusCounter" + failures: "PrometheusCounter" + + +@lru_cache(maxsize=1) +def _redis_breaker_collectors() -> _RedisBreakerCollectors: + from prometheus_client import Counter as PrometheusCounter + from prometheus_client import Gauge + + return _RedisBreakerCollectors( + state=Gauge( + "litellm_redis_circuit_breaker_state", + "Number of Redis circuit breakers currently in each state", + labelnames=("state",), + multiprocess_mode="livesum", + ), + transitions=PrometheusCounter( + "litellm_redis_circuit_breaker_transitions", + "Redis circuit breaker state transitions", + labelnames=("state",), + ), + failures=PrometheusCounter( + "litellm_redis_circuit_breaker_failures", + "Redis health failures counted by the circuit breaker", + labelnames=("failure_class",), + ), + ) + class PrometheusServicesLogger: # Class variables or attributes @@ -41,6 +79,11 @@ class PrometheusServicesLogger: self.Counter = Counter self.Gauge = Gauge self.REGISTRY = REGISTRY + breaker_collectors: Final = _redis_breaker_collectors() + self.redis_breaker_state = breaker_collectors.state + self.redis_breaker_transitions = breaker_collectors.transitions + self.redis_breaker_failures = breaker_collectors.failures + self._refresh_redis_breaker_gauge() verbose_logger.debug("in init prometheus services metrics") @@ -98,6 +141,20 @@ class PrometheusServicesLogger: return DEFAULT_METRICS return metrics + def _refresh_redis_breaker_gauge(self) -> None: + from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts + + counts: Final = redis_breaker_state_counts() + for state in (RedisCircuitBreaker.CLOSED, RedisCircuitBreaker.OPEN, RedisCircuitBreaker.HALF_OPEN): + self.redis_breaker_state.labels(state).set(counts.get(state, 0)) + + def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None: + self._refresh_redis_breaker_gauge() + if transitions_to is not None: + self.redis_breaker_transitions.labels(transitions_to).inc() + if failure_class is not None: + self.redis_breaker_failures.labels(failure_class).inc() + def is_metric_registered(self, metric_name) -> bool: # Use _names_to_collectors (O(1)) instead of REGISTRY.collect() (O(n)) to avoid # perf regression when a new Router is created per request (e.g. router_settings in DB). diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 6f16e1e3fb4..ad33c47fc51 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -6,7 +6,7 @@ import traceback from collections.abc import Mapping, Sequence from typing import Final, Literal, cast -from pydantic import ConfigDict, TypeAdapter, with_config +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from typing_extensions import ReadOnly, TypedDict import litellm @@ -222,6 +222,9 @@ async def convert_to_streaming_response_async( role=choice["message"]["role"], function_call=choice["message"].get("function_call", None), tool_calls=choice["message"].get("tool_calls", None), + provider_specific_fields=_cached_provider_specific_fields( + _CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(choice["message"]) + ), ) finish_reason = choice.get("finish_reason", None) @@ -243,6 +246,9 @@ async def convert_to_streaming_response_async( completion_tokens=response_object["usage"].get("completion_tokens", 0), prompt_tokens=response_object["usage"].get("prompt_tokens", 0), total_tokens=response_object["usage"].get("total_tokens", 0), + completion_tokens_details=_cached_completion_tokens_details( + _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + ), ), ) @@ -334,6 +340,9 @@ def convert_to_streaming_response( model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) + model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details( + _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + ) if "id" in response_object: model_response_object.id = response_object["id"] @@ -390,6 +399,35 @@ class _ParallelToolUseArgs(TypedDict): _PARALLEL_TOOL_USE_ARGS: Final = TypeAdapter(_ParallelToolUseArgs) +class _CachedMessageProviderFields(TypedDict, total=False): + provider_specific_fields: ReadOnly[object] + + +class _CachedUsageDetails(TypedDict, total=False): + completion_tokens_details: ReadOnly[object] + + +_CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields) +_CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails) +_FIELD_MAP: Final = TypeAdapter(dict[str, object]) +_COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper) + + +def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None: + fields: Final = message.get("provider_specific_fields") + return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else None + + +def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionTokensDetailsWrapper | None: + details: Final = usage.get("completion_tokens_details") + if not isinstance(details, dict): + return None + try: + return _COMPLETION_TOKENS_DETAILS.validate_python(details) + except ValidationError: + return None + + def handle_invalid_parallel_tool_calls( tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall], ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 236286a493b..484978bae04 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -280,6 +280,7 @@ class CustomStreamWrapper: self._response_headers = _response_headers self.response_id: str | None = None + self.received_model_id: str | None = None self.logging_loop = None self.rules = Rules() self.stream_options = stream_options or getattr(logging_obj, "stream_options", None) @@ -838,6 +839,7 @@ class CustomStreamWrapper: self.response_id = id if id and isinstance(id, str) and id.strip(): + self.received_model_id = id model_response.hidden_params["received_model_id"] = id if self.response_id is not None and isinstance(self.response_id, str): @@ -1756,7 +1758,9 @@ class CustomStreamWrapper: self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) def finish_reason_handler(self): - model_response: Final = self.model_response_creator() + model_response: Final = self.model_response_creator( + hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} + ) _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason if _finish_reason is not None: model_response.choices[0].finish_reason = _finish_reason diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 7c2e7b39d57..9558eb6036d 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -66,7 +66,7 @@ from litellm.types.utils import ( ) from ...base import BaseLLM -from ..common_utils import AnthropicError, process_anthropic_headers +from ..common_utils import AnthropicError, process_anthropic_headers, refusal_provider_specific_fields from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig, JsonModeRequest if TYPE_CHECKING: @@ -1104,7 +1104,8 @@ class ModelResponseIterator: # These are automatically handled by Anthropic API, we just pass them through pass elif type_chunk == "message_delta": - finish_reason, usage, container = self._handle_message_delta(chunk) + finish_reason, usage, container, refusal_fields = self._handle_message_delta(chunk) + provider_specific_fields.update(refusal_fields) text = self._release_answer_close() + self._drain_split_answer_blocks() if container: provider_specific_fields["container"] = container @@ -1299,7 +1300,9 @@ class ModelResponseIterator: return "" return _split_answer_text(blocks, self.response_format_schema) - def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]: + def _handle_message_delta( + self, chunk: dict + ) -> tuple[str, Usage | None, dict[str, Any] | None, Mapping[str, object]]: """ Handle message_delta event for finish_reason, usage, and container. @@ -1321,7 +1324,7 @@ class ModelResponseIterator: self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None ) container: Final = message_delta["delta"].get("container") - return finish_reason, usage, container + return finish_reason, usage, container, refusal_provider_specific_fields(message_delta["delta"]) def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None: """ diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e4d3fac753e..932de83c0c9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -104,10 +104,12 @@ from litellm.utils import ( ) from ..common_utils import ( + REFUSAL_STOP_FIELDS, AnthropicError, AnthropicModelInfo, eager_input_streaming_flag, process_anthropic_headers, + refusal_provider_specific_fields, requires_native_compaction_beta, strip_advisor_blocks_from_messages, ) @@ -294,7 +296,6 @@ DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING: Final = ( DROP_UNSUPPORTED_SPEED_WARNING: Final = ( "Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models." ) - SERVICE_TIER_PROVIDERS: Final = frozenset({"anthropic", "azure_ai"}) OPENAI_TO_ANTHROPIC_SERVICE_TIER: Final = MappingProxyType( @@ -306,6 +307,12 @@ DROP_UNSUPPORTED_SERVICE_TIER_WARNING: Final = ( "Anthropic only accepts 'auto' and 'standard_only'." ) +MESSAGES_SERVICE_TIER_PROVIDERS: Final = SERVICE_TIER_PROVIDERS | {"bedrock_mantle"} + +DROP_SERVICE_TIER_PROVIDER_WARNING: Final = ( + "Dropping `service_tier` for model=%s on provider=%s. Only the Anthropic API routes accept it." +) + class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ @@ -539,6 +546,24 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) -> None: return cls._maybe_drop_speed_param(model, optional_params, drop_params, custom_llm_provider) + @staticmethod + def provider_supports_service_tier(custom_llm_provider: str | None) -> bool: + """Anthropic's ``service_tier`` is accepted by the direct API and the Foundry and Mantle gateways in + front of it. Vertex rejects the field and Bedrock Invoke only takes its own tier names.""" + return custom_llm_provider is None or custom_llm_provider in MESSAGES_SERVICE_TIER_PROVIDERS + + @classmethod + def maybe_drop_service_tier_param( + cls, + model: str, + optional_params: dict[str, object], # mutable-ok: mirrors override contract + custom_llm_provider: str | None = None, + ) -> None: + if "service_tier" not in optional_params or cls.provider_supports_service_tier(custom_llm_provider): + return + litellm.verbose_logger.warning(DROP_SERVICE_TIER_PROVIDER_WARNING, model, custom_llm_provider) + optional_params.pop("service_tier", None) + @staticmethod def _raise_invalid_reasoning_effort(model: str, value: object, llm_provider: str) -> NoReturn: """Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``. @@ -717,12 +742,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Already handled above continue - if key == "properties" and isinstance(value, dict): - result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()} + if key in ("properties", "$defs") and isinstance(value, dict): + nested_schemas: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict + result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in nested_schemas.items()} elif key == "items" and isinstance(value, dict): - result[key] = AnthropicConfig.filter_anthropic_output_schema(value) - elif key == "$defs" and isinstance(value, dict): - result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()} + nested_schema: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict + result[key] = AnthropicConfig.filter_anthropic_output_schema(nested_schema) elif key == "anyOf" and isinstance(value, list): result["anyOf"] = result.get("anyOf", []) + [ AnthropicConfig.filter_anthropic_output_schema(item) for item in value @@ -2700,7 +2725,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if compaction_blocks is not None: provider_specific_fields["compaction_blocks"] = compaction_blocks - return provider_specific_fields + refusal_fields: Final = refusal_provider_specific_fields( + REFUSAL_STOP_FIELDS.validate_python(completion_response) + ) + return {**provider_specific_fields, **refusal_fields} def transform_parsed_response( self, @@ -2783,7 +2811,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _message.provider_specific_fields = provider_specific_fields if json_mode_message is not None: + refusal_fields: Final = refusal_provider_specific_fields( + REFUSAL_STOP_FIELDS.validate_python(completion_response) + ) completion_response["stop_reason"] = "stop" + if refusal_fields: + json_mode_message.provider_specific_fields = dict(refusal_fields) _message = json_mode_message model_response.choices[0].message = _message diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 47a82aa097c..dbcb7ea4962 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -8,11 +8,12 @@ import re from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import Any, ClassVar, Final, Literal, TypeVar +from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar from urllib.parse import quote, urlparse import httpx -from pydantic import ConfigDict, Field, StrictBool, TypeAdapter, ValidationError +from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -2150,3 +2151,35 @@ def create_anthropic_model_list_response( "first_id": data[0]["id"] if data else None, "last_id": data[-1]["id"] if data else None, } + + +class RefusalStopFields(TypedDict, total=False): + stop_reason: ReadOnly[object] + stop_details: ReadOnly[object] + + +def _is_dict(value: object) -> bool: + return isinstance(value, dict) + + +def _no_stop_fields() -> Mapping[str, object]: + return {} + + +def _dict_or_empty(value: object) -> object: + return value if _is_dict(value) else _no_stop_fields() + + +REFUSAL_STOP_FIELDS: Final[TypeAdapter[RefusalStopFields]] = TypeAdapter( + Annotated[RefusalStopFields, BeforeValidator(_dict_or_empty)] +) + + +def refusal_provider_specific_fields(message: Mapping[str, object]) -> Mapping[str, object]: + stop_reason: Final = message.get("stop_reason") + stop_details: Final = message.get("stop_details") + if stop_details is not None: + return {"stop_reason": stop_reason, "stop_details": stop_details} + if stop_reason == "refusal": + return {"stop_reason": stop_reason} + return {} diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index f1eb30b63c1..4c667574d40 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -98,6 +98,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): "output_format", "inference_geo", "speed", + "service_tier", "output_config", "reasoning_effort", "safeguards", diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index 8371615baea..36bd10ab8e6 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -205,6 +205,11 @@ class AnthropicMessagesRequestUtils: drop_params=drop_params, custom_llm_provider=custom_llm_provider, ) + AnthropicConfig.maybe_drop_service_tier_param( + model=model, + optional_params=filtered_params, + custom_llm_provider=custom_llm_provider, + ) for param in ("temperature", "top_p", "top_k"): if param in filtered_params: AnthropicModelInfo._apply_sampling_param( # pyright: ignore[reportPrivateUsage] # same gating the /chat/completions path applies; forking it would drift diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9722f9dccb2..678913ad261 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1191,8 +1191,10 @@ def bedrock_model_accepts_cache_points(model: str | None) -> bool: never silently degrade. ``litellm.utils.supports_prompt_caching`` is not reusable here: it returns False for unmapped models, the opposite polarity. """ - if model is None: - return True + return model is None or bedrock_declared_cache_point_support(model) is not False + + +def bedrock_declared_cache_point_support(model: str) -> bool | None: if _OPENAI_FAMILY_MODEL_RE.search(model): return False map_keys: Final = (model, get_bedrock_base_model(model)) @@ -1209,7 +1211,7 @@ def bedrock_model_accepts_cache_points(model: str | None) -> bool: if explicit_marker_support is not None: return explicit_marker_support if not any(key in litellm.model_cost for key in map_keys): - return True + return None return any( entry.get("supports_prompt_caching") is True for key in map_keys diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index 0ee5dcbe557..ca93cc3e617 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -56,6 +56,10 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): model ID in the request body (unlike Bedrock Invoke which puts it in the URL). """ + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = ( + AmazonAnthropicClaudeMessagesConfig.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS | {"service_tier"} + ) + @property def beta_headers_provider(self) -> str: return "bedrock_mantle" diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 7163cc5d3e0..3e7db2981ef 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -226,7 +226,8 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): def _extract_error_message(self, parsed_chunk: dict[str, Any]) -> str | None: error_obj: Final = parsed_chunk.get("error") or (parsed_chunk.get("response") or {}).get("error") if error_obj is None: - return None + top_level_message: Final = parsed_chunk.get("message") + return top_level_message if isinstance(top_level_message, str) else None if isinstance(error_obj, dict): return error_obj.get("message") or str(error_obj) return str(error_obj) diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py index d8c1b6a368e..22d3f5b6dd1 100644 --- a/litellm/llms/litellm_proxy/skills/skill_search.py +++ b/litellm/llms/litellm_proxy/skills/skill_search.py @@ -107,7 +107,11 @@ class SkillSearchIndex: if isinstance(scores, EmbeddingFailed): return SkillSearchEmbeddingFailed(reason=scores.reason) ranked: Final = sorted( - (SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)), + ( + SkillSearchHit(skill=skill, score=score) + for skill, score in zip(skills, scores, strict=True) + if score is not None + ), key=lambda hit: hit.score, reverse=True, ) diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index 9aef15f15a3..9aaf552e79a 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -40,6 +40,11 @@ class CacheControlSupportedModels(str, Enum): ZAI = "z-ai" +def openrouter_model_supports_cache_control(model: str) -> bool: + model_lower: Final = model.lower() + return any(supported_model.value in model_lower for supported_model in CacheControlSupportedModels) + + class OpenrouterConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list: """ @@ -82,8 +87,7 @@ class OpenrouterConfig(OpenAIGPTConfig): Returns: bool: True if model supports cache_control (Claude or Gemini models) """ - model_lower: Final = model.lower() - return any(supported_model.value in model_lower for supported_model in CacheControlSupportedModels) + return openrouter_model_supports_cache_control(model) def remove_cache_control_flag_from_messages_and_tools( self, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index a6c060732c5..bf2a7d7da15 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -96,7 +96,7 @@ def _decoded_messages(response: _MessagesResponse) -> _AnthropicMessagesResponse return response.json() -def _is_claude_model(model: str) -> bool: +def is_snowflake_claude_model(model: str) -> bool: """Return True if model name (after stripping snowflake/ prefix) is a Claude model.""" name: Final = model.lower().removeprefix("snowflake/") return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES) @@ -222,7 +222,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): "tools", "tool_choice", ] - if _is_claude_model(model): + if is_snowflake_claude_model(model): params.append("thinking") return params @@ -236,7 +236,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): stream: bool | None = None, ) -> str: api_base = self._get_api_base(api_base, optional_params) - if _is_claude_model(model): + if is_snowflake_claude_model(model): return f"{api_base}/cortex/v1/messages" return f"{api_base}/cortex/v1/chat/completions" @@ -259,7 +259,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): api_key=api_key, api_base=api_base, ) - if _is_claude_model(model): + if is_snowflake_claude_model(model): headers["anthropic-version"] = ANTHROPIC_VERSION return headers @@ -406,7 +406,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): stream: Final[bool] = optional_params.pop("stream", False) or False extra_body: Final = optional_params.pop("extra_body", {}) - if _is_claude_model(model): + if is_snowflake_claude_model(model): return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) return self._transform_request_openai(model, messages, optional_params, stream, extra_body) @@ -422,7 +422,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): litellm_params: dict[str, object], # mutable-ok: BaseConfig signature headers: dict[str, object], # mutable-ok: BaseConfig signature ) -> dict[str, object]: # mutable-ok: BaseConfig signature - inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages + inlined_messages: Final = ( + await async_inline_remote_media(messages) if is_snowflake_claude_model(model) else messages + ) return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers) def _transform_request_openai( @@ -529,7 +531,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - if _is_claude_model(model): + if is_snowflake_claude_model(model): return self._transform_response_anthropic( model, raw_response, model_response, logging_obj, request_data, messages ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9011e4511d7..04f391d4d20 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -72571,10 +72571,12 @@ "gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, + "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -72585,6 +72587,7 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -72617,10 +72620,12 @@ "vertex_ai/gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, + "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -72631,6 +72636,7 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", diff --git a/litellm/models/access_group.py b/litellm/models/access_group.py index 513a60aa9ba..fb0ab211bbc 100644 --- a/litellm/models/access_group.py +++ b/litellm/models/access_group.py @@ -5,6 +5,7 @@ Canonical definition for ``litellm_accessgrouptable``. Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ +from collections.abc import Mapping, Sequence from datetime import datetime from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -16,7 +17,9 @@ class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase): description: str | None = None access_model_names: list[str] = [] access_mcp_server_ids: list[str] = [] + mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] = [] + access_skill_names: tuple[str, ...] = () assigned_team_ids: list[str] = [] assigned_key_ids: list[str] = [] created_at: datetime | None = None diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 18b7c279606..43537d48f5e 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,3 +1,4 @@ +import functools import re import secrets from collections.abc import Awaitable, Callable, Mapping, Sequence @@ -42,6 +43,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent ) from litellm.proxy._types import ( UI_TEAM_ID, + LiteLLM_AccessGroupTable, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, ProxyException, @@ -53,8 +55,10 @@ from litellm.proxy._types import ( user_api_key_has_admin_view, ) from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, CeilingResolver, resolve_agent_access_group_ceiling, + resolve_managed_agent_ceilings, ) from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy @@ -126,6 +130,22 @@ class UnloadableEntitlementError(Exception): it places no ceiling; denying there would refuse MCP to every caller during a cold-cache fault.""" +@dataclass(frozen=True, slots=True) +class KeyGroupTools: + """``owned``: tools the key's owned groups grant on a server, ``None`` when they place no bound. + ``owns_server``: whether an owned group grants the server at all, so an unbounded owned grant is told + apart from no owned grant. ``attached``: the bound every attached group, owned or not, places there, + applied only when the key holds no grant of its own on the server.""" + + owned: Sequence[str] | None + owns_server: bool + attached: Sequence[str] | None + + +_NO_KEY_GROUP_TOOLS: Final = KeyGroupTools(owned=None, owns_server=False, attached=None) +_UNREADABLE_KEY_GROUP_TOOLS: Final = KeyGroupTools(owned=(), owns_server=False, attached=()) + + def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: list[str] | None = None) -> list[str] | None: """Resolve the single MCP server name a cold-start passthrough bypass may target. Delegates parsing to @@ -251,6 +271,15 @@ def _agent_capped_servers( ) +def _intersect_tool_ceiling(own: Sequence[str] | None, ceiling: Sequence[str] | None) -> Sequence[str] | None: + """The agent's own tool grants capped by its access-group ceiling, ``None`` only when neither bounds.""" + if ceiling is None: + return own + if own is None: + return ceiling + return tuple(set(own) & set(ceiling)) + + def is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool: """True when this auth is a keyless subject admitted by the gateway session / bridge user path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``. @@ -2363,16 +2392,177 @@ class MCPRequestHandler: def _union_tool_grants( direct: Sequence[str] | None, via_toolsets: Sequence[str] | None, + via_access_groups: Sequence[str] | None = None, ) -> Sequence[str] | None: - """Union of one level's direct tool grants and its toolset-granted tools on one server, - ``None`` when neither source restricts (allow-all from this level). A direct grant - containing ``MCP_ALL_TOOLS_WILDCARD`` makes the level unrestricted, so it returns - ``None`` whatever the toolsets name.""" + """Union of one level's direct tool grants, its toolset-granted tools and its unified + access-group tools on one server, ``None`` when no source restricts (allow-all from this + level). A direct grant containing ``MCP_ALL_TOOLS_WILDCARD`` makes the level unrestricted, + so it returns ``None`` whatever the toolsets or groups name.""" if direct is not None and MCP_ALL_TOOLS_WILDCARD in direct: return None - if direct is None and via_toolsets is None: + if direct is None and via_toolsets is None and via_access_groups is None: return None - return tuple({*(direct or ()), *(via_toolsets or ())}) + return tuple({*(direct or ()), *(via_toolsets or ()), *(via_access_groups or ())}) + + @staticmethod + def access_group_tools_for_server( + groups: Sequence[LiteLLM_AccessGroupTable], server_id: str + ) -> Sequence[str] | None: + """Tools the unified access groups grant on ``server_id``: the union of each granting group's + ``mcp_tool_permissions`` entry for it. A granting group with no entry grants every tool, so + ``None`` wins. A group that does not grant the server contributes nothing, and an explicit + ``[]`` entry on a granting group grants no tools. ``None`` when no group grants the server.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + def granted_tools(group: LiteLLM_AccessGroupTable) -> Sequence[str] | None: + if not MCPRequestHandler._group_grants_server(group, server_id): + return None + permissions: Final = { + server: list(allowed) for server, allowed in (group.mcp_tool_permissions or {}).items() + } + tools: Final = global_mcp_server_manager.expand_tool_permissions(permissions) + return tools.get(server_id, (MCP_ALL_TOOLS_WILDCARD,)) + + grants: Final = tuple(granted_tools(group) for group in groups) + named: Final = tuple(grant for grant in grants if grant is not None) + if not named or any(MCP_ALL_TOOLS_WILDCARD in grant for grant in named): + return None + return tuple({tool for grant in named for tool in grant}) # comprehension-ok: flatten + + @staticmethod + def _group_grants_server(group: LiteLLM_AccessGroupTable, server_id: str) -> bool: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return server_id in global_mcp_server_manager.expand_permission_list(list(group.access_mcp_server_ids)) + + @staticmethod + def _key_tool_scope( + direct: Sequence[str] | None, + via_toolsets: Sequence[str] | None, + groups: KeyGroupTools, + grants_server: bool, + ) -> Sequence[str] | None: + """The key's own tool scope on a server when it holds any grant there (a direct server grant with + or without a tool entry, a toolset or an owned group), so an unrestricted grant stays unrestricted. + Only a key with no grant of its own takes the bound its attached groups place.""" + if grants_server or direct is not None or via_toolsets is not None or groups.owns_server: + return MCPRequestHandler._union_tool_grants(direct, via_toolsets, groups.owned) + return groups.attached + + @staticmethod + def _group_authorizes_key(group: LiteLLM_AccessGroupTable, team_id: str | None, token: str | None) -> bool: + """Mirror of ``get_authorized_resources_from_key_access_groups``: a group a key merely names + binds only when the group lists the key's team or the key itself as an owner.""" + return bool(team_id and team_id in group.assigned_team_ids) or bool(token and token in group.assigned_key_ids) + + @staticmethod + def _loaded_group_tools_for_server( + groups: Sequence[LiteLLM_AccessGroupTable], requested: Sequence[str], server_id: str + ) -> Sequence[str] | None: + """Tool grants of the agent groups that loaded. A named agent group that did not load exists but + could not be read (a group that does not exist never reaches the ceiling), so the level grants no + tools on ``server_id`` rather than shedding a restriction while the fault lasts.""" + if {group.access_group_id for group in groups} != set(requested): + return () + return MCPRequestHandler.access_group_tools_for_server(groups, server_id) + + @staticmethod + async def _key_direct_servers( + key_obj_perm: LiteLLM_ObjectPermissionTable | None, *, requires_fresh_policy: bool + ) -> AbstractSet[str]: + """Servers the key's object permission grants on its own: ``mcp_servers`` by id, name or alias and + legacy ``mcp_access_groups`` by the tag the server carries, exactly as server admission resolves + them. A tag never doubles as a server name, so a tag that collides with one grants nothing by itself.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + if key_obj_perm is None: + return frozenset() + try: + tagged: Final = ( + await MCPRequestHandler.get_mcp_servers_from_access_groups( + key_obj_perm.mcp_access_groups, requires_fresh_policy=requires_fresh_policy + ) + if key_obj_perm.mcp_access_groups + else () + ) + except Exception as e: # noqa: BLE001 # a key naming tags we cannot resolve must deny, not widen + raise UnloadableEntitlementError(f"key access group tags unreadable: {e}") from e + return frozenset(global_mcp_server_manager.expand_permission_list(key_obj_perm.mcp_servers or [])) | frozenset( + tagged + ) + + @staticmethod + async def _key_access_group_tools_for_server(user_api_key_auth: UserAPIKeyAuth, server_id: str) -> KeyGroupTools: + """Tool bounds the key's attached groups place on ``server_id``. Attaching a group to a key admits + the group's servers (``_get_key_access_group_mcp_server_extras``), so an attached group the key + does not own still carries its own tool permissions onto the servers it admits, it just cannot + add tools to a grant the key already holds there. Only owned groups are grant sources. A named + group that does not exist grants and bounds nothing, one that exists but cannot be read denies.""" + from litellm.proxy.auth.auth_checks import get_access_group_tables + + if not user_api_key_auth.access_group_ids: + return _NO_KEY_GROUP_TOOLS + try: + groups: Final = await get_access_group_tables( + access_group_ids=user_api_key_auth.access_group_ids, + check_db_only=bool(user_api_key_auth.requires_fresh_policy), + unreadable="raise", + ) + except Exception as e: # noqa: BLE001 # a key naming groups we cannot read must deny, not widen + if user_api_key_auth.requires_fresh_policy: + raise UnloadableEntitlementError(f"key access groups unreadable: {e}") from e + return _UNREADABLE_KEY_GROUP_TOOLS + owned: Final = tuple( + group + for group in groups + if MCPRequestHandler._group_authorizes_key(group, user_api_key_auth.team_id, user_api_key_auth.token) + ) + return KeyGroupTools( + owned=MCPRequestHandler.access_group_tools_for_server(owned, server_id), + owns_server=any(MCPRequestHandler._group_grants_server(group, server_id) for group in owned), + attached=MCPRequestHandler.access_group_tools_for_server(groups, server_id), + ) + + @staticmethod + async def _team_access_group_tools_for_server( + user_api_key_auth: UserAPIKeyAuth, server_id: str + ) -> Sequence[str] | None: + from litellm.proxy.auth.auth_checks import get_access_group_tables, get_team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if not user_api_key_auth.team_id or user_api_key_auth.team_id == UI_TEAM_ID or prisma_client is None: + return None + try: + team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object( + team_id=user_api_key_auth.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth.requires_fresh_policy), + ) + except Exception as e: # noqa: BLE001 # a team we cannot read must deny, not widen + raise UnloadableEntitlementError(f"team unreadable: {e}") from e + if not team_obj.access_group_ids: + return None + try: + groups: Final = await get_access_group_tables( + access_group_ids=team_obj.access_group_ids, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth.requires_fresh_policy), + unreadable="raise", + ) + except Exception as e: # noqa: BLE001 # team groups we cannot read grant no tools, they never widen + if user_api_key_auth.requires_fresh_policy: + raise UnloadableEntitlementError(f"team access groups unreadable: {e}") from e + return () + return MCPRequestHandler.access_group_tools_for_server(groups, server_id) @staticmethod async def _key_object_permission_hydrated( @@ -2493,7 +2683,17 @@ class MCPRequestHandler: else None ) - key_tools: Final = _as_list(MCPRequestHandler._union_tool_grants(key_direct_tools, key_toolset_tools)) + key_group_tools: Final = await MCPRequestHandler._key_access_group_tools_for_server( + user_api_key_auth, server_id + ) + key_direct_servers: Final = await MCPRequestHandler._key_direct_servers( + key_obj_perm, requires_fresh_policy=bool(user_api_key_auth.requires_fresh_policy) + ) + key_tools: Final = _as_list( + MCPRequestHandler._key_tool_scope( + key_direct_tools, key_toolset_tools, key_group_tools, server_id in key_direct_servers + ) + ) team_direct_tools: Final = ( global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id) if team_obj_perm @@ -2505,11 +2705,17 @@ class MCPRequestHandler: team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( team_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy ) - team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools) + team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server( + user_api_key_auth, server_id + ) + team_tools: Final = MCPRequestHandler._union_tool_grants( + team_direct_tools, team_toolset_tools, team_group_tools + ) - # Apply same inheritance logic as get_allowed_mcp_servers - if team_tools: - if key_tools: + # Apply same inheritance logic as get_allowed_mcp_servers. ``is not None`` so an explicit + # empty allowlist stays a deny instead of reading as "no restriction". + if team_tools is not None: + if key_tools is not None: # Both have restrictions → intersection allowed_tools = list(set(team_tools) & set(key_tools)) else: @@ -3525,6 +3731,9 @@ class MCPRequestHandler: team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( team_obj_perm, server_id, requires_fresh_policy=caller_auth.requires_fresh_policy ) + team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server( + caller_auth, server_id + ) except Exception as e: # noqa: BLE001 # an unresolved caller team must deny, not widen verbose_logger.warning( "MCP agent caller team tool ceiling unresolvable, denying tools on %r: %s", server_id, e @@ -3535,7 +3744,9 @@ class MCPRequestHandler: if team_obj_perm else None ) - team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools) + team_tools: Final = MCPRequestHandler._union_tool_grants( + team_direct_tools, team_toolset_tools, team_group_tools + ) team_capped: Final = ( allowed_tools if team_tools is None @@ -3729,12 +3940,50 @@ class MCPRequestHandler: global_mcp_server_manager, ) + ceilings: Final = await MCPRequestHandler._agent_access_group_ceilings(user_api_key_auth, resolve_ceiling) + if not ceilings: + return None + expanded: Final[tuple[frozenset[str], ...]] = tuple( + frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) + for ceiling in ceilings + ) + return functools.reduce(lambda acc, servers: acc & servers, expanded[1:], expanded[0]) + + @staticmethod + async def _get_agent_access_group_tool_ceiling( + user_api_key_auth: UserAPIKeyAuth, + server_id: str, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + ) -> Sequence[str] | None: + """Ceiling the agent's attached unified access groups place on ``server_id`` tools, the tools axis + twin of ``_get_agent_access_group_server_ceiling``. None when the agent has no groups attached or + no attached group bounds tools on the server. The agent's own grants are intersected with it.""" + ceilings: Final = await MCPRequestHandler._agent_access_group_ceilings(user_api_key_auth, resolve_ceiling) + if not ceilings: + return None + bounds: Final = tuple( + MCPRequestHandler._loaded_group_tools_for_server(ceiling.access_groups, ceiling.access_group_ids, server_id) + for ceiling in ceilings + ) + return functools.reduce(_intersect_tool_ceiling, bounds[1:], bounds[0]) + + @staticmethod + async def _agent_access_group_ceilings( + user_api_key_auth: UserAPIKeyAuth, + resolve_ceiling: CeilingResolver, + ) -> tuple[AgentAccessGroupCeiling, ...]: + """A managed agent is bounded by the groups on its admitted policy, the same source server admission + reads, never by the registry row that policy may have superseded. Empty means no groups attached.""" if not user_api_key_auth.agent_id: - return None - ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) - if ceiling is None: - return None - return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) + return () + managed: Final = managed_agent_policy(user_api_key_auth) + try: + if managed is not None: + return await resolve_managed_agent_ceilings(managed) + ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) + except Exception as e: # noqa: BLE001 # agent groups we cannot read must deny, not widen + raise UnloadableEntitlementError(f"agent access groups unreadable: {e}") from e + return (ceiling,) if ceiling is not None else () @staticmethod async def get_agent_tool_permissions_for_server( @@ -3764,14 +4013,17 @@ class MCPRequestHandler: if agent_object_permission is not None else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) ) - if obj_perm is None: - return None try: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + group_tools: Final = await MCPRequestHandler._get_agent_access_group_tool_ceiling( + user_api_key_auth, server_id + ) + if obj_perm is None: + return _as_list(group_tools) direct_tools: Final = ( global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id) if obj_perm.mcp_tool_permissions @@ -3780,7 +4032,8 @@ class MCPRequestHandler: toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy ) - agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) + own_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) + agent_tools: Final = _intersect_tool_ceiling(own_tools, group_tools) return list(agent_tools) if agent_tools is not None else None except Exception as e: if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError): diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 23a63c00bd0..16ab77c2c2d 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -258,9 +258,11 @@ def _split_core_tools(tools: Sequence[Tool], core_tools: Sequence[str]) -> tuple def _top_hits( - tools: Sequence[Tool], scores: Sequence[float], minimum: float, limit: int + tools: Sequence[Tool], scores: Sequence[float | None], minimum: float, limit: int ) -> tuple[tuple[float, Tool], ...]: - hits: Final = ((score, tool) for score, tool in zip(scores, tools, strict=True) if score >= minimum) + hits: Final = ( + (score, tool) for score, tool in zip(scores, tools, strict=True) if score is not None and score >= minimum + ) return tuple(sorted(hits, key=lambda hit: hit[0], reverse=True)[:limit]) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index effb30fe8cb..9086c8417ae 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -648,6 +648,20 @@ ], "title": "Access Model Names" }, + "access_skill_names": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Skill Names" + }, "assigned_key_ids": { "anyOf": [ { @@ -686,6 +700,23 @@ } ], "title": "Description" + }, + "mcp_tool_permissions": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "string" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Mcp Tool Permissions" } }, "required": [ @@ -812,6 +843,14 @@ "title": "Access Model Names", "type": "array" }, + "access_skill_names": { + "default": [], + "items": { + "type": "string" + }, + "title": "Access Skill Names", + "type": "array" + }, "assigned_key_ids": { "items": { "type": "string" @@ -867,6 +906,23 @@ ], "title": "Description" }, + "mcp_tool_permissions": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "string" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Mcp Tool Permissions" + }, "updated_at": { "format": "date-time", "title": "Updated At", @@ -957,6 +1013,20 @@ ], "title": "Access Model Names" }, + "access_skill_names": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Skill Names" + }, "assigned_key_ids": { "anyOf": [ { @@ -995,6 +1065,23 @@ } ], "title": "Description" + }, + "mcp_tool_permissions": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "string" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Mcp Tool Permissions" } }, "title": "AccessGroupUpdateRequest", @@ -7644,7 +7731,7 @@ "paths": { "/claude-code/marketplace.json": { "get": { - "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```", + "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add \n- claude plugin install @\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` or an access group on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```", "operationId": "get_marketplace_claude_code_marketplace_json_get", "parameters": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 252d03b3983..5d089d5be36 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4260,6 +4260,7 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): models: tuple[str, ...] mcp_server_ids: tuple[str, ...] = () agent_ids: tuple[str, ...] = () + skill_names: tuple[str, ...] = () class TeamInfoMember(Member): @@ -4297,6 +4298,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None + access_group_skill_names: tuple[str, ...] | None = None + # Listed group ids with no row; skill grants are suspended while this is non-empty. + unresolved_access_group_ids: tuple[str, ...] | None = None access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None # Parent org's model ceiling, reported only to callers who can manage the team. # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling. diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 1d1778a960b..af2baa2bd89 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -116,7 +116,11 @@ class AgentSearchIndex: if isinstance(scores, EmbeddingFailed): return AgentSearchEmbeddingFailed(reason=scores.reason) ranked: Final = sorted( - (AgentSearchHit(agent=agent, score=score) for agent, score in zip(agents, scores, strict=True)), + ( + AgentSearchHit(agent=agent, score=score) + for agent, score in zip(agents, scores, strict=True) + if score is not None + ), key=lambda hit: hit.score, reverse=True, ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index db661fbea30..557d1c7bc0a 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -25,11 +25,19 @@ class AgentAccessGroupCeiling: models: frozenset[str] mcp_server_ids: frozenset[str] agent_ids: frozenset[str] + access_groups: tuple[LiteLLM_AccessGroupTable, ...] = () CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] +class _MissingGroup: + __slots__ = () + + +_MISSING: Final = _MissingGroup() + + async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through @@ -38,7 +46,7 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup: - from litellm.proxy.auth.auth_checks import get_access_object + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache if prisma_client is None: @@ -52,6 +60,8 @@ async def _load_access_group(access_group_id: str, *, check_db_only: bool = Fals proxy_logging_obj=proxy_logging_obj, check_db_only=check_db_only, ) + except AccessGroupNotFoundError: + raise except HTTPException as e: if check_db_only: raise @@ -68,23 +78,35 @@ async def resolve_agent_access_group_ceiling( *, check_db_only: bool = False, ) -> AgentAccessGroupCeiling | None: - """``None`` when the agent has no access groups attached, so nothing is capped.""" - access_group_ids: Final = await load_access_group_ids(agent_id) - if not access_group_ids: + """``None`` when the agent has no access groups attached, so nothing is capped. A named group that + provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group + that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once + any group is named, so an agent whose every group is missing is still capped to nothing.""" + named_ids: Final = await load_access_group_ids(agent_id) + if not named_ids: return None - loaded: Final = await asyncio.gather( - *( - _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id) - for group_id in access_group_ids - ) + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup: + try: + return await ( + _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id) + ) + except AccessGroupNotFoundError: + return _MISSING + + results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids)) + access_group_ids: Final = tuple( + group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup) ) - groups: Final = tuple(group for group in loaded if group is not None) + groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable)) return AgentAccessGroupCeiling( access_group_ids=access_group_ids, models=frozenset(model for group in groups for model in group.access_model_names), mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), + access_groups=groups, ) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 806f3ff7535..8e348423837 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -33,7 +33,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( SkillVisibility, - skill_visibility, + resolve_skill_visibility, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.resource_ownership import is_proxy_admin @@ -122,7 +122,7 @@ async def get_marketplace(request: Request, key: str | None = None): Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...` the key is authenticated and the catalog also holds the disabled plugins granted - to it through `object_permission.skills` on the key or its team. + to it through `object_permission.skills` or an access group on the key or its team. Returns: Marketplace catalog with list of available plugins and their git sources. @@ -140,7 +140,7 @@ async def get_marketplace(request: Request, key: str | None = None): caller: Final[UserAPIKeyAuth | None] = ( await user_api_key_auth(request=request, api_key=f"Bearer {key}") if key else None ) - visibility: Final[SkillVisibility] = skill_visibility(caller) + visibility: Final[SkillVisibility] = await resolve_skill_visibility(caller) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( where=visibility.where() ) @@ -464,7 +464,7 @@ async def list_plugins( try: prisma_client: Final = await get_prisma_client() - visibility: Final[SkillVisibility] = skill_visibility(user_api_key_dict) + visibility: Final[SkillVisibility] = await resolve_skill_visibility(user_api_key_dict) plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( where={"enabled": True} if enabled_only else visibility.where() ) @@ -545,7 +545,7 @@ async def get_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - if not skill_visibility(user_api_key_dict).allows(plugin): + if not (await resolve_skill_visibility(user_api_key_dict)).allows(plugin): raise HTTPException( status_code=403, detail={"error": f"Plugin '{plugin_name}' is not granted to this key"}, diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py index e1e1ee6f160..9e5b55e7e36 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py @@ -1,18 +1,22 @@ """ Claude Code marketplace visibility: enabled plugins are public, disabled plugins are private and resolve only for proxy admins or keys granted them via -``object_permission.skills``. +``object_permission.skills`` or via ``access_skill_names`` on an access group +attached to the key or its team. """ +from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias -from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth +from litellm.proxy._types import UI_TEAM_ID, LiteLLM_ObjectPermissionTable, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import is_proxy_admin if TYPE_CHECKING: from prisma.types import LiteLLM_ClaudeCodePluginTableWhereInput + from litellm.proxy._types import LiteLLM_TeamTable + class _SkillRecord(Protocol): name: str @@ -23,14 +27,28 @@ def _skills_of(permission: LiteLLM_ObjectPermissionTable | None) -> frozenset[st return frozenset(permission.skills or ()) if permission is not None else frozenset() -def granted_skills(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: +@dataclass(frozen=True, slots=True) +class AccessGroupSkills: + """Plugin names inherited from access groups, split by the scope the group is attached to.""" + + key: frozenset[str] = frozenset() + team: frozenset[str] = frozenset() + unresolved: bool = False + """A listed group returned no row, so the policy is incomplete and only public plugins show.""" + + +NO_ACCESS_GROUP_SKILLS: Final = AccessGroupSkills() + + +def granted_skills( + user_api_key_dict: UserAPIKeyAuth, inherited: AccessGroupSkills = NO_ACCESS_GROUP_SKILLS +) -> frozenset[str]: """Key grant intersected with the team grant when both are non-empty; either alone applies as is. - An empty list is the Prisma column default for every object-permission row, so it means - "no private grants configured here" and defers to the other scope, same as the agents check. + An empty list is the Prisma column default, so it means "no private grants configured here". """ - key_skills: Final = _skills_of(user_api_key_dict.object_permission) - team_skills: Final = _skills_of(user_api_key_dict.team_object_permission) + key_skills: Final = _skills_of(user_api_key_dict.object_permission) | inherited.key + team_skills: Final = _skills_of(user_api_key_dict.team_object_permission) | inherited.team match (bool(key_skills), bool(team_skills)): case (True, True): return key_skills & team_skills @@ -59,9 +77,87 @@ class SkillVisibility: PUBLIC_ONLY: Final = SkillVisibility(granted=frozenset(), sees_private=False) -def skill_visibility(user_api_key_dict: UserAPIKeyAuth | None) -> SkillVisibility: +def skill_visibility( + user_api_key_dict: UserAPIKeyAuth | None, inherited: AccessGroupSkills = NO_ACCESS_GROUP_SKILLS +) -> SkillVisibility: if user_api_key_dict is None: return PUBLIC_ONLY if is_proxy_admin(user_api_key_dict): return SkillVisibility(granted=frozenset(), sees_private=True) - return SkillVisibility(granted=granted_skills(user_api_key_dict), sees_private=False) + if inherited.unresolved: + return PUBLIC_ONLY + return SkillVisibility(granted=granted_skills(user_api_key_dict, inherited), sees_private=False) + + +async def _team_object(user_api_key_dict: UserAPIKeyAuth, check_db_only: bool) -> "LiteLLM_TeamTable | None": + from litellm.proxy.auth.auth_checks import TeamNotFoundError, get_team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + team_id: Final = user_api_key_dict.team_id + if not team_id or team_id == UI_TEAM_ID or prisma_client is None: + return None + try: + team_obj: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + except TeamNotFoundError: + return None + return None if team_obj.blocked else team_obj + + +async def access_group_skills(user_api_key_dict: UserAPIKeyAuth) -> AccessGroupSkills: + """Resolve ``access_skill_names`` of the groups attached to the key and to its team. + + Key-side groups count only when the group lists the key or its team as an owner, same as + the model and MCP checks, so a key cannot claim a group it was never assigned to. + """ + from litellm.proxy.auth.auth_checks import ( + AccessGroupNotFoundError, + get_authorized_resources_from_key_access_groups, + get_skill_names_from_access_groups, + ) + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + return NO_ACCESS_GROUP_SKILLS + check_db_only: Final = bool(user_api_key_dict.requires_fresh_policy) + team_obj: Final = await _team_object(user_api_key_dict, check_db_only) + try: + key_skills: Final = frozenset( + await get_authorized_resources_from_key_access_groups( + valid_token=user_api_key_dict, + team_object=team_obj, + resource_field="access_skill_names", + check_db_only=check_db_only, + strict_errors=True, + ) + ) + team_skills: Final = await get_skill_names_from_access_groups( + access_group_ids=tuple(team_obj.access_group_ids or ()) if team_obj is not None else (), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + strict_errors=True, + ) + return AccessGroupSkills(key=key_skills, team=team_skills) + except AccessGroupNotFoundError: + return AccessGroupSkills(unresolved=True) + + +AccessGroupSkillsLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[AccessGroupSkills]] + + +async def resolve_skill_visibility( + user_api_key_dict: UserAPIKeyAuth | None, + load_access_group_skills: AccessGroupSkillsLoader = access_group_skills, +) -> SkillVisibility: + """``skill_visibility`` with the access-group grants of the key and its team folded in.""" + if user_api_key_dict is None or is_proxy_admin(user_api_key_dict): + return skill_visibility(user_api_key_dict) + return skill_visibility(user_api_key_dict, await load_access_group_skills(user_api_key_dict)) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5568ae2995c..93dbfa9a9d2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3411,6 +3411,19 @@ class TeamNotFoundError(HTTPException): super().__init__(status_code=404, detail=detail) +class AccessGroupNotFoundError(HTTPException): + """The access group row is provably absent, as opposed to merely unreadable. ``get_access_object`` + reports every failure as a 404, so a group that was never created and a database that would not + answer look the same to its callers. Policy code that must not treat a degraded read as a definitive + answer keys on this subclass. It stays a 404 with the same detail, so every other caller is unaffected.""" + + def __init__(self, access_group_id: str) -> None: + super().__init__( + status_code=404, + detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."}, + ) + + @log_db_metrics async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False @@ -3667,10 +3680,7 @@ async def get_access_object( ).find_unique(where={"access_group_id": access_group_id}) if response is None: - raise HTTPException( - status_code=404, - detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."}, - ) + raise AccessGroupNotFoundError(access_group_id=access_group_id) _response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict()) @@ -4532,67 +4542,99 @@ async def get_org_object_for_request( return org -async def _get_resources_from_access_groups( +async def get_access_group_tables( access_group_ids: Sequence[str], - resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], prisma_client: DatabaseClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, check_db_only: bool = False, -) -> list[str]: - """ - Fetch access groups by their IDs (from cache or DB) and collect - the specified resource field across all of them. - - Args: - access_group_ids: List of access group IDs to fetch - resource_field: Which resource list to extract from each access group - - "access_model_names": model names (for model access checks) - - "access_mcp_server_ids": MCP server IDs (for MCP access checks) - - "access_agent_ids": agent IDs (for agent access checks) - prisma_client: Optional PrismaClient (lazy-imported from proxy_server if None) - user_api_key_cache: Optional DualCache (lazy-imported from proxy_server if None) - proxy_logging_obj: Optional ProxyLogging (lazy-imported from proxy_server if None) - - Returns: - Deduplicated list of resource identifiers from all resolved access groups. - """ + *, + unreadable: Literal["skip", "raise"] = "skip", + missing: Literal["skip", "raise"] = "skip", +) -> tuple[LiteLLM_AccessGroupTable, ...]: + """Load unified access groups by id (cache first, then DB). A group that provably does not exist is + skipped, it grants nothing, unless ``missing`` is ``"raise"`` and the caller wants to see + ``AccessGroupNotFoundError`` itself. A group that exists but cannot be read is skipped too unless + ``check_db_only`` is set or ``unreadable`` is ``"raise"``, in which case the failure propagates so a + caller enforcing policy can deny instead of shedding a restriction it could not read.""" if not access_group_ids: - return [] + return () - # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache - prisma_client = prisma_client or _prisma_client - user_api_key_cache = user_api_key_cache or _user_api_key_cache - proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj + prisma_client = prisma_client or _prisma_client # rebind-ok: lazy proxy_server fallback + user_api_key_cache = user_api_key_cache or _user_api_key_cache # rebind-ok: lazy proxy_server fallback + proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj # rebind-ok: lazy proxy_server fallback if user_api_key_cache is None: - return [] + return () - resources: Final[list[str]] = [] - for ag_id in access_group_ids: + async def load(ag_id: str) -> LiteLLM_AccessGroupTable | None: try: - ag = await get_access_object( + return await get_access_object( access_group_id=ag_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, check_db_only=check_db_only, ) - resources.extend(getattr(ag, resource_field, [])) - except Exception: - if check_db_only: + except AccessGroupNotFoundError: + if missing == "raise": raise - verbose_proxy_logger.debug( - "Could not fetch access group %s for resource field %s", - ag_id, - resource_field, - ) - return list(set(resources)) + verbose_proxy_logger.debug("One of the requested access groups does not exist, it grants nothing") + return None + except Exception: + if check_db_only or unreadable == "raise": + raise + verbose_proxy_logger.debug("Could not fetch one of the requested access groups") + return None + + loaded: Final = [await load(ag_id) for ag_id in access_group_ids] + return tuple(ag for ag in loaded if ag is not None) + + +async def _get_resources_from_access_groups( + access_group_ids: Sequence[str], + resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names"], + prisma_client: DatabaseClient | None = None, + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, + check_db_only: bool = False, + strict_errors: bool = False, +) -> list[str]: + """Deduplicated ``resource_field`` values across the resolved access groups. With ``strict_errors`` a + missing or unreadable group raises instead of being skipped, so policy code never loses a ceiling.""" + strictness: Final[Literal["skip", "raise"]] = "raise" if strict_errors else "skip" + groups: Final = await get_access_group_tables( + access_group_ids=access_group_ids, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + unreadable=strictness, + missing=strictness, + ) + return list( + {resource for ag in groups for resource in _access_group_field(ag, resource_field)} # comprehension-ok: flatten + ) + + +def _access_group_field( + group: LiteLLM_AccessGroupTable, + resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names"], +) -> list[str]: + match resource_field: + case "access_model_names": + return group.access_model_names + case "access_mcp_server_ids": + return group.access_mcp_server_ids + case "access_agent_ids": + return group.access_agent_ids + case "access_skill_names": + return list(group.access_skill_names or ()) async def get_models_from_access_groups( @@ -4665,6 +4707,28 @@ async def get_agent_ids_from_access_groups( _get_agent_ids_from_access_groups: Final = get_agent_ids_from_access_groups +async def get_skill_names_from_access_groups( + access_group_ids: Sequence[str], + prisma_client: PrismaClient | None = None, + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, + check_db_only: bool = False, + strict_errors: bool = False, +) -> frozenset[str]: + """Claude Code plugin names granted by unified access groups, matched by plugin name.""" + return frozenset( + await _get_resources_from_access_groups( + access_group_ids=access_group_ids, + resource_field="access_skill_names", + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + strict_errors=strict_errors, + ) + ) + + def _resolve_all_team_model_sentinel_for_auth_check( models: list[str], llm_router: Router | None, @@ -5584,7 +5648,9 @@ async def can_team_access_model( async def get_authorized_resources_from_key_access_groups( valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, - resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], + resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names"], + check_db_only: bool = False, + strict_errors: bool = False, ) -> list[str]: """ For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row @@ -5618,8 +5684,15 @@ async def get_authorized_resources_from_key_access_groups( prisma_client=_prisma_client, user_api_key_cache=_user_api_key_cache, proxy_logging_obj=_proxy_logging_obj, + check_db_only=check_db_only, ) + except AccessGroupNotFoundError: + if check_db_only or strict_errors: + raise + continue except Exception: + if check_db_only or strict_errors: + raise continue team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 7450ef2b509..c2829f28e61 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -3,7 +3,8 @@ from __future__ import annotations import math -from collections.abc import Awaitable, Mapping, Sequence +import time +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from itertools import chain, islice from types import MappingProxyType @@ -13,7 +14,7 @@ from fastapi import HTTPException from openai import OpenAIError from pydantic import ConfigDict -from litellm.exceptions import BudgetExceededError +from litellm.exceptions import BudgetExceededError, ContextWindowExceededError from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: @@ -25,6 +26,7 @@ Vector: TypeAlias = tuple[float, ...] DEFAULT_MAX_CACHED_VECTORS: Final = 5000 """Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first.""" +DEFAULT_TOO_LONG_TTL_SECONDS: Final = 3600.0 class Embedder(Protocol): @@ -36,6 +38,18 @@ class EmbeddingFailed: reason: str +@dataclass(frozen=True, slots=True) +class _InputTooLong: + reason: str + + +@dataclass(frozen=True, slots=True) +class _TextVectors: + vectors: Mapping[str, Vector] + too_long: frozenset[str] + failure: EmbeddingFailed | None = None + + class _EmbeddingItem(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -102,11 +116,13 @@ def router_embedder( _CacheKey: TypeAlias = tuple[str, str] -async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed: +async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed | _InputTooLong: try: vectors: Final = tuple(await embed(texts)) except HTTPException: raise + except ContextWindowExceededError as exc: + return _InputTooLong(reason=f"embedding the search query failed: {exc}") except (OpenAIError, ValueError, BudgetExceededError) as exc: return EmbeddingFailed(reason=f"embedding the search query failed: {exc}") if len(vectors) != len(texts): @@ -114,32 +130,97 @@ async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ... return vectors +async def _embed_texts(embed: Embedder, texts: Sequence[str]) -> _TextVectors: + if not texts: + return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset()) + embedded: Final = await _embed_all(embed, texts) + if isinstance(embedded, EmbeddingFailed): + return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset(), failure=embedded) + if not isinstance(embedded, _InputTooLong): + return _TextVectors(vectors=MappingProxyType(dict(zip(texts, embedded, strict=True))), too_long=frozenset()) + if len(texts) == 1: + return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset({texts[0]})) + middle: Final = len(texts) // 2 + left: Final = await _embed_texts(embed, texts[:middle]) + if left.failure is not None: + return left + right: Final = await _embed_texts(embed, texts[middle:]) + return _TextVectors( + vectors=MappingProxyType({**left.vectors, **right.vectors}), + too_long=left.too_long | right.too_long, + failure=right.failure, + ) + + +async def _embed_query_with( + embed: Embedder, query: str, texts: Sequence[str] +) -> tuple[Vector, _TextVectors] | EmbeddingFailed: + if not texts: + embedded_query: Final = await _embed_all(embed, (query,)) + if isinstance(embedded_query, EmbeddingFailed): + return embedded_query + if isinstance(embedded_query, _InputTooLong): + return EmbeddingFailed(reason=embedded_query.reason) + return embedded_query[0], _TextVectors(vectors=MappingProxyType({}), too_long=frozenset()) + + embedded: Final = await _embed_all(embed, (query, *texts)) + if isinstance(embedded, EmbeddingFailed): + return embedded + if not isinstance(embedded, _InputTooLong): + return embedded[0], _TextVectors( + vectors=MappingProxyType(dict(zip(texts, embedded[1:], strict=True))), too_long=frozenset() + ) + query_embedding: Final = await _embed_all(embed, (query,)) + if isinstance(query_embedding, EmbeddingFailed): + return query_embedding + if isinstance(query_embedding, _InputTooLong): + return EmbeddingFailed(reason=query_embedding.reason) + vectors: Final = await _embed_texts(embed, texts) + return query_embedding[0], vectors + + @dataclass(frozen=True, slots=True) class _Embedded: query_vector: Vector vectors: Mapping[str, Vector] + too_long: frozenset[str] + failure: EmbeddingFailed | None def _same_dimension(query_vector: Vector, vectors: Mapping[str, Vector], texts: Sequence[str]) -> bool: - return all(len(vectors[text]) == len(query_vector) for text in texts) + return all(len(vectors[text]) == len(query_vector) for text in texts if text in vectors) async def _embed_query_and_texts( - embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector] + embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector], too_long: frozenset[str] ) -> _Embedded | EmbeddingFailed: - missing: Final = tuple(dict.fromkeys(text for text in texts if text not in cached)) - embedded: Final = await _embed_all(embed, (query, *missing)) - if isinstance(embedded, EmbeddingFailed): + pending: Final = (text for text in texts if text not in cached and text not in too_long) + missing: Final = tuple(dict.fromkeys(pending)) + query_embedding: Final = await _embed_query_with(embed, query, missing) + if isinstance(query_embedding, EmbeddingFailed): + return query_embedding + query_vector, new_text_vectors = query_embedding + vectors: Final = MappingProxyType(dict(chain(cached.items(), new_text_vectors.vectors.items()))) + embedded: Final = _Embedded( + query_vector=query_vector, + vectors=vectors, + too_long=new_text_vectors.too_long, + failure=new_text_vectors.failure, + ) + if embedded.failure is not None: return embedded - vectors: Final = MappingProxyType(dict(chain(cached.items(), zip(missing, embedded[1:], strict=True)))) - if _same_dimension(embedded[0], vectors, texts): - return _Embedded(query_vector=embedded[0], vectors=vectors) - unique: Final = tuple(dict.fromkeys(texts)) - reembedded: Final = await _embed_all(embed, (query, *unique)) + if _same_dimension(query_vector, vectors, texts): + return embedded + unique: Final = tuple(dict.fromkeys(text for text in texts if text in vectors)) + reembedded: Final = await _embed_query_with(embed, query, unique) if isinstance(reembedded, EmbeddingFailed): return reembedded + reembedded_query, reembedded_text_vectors = reembedded return _Embedded( - query_vector=reembedded[0], vectors=MappingProxyType(dict(zip(unique, reembedded[1:], strict=True))) + query_vector=reembedded_query, + vectors=reembedded_text_vectors.vectors, + too_long=new_text_vectors.too_long | reembedded_text_vectors.too_long, + failure=reembedded_text_vectors.failure, ) @@ -148,18 +229,49 @@ class SemanticTextIndex: Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first.""" - def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None: + def __init__( + self, + max_entries: int = DEFAULT_MAX_CACHED_VECTORS, + too_long_ttl_seconds: float = DEFAULT_TOO_LONG_TTL_SECONDS, + clock: Callable[[], float] = time.monotonic, + ) -> None: self._max_entries: Final = max_entries + self._too_long_ttl_seconds: Final = too_long_ttl_seconds + self._clock: Final = clock self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({}) + self._too_long: Mapping[_CacheKey, float] = MappingProxyType({}) def _cached(self, embedding_model: str) -> Mapping[str, Vector]: return MappingProxyType( {text: vector for (model, text), vector in self._vectors.items() if model == embedding_model} ) + def _too_long_for(self, embedding_model: str, now: float) -> frozenset[str]: + return frozenset( + text + for (model, text), seen_at in self._too_long.items() + if model == embedding_model and now - seen_at < self._too_long_ttl_seconds + ) + + def _merged_too_long( + self, embedding_model: str, embedded: _Embedded, texts: Sequence[str], now: float + ) -> Mapping[_CacheKey, float]: + fresh: Final = MappingProxyType( + {key: seen_at for key, seen_at in self._too_long.items() if now - seen_at < self._too_long_ttl_seconds} + ) + newly_too_long: Final = ( + (embedding_model, text) + for text in texts + if text in embedded.too_long and (embedding_model, text) not in fresh + ) + ordered: Final = MappingProxyType({**fresh, **dict.fromkeys(newly_too_long, now)}) + return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None))) + def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]: dimension: Final = len(embedded.query_vector) - touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts}) + touched: Final = MappingProxyType( + {(embedding_model, text): embedded.vectors[text] for text in texts if text in embedded.vectors} + ) untouched: Final = MappingProxyType( { key: vector @@ -175,14 +287,27 @@ class SemanticTextIndex: async def scores( self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str - ) -> tuple[float, ...] | EmbeddingFailed: - """Cosine similarity of `query` to each entry of `texts`, in the same order.""" + ) -> tuple[float | None, ...] | EmbeddingFailed: + """Cosine similarity for embeddable entries of `texts`, with None for omitted entries.""" if not texts: return () - embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model)) + now: Final = self._clock() + embedded: Final = await _embed_query_and_texts( + embed, query, texts, self._cached(embedding_model), self._too_long_for(embedding_model, now) + ) if isinstance(embedded, EmbeddingFailed): return embedded - if not _same_dimension(embedded.query_vector, embedded.vectors, texts): + same_dimension: Final = _same_dimension(embedded.query_vector, embedded.vectors, texts) + if embedded.failure is not None: + if same_dimension: + self._vectors = self._merged(embedding_model, embedded, texts) + self._too_long = self._merged_too_long(embedding_model, embedded, texts, now) + return embedded.failure + if not same_dimension: return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions") self._vectors = self._merged(embedding_model, embedded, texts) - return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts) + self._too_long = self._merged_too_long(embedding_model, embedded, texts, now) + return tuple( + cosine_similarity(embedded.query_vector, embedded.vectors[text]) if text in embedded.vectors else None + for text in texts + ) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 593f7c0503f..0cbce6339fd 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -11,9 +11,9 @@ import subprocess import time import urllib import urllib.parse -from collections.abc import Callable +from collections.abc import Awaitable, Callable, Coroutine from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy.db.db_span import db_span @@ -96,10 +96,20 @@ class _PrismaDrainTracker: await self._drained.wait() +_EngineResponseT: Final = TypeVar("_EngineResponseT") + + class _TrackedPrismaEngine: + """Engine wrapper that counts in-flight work and never lets a cancelled caller drop an engine request. + + A dropped request leaves the engine holding a busy pool connection and can crash it once that statement + finishes, so every request runs to completion in its own task and only the caller is cancelled. + """ + def __init__(self, engine: _PrismaEngine, tracker: _PrismaDrainTracker) -> None: self._engine = engine self.tracker = tracker + self._detached: frozenset[asyncio.Task[object]] = frozenset() @property def process(self) -> _PrismaProcess: @@ -109,7 +119,60 @@ class _TrackedPrismaEngine: return getattr(self._engine, name) async def query(self, content: str, *, tx_id: str | None) -> object: + return await self._finish_even_if_cancelled(self._tracked_query(content, tx_id)) + + async def start_transaction(self, *, content: str) -> str: + return await self._finish_even_if_cancelled( + self._tracked_start_transaction(content), abandoned=self.rollback_transaction + ) + + async def commit_transaction(self, tx_id: str) -> None: + await self._finish_even_if_cancelled(self._tracked_commit_transaction(tx_id)) + + async def rollback_transaction(self, tx_id: str) -> None: + await self._finish_even_if_cancelled(self._tracked_rollback_transaction(tx_id)) + + async def _finish_even_if_cancelled( + self, + request: Coroutine[object, object, _EngineResponseT], + abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None = None, + ) -> _EngineResponseT: self.tracker.begin_operation() + task: Final = asyncio.ensure_future(request) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + self._detached = self._detached.union((task,)) + task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned)) + raise + + def _detached_request_finished( + self, + done: "asyncio.Task[_EngineResponseT]", + abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None, + ) -> None: + self._detached = self._detached.difference((done,)) + if done.cancelled(): + return + failure: Final = done.exception() + if failure is not None: + verbose_proxy_logger.debug("Prisma request abandoned by a cancelled caller failed: %s", failure) + return + if abandoned is None: + return + settle: Final = asyncio.ensure_future(abandoned(done.result())) + self._detached = self._detached.union((settle,)) + settle.add_done_callback(self._abandoned_transaction_settled) + + def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]") -> None: + self._detached = self._detached.difference((settle,)) + if settle.cancelled() or settle.exception() is None: + return + verbose_proxy_logger.debug( + "Rolling back a transaction abandoned by a cancelled caller failed: %s", settle.exception() + ) + + async def _tracked_query(self, content: str, tx_id: str | None) -> object: try: if db_io_claimed(): record_db_io() @@ -121,9 +184,8 @@ class _TrackedPrismaEngine: finally: self.tracker.end_operation() - async def start_transaction(self, *, content: str) -> str: + async def _tracked_start_transaction(self, content: str) -> str: record_db_io() - self.tracker.begin_operation() try: transaction_id: Final = await self._engine.start_transaction(content=content) except (Exception, asyncio.CancelledError): @@ -132,18 +194,16 @@ class _TrackedPrismaEngine: self.tracker.transaction_started(transaction_id) return transaction_id - async def commit_transaction(self, tx_id: str) -> None: + async def _tracked_commit_transaction(self, tx_id: str) -> None: record_db_io() - self.tracker.begin_operation() try: await self._engine.commit_transaction(tx_id) finally: self.tracker.end_operation() self.tracker.transaction_finished(tx_id) - async def rollback_transaction(self, tx_id: str) -> None: + async def _tracked_rollback_transaction(self, tx_id: str) -> None: record_db_io() - self.tracker.begin_operation() try: await self._engine.rollback_transaction(tx_id) finally: diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 4c85d6148c0..86e77596e70 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -33,9 +33,15 @@ async def get_ui_config(): is_control_plane: Final = len(proxy_config.worker_registry) > 0 + raw_gateway_base_url: Final[str | None] = os.getenv("LITELLM_GATEWAY_URL") + gateway_base_url: Final[str | None] = ( + raw_gateway_base_url.strip().rstrip("/") or None if raw_gateway_base_url is not None else None + ) + return UiDiscoveryEndpoints( server_root_path=get_server_root_path(), proxy_base_url=get_proxy_base_url(), + gateway_base_url=gateway_base_url, auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso, admin_ui_disabled=admin_ui_disabled, sso_configured=sso_configured, diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py new file mode 100644 index 00000000000..edff0f8af45 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py @@ -0,0 +1,32 @@ +from typing import TYPE_CHECKING, Final + +from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail +from litellm.types.guardrails import SupportedGuardrailIntegrations + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> AiriaGuardrail: + import litellm + + _airia_callback: Final = AiriaGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + timeout=litellm_params.timeout, + guardrail_name=guardrail["guardrail_name"], + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_airia_callback) + + return _airia_callback + + +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.AIRIA.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.AIRIA.value: AiriaGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py new file mode 100644 index 00000000000..442eca3f9a4 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -0,0 +1,246 @@ +import os +import uuid +from collections.abc import Mapping +from typing import ( + Any, # noqa: TID251 # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs + Final, + Literal, +) + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # the decorator is untyped upstream +) +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # the factory's params are untyped upstream + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel +from litellm.types.utils import GenericGuardrailAPIInputs + +GUARDRAIL_PATH: Final = "/v1/guardrails/litellm" + +ACTION_NONE: Final = "NONE" +ACTION_BLOCKED: Final = "BLOCKED" +ACTION_INTERVENED: Final = "GUARDRAIL_INTERVENED" + +DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy." + +SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call) + + +def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool: + """Whether this is a streamed response on a route the framework can only block, not rewrite. + + ``incremental_diff``, the mode that writes a rewrite into the stream, exists only for OpenAI chat + completions: ``UnifiedLLMGuardrails._resolve_transform_call_type`` sends every other streamed route + to ``block_only``, which drops the rewrite and replays the buffered original. On those routes a + REDACT verdict can only be honoured as a block. The route check mirrors that resolver so the two + cannot disagree, and a route it cannot resolve fails closed the same way the resolver does. + """ + if request_data is None or not _is_streamed(request_data): + return False + route: Final = _request_route(request_data) + call_types: Final = get_call_types_for_route(route) if route else None + if not call_types: + return True + # Deferred like the resolver's own import: the handler package imports the proxy. + from litellm.llms import load_guardrail_translation_mappings # noqa: PLC0415 # circular import + from litellm.llms.openai.chat.guardrail_translation.handler import ( # noqa: PLC0415 # circular import + OpenAIChatCompletionsHandler, + ) + + handler: Final[object] = load_guardrail_translation_mappings().get(call_types[0]) + return not (isinstance(handler, type) and issubclass(handler, OpenAIChatCompletionsHandler)) + + +def _is_streamed(request_data: Mapping[str, object]) -> bool: + if request_data.get("stream") is True: + return True + proxy_request: Final = _as_object(request_data.get("proxy_server_request")) + body: Final = _as_object(proxy_request.get("body")) if proxy_request is not None else None + return body is not None and body.get("stream") is True + + +def _request_route(request_data: Mapping[str, object]) -> str | None: + """The route the proxy stamps as ``user_api_key_request_route``, in whichever metadata bucket the + route uses: ``metadata`` for chat completions, ``litellm_metadata`` for Responses and Anthropic.""" + for bucket_name in ("metadata", "litellm_metadata"): + bucket = _as_object(request_data.get(bucket_name)) + route = bucket.get("user_api_key_request_route") if bucket is not None else None + if isinstance(route, str) and route: + return route + return None + + +def _as_object(value: object) -> Mapping[str, object] | None: + """`value` as a JSON object, or None. States the key type once so the checker does not see + `dict[Unknown, Unknown]` at every read.""" + return value if isinstance(value, dict) else None # pyright: ignore[reportUnknownVariableType] # isinstance narrows to dict[Unknown, Unknown]; JSON keys are str + + +class AiriaGuardrail(CustomGuardrail): + """Evaluates prompts and responses against your Airia guardrail policy.""" + + def __init__( + self, + api_base: str | None = None, + api_key: str | None = None, + timeout: float | None = None, + supported_event_hooks: list[GuardrailEventHooks] | None = None, # mutable-ok: matches CustomGuardrail.__init__ + async_handler: AsyncHTTPHandler | None = None, # dependency-injected in tests; defaults to a real client + **kwargs: Any, # kwargs-ok: passed straight through to CustomGuardrail.__init__ + ) -> None: + resolved_timeout: Final = timeout or float(os.getenv("AIRIA_TIMEOUT", "10")) + + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + params={"timeout": httpx.Timeout(timeout=resolved_timeout, connect=5.0)}, + ) + + self.api_base = (api_base or os.getenv("AIRIA_GATEWAY_URL", "")).rstrip("/") + self.api_key = api_key or os.getenv("AIRIA_API_KEY") + self.streaming_transform_mode: Final[Literal["block_only", "incremental_diff"]] = "incremental_diff" + self.streaming_end_of_stream_only: Final = True + # incremental_diff exists only for streamed chat completions. Every other streamed route falls + # back to block_only, where the framework releases the original text live unless told to hold + # it; holding it keeps the text behind the end-of-stream verdict, and apply_guardrail turns a + # REDACT that block_only cannot deliver into a block. Never read on the incremental_diff path. + self.streaming_buffer_until_moderated: Final = True + + if not self.api_base: + raise ValueError("AiriaGuardrail requires api_base, or the AIRIA_GATEWAY_URL environment variable.") + if not self.api_key: + raise ValueError("AiriaGuardrail requires api_key, or the AIRIA_API_KEY environment variable.") + + self.optional_params = kwargs + super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped **kwargs + supported_event_hooks=supported_event_hooks or [*SUPPORTED_EVENT_HOOKS], + **kwargs, + ) + verbose_proxy_logger.info("AiriaGuardrail initialized with gateway: %s", self.api_base) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: the base class declares this parameter as a dict + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + call_id: Final = ( + logging_obj.litellm_call_id + if logging_obj + else (request_data.get("litellm_call_id") if request_data else None) + ) or str(uuid.uuid4()) + + try: + response: Final = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # the handler's `data` param is an untyped dict upstream + f"{self.api_base}{GUARDRAIL_PATH}", + json={ + "input_type": input_type, + "texts": inputs.get("texts") or (), + "images": inputs.get("images") or (), + "structured_messages": inputs.get("structured_messages") or (), + "tools": inputs.get("tools") or (), + "tool_calls": inputs.get("tool_calls") or (), + "model": inputs.get("model"), + "litellm_call_id": call_id, + }, + headers={"Authorization": f"Bearer {self.api_key}"}, + ) + response.raise_for_status() + raw_body: Final[object] = response.json() + body: Final = _as_object(raw_body) + if body is None: + raise TypeError(f"expected a JSON object in the response body, got {type(raw_body).__name__}") + except Exception as error: + verbose_proxy_logger.error( + "Airia guardrail could not evaluate the request (litellm_call_id=%s, input_type=%s): %s", + call_id, + input_type, + error, + ) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=f"Airia guardrail could not evaluate the request: {error}", + blocked_content=False, + ) from error + + action: Final = body.get("action") + + if action == ACTION_BLOCKED: + reason: Final = body.get("blocked_reason") + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=reason if isinstance(reason, str) and reason else DEFAULT_BLOCKED_MESSAGE, + blocked_content=True, + ) + + if action == ACTION_INTERVENED: + if input_type == "response" and _stream_cannot_carry_a_rewrite(request_data): + raise self._blocked() + return self._rewritten(body, inputs) + + if action != ACTION_NONE: + raise self._blocked() + + return inputs + + def _rewritten( + self, + body: Mapping[str, object], + inputs: GenericGuardrailAPIInputs, + ) -> GenericGuardrailAPIInputs: + texts: Final = body.get("texts") + structured_messages: Final = body.get("structured_messages") + if not self._is_applicable_rewrite(texts, inputs.get("texts")): + raise self._blocked() + if not self._is_applicable_rewrite(structured_messages, inputs.get("structured_messages")): + raise self._blocked() + # An intervention that rewrites nothing would hand the original content through under a + # verdict that said it must not pass. + if texts is None and structured_messages is None: + raise self._blocked() + + rewritten: Final[GenericGuardrailAPIInputs] = {**inputs} + if isinstance(texts, list): + rewritten["texts"] = texts + if isinstance(structured_messages, list): + rewritten["structured_messages"] = structured_messages + return rewritten + + @staticmethod + def _is_applicable_rewrite(rewrite: object, original: object) -> bool: + """A rewrite is safe to apply positionally only when it is absent (this field was not + touched) or a list the same length as what was sent. A different count would misalign + the positional write-back downstream, dropping content or leaving other fields unchanged. + + A rewrite for a field that was never sent is a mismatch too, not a no-op: post_call sends + only `texts`, so `structured_messages: []` would otherwise count as an applied rewrite + while the `texts` that carried the content passed through untouched. + """ + if rewrite is None: + return True + if not isinstance(original, list) or not isinstance(rewrite, list): + return False + return len(rewrite) == len(original) # pyright: ignore[reportUnknownArgumentType] # only the counts matter; the items are the gateway's + + def _blocked(self) -> GuardrailRaisedException: + return GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=DEFAULT_BLOCKED_MESSAGE, + blocked_content=True, + ) + + @staticmethod + def get_config_model() -> type[AiriaGuardrailConfigModel] | None: + return AiriaGuardrailConfigModel diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 3f05c3a407a..5f31aa788e2 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -2,7 +2,7 @@ import asyncio from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, status from typing_extensions import ReadOnly, TypedDict @@ -45,6 +45,9 @@ from litellm.types.access_group import ( AccessGroupUpdateRequest, ) +if TYPE_CHECKING: + from prisma import Json + router: Final = APIRouter( tags=["access group management"], ) @@ -57,6 +60,9 @@ class _AccessGroupRecord(Protocol): @property def access_mcp_server_ids(self) -> Sequence[str] | None: ... + @property + def mcp_tool_permissions(self) -> "Mapping[str, Sequence[str]] | Json | None": ... + @property def access_agent_ids(self) -> Sequence[str] | None: ... @@ -179,6 +185,12 @@ def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None: ) +def _tool_permissions_json(value: Mapping[str, Sequence[str]] | None) -> "Json": + from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools + + return Json(None) if value is None else Json({server: list(tools) for server, tools in value.items()}) + + @dataclass(frozen=True, slots=True) class _ResourceNames: mcp_servers: Mapping[str, str] @@ -544,7 +556,9 @@ async def create_access_group( "description": data.description, "access_model_names": data.access_model_names or [], "access_mcp_server_ids": data.access_mcp_server_ids or [], + "mcp_tool_permissions": _tool_permissions_json(data.mcp_tool_permissions), "access_agent_ids": data.access_agent_ids or [], + "access_skill_names": list(data.access_skill_names or ()), "assigned_team_ids": data.assigned_team_ids or [], "assigned_key_ids": data.assigned_key_ids or [], "created_by": user_api_key_dict.user_id, @@ -636,19 +650,18 @@ async def update_access_group( update_fields: Final = data.model_dump(exclude_unset=True) update_data: Final[dict] = {"updated_by": user_api_key_dict.user_id} for field, value in update_fields.items(): - if ( - field - in ( - "assigned_team_ids", - "assigned_key_ids", - "access_model_names", - "access_mcp_server_ids", - "access_agent_ids", - ) - and value is None + if field in ( + "assigned_team_ids", + "assigned_key_ids", + "access_model_names", + "access_mcp_server_ids", + "access_agent_ids", + "access_skill_names", ): - value = [] - update_data[field] = value + value = list(value or ()) + update_data[field] = ( + _tool_permissions_json(data.mcp_tool_permissions) if field == "mcp_tool_permissions" else value + ) # Initialize delta lists before the try block so they remain accessible # for cache updates after the transaction, even if an error path is added later. diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 04886072856..a4ffa5cfd4d 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -9,9 +9,9 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass from functools import partial -from itertools import chain +from itertools import chain, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, overload +from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypeAlias, overload from fastapi import ( APIRouter, @@ -29,6 +29,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.constants import SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.models.user import SCIMPlaceholder from litellm.proxy._types import ( @@ -70,6 +71,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports handle_exception_on_proxy, premium_user_check, ) +from litellm.repositories.chunked_in import find_many_in from litellm.repositories.table_repositories import ( InvitationLinkRepository, OrganizationMembershipRepository, @@ -565,6 +567,111 @@ class _CaseInsensitiveMatch(TypedDict): mode: ReadOnly[str] +_InFilter: Final = TypedDict("_InFilter", {"in": ReadOnly[Sequence[str]]}) +_CaseInsensitiveInFilter: Final = TypedDict( + "_CaseInsensitiveInFilter", {"in": ReadOnly[Sequence[str]], "mode": ReadOnly[str]} +) + + +class _UserIdIn(TypedDict): + user_id: ReadOnly[_InFilter] + + +class _SsoUserIdIn(TypedDict): + sso_user_id: ReadOnly[_InFilter] + + +class _UserEmailIn(TypedDict): + user_email: ReadOnly[_CaseInsensitiveInFilter] + + +_IdentityRowsClause: TypeAlias = _UserIdIn | _SsoUserIdIn | _UserEmailIn + + +class _IdentityRowsWhere(TypedDict): + OR: ReadOnly[Sequence[_IdentityRowsClause]] + + +@dataclass(frozen=True, slots=True) +class _IdentityIndex: + user_ids: frozenset[str] + by_sso_user_id: Mapping[str, tuple[str, ...]] + by_email: Mapping[str, tuple[str, ...]] + + +async def _identity_rows_for_chunk(chunk: Sequence[str], prisma_client: PrismaClient) -> tuple[LiteLLM_UserTable, ...]: + subjects: Final = list(dict.fromkeys(value.strip() for value in chunk)) + where: Final[_IdentityRowsWhere] = { + "OR": [ + {"user_id": {"in": list(chunk)}}, # bounded-ok: chunk cap is 10,000 + {"sso_user_id": {"in": subjects}}, # bounded-ok: subjects are chunk-bounded + {"user_email": {"in": subjects, "mode": "insensitive"}}, # bounded-ok: subjects are chunk-bounded + ] + } + return tuple(await _table(UserRepository(prisma_client)).find_many(where=where)) + + +async def _identity_rows_for_values( + values: Sequence[str], prisma_client: PrismaClient +) -> tuple[LiteLLM_UserTable, ...]: + unique_values: Final = tuple(dict.fromkeys(values)) + if not unique_values: + return () + + batch_size: Final = SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE + chunks: Final = tuple( + unique_values[start : start + batch_size] for start in range(0, len(unique_values), batch_size) + ) + chunk_rows: Final = tuple([await _identity_rows_for_chunk(chunk, prisma_client) for chunk in chunks]) + rows: Final = tuple(chain.from_iterable(chunk_rows)) + return tuple({row.user_id: row for row in rows}.values()) + + +def _identity_index(rows: Sequence[LiteLLM_UserTable]) -> _IdentityIndex: + sso_entries: Final = tuple(sorted((row.sso_user_id, row.user_id) for row in rows if row.sso_user_id is not None)) + email_entries: Final = tuple( + sorted((row.user_email.lower(), row.user_id) for row in rows if row.user_email is not None) + ) + by_sso_user_id: Final = MappingProxyType( + { + subject: tuple(sorted({user_id for _, user_id in matches})) + for subject, matches in groupby(sso_entries, key=lambda entry: entry[0]) + } + ) + by_email: Final = MappingProxyType( + { + email: tuple(sorted({user_id for _, user_id in matches})) + for email, matches in groupby(email_entries, key=lambda entry: entry[0]) + } + ) + return _IdentityIndex( + user_ids=frozenset(row.user_id for row in rows), + by_sso_user_id=by_sso_user_id, + by_email=by_email, + ) + + +def _users_named_by(value: str, index: _IdentityIndex) -> tuple[str, ...]: + subject: Final = value.strip() + named: Final = frozenset( + chain( + index.by_sso_user_id.get(subject, ()), + index.by_email.get(subject.lower(), ()), + ) + ) + return tuple(sorted(named)) + + +def _accounts_named_by(value: str, index: _IdentityIndex) -> tuple[str, ...]: + named: Final = frozenset( + chain( + _users_named_by(value, index), + (value,) if value in index.user_ids else (), + ) + ) + return tuple(sorted(named)) + + async def _users_named_by_member_value( value: str, prisma_client: PrismaClient, *, take: int | None = 2 ) -> tuple[str, ...]: @@ -597,50 +704,35 @@ async def _users_named_by_member_value( return tuple(dict.fromkeys(row.user_id for row in rows)) -async def _accounts_named_by_member_value(value: str, prisma_client: PrismaClient) -> tuple[str, ...]: - """Every user id this member value names, by user id, SSO identity or email. +async def _scim_team_ids_among(values: Sequence[str], prisma_client: PrismaClient) -> frozenset[str]: + unique_values: Final = tuple(dict.fromkeys(values)) + if not unique_values: + return frozenset() - Classification needs to know whether the value is one account's ``user_id`` and - whether it names any other account, so all three fields are read in one pass. The - id is compared exactly and unstripped, as a primary key lookup would; the - identities compare as ``_users_named_by_member_value`` describes. Two rows are - enough to tell one account from several, so the read stops there. Only a full - read that lacks the row keyed by the value leaves that row's existence open, and - only then is the id read on its own. - """ - subject: Final = value.strip() - email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"} - users: Final = _table(UserRepository(prisma_client)) - rows: Final = await users.find_many( - where={ - "OR": [ - {"user_id": value}, - {"sso_user_id": subject}, - {"user_email": email}, - ], - }, - take=2, + teams: Final = await find_many_in( + _table(TeamRepository(prisma_client)), + "team_id", + unique_values, + chunk_size=SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE, ) - named: Final = tuple(dict.fromkeys(row.user_id for row in rows)) - if len(named) < 2 or value in named: - return named - keyed: Final = await users.find_unique(where={"user_id": value}) - return named if keyed is None else (value, *named) + return frozenset(team.team_id for team in teams if _team_metadata_has_scim_provenance(team.metadata)) -async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember: +def _classify_group_member( + member: SCIMMember, named: tuple[str, ...], scim_team_ids: frozenset[str] +) -> _ClassifiedGroupMember: """ Decide what a single SCIM group member refers to. A LiteLLM team only holds users, so a member is dropped when it declares a type other than ``User`` or when its id names an existing team. Both of those checks - are placed around the user lookup rather than before it, because the id of a + are placed around identity resolution rather than before it, because the id of a real user is the one thing that outranks them: - ``"type": "Group"`` (what Entra sends for a nested group) is dropped without - a lookup. This bug provisioned nested group GUIDs as users, so those rows + identity resolution. This bug provisioned nested group GUIDs as users, so those rows exist in the wild and would otherwise resolve as members all over again. - - any other unrecognized type is dropped only after the user lookup misses. + - any other unrecognized type is dropped only after identity resolution misses. Clients do send non-canonical types on real members (RFC 7643 defines ``direct`` for ``User.groups``), and dropping a live user over one would revoke that user's team access on the next full sync. @@ -650,8 +742,8 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient even if some team happens to share the id, and a team created here rather than through SCIM is not evidence of anything about the member. - When those checks miss on an otherwise user-shaped member, its value is looked - up as an SSO identity or an email, and a match resolves to that user's + When those checks miss on an otherwise user-shaped member, its value is matched + against SSO identities or emails, and a match resolves to that user's ``user_id``. A value that names more than one account is ambiguous rather than unknown: it names a real person we cannot identify, so it is neither guessed at nor provisioned. @@ -670,7 +762,6 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if member_type == "group": return _SkippedGroupMember(value=value, reason="nested_group") - named: Final = await _accounts_named_by_member_value(value, prisma_client) if value in named: shared_with: Final = tuple(other for other in named if other != value) if shared_with: @@ -687,10 +778,8 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if member_type is not None and member_type != "user": return _SkippedGroupMember(value=value, reason="non_user_type") - if member_type is None: - team: Final = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": value}) - if team is not None and _team_metadata_has_scim_provenance(team.metadata): - return _SkippedGroupMember(value=value, reason="existing_team") + if member_type is None and value in scim_team_ids: + return _SkippedGroupMember(value=value, reason="existing_team") if len(named) == 1: verbose_proxy_logger.info( @@ -837,16 +926,10 @@ async def _member_ids_to_drop( HTTPException: 400 when a member id names more than one current member. """ written: Final = frozenset(_member_value(member) for member in members) + rows: Final = await _identity_rows_for_values(sorted(written), prisma_client) + index: Final = _identity_index(rows) matched: Final = tuple( - [ - ( - value, - _roster_entries_named_by( - value, roster, await _users_named_by_member_value(value, prisma_client, take=None) - ), - ) - for value in sorted(written) - ] + (value, _roster_entries_named_by(value, roster, _users_named_by(value, index))) for value in sorted(written) ) undecidable: Final = tuple(value for value, entries in matched if len(entries) > 1) if undecidable: @@ -881,7 +964,22 @@ async def _resolve_group_member_ids( other than a user. 500 when a member's user row can neither be created nor found. """ - classified: Final = tuple([await _classify_group_member(member, prisma_client) for member in members]) + values: Final = tuple(_member_value(member) for member in members) + identity_values: Final = tuple( + value for member, value in zip(members, values) if _normalized_member_type(member) != "group" + ) + rows: Final = await _identity_rows_for_values(identity_values, prisma_client) + index: Final = _identity_index(rows) + named_by: Final = MappingProxyType({value: _accounts_named_by(value, index) for value in dict.fromkeys(values)}) + team_candidates: Final = tuple( + value + for member, value in zip(members, values) + if _normalized_member_type(member) is None and value not in named_by[value] + ) + scim_teams: Final = await _scim_team_ids_among(team_candidates, prisma_client) + classified: Final = tuple( + _classify_group_member(member, named_by[value], scim_teams) for member, value in zip(members, values) + ) partition: Final = _partition_classified_members(classified) for skipped in partition.skipped: diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 67cbb0be45d..3eb54c2147a 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -17,6 +17,7 @@ from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never from litellm._logging import verbose_proxy_logger +from litellm.models.budget import LiteLLM_BudgetTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.verification_token import LiteLLM_VerificationToken from litellm.proxy._types import ( @@ -32,7 +33,16 @@ TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_field # TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset( - {"tpm_limit", "rpm_limit", "max_budget", "soft_budget"} + { + "tpm_limit", + "rpm_limit", + "max_budget", + "soft_budget", + "team_member_budget", + "team_member_budget_duration", + "team_member_tpm_limit", + "team_member_rpm_limit", + } ) TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION: Final = "raise_max_budget" TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects" @@ -66,6 +76,19 @@ _NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"}) _SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields" +@dataclass(frozen=True, slots=True) +class TeamMemberDefaults: + team_member_budget: float | None = None + team_member_budget_duration: str | None = None + team_member_tpm_limit: int | None = None + team_member_rpm_limit: int | None = None + + +TEAM_MEMBER_DEFAULT_FIELDS: Final[frozenset[str]] = frozenset( + {"team_member_budget", "team_member_budget_duration", "team_member_tpm_limit", "team_member_rpm_limit"} +) + + @dataclass(frozen=True, slots=True) class TeamAdminEditAllowed: request: UpdateTeamRequest @@ -250,34 +273,76 @@ def _metadata_changes( ) +def stored_team_member_budget_id(team: LiteLLM_TeamTable) -> str | None: + budget_id: Final = _stored_metadata(_JSON_OBJECT.validate_json(team.model_dump_json())).get("team_member_budget_id") + return budget_id if isinstance(budget_id, str) else None + + +def team_member_defaults(budget: LiteLLM_BudgetTable | None) -> TeamMemberDefaults | None: + if budget is None: + return None + return TeamMemberDefaults( + team_member_budget=budget.max_budget, + team_member_budget_duration=budget.budget_duration, + team_member_tpm_limit=budget.tpm_limit, + team_member_rpm_limit=budget.rpm_limit, + ) + + def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]: table: Final = existing_row.litellm_model_table return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY def _column_changed( - field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable + field: str, + submitted: Mapping[str, object], + existing: Mapping[str, object], + existing_row: LiteLLM_TeamTable, + member_defaults: Mapping[str, object], ) -> bool: if field == "model_aliases": return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row) + if field in member_defaults: + return submitted.get(field) != member_defaults[field] if field in LiteLLM_TeamTable.model_fields: return submitted.get(field) != existing.get(field) return True -def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]: +def _member_defaults_by_field(member_defaults: TeamMemberDefaults | None) -> Mapping[str, object]: + stored: Final = member_defaults or TeamMemberDefaults() + return MappingProxyType( + { + "team_member_budget": stored.team_member_budget, + "team_member_budget_duration": stored.team_member_budget_duration, + "team_member_tpm_limit": stored.team_member_tpm_limit, + "team_member_rpm_limit": stored.team_member_rpm_limit, + } + ) + + +def changed_team_fields( + data: UpdateTeamRequest, + existing_row: LiteLLM_TeamTable, + member_defaults: TeamMemberDefaults | None = None, +) -> frozenset[str]: """Logical field names whose stored value the request would change. Request and stored row are compared as JSON values so both sides share one representation. Fields the server folds into metadata are attributed to their own name whether they arrive top-level or inside - ``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored - counterpart on the team row count as changed whenever they are sent. + ``metadata``; anything else in ``metadata`` is attributed to ``metadata``. The ``team_member_*`` defaults + compare against the team's member budget row (``None`` when it has none, read as all null). Fields with no + stored counterpart count as changed whenever they are sent. """ submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True)) existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json()) column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS + stored_member_defaults: Final = _member_defaults_by_field(member_defaults) column_changes: Final = frozenset( - field for field in column_fields if _column_changed(field, submitted, existing, existing_row) + field + for field in column_fields + if _column_changed(field, submitted, existing, existing_row, stored_member_defaults) ) return column_changes | _metadata_changes(data, submitted, existing) @@ -295,14 +360,21 @@ def team_admin_edit_verdict( data: UpdateTeamRequest, existing: LiteLLM_TeamTable, permitted: frozenset[str], + member_defaults: TeamMemberDefaults | None = None, ) -> TeamAdminEditVerdict: if not permitted: return TeamAdminEditingDisabled() - changed: Final = changed_team_fields(data, existing) + changed: Final = changed_team_fields(data, existing, member_defaults) blocked: Final = sorted(changed - permitted) if blocked: return TeamAdminFieldNotPermitted(field=blocked[0]) - return TeamAdminEditAllowed(request=_only_changes(data, changed)) + creates_member_budget: Final = member_defaults is None and bool(changed & TEAM_MEMBER_DEFAULT_FIELDS) + sent_for_new_member_budget: Final = ( + frozenset(data.model_fields_set) & TEAM_MEMBER_DEFAULT_FIELDS & permitted + if creates_member_budget + else frozenset[str]() + ) + return TeamAdminEditAllowed(request=_only_changes(data, changed | sent_for_new_member_budget)) def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 9cc1afd7beb..a76fc9254b8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -18,6 +18,7 @@ from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timezone +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -164,10 +165,14 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( ) from litellm.proxy.management_endpoints.team_admin_field_permissions import ( SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, + TEAM_MEMBER_DEFAULT_FIELDS, + TeamMemberDefaults, resolve_team_admin_editable_fields, + stored_team_member_budget_id, team_admin_edit_verdict, team_admin_may_raise_max_budget, team_admin_request_or_raise, + team_member_defaults, ) from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, @@ -2365,12 +2370,19 @@ async def update_team( if access_role is None: team_access_denied() if access_role == "team_admin": + permitted: Final = resolve_team_admin_editable_fields( + _general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + ) + reads_member_defaults: Final = bool(permitted) and bool(data.model_fields_set & TEAM_MEMBER_DEFAULT_FIELDS) data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below team_admin_edit_verdict( data=data, existing=existing_team, - permitted=resolve_team_admin_editable_fields( - _general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + permitted=permitted, + member_defaults=( + await _stored_team_member_defaults(prisma_client, existing_team) + if reads_member_defaults + else None ), ) ) @@ -4899,6 +4911,15 @@ async def validate_membership(user_api_key_dict: UserAPIKeyAuth, team_table: Lit ) +async def _stored_team_member_defaults( + prisma_client: PrismaClient, team: LiteLLM_TeamTable +) -> TeamMemberDefaults | None: + budget_id: Final = stored_team_member_budget_id(team) + if budget_id is None: + return None + return team_member_defaults(await BudgetRepository(prisma_client).find_by_id(budget_id)) + + async def _add_team_member_budget_table( team_member_budget_id: str, prisma_client: PrismaClient, @@ -4952,6 +4973,16 @@ def _parent_organization_models(team_row: BaseModel) -> list[str] | None: return organization.models if organization is not None else None +def _effective_access_group_skill_names( + resolved_groups: "Sequence[prisma_models.LiteLLM_AccessGroupTable]", unresolved_ids: Sequence[str] +) -> tuple[str, ...]: + """Skill names the team's keys can actually use: none while a listed group has no row, + because the skill check then clamps the team to public plugins (see skill_visibility).""" + if unresolved_ids: + return () + return tuple(frozenset(chain.from_iterable(group.access_skill_names or () for group in resolved_groups))) + + async def _resolve_team_access_group_resources( _team_info: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: @@ -4963,6 +4994,9 @@ async def _resolve_team_access_group_resources( resolved_groups: Final = tuple( ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup ) + unresolved_ids: Final = tuple( + ag_id for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id not in ag_lookup + ) return _team_info.model_copy( update={ "access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}), @@ -4970,6 +5004,8 @@ async def _resolve_team_access_group_resources( {s for group in resolved_groups for s in (group.access_mcp_server_ids or [])} ), "access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}), + "access_group_skill_names": _effective_access_group_skill_names(resolved_groups, unresolved_ids), + "unresolved_access_group_ids": unresolved_ids, "access_group_details": tuple( TeamAccessGroupModelGrant( access_group_id=group.access_group_id, @@ -4977,6 +5013,7 @@ async def _resolve_team_access_group_resources( models=tuple(group.access_model_names or ()), mcp_server_ids=tuple(group.access_mcp_server_ids or ()), agent_ids=tuple(group.access_agent_ids or ()), + skill_names=tuple(group.access_skill_names or ()), ) for group in resolved_groups ), @@ -5629,6 +5666,7 @@ def _convert_teams_to_response_models( **team_dict, members_count=members_count, keys_count=keys_count, + access_group_skill_names=None, ) ) return team_list @@ -5916,6 +5954,9 @@ async def list_team_v2( team_item.access_group_agent_ids = list( {a for group in team_groups for a in (group.access_agent_ids or [])} ) + team_item.access_group_skill_names = _effective_access_group_skill_names( + team_groups, tuple(ag_id for ag_id in (team_item.access_group_ids or []) if ag_id not in ag_lookup) + ) return { "teams": team_list, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index ca58247a2c9..12438614755 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable { // Resource memberships - explicit arrays per type access_model_names String[] @default([]) access_mcp_server_ids String[] @default([]) + mcp_tool_permissions Json? access_agent_ids String[] @default([]) + access_skill_names String[] @default([]) assigned_team_ids String[] @default([]) assigned_key_ids String[] @default([]) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 477d3a64339..0ab2052abb0 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -5,6 +5,8 @@ Handler for transforming responses api requests to litellm.completion requests from collections.abc import Coroutine, Mapping from typing import Final +from pydantic import TypeAdapter + import litellm from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.litellm_completion_transformation.streaming_iterator import ( @@ -21,6 +23,8 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ModelResponse +_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) + class LiteLLMCompletionTransformationHandler: def response_api_handler( @@ -107,6 +111,7 @@ class LiteLLMCompletionTransformationHandler: previous_response_id=previous_response_id, litellm_completion_request=litellm_completion_request, instructions=responses_api_request.get("instructions"), + request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 0d9e1a7dd38..732c32ff2cb 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -51,6 +51,7 @@ class ResponsesSessionHandler: @staticmethod async def get_chat_completion_message_history_for_previous_response_id( previous_response_id: str, + item_cache_control: bool = False, ) -> ChatCompletionSession: """ Return the chat completion message history for a previous response id @@ -86,6 +87,7 @@ class ResponsesSessionHandler: spend_log=spend_log, proxy_server_request_dict=proxy_server_request_dict, chat_completion_message_history=chat_completion_message_history, + item_cache_control=item_cache_control, ) ) @@ -117,6 +119,7 @@ class ResponsesSessionHandler: | ChatCompletionResponseMessage | Message ], + item_cache_control: bool = False, ): """ Extend the chat completion message history with the spend log payload @@ -148,6 +151,7 @@ class ResponsesSessionHandler: input=response_input_param, responses_api_request={}, replay_reasoning=True, + item_cache_control=item_cache_control, ) ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index c2cc876970c..d7235bd3ca8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -5,8 +5,10 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re import uuid -from collections.abc import Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass +from functools import partial +from itertools import accumulate, chain, islice from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -15,6 +17,7 @@ from typing import ( Literal, Protocol, TypeAlias, + TypeVar, cast, runtime_checkable, ) @@ -33,9 +36,12 @@ from openai.types.responses.tool_param import FunctionToolParam from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +import litellm from litellm._logging import verbose_logger from litellm.caching import InMemoryCache from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) @@ -45,6 +51,11 @@ from litellm.litellm_core_utils.hidden_params import ( set_hidden_params, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import bedrock_declared_cache_point_support +from litellm.llms.openrouter.chat.transformation import openrouter_model_supports_cache_control +from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) @@ -52,6 +63,7 @@ from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, + ChatCompletionCachedContent, ChatCompletionImageObject, ChatCompletionImageUrlObject, ChatCompletionRedactedThinkingBlock, @@ -186,6 +198,7 @@ TOOL_CALLS_CACHE: Final = InMemoryCache() _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) _RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) @@ -213,6 +226,14 @@ class _ResponsesToolCallItem(Protocol): def get(self, key: str, /) -> object: ... +class _MarkedToolCallChunk(ChatCompletionToolCallChunk): + cache_control: ReadOnly[ChatCompletionCachedContent] + + +class _MarkedToolMessage(ChatCompletionToolMessage): + cache_control: ReadOnly[ChatCompletionCachedContent] + + class _ToolFunctionDefinition(TypedDict, total=False): name: ReadOnly[str] description: ReadOnly[str] @@ -230,6 +251,258 @@ def _input_item_role(input_item: Mapping[str, object]) -> str: return cast(str, input_item.get("role") or "user") # cast-ok: client-supplied role forwarded verbatim, unvalidated +def _input_item_cache_control(input_item: Mapping[str, object]) -> ChatCompletionCachedContent | None: + cache_control: Final = input_item.get("cache_control") + if not cache_control: + return None + return cast( # cast-ok: forwarded as the client sent it, like a message mark on chat completions + ChatCompletionCachedContent, cache_control + ) + + +def _validated_cache_control(value: object) -> ChatCompletionCachedContent | None: + try: + return _CACHED_CONTENT_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _dropped_cache_control(_: object) -> None: + return None + + +def drops_unsupported_params(drop_params: object) -> bool: + return litellm.drop_params is True or normalize_drop_params(drop_params) is True + + +def _without_cache_control_key(item: object) -> object: + if not isinstance(item, dict): + return item + return {key: value for key, value in _STR_KEY_DICT_ADAPTER.validate_python(item).items() if key != "cache_control"} + + +def without_item_cache_control(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list): + return input + return cast( # cast-ok: the same input items, each minus its own cache_control key + ResponseInputParam, [_without_cache_control_key(item) for item in input] + ) + + +_PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) + + +def _bedrock_target_keeps_item_marks(model: str, drop_params: bool) -> bool: + declared: Final = bedrock_declared_cache_point_support(model) + return declared is True or (declared is None and not drop_params) + + +def _target_keeps_item_marks(model: str, custom_llm_provider: str | None, drop_params: bool) -> bool: + if supports_anthropic_cache_control(model, custom_llm_provider): + return True + if custom_llm_provider in _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: + return True + if custom_llm_provider == "openrouter": + return openrouter_model_supports_cache_control(model) + if custom_llm_provider == "snowflake": + return is_snowflake_claude_model(model) + if custom_llm_provider == "databricks": + return "claude" in model.lower() + return custom_llm_provider == "bedrock" and _bedrock_target_keeps_item_marks(model, drop_params) + + +def _refuses_marks_on_tool_blocks(model: str, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == "bedrock" and BaseAWSLLM.get_bedrock_invoke_provider(model) == "nova" + + +def _with_cache_control_on_last_block( + blocks: Sequence[object], cache_control: ChatCompletionCachedContent | None +) -> Sequence[object]: + if cache_control is None or not blocks or not isinstance(blocks[-1], dict): + return blocks + last_block: Final = _STR_KEY_DICT_ADAPTER.validate_python(blocks[-1]) + if last_block.get("cache_control") is not None: + return blocks + return [*blocks[:-1], {**last_block, "cache_control": cache_control}] + + +def _content_blocks_carrying(content: object, cache_control: ChatCompletionCachedContent | None) -> Sequence[object]: + if isinstance(content, str): + return ( + ChatCompletionTextObject(type="text", text=content) + if cache_control is None + else ChatCompletionTextObject(type="text", text=content, cache_control=cache_control), + ) + return _with_cache_control_on_last_block(_OBJECT_LIST_ADAPTER.validate_python(content), cache_control) + + +_MARK_CARRYING_MESSAGE_FIELDS: Final = ("content", "tool_calls") + + +def _mark_carrying_blocks(fields: Mapping[str, object]) -> Sequence[object]: + lists: Final = tuple(fields.get(key) for key in _MARK_CARRYING_MESSAGE_FIELDS) + return tuple( + chain.from_iterable(_OBJECT_LIST_ADAPTER.validate_python(value) for value in lists if isinstance(value, list)) + ) + + +def _block_carries_cache_control(block: object) -> bool: + return isinstance(block, dict) and _STR_KEY_DICT_ADAPTER.validate_python(block).get("cache_control") is not None + + +def _cache_control_count(message: object) -> int: + if not isinstance(message, dict): + return 0 + fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + block_marks: Final = sum(1 for block in _mark_carrying_blocks(fields) if _block_carries_cache_control(block)) + return int(fields.get("cache_control") is not None) + block_marks + + +_CacheControlRewrite: TypeAlias = Callable[[object], ChatCompletionCachedContent | None] + + +def _fields_with_cache_control(fields: Mapping[str, object], rewrite: _CacheControlRewrite) -> Mapping[str, object]: + others: Final = {key: value for key, value in fields.items() if key != "cache_control"} + if "cache_control" not in fields: + return others + cache_control: Final = rewrite(fields["cache_control"]) + return others if cache_control is None else {**others, "cache_control": cache_control} + + +def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> object: + if not isinstance(block, dict): + return block + return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite) + + +_MessageT = TypeVar("_MessageT") + + +def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: + if not isinstance(message, dict): + return message + fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + blocks_rewritten: Final = { + key: ( + [_block_with_cache_control(block, rewrite) for block in _OBJECT_LIST_ADAPTER.validate_python(value)] + if key in _MARK_CARRYING_MESSAGE_FIELDS and isinstance(value, list) + else value + ) + for key, value in fields.items() + } + return cast( # cast-ok: the same message shape with each cache marker rewritten or dropped + _MessageT, _fields_with_cache_control(blocks_rewritten, rewrite) + ) + + +def _kept_within_budget(state: tuple[int, bool], marks: int) -> tuple[int, bool]: + remaining, _ = state + if marks > remaining: + return remaining, False + return remaining - marks, True + + +def _messages_within_cache_control_limit(messages: Sequence[_MessageT], reserved_marks: int) -> Sequence[_MessageT]: + budget: Final = max(MAX_CACHE_CONTROL_BLOCKS - reserved_marks, 0) + newest_first: Final = tuple(reversed(messages)) + budget_after_each: Final = accumulate( + (_cache_control_count(message) for message in newest_first), _kept_within_budget, initial=(budget, True) + ) + kept_flags: Final = tuple(kept for _, kept in islice(budget_after_each, 1, None)) + kept_newest_first: Final = tuple( + message if kept else _with_cache_control(message, _dropped_cache_control) + for message, kept in zip(newest_first, kept_flags) + ) + return tuple(reversed(kept_newest_first)) + + +_SYSTEM_MESSAGE_FIELDS: Final = frozenset(("content", "cache_control")) + + +def _marked_developer_message_fields(message: object) -> Mapping[str, object] | None: + if not isinstance(message, dict): + return None + fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + return fields if fields.get("role") == "developer" and "cache_control" in fields else None + + +def _developer_message_as_system_message(message: _MessageT) -> _MessageT: + fields: Final = _marked_developer_message_fields(message) + if fields is None: + return message + return cast( # cast-ok: the system message the chat route maps a developer message to, mark included + _MessageT, {"role": "system", **{key: value for key, value in fields.items() if key in _SYSTEM_MESSAGE_FIELDS}} + ) + + +def _message_fields(message: object) -> Mapping[str, object] | None: + return _STR_KEY_DICT_ADAPTER.validate_python(message) if isinstance(message, dict) else None + + +def _without_tool_block_marks(message: _MessageT) -> _MessageT: + fields: Final = _message_fields(message) + if fields is None: + return message + if fields.get("role") == "tool": + return _with_cache_control(message, _dropped_cache_control) + tool_calls: Final = fields.get("tool_calls") + if not isinstance(tool_calls, list): + return message + return cast( # cast-ok: the same assistant message with its tool calls unmarked + _MessageT, + { + **fields, + "tool_calls": [ + _block_with_cache_control(tool_call, _dropped_cache_control) + for tool_call in _OBJECT_LIST_ADAPTER.validate_python(tool_calls) + ], + }, + ) + + +@dataclass(frozen=True, slots=True) +class _ItemMarkRules: + drop_params: bool + refuses_tool_block_marks: bool + external_marks: int + + def sanitized(self, message: _MessageT) -> _MessageT: + valid: Final = _with_cache_control(message, _validated_cache_control) + return _without_tool_block_marks(valid) if self.refuses_tool_block_marks else valid + + def checked(self, message: _MessageT) -> _MessageT: + return self.sanitized(message) if self.drop_params else message + + +def _item_mark_rules( + model: str, custom_llm_provider: str | None, tools: Iterable[object], request_kwargs: Mapping[str, object] +) -> _ItemMarkRules: + return _ItemMarkRules( + drop_params=drops_unsupported_params(request_kwargs.get("drop_params")), + refuses_tool_block_marks=_refuses_marks_on_tool_blocks(model, custom_llm_provider), + external_marks=AnthropicCacheControlHook.count_external_cache_breakpoints( + tools, request_kwargs.get("cache_control"), request_kwargs + ), + ) + + +def _request_messages_for_cache_points(messages: Sequence[_MessageT], rules: _ItemMarkRules) -> Sequence[_MessageT]: + as_system: Final = tuple(_developer_message_as_system_message(message) for message in messages) + if not rules.drop_params: + return as_system + return _messages_within_cache_control_limit(tuple(map(rules.checked, as_system)), rules.external_marks) + + +def _replayed_messages_for_cache_points( + replayed_messages: Sequence[_MessageT], current_messages: Sequence[object], rules: _ItemMarkRules +) -> Sequence[_MessageT]: + as_system: Final = tuple(_developer_message_as_system_message(message) for message in replayed_messages) + current_marks: Final = sum(_cache_control_count(message) for message in current_messages) + return _messages_within_cache_control_limit( + tuple(map(rules.sanitized, as_system)), current_marks + rules.external_marks + ) + + class ChatCompletionSession(TypedDict, total=False): messages: list[ AllMessageValues @@ -482,11 +755,27 @@ class LiteLLMCompletionResponsesConfig: api_base=kwargs.get("api_base"), ) + request_kwargs: Final = _STR_KEY_DICT_ADAPTER.validate_python(kwargs) + keeps_item_marks: Final = _target_keeps_item_marks( + model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params")) + ) + bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input, + responses_api_request=responses_api_request, + replay_reasoning=True, + item_cache_control=keeps_item_marks, + ) + litellm_completion_request: dict = { - "messages": LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input, - responses_api_request=responses_api_request, - replay_reasoning=True, + "messages": ( + list( + _request_messages_for_cache_points( + bridged_messages, + _item_mark_rules(model, custom_llm_provider, tools, request_kwargs), + ) + ) + if keeps_item_marks + else bridged_messages ), "model": model, "tool_choice": LiteLLMCompletionResponsesConfig._transform_tool_choice( @@ -535,6 +824,7 @@ class LiteLLMCompletionResponsesConfig: input: str | ResponseInputParam, responses_api_request: ResponsesAPIOptionalRequestParams | dict, replay_reasoning: bool = False, + item_cache_control: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -549,11 +839,13 @@ class LiteLLMCompletionResponsesConfig: sent to a model: prior-turn ``reasoning`` items are then rebuilt as assistant ``reasoning_content`` and signed ``thinking_blocks`` so the provider gets its own chain-of-thought back instead of reading it as - visible text. + visible text. ``item_cache_control`` belongs to the same callers, when + the target keeps an input item's ``cache_control`` on the message. Callers that only inspect the messages (token counting, rate limiting, - guardrail scanning) leave it off, because they need every piece of text - in the request to stay readable as message ``content``. + guardrail scanning) leave both off, because they need every piece of text + in the request to stay readable as message ``content``, with no item mark + a strict message schema could refuse. """ messages: list[ AllMessageValues @@ -571,7 +863,7 @@ class LiteLLMCompletionResponsesConfig: messages.extend( LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( - input=input, + input=input if item_cache_control else without_item_cache_control(input), replay_reasoning=replay_reasoning, ) ) @@ -583,19 +875,44 @@ class LiteLLMCompletionResponsesConfig: previous_response_id: str, litellm_completion_request: dict, instructions: str | None, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, + request_kwargs: Mapping[str, object] = MappingProxyType({}), ) -> dict: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages """ - chat_completion_session = ChatCompletionSession(messages=[], litellm_session_id=None) - if previous_response_id: - chat_completion_session = ( - await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - previous_response_id=previous_response_id, - ) - ) + request_fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(litellm_completion_request) + requested_provider: Final = request_fields.get("custom_llm_provider") + model: Final = str(request_fields.get("model") or "") + provider: Final = requested_provider if isinstance(requested_provider, str) else None + keeps_item_marks: Final = _target_keeps_item_marks( + model, provider, drops_unsupported_params(request_kwargs.get("drop_params")) + ) + lookup: Final = session_lookup or partial( + ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id, + item_cache_control=keeps_item_marks, + ) + chat_completion_session: Final = ( + await lookup(previous_response_id) + if previous_response_id + else ChatCompletionSession(messages=[], litellm_session_id=None) + ) _messages: Final = litellm_completion_request.get("messages") or [] - session_messages: Final = chat_completion_session.get("messages") or [] + replayed_messages: Final = chat_completion_session.get("messages") or [] + session_messages: Final = ( + _replayed_messages_for_cache_points( + replayed_messages, + _OBJECT_LIST_ADAPTER.validate_python(_messages), + _item_mark_rules( + model, + provider, + _OBJECT_LIST_ADAPTER.validate_python(request_fields.get("tools") or []), + request_kwargs, + ), + ) + if keeps_item_marks + else replayed_messages + ) instructions_end: Final = 1 if instructions else 0 carried_instructions: Final = None if instructions else chat_completion_session.get("instructions") leading_system_messages: Final = ( @@ -609,7 +926,7 @@ class LiteLLMCompletionResponsesConfig: # Store original _messages before combining for safety check original_new_messages: Final = _messages.copy() if _messages else [] - combined_messages = leading_system_messages + session_messages + _messages[instructions_end:] + combined_messages = [*leading_system_messages, *session_messages, *_messages[instructions_end:]] # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache @@ -631,7 +948,7 @@ class LiteLLMCompletionResponsesConfig: elif session_messages: # If we had session messages but they got filtered out, # restore them - combined_messages = session_messages + combined_messages = list(session_messages) else: # Both are empty - this likely means function_call_output had empty/invalid call_id # Provide a helpful error message @@ -955,22 +1272,23 @@ class LiteLLMCompletionResponsesConfig: if new_content is None: return None previous_content: Final = last_message.get("content") + previous_cache_control: Final = _input_item_cache_control(last_message) + new_cache_control: Final = _input_item_cache_control(new_message) content: Final = ( new_content - if not previous_content + if not previous_content and isinstance(new_content, str) else [ - block - for value in (previous_content, new_content) - for block in ( - (ChatCompletionTextObject(type="text", text=value),) - if isinstance(value, str) - else _OBJECT_LIST_ADAPTER.validate_python(value) - ) + *(_content_blocks_carrying(previous_content, previous_cache_control) if previous_content else ()), + *_content_blocks_carrying(new_content, new_cache_control), ] ) + message_cache_control: Final = ( + None if isinstance(content, list) else new_cache_control or previous_cache_control + ) merged: Final = { - **last_message, + **{key: value for key, value in last_message.items() if key != "cache_control"}, "content": content, + **({} if message_cache_control is None else {"cache_control": message_cache_control}), } return cast( # cast-ok: preserves the assistant fields and content blocks ChatCompletionAssistantMessage, merged @@ -1468,14 +1786,24 @@ class LiteLLMCompletionResponsesConfig: # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: return [] - return [ - GenericChatCompletionMessage( - role=_input_item_role(input_item), - content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ), - ) - ] + return [LiteLLMCompletionResponsesConfig._transform_responses_api_message_item(input_item, content)] + + @staticmethod + def _transform_responses_api_message_item( + input_item: Mapping[str, object], content: object + ) -> GenericChatCompletionMessage: + role: Final = _input_item_role(input_item) + content_value: Final = ( + LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) + ) + cache_control: Final = _input_item_cache_control(input_item) + if cache_control is None: + return GenericChatCompletionMessage(role=role, content=content_value) + if isinstance(content_value, str): + return GenericChatCompletionMessage(role=role, content=content_value, cache_control=cache_control) + return GenericChatCompletionMessage( + role=role, content=list(_with_cache_control_on_last_block(content_value, cache_control)) + ) @staticmethod def _reasoning_text_from_content(input_item: Mapping[str, object]) -> str | None: @@ -1665,10 +1993,14 @@ class LiteLLMCompletionResponsesConfig: except Exception: return str(output) - tool_output_message: Final = ChatCompletionToolMessage( - role="tool", - content=_normalize_function_call_output_to_tool_content(tool_call_output.get("output")), - tool_call_id=str(call_id), + tool_content: Final = _normalize_function_call_output_to_tool_content(tool_call_output.get("output")) + output_cache_control: Final = _input_item_cache_control(tool_call_output) + tool_output_message: Final = ( + ChatCompletionToolMessage(role="tool", content=tool_content, tool_call_id=str(call_id)) + if output_cache_control is None + else _MarkedToolMessage( + role="tool", content=tool_content, tool_call_id=str(call_id), cache_control=output_cache_control + ) ) _tool_use_definition: Final = TOOL_CALLS_CACHE.get_cache( @@ -1750,24 +2082,26 @@ class LiteLLMCompletionResponsesConfig: raw_name: Final = function_call.get("name") or "" namespace: Final = function_call.get("namespace") or "" qualify: Final = bool(namespace) and function_call.get("type") != "custom_tool_call" - tool_call: Final = ChatCompletionToolCallChunk( - id=function_call.get("call_id") or function_call.get("id") or "", - type="function", - function=ChatCompletionToolCallFunctionChunk( - name=f"{namespace}__{raw_name}" if qualify else raw_name, - arguments=serialize_tool_call_arguments(raw_arguments), - ), - index=0, + tool_call_id: Final = function_call.get("call_id") or function_call.get("id") or "" + tool_call_function: Final = ChatCompletionToolCallFunctionChunk( + name=f"{namespace}__{raw_name}" if qualify else raw_name, + arguments=serialize_tool_call_arguments(raw_arguments), + ) + function_call_cache_control: Final = _input_item_cache_control(function_call) + tool_call: Final = ( + ChatCompletionToolCallChunk(id=tool_call_id, type="function", function=tool_call_function, index=0) + if function_call_cache_control is None + else _MarkedToolCallChunk( + id=tool_call_id, + type="function", + function=tool_call_function, + index=0, + cache_control=function_call_cache_control, + ) ) # Create an assistant message with the tool call - chat_completion_response_message: Final = ChatCompletionResponseMessage( - tool_calls=[tool_call], - role="assistant", - content=None, # Function calls don't have content - ) - - return [chat_completion_response_message] + return [ChatCompletionResponseMessage(tool_calls=[tool_call], role="assistant", content=None)] @staticmethod def _resolve_file_id(item: Mapping[str, object]) -> object: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f6bd4bfba54..c90455e8e75 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -10,7 +10,10 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeA import httpx from pydantic import BaseModel, TypeAdapter, ValidationError -from typing_extensions import assert_never +from typing_extensions import ( + TypeIs, # noqa: TID251 # narrows a client input item before its mark is read + assert_never, +) import litellm from litellm._logging import verbose_logger @@ -568,12 +571,54 @@ def _will_bridge_to_chat_completions( ) -def _stamp_injection_points_with_dialect( +def _is_a_client_input_item( + item: object, +) -> TypeIs[Mapping[str, object]]: # guard-ok: bounds a client input item before its mark is read + return isinstance(item, Mapping) + + +_ITEM_TYPES_THE_BRIDGE_FORWARDS_MARKS_FROM: Final = frozenset(("function_call", "function_call_output")) + + +def _item_carries_a_mark_outside_the_provisional_messages(item: object) -> bool: + return ( + _is_a_client_input_item(item) + and item.get("type") in _ITEM_TYPES_THE_BRIDGE_FORWARDS_MARKS_FROM + and bool(item.get("cache_control")) + ) + + +def _function_tool_carries_a_mark(tool: Mapping[str, object]) -> bool: + return isinstance(tool, dict) and tool.get("type") == "function" and bool(tool.get("cache_control")) + + +def _cache_control_marks_outside_the_provisional_messages( + input: str | ResponseInputParam | None, tools: Iterable[Mapping[str, object]] | None +) -> int: + """Client marks on the request parts the provisional message list leaves out. + + A ``function_call`` or ``function_call_output`` item has no role, so the hook's pass here + never sees its mark, while the bridge forwards it onto the tool call or tool message it + builds, where the provider counts it against the cap. A ``function`` tool's mark reaches + the provider the same way, on the chat tool the bridge builds from it. + """ + item_marks: Final = ( + 0 + if input is None or isinstance(input, str) + else sum(1 for item in input if _item_carries_a_mark_outside_the_provisional_messages(item)) + ) + tool_marks: Final = sum(1 for tool in tools or () if _function_tool_carries_a_mark(tool)) + return item_marks + tool_marks + + +def _stamp_injection_points_for_the_hook( kwargs: dict[str, object], # mutable-ok: the points are rewritten in the caller's own kwargs for the hook to read model: str, custom_llm_provider: str | None, + input: str | ResponseInputParam | None, + tools: Iterable[Mapping[str, object]] | None, ) -> None: - """Carry the provider this layer resolved onto the points. + """Carry onto the points the provider this layer resolved and the marks its pass cannot see. The hook reads ``custom_llm_provider`` from the request kwargs, which never hold the one resolved here, and resolving the model name alone reads a Foundry deployment of an OpenAI @@ -582,8 +627,9 @@ def _stamp_injection_points_with_dialect( points: Final = configured_injection_points(kwargs.get("cache_control_injection_points")) if not points: return - kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_with_dialect( + kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_for_prompt_hook( points, + _cache_control_marks_outside_the_provisional_messages(input, tools), model, custom_llm_provider, kwargs.get("api_base") or kwargs.get("base_url"), @@ -704,7 +750,7 @@ async def aresponses( _api_base_kwarg(kwargs), ), ): - _stamp_injection_points_with_dialect(kwargs, model, custom_llm_provider) + _stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools) ( model, merged_input, @@ -841,6 +887,7 @@ def _resolve_prompt_swapped_provider( def _apply_prompt_management_to_responses_call( input: str | ResponseInputParam, + tools: Iterable[ToolParam] | None, model: str, custom_llm_provider: str | None, litellm_logging_obj: LiteLLMLoggingObj | None, @@ -875,7 +922,7 @@ def _apply_prompt_management_to_responses_call( _api_base_kwarg(kwargs), ), ): - _stamp_injection_points_with_dialect(kwargs, model, custom_llm_provider) + _stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools) ( model, merged_input, @@ -1274,6 +1321,7 @@ def responses( input, model, custom_llm_provider, deployment_model_info = _apply_prompt_management_to_responses_call( input=input, + tools=tools, model=model, custom_llm_provider=custom_llm_provider, litellm_logging_obj=litellm_logging_obj, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 36f433a4741..6ea73970aaa 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -282,6 +282,21 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro return mapped_exception +def _event_field(event: object, name: str) -> object: + return event.get(name) if _is_json_object(event) else getattr(event, name, None) + + +def _is_error_field_value(value: object) -> bool: + return isinstance(value, (str, int)) and not isinstance(value, bool) + + +def _error_event_error_object(event: object) -> object: + nested_error: Final = _event_field(event, "error") + if nested_error is not None: + return nested_error + return {field: value for field in ("message", "code") if _is_error_field_value(value := _event_field(event, field))} + + def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: if isinstance(mapped_exception, litellm.ContentPolicyViolationError): return True @@ -659,7 +674,7 @@ class BaseResponsesAPIStreamingIterator: error_obj: Final[object] = ( getattr(getattr(result, "response", None), "error", None) if chunk_type == "response.failed" - else getattr(result, "error", None) + else _error_event_error_object(result) ) mapped_exception: Final = self._map_error_event_exception(error_obj) @@ -1837,7 +1852,7 @@ _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item. def _ws_event_error(event: Mapping[str, object]) -> object: if event.get("type") == "error": - return event.get("error") + return _error_event_error_object(event) response: Final = event.get("response") return response.get("error") if _is_json_object(response) else None diff --git a/litellm/router.py b/litellm/router.py index ee8a58f21e2..9066311451b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -23,6 +23,7 @@ import weakref from collections import defaultdict from collections.abc import ( AsyncGenerator, + AsyncIterable, AsyncIterator, Awaitable, Callable, @@ -218,6 +219,7 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, AttemptedFallbackTargets, + attempted_fallback_targets_from, attempted_retries_for_request, carry_over_pre_routing_selection, carry_over_routed_deployment, @@ -229,6 +231,7 @@ from litellm.router_utils.fallback_event_handlers import ( get_fallback_model_group_for_lookup_groups, get_pre_routing_selection, has_unattempted_fallback_target, + is_bad_request_without_fallback, mid_stream_fallback_hop_kwargs, mid_stream_fallback_snapshot_kwargs, mid_stream_retry_kwargs, @@ -691,6 +694,67 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: return error.original_exception if error.original_exception is not None else error +@dataclass(frozen=True, slots=True) +class _FirstChunk: + """The chunk Router._acompletion / _completion already read off a stream; None means the provider closed it empty.""" + + chunk: ModelResponseStream | None + + +async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk: + """ + Read the first chunk while Router._acompletion still holds the request, so a + stream that fails before any token raises there, inside async_function_with_retries + and the fallback chain, instead of surfacing to the caller mid-iteration. + """ + from litellm.exceptions import MidStreamFallbackError + + try: + return _FirstChunk(chunk=await stream.__anext__()) + except StopAsyncIteration: + return _FirstChunk(chunk=None) + except MidStreamFallbackError as e: + raise _mid_stream_retry_trigger(e) from e + + +def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk: + """Sync twin of _aread_first_chunk for Router._completion.""" + from litellm.exceptions import MidStreamFallbackError + + try: + return _FirstChunk(chunk=next(stream)) + except StopIteration: + return _FirstChunk(chunk=None) + except MidStreamFallbackError as e: + raise _mid_stream_retry_trigger(e) from e + + +def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk | None) -> Iterator[object]: + """The stream with the chunk _read_first_chunk already read put back in front (None: the stream as is).""" + if taken is None: + yield from model_response + return + if taken.chunk is None: + return + yield taken.chunk + yield from model_response + + +async def _achunks_with_first( + model_response: AsyncIterable[object], taken: _FirstChunk | None +) -> AsyncIterator[object]: + """Async twin of _chunks_with_first.""" + if taken is None: + async for item in model_response: + yield item + return + if taken.chunk is None: + return + yield taken.chunk + async for item in model_response: + yield item + + def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: """ Whether `chunk` should make Router._aanthropic_messages_streaming_iterator @@ -960,6 +1024,7 @@ class Router: allowed_fails_policy: AllowedFailsPolicy | None = None, # set custom allowed fails policy cooldown_time: float | None = None, # (seconds) time to cooldown a deployment after failure disable_cooldowns: bool | None = None, + fallback_on_unsupported_params: bool = True, routing_strategy: RoutingStrategyName = "simple-shuffle", optional_pre_call_checks: OptionalPreCallChecks | None = None, routing_strategy_args: dict = {}, # just for latency-based @@ -1202,6 +1267,7 @@ class Router: self.cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS self.cooldown_cache = CooldownCache(cache=self.cache, default_cooldown_time=self.cooldown_time) self.disable_cooldowns = disable_cooldowns + self.fallback_on_unsupported_params = fallback_on_unsupported_params self.enable_health_check_routing = enable_health_check_routing self.enable_weighted_failover = enable_weighted_failover self.health_check_ignore_transient_errors = health_check_ignore_transient_errors @@ -2791,6 +2857,7 @@ class Router: model_response=response, messages=messages, initial_kwargs=input_kwargs_for_streaming_fallback, + first_chunk=_read_first_chunk(response), ) return response @@ -3050,6 +3117,7 @@ class Router: messages: list[dict[str, str]], initial_kwargs: dict, deployment_slot: contextlib.AsyncExitStack | None = None, + first_chunk: _FirstChunk | None = None, ) -> CustomStreamWrapper: """ Helper to iterate over a streaming response. @@ -3058,6 +3126,8 @@ class Router: `deployment_slot` holds the deployment's max_parallel_requests semaphore; it is released when the stream is exhausted, closed, or falls back to another deployment + + `first_chunk` carries the chunk _acompletion already read (None: iterate from the start) """ from litellm.exceptions import MidStreamFallbackError @@ -3102,7 +3172,7 @@ class Router: async def stream_with_fallbacks(): fallback_response = None # Track for cleanup in finally try: - async for item in model_response: + async for item in _achunks_with_first(model_response, first_chunk): yield item except MidStreamFallbackError as e: with anyio.CancelScope(shield=True): @@ -3204,7 +3274,7 @@ class Router: wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) # weak, so the generator closing over it does not keep the wrapper out of # refcount teardown and delay the `finally` that releases the deployment slot - wrapper_ref: Final = weakref.ref(wrapped_response) + wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) return wrapped_response @staticmethod @@ -3677,6 +3747,7 @@ class Router: model_response: CustomStreamWrapper, messages: list[dict[str, str]], initial_kwargs: dict, + first_chunk: _FirstChunk | None = None, ) -> CustomStreamWrapper: """ Sync equivalent of _acompletion_streaming_iterator. @@ -3712,7 +3783,7 @@ class Router: def stream_with_fallbacks(): fallback_response = None try: - for item in model_response: + for item in _chunks_with_first(model_response, first_chunk): yield item except MidStreamFallbackError as e: if fallbacks_disabled_for_request(initial_kwargs) or ( @@ -3808,7 +3879,7 @@ class Router: wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) # weak, for the same reason as the async twin - wrapper_ref: Final = weakref.ref(wrapped_response) + wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) return wrapped_response async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): @@ -3959,6 +4030,9 @@ class Router: and response.make_call is not None ): await response.fetch_stream() + first_chunk: Final = ( + await _aread_first_chunk(response) if isinstance(response, CustomStreamWrapper) else None + ) self.success_calls[model_name] += 1 verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) @@ -3975,6 +4049,7 @@ class Router: messages=messages, initial_kwargs=input_kwargs_for_streaming_fallback, deployment_slot=deployment_slot.pop_all(), + first_chunk=first_chunk, ) return response @@ -7649,6 +7724,7 @@ class Router: original_exception: Final = e fallback_model_group = None original_model_group: Final[str | None] = kwargs.get("model") + fallback_walk: Final = attempted_fallback_targets_from(kwargs) # pyright: ignore[reportUnknownArgumentType] # keep legacy kwargs untyped # A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier # behind the router name, and fallbacks are configured per tier, not per router. lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group) @@ -7657,13 +7733,22 @@ class Router: hop_depth: Final = kwargs.get("fallback_depth") nested_fallback_hop: Final = isinstance(hop_depth, int) and hop_depth > 0 - if disable_fallbacks is True or original_model_group is None or is_guardrail_intervention(e): + unsupported_params_without_fallback: Final = is_bad_request_without_fallback( + e, self.fallback_on_unsupported_params + ) + if ( + disable_fallbacks is True + or original_model_group is None + or is_guardrail_intervention(e) + or unsupported_params_without_fallback + ): raise e input_kwargs: Final = { "litellm_router": self, "original_exception": original_exception, **kwargs, + "attempted_targets": fallback_walk, } if "max_fallbacks" not in input_kwargs: @@ -7893,7 +7978,10 @@ class Router: and not nested_fallback_hop ): original_exception.message += format_fallback_outcome_message( - model_group, attempted_fallback_group, fallback_failure_exception_str + model_group, + attempted_fallback_group, + fallback_failure_exception_str, + skipped=fallback_walk.skipped, ) raise original_exception @@ -12731,6 +12819,7 @@ class Router: "enable_weighted_failover", "enable_tag_filtering", "tag_routing_prefix", + "fallback_on_unsupported_params", ] for var in vars_to_include: diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 6976a64ef39..7623a2bd343 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -15,7 +15,7 @@ from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure -from litellm.types.router import CredentialLiteLLMParams +from litellm.types.router import CredentialLiteLLMParams, FallbackSkipReason, SkippedFallbackTarget from litellm.types.utils import LlmProviders _V = TypeVar("_V") @@ -134,15 +134,61 @@ def format_fallback_outcome_message( model_group: str | None, fallback_model_group: Sequence[object] | None, fallback_failure_detail: str, + skipped: Sequence[SkippedFallbackTarget] = (), ) -> str: """User-facing explanation appended when the fallback orchestrator gives up and re-raises the primary error.""" lead: Final = f"\n\nLiteLLM: model group '{model_group}' failed with the error above." if not fallback_model_group: return f"{lead} No fallback was attempted." + + def fallback_target_model_group(target: object) -> str | None: + if isinstance(target, str): + return target + if not isinstance(target, Mapping): + return None + model: Final = target.get("model") + return model if isinstance(model, str) else None + + skipped_targets: Final = frozenset(target.target for target in skipped) + attempted_targets: Final = tuple( + target for target in fallback_model_group if fallback_target_model_group(target) not in skipped_targets + ) + not_authorized_targets: Final = tuple( + dict.fromkeys(target.target for target in skipped if target.reason is FallbackSkipReason.NOT_AUTHORIZED) + ) + over_budget_targets: Final = tuple( + dict.fromkeys(target.target for target in skipped if target.reason is FallbackSkipReason.OVER_BUDGET) + ) + attempted_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in attempted_targets) + not_authorized_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in not_authorized_targets) + over_budget_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in over_budget_targets) + attempted_failure: Final = ( + (f"Fallback to {attempted_text} also failed: {fallback_failure_detail}",) + if fallback_failure_detail and attempted_targets + else () + ) + not_authorized_sentence: Final = ( + ( + f"Fallback to {not_authorized_text} was skipped because the caller is not authorized to call " + f"{'it' if len(not_authorized_targets) == 1 else 'them'}.", + ) + if not_authorized_targets + else () + ) + over_budget_sentence: Final = ( + ( + f"Fallback to {over_budget_text} was skipped because the caller is over budget for " + f"{'it' if len(over_budget_targets) == 1 else 'them'}.", + ) + if over_budget_targets + else () + ) + sentences: Final = attempted_failure + not_authorized_sentence + over_budget_sentence + if sentences: + return f"{lead} {' '.join(sentences)}" + targets: Final = ", ".join(str(mask_sensitive_structure(target)) for target in fallback_model_group) - if not fallback_failure_detail: - return f"{lead} Fallback model group(s) configured: {targets}." - return f"{lead} Fallback to {targets} also failed: {fallback_failure_detail}" + return f"{lead} Fallback model group(s) configured: {targets}." def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e1de57c91a0..4a2ec6ed1f8 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,7 +1,7 @@ import hashlib import json from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from enum import Enum from types import MappingProxyType @@ -29,7 +29,12 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, ) -from litellm.types.router import LiteLLMParamsTypedDict, reject_server_owned_wif_params +from litellm.types.router import ( + FallbackSkipReason, + LiteLLMParamsTypedDict, + SkippedFallbackTarget, + reject_server_owned_wif_params, +) if TYPE_CHECKING: from litellm.router import Router as _Router @@ -202,6 +207,7 @@ class AttemptedFallbackTargets: """ keys: frozenset[str] = frozenset() + _skip_reasons: dict[str, FallbackSkipReason] = field(default_factory=dict) # mutable-ok: O(1) per caller-sized list def __contains__(self, key: str) -> bool: return key in self.keys @@ -209,6 +215,20 @@ class AttemptedFallbackTargets: def record(self, key: str) -> None: self.keys = self.keys | frozenset((key,)) + @property + def skipped(self) -> tuple[SkippedFallbackTarget, ...]: + return tuple( + SkippedFallbackTarget(target=target, reason=reason) for target, reason in self._skip_reasons.items() + ) + + def record_skip(self, target: str, reason: FallbackSkipReason) -> None: + self._skip_reasons.setdefault(target, reason) + + +def attempted_fallback_targets_from(kwargs: Mapping[str, object]) -> AttemptedFallbackTargets: + carried_targets: Final[object] = kwargs.get("attempted_targets") + return carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() + def has_unattempted_fallback_target( fallback_model_group: Sequence[object] | None, kwargs: Mapping[str, object] @@ -222,6 +242,16 @@ def has_unattempted_fallback_target( return any((key := fallback_attempt_key(target)) is None or key not in attempted for target in fallback_model_group) +def is_bad_request_without_fallback(exception: Exception, fallback_on_unsupported_params: bool) -> bool: + if fallback_on_unsupported_params or not isinstance(exception, litellm.BadRequestError): + return False + if isinstance(exception, litellm.UnsupportedParamsError): + return True + if isinstance(exception, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError)): + return False + return exception.status_code == 400 and bool(getattr(exception, "llm_provider", "")) + + def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: """ Handles wildcard routing scenario @@ -514,7 +544,7 @@ def fallback_lookup_groups(kwargs: Mapping[str, object], model_group: str | None def _resolved_a_specific_chain( - fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract ) -> bool: resolved, generic_idx = result @@ -524,7 +554,7 @@ def _resolved_a_specific_chain( def get_fallback_model_group_for_lookup_groups( - fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape lookup_groups: tuple[str, ...], ) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract """ @@ -538,7 +568,10 @@ def get_fallback_model_group_for_lookup_groups( return next((result for result in results if result[0] is not None), (None, None)) -def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]: +def get_fallback_model_group( + fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape + model_group: str, +) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract """ Returns: - fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"] @@ -607,6 +640,8 @@ async def _is_fallback_target_authorized( fallback_entry: str | Mapping[str, object], original_model_group: str, kwargs: Mapping[str, object], + *, + attempted: AttemptedFallbackTargets, ) -> bool: access_check: Final = litellm_router.fallback_access_check target: Final = _get_fallback_target_model_group(fallback_entry) @@ -623,6 +658,7 @@ async def _is_fallback_target_authorized( "Skipping fallback to model_group = %s: caller is not authorized to call it", mask_sensitive_structure(fallback_entry), ) + attempted.record_skip(target, FallbackSkipReason.NOT_AUTHORIZED) return False @@ -631,6 +667,8 @@ async def _is_fallback_target_within_budget( fallback_entry: str | Mapping[str, object], original_model_group: str, kwargs: Mapping[str, object], + *, + attempted: AttemptedFallbackTargets, ) -> bool: budget_check: Final = litellm_router.fallback_budget_check target: Final = _get_fallback_target_model_group(fallback_entry) @@ -642,6 +680,7 @@ async def _is_fallback_target_within_budget( "Skipping fallback to model_group = %s: caller is over budget", mask_sensitive_structure(fallback_entry), ) + attempted.record_skip(target, FallbackSkipReason.OVER_BUDGET) return False @@ -730,13 +769,7 @@ async def run_async_fallback( same_model_group_only: Final = references_provider_scoped_resource(kwargs) or creates_provider_scoped_resource( kwargs ) - # Read out of kwargs and narrowed here rather than declared as a parameter: every caller - # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter - # would carry an annotation that no call site can actually be checked against. - carried_targets: Final = kwargs.get("attempted_targets") - attempted: Final = ( - carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() - ) + attempted: Final = attempted_fallback_targets_from(kwargs) # pyright: ignore[reportUnknownArgumentType] # keep legacy kwargs untyped failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group attempted.record(failed_model_group) @@ -758,9 +791,13 @@ async def run_async_fallback( original_model_group, ) continue - if not await _is_fallback_target_authorized(litellm_router, mg, original_model_group, kwargs): + if not await _is_fallback_target_authorized( + litellm_router, mg, original_model_group, kwargs, attempted=attempted + ): continue - if not await _is_fallback_target_within_budget(litellm_router, mg, original_model_group, kwargs): + if not await _is_fallback_target_within_budget( + litellm_router, mg, original_model_group, kwargs, attempted=attempted + ): continue attempt_key = fallback_attempt_key(mg) if attempt_key is not None: diff --git a/litellm/types/access_group.py b/litellm/types/access_group.py index b4bf345ecbe..2fa8cb8e592 100644 --- a/litellm/types/access_group.py +++ b/litellm/types/access_group.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from datetime import datetime from litellm.types.llms.base import LiteLLMBaseModel @@ -8,7 +9,9 @@ class AccessGroupCreateRequest(LiteLLMBaseModel): description: str | None = None access_model_names: list[str] | None = None access_mcp_server_ids: list[str] | None = None + mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] | None = None + access_skill_names: tuple[str, ...] | None = None assigned_team_ids: list[str] | None = None assigned_key_ids: list[str] | None = None @@ -18,7 +21,9 @@ class AccessGroupUpdateRequest(LiteLLMBaseModel): description: str | None = None access_model_names: list[str] | None = None access_mcp_server_ids: list[str] | None = None + mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] | None = None + access_skill_names: tuple[str, ...] | None = None assigned_team_ids: list[str] | None = None assigned_key_ids: list[str] | None = None @@ -36,7 +41,9 @@ class AccessGroupResponse(LiteLLMBaseModel): description: str | None = None access_model_names: list[str] access_mcp_server_ids: list[str] + mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] + access_skill_names: tuple[str, ...] = () assigned_team_ids: list[str] assigned_key_ids: list[str] access_mcp_servers: tuple[AccessGroupResource, ...] diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6c86b058f0f..f20a3fa3d7c 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -156,6 +156,7 @@ class SupportedGuardrailIntegrations(Enum): AGENT_365 = "agent_365" LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" + AIRIA = "airia" WINGBACK = "wingback" DECISION_MODEL = "decision_model" diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 25f55ddb84b..712f41ecc61 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -430,6 +430,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): cache_control: dict[str, Any] | None # Automatic prompt caching reasoning_effort: str | None safeguards: ReadOnly[list[dict[str, object]] | None] + service_tier: ReadOnly[str | None] class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 26f65fee4fe..d25db0d50a8 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -947,6 +947,7 @@ class ChatCompletionDeveloperMessage(OpenAIChatCompletionDeveloperMessage, total class GenericChatCompletionMessage(TypedDict, total=False): role: Required[str] content: Required[str | list] + cache_control: ReadOnly[ChatCompletionCachedContent] ValidUserMessageContentTypes = [ @@ -1931,7 +1932,10 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject): class ErrorEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.ERROR] sequence_number: int - error: ErrorEventError + error: ErrorEventError | None = None + code: str | int | None = None + message: str | None = None + param: str | Mapping[str, object] | None = None class GenericEvent(BaseLiteLLMOpenAIResponseObject): diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index c30fb7caf2e..e30ad26bfe5 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -7,6 +7,7 @@ from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry class UiDiscoveryEndpoints(LiteLLMBaseModel): server_root_path: str proxy_base_url: str | None + gateway_base_url: str | None = None auto_redirect_to_sso: bool admin_ui_disabled: bool sso_configured: bool diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/airia.py b/litellm/types/proxy/guardrails/guardrail_hooks/airia.py new file mode 100644 index 00000000000..c5beac110f9 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/airia.py @@ -0,0 +1,24 @@ +from pydantic import Field + +from .base import GuardrailConfigModel + + +class AiriaGuardrailConfigModel(GuardrailConfigModel): # pyright: ignore[reportMissingTypeArgument] # parameterizing adds an optional_params field to the UI form + api_base: str | None = Field( + default=None, + description="Base URL of the Airia AI Gateway, e.g. https://gateway.airia.ai. If not " + "provided, the `AIRIA_GATEWAY_URL` environment variable is checked.", + ) + api_key: str | None = Field( + default=None, + description="An Airia API key. If not provided, the `AIRIA_API_KEY` environment variable is checked.", + ) + timeout: float | None = Field( + default=None, + description="Request timeout in seconds. If not provided, the `AIRIA_TIMEOUT` environment " + "variable is checked, defaulting to 10.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Airia Guardrail" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 065205dcad8..743b2d7c963 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -92,6 +92,7 @@ class TeamListItem(LiteLLM_TeamTable): access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None + access_group_skill_names: tuple[str, ...] | None = None class TeamListResponse(LiteLLMBaseModel): diff --git a/litellm/types/router.py b/litellm/types/router.py index 273edf0f913..de993d46e44 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -56,6 +56,17 @@ from .utils import ( AllowedModelRegion: TypeAlias = Literal["eu", "us"] +class FallbackSkipReason(str, enum.Enum): + NOT_AUTHORIZED = "not_authorized" + OVER_BUDGET = "over_budget" + + +@dataclass(frozen=True, slots=True) +class SkippedFallbackTarget: + target: str + reason: FallbackSkipReason + + class ConfigurableClientsideParamsCustomAuth(TypedDict): api_base: str @@ -188,6 +199,7 @@ class UpdateRouterConfig(LiteLLMBaseModel): weights: RouterWeights | None = None tag_routing_prefix: str | None = None optional_pre_call_checks: OptionalPreCallChecks | None = None + fallback_on_unsupported_params: bool | None = None model_config = ConfigDict(protected_namespaces=()) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9011e4511d7..04f391d4d20 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -72571,10 +72571,12 @@ "gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, + "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -72585,6 +72587,7 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", @@ -72617,10 +72620,12 @@ "vertex_ai/gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, "cache_read_input_token_cost_batches": 7.5e-08, + "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_image": 0.00168, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 131072, @@ -72631,6 +72636,7 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", diff --git a/ruff-strict.toml b/ruff-strict.toml index 9e0b11a1af4..28102da3ccd 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -30,6 +30,7 @@ external = [ # grows over time; typing it concretely (`object`) broke that forwarding call outright — # basedpyright turned every named param into a reportArgumentType error. Any is correct here. "litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"] +"litellm/proxy/guardrails/guardrail_hooks/airia/airia.py" = ["ANN401"] # Same reason: `**kwargs` forwards verbatim to CustomGuardrail.__init__, and the lifecycle # hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here # would break the override rather than describe it. diff --git a/schema.prisma b/schema.prisma index ca58247a2c9..12438614755 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable { // Resource memberships - explicit arrays per type access_model_names String[] @default([]) access_mcp_server_ids String[] @default([]) + mcp_tool_permissions Json? access_agent_ids String[] @default([]) + access_skill_names String[] @default([]) assigned_team_ids String[] @default([]) assigned_key_ids String[] @default([]) diff --git a/scripts/benchmark_ocr_callbacks.py b/scripts/benchmark_ocr_callbacks.py index 5db182c3d0e..85cd4ffae49 100644 --- a/scripts/benchmark_ocr_callbacks.py +++ b/scripts/benchmark_ocr_callbacks.py @@ -23,7 +23,7 @@ from collections.abc import Sequence from dataclasses import asdict, dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from typing import Final, cast +from typing import Final SIZES: Final = ( 1024, @@ -50,15 +50,12 @@ class Server(ThreadingHTTPServer): def __init__(self) -> None: super().__init__(("127.0.0.1", 0), Handler) - self.user_agents: set[str] = set() class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def do_POST(self) -> None: - server: Final = cast(Server, self.server) - server.user_agents.add(self.headers.get("User-Agent", "")) length: Final = int(self.headers["Content-Length"]) body: Final = self.rfile.read(length) request: Final = json.loads(body) @@ -105,10 +102,16 @@ def percentile(values: Sequence[float], quantile: float) -> float: return ordered[index] -def verify(response: object) -> None: +def verify(response: object, expected_transport: str) -> None: + from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + pages: Final = getattr(response, "pages", ()) if len(pages) != 1 or getattr(pages[0], "markdown", None) != EXPECTED_MARKDOWN: raise RuntimeError(f"unexpected OCR response: {response!r}") + additional_headers: Final = get_hidden_params_dict(response).get("additional_headers") + rust_served: Final = isinstance(additional_headers, dict) and additional_headers.get("x-litellm-rust") == "true" + if rust_served != (expected_transport == "rust"): + raise RuntimeError(f"unexpected transport for {expected_transport}: rust_marker={rust_served}") def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Result: @@ -125,7 +128,9 @@ def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Res ) -def sync_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]: +def sync_samples( + litellm: object, url: str, request_document: dict[str, str], count: int, expected_transport: str +) -> tuple[float, ...]: samples: list[float] = [] for _ in range(count): started: Final = time.perf_counter() @@ -133,11 +138,13 @@ def sync_samples(litellm: object, url: str, request_document: dict[str, str], co model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30 ) samples.append(time.perf_counter() - started) - verify(response) + verify(response, expected_transport) return tuple(samples) -async def async_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]: +async def async_samples( + litellm: object, url: str, request_document: dict[str, str], count: int, expected_transport: str +) -> tuple[float, ...]: samples: list[float] = [] for _ in range(count): started: Final = time.perf_counter() @@ -145,7 +152,7 @@ async def async_samples(litellm: object, url: str, request_document: dict[str, s model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30 ) samples.append(time.perf_counter() - started) - verify(response) + verify(response, expected_transport) return tuple(samples) @@ -218,17 +225,20 @@ async def main() -> int: try: for size in args.sizes: request_document: Final = document(size) - sync_samples(litellm, url, request_document, args.warmup) + sync_samples(litellm, url, request_document, args.warmup, args.expected_transport) sync_result: Final = summarize( - args.label, "sync", size, sync_samples(litellm, url, request_document, args.iterations) + args.label, + "sync", + size, + sync_samples(litellm, url, request_document, args.iterations, args.expected_transport), ) results.append(sync_result) - await async_samples(litellm, url, request_document, args.warmup) + await async_samples(litellm, url, request_document, args.warmup, args.expected_transport) async_result: Final = summarize( args.label, "async", size, - await async_samples(litellm, url, request_document, args.iterations), + await async_samples(litellm, url, request_document, args.iterations, args.expected_transport), ) results.append(async_result) sys.stdout.write(json.dumps(asdict(sync_result)) + "\n") @@ -256,12 +266,6 @@ async def main() -> int: raise RuntimeError("callback registrations appeared in the no-callback case") await GLOBAL_LOGGING_WORKER.stop() - user_agents: Final = tuple(sorted(server.user_agents)) - python_transport: Final = any( - value.startswith("python-httpx") or value.startswith("litellm/") for value in user_agents - ) - if (args.expected_transport == "python") != python_transport: - raise RuntimeError(f"unexpected transport for {args.label}: user_agents={user_agents}") artifact: Final = { "label": args.label, "callbacks": args.callbacks, @@ -275,11 +279,10 @@ async def main() -> int: "package": str(package), "native": native_path, "native_sha256": native_sha256, - "user_agents": user_agents, "results": tuple(asdict(result) for result in results), } args.output.write_text(json.dumps(artifact, indent=2) + "\n") - sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package", "user_agents")}) + "\n") + sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package")}) + "\n") sys.stdout.write(f"results={args.output}\n") sys.stdout.flush() return 0 diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index b948db2b231..d489a023bf2 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **team**: Optional `team_member_model_max_budget` argument on `litellm_team`, a JSON string of per-model budgets applied to every team member and sent as `team_member_model_max_budget` on `/team/new` and `/team/update`; read back from `team_member_budget_table.model_max_budget` on `/team/info` so out-of-band changes show as drift, with `budget_limit`/`time_period` treated as equal to the `max_budget`/`budget_duration` spelling the proxy stores; removing it sends `{}`, which clears every per-model member budget; `tpm_limit`/`rpm_limit` are rejected at plan time because the proxy rejects them on team budgets. Needs a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887): v1.105.0 and v1.106.0-rc.1 accept the field but silently drop it, so every plan shows a diff - **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are persisted through `/model/{id}/update` since `/model/update` ignores `model_info`; also exported by the `litellm_model` and `litellm_models` data sources - **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 19575907269..66511be8317 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -98,6 +98,24 @@ resource "litellm_team" "model_dependent_team" { } ``` +### Team with Per-Member Model Budgets + +```hcl +resource "litellm_team" "engineering" { + team_alias = "engineering" + models = ["claude-opus", "claude-sonnet"] + + # Each member may spend $500 a day in total + team_member_budget = 500 + team_member_budget_duration = "1d" + + # ...of which at most $400 on Opus. Past that, Sonnet keeps working. + team_member_model_max_budget = jsonencode({ + "claude-opus" = { budget_limit = 400, time_period = "1d" } + }) +} +``` + ## Argument Reference The following arguments are supported: @@ -150,6 +168,8 @@ The following arguments are supported: * `team_member_key_duration` - (Optional) Lifetime for keys created by team members (e.g. `1d`, `1w`). +* `team_member_model_max_budget` - (Optional) JSON string of per-model budgets applied to each individual team member, e.g. `jsonencode({"claude-opus" = {budget_limit = 400, time_period = "1d"}})`. Each model maps to an object with `budget_limit` (or `max_budget`) and `time_period` (or `budget_duration`). The proxy rejects `tpm_limit` and `rpm_limit` on team budgets, so the provider rejects them at plan time. A member who reaches a model's cap is blocked on that model only and can keep calling the team's other models until `team_member_budget` runs out. The key is matched against the model name on the request, so for a router alias the cap applies to the alias, not the model the router picks. Remove the attribute to clear every per-model member budget. Requires a LiteLLM Enterprise license on the proxy, and a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887). Older proxies, including v1.105.0 and v1.106.0-rc.1, accept the field but silently drop it, so every plan shows a diff. + * `model_rpm_limit` - (Optional) Map of model name to requests per minute limit for that model. * `model_tpm_limit` - (Optional) Map of model name to tokens per minute limit for that model. diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 622fec7d0a0..3f256da36d4 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -6,6 +6,7 @@ import ( "io" "log" "net/http" + "reflect" "github.com/google/uuid" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" @@ -135,6 +136,13 @@ func ResourceLiteLLMTeam() *schema.Resource { Type: schema.TypeString, Optional: true, }, + "team_member_model_max_budget": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validateTeamMemberModelMaxBudget, + DiffSuppressFunc: teamMemberModelMaxBudgetEquivalent, + Description: "JSON string of per-model budgets applied to every team member (e.g. '{\"claude-opus\": {\"budget_limit\": 400, \"time_period\": \"1d\"}}'). Each entry takes budget_limit (or max_budget) and time_period (or budget_duration). Caps match the model name on the request, so behind a router alias the cap must be keyed on the alias, not the model the router picks. Requires a LiteLLM Enterprise license and a proxy that includes https://github.com/BerriAI/litellm/pull/43887; older proxies drop the value, so every plan shows a diff", + }, "model_rpm_limit": { Type: schema.TypeMap, Optional: true, @@ -280,6 +288,11 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { if teamResp.AllowedPassthroughRoutes != nil { d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes) } + var memberModelMaxBudget map[string]interface{} + if teamResp.TeamMemberBudgetTable != nil { + memberModelMaxBudget = teamResp.TeamMemberBudgetTable.ModelMaxBudget + } + d.Set("team_member_model_max_budget", keyModelMaxBudgetJSON(memberModelMaxBudget)) // Explicitly fetch the current permissions from the API permResp, err := getTeamPermissions(client, d.Id()) @@ -362,6 +375,75 @@ func resourceLiteLLMTeamDelete(d *schema.ResourceData, m interface{}) error { return nil } +// The proxy rejects tpm_limit/rpm_limit on team budgets, so only spend fields are allowed here. +var teamMemberModelBudgetFields = map[string]bool{ + "budget_limit": true, + "max_budget": true, + "time_period": true, + "budget_duration": true, +} + +func validateTeamMemberModelMaxBudget(v interface{}, k string) ([]string, []error) { + var parsed map[string]json.RawMessage + if err := json.Unmarshal([]byte(v.(string)), &parsed); err != nil || parsed == nil { + return nil, []error{fmt.Errorf("%q must be a JSON object keyed by model name, got %s", k, v)} + } + for model, cfg := range parsed { + var budget map[string]json.RawMessage + if err := json.Unmarshal(cfg, &budget); err != nil || len(budget) == 0 { + return nil, []error{fmt.Errorf("%q[%q] must be a budget object such as {\"budget_limit\": 400, \"time_period\": \"1d\"}, got %s", k, model, cfg)} + } + for field := range budget { + if !teamMemberModelBudgetFields[field] { + return nil, []error{fmt.Errorf("%q[%q] has unsupported field %q; team member budgets accept budget_limit, max_budget, time_period and budget_duration only", k, model, field)} + } + } + } + return nil, nil +} + +// The proxy stores budget_limit/time_period as max_budget/budget_duration, so both spellings must compare equal. +var modelBudgetFieldAliases = map[string]string{ + "budget_limit": "max_budget", + "time_period": "budget_duration", +} + +func normalizeModelMaxBudget(raw string) (map[string]interface{}, bool) { + if raw == "" { + return map[string]interface{}{}, true + } + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(raw), &parsed); err != nil { + return nil, false + } + normalized := make(map[string]interface{}, len(parsed)) + for model, cfg := range parsed { + budget, ok := cfg.(map[string]interface{}) + if !ok { + normalized[model] = cfg + continue + } + entry := make(map[string]interface{}, len(budget)) + for field, value := range budget { + if value == nil { + continue + } + if canonical, isAlias := modelBudgetFieldAliases[field]; isAlias { + field = canonical + } + entry[field] = value + } + normalized[model] = entry + } + return normalized, true +} + +func teamMemberModelMaxBudgetEquivalent(k, oldValue, newValue string, d *schema.ResourceData) bool { + oldParsed, okOld := normalizeModelMaxBudget(oldValue) + newParsed, okNew := normalizeModelMaxBudget(newValue) + return okOld && okNew && reflect.DeepEqual(oldParsed, newParsed) +} + func resolveTeamID(d *schema.ResourceData) string { if v, ok := d.GetOk("team_id"); ok { return v.(string) @@ -402,6 +484,14 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} teamData["tags"] = v } + // The proxy treats a missing field as "leave unchanged", so an empty map is sent only to clear a value that was set. + memberModelMaxBudget := parseKeyModelMaxBudget(d.Get("team_member_model_max_budget").(string)) + if len(memberModelMaxBudget) > 0 { + teamData["team_member_model_max_budget"] = memberModelMaxBudget + } else if old, _ := d.GetChange("team_member_model_max_budget"); len(parseKeyModelMaxBudget(old.(string))) > 0 { + teamData["team_member_model_max_budget"] = memberModelMaxBudget + } + if metadata := buildTeamMetadata(d); metadata != nil { teamData["metadata"] = metadata } diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go index 2aa8043ad89..d9d09c13dda 100644 --- a/terraform/provider/litellm/resource_team_test.go +++ b/terraform/provider/litellm/resource_team_test.go @@ -472,3 +472,238 @@ func TestHandleResponseAcceptsFullSuccessRange(t *testing.T) { }) } } + +func TestTeamCreateSendsTeamMemberModelMaxBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`, + }) + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + want := map[string]interface{}{"claude-opus": map[string]interface{}{"budget_limit": 400.0, "time_period": "1d"}} + if got := captured["team_member_model_max_budget"]; !reflect.DeepEqual(got, want) { + t.Fatalf("payload team_member_model_max_budget = %v, want %v", got, want) + } +} + +func TestTeamReadMapsTeamMemberModelMaxBudgetFromBudgetTable(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "eng", + "team_member_budget_table": { + "max_budget": 500.0, + "model_max_budget": {"claude-opus": {"max_budget": 400.0, "budget_duration": "1d"}} + } + } + }`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"}) + d.SetId("team-1") + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read returned error: %v", err) + } + + var got map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("team_member_model_max_budget").(string)), &got); err != nil { + t.Fatalf("team_member_model_max_budget in state is not JSON: %v", err) + } + want := map[string]interface{}{"claude-opus": map[string]interface{}{"max_budget": 400.0, "budget_duration": "1d"}} + if !reflect.DeepEqual(got, want) { + t.Errorf("team_member_model_max_budget = %v, want %v", got, want) + } +} + +func TestTeamReadClearsTeamMemberModelMaxBudgetWhenAbsent(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng","team_member_budget_table":{"model_max_budget":null}}}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`, + }) + d.SetId("team-1") + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read returned error: %v", err) + } + if got := d.Get("team_member_model_max_budget").(string); got != "" { + t.Errorf("team_member_model_max_budget = %q, want empty so drift shows up", got) + } +} + +func TestTeamUpdateClearsRemovedTeamMemberModelMaxBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_alias": "eng", + "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`, + }) + priorData.SetId("team-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{"team_alias": "eng"}) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(prior, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + if got, ok := captured["team_member_model_max_budget"]; !ok || !reflect.DeepEqual(got, map[string]interface{}{}) { + t.Errorf("payload team_member_model_max_budget = %v (present=%v), want explicit empty object", got, ok) + } +} + +func TestTeamMemberModelMaxBudgetAliasesAreEquivalent(t *testing.T) { + configured := `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}` + stored := `{"claude-opus":{"max_budget":400,"budget_duration":"1d"}}` + if !teamMemberModelMaxBudgetEquivalent("", stored, configured, nil) { + t.Fatalf("budget_limit/time_period should match the stored max_budget/budget_duration spelling") + } + if teamMemberModelMaxBudgetEquivalent("", stored, `{"claude-opus": {"budget_limit": 300, "time_period": "1d"}}`, nil) { + t.Fatalf("a changed limit must not be suppressed") + } +} + +func TestTeamMemberModelMaxBudgetRejectsUnknownFields(t *testing.T) { + schema := ResourceLiteLLMTeam().Schema["team_member_model_max_budget"] + for _, bad := range []string{ + `{"claude-opus": {"limit": 400}}`, + `{"claude-opus": {"budget_limit": 400, "time_period": "1d", "tpm_limit": 1000}}`, + `{"claude-opus": {"budget_limit": 400, "time_period": "1d", "rpm_limit": 10}}`, + } { + if _, errs := schema.ValidateFunc(bad, "team_member_model_max_budget"); len(errs) == 0 { + t.Errorf("expected %s to be rejected", bad) + } + } + if _, errs := schema.ValidateFunc(`{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`, "team_member_model_max_budget"); len(errs) != 0 { + t.Fatalf("valid config rejected: %v", errs) + } +} + +func TestTeamCreateOmitsTeamMemberModelMaxBudgetWhenUnset(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"}) + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + if got, ok := captured["team_member_model_max_budget"]; ok { + t.Fatalf("payload team_member_model_max_budget = %v, want key absent", got) + } +} + +func TestTeamCreateOmitsEmptyTeamMemberModelMaxBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "team_member_model_max_budget": `{}`, + }) + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + if got, ok := captured["team_member_model_max_budget"]; ok { + t.Fatalf("payload team_member_model_max_budget = %v, want key absent so no empty member budget row is created", got) + } +} + +func teamUpdateData(t *testing.T, prior, config map[string]interface{}) *schema.ResourceData { + t.Helper() + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, prior) + priorData.SetId("team-1") + state := priorData.State() + diff, err := res.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(state, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + return d +} + +func TestTeamUpdateSendsChangedTeamMemberModelMaxBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + d := teamUpdateData(t, + map[string]interface{}{"team_alias": "eng", "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`}, + map[string]interface{}{"team_alias": "eng", "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 300, "time_period": "1d"}}`}, + ) + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + want := map[string]interface{}{"claude-opus": map[string]interface{}{"budget_limit": 300.0, "time_period": "1d"}} + if got := captured["team_member_model_max_budget"]; !reflect.DeepEqual(got, want) { + t.Fatalf("payload team_member_model_max_budget = %v, want %v", got, want) + } +} + +func TestTeamUpdateSendsMemberBudgetWithPerModelBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + d := teamUpdateData(t, + map[string]interface{}{"team_alias": "eng"}, + map[string]interface{}{ + "team_alias": "eng", + "team_member_budget": 500.0, + "team_member_budget_duration": "1d", + "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`, + }, + ) + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + if got := captured["team_member_budget"]; got != 500.0 { + t.Errorf("payload team_member_budget = %v, want 500", got) + } + if got := captured["team_member_budget_duration"]; got != "1d" { + t.Errorf("payload team_member_budget_duration = %v, want 1d", got) + } + want := map[string]interface{}{"claude-opus": map[string]interface{}{"budget_limit": 400.0, "time_period": "1d"}} + if got := captured["team_member_model_max_budget"]; !reflect.DeepEqual(got, want) { + t.Errorf("payload team_member_model_max_budget = %v, want %v", got, want) + } +} + +func TestTeamReadWithoutMemberBudgetTableLeavesFieldEmpty(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`) + defer srv.Close() + + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"}) + d.SetId("team-1") + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read returned error: %v", err) + } + if got := d.Get("team_member_model_max_budget").(string); got != "" { + t.Errorf("team_member_model_max_budget = %q, want empty", got) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 8baf88edee1..03cb145dba2 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -73,6 +73,12 @@ type TeamResponse struct { ModelRPMLimit map[string]interface{} `json:"model_rpm_limit,omitempty"` ModelTPMLimit map[string]interface{} `json:"model_tpm_limit,omitempty"` AllowedPassthroughRoutes []string `json:"allowed_passthrough_routes,omitempty"` + TeamMemberBudgetTable *TeamMemberBudgetTable `json:"team_member_budget_table,omitempty"` +} + +// TeamMemberBudgetTable is the team's default member budget row, the only place /team/info returns team_member_model_max_budget. +type TeamMemberBudgetTable struct { + ModelMaxBudget map[string]interface{} `json:"model_max_budget,omitempty"` } // OrganizationResponse represents a response from the API containing organization information. diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index 9e25d78b6cd..cdd194b267f 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -39,6 +39,14 @@ assertions: [denied_without_permission] source: "rest_endpoints.py:305-386" rationale: Tool-level permission guard; multi-tenant safety +- id: mcp.tool_search.api_key.oversized_description_omitted + module: mcp + tier: P0 + operation: tool_search + auth_family: api_key + assertions: [oversized_description_omitted] + source: "test_mcp_tool_search_e2e.py" + rationale: Native semantic search omits tools whose descriptions exceed the embedding input limit - id: mcp.list_tools.bearer.succeeds module: mcp tier: P1 diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index 63194af221e..b712f5ec45d 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Protocol from e2e_config import datadog_mcp_url, unique_marker @@ -44,6 +44,7 @@ def register_datadog_mcp( *, mcp_access_groups: list[str] | None = None, allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), + tool_name_to_description: Mapping[str, str] | None = None, ) -> str: """Register the core Datadog toolset with its credentials from the env. By default the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose @@ -61,6 +62,7 @@ def register_datadog_mcp( }, allowed_tools=None if allowed_tools is None else list(allowed_tools), mcp_access_groups=mcp_access_groups, + tool_name_to_description=tool_name_to_description, ) resources.defer(lambda: client.delete_server(server_id)) return server_id diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index e373a8e31ae..26a1264a0f1 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -42,6 +42,7 @@ class McpServerNewBody(BaseModel): static_headers: dict[str, str] | None = None allowed_tools: list[str] | None = None mcp_access_groups: list[str] | None = None + tool_name_to_description: dict[str, str] | None = None class McpServerNewResponse(BaseModel): @@ -131,6 +132,34 @@ class McpCallToolBody(BaseModel): server_id: str +class McpVirtualToolCallBody(BaseModel): + name: str + arguments: dict[str, McpToolArg] + + +class McpToolsListParams(BaseModel): + server_id: str | None = None + + +class McpToolSearchSettings(BaseModel): + model_config = ConfigDict(frozen=True) + + embedding_model: str | None = None + top_k: int = 5 + similarity_threshold: float = 0.0 + core_tools: tuple[str, ...] = () + + +class McpToolSearchSettingsResponse(BaseModel): + values: McpToolSearchSettings + + +class McpToolSearchSettingsUpdateResponse(BaseModel): + message: str + status: str + settings: McpToolSearchSettings + + class McpCallContent(BaseModel): type: str | None = None text: str | None = None @@ -166,6 +195,7 @@ class McpClient: static_headers: dict[str, str] | None = None, allowed_tools: list[str] | None = None, mcp_access_groups: list[str] | None = None, + tool_name_to_description: Mapping[str, str] | None = None, ) -> str: return unwrap( self.proxy.transport.post( @@ -180,6 +210,9 @@ class McpClient: static_headers=static_headers, allowed_tools=allowed_tools, mcp_access_groups=mcp_access_groups, + tool_name_to_description=( + dict(tool_name_to_description) if tool_name_to_description is not None else None + ), ), response_type=McpServerNewResponse, ) @@ -244,14 +277,21 @@ class McpClient: mcp_access_groups: list[str] | None = None, mcp_toolsets: list[str] | None = None, models: list[str] | None = None, + mcp_tool_search_enabled: bool | None = None, ) -> str: object_permission = ( ObjectPermission( mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups, mcp_toolsets=mcp_toolsets, + mcp_tool_search_enabled=mcp_tool_search_enabled, + ) + if ( + mcp_servers is not None + or mcp_access_groups is not None + or mcp_toolsets is not None + or mcp_tool_search_enabled is not None ) - if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None else None ) return self.proxy.generate_key( @@ -263,14 +303,37 @@ class McpClient: ) @step("List the MCP tools the key can see from /mcp-rest/tools/list") - def list_tools(self, key: str) -> Result[McpToolsListResponse]: + def list_tools(self, key: str, server_id: str | None = None) -> Result[McpToolsListResponse]: return self.proxy.transport.get( "/mcp-rest/tools/list", headers=ApiKeyHeaders(x_litellm_api_key=key), - params=NoBody(), + params=McpToolsListParams(server_id=server_id), response_type=McpToolsListResponse, ) + @step("Read the MCP tool search settings from /get/mcp_tool_search_settings") + def get_mcp_tool_search_settings(self) -> McpToolSearchSettings: + return unwrap( + self.proxy.transport.get( + "/get/mcp_tool_search_settings", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=McpToolSearchSettingsResponse, + ) + ).values + + @step("Update the MCP tool search settings with /update/mcp_tool_search_settings") + def update_mcp_tool_search_settings(self, settings: McpToolSearchSettings) -> None: + _ = unwrap( + self.proxy.transport.patch( + "/update/mcp_tool_search_settings", + headers=self.proxy.transport.master, + json=settings, + response_type=McpToolSearchSettingsUpdateResponse, + ) + ) + settle_propagation(time.monotonic()) + @step('Wait for /mcp-rest/tools/list to show the MCP server\'s tool matching "{needle}"') def await_tool(self, key: str, server_id: str, needle: str) -> str: """Poll tools/list until `server_id` serves a tool matching `needle`, and @@ -419,6 +482,21 @@ class McpClient: response_type=McpCallToolResponse, ) + @step("Call the MCP virtual tool {name} through /mcp-rest/tools/call") + def call_virtual_tool( + self, + key: str, + *, + name: str, + arguments: McpToolArguments, + ) -> Result[McpCallToolResponse]: + return self.proxy.transport.post( + "/mcp-rest/tools/call", + headers=ApiKeyHeaders(x_litellm_api_key=key), + json=McpVirtualToolCallBody(name=name, arguments=dict(arguments)), + response_type=McpCallToolResponse, + ) + def _is_mcp_not_synced( result: Result[McpCallToolResponse], diff --git a/tests/e2e/mcp/test_mcp_tool_search_e2e.py b/tests/e2e/mcp/test_mcp_tool_search_e2e.py new file mode 100644 index 00000000000..847239e8a66 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_tool_search_e2e.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import provider_edge_base, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient, McpToolSearchSettings +from models import LiteLLMParamsBody +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +MCP_TOOL_SEARCH_TOOL_NAME: Final = "mcp_tool_search" +OVERSIZED_DESCRIPTION: Final = f"OVERSIZED {'weather ' * 12_000}" + + +class TestMcpToolSearch: + @pytest.mark.covers("mcp.tool_search.api_key.oversized_description_omitted") + def test_native_tool_search_omits_oversized_description( + self, proxy: ProxyClient, client: McpClient, resources: ResourceManager + ) -> None: + """Exercise shared ranking through the native virtual tool because proxy mode does not apply description overrides.""" + model_name: Final = f"e2e-mcp-tool-search-{unique_marker()}" + embedding_api_base: Final = provider_edge_base("openai") + model_id: Final = proxy.create_model( + model_name, + LiteLLMParamsBody( + model="openai/text-embedding-3-large", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if embedding_api_base is None else f"{embedding_api_base}/v1", + ), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + + previous_settings: Final = client.get_mcp_tool_search_settings() + resources.defer(lambda: client.update_mcp_tool_search_settings(previous_settings)) + client.update_mcp_tool_search_settings( + McpToolSearchSettings( + embedding_model=model_name, + top_k=100, + similarity_threshold=0.0, + core_tools=previous_settings.core_tools, + ) + ) + + server_id: Final = register_datadog_mcp( + client, + resources, + allowed_tools=None, + tool_name_to_description={SEARCH_LOGS_TOOL: OVERSIZED_DESCRIPTION}, + ) + client.await_registered(server_id) + key: Final = client.generate_key( + user_id=f"e2e-mcp-tool-search-{unique_marker()}", + mcp_servers=[server_id], + models=[model_name], + mcp_tool_search_enabled=True, + ) + resources.defer(lambda: proxy.delete_key(key)) + + tools: Final = unwrap(client.list_tools(key, server_id=server_id)) + assert len(tools.tool_names_for_server(server_id)) >= 2 + oversized_tool_name: Final = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) + assert oversized_tool_name is not None + oversized_tool: Final = next((tool for tool in tools.tools if tool.name == oversized_tool_name), None) + assert oversized_tool is not None + assert oversized_tool.description == OVERSIZED_DESCRIPTION + normal_tool_name: Final = next( + (name for name in tools.tool_names_for_server(server_id) if name != oversized_tool_name), None + ) + assert normal_tool_name is not None + + result: Final = unwrap( + client.call_virtual_tool( + key, + name=MCP_TOOL_SEARCH_TOOL_NAME, + arguments={"query": normal_tool_name, "top_k": 100}, + ) + ) + + assert result.is_error is not True, result.all_text + assert normal_tool_name in result.all_text + assert oversized_tool_name not in result.all_text diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 5259166e9a9..7eb9252acfb 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -67,6 +67,7 @@ class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None mcp_toolsets: list[str] | None = None + mcp_tool_search_enabled: bool | None = None class KeyGenerateBody(BaseModel): diff --git a/tests/integration/_support/anthropic_refusal.py b/tests/integration/_support/anthropic_refusal.py new file mode 100644 index 00000000000..6d5fd236074 --- /dev/null +++ b/tests/integration/_support/anthropic_refusal.py @@ -0,0 +1,533 @@ +import hashlib +import json +import os +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final, TypeAlias + +import yaml +from integration._support.anthropic_thinking import ( + JSON_LIST, + JSON_OBJECT, + Event, + marker_of, + stream_reply, + streams, +) +from integration._support.client import eventually +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response" +DOCS_STREAMING_REFUSAL_URL: Final = ( + "https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals" +) + +# Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12. +DOCS_REFUSAL: Final[Event] = { + "id": "msg_01XFUDYJgAACzvnptvVoYEL", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", + }, + "usage": {"input_tokens": 412, "output_tokens": 0}, +} +DOCS_PARTIAL_REFUSAL: Final[Event] = { + "role": "assistant", + "content": [{"type": "text", "text": "Hello.."}], + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", + }, +} + +# Captured from api.anthropic.com (claude-sonnet-5-5, prompt "Reply with the single word: ready") on 2026-10-13. +CAPTURED_TURN: Final[Event] = { + "model": "claude-sonnet-5-5", + "id": "msg_011CfyXFGze1cDkk8o7gDRRS", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ready"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, +} +CAPTURED_EVENTS: Final[tuple[Event, ...]] = ( + { + "type": "message_start", + "message": { + "model": "claude-sonnet-5-5", + "id": "msg_011CfyXFLwXUxsp7nNAm3gcu", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 4, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ready"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 4, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, +) + +DOCS_STOP_DETAILS: Final = DOCS_REFUSAL["stop_details"] +REFUSAL_FIELDS: Final[Event] = {"stop_reason": "refusal", "stop_details": DOCS_STOP_DETAILS} +STOP_FIELDS: Final = ("stop_reason", "stop_details") +LONG_STOP_DETAILS: Final = "d" * 5120 +MALFORMED_STOP_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"int": 7, "list": ["refusal"], "empty": "", "long": LONG_STOP_DETAILS} +) +_FIXTURE: Final = re.compile(r"refusal-fixture-([a-z]+)") + +Location: TypeAlias = tuple[str | int, ...] + + +@dataclass(frozen=True, slots=True) +class Fixture: + status: int + body: bytes + events: tuple[Event, ...] = () + + +def _text_block(index: int, block: JsonValue) -> tuple[Event, ...]: + text: Final = JSON_OBJECT.validate_python(block)["text"] + return ( + {"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": index}, + ) + + +def refusal_events(message: Mapping[str, JsonValue]) -> tuple[Event, ...]: + usage: Final = JSON_OBJECT.validate_python(message["usage"]) + start: Final[Event] = { + "type": "message_start", + "message": { + "id": message["id"], + "type": "message", + "role": "assistant", + "model": message["model"], + "content": [], + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": usage, + }, + } + blocks: Final = chain.from_iterable( + _text_block(index, block) for index, block in enumerate(JSON_LIST.validate_python(message["content"])) + ) + delta: Final[Event] = { + "type": "message_delta", + "delta": { + "stop_reason": message["stop_reason"], + "stop_sequence": None, + "stop_details": message["stop_details"], + }, + "usage": {name: usage[name] for name in ("output_tokens", "output_tokens_details") if name in usage}, + } + return (start, *blocks, delta, {"type": "message_stop"}) + + +def _without_stop_details(value: Mapping[str, JsonValue]) -> Event: + return {name: item for name, item in value.items() if name != "stop_details"} + + +def _captured_event(event: Event, stop_details: JsonValue, *, keep: bool) -> Event: + if event["type"] == "message_start": + message: Final = JSON_OBJECT.validate_python(event["message"]) + return {**event, "message": message if keep else _without_stop_details(message)} + if event["type"] != "message_delta": + return event + delta: Final = JSON_OBJECT.validate_python(event["delta"]) + replaced: Final = {**delta, "stop_details": stop_details} if keep else _without_stop_details(delta) + return {**event, "delta": replaced} + + +def _captured(stop_details: JsonValue, *, keep: bool = True) -> Fixture: + body: Final = {**CAPTURED_TURN, "stop_details": stop_details} if keep else _without_stop_details(CAPTURED_TURN) + return Fixture( + 200, + json.dumps(body).encode(), + tuple(_captured_event(event, stop_details, keep=keep) for event in CAPTURED_EVENTS), + ) + + +def _duplicated_stop_details() -> bytes: + pair: Final = f'"stop_details": {json.dumps(DOCS_STOP_DETAILS)}' + return json.dumps(DOCS_REFUSAL).replace(pair, f"{pair}, {pair}").encode() + + +def _anthropic_error(kind: str, message: str) -> bytes: + return json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode() + + +_PARTIAL: Final[Event] = {**DOCS_REFUSAL, **DOCS_PARTIAL_REFUSAL} +_REASONED: Final[Event] = { + **_PARTIAL, + "usage": {"input_tokens": 412, "output_tokens": 10, "output_tokens_details": {"thinking_tokens": 4}}, +} + +JSON_TOOL_NAME: Final = "json_tool_call" +JSON_ANSWER_SCHEMA: Final[Event] = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + "additionalProperties": False, +} +JSON_ANSWER_FORMAT: Final[Event] = { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": JSON_ANSWER_SCHEMA, "strict": True}, +} +DOCS_PARTIAL_TEXT: Final = "Hello.." +_DECLARED_INPUT: Final[Event] = {"answer": DOCS_PARTIAL_TEXT} +_ENVELOPED_INPUT: Final[Event] = {"parameters": _DECLARED_INPUT} +DECLARED_PARTIAL_JSON: Final = '{"answer": "Hello..' +ENVELOPED_PARTIAL_JSON: Final = '{"parameters": {"answer": "Hello..' +_JSON_TOOL_USE_ID: Final = "toolu_01RefusalFixtureJsonTool" + + +def _json_tool_refusal(tool_input: Event, partial_json: str) -> Fixture: + message: Final[Event] = { + **DOCS_REFUSAL, + "content": [{"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": tool_input}], + } + start, delta, stop = refusal_events(DOCS_REFUSAL) + tool_block: Final[tuple[Event, ...]] = ( + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": {}}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": partial_json}, + }, + {"type": "content_block_stop", "index": 0}, + ) + return Fixture(200, json.dumps(message).encode(), (start, *tool_block, delta, stop)) + + +FIXTURES: Final[Mapping[str, Fixture]] = MappingProxyType( + { + "refusal": Fixture(200, json.dumps(DOCS_REFUSAL).encode(), refusal_events(DOCS_REFUSAL)), + "partial": Fixture(200, json.dumps(_PARTIAL).encode(), refusal_events(_PARTIAL)), + "reasoned": Fixture(200, json.dumps(_REASONED).encode(), refusal_events(_REASONED)), + "jsontool": _json_tool_refusal(_DECLARED_INPUT, DECLARED_PARTIAL_JSON), + "jsonheld": _json_tool_refusal(_ENVELOPED_INPUT, ENVELOPED_PARTIAL_JSON), + "duplicate": Fixture(200, _duplicated_stop_details()), + "null": _captured(None), + "missing": _captured(None, keep=False), + **{name: _captured(value) for name, value in MALFORMED_STOP_DETAILS.items()}, + "nondict": Fixture(200, b"[]"), + "unauthorized": Fixture(401, _anthropic_error("authentication_error", "invalid x-api-key")), + "invalid": Fixture(400, _anthropic_error("invalid_request_error", "max_tokens: Field required")), + } +) + + +def refusal_prompt(fixture: str, marker: str) -> str: + return f"refusal-fixture-{fixture} marker-{marker}" + + +CITED_FIELDS: Final[Event] = {"citations": ["https://example.com/source"]} +_PLAIN_DELTA: Final[Event] = {"role": "assistant", "content": "ready"} +_CITED_DELTA: Final[Event] = {**_PLAIN_DELTA, "provider_specific_fields": CITED_FIELDS} +_OPENAI_USAGE: Final[Event] = { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + "completion_tokens_details": {"reasoning_tokens": 4, "text_tokens": 3}, +} +UNVALIDATED_USAGE_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "text": {"reasoning_tokens": "abc"}, + "fraction": {"reasoning_tokens": 1.5}, + "object": {"reasoning_tokens": {"a": 1}}, + } +) + + +@dataclass(frozen=True, slots=True) +class OpenAIFixture: + first_delta: Event + usage: tuple[Event, ...] = () + + +_OPENAI_FIXTURES: Final[Mapping[str, OpenAIFixture]] = MappingProxyType( + { + "plain": OpenAIFixture(_PLAIN_DELTA), + "cited": OpenAIFixture(_CITED_DELTA), + "plainusage": OpenAIFixture(_PLAIN_DELTA, (_OPENAI_USAGE,)), + "citedusage": OpenAIFixture(_CITED_DELTA, (_OPENAI_USAGE,)), + } +) + + +def openai_id(marker: str) -> str: + return f"chatcmpl-{marker}" + + +def _openai_frame(marker: str, body: Event) -> bytes: + chunk: Final[Event] = { + "id": openai_id(marker), + "object": "chat.completion.chunk", + "created": 1760000000, + "model": "gpt-4o-mini", + **body, + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +def _openai_choice(delta: Event, finish_reason: str | None) -> Event: + return {"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]} + + +def _openai_stream(name: str, request: Request) -> Reply: + assert streams(request), request.body + marker: Final = marker_of(request) + fixture: Final = _OPENAI_FIXTURES[name] + return Reply( + content_type="text/event-stream", + chunks=( + _openai_frame(marker, _openai_choice(fixture.first_delta, None)), + _openai_frame(marker, _openai_choice({"content": " now"}, "stop")), + *(_openai_frame(marker, {"choices": [], "usage": usage}) for usage in fixture.usage), + b"data: [DONE]\n\n", + ), + ) + + +def respond(request: Request) -> Reply: + found: Final[Sequence[str]] = _FIXTURE.findall(request.body.decode()) + assert found, request.body + if request.target.endswith("/chat/completions"): + return _openai_stream(found[-1], request) + fixture: Final = FIXTURES[found[-1]] + if fixture.status != 200 or not fixture.events or not streams(request): + return Reply(status=fixture.status, body=fixture.body) + return stream_reply(request, fixture.events) + + +def _stop_fields(fields: JsonValue) -> Event: + if not isinstance(fields, dict): + return {} + return {name: fields[name] for name in STOP_FIELDS if name in fields} + + +def relayed(value: JsonValue, location: Location = ()) -> tuple[tuple[Location, Event], ...]: + if isinstance(value, list): + return tuple(chain.from_iterable(relayed(item, (*location, index)) for index, item in enumerate(value))) + if not isinstance(value, dict): + return () + found: Final = _stop_fields(value.get("provider_specific_fields")) + here: Final = (((*location, "provider_specific_fields"), found),) if found else () + return (*here, *chain.from_iterable(relayed(item, (*location, name)) for name, item in value.items())) + + +_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") +_STORED: Final = TypeAdapter(bytes | None) +STALE_ID_SUFFIX: Final = "-stale" + + +def _stored(store: Redis, key: str) -> bytes: + if not _CACHE_KEY.fullmatch(key) or store.type(key) != b"string": + return b"" + return _STORED.validate_python(store.get(key)) or b"" + + +def _entry_key(store: Redis, needle: str) -> str | None: + keys: Final = (raw.decode() for raw in store.scan_iter(count=1000)) + return next((key for key in keys if needle.encode() in _stored(store, key)), None) + + +def _with_message_fields(response: Event, fields: JsonValue) -> Event: + choice: Final = JSON_OBJECT.validate_python(JSON_LIST.validate_python(response["choices"])[0]) + message: Final = JSON_OBJECT.validate_python(choice["message"]) + return {**response, "choices": [{**choice, "message": {**message, "provider_specific_fields": fields}}]} + + +def _with_usage_details(response: Event, details: JsonValue) -> Event: + usage: Final = JSON_OBJECT.validate_python(response["usage"]) + return {**response, "usage": {**usage, "completion_tokens_details": details}} + + +def _rewrite_cached_response(needle: str, rewrite: Callable[[Event], Event]) -> str: + store: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + key: Final = eventually(lambda: _entry_key(store, needle), lambda found: found is not None) + assert key is not None, needle + entry: Final = JSON_OBJECT.validate_json(_stored(store, key)) + stored: Final = entry["response"] + encoded: Final = isinstance(stored, str) + response: Final = ( + JSON_OBJECT.validate_json(stored) if isinstance(stored, str) else JSON_OBJECT.validate_python(stored) + ) + rewritten: Final = {**rewrite(response), "id": f"{response['id']}{STALE_ID_SUFFIX}"} + store.set(key, json.dumps({**entry, "response": json.dumps(rewritten) if encoded else rewritten}), ex=60) + return str(rewritten["id"]) + + +def rewrite_cached_message_fields(needle: str, fields: JsonValue) -> str: + return _rewrite_cached_response(needle, lambda response: _with_message_fields(response, fields)) + + +def rewrite_cached_usage_details(needle: str, details: JsonValue) -> str: + return _rewrite_cached_response(needle, lambda response: _with_usage_details(response, details)) + + +def key_hash(key: str) -> str: + return hashlib.sha256(key.encode()).hexdigest() + + +def _key_hash_of(event: Mapping[str, JsonValue]) -> JsonValue: + return JSON_OBJECT.validate_python(event.get("metadata") or {}).get("user_api_key_hash") + + +class Ledger: + def __init__(self, wire: Wire) -> None: + self.wire: Final = wire + self.batches: tuple[Request, ...] = () + + def events(self) -> tuple[Event, ...]: + self.batches = (*self.batches, *self.wire.drain()) + decoded: Final = chain.from_iterable(JSON_LIST.validate_json(batch.body) for batch in self.batches) + return tuple(JSON_OBJECT.validate_python(event) for event in decoded) + + def for_key(self, key: str) -> tuple[Event, ...]: + hashed: Final = key_hash(key) + return tuple(event for event in self.events() if _key_hash_of(event) == hashed) + + def await_key(self, key: str, count: int = 1, seconds: float = 30) -> tuple[Event, ...]: + return eventually(lambda: self.for_key(key), lambda found: len(found) >= count, seconds=seconds) + + +_HIDDEN_PARAMS_PROBE: Final = """ +import httpx +from litellm.integrations.custom_logger import CustomLogger + + +class HiddenParamsProbe(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + payload = kwargs.get("standard_logging_object") or {{}} + hidden = getattr(response_obj, "_hidden_params", None) or {{}} + event = {{ + "metadata": {{"user_api_key_hash": (payload.get("metadata") or {{}}).get("user_api_key_hash")}}, + "response_id": getattr(response_obj, "id", None), + "received_model_id": hidden.get("received_model_id"), + }} + async with httpx.AsyncClient(trust_env=False) as client: + await client.post({url!r}, json=[event]) + + +instance = HiddenParamsProbe() +""" + + +def hidden_params_probe(directory: Path, url: str) -> str: + (directory / "hidden_params_probe.py").write_text(_HIDDEN_PARAMS_PROBE.format(url=url)) + return "hidden_params_probe.instance" + + +def sink_config( + directory: Path, + sinks: Mapping[str, str], + *, + retries: int = 0, + extra_callbacks: Sequence[str] = (), + store_prompts: bool = False, +) -> Path: + shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = { + **JSON_OBJECT.validate_python(shared["litellm_settings"]), + "callbacks": [*sinks, *extra_callbacks], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + } + callback_settings: Final = { + name: { + "callback_type": "generic_api", + "endpoint": url, + "headers": {"x-refusal-sink": name}, + "max_retries": retries, + "retry_delay": 0.5, + } + for name, url in sinks.items() + } + general_settings: Final = { + **JSON_OBJECT.validate_python(shared["general_settings"]), + "store_prompts_in_spend_logs": store_prompts, + } + path: Final = directory / "anthropic-refusal-sinks.yaml" + path.write_text( + yaml.safe_dump( + { + **shared, + "general_settings": general_settings, + "litellm_settings": settings, + "callback_settings": callback_settings, + } + ) + ) + return path + + +def frames(text: str) -> tuple[Event, ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: {") + ) + + +def frame_types(events: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + return tuple(event.get("type") for event in events) diff --git a/tests/integration/_support/redis_breaker_worker.py b/tests/integration/_support/redis_breaker_worker.py index f551baacf24..02d9c0ea2d5 100644 --- a/tests/integration/_support/redis_breaker_worker.py +++ b/tests/integration/_support/redis_breaker_worker.py @@ -11,10 +11,12 @@ import sys from collections.abc import Iterator, Mapping from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import Final, Literal, TypeAlias from pydantic import TypeAdapter +Scenario: TypeAlias = Literal["callbacks_before_cache", "callbacks_after_cache", "no_callbacks"] + STATE_GAUGE: Final = "litellm_redis_circuit_breaker_state" WORKER_TIMEOUT_SECONDS: Final = 120 @@ -29,9 +31,11 @@ class WorkerReport: _REPORT: Final = TypeAdapter(WorkerReport) -def run_breaker_worker(api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ) -> WorkerReport: +def run_breaker_worker( + scenario: Scenario, api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ +) -> WorkerReport: completed: Final = subprocess.run( - (sys.executable, "-I", __file__, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]), + (sys.executable, "-I", __file__, scenario, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]), env={**environment, "PROMETHEUS_MULTIPROC_DIR": str(prom_dir)}, capture_output=True, text=True, @@ -56,11 +60,15 @@ def _series(families: Iterator[object]) -> Iterator[tuple[tuple[tuple[str, str], yield tuple(sorted(sample.labels.items())), float(sample.value) -def main(api_base: str, redis_host: str, redis_port: str) -> None: +def main(scenario: str, api_base: str, redis_host: str, redis_port: str) -> None: import litellm from litellm.caching.caching import Cache + if scenario == "callbacks_before_cache": + litellm.service_callback = ["prometheus_system"] # test-quality-ok: this worker process exits right after the report litellm.cache = Cache(type="redis", host=redis_host, port=redis_port) # test-quality-ok: throwaway worker process + if scenario == "callbacks_after_cache": + litellm.service_callback = ["prometheus_system"] # test-quality-ok: throwaway worker process response: Final = litellm.completion( model="openai/gpt-5.6", api_base=api_base, @@ -76,4 +84,4 @@ def main(api_base: str, redis_host: str, redis_port: str) -> None: if __name__ == "__main__": - main(*sys.argv[1:4]) + main(*sys.argv[1:5]) diff --git a/tests/integration/_support/responses_stream.py b/tests/integration/_support/responses_stream.py index 1bcfebfba90..b28043e680f 100644 --- a/tests/integration/_support/responses_stream.py +++ b/tests/integration/_support/responses_stream.py @@ -2,13 +2,17 @@ import json from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Final -from integration._support.client import object_value, string_value +import httpx +from integration._support.client import JSON_OBJECT, eventually, object_value, string_value +from integration._support.database import read_rows from integration._support.wire import Reply, Request from pydantic import JsonValue AZURE_TARGET: Final = "/openai/v1/responses?api-version=" OPENAI_TARGET: Final = "/responses" RATE_LIMIT_MESSAGE: Final = "Your requests to gpt-6 have exceeded token rate limit." +RATE_LIMIT_PREFIX: Final = "litellm.RateLimitError: " +SENTINEL_PREFIX: Final = "litellm.MidStreamFallbackError: " def frame(event: Mapping[str, JsonValue]) -> bytes: @@ -36,6 +40,14 @@ def error_event(error: Mapping[str, JsonValue] | None) -> Mapping[str, JsonValue return {"type": "error", "sequence_number": 1, **({} if error is None else {"error": dict(error)})} +def openai_error_event(code: str, message: str) -> Mapping[str, JsonValue]: + return {"type": "error", "sequence_number": 1, "code": code, "message": message, "param": None} + + +def openai_error_stream(identity: str, code: str, message: str) -> tuple[bytes, ...]: + return (frame(created(identity)), frame(openai_error_event(code, message))) + + def azure_rate_limit() -> Mapping[str, JsonValue]: return { "type": "too_many_requests", @@ -126,3 +138,78 @@ def chat_content(frames: Sequence[Mapping[str, JsonValue]]) -> str: yield string_value(object_value(object_value(choice)["delta"]).get("content") or "") return "".join(deltas()) + + +def chat_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "stream": True, + "messages": [{"role": "user", "content": marker}], + "tools": function_tools(), + "num_retries": 0, + "cache": {"no-cache": True}, + **extra, + } + + +def messages_body(model: str, marker: str) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "stream": True, + "messages": [{"role": "user", "content": marker}], + "tools": [ + { + "name": "get_weather", + "description": "Weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ], + "num_retries": 0, + "cache": {"no-cache": True}, + } + + +def error_body(response: httpx.Response) -> Mapping[str, JsonValue]: + return object_value(JSON_OBJECT.validate_json(response.content)["error"]) + + +def data_frames(text: str) -> tuple[Mapping[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def sse_events(text: str) -> tuple[tuple[str, Mapping[str, JsonValue]], ...]: + def parse(block: str) -> tuple[str, Mapping[str, JsonValue]]: + lines: Final = block.splitlines() + event: Final = next(line.removeprefix("event: ") for line in lines if line.startswith("event: ")) + data: Final = next(line.removeprefix("data: ") for line in lines if line.startswith("data: ")) + return event, JSON_OBJECT.validate_json(data) + + return tuple(parse(block) for block in text.strip().split("\n\n") if "event: " in block) + + +def spend_row(call_id: str) -> Mapping[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT litellm_call_id, status, model_group, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', + (call_id,), + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + return rows[0] + + +def assert_provider_typed_rate_limit( + error: Mapping[str, JsonValue], provider_message: str = RATE_LIMIT_MESSAGE +) -> None: + assert error["type"] == "throttling_error", error + assert str(error["code"]) == "429", error + message: Final = string_value(error["message"]) + assert message.startswith(RATE_LIMIT_PREFIX) and provider_message in message, message + assert SENTINEL_PREFIX not in message, message diff --git a/tests/integration/caching/response_cache_case.py b/tests/integration/caching/response_cache_case.py index 9df34721202..8dfb6a0a479 100644 --- a/tests/integration/caching/response_cache_case.py +++ b/tests/integration/caching/response_cache_case.py @@ -336,6 +336,29 @@ def flip_entry_to_empty(store: Redis, key: str, field: str) -> str: return stale_id +def rewrite_usage_details(store: Redis, key: str, details: JsonValue) -> str: + """Rewrite the stored usage's ``completion_tokens_details`` to ``details`` and stamp the id ``-stale``, keeping the + encoding, and return the stale id.""" + raw: Final = store.get(key) + assert raw is not None, key + entry, is_json = _decode_entry(raw) + response: Final = entry["response"] + stored: Final = object_value(json.loads(response) if isinstance(response, str) else response) + stale_id: Final = f"{string_value(stored['id'])}{STALE_ID_SUFFIX}" + usage: Final = object_value(stored["usage"]) + rewritten_response: Final[dict[str, JsonValue]] = { + **stored, + "usage": {**usage, "completion_tokens_details": details}, + "id": stale_id, + } + rewritten: Final[dict[str, JsonValue]] = { + **entry, + "response": json.dumps(rewritten_response) if isinstance(response, str) else rewritten_response, + } + store.set(key, json.dumps(rewritten) if is_json else str(rewritten), ex=STALE_ENTRY_TTL_SECONDS) + return stale_id + + def chat_body(model: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]: return {"model": model, "messages": [{"role": "user", "content": text}], **extra} diff --git a/tests/integration/database/test_prisma_migration_lock.py b/tests/integration/database/test_prisma_migration_lock.py new file mode 100644 index 00000000000..0c6365b2b5b --- /dev/null +++ b/tests/integration/database/test_prisma_migration_lock.py @@ -0,0 +1,58 @@ +import os +import secrets +from typing import Final + +import psycopg +import pytest +from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, PRISMA_ADVISORY_LOCK_KEY, prisma_migration_lock + +pytestmark = pytest.mark.skipif("DATABASE_URL" not in os.environ, reason="needs a live Postgres") + + +def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, ...]]) -> bool: + row: Final = connection.execute("SELECT pg_try_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)).fetchone() + taken: Final = row is not None and row[0] is True + if taken: + connection.execute("SELECT pg_advisory_unlock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)) + return taken + + +def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]: + rows: Final = connection.execute( + "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted" + ).fetchall() + return frozenset(int(row[0]) for row in rows) + + +def _index_builds_concurrently_now(connection: psycopg.Connection[tuple[object, ...]]) -> bool: + table: Final = f'"lit_cic_{secrets.token_hex(4)}"' + connection.execute(f"CREATE TABLE {table} (a int)") + try: + connection.execute("SET lock_timeout = '5s'") + connection.execute(f"CREATE INDEX CONCURRENTLY ON {table} (a)") + except psycopg.errors.LockNotAvailable: + return False + finally: + connection.execute("RESET lock_timeout") + connection.execute(f"DROP TABLE {table}") + return True + + +def test_prisma_migration_lock_does_not_stall_a_concurrent_index_build(): + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer: + with prisma_migration_lock(os.environ["DATABASE_URL"]): + assert _index_builds_concurrently_now(peer) + + +def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leaves_nothing_behind(): + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer: + assert _can_take_prismas_session_lock(peer) + + with prisma_migration_lock(os.environ["DATABASE_URL"]): + held: Final = _granted_advisory_keys(peer) + assert PRISMA_ADVISORY_LOCK_KEY in held + assert MIGRATION_LOCK_KEY in held + assert not _can_take_prismas_session_lock(peer) + + assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY}) + assert _can_take_prismas_session_lock(peer) diff --git a/tests/integration/management/test_team_update.py b/tests/integration/management/test_team_update.py index 0414d0bb77f..e1d9c43f318 100644 --- a/tests/integration/management/test_team_update.py +++ b/tests/integration/management/test_team_update.py @@ -162,6 +162,87 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context( assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text +_MEMBER_DEFAULTS = [ + pytest.param("team_member_budget", 25.0, "max_budget", id="team_member_budget"), + pytest.param("team_member_budget_duration", "7d", "budget_duration", id="team_member_budget_duration"), + pytest.param("team_member_tpm_limit", 200, "tpm_limit", id="team_member_tpm_limit"), + pytest.param("team_member_rpm_limit", 9, "rpm_limit", id="team_member_rpm_limit"), +] + + +@pytest.mark.parametrize("field,value,budget_column", _MEMBER_DEFAULTS) +async def test_team_update_team_admin_member_default_follows_the_proxy_admin_grant( + field: str, value: str | float, budget_column: str, gateway_client, db, resources, actors +): + callers = await _callers(db, actors, resources.prefix) + await _seed_target(db, actors, callers, "alpha", resources.prefix) + headers = {"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"} + body = {"team_id": resources.prefix, field: value} + + with team_admin_permissions(gateway_client.gateway, ["tpm_limit"]): + denied = await gateway_client.post("/team/update", headers=headers, json=body) + assert denied.status_code == 403, denied.text + assert f"'{field}'" in denied.text and "Team admin editable fields" in denied.text, denied.text + + with team_admin_permissions(gateway_client.gateway, [field]): + allowed = await gateway_client.post("/team/update", headers=headers, json=body) + assert allowed.status_code == 200, allowed.text + + info = await gateway_client.get("/team/info", params={"team_id": resources.prefix}) + assert info.status_code == 200, info.text + assert info.json()["team_info"]["team_member_budget_table"][budget_column] == value + + +async def test_team_update_team_admin_resending_member_defaults_leaves_the_budget_row_alone( + gateway_client, db, resources, actors +): + callers = await _callers(db, actors, resources.prefix) + await _seed_target(db, actors, callers, "alpha", resources.prefix) + stored = { + "team_member_budget": 10.0, + "team_member_budget_duration": "30d", + "team_member_tpm_limit": 100, + "team_member_rpm_limit": 5, + } + seeded = await gateway_client.post("/team/update", json={"team_id": resources.prefix, **stored}) + assert seeded.status_code == 200, seeded.text + team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix}) + assert team is not None + budget_id = team.metadata["team_member_budget_id"] + before = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id}) + assert before is not None + + with team_admin_permissions(gateway_client.gateway, ["tpm_limit"]): + resp = await gateway_client.post( + "/team/update", + headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"}, + json={"team_id": resources.prefix, "tpm_limit": 1000, **stored}, + ) + assert resp.status_code == 200, resp.text + + after = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id}) + assert after is not None and after.updated_at == before.updated_at + team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix}) + assert team is not None and team.tpm_limit == 1000 + + +async def test_team_update_team_admin_cannot_change_member_defaults_of_another_team( + gateway_client, db, resources, actors +): + callers = await _callers(db, actors, resources.prefix) + await _seed_target(db, actors, callers, "beta", resources.prefix) + + with team_admin_permissions(gateway_client.gateway, ["team_member_budget"]): + resp = await gateway_client.post( + "/team/update", + headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"}, + json={"team_id": resources.prefix, "team_member_budget": 25.0}, + ) + assert resp.status_code == 403, resp.text + team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix}) + assert team is not None and "team_member_budget_id" not in (team.metadata or {}), "denied but team mutated" + + async def test_team_update_team_admin_soft_budget_follows_the_proxy_admin_grant(gateway_client, db, resources, actors): callers = await _callers(db, actors, resources.prefix) await _seed_target(db, actors, callers, "alpha", resources.prefix) diff --git a/tests/integration/providers/_cache_control_marks_support.py b/tests/integration/providers/_cache_control_marks_support.py index 5194f7d7ca3..2f6071f9e45 100644 --- a/tests/integration/providers/_cache_control_marks_support.py +++ b/tests/integration/providers/_cache_control_marks_support.py @@ -20,6 +20,7 @@ BEDROCK_CACHE_CHECKPOINT_CAP: Final = 4 ANTHROPIC_MODEL: Final = "claude-opus-5-5" BEDROCK_MODEL: Final = "anthropic.claude-opus-5-5" +UNMAPPED_BEDROCK_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/scripted" PROVIDER_KEY: Final = "synthetic-provider-key" SYSTEM: Final = "Use the provided tool results to answer the user." ASK: Final = "Look up the weather in London, Paris, and Tokyo." @@ -409,6 +410,23 @@ def bedrock_peer(request: Request) -> Reply: ) +def _without_cache_points(value: JsonValue) -> JsonValue: + if isinstance(value, dict): + return {key: _without_cache_points(inner) for key, inner in value.items()} + if isinstance(value, list): + return [ + _without_cache_points(inner) for inner in value if not (isinstance(inner, dict) and "cachePoint" in inner) + ] + return value + + +def sent_without_cache_points(requests: Sequence[Request], marker: str) -> list[JsonValue]: + return [ + _without_cache_points(_JSON_OBJECT.validate_json(request.body.replace(marker.encode(), b"marker"))) + for request in requests + ] + + def gateway_injected(response_id: str) -> bool: rows: Final = eventually( lambda: read_rows( diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py new file mode 100644 index 00000000000..dd4b7602b96 --- /dev/null +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py @@ -0,0 +1,337 @@ +import asyncio +import json +import re +import threading +import time +import uuid +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.anthropic_refusal import ( + REFUSAL_FIELDS, + Ledger, + Location, + refusal_prompt, + relayed, + respond, + sink_config, +) +from integration._support.anthropic_thinking import JSON_OBJECT, MODEL, Event, chunks_of +from integration._support.client import Gateway, eventually, held +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.timing import assert_not_delayed_by +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark: Final = [ + pytest.mark.timeout(2 * graceful_stop_seconds() + 120), + pytest.mark.xdist_group("anthropic_refusal_stop_details_chaos"), +] + +_MODEL_NAME: Final = "anthropic-refusal-chaos" +_SLOW_SINK_SECONDS: Final = 4 +_SINK: Final = "refusal_chaos_sink" +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields") +_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields") +_PATHS: Final[Mapping[str, str]] = MappingProxyType( + {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"} +) + +Surface: TypeAlias = Literal["chat", "messages", "responses"] +_RecordProperty: TypeAlias = Callable[[str, object], None] + + +@dataclass(frozen=True, slots=True) +class _Call: + surface: Surface + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +@dataclass(frozen=True, slots=True) +class _Held: + release: threading.Event + markers: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + self.markers.put(_markers_in(request.body.decode())[0]) + assert self.release.wait(timeout=90), "The burst was never released" + return respond(request) + + +def _markers_in(text: str) -> tuple[str, ...]: + return tuple(dict.fromkeys(_MARKER.findall(text))) + + +def _calls(count: int) -> tuple[_Call, ...]: + surfaces: Final[tuple[Surface, ...]] = ("chat", "messages", "responses") + return tuple( + _Call(surfaces[index % 3], stream=(index // 3) % 2 == 0, marker=uuid.uuid4().hex) for index in range(count) + ) + + +def _body(call: _Call) -> Event: + text: Final = refusal_prompt("refusal", call.marker) + shared: Final[Event] = {"model": _MODEL_NAME, "stream": call.stream, "cache": {"no-cache": True}} + match call.surface: + case "chat": + return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64} + case "messages": + return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64} + case "responses": + return {**shared, "input": text, "max_output_tokens": 64} + + +async def _send(client: httpx.AsyncClient, key: str, call: _Call) -> _Served: + try: + response: Final = await client.post( + _PATHS[call.surface], json=_body(call), headers={"Authorization": f"Bearer {key}"} + ) + except httpx.TransportError as error: + return _Served(call, 0, repr(error)) + return _Served(call, response.status_code, response.text) + + +async def _burst(gateway: Gateway, calls: Sequence[_Call]) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client: + return tuple(await asyncio.gather(*(_send(client, gateway.key, call) for call in calls))) + + +def _chaos_config(directory: Path, upstream: Wire, sink: str, *, retries: int) -> Path: + base: Final = sink_config(directory, {_SINK: sink}, retries=retries) + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(base.read_text())) + model_list: Final = [ + { + "model_name": _MODEL_NAME, + "litellm_params": {"model": f"anthropic/{MODEL}", "api_base": upstream.url, "api_key": "scripted-key"}, + } + ] + path: Final = directory / "anthropic-refusal-chaos.yaml" + path.write_text(yaml.safe_dump({**config, "model_list": model_list})) + return path + + +def _ledger_events(ledgers: Sequence[Ledger]) -> tuple[Event, ...]: + return tuple(chain.from_iterable(ledger.events() for ledger in ledgers)) + + +def _marker_counts(ledgers: Sequence[Ledger], markers: frozenset[str]) -> Mapping[str, int]: + found: Final = chain.from_iterable(_markers_in(json.dumps(event)) for event in _ledger_events(ledgers)) + return MappingProxyType(Counter(marker for marker in found if marker in markers)) + + +def _events_for(ledgers: Sequence[Ledger], marker: str) -> tuple[Event, ...]: + return tuple(event for event in _ledger_events(ledgers) if marker in json.dumps(event)) + + +def _assert_answered(item: _Served) -> None: + assert item.status == 200, (item.call, item.text) + if item.call.surface != "chat": + return + if item.call.stream: + found: Final = tuple(chain.from_iterable(relayed(chunk) for chunk in chunks_of(item.text))) + assert found == ((_DELTA_FIELDS, REFUSAL_FIELDS),), item.text + return + assert relayed(JSON_OBJECT.validate_json(item.text)) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), item.text + + +def _assert_logged_chat_relays(ledgers: Sequence[Ledger], calls: Sequence[_Call]) -> None: + for call in calls: + if call.surface != "chat": + continue + (event,) = _events_for(ledgers, call.marker) + assert relayed(event["response"]) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), event + + +def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: + return eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _port(wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + assert port is not None, wire.url + return port + + +def _accept(_request: Request) -> Reply: + return Reply() + + +async def _timed_chat(gateway: Gateway) -> float: + started: Final = time.monotonic() + (served,) = await _burst(gateway, (_Call("chat", stream=False, marker=uuid.uuid4().hex),)) + _assert_answered(served) + return time.monotonic() - started + + +async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(30) + markers: Final = frozenset(call.marker for call in calls) + with wire_server(respond) as upstream, ExitStack() as first_sink: + sink: Final = first_sink.enter_context(wire_server(_accept)) + first: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=6) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + _worker_pids(owned) + served_first: Final = await _burst(owned.gateway, calls[:10]) + await asyncio.to_thread( + eventually, lambda: sum(_marker_counts((first,), markers).values()), lambda total: total >= 1, 30 + ) + first.events() + first_sink.close() + served_during: Final = await _burst(owned.gateway, calls[10:20]) + with wire_server(_accept, port=_port(sink)) as restarted: + second: Final = Ledger(restarted) + served_after: Final = await _burst(owned.gateway, calls[20:]) + for item in (*served_first, *served_during, *served_after): + _assert_answered(item) + counts: Final = await asyncio.to_thread( + held, + lambda: _marker_counts((first, second), markers), + lambda found: set(found) == markers and set(found.values()) == {1}, + holding=3, + seconds=60, + ) + assert sum(counts.values()) == 30, counts + _assert_logged_chat_relays((first, second), calls) + assert len(upstream.drain()) == 30 + + +async def test_slow_sink_neither_deadlocks_nor_uploads_a_refusal_event_twice(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(30) + markers: Final = frozenset(call.marker for call in calls) + + def slow(_request: Request) -> Reply: + time.sleep(_SLOW_SINK_SECONDS) + return Reply() + + with wire_server(respond) as upstream, wire_server(slow) as sink: + ledger: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + _worker_pids(owned) + control: Final = await _timed_chat(owned.gateway) + served: Final = await _burst(owned.gateway, calls) + for item in served: + _assert_answered(item) + follow_up: Final = await _timed_chat(owned.gateway) + assert_not_delayed_by((follow_up,), control=(control,), injected_delay=_SLOW_SINK_SECONDS) + counts: Final = await asyncio.to_thread( + held, + lambda: _marker_counts((ledger,), markers), + lambda found: set(found) == markers and set(found.values()) == {1}, + holding=5, + seconds=90, + ) + assert sum(counts.values()) == 30, counts + _assert_logged_chat_relays((ledger,), calls) + + +async def test_worker_sigkill_mid_burst_leaves_the_sibling_relaying_and_logging_once( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(20) + hold: Final = _Held(threading.Event(), SimpleQueue()) + with wire_server(hold.respond) as upstream, wire_server(_accept) as sink: + ledger: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + workers: Final = _worker_pids(owned) + burst: Final = asyncio.create_task(_burst(owned.gateway, calls)) + await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, upstream.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + hold.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + completed: Final = tuple(item for item in served if item.status == 200) + assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served]) + for item in completed: + _assert_answered(item) + follow_up: Final = _Call("chat", stream=True, marker=uuid.uuid4().hex) + (answered,) = await _burst(owned.gateway, (follow_up,)) + _assert_answered(answered) + survivors: Final = frozenset((*(item.call.marker for item in completed), follow_up.marker)) + counts: Final = await asyncio.to_thread( + held, + lambda: _marker_counts((ledger,), frozenset(call.marker for call in (*calls, follow_up))), + lambda found: set(found) == survivors and set(found.values()) == {1}, + holding=3, + seconds=60, + ) + assert sum(counts.values()) == len(survivors), counts + _assert_logged_chat_relays((ledger,), tuple(item.call for item in completed)) + + +async def test_proxy_restart_mid_burst_never_logs_a_refusal_twice( + gateway: Gateway, tmp_path: Path, record_property: _RecordProperty +) -> None: + calls: Final = _calls(20) + markers: Final = frozenset(call.marker for call in calls) + hold: Final = _Held(threading.Event(), SimpleQueue()) + with wire_server(hold.respond) as upstream, wire_server(_accept) as sink: + ledger: Final = Ledger(sink) + config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first: + _worker_pids(first) + burst: Final = asyncio.create_task(_burst(first.gateway, calls)) + await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60) + first.begin_shutdown() + hold.release.set() + served: Final = await burst + for item in served: + _assert_answered(item) + second_directory: Final = tmp_path / "second" + second_directory.mkdir() + with owned_proxy_process(gateway, second_directory, {}, config=config, workers=2) as second: + follow_up: Final = _Call("chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(second.gateway, (follow_up,)) + _assert_answered(answered) + await asyncio.to_thread( + eventually, + lambda: _marker_counts((ledger,), frozenset({follow_up.marker})), + lambda found: found[follow_up.marker] == 1, + 30, + ) + counts: Final = _marker_counts((ledger,), markers) + assert set(counts.values()) <= {1}, counts + _assert_logged_chat_relays((ledger,), tuple(call for call in calls if call.marker in counts)) + record_property("served", len(served)) + record_property("landed", sum(counts.values())) diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py new file mode 100644 index 00000000000..461839ccbcc --- /dev/null +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py @@ -0,0 +1,1356 @@ +import asyncio +import json +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from typing import Final, Literal, TypeAlias +from urllib.parse import unquote, urlsplit + +import anthropic +import httpx +import openai +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.anthropic_refusal import ( + CITED_FIELDS, + DECLARED_PARTIAL_JSON, + DOCS_PARTIAL_TEXT, + DOCS_REFUSAL, + DOCS_STOP_DETAILS, + ENVELOPED_PARTIAL_JSON, + FIXTURES, + JSON_ANSWER_FORMAT, + JSON_TOOL_NAME, + MALFORMED_STOP_DETAILS, + REFUSAL_FIELDS, + UNVALIDATED_USAGE_DETAILS, + Ledger, + Location, + frame_types, + frames, + hidden_params_probe, + openai_id, + refusal_events, + refusal_prompt, + relayed, + respond, + rewrite_cached_message_fields, + rewrite_cached_usage_details, + sink_config, +) +from integration._support.anthropic_thinking import ( + BEDROCK_MODEL, + JSON_LIST, + JSON_OBJECT, + MODEL, + NO_CACHE, + Event, + aws_chunks, + chunks_of, + sse_chunks, +) +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, held +from integration._support.forward_proxy import tunnelling_forward_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Wire, wire_server +from pydantic import BaseModel, JsonValue + +pytestmark: Final = [ + pytest.mark.timeout(2 * graceful_stop_seconds() + 120), + pytest.mark.xdist_group("anthropic_refusal_stop_details_wire"), +] + +_ANTHROPIC_BASE: Final = "http://api.anthropic.com" +_ANTHROPIC_KEY: Final = "scripted-anthropic-key" +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_HOST: Final = f"{_VERTEX_LOCATION}-aiplatform.googleapis.com" +_VERTEX_MODEL_PATH: Final = ( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{MODEL}" +) +_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields") +_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields") +_RESPONSE_FIELDS: Final[Location] = ("provider_specific_fields",) +_STREAMED_RESPONSE_FIELDS: Final[Location] = ("response", "provider_specific_fields") + +Client: TypeAlias = Literal["sdk", "async_sdk", "httpx"] +Provider: TypeAlias = Literal[ + "anthropic", + "bedrock_invoke", + "claude_platform", + "vertex_ai", + "azure_ai", + "bedrock_mantle", + "snowflake", + "openai_compatible", +] +Surface: TypeAlias = Literal["chat", "chat_stream", "messages", "messages_stream", "responses", "responses_stream"] + +CLIENTS: Final[tuple[Client, ...]] = ("sdk", "async_sdk", "httpx") +FAMILY: Final[tuple[Provider, ...]] = ( + "bedrock_invoke", + "claude_platform", + "vertex_ai", + "azure_ai", + "bedrock_mantle", + "snowflake", +) +_STREAM_TARGETS: Final[Mapping[Provider, tuple[str, str]]] = { + "anthropic": ("/v1/messages", "/v1/messages"), + "bedrock_invoke": (f"/model/{BEDROCK_MODEL}/invoke", f"/model/{BEDROCK_MODEL}/invoke-with-response-stream"), + "claude_platform": ("/v1/messages", "/v1/messages"), + "vertex_ai": (f"{_VERTEX_MODEL_PATH}:rawPredict", f"{_VERTEX_MODEL_PATH}:streamRawPredict"), + "azure_ai": ("/anthropic/v1/messages", "/anthropic/v1/messages"), + "bedrock_mantle": ("/anthropic/v1/messages", "/anthropic/v1/messages"), + "snowflake": ("/api/v2/cortex/v1/messages", "/api/v2/cortex/v1/messages"), +} + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + wire: Wire + vertex: Wire + sink: Ledger + forbidden: Ledger + absent: Ledger + hidden: Ledger + credentials: str + + +def _service_account_json(token_url: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": _VERTEX_PROJECT, + "private_key_id": "scripted", + "private_key": private_key, + "client_email": f"scripted@{_VERTEX_PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{token_url}/_oauth/authorize", + "token_uri": f"{token_url}/_oauth/token", + } + ) + + +def _port(wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + assert port is not None, wire.url + return port + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("anthropic-refusal-stop-details") + certificate: Final = write_self_signed_cert(directory, (_VERTEX_HOST,)) + with ( + gateway_from_environment() as parent, + wire_server(respond) as wire, + wire_server(respond, tls=server_context(*certificate)) as vertex, + tunnelling_forward_proxy(_port(vertex)) as tunnel, + wire_server(lambda _request: Reply()) as sink, + wire_server(lambda _request: Reply(status=403)) as forbidden, + wire_server(lambda _request: Reply(status=404)) as absent, + wire_server(lambda _request: Reply()) as hidden, + ): + credentials: Final = _service_account_json(parent.upstream_url.rstrip("/")) + credentials_path: Final = directory / "vertex-service-account.json" + credentials_path.write_text(credentials) + config: Final = sink_config( + directory, + {"refusal_sink": sink.url, "refusal_sink_403": forbidden.url, "refusal_sink_404": absent.url}, + extra_callbacks=(hidden_params_probe(directory, hidden.url),), + store_prompts=True, + ) + overrides: Final = { + "ANTHROPIC_API_BASE": _ANTHROPIC_BASE, + "ANTHROPIC_API_KEY": _ANTHROPIC_KEY, + "AIOHTTP_TRUST_ENV": "True", + "HTTP_PROXY": wire.url, + "http_proxy": wire.url, + "HTTPS_PROXY": tunnel.url, + "https_proxy": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(certificate[0]), + "DEFAULT_VERTEXAI_PROJECT": _VERTEX_PROJECT, + "DEFAULT_VERTEXAI_LOCATION": _VERTEX_LOCATION, + "DEFAULT_GOOGLE_APPLICATION_CREDENTIALS": str(credentials_path), + } + with owned_proxy(parent, directory, overrides, config=config, workers=2) as owned: + yield Rig(owned, wire, vertex, Ledger(sink), Ledger(forbidden), Ledger(absent), Ledger(hidden), credentials) + + +@pytest.fixture(autouse=True) +def _drained_upstreams(rig: Rig) -> None: + rig.wire.drain() + rig.vertex.drain() + + +def _marker() -> str: + return uuid.uuid4().hex + + +def _deployment(scenario: Scenario, provider: Provider, rig: Rig) -> str: + wire_url: Final = rig.wire.url + match provider: + case "anthropic": + return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key=_ANTHROPIC_KEY) + case "bedrock_invoke": + return scenario.model( + model=f"bedrock/invoke/{BEDROCK_MODEL}", + api_base=wire_url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire_url, + ) + case "claude_platform": + return scenario.model( + model=f"bedrock/claude_platform/{MODEL}", + api_base=wire_url, + api_key="scripted-platform-key", + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + case "vertex_ai": + return scenario.model( + model=f"vertex_ai/{MODEL}", + api_base=f"{wire_url}{_VERTEX_MODEL_PATH}", + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=rig.credentials, + ) + case "azure_ai": + return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key") + case "bedrock_mantle": + return scenario.model( + model=f"bedrock_mantle/{BEDROCK_MODEL}", + api_base=wire_url, + api_key="scripted-mantle-key", + aws_region_name="us-east-1", + ) + case "snowflake": + return scenario.model(model=f"snowflake/{MODEL}", api_base=wire_url, api_key="scripted-snowflake-key") + case "openai_compatible": + return scenario.model(model="openai/gpt-4o-mini", api_base=f"{wire_url}/v1", api_key="scripted-openai-key") + + +def _base_url(rig: Rig) -> str: + return str(rig.gateway.client.base_url).rstrip("/") + + +def _dump(model: BaseModel) -> Event: + return JSON_OBJECT.validate_python(model.model_dump()) + + +def _post(rig: Rig, path: str, body: Mapping[str, JsonValue], key: str) -> bytes: + response: Final = rig.gateway.request("POST", path, body, key=key) + assert response.status_code == 200, response.text + return response.content + + +def _chat_request(model: str, text: str, *, stream: bool, cached: bool) -> Event: + request: Final[Event] = { + "model": model, + "messages": [{"role": "user", "content": text}], + "stream": stream, + "max_tokens": 64, + } + return request if cached else {**request, **NO_CACHE} + + +def _chat(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event: + match client: + case "httpx": + return JSON_OBJECT.validate_json( + _post(rig, "/v1/chat/completions", _chat_request(model, text, stream=False, cached=cached), key) + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump( + sdk.chat.completions.create( + model=model, messages=[{"role": "user", "content": text}], max_tokens=64, extra_body=NO_CACHE + ) + ) + case "async_sdk": + + async def call() -> Event: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump( + await sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": text}], + max_tokens=64, + extra_body=NO_CACHE, + ) + ) + + return asyncio.run(call()) + + +def _chat_stream( + rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False +) -> tuple[Event, ...]: + match client: + case "httpx": + return chunks_of( + _post(rig, "/v1/chat/completions", _chat_request(model, text, stream=True, cached=cached), key).decode() + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return tuple( + _dump(chunk) + for chunk in sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": text}], + max_tokens=64, + stream=True, + extra_body=NO_CACHE, + ) + ) + case "async_sdk": + + async def call() -> tuple[Event, ...]: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + stream: Final = await sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": text}], + max_tokens=64, + stream=True, + extra_body=NO_CACHE, + ) + return tuple([_dump(chunk) async for chunk in stream]) + + return asyncio.run(call()) + + +def _messages_request(model: str, text: str, *, stream: bool, cached: bool) -> Event: + request: Final[Event] = { + "model": model, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": text}], + } + return request if cached else {**request, **NO_CACHE} + + +def _messages(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event: + match client: + case "httpx": + return JSON_OBJECT.validate_json( + _post(rig, "/v1/messages", _messages_request(model, text, stream=False, cached=cached), key) + ) + case "sdk": + with anthropic.Anthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + return _dump( + sdk.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": text}], extra_body=NO_CACHE + ) + ) + case "async_sdk": + + async def call() -> Event: + async with anthropic.AsyncAnthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + return _dump( + await sdk.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": text}], + extra_body=NO_CACHE, + ) + ) + + return asyncio.run(call()) + + +def _messages_stream( + rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False +) -> tuple[Event, ...]: + match client: + case "httpx": + return frames( + _post(rig, "/v1/messages", _messages_request(model, text, stream=True, cached=cached), key).decode() + ) + case "sdk": + with anthropic.Anthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + return tuple( + _dump(event) + for event in sdk.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": text}], + stream=True, + extra_body=NO_CACHE, + ) + ) + case "async_sdk": + + async def call() -> tuple[Event, ...]: + async with anthropic.AsyncAnthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk: + stream: Final = await sdk.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": text}], + stream=True, + extra_body=NO_CACHE, + ) + return tuple([_dump(event) async for event in stream]) + + return asyncio.run(call()) + + +def _responses_request(model: str, text: str, *, stream: bool, cached: bool) -> Event: + request: Final[Event] = {"model": model, "input": text, "max_output_tokens": 64, "stream": stream} + return request if cached else {**request, **NO_CACHE} + + +def _responses(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event: + match client: + case "httpx": + return JSON_OBJECT.validate_json( + _post(rig, "/v1/responses", _responses_request(model, text, stream=False, cached=cached), key) + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump(sdk.responses.create(model=model, input=text, max_output_tokens=64, extra_body=NO_CACHE)) + case "async_sdk": + + async def call() -> Event: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return _dump( + await sdk.responses.create(model=model, input=text, max_output_tokens=64, extra_body=NO_CACHE) + ) + + return asyncio.run(call()) + + +def _responses_stream(rig: Rig, key: str, model: str, text: str, client: Client = "httpx") -> tuple[Event, ...]: + match client: + case "httpx": + return frames( + _post(rig, "/v1/responses", _responses_request(model, text, stream=True, cached=False), key).decode() + ) + case "sdk": + with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + return tuple( + _dump(event) + for event in sdk.responses.create( + model=model, input=text, max_output_tokens=64, stream=True, extra_body=NO_CACHE + ) + ) + case "async_sdk": + + async def call() -> tuple[Event, ...]: + async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk: + stream: Final = await sdk.responses.create( + model=model, input=text, max_output_tokens=64, stream=True, extra_body=NO_CACHE + ) + return tuple([_dump(event) async for event in stream]) + + return asyncio.run(call()) + + +def _relayed_in(events: Sequence[JsonValue]) -> tuple[tuple[Location, Event], ...]: + return tuple(chain.from_iterable(relayed(event) for event in events)) + + +def _choice(chunk: Mapping[str, JsonValue]) -> Event: + choices: Final = JSON_LIST.validate_python(chunk.get("choices") or []) + return JSON_OBJECT.validate_python(choices[0]) if choices else {} + + +def _finish_reasons(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + return tuple(_choice(chunk)["finish_reason"] for chunk in chunks if _choice(chunk).get("finish_reason")) + + +def _streamed_text(chunks: Sequence[Mapping[str, JsonValue]]) -> str: + return "".join( + str(JSON_OBJECT.validate_python(_choice(chunk).get("delta") or {}).get("content") or "") for chunk in chunks + ) + + +def _message_text(body: Mapping[str, JsonValue]) -> JsonValue: + return JSON_OBJECT.validate_python(_choice(body)["message"]).get("content") + + +def _native_stops(events: Sequence[Mapping[str, JsonValue]]) -> tuple[Event, ...]: + deltas: Final = tuple( + JSON_OBJECT.validate_python(event["delta"]) for event in events if event.get("type") == "message_delta" + ) + return tuple({name: delta.get(name) for name in ("stop_reason", "stop_details")} for delta in deltas) + + +def _native_deltas(events: Sequence[Mapping[str, JsonValue]]) -> tuple[Event, ...]: + deltas: Final = tuple( + JSON_OBJECT.validate_python(event["delta"]) for event in events if event.get("type") == "message_delta" + ) + return tuple({name: delta[name] for name in ("stop_reason", "stop_details") if name in delta} for delta in deltas) + + +def _targets(wire: Wire) -> tuple[str, ...]: + return tuple(unquote(request.target) for request in wire.drain()) + + +def _paths(wire: Wire) -> tuple[str, ...]: + return tuple(urlsplit(target).path for target in _targets(wire)) + + +def _logged(rig: Rig, key: str, count: int = 1) -> tuple[Event, ...]: + events: Final = rig.sink.await_key(key, count) + assert len(events) == count, events + return events + + +def _logged_response(rig: Rig, key: str) -> JsonValue: + return _logged(rig, key)[0]["response"] + + +@pytest.mark.parametrize("client", CLIENTS) +def test_chat_refusal_relays_the_docs_stop_details_to_the_caller_and_the_log(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body + assert _choice(body)["finish_reason"] == "content_filter", body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_chat_stream_refusal_relays_the_docs_stop_details_in_one_chunk_and_the_log(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks + assert _finish_reasons(chunks) == ("content_filter",), chunks + assert _streamed_text(chunks) == "", chunks + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_messages_refusal_returns_the_docs_body_and_logs_the_stop_details(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _messages(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert (body["stop_reason"], body["stop_details"], body["content"]) == ("refusal", DOCS_STOP_DETAILS, []), body + assert relayed(body) == (), body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_messages_stream_refusal_keeps_the_client_frames_and_logs_the_stop_details(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert frame_types(events) == frame_types(FIXTURES["refusal"].events), events + assert _native_stops(events) == (REFUSAL_FIELDS,), events + assert _relayed_in(events) == (), events + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_responses_refusal_relays_the_docs_stop_details_at_the_top_level(rig: Rig, client: Client) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _responses(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert relayed(body) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), body + assert (body["status"], body["incomplete_details"]) == ("incomplete", {"reason": "content_filter"}), body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("client", CLIENTS) +def test_responses_stream_refusal_relays_the_docs_stop_details_on_the_incomplete_event( + rig: Rig, client: Client +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + events: Final = _responses_stream(rig, key, model, refusal_prompt("refusal", _marker()), client) + assert _relayed_in(events) == ((_STREAMED_RESPONSE_FIELDS, REFUSAL_FIELDS),), events + assert frame_types(events)[-1] == "response.incomplete", events + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +def test_anthropic_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + body: Final = { + "model": MODEL, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + content: Final = _post(rig, "/anthropic/v1/messages", body, key) + if stream: + assert content == b"".join(sse_chunks(refusal_events(DOCS_REFUSAL))), content + else: + assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content + assert _targets(rig.wire) == (f"{_ANTHROPIC_BASE}/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("provider", FAMILY) +def test_family_chat_refusal_relays_the_docs_stop_details(rig: Rig, provider: Provider) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker())) + expected: Final = () if provider == "snowflake" else ((_CHAT_FIELDS, REFUSAL_FIELDS),) + assert relayed(body) == expected, body + assert _choice(body)["finish_reason"] == ("stop" if provider == "snowflake" else "content_filter"), body + assert _paths(rig.wire) == (_STREAM_TARGETS[provider][0],) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == expected, logged + + +@pytest.mark.parametrize("provider", FAMILY) +def test_family_chat_stream_refusal_relays_the_docs_stop_details(rig: Rig, provider: Provider) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, rig) + key: Final = scenario.key() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt("refusal", _marker())) + assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks + assert _finish_reasons(chunks) == ("content_filter",), chunks + assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("provider", FAMILY) +def test_family_messages_stream_refusal_keeps_frames_and_logs_the_stop_details(rig: Rig, provider: Provider) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, provider, rig) + key: Final = scenario.key() + events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker())) + native: Final = {"stop_reason": "refusal"} if provider == "snowflake" else REFUSAL_FIELDS + assert _native_deltas(events) == (native,), events + assert _relayed_in(events) == (), events + assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _bedrock_passthrough_deployment(scenario: Scenario, rig: Rig) -> str: + return scenario.model( + model=f"bedrock/{BEDROCK_MODEL}", + api_base=rig.wire.url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=rig.wire.url, + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["invoke", "invoke-with-response-stream"]) +def test_bedrock_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None: + with rig.gateway.scenario() as scenario: + deployment: Final = _bedrock_passthrough_deployment(scenario, rig) + key: Final = scenario.key() + action: Final = "invoke-with-response-stream" if stream else "invoke" + body: Final = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + content: Final = _post(rig, f"/bedrock/model/{deployment}/{action}", body, key) + if stream: + assert content == b"".join(aws_chunks(refusal_events(DOCS_REFUSAL))), content + else: + assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content + assert _paths(rig.wire) == (f"/model/{BEDROCK_MODEL}/{action}",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +@pytest.mark.parametrize("stream", [False, True], ids=["rawPredict", "streamRawPredict"]) +def test_vertex_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + action: Final = "streamRawPredict" if stream else "rawPredict" + body: Final = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + content: Final = _post(rig, f"/vertex_ai{_VERTEX_MODEL_PATH}:{action}", body, key) + if stream: + assert content == b"".join(sse_chunks(refusal_events(DOCS_REFUSAL))), content + else: + assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content + assert _paths(rig.vertex) == (f"{_VERTEX_MODEL_PATH}:{action}",) + assert _targets(rig.wire) == () + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _surface_call(rig: Rig, surface: Surface, key: str, model: str, text: str) -> JsonValue: + match surface: + case "chat": + return _chat(rig, key, model, text, cached=True) + case "chat_stream": + return list(_chat_stream(rig, key, model, text, cached=True)) + case "messages": + return _messages(rig, key, model, text, cached=True) + case "messages_stream": + return list(_messages_stream(rig, key, model, text, cached=True)) + case "responses": + return _responses(rig, key, model, text, cached=True) + case "responses_stream": + return list(_responses_stream(rig, key, model, text)) + + +@dataclass(frozen=True, slots=True) +class CacheTwin: + fields: tuple[tuple[Location, Event], ...] + served_from_cache: tuple[bool, bool] + + +_CACHE_TWINS: Final[Mapping[Surface, CacheTwin]] = { + "chat": CacheTwin(((_CHAT_FIELDS, REFUSAL_FIELDS),), (False, True)), + "chat_stream": CacheTwin(((_DELTA_FIELDS, REFUSAL_FIELDS),), (False, True)), + "messages": CacheTwin((), (False, False)), + "messages_stream": CacheTwin((), (False, False)), +} + + +def _relayed_anywhere(value: JsonValue) -> tuple[tuple[Location, Event], ...]: + return _relayed_in(JSON_LIST.validate_python(value)) if isinstance(value, list) else relayed(value) + + +@pytest.mark.parametrize("surface", tuple(_CACHE_TWINS)) +def test_cache_hit_twin_carries_the_same_stop_details_as_the_miss(rig: Rig, surface: Surface) -> None: + twin: Final = _CACHE_TWINS[surface] + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("refusal", _marker()) + first: Final = _surface_call(rig, surface, key, model, text) + _logged(rig, key) + second: Final = _surface_call(rig, surface, key, model, text) + logged: Final = _logged(rig, key, 2) + upstream: Final = _paths(rig.wire) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "upstream": upstream, "first": first, "second": second} + assert (served_from_cache, len(upstream)) == (twin.served_from_cache, twin.served_from_cache.count(False)), ( + evidence + ) + assert (_relayed_anywhere(first), _relayed_anywhere(second)) == (twin.fields, twin.fields), evidence + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, ( + served_from_cache, + logged, + ) + + +def test_responses_cache_hit_logs_the_same_stop_details_as_the_miss(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("refusal", _marker()) + first: Final = _surface_call(rig, "responses", key, model, text) + _logged(rig, key) + _surface_call(rig, "responses", key, model, text) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + assert (len(_paths(rig.wire)), served_from_cache.count(True)) == (1, 1), (served_from_cache, logged) + assert relayed(first) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), first + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, ( + served_from_cache, + logged, + ) + + +def test_passthrough_is_not_served_from_the_response_cache(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + body: Final = { + "model": MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}], + } + first: Final = _post(rig, "/anthropic/v1/messages", body, key) + second: Final = _post(rig, "/anthropic/v1/messages", body, key) + assert JSON_OBJECT.validate_json(first) == JSON_OBJECT.validate_json(second) == DOCS_REFUSAL + assert _targets(rig.wire) == (f"{_ANTHROPIC_BASE}/v1/messages",) * 2 + logged: Final = _logged(rig, key, 2) + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2 + + +@pytest.mark.parametrize("variant", tuple(MALFORMED_STOP_DETAILS)) +def test_malformed_stop_details_are_relayed_as_sent_without_breaking_the_turn(rig: Rig, variant: str) -> None: + value: Final = MALFORMED_STOP_DETAILS[variant] + expected: Final[Event] = {"stop_reason": "end_turn", "stop_details": value} + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt(variant, _marker())) + assert relayed(body) == ((_CHAT_FIELDS, expected),), body + assert (_choice(body)["finish_reason"], _message_text(body)) == ("stop", "ready"), body + chunks: Final = _chat_stream(rig, key, model, refusal_prompt(variant, _marker())) + assert _relayed_in(chunks) == ((_DELTA_FIELDS, expected),), chunks + assert (_finish_reasons(chunks), _streamed_text(chunks)) == (("stop",), "ready"), chunks + assert len(_paths(rig.wire)) == 2 + logged: Final = _logged(rig, key, 2) + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, expected),),) * 2, logged + assert rig.gateway.client.get("/health/liveliness").status_code == 200 + + +def test_duplicated_stop_details_key_relays_the_docs_object_once(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("duplicate", _marker())) + assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def test_non_object_upstream_body_fails_the_call_and_leaves_other_keys_serving(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + broken_key: Final = scenario.key() + healthy_key: Final = scenario.key() + broken: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + _chat_request(model, refusal_prompt("nondict", _marker()), stream=False, cached=False), + key=broken_key, + ) + assert broken.status_code == 500, broken.text + assert "error" in JSON_OBJECT.validate_json(broken.content), broken.text + healthy: Final = _chat(rig, healthy_key, model, refusal_prompt("refusal", _marker())) + assert relayed(healthy) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), healthy + assert len(_paths(rig.wire)) == 2 + failure: Final = _logged(rig, broken_key)[0] + assert (failure["status"], relayed(failure["response"])) == ("failure", ()), failure + + +@pytest.mark.parametrize( + ("fixture", "status", "message"), + [("unauthorized", 401, "invalid x-api-key"), ("invalid", 400, "max_tokens: Field required")], +) +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +def test_provider_errors_reach_the_caller_and_log_a_failure_without_stop_fields( + rig: Rig, fixture: str, status: int, message: str, stream: bool +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + _chat_request(model, refusal_prompt(fixture, _marker()), stream=stream, cached=False), + key=key, + ) + assert response.status_code == status, response.text + assert message in response.text, response.text + assert _paths(rig.wire) == ("/v1/messages",) + failure: Final = _logged(rig, key)[0] + assert (failure["status"], relayed(failure["response"])) == ("failure", ()), failure + + +def test_failing_sinks_get_one_attempt_each_and_the_proxy_stays_healthy(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker())) + assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body + for ledger in (rig.sink, rig.forbidden, rig.absent): + events = ledger.await_key(key) + assert tuple(relayed(event["response"]) for event in events) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) + held( + lambda: tuple(len(ledger.for_key(key)) for ledger in (rig.sink, rig.forbidden, rig.absent)), + lambda counts: counts == (1, 1, 1), + holding=3, + ) + readiness: Final = rig.gateway.client.get("/health/readiness") + assert (readiness.status_code, readiness.json()["status"]) == (200, "healthy"), readiness.text + + +def test_unauthenticated_refusal_request_never_reaches_the_upstream_or_the_sink(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + stranger: Final = "sk-not-a-virtual-key-" + _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + _chat_request(model, refusal_prompt("refusal", _marker()), stream=False, cached=False), + key=stranger, + ) + assert response.status_code == 401, response.text + assert _targets(rig.wire) == () + assert rig.sink.for_key(stranger) == () + + +@pytest.mark.parametrize("fixture", ["null", "missing"]) +def test_normal_turn_adds_no_stop_fields_on_any_surface(rig: Rig, fixture: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt(fixture, _marker())) + assert relayed(body) == (), body + assert (_choice(body)["finish_reason"], _message_text(body)) == ("stop", "ready"), body + chunks: Final = _chat_stream(rig, key, model, refusal_prompt(fixture, _marker())) + assert _relayed_in(chunks) == (), chunks + assert (_finish_reasons(chunks), _streamed_text(chunks)) == (("stop",), "ready"), chunks + events: Final = _messages_stream(rig, key, model, refusal_prompt(fixture, _marker())) + expected_native: Final = {"stop_reason": "end_turn", "stop_details": None} + assert _native_stops(events) == (expected_native,), events + assert len(_paths(rig.wire)) == 3 + logged: Final = _logged(rig, key, 3) + assert tuple(relayed(event["response"]) for event in logged) == ((), (), ()), logged + + +def test_refusal_after_partial_text_keeps_the_text_and_relays_the_stop_details(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + chat_key: Final = scenario.key() + messages_key: Final = scenario.key() + chunks: Final = _chat_stream(rig, chat_key, model, refusal_prompt("partial", _marker())) + assert _streamed_text(chunks) == "Hello..", chunks + assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks + assert _finish_reasons(chunks) == ("content_filter",), chunks + logged_chat: Final = _logged_response(rig, chat_key) + assert relayed(logged_chat) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged_chat + assert _message_text(JSON_OBJECT.validate_python(logged_chat)) == "Hello..", logged_chat + events: Final = _messages_stream(rig, messages_key, model, refusal_prompt("partial", _marker())) + assert _native_stops(events) == (REFUSAL_FIELDS,), events + assert frame_types(events) == frame_types(FIXTURES["partial"].events), events + logged_messages: Final = _logged_response(rig, messages_key) + assert relayed(logged_messages) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged_messages + assert len(_paths(rig.wire)) == 2 + + +def test_repeated_identical_refusals_log_once_per_request(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("refusal", _marker()) + bodies: Final = (_chat(rig, key, model, text), _chat(rig, key, model, text)) + assert tuple(relayed(body) for body in bodies) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, bodies + assert len(_paths(rig.wire)) == 2 + logged: Final = _logged(rig, key, 2) + assert len({event["litellm_call_id"] for event in logged}) == 2, logged + assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2 + + +@pytest.mark.parametrize("fixture", ["refusal", "int", "null"]) +def test_chat_stream_logged_response_keeps_the_received_model_id(rig: Rig, fixture: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt(fixture, _marker())) + assert len({chunk["id"] for chunk in chunks}) == 1, chunks + probed: Final = rig.hidden.await_key(key) + assert len(probed) == 1, probed + assert (probed[0]["response_id"], probed[0]["received_model_id"]) == (chunks[0]["id"], chunks[0]["id"]), probed + + +def _delta_fields(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + deltas: Final = tuple(JSON_OBJECT.validate_python(_choice(chunk).get("delta") or {}) for chunk in chunks) + return tuple( + delta["provider_specific_fields"] for delta in deltas if delta.get("provider_specific_fields") is not None + ) + + +def _message_fields(response: JsonValue) -> tuple[JsonValue, ...]: + message: Final = JSON_OBJECT.validate_python(_choice(JSON_OBJECT.validate_python(response))["message"]) + fields: Final = message.get("provider_specific_fields") + return () if fields is None else (fields,) + + +def test_openai_compatible_stream_logged_response_keeps_the_received_model_id(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + marker: Final = _marker() + chunks: Final = _chat_stream(rig, key, model, refusal_prompt("plain", marker)) + assert (_streamed_text(chunks), _finish_reasons(chunks), _delta_fields(chunks)) == ("ready now", ("stop",), ()) + assert {chunk["id"] for chunk in chunks} == {openai_id(marker)}, chunks + assert _paths(rig.wire) == ("/v1/chat/completions",) + probed: Final = rig.hidden.await_key(key) + assert len(probed) == 1, probed + assert (probed[0]["response_id"], probed[0]["received_model_id"]) == (openai_id(marker),) * 2, probed + + +@pytest.mark.parametrize(("fixture", "fields"), [("plain", ()), ("cited", (CITED_FIELDS,))], ids=["plain", "cited"]) +def test_openai_compatible_stream_cache_hit_replays_the_provider_fields_of_the_miss( + rig: Rig, fixture: str, fields: tuple[JsonValue, ...] +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + text: Final = refusal_prompt(fixture, _marker()) + first: Final = _chat_stream(rig, key, model, text, cached=True) + _logged(rig, key) + second: Final = _chat_stream(rig, key, model, text, cached=True) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence + streamed: Final = tuple( + (_streamed_text(chunks), _finish_reasons(chunks), _delta_fields(chunks)) for chunks in (first, second) + ) + assert streamed == (("ready now", ("stop",), fields),) * 2, evidence + assert tuple(_message_fields(event["response"]) for event in logged) == (fields, fields), evidence + + +def _spend_rows(rig: Rig, key: str) -> tuple[Event, ...]: + response: Final = rig.gateway.request("GET", "/spend/logs", params={"api_key": key}) + assert response.status_code == 200, response.text + return tuple(JSON_OBJECT.validate_python(row) for row in JSON_LIST.validate_json(response.content)) + + +def test_spend_logs_endpoint_reports_the_stop_details_of_a_refusal(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker())) + rows: Final = eventually(lambda: _spend_rows(rig, key), lambda found: len(found) >= 1, seconds=70) + assert len(rows) == 1, rows + reported: Final = (rows[0]["request_id"], relayed(rows[0]["response"])) + expected: Final = (body["id"], ((_CHAT_FIELDS, REFUSAL_FIELDS),)) + assert ((body["id"], relayed(body)), reported) == (expected, expected), rows + + +def _usage_counts(usage: JsonValue, total: str, details: str) -> tuple[JsonValue, JsonValue]: + counted: Final = JSON_OBJECT.validate_python(usage) + return (counted[total], JSON_OBJECT.validate_python(counted[details])["reasoning_tokens"]) + + +def _chat_usage_counts(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]: + return tuple( + _usage_counts(chunk["usage"], "completion_tokens", "completion_tokens_details") + for chunk in chunks + if chunk.get("usage") + ) + + +def _cached_stream_usage_counts( + rig: Rig, surface: Literal["chat_stream", "responses_stream"], key: str, model: str, text: str +) -> tuple[tuple[JsonValue, JsonValue], ...]: + match surface: + case "chat_stream": + request: Final = { + **_chat_request(model, text, stream=True, cached=True), + "stream_options": {"include_usage": True}, + } + return _chat_usage_counts(chunks_of(_post(rig, "/v1/chat/completions", request, key).decode())) + case "responses_stream": + events: Final = frames( + _post(rig, "/v1/responses", _responses_request(model, text, stream=True, cached=True), key).decode() + ) + responses: Final = tuple( + JSON_OBJECT.validate_python(event["response"]) + for event in events + if isinstance(event.get("response"), dict) + ) + return tuple( + _usage_counts(response["usage"], "output_tokens", "output_tokens_details") + for response in responses + if response.get("usage") + ) + + +def _logged_usage_counts(logged: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]: + return tuple( + _usage_counts( + JSON_OBJECT.validate_python(event["response"])["usage"], "completion_tokens", "completion_tokens_details" + ) + for event in logged + ) + + +@pytest.mark.parametrize("fixture", ["plainusage", "citedusage"], ids=["plain", "cited"]) +def test_openai_compatible_stream_cache_hit_replays_the_usage_details_of_the_miss(rig: Rig, fixture: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + text: Final = refusal_prompt(fixture, _marker()) + first: Final = _cached_stream_usage_counts(rig, "chat_stream", key, model, text) + _logged(rig, key) + second: Final = _cached_stream_usage_counts(rig, "chat_stream", key, model, text) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence + assert (first, second, _logged_usage_counts(logged)) == (((7, 4),), ((7, 4),), ((7, 4), (7, 4))), evidence + + +def test_text_completion_stream_cache_hit_replays_the_usage_details_of_the_miss(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + request: Final[Event] = { + "model": model, + "prompt": refusal_prompt("plainusage", _marker()), + "max_tokens": 64, + "stream": True, + "stream_options": {"include_usage": True}, + } + first: Final = _chat_usage_counts(chunks_of(_post(rig, "/v1/completions", request, key).decode())) + _logged(rig, key) + second: Final = _chat_usage_counts(chunks_of(_post(rig, "/v1/completions", request, key).decode())) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence + assert (first, second, _logged_usage_counts(logged)) == (((7, 4),), ((7, 4),), ((7, 4), (7, 4))), evidence + + +@pytest.mark.parametrize("surface", ["chat_stream", "responses_stream"]) +def test_anthropic_stream_cache_hit_keeps_the_usage_details_of_the_miss( + rig: Rig, surface: Literal["chat_stream", "responses_stream"] +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + text: Final = refusal_prompt("reasoned", _marker()) + first: Final = _cached_stream_usage_counts(rig, surface, key, model, text) + _logged(rig, key) + second: Final = _cached_stream_usage_counts(rig, surface, key, model, text) + logged: Final = _logged(rig, key, 2) + served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) + evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged} + assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/messages",)), evidence + assert (first, second, _logged_usage_counts(logged)) == (((10, 4),), ((10, 4),), ((10, 4), (10, 4))), evidence + + +def _response_id(event: Mapping[str, JsonValue]) -> JsonValue: + return JSON_OBJECT.validate_python(event["response"]).get("id") + + +@pytest.mark.parametrize("variant", tuple(MALFORMED_STOP_DETAILS)) +def test_stream_cache_hit_skips_cached_provider_fields_that_are_not_an_object(rig: Rig, variant: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + marker: Final = _marker() + text: Final = refusal_prompt("plain", marker) + filled: Final = chunks_of( + _post( + rig, + "/v1/chat/completions", + {**_chat_request(model, text, stream=True, cached=True), "cache": {"ttl": 2}}, + key, + ).decode() + ) + assert _streamed_text(filled) == "ready now", filled + _logged(rig, key) + stale_id: Final = rewrite_cached_message_fields(openai_id(marker), MALFORMED_STOP_DETAILS[variant]) + replayed: Final = eventually( + lambda: _chat_stream(rig, key, model, text, cached=True), + lambda chunks: {chunk["id"] for chunk in chunks} == {stale_id}, + seconds=30, + ) + streamed: Final = (_streamed_text(replayed), _finish_reasons(replayed), _delta_fields(replayed)) + assert streamed == ("ready now", ("stop",), ()), replayed + assert _paths(rig.wire) == ("/v1/chat/completions",) + stale_logs: Final = eventually( + lambda: tuple(event for event in rig.sink.for_key(key) if _response_id(event) == stale_id), + lambda found: len(found) >= 1, + ) + assert tuple(_message_fields(event["response"]) for event in stale_logs) == ((),) * len(stale_logs), stale_logs + + +JsonModeRoute: TypeAlias = Literal["anthropic_native", "anthropic_tool", "bedrock_invoke"] +_UNMAPPED_MODEL: Final = "claude-scripted-9" +_JSON_ANSWER: Final = json.dumps({"answer": DOCS_PARTIAL_TEXT}) + + +@dataclass(frozen=True, slots=True) +class JsonModeCase: + route: JsonModeRoute + fixture: str + streamed_text: str + message_text: str + streamed_finish_reason: str + message_finish_reason: str + requested: Literal["json_tool", "native"] + + +_JSON_MODE_CASES: Final[Mapping[str, JsonModeCase]] = { + "anthropic_native": JsonModeCase( + "anthropic_native", + "partial", + DOCS_PARTIAL_TEXT, + DOCS_PARTIAL_TEXT, + "content_filter", + "content_filter", + "native", + ), + "anthropic_tool_declared": JsonModeCase( + "anthropic_tool", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), + "anthropic_tool_enveloped": JsonModeCase( + "anthropic_tool", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), + "bedrock_invoke_declared": JsonModeCase( + "bedrock_invoke", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), + "bedrock_invoke_enveloped": JsonModeCase( + "bedrock_invoke", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool" + ), +} + + +def _json_mode_deployment(scenario: Scenario, route: JsonModeRoute, rig: Rig) -> tuple[str, Provider]: + match route: + case "anthropic_native": + return _deployment(scenario, "anthropic", rig), "anthropic" + case "anthropic_tool": + return ( + scenario.model(model=f"anthropic/{_UNMAPPED_MODEL}", api_base=rig.wire.url, api_key=_ANTHROPIC_KEY), + "anthropic", + ) + case "bedrock_invoke": + return _deployment(scenario, "bedrock_invoke", rig), "bedrock_invoke" + + +def _requested_format(body: Mapping[str, JsonValue]) -> tuple[bool, bool]: + tools: Final = JSON_LIST.validate_python(body.get("tools") or []) + output_config: Final = JSON_OBJECT.validate_python(body.get("output_config") or {}) + json_tool: Final = any(JSON_OBJECT.validate_python(tool).get("name") == JSON_TOOL_NAME for tool in tools) + return json_tool, "output_format" in body or "format" in output_config + + +def _received(wire: Wire) -> tuple[tuple[str, tuple[bool, bool]], ...]: + return tuple( + (urlsplit(unquote(request.target)).path, _requested_format(JSON_OBJECT.validate_json(request.body))) + for request in wire.drain() + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +@pytest.mark.parametrize("case", tuple(_JSON_MODE_CASES)) +def test_json_mode_refusal_relays_the_docs_stop_details_through_the_response_format_path( + rig: Rig, case: str, stream: bool +) -> None: + expected: Final = _JSON_MODE_CASES[case] + with rig.gateway.scenario() as scenario: + model, provider = _json_mode_deployment(scenario, expected.route, rig) + key: Final = scenario.key() + request: Final = { + **_chat_request(model, refusal_prompt(expected.fixture, _marker()), stream=stream, cached=False), + "response_format": JSON_ANSWER_FORMAT, + } + content: Final = _post(rig, "/v1/chat/completions", request, key) + requested: Final = (expected.requested == "json_tool", expected.requested == "native") + assert _received(rig.wire) == ((_STREAM_TARGETS[provider][int(stream)], requested),) + if stream: + chunks: Final = chunks_of(content.decode()) + streamed: Final = (_streamed_text(chunks), _finish_reasons(chunks), _relayed_in(chunks)) + assert streamed == ( + expected.streamed_text, + (expected.streamed_finish_reason,), + ((_DELTA_FIELDS, REFUSAL_FIELDS),), + ), chunks + else: + body: Final = JSON_OBJECT.validate_json(content) + answered: Final = (_message_text(body), _choice(body)["finish_reason"], relayed(body)) + assert answered == ( + expected.message_text, + expected.message_finish_reason, + ((_CHAT_FIELDS, REFUSAL_FIELDS),), + ), body + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _serialized_choice(body: Mapping[str, JsonValue]) -> Event: + choice: Final = _choice(body) + return {**choice, "message": JSON_OBJECT.validate_python(choice["message"])} + + +_ANTHROPIC_MESSAGE_FIELDS: Final[Event] = {"citations": None, "thinking_blocks": None} + + +@pytest.mark.parametrize(("fixture", "content"), [("refusal", None), ("partial", DOCS_PARTIAL_TEXT)]) +def test_chat_refusal_body_serializes_the_stop_details_beside_the_existing_message_fields( + rig: Rig, fixture: str, content: str | None +) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "anthropic", rig) + key: Final = scenario.key() + body: Final = _chat(rig, key, model, refusal_prompt(fixture, _marker())) + expected: Final[Event] = { + "finish_reason": "content_filter", + "index": 0, + "message": { + "content": content, + "role": "assistant", + "provider_specific_fields": {**_ANTHROPIC_MESSAGE_FIELDS, **REFUSAL_FIELDS}, + }, + } + assert _serialized_choice(body) == expected, body + assert _paths(rig.wire) == ("/v1/messages",) + logged: Final = _logged_response(rig, key) + assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged + + +def _usage_of(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]: + usages: Final = tuple(JSON_OBJECT.validate_python(chunk["usage"]) for chunk in chunks if chunk.get("usage")) + return tuple((usage["completion_tokens"], usage.get("completion_tokens_details")) for usage in usages) + + +def _served_ids(response: httpx.Response) -> frozenset[JsonValue]: + return frozenset(chunk["id"] for chunk in chunks_of(response.text)) if response.status_code == 200 else frozenset() + + +@pytest.mark.parametrize("variant", tuple(UNVALIDATED_USAGE_DETAILS)) +def test_stream_cache_hit_skips_cached_usage_details_that_fail_validation(rig: Rig, variant: str) -> None: + with rig.gateway.scenario() as scenario: + model: Final = _deployment(scenario, "openai_compatible", rig) + key: Final = scenario.key() + marker: Final = _marker() + request: Final = { + **_chat_request(model, refusal_prompt("plainusage", marker), stream=True, cached=True), + "stream_options": {"include_usage": True}, + } + filled: Final = chunks_of(_post(rig, "/v1/chat/completions", {**request, "cache": {"ttl": 2}}, key).decode()) + assert (_streamed_text(filled), _chat_usage_counts(filled)) == ("ready now", ((7, 4),)), filled + _logged(rig, key) + stale_id: Final = rewrite_cached_usage_details(openai_id(marker), UNVALIDATED_USAGE_DETAILS[variant]) + replayed: Final = eventually( + lambda: rig.gateway.request("POST", "/v1/chat/completions", request, key=key), + lambda response: response.status_code != 200 or _served_ids(response) == {stale_id}, + seconds=30, + ) + assert replayed.status_code == 200, replayed.text + chunks: Final = chunks_of(replayed.text) + streamed: Final = (_streamed_text(chunks), _finish_reasons(chunks), _usage_of(chunks)) + assert streamed == ("ready now", ("stop",), ((7, {"reasoning_tokens": 0, "text_tokens": 7}),)), chunks + assert _paths(rig.wire) == ("/v1/chat/completions",) + stale_logs: Final = eventually( + lambda: tuple(event for event in rig.sink.for_key(key) if _response_id(event) == stale_id), + lambda found: len(found) >= 1, + ) + assert tuple(event["status"] for event in stale_logs) == ("success",) * len(stale_logs), stale_logs diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py index bc0a6f31e74..27e1a819e06 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py @@ -15,10 +15,9 @@ import psutil import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows -from integration._support.process import owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from integration.providers._cache_control_marks_support import ( - ASK_LABEL, CITIES, POINTS, SYSTEM_LABEL, @@ -40,6 +39,8 @@ from integration.providers._cache_control_marks_support import ( ) from pydantic import JsonValue +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _SURFACES: Final = ("chat", "chat-stream", "chat-unmarked", "messages", "responses") _MODEL: Final = "capped-claude" @@ -76,8 +77,6 @@ def _surface_request(surface: str, marker: str) -> tuple[str, dict[str, JsonValu def _expected_labels(item: _Sent) -> list[str]: - if item.surface == "responses": - return [SYSTEM_LABEL, ASK_LABEL] if item.surface == "chat-unmarked": return [SYSTEM_LABEL, final_label(item.marker)] return client_marked() @@ -148,7 +147,6 @@ def _single_spend_row(item: _Sent) -> None: assert len(rows) == 1, (item.surface, item.call_id) -@pytest.mark.timeout(240) async def test_capped_burst_rides_out_a_provider_outage_and_logs_every_request_once( gateway: Gateway, tmp_path: Path ) -> None: @@ -183,7 +181,6 @@ async def test_capped_burst_rides_out_a_provider_outage_and_logs_every_request_o _single_spend_row(item) -@pytest.mark.timeout(240) async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_capped_requests( gateway: Gateway, tmp_path: Path ) -> None: @@ -227,7 +224,6 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_capped_reques _single_spend_row(item) -@pytest.mark.timeout(240) def test_yaml_auto_caching_stands_down_for_tool_call_marks_and_outranks_a_key_opt_out( gateway: Gateway, tmp_path: Path ) -> None: @@ -262,7 +258,7 @@ def test_yaml_auto_caching_stands_down_for_tool_call_marks_and_outranks_a_key_op ] -@pytest.mark.timeout(360) +@pytest.mark.timeout(4 * graceful_stop_seconds() + 120) async def test_proxy_restart_mid_burst_serves_capped_requests_after_the_reboot( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py index 21fee0554ea..8c698817ccf 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py @@ -20,7 +20,7 @@ from integration._support.client import ( string_value, ) from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from integration.providers._cache_control_marks_support import ( ANTHROPIC_MODEL, @@ -61,6 +61,7 @@ from pydantic import JsonValue, TypeAdapter from litellm.utils import get_prompt_cache_min_tokens +_OWNED_PROXY_TIMEOUT: Final = 2 * graceful_stop_seconds() + 120 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _LANE_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _CLIENT_MARKED: Final = client_marked() @@ -505,17 +506,17 @@ def test_messages_endpoint_points_skip_injection_when_tool_use_marks_fill_the_ca assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED -def test_responses_bridge_keeps_system_and_user_marks_within_the_cap(gateway: Gateway) -> None: +def test_responses_bridge_points_skip_injection_when_function_call_marks_fill_the_cap(gateway: Gateway) -> None: marker: Final = new_marker() with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: model: Final = _anthropic_deployment(scenario, wire, cache_control_injection_points=POINTS) response: Final = gateway.request("POST", "/v1/responses", responses_body(model, marker)) assert response.status_code == 200, response.text assert _JSON_OBJECT.validate_json(response.content)["status"] == "completed", response.text - assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL, ASK_LABEL] + assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED -@pytest.mark.timeout(240) +@pytest.mark.timeout(_OWNED_PROXY_TIMEOUT) def test_response_cache_serves_the_capped_request_once(uncapped_gateway: Gateway) -> None: marker: Final = new_marker() with wire_server(anthropic_peer) as wire, uncapped_gateway.scenario() as scenario: @@ -567,14 +568,16 @@ def test_assistant_without_tool_calls_keeps_configured_points(gateway: Gateway, assert labels == [SYSTEM_LABEL, ASK_LABEL, final_label(marker)] -def test_responses_stream_bridge_keeps_system_and_user_marks_within_the_cap(gateway: Gateway) -> None: +def test_responses_stream_bridge_points_skip_injection_when_function_call_marks_fill_the_cap( + gateway: Gateway, +) -> None: marker: Final = new_marker() with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: model: Final = _anthropic_deployment(scenario, wire, cache_control_injection_points=POINTS) status, text = _stream(gateway, "/v1/responses", {**responses_body(model, marker), "stream": True}) assert status == 200, text assert '"type":"response.completed"' in text, text - assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL, ASK_LABEL] + assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED def _openai_client(gateway: Gateway) -> openai.OpenAI: @@ -698,7 +701,9 @@ async def test_anthropic_sdk_async_messages_stream_keeps_the_capped_request_with assert gateway_injected(f"msg_{marker}") is False -async def test_openai_sdk_async_responses_keeps_system_and_user_marks_within_the_cap(gateway: Gateway) -> None: +async def test_openai_sdk_async_responses_points_skip_injection_when_function_call_marks_fill_the_cap( + gateway: Gateway, +) -> None: marker: Final = new_marker() with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: model: Final = _anthropic_deployment(scenario, wire, cache_control_injection_points=POINTS) @@ -712,7 +717,7 @@ async def test_openai_sdk_async_responses_keeps_system_and_user_marks_within_the input=cast(ResponseInputParam, body["input"]), ) assert response.status == "completed", response.model_dump() - assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL, ASK_LABEL] + assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED def test_request_level_points_skip_injection_when_client_marks_fill_the_cap(gateway: Gateway) -> None: @@ -859,7 +864,7 @@ def _cache_hit_rows(response_id: str) -> list[dict[str, JsonValue]]: ) -@pytest.mark.timeout(240) +@pytest.mark.timeout(_OWNED_PROXY_TIMEOUT) def test_response_cache_hit_records_the_injection_on_the_first_row_only(uncapped_gateway: Gateway) -> None: marker: Final = new_marker() unmarked: Final = [tool_call(city) for city in CITIES] @@ -877,7 +882,7 @@ def test_response_cache_hit_records_the_injection_on_the_first_row_only(uncapped assert hit_rows[0]["injected"] is None, hit_rows -@pytest.mark.timeout(240) +@pytest.mark.timeout(_OWNED_PROXY_TIMEOUT) @pytest.mark.parametrize("path", ("/v1/messages", "/v1/responses")) def test_response_cache_twins_on_messages_and_responses_stay_within_the_cap( uncapped_gateway: Gateway, path: str @@ -891,8 +896,7 @@ def test_response_cache_twins_on_messages_and_responses_stay_within_the_cap( responses: Final = tuple(uncapped_gateway.request("POST", path, body) for _ in range(2)) received: Final = wire.drain() assert [response.status_code for response in responses] == [200, 200], [response.text for response in responses] - expected: Final = _CLIENT_MARKED if path == "/v1/messages" else [SYSTEM_LABEL, ASK_LABEL] - assert [anthropic_labels(request) for request in received] == [expected] * len(received) + assert [anthropic_labels(request) for request in received] == [_CLIENT_MARKED] * len(received) assert len(received) == 1, [request.target for request in received] diff --git a/tests/integration/providers/test_chatgpt_responses_top_level_error_wire.py b/tests/integration/providers/test_chatgpt_responses_top_level_error_wire.py new file mode 100644 index 00000000000..7bef4ad6b3a --- /dev/null +++ b/tests/integration/providers/test_chatgpt_responses_top_level_error_wire.py @@ -0,0 +1,85 @@ +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest +from integration._support import codex_vendor as cv +from integration._support.client import Gateway, gateway_from_environment, object_value, string_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_stream import ( + RATE_LIMIT_MESSAGE, + RATE_LIMIT_PREFIX, + data_frames, + error_body, + openai_error_stream, +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL: Final = "chatgpt/gpt-6.1-sol" +_TOP_LEVEL_RATE_LIMIT: Final = "codex-top-level-rate-limit" +_VENDOR: Final = cv.CodexVendor() + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + wire: Wire + + +def _respond(request: Request) -> Reply: + if _TOP_LEVEL_RATE_LIMIT.encode() not in request.body: + return _VENDOR.respond(request) + stream: Final = openai_error_stream("resp_" + uuid.uuid4().hex, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + return Reply(content_type="text/event-stream", chunks=stream) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("chatgpt-top-level-error-rig") + with gateway_from_environment() as gateway, wire_server(_respond) as wire: + overrides: Final = {"CHATGPT_TOKEN_DIR": str(cv.login(directory)), "CHATGPT_API_BASE": wire.url} + config: Final = cv.proxy_config(directory, probe=False) + with owned_proxy_process(gateway, directory, overrides, config=config, workers=2) as owned: + yield _Rig(owned.gateway, wire) + + +def _rate_limited(rig: _Rig, *, stream: bool) -> httpx.Response: + rig.wire.drain() + marker: Final = uuid.uuid4().hex + with rig.gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=rig.wire.url, api_key=None) + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "input": [{"role": "user", "content": f"Say marker-{marker} {_TOP_LEVEL_RATE_LIMIT}"}], + "stream": stream, + "cache": {"no-cache": True}, + } + response: Final = rig.gateway.request("POST", "/v1/responses", body) + (request,) = rig.wire.drain() + cv.forwarded(request, marker) + return response + + +def test_chatgpt_stream_top_level_rate_limit_fails_with_the_provider_code_and_message(rig: _Rig) -> None: + response: Final = _rate_limited(rig, stream=True) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.headers + frames: Final = data_frames(response.text) + assert [frame["type"] for frame in frames] == ["response.created", "response.failed"], response.text + error: Final = object_value(object_value(frames[-1]["response"])["error"]) + assert error == {"code": "rate_limit_exceeded", "message": RATE_LIMIT_MESSAGE}, error + + +def test_chatgpt_aggregated_top_level_rate_limit_answers_429_with_the_provider_message(rig: _Rig) -> None: + response: Final = _rate_limited(rig, stream=False) + assert response.status_code == 429, (response.status_code, response.text) + assert response.headers["content-type"].startswith("application/json"), response.headers + message: Final = string_value(error_body(response)["message"]) + assert message.startswith(RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message + assert "event:" not in message, message + assert rig.gateway.request("GET", "/health/liveliness", None).status_code == 200 diff --git a/tests/integration/providers/test_openai_responses_websocket_wire.py b/tests/integration/providers/test_openai_responses_websocket_wire.py index fb3674219f7..18be6bb9abb 100644 --- a/tests/integration/providers/test_openai_responses_websocket_wire.py +++ b/tests/integration/providers/test_openai_responses_websocket_wire.py @@ -1,30 +1,35 @@ import asyncio +import hashlib import itertools import json import ssl import threading import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue -from typing import Final +from typing import Final, TypeAlias import pytest import websockets -from integration._support.client import Gateway, gateway_from_environment -from integration._support.process import owned_proxy +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_stream import RATE_LIMIT_MESSAGE, openai_error_event from integration._support.tls import server_context, write_self_signed_cert from pydantic import JsonValue from websockets.asyncio.server import ServerConnection, serve -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) PROVIDER_MODEL: Final = "ws-peer-model" USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) +_Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] + @dataclass(frozen=True, slots=True) class Peer: @@ -58,8 +63,20 @@ def _events(response_id: str, text: str) -> tuple[dict[str, JsonValue], ...]: ) +def _seven(response_id: str) -> tuple[dict[str, JsonValue], ...]: + return _events(response_id, "seven") + + +def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]: + created: Final = _events(response_id, "")[0] + return (created, dict(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE))) + + async def _answer( - connection: ServerConnection, paths: SimpleQueue[str], frames: SimpleQueue[dict[str, JsonValue]] + connection: ServerConnection, + paths: SimpleQueue[str], + frames: SimpleQueue[dict[str, JsonValue]], + script: _Script, ) -> None: paths.put(connection.request.path if connection.request is not None else "") turns: Final = itertools.count(1) @@ -68,7 +85,7 @@ async def _answer( frames.put(frame) if frame.get("type") != "response.create": continue - for event in _events(f"resp_peer_{next(turns)}", "seven"): + for event in script(f"resp_peer_{next(turns)}"): await connection.send(json.dumps(event)) @@ -78,21 +95,24 @@ async def _serve( frames: SimpleQueue[dict[str, JsonValue]], ports: SimpleQueue[int], stop: asyncio.Event, + script: _Script, ) -> None: - async with serve(lambda connection: _answer(connection, paths, frames), "127.0.0.1", 0, ssl=tls) as server: + async with serve(lambda connection: _answer(connection, paths, frames, script), "127.0.0.1", 0, ssl=tls) as server: ports.put(next(iter(server.sockets)).getsockname()[1]) await stop.wait() @contextmanager -def responses_peer(cert: tuple[Path, Path]) -> Iterator[Peer]: +def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterator[Peer]: loop: Final = asyncio.new_event_loop() stop: Final = asyncio.Event() paths: Final = SimpleQueue[str]() frames: Final = SimpleQueue[dict[str, JsonValue]]() ports: Final = SimpleQueue[int]() thread: Final = threading.Thread( - target=loop.run_until_complete, args=(_serve(server_context(*cert), paths, frames, ports, stop),), daemon=True + target=loop.run_until_complete, + args=(_serve(server_context(*cert), paths, frames, ports, stop, script),), + daemon=True, ) thread.start() try: @@ -207,3 +227,29 @@ def test_previous_response_id_from_the_first_turn_reaches_the_provider_as_its_ow assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == list(texts) assert "previous_response_id" not in forwarded[0] assert forwarded[1]["previous_response_id"] == "resp_peer_1" + + +def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_and_message( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert, _top_level_rate_limit) as peer, candidate.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) + key: Final = scenario.key(models=[model]) + text: Final = f"say seven {uuid.uuid4().hex}" + (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (text,))) + assert [event["type"] for event in events] == ["response.created", "error"], events + assert (events[-1]["code"], events[-1]["message"]) == ("rate_limit_exceeded", RATE_LIMIT_MESSAGE), events + (forwarded,) = _drain(peer.frames) + assert forwarded["input"][0]["content"][0]["text"] == text + (row,) = eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (hashlib.sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + assert row["status"] == "failure", row + error: Final = object_value(object_value(row["metadata"])["error_information"]) + assert str(error["error_code"]) == "429", error + assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error diff --git a/tests/integration/providers/test_responses_bridge_item_marks_guardrail_wire.py b/tests/integration/providers/test_responses_bridge_item_marks_guardrail_wire.py new file mode 100644 index 00000000000..522ba8b1e81 --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_item_marks_guardrail_wire.py @@ -0,0 +1,229 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + ASK, + ASK_LABEL, + EPHEMERAL, + SYSTEM, + SYSTEM_LABEL, + Mark, + anthropic_deployment, + anthropic_labels, + anthropic_marks, + anthropic_peer, + final_text, + new_marker, +) +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GUARDED: Final = "bridge-mark-guarded" +_MODES: Final = ("none", "texts", "rows") +_CITY: Final = "Tokyo" +_REDACTED_ASK: Final = ASK.replace(_CITY, "[CITY]") +_REDACTED_ASK_LABEL: Final = f"user:text:{_REDACTED_ASK}" + + +@dataclass(frozen=True, slots=True) +class _GuardedRig: + gateway: Gateway + provider: Wire + guardrail: Wire + + +@dataclass(frozen=True, slots=True) +class _Exchange: + guardrail: tuple[dict[str, JsonValue], ...] + provider: tuple[Request, ...] + + +def _redacted(value: JsonValue) -> JsonValue: + return value.replace(_CITY, "[CITY]") if isinstance(value, str) else value + + +def _redacted_part(part: JsonValue) -> JsonValue: + if not isinstance(part, dict) or "text" not in part: + return part + return {**part, "text": _redacted(part["text"])} + + +def _redacted_row(row: JsonValue) -> JsonValue: + message: Final = object_value(row) + content: Final = message.get("content") + if isinstance(content, list): + return {**message, "content": [_redacted_part(part) for part in content]} + return {**message, "content": _redacted(content)} + + +def _listed(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value + + +def _verdict(mode: str, payload: dict[str, JsonValue]) -> dict[str, JsonValue]: + match mode: + case "texts": + return {"action": "GUARDRAIL_INTERVENED", "texts": [_redacted(text) for text in _listed(payload["texts"])]} + case "rows": + rows: Final = _listed(payload["structured_messages"]) + return {"action": "GUARDRAIL_INTERVENED", "structured_messages": [_redacted_row(row) for row in rows]} + case _: + return {"action": "NONE"} + + +def _guardrail_peer(request: Request) -> Reply: + mode: Final = request.target.split("/")[1] + return Reply(body=json.dumps(_verdict(mode, _JSON_OBJECT.validate_json(request.body))).encode()) + + +def _guardrail(mode: str, api_base: str) -> dict[str, JsonValue]: + return { + "guardrail_name": f"bridge-mark-{mode}", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{api_base}/{mode}", + "api_key": "synthetic-guardrail-key", + }, + } + + +def _guarded_config(directory: Path, provider_url: str, guardrail_url: str) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [anthropic_deployment(_GUARDED, provider_url)], + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "guardrails": [_guardrail(mode, guardrail_url) for mode in _MODES], + } + path: Final = directory / "bridge-mark-guarded.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def guarded(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_GuardedRig]: + directory: Final = tmp_path_factory.mktemp("bridge-mark-guarded") + with ( + wire_server(anthropic_peer) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _guarded_config(directory, provider.url, guardrail.url) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _GuardedRig(owned, provider, guardrail) + + +def _items(marker: str, mark: JsonValue = EPHEMERAL) -> list[JsonValue]: + return [ + {"role": "system", "content": SYSTEM, "cache_control": EPHEMERAL}, + {"role": "user", "content": ASK, "cache_control": mark}, + {"role": "user", "content": final_text(marker)}, + ] + + +def _send(rig: _GuardedRig, mode: str, items: Sequence[JsonValue], *, stream: bool = False) -> str: + body: Final[dict[str, JsonValue]] = { + "model": _GUARDED, + "max_output_tokens": 64, + "input": list(items), + "stream": stream, + "guardrails": [f"bridge-mark-{mode}"], + } + headers: Final = {"Authorization": f"Bearer {rig.gateway.key}"} + with rig.gateway.client.stream("POST", "/v1/responses", json=body, headers=headers) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + assert "sunny" in text, text + if stream: + assert '"type":"response.completed"' in text, text + else: + assert _JSON_OBJECT.validate_json(text)["status"] == "completed", text + return text + + +def _exchange(rig: _GuardedRig, marker: str) -> _Exchange: + needle: Final = marker.encode() + return _Exchange( + guardrail=tuple( + _JSON_OBJECT.validate_json(request.body) for request in rig.guardrail.drain() if needle in request.body + ), + provider=tuple(request for request in rig.provider.drain() if needle in request.body), + ) + + +def _assert_read_without_marks(payload: dict[str, JsonValue], marker: str) -> None: + assert "cache_control" not in json.dumps(payload), payload + assert ASK in _listed(payload["texts"]), payload["texts"] + assert final_text(marker) in _listed(payload["texts"]), payload["texts"] + rows: Final = [object_value(row) for row in _listed(payload["structured_messages"])] + assert [row.get("content") for row in rows] == [SYSTEM, ASK, final_text(marker)], rows + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_guardrail_reads_the_input_without_item_marks_while_the_provider_keeps_them( + guarded: _GuardedRig, stream: bool +) -> None: + marker: Final = new_marker() + _send(guarded, "none", _items(marker), stream=stream) + seen: Final = _exchange(guarded, marker) + assert len(seen.guardrail) == 1, seen.guardrail + _assert_read_without_marks(seen.guardrail[0], marker) + assert [anthropic_labels(request) for request in seen.provider] == [[SYSTEM_LABEL, ASK_LABEL]] + + +def test_guardrail_on_a_response_cache_hit_reads_the_input_without_item_marks(guarded: _GuardedRig) -> None: + marker: Final = new_marker() + items: Final = _items(marker) + for _ in range(2): + _send(guarded, "none", items) + seen: Final = _exchange(guarded, marker) + assert len(seen.guardrail) == 2, seen.guardrail + for payload in seen.guardrail: + _assert_read_without_marks(payload, marker) + assert [anthropic_labels(request) for request in seen.provider] == [[SYSTEM_LABEL, ASK_LABEL]] + + +def test_guardrail_text_rewrite_keeps_the_item_mark(guarded: _GuardedRig) -> None: + marker: Final = new_marker() + _send(guarded, "texts", _items(marker)) + seen: Final = _exchange(guarded, marker) + assert len(seen.guardrail) == 1, seen.guardrail + _assert_read_without_marks(seen.guardrail[0], marker) + assert [anthropic_labels(request) for request in seen.provider] == [[SYSTEM_LABEL, _REDACTED_ASK_LABEL]] + + +@pytest.mark.parametrize( + "mark", + ( + pytest.param(EPHEMERAL, id="ephemeral"), + pytest.param({"type": "ephemeral", "ttl": "1h"}, id="ttl-1h"), + pytest.param({"type": "ephemeral", "ttl": "30m"}, id="ttl-30m"), + pytest.param({"type": "persistent"}, id="unknown-type"), + ), +) +def test_guardrail_row_rewrite_keeps_the_item_mark_as_sent(guarded: _GuardedRig, mark: dict[str, JsonValue]) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + {"role": "user", "content": ASK, "cache_control": mark}, + {"role": "user", "content": final_text(marker)}, + ] + _send(guarded, "rows", items) + seen: Final = _exchange(guarded, marker) + assert len(seen.guardrail) == 1, seen.guardrail + assert "cache_control" not in json.dumps(seen.guardrail[0]), seen.guardrail[0] + ttl: Final = mark.get("ttl") + expected: Final = (Mark(_REDACTED_ASK_LABEL, ttl if isinstance(ttl, str) else None),) + assert [anthropic_marks(request) for request in seen.provider] == [expected] diff --git a/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py b/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py new file mode 100644 index 00000000000..ea55552c331 --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py @@ -0,0 +1,1219 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from itertools import chain +from typing import Final, cast + +import httpx +import openai +import pytest +from integration._support import bedrock_runtime_peer +from integration._support import prompt_cache_breakpoint as pcb +from integration._support.client import Gateway, Scenario, object_value +from integration._support.openai_wire import answering_model_discovery, chat_reply +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + ANTHROPIC_MODEL, + ASK, + ASK_LABEL, + BEDROCK_MODEL, + CITIES, + EPHEMERAL, + POINTS, + PROVIDER_KEY, + SYSTEM, + SYSTEM_LABEL, + UNMAPPED_BEDROCK_ARN, + anthropic_labels, + anthropic_marks, + anthropic_peer, + bedrock_labels, + bedrock_peer, + call_id, + client_marked, + final_label, + final_text, + marker_of, + new_marker, + responses_body, + sent_without_cache_points, + tool_use_label, +) +from openai.types.responses import ResponseInputParam +from pydantic import JsonValue, TypeAdapter + +from litellm.utils import get_prompt_cache_min_tokens + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_HOUR: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} +_CHECK: Final = "I will check." +_ASSISTANT_LABEL: Final = f"assistant:text:{_CHECK}" +_RESULT_LABEL: Final = f"user:tool_result:{call_id(CITIES[0])}" +_BEDROCK_SYSTEM_LABEL: Final = f"system:text:{SYSTEM}" +_BEDROCK_RESULT_LABEL: Final = f"user:toolResult:{call_id(CITIES[0])}" +_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "name": "lookup_weather", "parameters": {"type": "object"}} +_GEMINI: Final = "gemini/gemini-3.8-flash" +_GEMINI_GENERATE_ONLY: Final = [("POST", "/models/gemini-3.8-flash:generateContent")] +_GEMINI_REPLY: Final = json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "sunny"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1600, "candidatesTokenCount": 1, "totalTokenCount": 1601}, + } +).encode() +_OPENAI_COMPATIBLE: Final = "openai/gpt-5.4" + + +@dataclass(frozen=True, slots=True) +class _Answer: + status: int + text: str + call_id: str + + +@pytest.fixture(scope="module") +def spend() -> Iterator[pcb.SpendLogs]: + with pcb.spend_logs() as logs: + yield logs + + +def _item(role: str, text: str, mark: JsonValue = EPHEMERAL) -> dict[str, JsonValue]: + return {"role": role, "content": text, "cache_control": mark} + + +def _responses(model: str, items: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "max_output_tokens": 64, "input": list(items), **fields} + + +def _marked_system(model: str, marker: str, mark: JsonValue = EPHEMERAL, **fields: JsonValue) -> dict[str, JsonValue]: + return _responses(model, [_item("system", SYSTEM, mark), {"role": "user", "content": final_text(marker)}], **fields) + + +def _role_items(role: str, marker: str) -> list[JsonValue]: + match role: + case "system" | "developer": + return [_item(role, SYSTEM), {"role": "user", "content": final_text(marker)}] + case "user": + return [_item("user", ASK), {"role": "user", "content": final_text(marker)}] + case _: + return [ + {"role": "user", "content": ASK}, + _item("assistant", _CHECK), + {"role": "user", "content": final_text(marker)}, + ] + + +def _tool_items(marker: str, *, output_mark: JsonValue = None, call_mark: JsonValue = None) -> list[JsonValue]: + call: Final[dict[str, JsonValue]] = { + "type": "function_call", + "call_id": call_id(CITIES[0]), + "name": "lookup_weather", + "arguments": json.dumps({"city": CITIES[0]}), + } + output: Final[dict[str, JsonValue]] = { + "type": "function_call_output", + "call_id": call_id(CITIES[0]), + "output": "sunny.", + } + return [ + {"role": "user", "content": ASK}, + {**call, "cache_control": call_mark} if call_mark is not None else call, + {**output, "cache_control": output_mark} if output_mark is not None else output, + {"role": "user", "content": final_text(marker)}, + ] + + +def _call_id_of(response: httpx.Response) -> str: + return response.headers["x-litellm-call-id"] if "x-litellm-call-id" in response.headers else "" + + +def _send(gateway: Gateway, body: dict[str, JsonValue], *, key: str | None = None) -> _Answer: + headers: Final = {"Authorization": f"Bearer {gateway.key if key is None else key}"} + with gateway.client.stream("POST", "/v1/responses", json=body, headers=headers) as response: + return _Answer(response.status_code, response.read().decode(), _call_id_of(response)) + + +def _send_raw(gateway: Gateway, raw: bytes) -> _Answer: + headers: Final = {"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"} + response: Final = gateway.client.post("/v1/responses", content=raw, headers=headers) + return _Answer(response.status_code, response.text, _call_id_of(response)) + + +def _assert_completed(answer: _Answer, *, stream: bool = False) -> None: + assert answer.status == 200, answer.text + assert "sunny" in answer.text, answer.text + if stream: + assert '"type":"response.completed"' in answer.text, answer.text + return + assert _JSON_OBJECT.validate_json(answer.text)["status"] == "completed", answer.text + + +def _only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, [request.target for request in received] + return received[0] + + +def _anthropic_deployment(scenario: Scenario, wire: Wire, *, inject: bool = False) -> str: + if inject: + return scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", + api_base=wire.url, + api_key=PROVIDER_KEY, + cache_control_injection_points=POINTS, + ) + return scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY) + + +def _bedrock_peer(request: Request) -> Reply: + if request.target.endswith("/converse-stream"): + return bedrock_runtime_peer.respond(request) + return bedrock_peer(request) + + +def _bedrock_deployment(scenario: Scenario, wire: Wire, route: str = "converse") -> str: + return scenario.model( + model=f"bedrock/{route}/{BEDROCK_MODEL}", + api_key=PROVIDER_KEY, + api_base=wire.url, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +def _unmapped_arn_deployment(scenario: Scenario, wire: Wire, **fields: JsonValue) -> str: + return scenario.model( + model=f"bedrock/converse/{UNMAPPED_BEDROCK_ARN}", + api_key=PROVIDER_KEY, + api_base=wire.url, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + **fields, + ) + + +def _blocks_of(body: dict[str, JsonValue]) -> Iterator[JsonValue]: + system: Final = body.get("system") + yield from system if isinstance(system, list) else [] + messages: Final = body.get("messages") + for message in messages if isinstance(messages, list) else []: + yield from _content_blocks(object_value(message).get("content")) + + +def _content_blocks(content: JsonValue) -> Iterator[JsonValue]: + yield from content if isinstance(content, list) else [] + + +def _bedrock_ttls(request: Request) -> list[JsonValue]: + blocks: Final = list(_blocks_of(_JSON_OBJECT.validate_json(request.body))) + return [ + object_value(block["cachePoint"]).get("ttl") + for block in blocks + if isinstance(block, dict) and "cachePoint" in block + ] + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +@pytest.mark.parametrize( + ("role", "expected"), + ( + pytest.param("system", SYSTEM_LABEL, id="system"), + pytest.param("developer", SYSTEM_LABEL, id="developer"), + pytest.param("user", ASK_LABEL, id="user"), + pytest.param("assistant", _ASSISTANT_LABEL, id="assistant"), + ), +) +def test_marked_item_reaches_anthropic_as_a_cache_point( + gateway: Gateway, spend: pcb.SpendLogs, role: str, expected: str, stream: bool +) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + answer: Final = _send(gateway, _responses(model, _role_items(role, marker), stream=stream)) + _assert_completed(answer, stream=stream) + assert anthropic_labels(_only_request(wire)) == [expected] + spend.landed(model, answer.call_id, None) + + +def test_list_content_message_mark_lands_on_the_last_block(gateway: Gateway) -> None: + marker: Final = new_marker() + content: Final[list[JsonValue]] = [{"type": "input_text", "text": "first"}, {"type": "input_text", "text": ASK}] + items: Final[list[JsonValue]] = [ + {"role": "user", "content": content, "cache_control": EPHEMERAL}, + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _responses(model, items))) + assert anthropic_labels(_only_request(wire)) == [ASK_LABEL] + + +def test_block_mark_outranks_the_message_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + { + "role": "user", + "content": [{"type": "input_text", "text": ASK, "cache_control": _HOUR}], + "cache_control": EPHEMERAL, + }, + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _responses(model, items))) + marks: Final = anthropic_marks(_only_request(wire)) + assert [(mark.label, mark.ttl) for mark in marks] == [(ASK_LABEL, "1h")] + + +def test_function_call_output_mark_reaches_anthropic_as_a_tool_result_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = _responses(model, _tool_items(marker, output_mark=EPHEMERAL), tools=[_TOOL]) + _assert_completed(_send(gateway, body)) + assert anthropic_labels(_only_request(wire)) == [_RESULT_LABEL] + + +def test_function_call_marks_reach_anthropic_as_tool_use_marks(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, responses_body(model, marker))) + assert anthropic_labels(_only_request(wire)) == client_marked() + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +@pytest.mark.parametrize( + ("role", "expected"), + ( + pytest.param("system", _BEDROCK_SYSTEM_LABEL, id="system"), + pytest.param("user", f"user:text:{ASK}", id="user"), + pytest.param("assistant", f"assistant:text:{_CHECK}", id="assistant"), + ), +) +def test_marked_item_reaches_bedrock_converse_as_a_cache_point( + gateway: Gateway, role: str, expected: str, stream: bool +) -> None: + marker: Final = new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_deployment(scenario, wire) + answer: Final = _send(gateway, _responses(model, _role_items(role, marker), stream=stream)) + assert answer.status == 200, answer.text + assert f"marker-{marker}" in answer.text or "sunny" in answer.text, answer.text + request: Final = _only_request(wire) + assert request.target.endswith("/converse-stream" if stream else "/converse"), request.target + assert bedrock_labels(request) == [expected] + + +def test_function_call_output_mark_reaches_bedrock_as_a_tool_result_cache_point(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_deployment(scenario, wire) + body: Final = _responses(model, _tool_items(marker, output_mark=EPHEMERAL), tools=[_TOOL]) + answer: Final = _send(gateway, body) + assert answer.status == 200, answer.text + assert bedrock_labels(_only_request(wire)) == [_BEDROCK_RESULT_LABEL] + + +@pytest.mark.parametrize("drop_params", (False, True), ids=("as-sent", "drop-params")) +def test_nova_tool_marks_follow_drop_params(gateway: Gateway, drop_params: bool) -> None: + marker: Final = new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="bedrock/converse/us.amazon.nova-2-lite-v1:0", + api_key=PROVIDER_KEY, + api_base=wire.url, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + drop_params=drop_params, + ) + items: Final = _tool_items(marker, output_mark=EPHEMERAL, call_mark=EPHEMERAL) + answer: Final = _send(gateway, _responses(model, items, tools=[_TOOL])) + assert answer.status == 200, answer.text + labels: Final = bedrock_labels(_only_request(wire)) + assert (_BEDROCK_RESULT_LABEL in labels) is not drop_params, labels + assert (len(labels) == 2) is not drop_params, labels + + +def test_one_hour_mark_keeps_its_ttl_on_anthropic(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _marked_system(model, marker, _HOUR))) + marks: Final = anthropic_marks(_only_request(wire)) + assert [(mark.label, mark.ttl) for mark in marks] == [(SYSTEM_LABEL, "1h")] + + +def test_one_hour_mark_keeps_its_ttl_on_bedrock(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_deployment(scenario, wire) + answer: Final = _send(gateway, _marked_system(model, marker, _HOUR)) + assert answer.status == 200, answer.text + request: Final = _only_request(wire) + assert bedrock_labels(request) == [_BEDROCK_SYSTEM_LABEL] + assert _bedrock_ttls(request) == ["1h"] + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def _sdk_input(marker: str) -> ResponseInputParam: + return cast(ResponseInputParam, _role_items("system", marker)) + + +def test_openai_sdk_sync_responses_carries_the_system_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + with _openai_client(gateway) as client: + response: Final = client.responses.create(model=model, input=_sdk_input(marker), max_output_tokens=64) + assert response.status == "completed", response.model_dump() + assert response.output_text == "sunny", response.model_dump() + assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL] + + +async def test_openai_sdk_async_stream_responses_carries_the_system_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + async with _async_openai_client(gateway) as client: + stream: Final = await client.responses.create( + model=model, input=_sdk_input(marker), max_output_tokens=64, stream=True + ) + events: Final = [event async for event in stream] + assert [event.type for event in events][-1] == "response.completed", events + assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL] + + +def _gemini_peer(request: Request) -> Reply: + if "cachedContents" in request.target and request.method == "GET": + return Reply(body=b'{"cachedContents":[]}') + if "cachedContents" in request.target: + return Reply( + body=json.dumps( + { + "name": "cachedContents/bridge", + "model": "models/gemini-3.8-flash", + "expireTime": "2099-01-01T00:00:00Z", + } + ).encode() + ) + return Reply(body=_GEMINI_REPLY) + + +@pytest.mark.parametrize( + "fields", + ( + pytest.param({}, id="bare"), + pytest.param({"instructions": SYSTEM, "tools": [_TOOL]}, id="instructions-and-tools"), + ), +) +def test_gemini_marked_items_never_start_a_context_cache(gateway: Gateway, fields: dict[str, JsonValue]) -> None: + marker: Final = new_marker() + long_system: Final = "lorem " * (2 * get_prompt_cache_min_tokens(_GEMINI)) + items: Final[list[JsonValue]] = [ + _item("system", long_system), + _item("developer", SYSTEM), + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(_gemini_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_GEMINI, api_base=wire.url, api_key=PROVIDER_KEY) + _assert_completed(_send(gateway, _responses(model, items, **fields))) + targets: Final = [(request.method, request.target.split("?")[0]) for request in wire.drain()] + assert targets == _GEMINI_GENERATE_ONLY + + +def test_azure_ai_claude_marked_system_item_reaches_the_anthropic_route(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"azure_ai/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY) + _assert_completed(_send(gateway, _marked_system(model, marker))) + request: Final = _only_request(wire) + assert request.target == "/anthropic/v1/messages", request.target + assert anthropic_labels(request) == [SYSTEM_LABEL] + + +def test_bedrock_invoke_marked_system_item_reaches_the_invoke_body(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_deployment(scenario, wire, route="invoke") + _assert_completed(_send(gateway, _marked_system(model, marker))) + request: Final = _only_request(wire) + assert request.target.endswith("/invoke"), request.target + assert anthropic_labels(request) == [SYSTEM_LABEL] + + +def _openai_message_marks(request: Request) -> list[str]: + body: Final = _JSON_OBJECT.validate_json(request.body) + messages: Final = body["messages"] if isinstance(body["messages"], list) else [] + return [ + f"cache_control@{object_value(message).get('role')}:message" + for message in messages + if "cache_control" in object_value(message) + ] + + +def _openai_block_marks(request: Request) -> list[JsonValue]: + return [object_value(block).get("cache_control") for block in _blocks_of(_JSON_OBJECT.validate_json(request.body))] + + +def test_databricks_claude_keeps_the_item_mark_like_its_chat_route(gateway: Gateway) -> None: + marker: Final = new_marker() + reply: Final = chat_reply(f"chatcmpl-{marker}", "databricks-claude-sonnet-4", "sunny", stream=False) + with wire_server(answering_model_discovery(lambda request: reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="databricks/databricks-claude-sonnet-4", api_base=wire.url, api_key=PROVIDER_KEY + ) + _assert_completed(_send(gateway, _marked_system(model, marker))) + posted: Final = [request for request in wire.drain() if request.method == "POST"] + assert [request.target for request in posted] == ["/chat/completions"], [request.target for request in posted] + assert _openai_block_marks(posted[0]) == [EPHEMERAL] + + +def test_openai_compatible_chat_bridge_leaves_the_item_mark_off(gateway: Gateway) -> None: + marker: Final = new_marker() + reply: Final = chat_reply(f"chatcmpl-{marker}", "gpt-5.4", "sunny", stream=False) + with wire_server(answering_model_discovery(lambda request: reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_OPENAI_COMPATIBLE, api_base=f"{wire.url}/v1", api_key=PROVIDER_KEY, use_chat_completions_api=True + ) + _assert_completed(_send(gateway, _marked_system(model, marker))) + posted: Final = [request for request in wire.drain() if request.method == "POST"] + assert [request.target for request in posted] == ["/v1/chat/completions"], [request.target for request in posted] + assert _openai_message_marks(posted[0]) == [] + + +def test_injection_points_keep_the_caller_one_hour_system_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire, inject=True) + _assert_completed(_send(gateway, _marked_system(model, marker, _HOUR))) + marks: Final = anthropic_marks(_only_request(wire)) + assert [(mark.label, mark.ttl) for mark in marks] == [(SYSTEM_LABEL, "1h"), (final_label(marker), None)] + + +@pytest.mark.parametrize( + "mark", + ( + pytest.param({}, id="empty-object"), + pytest.param(None, id="null"), + pytest.param("", id="empty-string"), + pytest.param(0, id="zero"), + pytest.param(False, id="false"), + ), +) +def test_falsy_marks_are_dropped(gateway: Gateway, mark: JsonValue) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _marked_system(model, marker, mark))) + assert anthropic_labels(_only_request(wire)) == [] + + +def test_missing_mark_is_dropped(gateway: Gateway) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + {"role": "system", "content": SYSTEM}, + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _responses(model, items))) + assert anthropic_labels(_only_request(wire)) == [] + + +def test_response_cache_serves_the_marked_request_once(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = _marked_system(model, marker) + answers: Final = tuple(_send(gateway, body) for _ in range(2)) + received: Final = wire.drain() + for answer in answers: + _assert_completed(answer) + assert [anthropic_labels(request) for request in received] == [[SYSTEM_LABEL]] + + +def _chat_with_message_mark(model: str, marker: str, mark: JsonValue = EPHEMERAL) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "messages": [ + {"role": "system", "content": SYSTEM, "cache_control": mark}, + {"role": "user", "content": final_text(marker)}, + ], + } + + +def _messages_with_system_mark(model: str, marker: str) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "system": [{"type": "text", "text": SYSTEM, "cache_control": EPHEMERAL}], + "messages": [{"role": "user", "content": final_text(marker)}], + } + + +@pytest.mark.parametrize("path", ("/v1/chat/completions", "/v1/messages")) +def test_message_marks_on_the_other_endpoints_are_unchanged(gateway: Gateway, path: str) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = ( + _chat_with_message_mark(model, marker) + if path == "/v1/chat/completions" + else _messages_with_system_mark(model, marker) + ) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 200, response.text + assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL] + + +def test_block_mark_on_responses_input_is_unchanged(gateway: Gateway) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + {"role": "user", "content": [{"type": "input_text", "text": ASK, "cache_control": EPHEMERAL}]}, + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _responses(model, items, instructions=SYSTEM))) + assert anthropic_labels(_only_request(wire)) == [ASK_LABEL] + + +def _non_object_marks(request: Request) -> list[JsonValue]: + blocks: Final = list(_blocks_of(_JSON_OBJECT.validate_json(request.body))) + return [ + block["cache_control"] + for block in blocks + if isinstance(block, dict) and "cache_control" in block and not isinstance(block["cache_control"], dict) + ] + + +def _strict_anthropic_peer(request: Request) -> Reply: + odd: Final = _non_object_marks(request) + if odd: + return Reply( + status=400, + body=json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"cache_control: Input should be a valid dictionary, got {json.dumps(odd[0])[:64]}", + }, + } + ).encode(), + ) + return anthropic_peer(request) + + +@pytest.mark.parametrize( + "mark", + ( + pytest.param("ephemeral", id="string"), + pytest.param(1, id="int"), + pytest.param(["ephemeral"], id="list"), + pytest.param("x" * 5120, id="5kb-string"), + ), +) +def test_non_object_marks_reach_anthropic_as_they_do_on_chat_completions(gateway: Gateway, mark: JsonValue) -> None: + marker: Final = new_marker() + with wire_server(_strict_anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + on_chat: Final = gateway.request("POST", "/v1/chat/completions", _chat_with_message_mark(model, marker, mark)) + on_responses: Final = _send(gateway, _marked_system(model, marker, mark)) + received: Final = wire.drain() + assert (on_chat.status_code, on_responses.status) == (400, 400), (on_chat.text, on_responses.text) + assert [_non_object_marks(request) for request in received] == [[mark], [mark]] + + +def test_string_mark_reaches_bedrock_as_it_does_on_chat_completions(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_deployment(scenario, wire) + on_chat: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_with_message_mark(model, marker, "ephemeral") + ) + assert on_chat.status_code == 200, on_chat.text + on_responses: Final = _send(gateway, _marked_system(model, marker, "ephemeral")) + assert on_responses.status == 200, on_responses.text + received: Final = wire.drain() + assert len(received) == 2, [request.target for request in received] + assert bedrock_labels(received[1]) == bedrock_labels(received[0]) + + +def test_unknown_mark_type_is_forwarded_like_a_block_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _marked_system(model, marker, {"type": "persistent"}))) + assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL] + + +def test_duplicate_mark_keys_keep_the_last_one(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + raw: Final = ( + f'{{"model":"{model}","max_output_tokens":64,"input":[' + f'{{"role":"system","content":{json.dumps(SYSTEM)},"cache_control":{{"type":"ephemeral"}},' + f'"cache_control":{{"type":"ephemeral","ttl":"1h"}}}},' + f'{{"role":"user","content":{json.dumps(final_text(marker))}}}]}}' + ).encode() + _assert_completed(_send_raw(gateway, raw)) + marks: Final = anthropic_marks(_only_request(wire)) + assert [(mark.label, mark.ttl) for mark in marks] == [(SYSTEM_LABEL, "1h")] + + +def _many_marks(marker: str, *, assistant_marked: bool) -> list[JsonValue]: + assistant: Final[dict[str, JsonValue]] = {"role": "assistant", "content": _CHECK} + return [ + _item("system", SYSTEM), + _item("user", ASK), + {**assistant, "cache_control": EPHEMERAL} if assistant_marked else assistant, + { + "type": "function_call", + "call_id": call_id(CITIES[0]), + "name": "lookup_weather", + "arguments": json.dumps({"city": CITIES[0]}), + }, + { + "type": "function_call_output", + "call_id": call_id(CITIES[0]), + "output": "sunny.", + "cache_control": EPHEMERAL, + }, + _item("user", final_text(marker)), + ] + + +def test_four_marks_reach_anthropic(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = _responses(model, _many_marks(marker, assistant_marked=False), tools=[_TOOL]) + _assert_completed(_send(gateway, body)) + labels: Final = anthropic_labels(_only_request(wire)) + assert labels == [SYSTEM_LABEL, ASK_LABEL, _RESULT_LABEL, final_label(marker)] + + +def test_five_marks_are_refused_by_the_vendor(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = _responses(model, _many_marks(marker, assistant_marked=True), tools=[_TOOL]) + answer: Final = _send(gateway, body) + labels: Final = anthropic_labels(_only_request(wire)) + assert answer.status == 400, answer.text + assert "A maximum of 4 blocks with cache_control" in answer.text, answer.text + assert len(labels) == 5, labels + + +def test_five_marks_keep_the_newest_four_under_drop_params(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY, drop_params=True + ) + body: Final = _responses(model, _many_marks(marker, assistant_marked=True), tools=[_TOOL]) + _assert_completed(_send(gateway, body)) + labels: Final = anthropic_labels(_only_request(wire)) + assert labels == [ASK_LABEL, _ASSISTANT_LABEL, _RESULT_LABEL, final_label(marker)] + + +def test_a_top_level_mark_takes_a_slot_under_drop_params(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY, drop_params=True + ) + body: Final = _responses( + model, _many_marks(marker, assistant_marked=False), tools=[_TOOL], cache_control=EPHEMERAL + ) + _assert_completed(_send(gateway, body)) + labels: Final = anthropic_labels(_only_request(wire)) + assert labels == [ASK_LABEL, _RESULT_LABEL, final_label(marker), "request"] + + +@pytest.mark.parametrize( + "mark", + ( + pytest.param("ephemeral", id="string"), + pytest.param({"type": "forever"}, id="unknown-type"), + ), +) +def test_malformed_mark_is_dropped_under_drop_params(gateway: Gateway, mark: JsonValue) -> None: + marker: Final = new_marker() + with wire_server(_strict_anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY, drop_params=True + ) + _assert_completed(_send(gateway, _marked_system(model, marker, mark))) + assert anthropic_labels(_only_request(wire)) == [] + + +def test_ttl_order_violation_is_the_vendors_verdict(gateway: Gateway) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + _item("system", SYSTEM), + _item("user", ASK, _HOUR), + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + answer: Final = _send(gateway, _responses(model, items)) + marks: Final = anthropic_marks(_only_request(wire)) + assert answer.status == 400, answer.text + assert "ttl='1h' cache_control block must not come after" in answer.text, answer.text + assert [(mark.label, mark.ttl) for mark in marks] == [(SYSTEM_LABEL, None), (ASK_LABEL, "1h")] + + +def test_unauthenticated_marked_request_never_reaches_the_provider(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + answer: Final = _send(gateway, _marked_system(model, marker), key=f"sk-not-a-key-{marker}") + received: Final = wire.drain() + assert answer.status == 401, answer.text + assert "error" in _JSON_OBJECT.validate_json(answer.text), answer.text + assert received == () + + +def _overloaded_peer(request: Request) -> Reply: + return Reply( + status=503, + body=json.dumps({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}).encode(), + ) + + +def test_provider_error_on_a_marked_request_reaches_the_caller(gateway: Gateway, spend: pcb.SpendLogs) -> None: + marker: Final = new_marker() + with wire_server(_overloaded_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + answer: Final = _send(gateway, _marked_system(model, marker)) + labels: Final = anthropic_labels(_only_request(wire)) + assert answer.status == 500, answer.text + assert "overloaded_error" in answer.text, answer.text + assert labels == [SYSTEM_LABEL] + spend.landed(model, answer.call_id, None, status="failure") + + +_SAP_DISCOVERY: Final = ("/oauth/token", "/v2/lm/") +_SAP_COMPLETION: Final = "/sap-deployment/v2/completion" + + +def _only_post(wire: Wire) -> Request: + posted: Final = [ + request + for request in wire.drain() + if request.method == "POST" and not request.target.startswith(_SAP_DISCOVERY) + ] + assert len(posted) == 1, [request.target for request in posted] + return posted[0] + + +def _message_mark_labels(message: JsonValue) -> list[str]: + fields: Final = object_value(message) + content: Final = fields.get("content") + blocks: Final = content if isinstance(content, list) else [] + block_labels: Final = [ + f"{fields.get('role')}:block" for block in blocks if object_value(block).get("cache_control") is not None + ] + return [f"{fields.get('role')}:message"] * int(fields.get("cache_control") is not None) + block_labels + + +def _openai_mark_labels(request: Request) -> list[str]: + messages: Final = _JSON_OBJECT.validate_json(request.body)["messages"] + return list(chain.from_iterable(map(_message_mark_labels, messages if isinstance(messages, list) else []))) + + +def _sap_mark_labels(request: Request) -> list[str]: + config: Final = object_value(_JSON_OBJECT.validate_json(request.body)["config"]) + prompt: Final = object_value(object_value(object_value(config["modules"])["prompt_templating"])["prompt"]) + template: Final = prompt["template"] + return list(chain.from_iterable(map(_message_mark_labels, template if isinstance(template, list) else []))) + + +def _route_mark_labels(request: Request) -> list[str]: + if "/converse" in request.target: + return bedrock_labels(request) + if request.target.endswith("/chat/completions"): + return _openai_mark_labels(request) + if request.target == _SAP_COMPLETION: + return _sap_mark_labels(request) + return anthropic_labels(request) + + +def _sap_peer(request: Request) -> Reply: + match request.target: + case "/oauth/token": + return Reply(body=json.dumps({"access_token": "scripted", "expires_in": 3600}).encode()) + case "/v2/lm/deployments": + deployment: Final = { + "scenarioId": "orchestration", + "configurationId": "orchestration-config", + "deploymentUrl": f"http://{request.headers['host']}/sap-deployment", + "createdAt": "2026-01-01T00:00:00Z", + } + return Reply(body=json.dumps({"resources": [deployment]}).encode()) + case "/v2/lm/configurations/orchestration-config": + return Reply(body=json.dumps({"executableId": "orchestration"}).encode()) + case _: + completion: Final = chat_reply(f"chatcmpl-{marker_of(request)}", "scripted", "sunny", stream=False) + return Reply(body=json.dumps({"final_result": json.loads(completion.body)}).encode()) + + +def _route_peer(request: Request) -> Reply: + if "/converse" in request.target: + return bedrock_peer(request) + if request.target.endswith("/chat/completions"): + return chat_reply(f"chatcmpl-{marker_of(request)}", "scripted", "sunny", stream=False) + if request.target.startswith((*_SAP_DISCOVERY, "/sap-deployment/")): + return _sap_peer(request) + return anthropic_peer(request) + + +def _route_deployment(gateway: Gateway, scenario: Scenario, wire: Wire, route: str) -> str: + match route: + case "vertex-claude": + return scenario.model( + model=f"vertex_ai/{ANTHROPIC_MODEL}", + api_key=None, + api_base=wire.url, + vertex_project="scripted-project", + vertex_location="us-east5", + vertex_credentials=service_account_json("scripted-project", gateway.upstream_url.rstrip("/")), + ) + case "bedrock-llama": + return scenario.model( + model="bedrock/converse/meta.llama3-3-70b-instruct-v1:0", + api_key=PROVIDER_KEY, + api_base=wire.url, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + case "sap-claude": + service_key: Final = { + "clientid": "scripted", + "clientsecret": PROVIDER_KEY, + "url": wire.url, + "serviceurls": {"AI_API_URL": wire.url}, + } + return scenario.model( + model="sap/anthropic--claude-4.5-sonnet", api_base=wire.url, api_key=json.dumps(service_key) + ) + case _ if route.startswith("openrouter/"): + return scenario.model(model=route, api_base=wire.url, api_key=PROVIDER_KEY, use_chat_completions_api=True) + case _: + return scenario.model(model=route, api_base=wire.url, api_key=PROVIDER_KEY) + + +def _chat_with_tool_mark(model: str, marker: str) -> dict[str, JsonValue]: + tool_call: Final[dict[str, JsonValue]] = { + "id": call_id(CITIES[0]), + "type": "function", + "function": {"name": "lookup_weather", "arguments": json.dumps({"city": CITIES[0]})}, + } + return { + "model": model, + "max_tokens": 64, + "tools": [{"type": "function", "function": {"name": "lookup_weather", "parameters": {"type": "object"}}}], + "messages": [ + {"role": "user", "content": ASK}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": call_id(CITIES[0]), "content": "sunny.", "cache_control": EPHEMERAL}, + {"role": "user", "content": final_text(marker)}, + ], + } + + +def _shaped_requests(model: str, shape: str) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]: + match shape: + case "tool": + bridged: Final = _responses(model, _tool_items(new_marker(), output_mark=EPHEMERAL), tools=[_TOOL]) + return _chat_with_tool_mark(model, new_marker()), bridged + case _: + return _chat_with_message_mark(model, new_marker()), _marked_system(model, new_marker()) + + +@pytest.mark.parametrize( + ("route", "shape", "keeps"), + ( + pytest.param("dashscope/qwen-plus", "system", True, id="dashscope"), + pytest.param("minimax/MiniMax-M2", "system", True, id="minimax"), + pytest.param("zai/glm-4.6", "system", True, id="zai"), + pytest.param(f"openrouter/anthropic/{ANTHROPIC_MODEL}", "system", True, id="openrouter-claude"), + pytest.param("openrouter/google/gemini-3.8-flash", "system", True, id="openrouter-gemini"), + pytest.param("snowflake/claude-4-sonnet", "tool", True, id="snowflake-claude-tool-output"), + pytest.param("vertex-claude", "system", True, id="vertex-claude"), + pytest.param("sap-claude", "system", True, id="sap-claude"), + pytest.param("anthropic/my-gateway-model", "system", True, id="anthropic-non-claude-name"), + pytest.param("snowflake/claude-4-sonnet", "system", False, id="snowflake-claude-string-system"), + pytest.param("openrouter/openai/gpt-5.4", "system", False, id="openrouter-gpt"), + pytest.param("snowflake/llama3.1-70b", "system", False, id="snowflake-llama"), + pytest.param("snowflake/llama-claude-distill", "system", False, id="snowflake-claude-inside-the-name"), + pytest.param("databricks/databricks-meta-llama-3-3-70b-instruct", "tool", False, id="databricks-llama"), + pytest.param("bedrock-llama", "system", False, id="bedrock-llama"), + ), +) +def test_item_mark_reaches_each_route_the_way_a_chat_message_mark_does( + gateway: Gateway, route: str, shape: str, keeps: bool +) -> None: + with wire_server(answering_model_discovery(_route_peer)) as wire, gateway.scenario() as scenario: + model: Final = _route_deployment(gateway, scenario, wire, route) + on_chat_body, bridged_body = _shaped_requests(model, shape) + on_chat: Final = gateway.request("POST", "/v1/chat/completions", on_chat_body) + assert on_chat.status_code == 200, on_chat.text + chat_request: Final = _only_post(wire) + _assert_completed(_send(gateway, bridged_body)) + bridged_request: Final = _only_post(wire) + bridged_labels: Final = _route_mark_labels(bridged_request) + assert bridged_request.target == chat_request.target + assert bridged_labels == (_route_mark_labels(chat_request) if keeps else []) + assert bool(bridged_labels) is keeps, bridged_labels + + +@pytest.mark.parametrize( + ("drop_params", "expected"), + ( + pytest.param(False, [_BEDROCK_SYSTEM_LABEL], id="as-sent"), + pytest.param(True, [], id="drop-params"), + ), +) +def test_unmapped_bedrock_arn_keeps_the_item_mark_unless_drop_params( + gateway: Gateway, drop_params: bool, expected: list[str] +) -> None: + marker: Final = new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _unmapped_arn_deployment(scenario, wire, drop_params=drop_params) + answer: Final = _send(gateway, _marked_system(model, marker)) + assert answer.status == 200, answer.text + assert bedrock_labels(_only_request(wire)) == expected + + +@pytest.mark.parametrize("request_drop_params", (True, "true"), ids=("bool", "string")) +def test_request_drop_params_trims_five_marks_to_the_newest_four( + gateway: Gateway, request_drop_params: JsonValue +) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + items: Final = _many_marks(marker, assistant_marked=True) + _assert_completed(_send(gateway, _responses(model, items, tools=[_TOOL], drop_params=request_drop_params))) + labels: Final = anthropic_labels(_only_request(wire)) + assert labels == [ASK_LABEL, _ASSISTANT_LABEL, _RESULT_LABEL, final_label(marker)] + + +@pytest.mark.parametrize( + ("deployment_fields", "request_drop_params"), + ( + pytest.param({}, None, id="request-null"), + pytest.param({}, "", id="request-empty-string"), + pytest.param({"drop_params": True}, False, id="request-false-over-deployment-true"), + ), +) +def test_request_without_drop_params_forwards_five_marks_as_sent( + gateway: Gateway, deployment_fields: dict[str, JsonValue], request_drop_params: JsonValue +) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY, **deployment_fields + ) + items: Final = _many_marks(marker, assistant_marked=True) + answer: Final = _send(gateway, _responses(model, items, tools=[_TOOL], drop_params=request_drop_params)) + labels: Final = anthropic_labels(_only_request(wire)) + assert answer.status == 400, answer.text + assert "A maximum of 4 blocks with cache_control" in answer.text, answer.text + assert len(labels) == 5, labels + + +def _marked_function_call_items(marker: str) -> list[JsonValue]: + calls: Final[list[JsonValue]] = [ + { + "type": "function_call", + "call_id": call_id(city), + "name": "lookup_weather", + "arguments": json.dumps({"city": city}), + "cache_control": EPHEMERAL, + } + for city in CITIES + ] + outputs: Final[list[JsonValue]] = [ + {"type": "function_call_output", "call_id": call_id(city), "output": "sunny."} for city in CITIES + ] + return [{"role": "user", "content": ASK}, *calls, *outputs, {"role": "user", "content": final_text(marker)}] + + +def test_injection_points_count_a_marked_function_tool_toward_the_cap(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire, inject=True) + body: Final = _responses( + model, + _marked_function_call_items(marker), + instructions=SYSTEM, + tools=[{**_TOOL, "cache_control": EPHEMERAL}], + ) + _assert_completed(_send(gateway, body)) + labels: Final = anthropic_labels(_only_request(wire)) + assert labels == ["tool:lookup_weather", *(tool_use_label(city) for city in CITIES)] + + +def test_assistant_text_after_a_function_call_keeps_its_mark(gateway: Gateway) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + {"role": "user", "content": ASK}, + { + "type": "function_call", + "call_id": call_id(CITIES[0]), + "name": "lookup_weather", + "arguments": json.dumps({"city": CITIES[0]}), + }, + _item("assistant", _CHECK), + {"type": "function_call_output", "call_id": call_id(CITIES[0]), "output": "sunny."}, + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + _assert_completed(_send(gateway, _responses(model, items, tools=[_TOOL]))) + labels: Final = anthropic_labels(_only_request(wire)) + assert labels == [_ASSISTANT_LABEL] + + +_COUNTED_TOKENS: Final = 42 + + +def _count_peer(request: Request) -> Reply: + if request.target.endswith("/count_tokens"): + return Reply(body=json.dumps({"input_tokens": _COUNTED_TOKENS}).encode()) + return anthropic_peer(request) + + +def test_input_tokens_counts_marked_items_without_their_marks(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(_count_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = {"model": model, "input": _many_marks(marker, assistant_marked=True), "tools": [_TOOL]} + response: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + counted: Final = _only_request(wire) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.text) == { + "object": "response.input_tokens", + "input_tokens": _COUNTED_TOKENS, + } + assert counted.target.endswith("/count_tokens"), counted.target + assert anthropic_labels(counted) == [] + + +_INPUT_SHAPES: Final = ("string", "list", "object") + + +def _shaped_input(shape: str, marker: str) -> JsonValue: + marked: Final = _item("user", final_text(marker)) + match shape: + case "string": + return final_text(marker) + case "list": + return [_item("system", SYSTEM), marked] + case _: + return {"type": "message", **marked} + + +def _shaped_body(model: str, shape: str, marker: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_output_tokens": 64, + "instructions": SYSTEM, + "input": _shaped_input(shape, marker), + **fields, + } + + +def _error_message(answer: _Answer) -> JsonValue: + if answer.status == 200: + return None + return object_value(_JSON_OBJECT.validate_json(answer.text)["error"]).get("message") + + +@pytest.mark.parametrize("shape", _INPUT_SHAPES) +def test_each_input_shape_reads_the_same_whether_or_not_the_route_keeps_item_marks( + gateway: Gateway, shape: str +) -> None: + kept_marker, stripped_marker = new_marker(), new_marker() + with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _unmapped_arn_deployment(scenario, wire) + kept: Final = _send(gateway, _shaped_body(model, shape, kept_marker)) + kept_sent: Final = sent_without_cache_points(wire.drain(), kept_marker) + stripped: Final = _send(gateway, _shaped_body(model, shape, stripped_marker, drop_params=True)) + stripped_sent: Final = sent_without_cache_points(wire.drain(), stripped_marker) + assert (stripped.status, _error_message(stripped)) == (kept.status, _error_message(kept)), ( + kept.text, + stripped.text, + ) + assert stripped_sent == kept_sent + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_object_input_on_a_mark_stripping_route_is_served(gateway: Gateway, stream: bool) -> None: + marker: Final = new_marker() + reply: Final = chat_reply(f"chatcmpl-{marker}", "gpt-5.4", "sunny", stream=stream) + with wire_server(answering_model_discovery(lambda request: reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_OPENAI_COMPATIBLE, api_base=f"{wire.url}/v1", api_key=PROVIDER_KEY, use_chat_completions_api=True + ) + _assert_completed(_send(gateway, _shaped_body(model, "object", marker, stream=stream)), stream=stream) + posted: Final = [request for request in wire.drain() if request.method == "POST"] + assert [request.target for request in posted] == ["/v1/chat/completions"], [request.target for request in posted] + assert _openai_message_marks(posted[0]) == [] + + +@pytest.mark.parametrize("shape", _INPUT_SHAPES) +def test_input_tokens_reads_each_input_shape_without_item_marks(gateway: Gateway, shape: str) -> None: + marker: Final = new_marker() + with wire_server(_count_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic_deployment(scenario, wire) + body: Final = { + key: value for key, value in _shaped_body(model, shape, marker).items() if key != "max_output_tokens" + } + response: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + counted: Final = wire.drain() + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.text) == { + "object": "response.input_tokens", + "input_tokens": _COUNTED_TOKENS, + } + assert [anthropic_labels(request) for request in counted] == [[]] + + +@pytest.mark.parametrize("shape", _INPUT_SHAPES) +def test_project_input_token_limit_reads_each_input_shape(gateway: Gateway, shape: str) -> None: + marker: Final = new_marker() + reply: Final = chat_reply(f"chatcmpl-{marker}", "gpt-5.4", "sunny", stream=False) + with wire_server(answering_model_discovery(lambda request: reply)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_OPENAI_COMPATIBLE, api_base=f"{wire.url}/v1", api_key=PROVIDER_KEY, use_chat_completions_api=True + ) + team: Final = scenario.team() + key: Final = scenario.key(team_id=team, project_id=scenario.project(team, model_itpm_limit={model: 1_000_000})) + _assert_completed(_send(gateway, _shaped_body(model, shape, marker), key=key)) + posted: Final = [request for request in wire.drain() if request.method == "POST"] + assert [request.target for request in posted] == ["/v1/chat/completions"], [request.target for request in posted] + assert _openai_message_marks(posted[0]) == [] diff --git a/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py b/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py new file mode 100644 index 00000000000..74ea3272503 --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py @@ -0,0 +1,575 @@ +import asyncio +import dataclasses +import threading +from collections import Counter +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, TypeAlias +from urllib.parse import urlsplit + +import httpx +import pytest +import yaml +from integration._support import prompt_cache_breakpoint as pcb +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.openai_wire import answering_model_discovery, chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + ANTHROPIC_MODEL, + ASK, + ASK_LABEL, + CITIES, + EPHEMERAL, + PROVIDER_KEY, + SYSTEM, + SYSTEM_LABEL, + anthropic_deployment, + anthropic_labels, + anthropic_peer, + bedrock_labels, + bedrock_peer, + call_id, + final_label, + final_text, + marker_of, + new_marker, + owned_config, + tool_use_label, +) +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_OWNED: Final = "bridge-mark-owned" +_OWNED_DROP: Final = "bridge-mark-owned-drop" +_NOVA: Final = "bridge-mark-nova" +_CHAT: Final = "bridge-mark-chat" +_COMPARE: Final = "Compare the three cities." +_CHECK: Final = "I will check." +_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "name": "lookup_weather", "parameters": {"type": "object"}} +_UNMARKED: Final = "unmarked" +_RecordProperty: TypeAlias = Callable[[str, object], None] + + +@dataclass(frozen=True, slots=True) +class _Call: + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + call_id: str + + +@dataclass(frozen=True, slots=True) +class _StoredRig: + wire: Wire + proxy: OwnedProxy + + @property + def gateway(self) -> Gateway: + return self.proxy.gateway + + +def _owned_config(directory: Path, model_list: Sequence[JsonValue]) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[Mapping[str, JsonValue]] = { + **stock, + "model_list": list(model_list), + "general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True}, + "router_settings": {**object_value(stock.get("router_settings") or {}), "num_retries": 0}, + } + path: Final = directory / "bridge-mark-owned.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _stored_peer(request: Request) -> Reply: + if "/converse" in request.target: + return bedrock_peer(request) + if request.target.endswith("/chat/completions"): + return chat_reply(f"chatcmpl-{marker_of(request)}", "gpt-5.4", "sunny", stream=False) + return anthropic_peer(request) + + +def _stored_models(api_base: str) -> list[JsonValue]: + return [ + anthropic_deployment(_OWNED, api_base), + anthropic_deployment(_OWNED_DROP, api_base, drop_params=True), + { + "model_name": _NOVA, + "litellm_params": { + "model": "bedrock/converse/us.amazon.nova-2-lite-v1:0", + "api_key": PROVIDER_KEY, + "api_base": api_base, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": api_base, + }, + }, + { + "model_name": _CHAT, + "litellm_params": { + "model": "openai/gpt-5.4", + "api_key": PROVIDER_KEY, + "api_base": f"{api_base}/v1", + "use_chat_completions_api": True, + }, + }, + ] + + +@pytest.fixture(scope="module") +def stored_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + directory: Final = tmp_path_factory.mktemp("bridge-mark-stored") + with wire_server(answering_model_discovery(_stored_peer)) as wire, gateway_from_environment() as gateway: + config: Final = _owned_config(directory, _stored_models(wire.url)) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned: + yield _StoredRig(wire, owned) + + +@pytest.fixture(scope="module") +def yaml_drop_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + directory: Final = tmp_path_factory.mktemp("bridge-mark-yaml-drop") + with wire_server(anthropic_peer) as wire, gateway_from_environment() as gateway: + config: Final = owned_config( + directory, [anthropic_deployment(_OWNED, wire.url)], litellm_settings={"drop_params": True} + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned: + yield _StoredRig(wire, owned) + + +@pytest.fixture(scope="module") +def spend() -> Iterator[pcb.SpendLogs]: + with pcb.spend_logs() as logs: + yield logs + + +def _body(model: str, call: _Call) -> Mapping[str, JsonValue]: + return { + "model": model, + "stream": call.stream, + "max_output_tokens": 64, + "cache": {"no-cache": True}, + "input": [ + {"role": "system", "content": SYSTEM, "cache_control": EPHEMERAL}, + {"role": "user", "content": final_text(call.marker)}, + ], + } + + +def _calls(count: int, *, stream: bool | None = None) -> tuple[_Call, ...]: + return tuple(_Call(index % 2 == 1 if stream is None else stream, new_marker()) for index in range(count)) + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", "/v1/responses", json=_body(model, call), headers={"Authorization": f"Bearer {key}"} + ) as response: + raw: Final = await response.aread() + return _Served(call, response.status_code, raw.decode(), response.headers["x-litellm-call-id"]) + + +async def _burst( + gateway: Gateway, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, gateway.key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_answered(served: _Served) -> None: + assert served.status == 200, served.text + assert "sunny" in served.text, served.text + assert served.text.startswith(("event:", "data:")) == served.call.stream, served.text + assert ("response.completed" in served.text) == served.call.stream, served.text + if not served.call.stream: + assert _JSON_OBJECT.validate_json(served.text)["status"] == "completed", served.text + + +def _marked_posts(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker_of(request) != _UNMARKED) + + +def _marked_once(posts: Sequence[Request], calls: Sequence[_Call]) -> None: + by_marker: Final = {marker_of(request): request for request in posts} + assert len(by_marker) == len(posts), [request.body for request in posts] + assert set(by_marker) == {call.marker for call in calls}, sorted(by_marker) + assert [anthropic_labels(by_marker[call.marker]) for call in calls] == [[SYSTEM_LABEL]] * len(calls) + + +def _assert_each_lands_once( + spend: pcb.SpendLogs, model: str, failed: Sequence[_Served], served: Sequence[_Served] +) -> None: + expected: Final = len(failed) + len(served) + rows: Final = eventually(lambda: spend.rows_for(model), lambda found: len(found) >= expected, seconds=70) + by_call: Final = {string_value(row["litellm_call_id"]): row for row in rows} + assert len(by_call) == len(rows) == expected, rows + for item in failed: + assert by_call[item.call_id]["status"] == "failure", (item.call_id, rows) + for item in served: + assert by_call[item.call_id]["status"] == "success", (item.call_id, rows) + + +def _health(gateway: Gateway, model: str) -> Mapping[str, JsonValue]: + response: Final = gateway.request("GET", f"/health?model={model}", None) + assert response.status_code in (200, 503), response.text + return _JSON_OBJECT.validate_json(response.text) + + +def _free_port() -> int: + with wire_server(anthropic_peer) as probe: + port: Final = urlsplit(probe.url).port + assert port is not None, probe.url + return port + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=api_base, api_key=PROVIDER_KEY) + + +async def test_mixed_burst_carries_the_system_mark_once_per_call(gateway: Gateway, spend: pcb.SpendLogs) -> None: + calls: Final = _calls(24) + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + served: Final = await _burst(gateway, model, calls) + assert len(served) == 24 + for item in served: + _assert_answered(item) + _marked_once(_marked_posts(wire), calls) + _assert_each_lands_once(spend, model, (), served) + + +async def test_peer_outage_fails_cleanly_and_the_restarted_peer_receives_the_marks( + gateway: Gateway, spend: pcb.SpendLogs +) -> None: + port: Final = _free_port() + while_down: Final = _calls(12) + after: Final = _calls(12) + with gateway.scenario() as scenario: + model: Final = _deployment(scenario, f"http://127.0.0.1:{port}") + failed: Final = await _burst(gateway, model, while_down) + assert len(failed) == 12 + for item in failed: + assert item.status >= 500, (item.status, item.text) + assert "sunny" not in item.text and "event:" not in item.text, item.text + down: Final = _health(gateway, model) + assert (down["healthy_count"], down["unhealthy_count"]) == (0, 1), down + with wire_server(anthropic_peer, port=port) as wire: + _health(gateway, model) + assert _marked_posts(wire) == () + served: Final = await _burst(gateway, model, after) + assert len(served) == 12 + for item in served: + _assert_answered(item) + _marked_once(_marked_posts(wire), after) + _assert_each_lands_once(spend, model, failed, served) + + +def _slow(request: Request) -> Reply: + reply: Final = anthropic_peer(request) + return dataclasses.replace(reply, pause_between_chunks=0.4) if reply.chunks else reply + + +async def test_concurrent_slow_streams_each_carry_the_mark_once(gateway: Gateway, spend: pcb.SpendLogs) -> None: + calls: Final = _calls(6, stream=True) + with wire_server(_slow) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + served: Final = await _burst(gateway, model, calls) + assert len(served) == 6 + for item in served: + _assert_answered(item) + _marked_once(_marked_posts(wire), calls) + _assert_each_lands_once(spend, model, (), served) + + +@dataclass(frozen=True, slots=True) +class _Held: + release: threading.Event + markers: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + marker: Final = marker_of(request) + if marker == _UNMARKED: + return anthropic_peer(request) + self.markers.put(marker) + if not self.release.wait(timeout=60): + return Reply(status=504, body=b'{"type":"error","error":{"type":"api_error","message":"never released"}}') + return anthropic_peer(request) + + +def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: + return eventually(lambda: pcb.started_worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=30) + + +async def _hold_burst( + held: _Held, candidate: Gateway, model: str, calls: tuple[_Call, ...] +) -> asyncio.Task[tuple[_Served, ...]]: + burst: Final = asyncio.create_task(_burst(candidate, model, calls, tolerate_transport_errors=True)) + await asyncio.to_thread(eventually, held.markers.qsize, lambda size: size == len(calls), 60) + return burst + + +def _follow_up(gateway: Gateway, wire: Wire, spend: pcb.SpendLogs) -> None: + call: Final = _Call(False, new_marker()) + response: Final = gateway.request("POST", "/v1/responses", dict(_body(_OWNED, call))) + served: Final = _Served(call, response.status_code, response.text, response.headers["x-litellm-call-id"]) + _assert_answered(served) + _marked_once(_marked_posts(wire), (call,)) + spend.landed(_OWNED, served.call_id, None) + + +async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering( + gateway: Gateway, tmp_path: Path, spend: pcb.SpendLogs +) -> None: + calls: Final = _calls(20, stream=False) + held: Final = _Held(threading.Event(), SimpleQueue()) + with wire_server(held.respond) as wire: + config: Final = _owned_config(tmp_path, [anthropic_deployment(_OWNED, wire.url)]) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = _worker_pids(owned) + burst: Final = await _hold_burst(held, candidate, _OWNED, calls) + held_by: Final = MappingProxyType({pid: pcb.open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill() + held.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered(item) + posts: Final = _marked_posts(wire) + assert {marker_of(request) for request in posts} == {call.marker for call in calls}, posts + assert [anthropic_labels(request) for request in posts] == [[SYSTEM_LABEL]] * len(posts) + _follow_up(candidate, wire, spend) + + +async def test_proxy_restart_mid_burst_never_lands_a_served_call_twice( + gateway: Gateway, tmp_path: Path, record_property: _RecordProperty, spend: pcb.SpendLogs +) -> None: + calls: Final = _calls(20, stream=False) + held: Final = _Held(threading.Event(), SimpleQueue()) + with wire_server(held.respond) as wire: + config: Final = _owned_config(tmp_path, [anthropic_deployment(_OWNED, wire.url)]) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first: + _worker_pids(first) + burst: Final = await _hold_burst(held, first.gateway, _OWNED, calls) + first.begin_shutdown() + held.release.set() + served: Final = await burst + for item in served: + _assert_answered(item) + wire.drain() + second_directory: Final = tmp_path / "second" + second_directory.mkdir() + with owned_proxy_process(gateway, second_directory, {}, config=config, workers=2) as second: + _follow_up(second.gateway, wire, spend) + counts: Final = Counter(string_value(row["litellm_call_id"]) for row in spend.rows_for(_OWNED)) + assert all(count == 1 for count in counts.values()), counts + landed: Final = sum(1 for item in served if item.call_id in counts) + record_property("served", len(served)) + record_property("landed", landed) + record_property("lost_responses", len(calls) - len(served)) + + +def test_previous_response_replay_keeps_the_system_mark(stored_rig: _StoredRig, spend: pcb.SpendLogs) -> None: + first_call: Final = _Call(False, new_marker()) + first: Final = stored_rig.gateway.request( + "POST", + "/v1/responses", + { + **_body(_OWNED, first_call), + "input": [ + {"role": "system", "content": SYSTEM, "cache_control": EPHEMERAL}, + {"role": "user", "content": final_text(first_call.marker)}, + ], + }, + ) + assert first.status_code == 200, first.text + first_id: Final = string_value(_JSON_OBJECT.validate_json(first.text)["id"]) + spend.landed(_OWNED, first.headers["x-litellm-call-id"], None) + assert [anthropic_labels(request) for request in _marked_posts(stored_rig.wire)] == [[SYSTEM_LABEL]] + second: Final = stored_rig.gateway.request( + "POST", + "/v1/responses", + { + "model": _OWNED, + "max_output_tokens": 64, + "cache": {"no-cache": True}, + "previous_response_id": first_id, + "input": [{"role": "user", "content": final_text(new_marker())}], + }, + ) + assert second.status_code == 200, second.text + assert _JSON_OBJECT.validate_json(second.text)["status"] == "completed", second.text + replayed: Final = _marked_posts(stored_rig.wire) + assert len(replayed) == 1, [request.target for request in replayed] + body: Final = _JSON_OBJECT.validate_json(replayed[0].body) + assert isinstance(body["messages"], list) and len(body["messages"]) == 3, body + assert anthropic_labels(replayed[0]) == [SYSTEM_LABEL] + + +def _respond(rig: _StoredRig, body: Mapping[str, JsonValue]) -> httpx.Response: + response: Final = rig.gateway.request("POST", "/v1/responses", dict(body)) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.text)["status"] == "completed", response.text + return response + + +def _stored_turn(rig: _StoredRig, spend: pcb.SpendLogs, model: str, items: Sequence[JsonValue]) -> tuple[str, Request]: + rig.wire.drain() + first: Final = _respond( + rig, {"model": model, "max_output_tokens": 64, "cache": {"no-cache": True}, "tools": [_TOOL], "input": items} + ) + spend.landed(model, first.headers["x-litellm-call-id"], None) + (posted,) = _marked_posts(rig.wire) + return string_value(_JSON_OBJECT.validate_json(first.text)["id"]), posted + + +def _replayed(rig: _StoredRig, model: str, previous: str, items: Sequence[JsonValue]) -> Request: + _respond( + rig, + { + "model": model, + "max_output_tokens": 64, + "cache": {"no-cache": True}, + "tools": [_TOOL], + "previous_response_id": previous, + "input": list(items), + }, + ) + (posted,) = _marked_posts(rig.wire) + return posted + + +def _marked(role: str, text: str, mark: JsonValue = EPHEMERAL) -> dict[str, JsonValue]: + return {"role": role, "content": text, "cache_control": mark} + + +def _weather_call(**fields: JsonValue) -> dict[str, JsonValue]: + return { + "type": "function_call", + "call_id": call_id(CITIES[0]), + "name": "lookup_weather", + "arguments": '{"city": "London"}', + **fields, + } + + +def _weather_output(**fields: JsonValue) -> dict[str, JsonValue]: + return {"type": "function_call_output", "call_id": call_id(CITIES[0]), "output": "sunny.", **fields} + + +def _marked_tool_round() -> list[JsonValue]: + return [ + {"role": "user", "content": ASK}, + _weather_call(cache_control=EPHEMERAL), + _weather_output(cache_control=EPHEMERAL), + ] + + +def test_replay_gives_the_current_marks_their_slots_first(stored_rig: _StoredRig, spend: pcb.SpendLogs) -> None: + first_marker: Final = new_marker() + second_marker: Final = new_marker() + previous, first = _stored_turn( + stored_rig, + spend, + _OWNED, + [_marked("system", SYSTEM), _marked("user", ASK), _marked("user", final_text(first_marker))], + ) + assert anthropic_labels(first) == [SYSTEM_LABEL, ASK_LABEL, final_label(first_marker)] + replayed: Final = _replayed( + stored_rig, _OWNED, previous, [_marked("user", _COMPARE), _marked("user", final_text(second_marker))] + ) + assert anthropic_labels(replayed) == [ + ASK_LABEL, + final_label(first_marker), + f"user:text:{_COMPARE}", + final_label(second_marker), + ] + + +def test_replay_never_forwards_a_mark_drop_params_dropped(stored_rig: _StoredRig, spend: pcb.SpendLogs) -> None: + previous, first = _stored_turn( + stored_rig, + spend, + _OWNED_DROP, + [_marked("system", SYSTEM, "ephemeral"), {"role": "user", "content": final_text(new_marker())}], + ) + assert b"cache_control" not in first.body, first.body + replayed: Final = _replayed(stored_rig, _OWNED, previous, [{"role": "user", "content": final_text(new_marker())}]) + assert b"cache_control" not in replayed.body, replayed.body + + +def test_replay_into_nova_drops_the_tool_block_marks(stored_rig: _StoredRig, spend: pcb.SpendLogs) -> None: + previous, first = _stored_turn( + stored_rig, spend, _OWNED, [*_marked_tool_round(), {"role": "user", "content": final_text(new_marker())}] + ) + assert anthropic_labels(first) == [tool_use_label(CITIES[0]), f"user:tool_result:{call_id(CITIES[0])}"] + replayed: Final = _replayed(stored_rig, _NOVA, previous, [{"role": "user", "content": final_text(new_marker())}]) + assert "/converse" in replayed.target, replayed.target + assert bedrock_labels(replayed) == [] + + +def test_replay_into_a_chat_route_without_marks_sends_none(stored_rig: _StoredRig, spend: pcb.SpendLogs) -> None: + previous, first = _stored_turn( + stored_rig, + spend, + _OWNED, + [_marked("system", SYSTEM), {"role": "user", "content": final_text(new_marker())}], + ) + assert anthropic_labels(first) == [SYSTEM_LABEL] + replayed: Final = _replayed(stored_rig, _CHAT, previous, [{"role": "user", "content": final_text(new_marker())}]) + assert replayed.target == "/v1/chat/completions", replayed.target + assert b"cache_control" not in replayed.body, replayed.body + + +def test_yaml_drop_params_keeps_the_newest_four_marks(yaml_drop_rig: _StoredRig) -> None: + marker: Final = new_marker() + yaml_drop_rig.wire.drain() + _respond( + yaml_drop_rig, + { + "model": _OWNED, + "max_output_tokens": 64, + "cache": {"no-cache": True}, + "tools": [_TOOL], + "input": [ + _marked("system", SYSTEM), + _marked("user", ASK), + _marked("assistant", _CHECK), + _weather_call(), + _weather_output(cache_control=EPHEMERAL), + _marked("user", final_text(marker)), + ], + }, + ) + (posted,) = _marked_posts(yaml_drop_rig.wire) + assert anthropic_labels(posted) == [ + ASK_LABEL, + f"assistant:text:{_CHECK}", + f"user:tool_result:{call_id(CITIES[0])}", + final_label(marker), + ] diff --git a/tests/integration/providers/test_voyage_mongodb_host_wire.py b/tests/integration/providers/test_voyage_mongodb_host_wire.py index 926bd5a7c69..273020d48a8 100644 --- a/tests/integration/providers/test_voyage_mongodb_host_wire.py +++ b/tests/integration/providers/test_voyage_mongodb_host_wire.py @@ -131,7 +131,12 @@ def _scrubbed() -> tuple[str, ...]: def _forward_environment(forward_url: str, **extra: str) -> dict[str, str]: - return {"HTTPS_PROXY": forward_url, "NO_PROXY": "127.0.0.1,localhost", **extra} + return { + "HTTPS_PROXY": forward_url, + "NO_PROXY": "127.0.0.1,localhost", + "PROXY_BATCH_POLLING_ENABLED": "false", + **extra, + } @pytest.fixture(scope="module") diff --git a/tests/integration/routing/test_team_model_tpm_limit.py b/tests/integration/routing/test_team_model_tpm_limit.py deleted file mode 100644 index 741c41c9285..00000000000 --- a/tests/integration/routing/test_team_model_tpm_limit.py +++ /dev/null @@ -1,89 +0,0 @@ -import json -import threading -import uuid -from concurrent.futures import Future, ThreadPoolExecutor -from typing import Final - -import httpx -import pytest -from integration._support.client import Gateway, eventually -from integration._support.wire import Reply, Request, wire_server - -PROVIDER_MODEL: Final = "gpt-4o-mini" -TEAM_MODEL_TPM: Final = 100 -MAX_TOKENS: Final = 60 -CONCURRENT_REQUESTS: Final = 3 -UPSTREAM_REPLY: Final = json.dumps( - { - "id": "chatcmpl-team-tpm-control", - "object": "chat.completion", - "created": 1, - "model": PROVIDER_MODEL, - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": "team tpm control"}, "finish_reason": "stop"} - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, - } -).encode() - - -@pytest.mark.covers("routing.team_model_tpm.concurrent_requests_over_the_limit_are_rejected_before_the_provider_call") -def test_concurrent_team_model_tpm_requests_reserve_tokens_before_reaching_the_provider(gateway: Gateway) -> None: - probe: Final = "team tpm probe " + uuid.uuid4().hex - release: Final = threading.Event() - - def respond(request: Request) -> Reply: - assert request.method == "POST" and request.target == "/v1/chat/completions" - assert request.headers["authorization"] == "Bearer synthetic-team-tpm-key" - body: Final = json.loads(request.body) - content: Final = body["messages"][0]["content"] - assert body == { - "model": PROVIDER_MODEL, - "messages": [{"role": "user", "content": content}], - "max_tokens": MAX_TOKENS, - } - assert content.startswith(probe), content - release.wait(timeout=10) - return Reply(body=UPSTREAM_REPLY) - - with wire_server(respond) as wire, gateway.scenario() as scenario: - model: Final = scenario.model( - model=f"openai/{PROVIDER_MODEL}", - api_base=wire.url + "/v1", - api_key="synthetic-team-tpm-key", - ) - team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_MODEL_TPM}}) - key: Final = scenario.key(team_id=team) - - def send(index: int) -> httpx.Response: - return gateway.request( - "POST", - "/v1/chat/completions", - { - "model": model, - "max_tokens": MAX_TOKENS, - "messages": [{"role": "user", "content": f"{probe} {index}"}], - }, - key=key, - ) - - with ThreadPoolExecutor(max_workers=CONCURRENT_REQUESTS) as pool: - futures: Final[tuple[Future[httpx.Response], ...]] = tuple( - pool.submit(send, index) for index in range(CONCURRENT_REQUESTS) - ) - eventually( - lambda: sum(future.done() for future in futures) + wire.received.qsize(), - lambda settled: settled >= CONCURRENT_REQUESTS, - seconds=10, - ) - release.set() - responses: Final = tuple(future.result(timeout=15) for future in futures) - statuses: Final = tuple(sorted(response.status_code for response in responses)) - assert statuses == (200, 429, 429), tuple(response.text for response in responses) - assert len(wire.drain()) == 1, statuses - served: Final = next(response for response in responses if response.status_code == 200) - assert served.json()["usage"] == {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, served.text - for rejected in (response for response in responses if response.status_code == 429): - error: Final = rejected.json()["error"] - assert (error["type"], error["code"], error["param"]) == ("throttling_error", "429", None), rejected.text - assert f"Limit type: tokens. Current limit: {TEAM_MODEL_TPM}," in error["message"], rejected.text diff --git a/tests/integration/sdk/test_cache_control_census_pydantic_messages_sdk.py b/tests/integration/sdk/test_cache_control_census_pydantic_messages_sdk.py index 327681010cb..44b2933639f 100644 --- a/tests/integration/sdk/test_cache_control_census_pydantic_messages_sdk.py +++ b/tests/integration/sdk/test_cache_control_census_pydantic_messages_sdk.py @@ -1,7 +1,10 @@ import json +from collections.abc import Callable from typing import Final -from integration._support.wire import Wire, wire_server +import pytest +from integration._support.openai_wire import responses_reply +from integration._support.wire import Reply, Request, Wire, wire_server from integration.providers._cache_control_marks_support import ( ANTHROPIC_MODEL, CITIES, @@ -10,14 +13,18 @@ from integration.providers._cache_control_marks_support import ( PROVIDER_KEY, SYSTEM, TOOL, + UNMAPPED_BEDROCK_ARN, anthropic_labels, anthropic_peer, ask, + bedrock_peer, call_id, client_marked, final_text, new_marker, + sent_without_cache_points, ) +from openai.types.responses import FunctionTool from pydantic import JsonValue import litellm @@ -78,3 +85,77 @@ async def test_pydantic_tool_call_marks_count_against_the_cap_in_process_async_s assert "".join(str(chunk.choices[0].delta.content or "") for chunk in chunks) == "sunny", chunks received: Final = wire.drain() assert [anthropic_labels(request) for request in received] == [_CLIENT_MARKED] + + +_OPENAI_MODEL: Final = "gpt-5.6" + + +def _openai_responses_peer(request: Request) -> Reply: + return responses_reply("resp_sdk_typed_tool", _OPENAI_MODEL, "sunny", stream=False) + + +def _outcome(call: Callable[[], object]) -> str: + try: + response: Final = call() + except Exception as error: + return f"{type(error).__name__}: {error}" + return f"{type(response).__name__}: {getattr(response, 'output_text', response)}" + + +def test_injection_points_leave_an_sdk_typed_function_tool_request_as_it_is_without_them() -> None: + tool: Final = FunctionTool(type="function", name="lookup_weather", parameters={"type": "object"}, strict=False) + with wire_server(_openai_responses_peer) as wire: + request: Final[dict[str, object]] = { + "model": f"openai/{_OPENAI_MODEL}", + "input": final_text(new_marker()), + "tools": [tool], + "api_base": wire.url, + "api_key": PROVIDER_KEY, + "num_retries": 0, + } + without_points: Final = _outcome(lambda: litellm.responses(**request)) # pyright: ignore[reportArgumentType] # the SDK-typed tool stands in for the dict shape the signature names + without_points_received: Final = len(wire.drain()) + with_points: Final = _outcome( + lambda: litellm.responses(**request, cache_control_injection_points=POINTS) # pyright: ignore[reportArgumentType] # the SDK-typed tool stands in for the dict shape the signature names + ) + with_points_received: Final = len(wire.drain()) + assert with_points == without_points + assert with_points_received == without_points_received + + +def _shaped_input(shape: str, marker: str) -> object: + system: Final = {"role": "system", "content": SYSTEM, "cache_control": EPHEMERAL} + marked: Final = {"type": "message", "role": "user", "content": final_text(marker), "cache_control": EPHEMERAL} + match shape: + case "string": + return final_text(marker) + case "list": + return [system, marked] + case "tuple": + return (system, marked) + case _: + return marked + + +@pytest.mark.parametrize("shape", ("string", "list", "tuple", "object")) +def test_each_input_shape_reads_the_same_whether_or_not_the_route_keeps_item_marks(shape: str) -> None: + kept_marker, stripped_marker = new_marker(), new_marker() + with wire_server(bedrock_peer) as wire: + request: Final[dict[str, object]] = { + "model": f"bedrock/converse/{UNMAPPED_BEDROCK_ARN}", + "instructions": SYSTEM, + "max_output_tokens": 64, + "api_base": wire.url, + "api_key": PROVIDER_KEY, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": wire.url, + "num_retries": 0, + } + kept: Final = _outcome(lambda: litellm.responses(input=_shaped_input(shape, kept_marker), **request)) # pyright: ignore[reportArgumentType] # the shapes a caller can pass, unvalidated + kept_sent: Final = sent_without_cache_points(wire.drain(), kept_marker) + stripped: Final = _outcome( + lambda: litellm.responses(input=_shaped_input(shape, stripped_marker), drop_params=True, **request) # pyright: ignore[reportArgumentType] # the shapes a caller can pass, unvalidated + ) + stripped_sent: Final = sent_without_cache_points(wire.drain(), stripped_marker) + assert stripped == kept + assert stripped_sent == kept_sent diff --git a/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py b/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py index d1582dd59b9..bd02f2a818b 100644 --- a/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py +++ b/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py @@ -1,5 +1,6 @@ """An SDK app under a prefork worker pool (Celery, gunicorn) with ``PROMETHEUS_MULTIPROC_DIR`` set: the Redis -circuit-breaker state gauge must collect as one series per state across every worker, never one per worker pid.""" +circuit-breaker state gauge must collect as one series per state across every worker, and it must stay off +the registry entirely for an app that never turned on LiteLLM's own Prometheus logging.""" from __future__ import annotations @@ -24,6 +25,12 @@ _CHAT_RESPONSE: Final = json.dumps( } ).encode() _CLOSED: Final = (("state", "closed"),) +_HALF_OPEN: Final = (("state", "half_open"),) +_OPEN: Final = (("state", "open"),) + + +def _all_closed(breakers: float) -> tuple[tuple[tuple[tuple[str, str], ...], float], ...]: + return ((_CLOSED, breakers), (_HALF_OPEN, 0.0), (_OPEN, 0.0)) def _reply(request: Request) -> Reply: @@ -52,6 +59,22 @@ def _worked(report: WorkerReport) -> WorkerReport: def test_two_worker_processes_collect_as_one_state_series_per_state(tmp_path: Path) -> None: with wire_server(_reply) as wire: - reports: Final = tuple(_worked(run_breaker_worker(wire.url, tmp_path)) for _ in range(2)) + reports: Final = tuple( + _worked(run_breaker_worker("callbacks_before_cache", wire.url, tmp_path)) for _ in range(2) + ) assert len({report.pid for report in reports}) == 2, reports - assert _collected_state_series(tmp_path) == ((_CLOSED, 2.0),), reports + assert _collected_state_series(tmp_path) == _all_closed(2.0), reports + + +def test_prometheus_logging_turned_on_after_the_cache_still_reports_the_closed_breaker(tmp_path: Path) -> None: + with wire_server(_reply) as wire: + report: Final = _worked(run_breaker_worker("callbacks_after_cache", wire.url, tmp_path)) + assert report.state_series == _all_closed(1.0), report + assert _collected_state_series(tmp_path) == _all_closed(1.0), report + + +def test_app_without_litellm_prometheus_logging_gets_no_breaker_series(tmp_path: Path) -> None: + with wire_server(_reply) as wire: + report: Final = _worked(run_breaker_worker("no_callbacks", wire.url, tmp_path)) + assert report.state_series == (), report + assert _collected_state_series(tmp_path) == (), report diff --git a/tests/integration/sdk/test_response_cache_sync_sdk.py b/tests/integration/sdk/test_response_cache_sync_sdk.py index 54d78adbaba..1fe73ec9012 100644 --- a/tests/integration/sdk/test_response_cache_sync_sdk.py +++ b/tests/integration/sdk/test_response_cache_sync_sdk.py @@ -9,17 +9,23 @@ from __future__ import annotations import os from collections.abc import Iterator +from itertools import chain from typing import Final import pytest +from pydantic import JsonValue, TypeAdapter import litellm from litellm.caching.caching import Cache +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.caching import LiteLLMCacheType from litellm.types.utils import ModelResponse, TextCompletionResponse +from tests.integration._support.client import JSON_OBJECT from tests.integration.caching.response_cache_case import ( CHAT, + CHAT_CHUNK, CHAT_MODEL, + CHAT_STREAM, PROVIDER_KEY, TEXT, TEXT_MODEL, @@ -31,8 +37,10 @@ from tests.integration.caching.response_cache_case import ( prompt, redis_store, rescript, + rewrite_usage_details, scripted_in_process, ) +from tests.integration.cost_calculation.cost_tracking_case import SseResponse @pytest.fixture @@ -123,3 +131,77 @@ def test_sync_text_completion_with_empty_choices_reaches_upstream_again() -> Non hit: Final = _text_completion(handle.api_base(), text) assert hit.id == refilled.id, hit.model_dump_json() assert upstream.calls(handle.scenario_id) == 2 + + +_REASONING_USAGE_FRAME: Final = ( + CHAT_CHUNK + '[],"usage":{"prompt_tokens":5,"completion_tokens":7,"total_tokens":12,' + '"completion_tokens_details":{"reasoning_tokens":4}}}' +) +_REASONING_STREAM: Final = SseResponse( + content_type="text/event-stream", + frames=(*CHAT_STREAM.frames[:-1], _REASONING_USAGE_FRAME, CHAT_STREAM.frames[-1]), +) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_UNVALIDATED_USAGE_DETAILS: Final[dict[str, JsonValue]] = { + "text": {"reasoning_tokens": "abc"}, + "fraction": {"reasoning_tokens": 1.5}, + "object": {"reasoning_tokens": {"a": 1}}, +} + + +def _stream(api_base: str, text: str) -> tuple[dict[str, JsonValue], ...]: + stream: Final = litellm.completion( + model=CHAT_MODEL, + messages=[{"role": "user", "content": text}], + api_base=api_base, + api_key=PROVIDER_KEY, + stream=True, + stream_options={"include_usage": True}, + ) + assert isinstance(stream, CustomStreamWrapper), type(stream) + return tuple(JSON_OBJECT.validate_python(chunk.model_dump(mode="json")) for chunk in stream) + + +def _replayed(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[frozenset[JsonValue], str, tuple[JsonValue, ...]]: + choices: Final = tuple( + JSON_OBJECT.validate_python(choice) + for choice in chain.from_iterable(_JSON_LIST.validate_python(chunk.get("choices") or []) for chunk in chunks) + ) + text: Final = "".join(str(JSON_OBJECT.validate_python(choice["delta"]).get("content") or "") for choice in choices) + usages: Final = tuple(JSON_OBJECT.validate_python(chunk["usage"]) for chunk in chunks if chunk.get("usage")) + reasoning: Final = tuple( + JSON_OBJECT.validate_python(usage.get("completion_tokens_details") or {}).get("reasoning_tokens") + for usage in usages + ) + return frozenset(chunk["id"] for chunk in chunks), text, reasoning + + +@pytest.mark.usefixtures("sdk_cache") +def test_sync_stream_cache_hit_replays_the_usage_details_of_the_miss() -> None: + upstream: Final = Upstream() + text: Final = prompt() + with scripted_in_process(_REASONING_STREAM) as handle, redis_store() as store: + first: Final = _replayed(_stream(handle.api_base(), text)) + (first_id,) = first[0] + assert first[1:] == ("streamed response", (4,)), first + await_entry(store, str(first_id)) + hit: Final = _replayed(_stream(handle.api_base(), text)) + assert hit == first, hit + assert upstream.calls(handle.scenario_id) == 1 + + +@pytest.mark.usefixtures("sdk_cache") +@pytest.mark.parametrize("variant", tuple(_UNVALIDATED_USAGE_DETAILS)) +def test_sync_stream_cache_hit_skips_cached_usage_details_that_fail_validation(variant: str) -> None: + upstream: Final = Upstream() + text: Final = prompt() + with scripted_in_process(_REASONING_STREAM) as handle, redis_store() as store: + first: Final = _replayed(_stream(handle.api_base(), text)) + (first_id,) = first[0] + assert first[1:] == ("streamed response", (4,)), first + stale_id: Final = rewrite_usage_details( + store, await_entry(store, str(first_id)), _UNVALIDATED_USAGE_DETAILS[variant] + ) + hit: Final = _replayed(_stream(handle.api_base(), text)) + assert hit == (frozenset({stale_id}), "streamed response", (0,)), hit + assert upstream.calls(handle.scenario_id) == 1 diff --git a/tests/integration/sdk/test_responses_bridge_stream_errors_sdk.py b/tests/integration/sdk/test_responses_bridge_stream_errors_sdk.py index 223d6a9eac6..7f13a0fcc46 100644 --- a/tests/integration/sdk/test_responses_bridge_stream_errors_sdk.py +++ b/tests/integration/sdk/test_responses_bridge_stream_errors_sdk.py @@ -6,6 +6,8 @@ from integration._support.responses_stream import ( AZURE_TARGET, RATE_LIMIT_MESSAGE, function_tools, + healthy_stream, + openai_error_stream, rate_limited_stream, serve, ) @@ -13,7 +15,7 @@ from integration._support.wire import Wire, wire_server import litellm from litellm import Router -from litellm.exceptions import MidStreamFallbackError, RateLimitError +from litellm.exceptions import APIError, MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper _MODEL: Final = "azure/gpt-6" @@ -135,3 +137,130 @@ async def test_router_async_stream_in_stream_rate_limit_with_fallbacks_disabled_ async for _ in response: pass _assert_surfaces_the_provider_exception(raised.value, wire, marker) + + +def _openai_error_stream(marker: str, code: str, message: str) -> tuple[bytes, ...]: + return openai_error_stream(f"resp_{marker}", code, message) + + +def test_sync_completion_stream_openai_error_event_raises_the_provider_code_and_message() -> None: + marker: Final = uuid.uuid4().hex + message: Final = f"The server had an error processing your request {marker}" + with wire_server(serve(_openai_error_stream(marker, "server_error", message), AZURE_TARGET)) as wire: + response: Final = litellm.completion( + model=_MODEL, + messages=_messages(marker), + tools=_TOOLS, + stream=True, + num_retries=0, + api_base=wire.url, + api_key=_API_KEY, + ) + with pytest.raises(MidStreamFallbackError) as raised: + for _ in response: + pass + assert type(raised.value.original_exception) is APIError, repr(raised.value.original_exception) + assert raised.value.status_code == 500 and message in str(raised.value), str(raised.value) + _assert_one_attempt(wire, marker) + + +async def test_router_async_stream_openai_error_event_rate_limit_surfaces_the_provider_exception() -> None: + marker: Final = uuid.uuid4().hex + stream: Final = _openai_error_stream(marker, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire: + response: Final = await _router(wire).acompletion( + model=_GROUP, messages=_messages(marker), tools=_TOOLS, stream=True + ) + with pytest.raises(RateLimitError) as raised: + async for _ in response: + pass + _assert_surfaces_the_provider_exception(raised.value, wire, marker) + + +async def test_async_completion_stream_top_level_rate_limit_wraps_the_provider_exception_once() -> None: + marker: Final = uuid.uuid4().hex + stream: Final = _openai_error_stream(marker, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire: + response: Final = await litellm.acompletion( + model=_MODEL, + messages=_messages(marker), + tools=_TOOLS, + stream=True, + num_retries=0, + api_base=wire.url, + api_key=_API_KEY, + ) + with pytest.raises(MidStreamFallbackError) as raised: + async for _ in response: + pass + _assert_wraps_the_provider_exception_once(raised.value, wire, marker) + + +def test_router_sync_stream_top_level_rate_limit_with_fallbacks_disabled_surfaces_the_provider_exception() -> None: + marker: Final = uuid.uuid4().hex + stream: Final = _openai_error_stream(marker, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire: + response: Final = _router(wire).completion( + model=_GROUP, messages=_messages(marker), tools=_TOOLS, stream=True, disable_fallbacks=True + ) + with pytest.raises(RateLimitError) as raised: + for _ in response: + pass + _assert_surfaces_the_provider_exception(raised.value, wire, marker) + + +async def test_router_async_stream_top_level_rate_limit_falls_back_to_the_healthy_deployment() -> None: + marker: Final = uuid.uuid4().hex + limited: Final = _openai_error_stream(marker, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with ( + wire_server(serve(limited, AZURE_TARGET)) as primary, + wire_server(serve(healthy_stream(f"resp_{marker}", "fallback answer"), AZURE_TARGET)) as spare, + ): + router: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": _MODEL, "api_base": wire.url, "api_key": _API_KEY}} + for name, wire in ((_GROUP, primary), ("bridged-spare", spare)) + ], + num_retries=0, + disable_cooldowns=True, + fallbacks=[{_GROUP: ["bridged-spare"]}], + ) + response: Final = await router.acompletion(model=_GROUP, messages=_messages(marker), tools=_TOOLS, stream=True) + content: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in response if chunk.choices]) + assert content == "fallback answer", content + _assert_one_attempt(primary, marker) + _assert_one_attempt(spare, marker) + + +def _provider_rate_limit(raised: Exception, message: str) -> None: + inner: Final = raised.original_exception if isinstance(raised, MidStreamFallbackError) else raised + assert type(inner) is RateLimitError, repr(inner) + assert getattr(raised, "status_code", None) == 429 and message in str(raised), str(raised) + + +def test_sync_responses_stream_top_level_rate_limit_raises_the_provider_rate_limit() -> None: + marker: Final = uuid.uuid4().hex + stream: Final = _openai_error_stream(marker, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire: + response: Final = litellm.responses( + model=_MODEL, input=marker, stream=True, api_base=wire.url, api_key=_API_KEY + ) + with pytest.raises((MidStreamFallbackError, RateLimitError)) as raised: + for _ in response: + pass + _provider_rate_limit(raised.value, RATE_LIMIT_MESSAGE) + _assert_one_attempt(wire, marker) + + +async def test_async_responses_stream_top_level_insufficient_quota_raises_a_rate_limit() -> None: + marker: Final = uuid.uuid4().hex + quota: Final = "You exceeded your current quota, please check your plan and billing details." + with wire_server(serve(_openai_error_stream(marker, "insufficient_quota", quota), AZURE_TARGET)) as wire: + response: Final = await litellm.aresponses( + model=_MODEL, input=marker, stream=True, api_base=wire.url, api_key=_API_KEY + ) + with pytest.raises((MidStreamFallbackError, RateLimitError)) as raised: + async for _ in response: + pass + _provider_rate_limit(raised.value, quota) + _assert_one_attempt(wire, marker) diff --git a/tests/integration/sdk/test_router_sync_stream_fallback_wire.py b/tests/integration/sdk/test_router_sync_stream_fallback_wire.py index fc5fb028123..36616eda50e 100644 --- a/tests/integration/sdk/test_router_sync_stream_fallback_wire.py +++ b/tests/integration/sdk/test_router_sync_stream_fallback_wire.py @@ -154,7 +154,8 @@ def test_walks_every_configured_fallback(client: str) -> None: router: Final = _router(wire, ("primary", "fb1", "fb2"), fallbacks=[{"primary": ["fb1", "fb2"]}]) streamed: Final = _stream(client, router) assert streamed.text == "answered by fb2", streamed - assert streamed.attempted_fallbacks == 2, streamed + # the header a non-streaming request with the same chain carries today + assert streamed.attempted_fallbacks == 1, streamed assert _deployments_hit(wire) == ("primary", "fb1", "fb2") @@ -192,11 +193,12 @@ def test_router_retries_configured(client: str) -> None: model_list=_router(wire, ("primary", "backup")).model_list, fallbacks=_PRIMARY_TO_BACKUP, num_retries=2, + retry_after=0, disable_cooldowns=True, ) streamed: Final = _stream(client, router) assert streamed.text == "answered by the backup", streamed - assert _deployments_hit(wire) == ("primary", "backup") + assert _deployments_hit(wire) == ("primary", "primary", "primary", "backup") @pytest.mark.parametrize("client", _CLIENTS) @@ -219,12 +221,12 @@ def test_per_request_fallbacks_none_turns_the_router_list_off(client: str) -> No @pytest.mark.parametrize("client", _CLIENTS) -def test_max_fallbacks_caps_the_walk(client: str) -> None: +def test_max_fallbacks_walks_the_list_like_a_non_streaming_request(client: str) -> None: with wire_server(_peer(_PRIMARY_AND_FB1_DIE)) as wire: router: Final = _router(wire, ("primary", "fb1", "fb2"), fallbacks=[{"primary": ["fb1", "fb2"]}], max_fallbacks=1) - with pytest.raises(litellm.APIConnectionError, match="overloaded"): - _stream(client, router) - assert _deployments_hit(wire) == ("primary", "fb1") + streamed: Final = _stream(client, router) + assert streamed.text == "answered by fb2", streamed + assert _deployments_hit(wire) == ("primary", "fb1", "fb2") def test_called_inside_a_running_loop() -> None: diff --git a/tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py b/tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py new file mode 100644 index 00000000000..4b9bc8588a7 --- /dev/null +++ b/tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import asyncio +import itertools +import json +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server + +import litellm +from litellm import Router + +_MODEL: Final = "gpt-5.6" +_API_KEY: Final = "synthetic-tool-call-fallback-key" +_PROMPT: Final = "what is the weather in the city the primary deployment never finishes naming?" +_TOOLS: Final = ( + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + }, +) +_ERROR_FRAME: Final = ( + b"data: " + json.dumps({"error": {"message": "overloaded", "type": "server_error", "code": 500}}).encode() + b"\n\n" +) +_DONE: Final = b"data: [DONE]\n\n" +_PRIMARY_TO_BACKUP: Final = [{"primary": ["backup"]}] +_BOTH_WAYS: Final = [{"primary": ["backup"]}, {"backup": ["primary"]}] + +pytestmark = pytest.mark.timeout(120) + + +def _chunk(delta: Mapping[str, object], finish_reason: str | None) -> bytes: + body: Final = { + "id": "chatcmpl-tool-call-fallback-wire", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": _MODEL, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _tool_call_open(call_id: str) -> bytes: + return _chunk( + { + "role": "assistant", + "tool_calls": [ + {"index": 0, "id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": ""}} + ], + }, + None, + ) + + +def _tool_call_arguments(fragment: str) -> bytes: + return _chunk({"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]}, None) + + +_FINISH: Final = _chunk({}, "tool_calls") + + +def _serves_tool_call(call_id: str, city: str) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=( + _tool_call_open(call_id), + _tool_call_arguments('{"city": "'), + _tool_call_arguments(f'{city}"}}'), + _FINISH, + _DONE, + ), + ) + + +def _dies_after_tool_call_opens(call_id: str) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(_tool_call_open(call_id), _tool_call_arguments('{"city": "'), _ERROR_FRAME, _DONE), + ) + + +_DIES_BEFORE_CONTENT: Final = Reply(content_type="text/event-stream", chunks=(_ERROR_FRAME, _DONE)) +_DROPS_BEFORE_CONTENT: Final = Reply(content_type="text/event-stream", chunks=(_DONE,), abort_after=0) +_REJECTS_BEFORE_STREAMING: Final = Reply( + status=500, body=json.dumps({"error": {"message": "overloaded", "type": "server_error", "code": 500}}).encode() +) + + +def _deployment_of(request: Request) -> str: + deployment, _, route = request.target.lstrip("/").partition("/") + assert request.method == "POST", request.method + assert route == "chat/completions", request.target + return deployment + + +def _peer(replies: Mapping[str, Reply]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + body: Final = json.loads(request.body) + assert body.get("stream") is True, body + assert [tool["function"]["name"] for tool in body["tools"]] == ["get_weather"], body + return replies[_deployment_of(request)] + + return respond + + +def _always(reply: Reply) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + _deployment_of(request) + return reply + + return respond + + +def _fails_then_serves(failures: int, served: Reply) -> Callable[[Request], Reply]: + attempts: Final = itertools.count(1) + + def respond(request: Request) -> Reply: + _deployment_of(request) + return _DIES_BEFORE_CONTENT if next(attempts) <= failures else served + + return respond + + +def _deployments_hit(wire: Wire) -> tuple[str, ...]: + return tuple(_deployment_of(request) for request in wire.drain()) + + +def _router(wire: Wire, deployments: tuple[str, ...], **settings: object) -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/{name}", "api_key": _API_KEY}, + } + for name in deployments + ], + disable_cooldowns=True, + **settings, + ) + + +@dataclass(frozen=True, slots=True) +class _ToolCallStream: + call_ids: tuple[str, ...] + arguments: str + finish_reasons: tuple[str, ...] + + +def _tool_call_deltas(chunk: object) -> tuple[object, ...]: + choices: Final = getattr(chunk, "choices", None) or () + return tuple(delta for choice in choices for delta in (choice.delta.tool_calls or ())) + + +def _collect(chunks: tuple[object, ...]) -> _ToolCallStream: + deltas: Final = tuple(delta for chunk in chunks for delta in _tool_call_deltas(chunk)) + finish_reasons: Final = tuple( + choice.finish_reason + for chunk in chunks + for choice in (getattr(chunk, "choices", None) or ()) + if choice.finish_reason + ) + return _ToolCallStream( + call_ids=tuple(delta.id for delta in deltas if delta.id), + arguments="".join(delta.function.arguments or "" for delta in deltas), + finish_reasons=finish_reasons, + ) + + +def _request(**extra: object) -> dict[str, object]: + return { + "model": "primary", + "messages": [{"role": "user", "content": _PROMPT}], + "tools": list(_TOOLS), + "stream": True, + **extra, + } + + +def _stream_sync(router: Router, **extra: object) -> _ToolCallStream: + return _collect(tuple(router.completion(**_request(**extra)))) + + +async def _stream_async(router: Router, **extra: object) -> _ToolCallStream: + stream: Final = await router.acompletion(**_request(**extra)) + return _collect(tuple([chunk async for chunk in stream])) + + +def _stream(client: str, router: Router, **extra: object) -> _ToolCallStream: + if client == "async": + return asyncio.run(_stream_async(router, **extra)) + return _stream_sync(router, **extra) + + +_CLIENTS: Final = ("sync", "async") +_BACKUP_ANSWERS: Final = _ToolCallStream( + call_ids=("call_backup",), arguments='{"city": "Lisbon"}', finish_reasons=("tool_calls",) +) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_primary_dies_before_the_tool_call_opens(client: str) -> None: + replies: Final = {"primary": _DIES_BEFORE_CONTENT, "backup": _serves_tool_call("call_backup", "Lisbon")} + with wire_server(_peer(replies)) as wire: + router: Final = _router(wire, ("primary", "backup"), num_retries=0, fallbacks=_PRIMARY_TO_BACKUP) + assert _stream(client, router) == _BACKUP_ANSWERS + assert _deployments_hit(wire) == ("primary", "backup") + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_primary_drops_the_connection_before_the_tool_call_opens(client: str) -> None: + replies: Final = {"primary": _DROPS_BEFORE_CONTENT, "backup": _serves_tool_call("call_backup", "Lisbon")} + with wire_server(_peer(replies)) as wire: + router: Final = _router(wire, ("primary", "backup"), num_retries=0, fallbacks=_PRIMARY_TO_BACKUP) + assert _stream(client, router) == _BACKUP_ANSWERS + assert _deployments_hit(wire) == ("primary", "backup") + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_primary_dies_after_tool_call_deltas_streamed(client: str) -> None: + replies: Final = { + "primary": _dies_after_tool_call_opens("call_primary"), + "backup": _serves_tool_call("call_backup", "Lisbon"), + } + with wire_server(_peer(replies)) as wire: + router: Final = _router(wire, ("primary", "backup"), num_retries=2, fallbacks=_PRIMARY_TO_BACKUP) + with pytest.raises(litellm.APIConnectionError, match="overloaded"): + _stream(client, router) + assert _deployments_hit(wire) == ("primary",) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_disable_fallbacks_holds_before_the_tool_call_opens(client: str) -> None: + replies: Final = {"primary": _DIES_BEFORE_CONTENT, "backup": _serves_tool_call("call_backup", "Lisbon")} + with wire_server(_peer(replies)) as wire: + router: Final = _router(wire, ("primary", "backup"), num_retries=0, fallbacks=_PRIMARY_TO_BACKUP) + with pytest.raises(litellm.APIConnectionError, match="overloaded"): + _stream(client, router, disable_fallbacks=True) + assert _deployments_hit(wire) == ("primary",) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_primary_dies_before_the_tool_call_opens_is_retried_within_the_group(client: str) -> None: + with wire_server(_fails_then_serves(1, _serves_tool_call("call_primary", "Lisbon"))) as wire: + router: Final = _router(wire, ("primary",), num_retries=1, retry_after=0) + assert _stream(client, router) == _ToolCallStream( + call_ids=("call_primary",), arguments='{"city": "Lisbon"}', finish_reasons=("tool_calls",) + ) + assert _deployments_hit(wire) == ("primary", "primary") + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_request_num_retries_overrides_the_router_budget(client: str) -> None: + with wire_server(_always(_DIES_BEFORE_CONTENT)) as wire: + router: Final = _router(wire, ("primary",), num_retries=3, retry_after=0) + with pytest.raises(litellm.APIConnectionError, match="overloaded"): + _stream(client, router, num_retries=1) + assert _deployments_hit(wire) == ("primary", "primary") + + +@pytest.mark.parametrize("client", _CLIENTS) +@pytest.mark.parametrize("failure", (_DIES_BEFORE_CONTENT, _DROPS_BEFORE_CONTENT), ids=("error-frame", "dropped")) +def test_every_stream_dies_before_content_spends_exactly_the_budget(client: str, failure: Reply) -> None: + num_retries: Final = 2 + with wire_server(_always(failure)) as wire: + router: Final = _router( + wire, ("primary", "backup"), num_retries=num_retries, fallbacks=_BOTH_WAYS, retry_after=0 + ) + with pytest.raises(litellm.APIConnectionError): + _stream(client, router) + assert _deployments_hit(wire) == ("primary",) * (num_retries + 1) + ("backup",) * (num_retries + 1) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_every_request_rejected_before_streaming_spends_a_bounded_budget(client: str) -> None: + num_retries: Final = 1 + with wire_server(_always(_REJECTS_BEFORE_STREAMING)) as wire: + router: Final = _router( + wire, ("primary", "backup"), num_retries=num_retries, fallbacks=_BOTH_WAYS, retry_after=0 + ) + with pytest.raises(litellm.InternalServerError, match="overloaded"): + _stream(client, router) + hit: Final = _deployments_hit(wire) + assert hit == ("primary",) * (num_retries + 1) + ("backup",) * (num_retries + 1), hit diff --git a/tests/integration/sdk/test_router_unsupported_params_fallback_wire.py b/tests/integration/sdk/test_router_unsupported_params_fallback_wire.py new file mode 100644 index 00000000000..e7675fb64ee --- /dev/null +++ b/tests/integration/sdk/test_router_unsupported_params_fallback_wire.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server + +import litellm +from litellm import Router + +_MODEL: Final = "gpt-5-wire-unmapped" +_API_KEY: Final = "synthetic-unsupported-params-key" +_PROMPT: Final = "which deployment answers when the requested model rejects the request?" +_CLIENTS: Final = ("sync", "async") +_PRIMARY_TO_BACKUP: Final = [{"primary": ["backup"]}] +_UNSUPPORTED_BY_PRIMARY: Final = {"reasoning_effort": "xhigh"} + + +def _answers(text: str) -> Reply: + body: Final = { + "id": "chatcmpl-bad-request-wire", + "object": "chat.completion", + "created": 1700000000, + "model": _MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + return Reply(body=json.dumps(body).encode()) + + +_UPSTREAM_500: Final = Reply(status=500, body=json.dumps({"error": {"message": "overloaded"}}).encode()) +_UPSTREAM_400: Final = Reply( + status=400, + body=json.dumps({"error": {"message": "Unsupported value: temperature", "type": "invalid_request_error"}}).encode(), +) +_UPSTREAM_CONTEXT_WINDOW: Final = Reply( + status=400, + body=json.dumps( + { + "error": { + "message": "This model's maximum context length is 8192 tokens", + "type": "invalid_request_error", + "code": "context_length_exceeded", + } + } + ).encode(), +) +_UPSTREAM_CONTENT_POLICY: Final = Reply( + status=400, + body=json.dumps( + { + "error": { + "message": "Your request was rejected", + "type": "invalid_request_error", + "code": "content_policy_violation", + } + } + ).encode(), +) +_BACKUP_ANSWERS: Final = _answers("answered by the backup") + + +def _peer(replies: Mapping[str, Reply]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST", request.method + deployment, _, route = request.target.lstrip("/").partition("/") + assert route == "chat/completions", request.target + return replies[deployment] + + return respond + + +def _deployments_hit(wire: Wire) -> tuple[str, ...]: + return tuple(request.target.lstrip("/").partition("/")[0] for request in wire.drain()) + + +_ALL_DAY_FOR_TEAM_A: Final = ( + {"start": "00:00", "end": "12:00", "timezone": "UTC", "team_ids": ["team-a"]}, + {"start": "12:00", "end": "00:00", "timezone": "UTC", "team_ids": ["team-a"]}, +) + + +def _router(wire: Wire, primary_model_info: Mapping[str, object] | None = None, **settings: object) -> Router: + return Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/primary", "api_key": _API_KEY}, + "model_info": dict(primary_model_info or {}), + }, + { + "model_name": "backup", + "litellm_params": { + "model": f"openai/{_MODEL}", + "api_base": f"{wire.url}/backup", + "api_key": _API_KEY, + "drop_params": True, + }, + }, + ], + num_retries=0, + disable_cooldowns=True, + **{"fallbacks": _PRIMARY_TO_BACKUP, **settings}, + ) + + +_MESSAGES: Final = [{"role": "user", "content": _PROMPT}] + + +def _content(response: litellm.ModelResponse) -> str: + choice: Final = response.choices[0] + assert isinstance(choice, litellm.Choices), choice + return str(choice.message.content) + + +async def _acomplete(router: Router, **request: object) -> str: + response: Final = await router.acompletion(model="primary", messages=_MESSAGES, **request) + assert isinstance(response, litellm.ModelResponse), response + return _content(response) + + +def _complete(client: str, router: Router, **request: object) -> str: + if client == "async": + return asyncio.run(_acomplete(router, **request)) + response: Final = router.completion(model="primary", messages=_MESSAGES, **request) + assert isinstance(response, litellm.ModelResponse), response + return _content(response) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_client_side_param_rejection_reaches_the_caller_without_trying_fallbacks(client: str) -> None: + with wire_server(_peer({"backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router(wire, fallback_on_unsupported_params=False) + with pytest.raises(litellm.UnsupportedParamsError, match="reasoning_effort"): + _complete(client, router, **_UNSUPPORTED_BY_PRIMARY) + assert _deployments_hit(wire) == () + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_client_side_param_rejection_still_falls_back_by_default(client: str) -> None: + with wire_server(_peer({"backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router(wire) + assert _complete(client, router, **_UNSUPPORTED_BY_PRIMARY) == "answered by the backup" + assert _deployments_hit(wire) == ("backup",) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_upstream_500_still_falls_back_when_bad_requests_raise(client: str) -> None: + with wire_server(_peer({"primary": _UPSTREAM_500, "backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router(wire, fallback_on_unsupported_params=False) + assert _complete(client, router) == "answered by the backup" + assert _deployments_hit(wire) == ("primary", "backup") + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_router_raised_400_still_falls_back_when_bad_requests_raise(client: str) -> None: + with wire_server(_peer({"backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router( + wire, primary_model_info={"access_windows": _ALL_DAY_FOR_TEAM_A}, fallback_on_unsupported_params=False + ) + answer: Final = _complete(client, router, metadata={"user_api_key_team_id": "team-b"}) + assert answer == "answered by the backup" + assert _deployments_hit(wire) == ("backup",) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_context_window_error_still_uses_its_own_fallbacks_when_bad_requests_raise(client: str) -> None: + with wire_server(_peer({"primary": _UPSTREAM_CONTEXT_WINDOW, "backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router( + wire, fallback_on_unsupported_params=False, fallbacks=None, context_window_fallbacks=_PRIMARY_TO_BACKUP + ) + assert _complete(client, router) == "answered by the backup" + assert _deployments_hit(wire) == ("primary", "backup") + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_content_policy_error_still_uses_its_own_fallbacks_when_bad_requests_raise(client: str) -> None: + with wire_server(_peer({"primary": _UPSTREAM_CONTENT_POLICY, "backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router( + wire, fallback_on_unsupported_params=False, fallbacks=None, content_policy_fallbacks=_PRIMARY_TO_BACKUP + ) + assert _complete(client, router) == "answered by the backup" + assert _deployments_hit(wire) == ("primary", "backup") + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_upstream_400_reaches_the_caller_without_trying_fallbacks(client: str) -> None: + with wire_server(_peer({"primary": _UPSTREAM_400, "backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router(wire, fallback_on_unsupported_params=False) + with pytest.raises(litellm.BadRequestError, match="Unsupported value: temperature"): + _complete(client, router) + assert _deployments_hit(wire) == ("primary",) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_flag_applied_through_update_settings_stops_fallbacks_on_upstream_400(client: str) -> None: + with wire_server(_peer({"primary": _UPSTREAM_400, "backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router(wire) + router.update_settings(fallback_on_unsupported_params=False) + with pytest.raises(litellm.BadRequestError, match="Unsupported value: temperature"): + _complete(client, router) + assert _deployments_hit(wire) == ("primary",) + + +@pytest.mark.parametrize("client", _CLIENTS) +def test_upstream_400_still_falls_back_by_default(client: str) -> None: + with wire_server(_peer({"primary": _UPSTREAM_400, "backup": _BACKUP_ANSWERS})) as wire: + router: Final = _router(wire) + assert _complete(client, router) == "answered by the backup" + assert _deployments_hit(wire) == ("primary", "backup") diff --git a/tests/integration/streaming/test_responses_bridge_stream_chaos.py b/tests/integration/streaming/test_responses_bridge_stream_chaos.py index 745be2e34b1..ede5a824f52 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_chaos.py +++ b/tests/integration/streaming/test_responses_bridge_stream_chaos.py @@ -27,9 +27,11 @@ from integration._support.database import read_rows from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.responses_stream import ( AZURE_TARGET, + RATE_LIMIT_MESSAGE, chat_content, function_tools, healthy_stream, + openai_error_stream, rate_limited_stream, ) from integration._support.wire import Reply, Request, wire_server @@ -39,15 +41,33 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) _MODEL: Final = "bridged-stream-chaos" _MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") -_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _SENTINEL_PREFIX: Final = "litellm.MidStreamFallbackError: " _RATE_LIMIT_PREFIX: Final = "litellm.RateLimitError: " _HEADERS: Final[Mapping[str, str]] = MappingProxyType({"anthropic-version": "2023-06-01"}) -Kind: TypeAlias = Literal["chat_limited", "chat_healthy", "messages_limited", "responses_limited"] +Kind: TypeAlias = Literal[ + "chat_limited", + "chat_healthy", + "messages_limited", + "responses_limited", + "chat_top_level", + "messages_top_level", + "responses_top_level", + "chat_dropped", +] _KINDS: Final[tuple[Kind, ...]] = ("chat_limited", "chat_healthy", "messages_limited", "responses_limited") +_TOP_LEVEL_KINDS: Final[tuple[Kind, ...]] = ( + "chat_top_level", + "chat_healthy", + "messages_top_level", + "responses_top_level", + "chat_dropped", +) _CHAT_KINDS: Final[tuple[Kind, ...]] = ("chat_limited", "chat_healthy") -_LOGGED_KINDS: Final[frozenset[Kind]] = frozenset({"chat_limited", "chat_healthy", "responses_limited"}) +_TOP_LEVEL_CHAT_KINDS: Final[tuple[Kind, ...]] = ("chat_top_level", "chat_healthy") +_LOGGED_KINDS: Final[frozenset[Kind]] = frozenset( + {"chat_limited", "chat_healthy", "responses_limited", "chat_top_level", "responses_top_level", "chat_dropped"} +) @dataclass(frozen=True, slots=True) @@ -128,16 +148,25 @@ def _respond(request: Request) -> Reply: identity: Final = f"resp_{uuid.uuid4().hex}" if b"chat_healthy" in request.body: return Reply(content_type="text/event-stream", chunks=healthy_stream(identity, f"answer marker-{marker}")) + if b"chat_dropped" in request.body: + return Reply( + content_type="text/event-stream", chunks=healthy_stream(identity, f"answer marker-{marker}"), abort_after=1 + ) + if b"_top_level" in request.body: + return Reply( + content_type="text/event-stream", + chunks=openai_error_stream(identity, "rate_limit_exceeded", RATE_LIMIT_MESSAGE), + ) return Reply(content_type="text/event-stream", chunks=rate_limited_stream(identity)) def _path(kind: Kind) -> str: match kind: - case "chat_limited" | "chat_healthy": + case "chat_limited" | "chat_healthy" | "chat_top_level" | "chat_dropped": return "/v1/chat/completions" - case "messages_limited": + case "messages_limited" | "messages_top_level": return "/v1/messages" - case "responses_limited": + case "responses_limited" | "responses_top_level": return "/v1/responses" @@ -150,9 +179,9 @@ def _body(call: _Call) -> Mapping[str, JsonValue]: "cache": {"no-cache": True}, } match call.kind: - case "chat_limited" | "chat_healthy": + case "chat_limited" | "chat_healthy" | "chat_top_level" | "chat_dropped": return {**common, "messages": [{"role": "user", "content": prompt}], "tools": function_tools()} - case "messages_limited": + case "messages_limited" | "messages_top_level": return { **common, "max_tokens": 64, @@ -165,7 +194,7 @@ def _body(call: _Call) -> Mapping[str, JsonValue]: } ], } - case "responses_limited": + case "responses_limited" | "responses_top_level": return {**common, "input": prompt} @@ -216,13 +245,14 @@ def _assert_answered_in_its_own_shape(served: _Served) -> None: case "chat_healthy": assert served.status == 200, served.text assert chat_content(_data_frames(served.text)) == f"answer marker-{served.call.marker}", served.text - case "chat_limited": + case "chat_limited" | "chat_top_level": assert served.status == 429, served.text error: Final = object_value(JSON_OBJECT.validate_json(served.text)["error"]) assert error["type"] == "throttling_error" and str(error["code"]) == "429", error message: Final = string_value(error["message"]) assert message.startswith(_RATE_LIMIT_PREFIX) and _SENTINEL_PREFIX not in message, message - case "messages_limited": + assert RATE_LIMIT_MESSAGE in message, message + case "messages_limited" | "messages_top_level": assert served.status == 200, served.text events: Final = _sse_events(served.text) assert events[0][0] == "message_start" and events[-1][0] == "error", events @@ -230,10 +260,21 @@ def _assert_answered_in_its_own_shape(served: _Served) -> None: assert frame_error["type"] == "rate_limit_error", frame_error frame_message: Final = string_value(frame_error["message"]) assert frame_message.startswith(_RATE_LIMIT_PREFIX) and _SENTINEL_PREFIX not in frame_message, frame_message + assert RATE_LIMIT_MESSAGE in frame_message, frame_message case "responses_limited": assert served.status == 200, served.text kinds: Final = [frame["type"] for frame in _data_frames(served.text)] assert kinds == ["response.created", "response.failed"], served.text + case "responses_top_level": + assert served.status == 200, served.text + frames: Final = _data_frames(served.text) + assert [frame["type"] for frame in frames] == ["response.created", "response.failed"], served.text + failure: Final = object_value(object_value(frames[-1]["response"])["error"]) + assert failure["code"] == "rate_limit_exceeded", failure + assert RATE_LIMIT_MESSAGE in string_value(failure["message"]), failure + case "chat_dropped": + assert served.status == 500, served.text + assert "error" in JSON_OBJECT.validate_json(served.text), served.text def _assert_forwarded(received: tuple[Request, ...], calls: tuple[_Call, ...]) -> None: @@ -263,12 +304,8 @@ def _assert_each_lands_once(served: tuple[_Served, ...]) -> None: assert by_call[item.call_id]["status"] == expected, (item.call_id, rows) -def _worker_pids(log: Path, count: int) -> tuple[int, ...]: - return eventually( - lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(log.read_text())), - lambda pids: len(pids) == count, - seconds=30, - ) +def _worker_pids(proxy: OwnedProxy, count: int) -> frozenset[int]: + return eventually(proxy.worker_pids, lambda pids: len(pids) == count, seconds=30) def _open_upstream_connections(pid: int, upstream: str) -> int: @@ -299,24 +336,28 @@ async def _held_burst(gateway: Gateway, calls: tuple[_Call, ...], held: _Held) - return burst +@pytest.mark.parametrize("kinds", [_KINDS, _TOP_LEVEL_KINDS], ids=["nested", "top_level"]) async def test_mixed_burst_of_bridged_streams_answers_each_call_in_its_own_shape_and_logs_each_once( - rig: _Rig, + rig: _Rig, kinds: tuple[Kind, ...] ) -> None: - calls: Final = _calls(24, _KINDS) + calls: Final = _calls(6 * len(kinds), kinds) with wire_server(_respond, port=rig.port) as wire: served: Final = await _burst(rig.gateway, calls) - assert len(served) == 24 + assert len(served) == len(calls) for item in served: _assert_answered_in_its_own_shape(item) _assert_forwarded(wire.drain(), calls) _assert_each_lands_once(served) -async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_the_bridged_streams(rig: _Rig) -> None: - calls: Final = _calls(20, _CHAT_KINDS) +@pytest.mark.parametrize("kinds", [_CHAT_KINDS, _TOP_LEVEL_CHAT_KINDS], ids=["nested", "top_level"]) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_the_bridged_streams( + rig: _Rig, kinds: tuple[Kind, ...] +) -> None: + calls: Final = _calls(20, kinds) held: Final = _Held(threading.Event(), SimpleQueue()) with wire_server(held.respond, port=rig.port) as wire: - workers: Final = _worker_pids(rig.proxy.log, 2) + workers: Final = _worker_pids(rig.proxy, 2) burst: Final = await _held_burst(rig.gateway, calls, held) held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) assert sum(held_by.values()) == 20, held_by diff --git a/tests/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index 162ddc07bad..d8d4a44cd68 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -13,104 +13,47 @@ import openai import pytest import yaml from integration._support.client import ( - JSON_OBJECT, Gateway, - eventually, gateway_from_environment, object_value, string_value, ) -from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.responses_stream import ( AZURE_TARGET, OPENAI_TARGET, RATE_LIMIT_MESSAGE, + RATE_LIMIT_PREFIX, + SENTINEL_PREFIX, + assert_provider_typed_rate_limit, azure_rate_limit, + chat_body, chat_content, created, + data_frames, delta, + error_body, error_event, failed, frame, function_tools, healthy_stream, + messages_body, + openai_error_event, + openai_error_stream, rate_limited_stream, serve, + spend_row, + sse_events, ) from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue -_RATE_LIMIT_PREFIX: Final = "litellm.RateLimitError: " -_SENTINEL_PREFIX: Final = "litellm.MidStreamFallbackError: " _PRIMARY: Final = "bridged-primary" _SPARE: Final = "bridged-spare" pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) -def chat_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: - return { - "model": model, - "stream": True, - "messages": [{"role": "user", "content": marker}], - "tools": function_tools(), - "num_retries": 0, - "cache": {"no-cache": True}, - **extra, - } - - -def messages_body(model: str, marker: str) -> dict[str, JsonValue]: - return { - "model": model, - "max_tokens": 64, - "stream": True, - "messages": [{"role": "user", "content": marker}], - "tools": [ - { - "name": "get_weather", - "description": "Weather for a city", - "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, - } - ], - "num_retries": 0, - "cache": {"no-cache": True}, - } - - -def error_body(response: httpx.Response) -> Mapping[str, JsonValue]: - return object_value(JSON_OBJECT.validate_json(response.content)["error"]) - - -def data_frames(text: str) -> tuple[Mapping[str, JsonValue], ...]: - return tuple( - JSON_OBJECT.validate_json(line.removeprefix("data: ")) - for line in text.splitlines() - if line.startswith("data: ") and line != "data: [DONE]" - ) - - -def spend_row(call_id: str) -> Mapping[str, JsonValue]: - rows: Final = eventually( - lambda: read_rows( - 'SELECT litellm_call_id, status, model_group FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', - (call_id,), - ), - lambda found: len(found) >= 1, - seconds=70, - ) - assert len(rows) == 1, rows - return rows[0] - - -def assert_provider_typed_rate_limit(error: Mapping[str, JsonValue]) -> None: - assert error["type"] == "throttling_error", error - assert str(error["code"]) == "429", error - message: Final = string_value(error["message"]) - assert message.startswith(_RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message - assert _SENTINEL_PREFIX not in message, message - - def assert_failed_once(wire: Wire, call_id: str, model: str, attempts: int = 1) -> tuple[Request, ...]: received: Final = wire.drain() assert len(received) == attempts, [request.target for request in received] @@ -237,6 +180,55 @@ def test_bridged_in_stream_rate_limit_whose_fallback_is_also_rate_limited_answer assert spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" +def test_bridged_top_level_rate_limit_falls_back_to_the_healthy_deployment(fallback_proxy: FallbackProxy) -> None: + identity: Final = "resp_" + uuid.uuid4().hex + limited: Final = openai_error_stream(identity, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with ( + wire_server(serve(limited, AZURE_TARGET), port=fallback_proxy.primary_port) as primary, + wire_server( + serve(healthy_stream(identity, "fallback answer"), AZURE_TARGET), port=fallback_proxy.spare_port + ) as spare, + ): + response: Final = _chat(fallback_proxy.gateway, chat_body(_PRIMARY, identity)) + assert response.status_code == 200, response.text + assert chat_content(data_frames(response.text)) == "fallback answer", response.text + assert len(primary.drain()) == 1 and len(spare.drain()) == 1 + row: Final = spend_row(response.headers["x-litellm-call-id"]) + assert row["status"] == "success" and row["model_group"] == _SPARE, row + + +_INVALID_PROMPT: Final = "Invalid prompt: your prompt was flagged." + + +@pytest.mark.parametrize( + "rejection", + [ + pytest.param( + error_event({"type": "invalid_request_error", "code": "invalid_prompt", "message": _INVALID_PROMPT}), + id="nested", + ), + pytest.param(openai_error_event("invalid_prompt", _INVALID_PROMPT), id="top_level"), + ], +) +def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallback( + fallback_proxy: FallbackProxy, rejection: Mapping[str, JsonValue] +) -> None: + identity: Final = "resp_" + uuid.uuid4().hex + with ( + wire_server( + serve((frame(created(identity)), frame(rejection)), AZURE_TARGET), port=fallback_proxy.primary_port + ) as primary, + wire_server( + serve(healthy_stream(identity, "fallback answer"), AZURE_TARGET), port=fallback_proxy.spare_port + ) as spare, + ): + response: Final = _chat(fallback_proxy.gateway, chat_body(_PRIMARY, identity)) + assert response.status_code == 400, response.text + message: Final = string_value(error_body(response)["message"]) + assert message.startswith("litellm.BadRequestError: ") and _INVALID_PROMPT in message, message + assert len(primary.drain()) == 1 and len(spare.drain()) == 0 + + def _in_stream_error_status(gateway: Gateway, stream: tuple[bytes, ...]) -> tuple[httpx.Response, str]: with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model="azure/gpt-6", api_base=wire.url, api_key="synthetic-azure-key") @@ -258,7 +250,7 @@ def test_bridged_in_stream_server_error_reaches_the_client_as_the_provider_error message: Final = string_value(error["message"]) assert str(error["code"]) == "500", error assert message.startswith("litellm.APIError: ") and "The server had an error" in message, message - assert _SENTINEL_PREFIX not in message, message + assert SENTINEL_PREFIX not in message, message def test_bridged_in_stream_invalid_prompt_is_a_bad_request_on_both_legs(gateway: Gateway) -> None: @@ -274,7 +266,7 @@ def test_bridged_in_stream_invalid_prompt_is_a_bad_request_on_both_legs(gateway: message: Final = string_value(error["message"]) assert str(error["code"]) == "400", error assert message.startswith("litellm.BadRequestError: ") and "Invalid prompt" in message, message - assert _SENTINEL_PREFIX not in message, message + assert SENTINEL_PREFIX not in message, message def test_bridged_error_event_without_an_error_object_is_a_provider_typed_internal_error(gateway: Gateway) -> None: @@ -286,7 +278,7 @@ def test_bridged_error_event_without_an_error_object_is_a_provider_typed_interna message: Final = string_value(error["message"]) assert str(error["code"]) == "500", error assert message.startswith("litellm.APIError: ") and "Response API in-stream error" in message, message - assert _SENTINEL_PREFIX not in message, message + assert SENTINEL_PREFIX not in message, message def test_bridged_error_event_with_a_numeric_code_is_a_throttling_error(gateway: Gateway) -> None: @@ -297,6 +289,14 @@ def test_bridged_error_event_with_a_numeric_code_is_a_throttling_error(gateway: assert_provider_typed_rate_limit(error_body(response)) +def test_bridged_openai_error_event_keeps_the_provider_code_and_message(gateway: Gateway) -> None: + identity: Final = "resp_" + uuid.uuid4().hex + stream: Final = (frame(created(identity)), frame(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE))) + response, _ = _in_stream_error_status(gateway, stream) + assert response.status_code == 429, response.text + assert_provider_typed_rate_limit(error_body(response)) + + def test_bridged_response_failed_without_an_error_event_is_a_throttling_error(gateway: Gateway) -> None: identity: Final = "resp_" + uuid.uuid4().hex stream: Final = (frame(created(identity)), frame(failed(identity, "rate_limit_exceeded", RATE_LIMIT_MESSAGE))) @@ -323,8 +323,8 @@ def test_bridged_rate_limit_after_output_is_a_provider_typed_error_frame_behind_ assert str(error["code"]) == "429", error assert error["type"] == "throttling_error", error message: Final = string_value(error["message"]) - assert message.startswith(_RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message - assert _SENTINEL_PREFIX not in message, message + assert message.startswith(RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message + assert SENTINEL_PREFIX not in message, message def test_bridged_transport_drop_after_response_created_is_a_500_without_the_sentinel(gateway: Gateway) -> None: @@ -346,7 +346,7 @@ def test_bridged_transport_drop_after_response_created_is_a_500_without_the_sent message: Final = string_value(error["message"]) assert str(error["code"]) == "500", error assert "never sent" not in response.text - assert _SENTINEL_PREFIX not in message, message + assert SENTINEL_PREFIX not in message, message assert_failed_once(wire, response.headers["x-litellm-call-id"], model) @@ -372,28 +372,18 @@ def test_plain_chat_http_rate_limit_is_a_throttling_error_on_both_legs(gateway: error: Final = object_value(raised.value.body) assert error["type"] == "throttling_error" and str(error["code"]) == "429", error message: Final = string_value(error["message"]) - assert message.startswith(_RATE_LIMIT_PREFIX) and "Rate limit reached" in message, message + assert message.startswith(RATE_LIMIT_PREFIX) and "Rate limit reached" in message, message assert_failed_once(wire, raised.value.response.headers["x-litellm-call-id"], model) -def sse_events(text: str) -> tuple[tuple[str, Mapping[str, JsonValue]], ...]: - def parse(block: str) -> tuple[str, Mapping[str, JsonValue]]: - lines: Final = block.splitlines() - event: Final = next(line.removeprefix("event: ") for line in lines if line.startswith("event: ")) - data: Final = next(line.removeprefix("data: ") for line in lines if line.startswith("data: ")) - return event, JSON_OBJECT.validate_json(data) - - return tuple(parse(block) for block in text.strip().split("\n\n") if "event: " in block) - - def assert_messages_errorframe(events: Sequence[tuple[str, Mapping[str, JsonValue]]]) -> None: assert events[0][0] == "message_start", events assert events[-1][0] == "error", events error: Final = object_value(events[-1][1]["error"]) assert error["type"] == "rate_limit_error", error message: Final = string_value(error["message"]) - assert message.startswith(_RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message - assert message.count(_RATE_LIMIT_PREFIX) == 1 and _SENTINEL_PREFIX not in message, message + assert message.startswith(RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message + assert message.count(RATE_LIMIT_PREFIX) == 1 and SENTINEL_PREFIX not in message, message def test_messages_over_the_bridged_stream_carry_the_provider_error_once_in_the_errorframe(gateway: Gateway) -> None: diff --git a/tests/integration/streaming/test_responses_top_level_error_events.py b/tests/integration/streaming/test_responses_top_level_error_events.py new file mode 100644 index 00000000000..63bdadceebc --- /dev/null +++ b/tests/integration/streaming/test_responses_top_level_error_events.py @@ -0,0 +1,407 @@ +import json +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, object_value, string_value +from integration._support.responses_stream import ( + AZURE_TARGET, + OPENAI_TARGET, + RATE_LIMIT_MESSAGE, + RATE_LIMIT_PREFIX, + SENTINEL_PREFIX, + assert_provider_typed_rate_limit, + chat_body, + chat_content, + created, + data_frames, + delta, + error_body, + frame, + function_tools, + healthy_stream, + messages_body, + openai_error_event, + openai_error_stream, + serve, + spend_row, + sse_events, +) +from integration._support.wire import Wire, wire_server +from openai.types.responses import ResponseFailedEvent, ResponseStreamEvent +from pydantic import JsonValue + +_AZURE_KEY: Final = "synthetic-azure-key" +_QUOTA_MESSAGE: Final = "You exceeded your current quota, please check your plan and billing details." +_SERVER_MESSAGE: Final = "The server had an error while processing your request. Sorry about that!" +_ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} + + +@dataclass(frozen=True, slots=True) +class _Served: + response: httpx.Response + wire: Wire + + +def _azure_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model(model="azure/gpt-6", api_base=wire.url, api_key=_AZURE_KEY) + + +def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue], stream: tuple[bytes, ...]) -> _Served: + with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = _azure_model(scenario, wire) + response: Final = gateway.request("POST", path, {**body, "model": model}, headers=_ANTHROPIC_HEADERS) + return _Served(response, wire) + + +def _chat(gateway: Gateway, stream: tuple[bytes, ...]) -> tuple[httpx.Response, str]: + marker: Final = uuid.uuid4().hex + served: Final = _post(gateway, "/v1/chat/completions", chat_body("", marker), stream) + _assert_one_attempt(served.wire, marker) + return served.response, marker + + +def _assert_one_attempt(wire: Wire, marker: str) -> None: + received: Final = wire.drain() + assert len(received) == 1 and marker.encode() in received[0].body, [request.target for request in received] + + +def _assert_failure_logged(call_id: str, logged_code: str, provider_message: str) -> None: + row: Final = spend_row(call_id) + assert row["status"] == "failure", row + error: Final = object_value(object_value(row["metadata"])["error_information"]) + assert str(error["error_code"]) == logged_code, error + assert provider_message in string_value(error["error_message"]), error + + +def _identity() -> str: + return "resp_" + uuid.uuid4().hex + + +def test_openai_sdk_chat_stream_top_level_rate_limit_is_a_throttling_error_with_the_provider_message( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = _azure_model(scenario, wire) + client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + with pytest.raises(openai.RateLimitError) as raised: + client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": marker}], + tools=function_tools(), + stream=True, + extra_body={"num_retries": 0, "cache": {"no-cache": True}}, + ) + assert raised.value.status_code == 429 + assert_provider_typed_rate_limit(object_value(raised.value.body)) + _assert_one_attempt(wire, marker) + _assert_failure_logged(raised.value.response.headers["x-litellm-call-id"], "429", RATE_LIMIT_MESSAGE) + + +async def test_async_openai_sdk_codex_chat_stream_top_level_rate_limit_is_a_throttling_error(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, OPENAI_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-5.3-codex", api_base=wire.url, api_key="synthetic-openai-key") + client: Final = openai.AsyncOpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + with pytest.raises(openai.RateLimitError) as raised: + await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": marker}], + stream=True, + extra_body={"num_retries": 0, "cache": {"no-cache": True}}, + ) + assert raised.value.status_code == 429 + assert_provider_typed_rate_limit(object_value(raised.value.body)) + _assert_one_attempt(wire, marker) + + +@pytest.mark.parametrize( + ("code", "message", "status"), + [ + pytest.param("insufficient_quota", _QUOTA_MESSAGE, 429, id="insufficient_quota"), + pytest.param("server_error", _SERVER_MESSAGE, 500, id="server_error"), + pytest.param("invalid_prompt", "Invalid prompt: your prompt was flagged.", 400, id="invalid_prompt"), + pytest.param( + "context_length_exceeded", + "Your input exceeds the context window of this model.", + 400, + id="context_length_exceeded", + ), + pytest.param("503", "The engine is currently overloaded.", 503, id="decimal_503"), + pytest.param("brand_new_error_code", "Something new went wrong upstream.", 500, id="unknown_code"), + pytest.param("", "An error with an empty code.", 500, id="empty_code"), + ], +) +def test_top_level_error_code_sets_the_chat_status_and_keeps_the_provider_message( + gateway: Gateway, code: str, message: str, status: int +) -> None: + response, _ = _chat(gateway, openai_error_stream(_identity(), code, message)) + assert response.status_code == status, response.text + error: Final = error_body(response) + assert str(error["code"]) == str(status), error + client_message: Final = string_value(error["message"]) + assert message in client_message and SENTINEL_PREFIX not in client_message, client_message + logged_code: Final = code if status == 400 else str(status) + _assert_failure_logged(response.headers["x-litellm-call-id"], logged_code, message) + + +def test_top_level_rate_limit_after_text_is_a_throttling_error_frame_behind_the_text(gateway: Gateway) -> None: + identity: Final = _identity() + stream: Final = ( + frame(created(identity)), + frame(delta(identity, "Hello")), + frame(delta(identity, " there")), + frame(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE)), + ) + response, _ = _chat(gateway, stream) + assert response.status_code == 200, response.text + frames: Final = data_frames(response.text) + assert chat_content(frames) == "Hello there", response.text + assert_provider_typed_rate_limit(object_value(frames[-1]["error"])) + + +def test_nested_error_object_wins_over_top_level_fields(gateway: Gateway) -> None: + event: Final = { + **openai_error_event("rate_limit_exceeded", "top-level says slow down"), + "error": {"type": "server_error", "code": "server_error", "message": "nested says the server broke"}, + } + response, _ = _chat(gateway, (frame(created(_identity())), frame(event))) + assert response.status_code == 500, response.text + message: Final = string_value(error_body(response)["message"]) + assert "nested says the server broke" in message and "top-level says" not in message, message + + +def _raw_event(body: str) -> bytes: + return f"event: error\ndata: {body}\n\n".encode() + + +_FIVE_KB_MESSAGE: Final = "Rate limit reached for gpt-6. " + "x" * 5120 + + +@pytest.mark.parametrize( + ("event", "provider_message"), + [ + pytest.param( + frame({**openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE), "error": None}), + RATE_LIMIT_MESSAGE, + id="error_null", + ), + pytest.param( + frame({"type": "error", "sequence_number": 1, "code": 429, "message": RATE_LIMIT_MESSAGE, "param": None}), + RATE_LIMIT_MESSAGE, + id="integer_code", + ), + pytest.param(frame(openai_error_event("rate_limit_exceeded", "")), "", id="empty_message"), + pytest.param( + frame({"type": "error", "sequence_number": 1, "code": "rate_limit_exceeded", "message": {"text": "x"}}), + "", + id="object_message", + ), + pytest.param(frame(openai_error_event("rate_limit_exceeded", _FIVE_KB_MESSAGE)), _FIVE_KB_MESSAGE, id="5kb"), + pytest.param( + _raw_event( + '{"type": "error", "sequence_number": 1, "code": "rate_limit_exceeded", ' + f'"code": "rate_limit_exceeded", "message": {json.dumps(RATE_LIMIT_MESSAGE)}, "param": null}}' + ), + RATE_LIMIT_MESSAGE, + id="duplicated_code", + ), + ], +) +def test_malformed_top_level_fields_still_answer_a_throttling_error( + gateway: Gateway, event: bytes, provider_message: str +) -> None: + response, _ = _chat(gateway, (frame(created(_identity())), event)) + assert response.status_code == 429, response.text + assert_provider_typed_rate_limit(error_body(response), provider_message) + healthy: Final = _post( + gateway, "/v1/chat/completions", chat_body("", "after"), healthy_stream(_identity(), "still serving") + ) + assert chat_content(data_frames(healthy.response.text)) == "still serving", healthy.response.text + + +def _cached_chat(gateway: Gateway, stream: tuple[bytes, ...]) -> tuple[tuple[httpx.Response, ...], int]: + marker: Final = uuid.uuid4().hex + body: Final = {key: value for key, value in chat_body("", marker).items() if key != "cache"} + with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = _azure_model(scenario, wire) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", {**body, "model": model}) for _ in range(2) + ) + return responses, len(wire.drain()) + + +def test_a_top_level_error_stream_is_never_served_from_the_response_cache(gateway: Gateway) -> None: + responses, attempts = _cached_chat( + gateway, openai_error_stream(_identity(), "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + ) + assert attempts == 2, attempts + for response in responses: + assert response.status_code == 429, response.text + assert_provider_typed_rate_limit(error_body(response)) + + +def _responses_frames(gateway: Gateway, stream: tuple[bytes, ...]) -> tuple[httpx.Response, Mapping[str, JsonValue]]: + marker: Final = uuid.uuid4().hex + served: Final = _post( + gateway, "/v1/responses", {"input": marker, "stream": True, "cache": {"no-cache": True}}, stream + ) + _assert_one_attempt(served.wire, marker) + assert served.response.status_code == 200, served.response.text + assert served.response.text.rstrip().endswith("data: [DONE]"), served.response.text + frames: Final = data_frames(served.response.text) + assert frames[-1]["type"] == "response.failed", served.response.text + return served.response, object_value(object_value(frames[-1]["response"])["error"]) + + +def _assert_rate_limit_failure(error: Mapping[str, JsonValue]) -> None: + assert error["code"] == "rate_limit_exceeded", error + assert RATE_LIMIT_MESSAGE in string_value(error["message"]), error + + +def test_native_responses_stream_fails_with_the_provider_code_and_message(gateway: Gateway) -> None: + identity: Final = _identity() + response, error = _responses_frames( + gateway, openai_error_stream(identity, "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + ) + assert [frame["type"] for frame in data_frames(response.text)] == ["response.created", "response.failed"] + _assert_rate_limit_failure(error) + _assert_failure_logged(response.headers["x-litellm-call-id"], "429", RATE_LIMIT_MESSAGE) + + +def test_native_responses_stream_error_after_text_fails_with_the_provider_code(gateway: Gateway) -> None: + identity: Final = _identity() + stream: Final = ( + frame(created(identity)), + frame(delta(identity, "partial")), + frame(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE)), + ) + response, error = _responses_frames(gateway, stream) + kinds: Final = [frame["type"] for frame in data_frames(response.text)] + assert kinds == ["response.created", "response.output_text.delta", "response.failed"], kinds + _assert_rate_limit_failure(error) + + +def _failed_error(events: list[ResponseStreamEvent]) -> tuple[str, str]: + assert [event.type for event in events] == ["response.created", "response.failed"], events + failed: Final = events[-1] + assert isinstance(failed, ResponseFailedEvent) and failed.response.error is not None, failed + return failed.response.error.code, failed.response.error.message + + +def test_openai_sdk_responses_stream_fails_with_insufficient_quota(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "insufficient_quota", _QUOTA_MESSAGE) + with wire_server(serve(stream, OPENAI_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-5.3-codex", api_base=wire.url, api_key="synthetic-openai-key") + client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + events: Final = list( + client.responses.create(model=model, input=marker, stream=True, extra_body={"cache": {"no-cache": True}}) + ) + code, message = _failed_error(events) + assert code == "insufficient_quota" and _QUOTA_MESSAGE in message, events + _assert_one_attempt(wire, marker) + + +async def test_async_openai_sdk_responses_stream_fails_with_the_server_error_message(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "server_error", _SERVER_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = _azure_model(scenario, wire) + client: Final = openai.AsyncOpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + response: Final = await client.responses.create( + model=model, input=marker, stream=True, extra_body={"cache": {"no-cache": True}} + ) + events: Final = [event async for event in response] + code, message = _failed_error(events) + assert code == "server_error" and _SERVER_MESSAGE in message, events + _assert_one_attempt(wire, marker) + + +def _assert_messages_error(error: Mapping[str, JsonValue], error_type: str, provider_message: str) -> None: + assert error["type"] == error_type, error + message: Final = string_value(error["message"]) + assert provider_message in message and SENTINEL_PREFIX not in message, message + + +def test_messages_stream_top_level_rate_limit_is_a_rate_limit_errorframe(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + served: Final = _post(gateway, "/v1/messages", messages_body("", marker), stream) + assert served.response.status_code == 200, served.response.text + events: Final = sse_events(served.response.text) + assert events[0][0] == "message_start" and events[-1][0] == "error", events + error: Final = object_value(events[-1][1]["error"]) + _assert_messages_error(error, "rate_limit_error", RATE_LIMIT_MESSAGE) + assert string_value(error["message"]).count(RATE_LIMIT_PREFIX) == 1, error + _assert_one_attempt(served.wire, marker) + + +def _anthropic_tools() -> list[anthropic.types.ToolParam]: # mutable-ok: the Anthropic SDK tools parameter is a list + return [ + { + "name": "get_weather", + "description": "Weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ] + + +async def _drain_async_messages_stream(client: anthropic.AsyncAnthropic, model: str, marker: str) -> None: + async with client.messages.stream( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + tools=_anthropic_tools(), + extra_body={"num_retries": 0, "cache": {"no-cache": True}}, + ) as events: + async for _ in events: + pass + + +def _drain_messages_stream(client: anthropic.Anthropic, model: str, marker: str) -> None: + with client.messages.stream( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + tools=_anthropic_tools(), + extra_body={"num_retries": 0, "cache": {"no-cache": True}}, + ) as events: + for _ in events: + pass + + +async def test_async_anthropic_sdk_stream_raises_the_top_level_rate_limit(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "rate_limit_exceeded", RATE_LIMIT_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = _azure_model(scenario, wire) + client: Final = anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) + with pytest.raises(anthropic.APIStatusError) as raised: + await _drain_async_messages_stream(client, model, marker) + _assert_messages_error( + object_value(object_value(raised.value.body)["error"]), "rate_limit_error", RATE_LIMIT_MESSAGE + ) + _assert_one_attempt(wire, marker) + + +def test_sync_anthropic_sdk_stream_raises_the_top_level_server_error_message(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + stream: Final = openai_error_stream(_identity(), "server_error", _SERVER_MESSAGE) + with wire_server(serve(stream, AZURE_TARGET)) as wire, gateway.scenario() as scenario: + model: Final = _azure_model(scenario, wire) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + with pytest.raises(anthropic.APIStatusError) as raised: + _drain_messages_stream(client, model, marker) + _assert_messages_error(object_value(object_value(raised.value.body)["error"]), "api_error", _SERVER_MESSAGE) + _assert_one_attempt(wire, marker) diff --git a/tests/integration/translation/chat_completions/bases/openrouter.py b/tests/integration/translation/chat_completions/bases/openrouter.py index 4823e01cc28..48c189b1d42 100644 --- a/tests/integration/translation/chat_completions/bases/openrouter.py +++ b/tests/integration/translation/chat_completions/bases/openrouter.py @@ -57,10 +57,11 @@ CLAUDE_SONNET_4_TEST_CASE: Final = TranslationTestCase( "provider_specific_fields": { "reasoning": CLAUDE_SONNET_4_MESSAGE["reasoning"], "reasoning_details": CLAUDE_SONNET_4_MESSAGE["reasoning_details"], - "refusal": CLAUDE_SONNET_4_MESSAGE["refusal"], }, + "refusal": CLAUDE_SONNET_4_MESSAGE["refusal"], }, "finish_reason": "stop", + "logprobs": CLAUDE_SONNET_4_CHOICE["logprobs"], "provider_specific_fields": { "native_finish_reason": CLAUDE_SONNET_4_CHOICE["native_finish_reason"], }, @@ -74,6 +75,7 @@ CLAUDE_SONNET_4_TEST_CASE: Final = TranslationTestCase( }, }, "provider": CLAUDE_SONNET_4_RESPONSE["provider"], + "system_fingerprint": CLAUDE_SONNET_4_RESPONSE["system_fingerprint"], }, ) @@ -132,10 +134,11 @@ GEMINI_3_1_FLASH_IMAGE_TEST_CASE: Final = TranslationTestCase( "provider_specific_fields": { "reasoning": GEMINI_3_1_FLASH_IMAGE_MESSAGE["reasoning"], "reasoning_details": GEMINI_3_1_FLASH_IMAGE_MESSAGE["reasoning_details"], - "refusal": GEMINI_3_1_FLASH_IMAGE_MESSAGE["refusal"], }, + "refusal": GEMINI_3_1_FLASH_IMAGE_MESSAGE["refusal"], }, "finish_reason": "stop", + "logprobs": GEMINI_3_1_FLASH_IMAGE_CHOICE["logprobs"], "provider_specific_fields": { "native_finish_reason": GEMINI_3_1_FLASH_IMAGE_CHOICE["native_finish_reason"], }, @@ -149,6 +152,7 @@ GEMINI_3_1_FLASH_IMAGE_TEST_CASE: Final = TranslationTestCase( }, }, "provider": GEMINI_3_1_FLASH_IMAGE_RESPONSE["provider"], + "system_fingerprint": GEMINI_3_1_FLASH_IMAGE_RESPONSE["system_fingerprint"], "service_tier": GEMINI_3_1_FLASH_IMAGE_RESPONSE["service_tier"], }, ) diff --git a/tests/integration/translation/messages/service_tier/__init__.py b/tests/integration/translation/messages/service_tier/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py new file mode 100644 index 00000000000..de453564e15 --- /dev/null +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py @@ -0,0 +1,29 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.messages.bases.anthropic import CLAUDE_SONNET_4_6_TEST_CASE +from integration.translation.runner import assert_translation + +CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( + CLAUDE_SONNET_4_6_TEST_CASE, + scenario="service_tier", + litellm_request={**CLAUDE_SONNET_4_6_TEST_CASE.litellm_request, "service_tier": "standard_only"}, + expected_provider_request={ + **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request, + "service_tier": "standard_only", + }, + mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1RG36pCcTLE26htFFqC"}, + expected_litellm_response={ + **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response, + "id": "msg_011Cg1RG36pCcTLE26htFFqC", + }, +) + + +@pytest.mark.parametrize("case", [CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id) +def test_messages_service_tier_anthropic(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py new file mode 100644 index 00000000000..13bf9dd888f --- /dev/null +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py @@ -0,0 +1,29 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.messages.bases.azure_ai import CLAUDE_SONNET_4_6_TEST_CASE +from integration.translation.runner import assert_translation + +CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( + CLAUDE_SONNET_4_6_TEST_CASE, + scenario="service_tier", + litellm_request={**CLAUDE_SONNET_4_6_TEST_CASE.litellm_request, "service_tier": "standard_only"}, + expected_provider_request={ + **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request, + "service_tier": "standard_only", + }, + mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1SnrstCdzZPwmTk2qQa"}, + expected_litellm_response={ + **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response, + "id": "msg_011Cg1SnrstCdzZPwmTk2qQa", + }, +) + + +@pytest.mark.parametrize("case", [CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id) +def test_messages_service_tier_azure_ai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py new file mode 100644 index 00000000000..261d34271e9 --- /dev/null +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py @@ -0,0 +1,31 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.messages.bases.bedrock_mantle import CLAUDE_SONNET_5_TEST_CASE +from integration.translation.runner import assert_translation + +CLAUDE_SONNET_5_SERVICE_TIER_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_TEST_CASE, + scenario="service_tier", + litellm_request={**CLAUDE_SONNET_5_TEST_CASE.litellm_request, "service_tier": "standard_only"}, + expected_provider_request={**CLAUDE_SONNET_5_TEST_CASE.expected_provider_request, "service_tier": "standard_only"}, + mock_provider_response={ + **CLAUDE_SONNET_5_TEST_CASE.mock_provider_response, + "id": "msg_bdrk_6oqdfuf56lp47vwldyos5f7l4lyyot337z2cv3h3lbaj3a3ei7vq", + }, + expected_litellm_response={ + **CLAUDE_SONNET_5_TEST_CASE.expected_litellm_response, + "id": "msg_bdrk_6oqdfuf56lp47vwldyos5f7l4lyyot337z2cv3h3lbaj3a3ei7vq", + }, +) + + +@pytest.mark.parametrize("case", [CLAUDE_SONNET_5_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id) +def test_messages_service_tier_bedrock_mantle( + case: TranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py new file mode 100644 index 00000000000..433c5bc4b2f --- /dev/null +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py @@ -0,0 +1,42 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.messages.bases.openai import GPT_5_4_TEST_CASE +from integration.translation.runner import assert_translation + +GPT_5_4_SERVICE_TIER_TEST_CASE: Final = replace( + GPT_5_4_TEST_CASE, + scenario="service_tier", + litellm_request={**GPT_5_4_TEST_CASE.litellm_request, "service_tier": "auto"}, + expected_provider_request={**GPT_5_4_TEST_CASE.expected_provider_request, "service_tier": "auto"}, + mock_provider_response={ + **GPT_5_4_TEST_CASE.mock_provider_response, + "id": "resp_0738fe6d3b086cd1006acef6c23a1487d2860e02c53af945a6", + "created_at": 1791948482, + "access_programs": None, + "completed_at": 1791948483, + "output": [ + { + "id": "msg_0738fe6d3b086cd1006acef6c310cc87d283ecc2cc1d6ff29c", + "type": "message", + "status": "completed", + "content": [{"type": "output_text", "annotations": [], "logprobs": [], "text": "Hello."}], + "phase": "final_answer", + "role": "assistant", + } + ], + }, + expected_litellm_response={ + **GPT_5_4_TEST_CASE.expected_litellm_response, + "id": "resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDo3NGI3ZDQzZjBmNjQ1MjlkYWY4NWQ4MjRmNjE3ZGYzNDRjZTdhNGI3NjI3YTYzMDMzZGUzMDVmNTE2NTRlZjc2O3Jlc3BvbnNlX2lkOnJlc3BfMDczOGZlNmQzYjA4NmNkMTAwNmFjZWY2YzIzYTE0ODdkMjg2MGUwMmM1M2FmOTQ1YTY=", + }, +) + + +@pytest.mark.parametrize("case", [GPT_5_4_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id) +def test_messages_service_tier_openai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index b3cd7bd01f9..6c317c3b9d9 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -5,7 +5,6 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import catalog from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -19,6 +18,7 @@ from tests.test_litellm_rust.support.requests import ( MESSAGES_RESPONSE, request_body, ) +from tests.test_litellm_rust.support.served import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -48,11 +48,6 @@ def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: } -def assert_served_natively(server: RecordingServer) -> None: - assert len(server.requests) == 1 - assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") - - @pytest.mark.asyncio async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response( messages_server: RecordingServer, @@ -63,7 +58,7 @@ async def test_native_messages_callbacks_see_the_provider_request_and_the_public **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success") ) - assert_served_natively(messages_server) + assert_served_natively(messages_server, response) assert response["content"] == MESSAGES_RESPONSE["content"] sent: Final = messages_server.requests[0] assert sent.path == "/v1/messages" @@ -107,7 +102,7 @@ async def test_native_messages_provider_error_reaches_caller_and_failure_callbac with pytest.raises(litellm.BadRequestError) as raised: await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()])) - assert_served_natively(messages_server) + assert len(messages_server.requests) == 1 assert [phase for phase, _ in observed] == ["sync", "async"] assert all(error is raised.value for _, error in observed) @@ -127,14 +122,13 @@ async def test_native_messages_stream_relays_provider_events_and_logs_success_on **arguments(messages_server, stream=True, callbacks=[recorder]) ) assert isinstance(stream, AsyncIterator) - assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true" first: Final = await anext(stream) await drain_logging() assert "async_log_success_event" not in recorder.names rest: Final = [chunk async for chunk in stream] assert first + b"".join(rest) == sse_payload() - assert_served_natively(messages_server) + assert_served_natively(messages_server, stream) assert messages_server.requests[0].body["stream"] is True success: Final = await recorder.wait_for_async("async_log_success_event") assert len(success) == 1 @@ -171,10 +165,9 @@ def test_native_sync_messages_stream_relays_provider_events_and_logs_success_onc stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) assert isinstance(stream, Iterator) - assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true" assert b"".join(stream) == sse_payload() - assert_served_natively(messages_server) + assert_served_natively(messages_server, stream) assert len(recorder.wait_for("async_log_success_event")) == 1 @@ -183,7 +176,7 @@ def test_native_sync_messages_returns_the_provider_message(messages_server: Reco response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder])) - assert_served_natively(messages_server) + assert_served_natively(messages_server, response) assert response["content"] == MESSAGES_RESPONSE["content"] assert len(recorder.wait_for("log_success_event")) == 1 @@ -216,7 +209,7 @@ async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(me response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()])) - assert_served_natively(messages_server) + assert_served_natively(messages_server, response) assert response["content"] == MESSAGES_RESPONSE["content"] @@ -259,9 +252,9 @@ async def test_native_messages_does_not_restore_keywords_deleted_by_a_pre_call_h stream=False, ) with rebound(litellm, "callbacks", [DropKeywords()]): - await litellm.anthropic.messages.acreate(**call_arguments) + response: Final = await litellm.anthropic.messages.acreate(**call_arguments) - assert_served_natively(messages_server) + assert_served_natively(messages_server, response) sent: Final = messages_server.requests[0] assert sent.body == { "model": MESSAGES_MODEL.removeprefix("anthropic/"), diff --git a/tests/test_litellm_rust/messages/test_request_shaping.py b/tests/test_litellm_rust/messages/test_request_shaping.py index 3584a2a8593..3ae17e0f614 100644 --- a/tests/test_litellm_rust/messages/test_request_shaping.py +++ b/tests/test_litellm_rust/messages/test_request_shaping.py @@ -17,6 +17,7 @@ from litellm.rust_bridge.configuration import Rollout from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_RESPONSE +from tests.test_litellm_rust.support.served import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -47,10 +48,9 @@ def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: } -def sent(server: RecordingServer) -> tuple[dict[str, object], dict[str, str]]: - assert len(server.requests) == 1 +def sent(server: RecordingServer, response: object) -> tuple[dict[str, object], dict[str, str]]: + assert_served_natively(server, response) request: Final = server.requests[0] - assert not request.headers.get("user-agent", "").startswith("python-httpx") assert isinstance(request.body, dict) return request.body, request.headers @@ -59,9 +59,9 @@ def sent(server: RecordingServer) -> tuple[dict[str, object], dict[str, str]]: async def test_reasoning_effort_becomes_adaptive_thinking_and_effort_on_the_wire( messages_server: RecordingServer, ) -> None: - await litellm.anthropic.messages.acreate(**arguments(messages_server, reasoning_effort="high")) + response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, reasoning_effort="high")) - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert "reasoning_effort" not in body assert body["thinking"] == {"type": "adaptive", "display": "summarized"} assert body["output_config"] == {"effort": "high"} @@ -71,7 +71,7 @@ async def test_reasoning_effort_becomes_adaptive_thinking_and_effort_on_the_wire async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for_a_legacy_model( messages_server: RecordingServer, ) -> None: - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments( messages_server, model=LEGACY_THINKING_MODEL, @@ -82,7 +82,7 @@ async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for ) ) - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert body["thinking"] == {"type": "enabled", "budget_tokens": 2999} assert "output_config" not in body assert "temperature" not in body @@ -90,11 +90,11 @@ async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for @pytest.mark.asyncio async def test_removed_sampling_params_are_dropped_under_drop_params(messages_server: RecordingServer) -> None: - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments(messages_server, temperature=0.2, top_p=0.9, top_k=5, drop_params=True) ) - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert not {"temperature", "top_p", "top_k"} & body.keys() @@ -130,9 +130,9 @@ async def test_replayed_history_is_sanitized_before_it_reaches_the_provider( {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]}, ] - await litellm.anthropic.messages.acreate(**arguments(messages_server, messages=history)) + response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, messages=history)) - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert body["messages"] == [ {"role": "user", "content": "run it"}, {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]}, @@ -142,7 +142,7 @@ async def test_replayed_history_is_sanitized_before_it_reaches_the_provider( @pytest.mark.asyncio async def test_feature_betas_merge_into_the_forwarded_beta_header(messages_server: RecordingServer) -> None: - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments( messages_server, output_format={"type": "json_schema", "schema": {"type": "object"}}, @@ -150,15 +150,17 @@ async def test_feature_betas_merge_into_the_forwarded_beta_header(messages_serve ) ) - _, headers = sent(messages_server) + _, headers = sent(messages_server, response) assert headers["anthropic-beta"] == "structured-outputs-2025-11-13,web-search-2025-03-05" @pytest.mark.asyncio async def test_oauth_token_authenticates_as_a_bearer_with_the_oauth_beta(messages_server: RecordingServer) -> None: - await litellm.anthropic.messages.acreate(**arguments(messages_server, api_key="sk-ant-oat01-token")) + response: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, api_key="sk-ant-oat01-token") + ) - _, headers = sent(messages_server) + _, headers = sent(messages_server, response) assert "x-api-key" not in headers assert headers["authorization"] == "Bearer sk-ant-oat01-token" assert headers["anthropic-beta"] == "oauth-2025-04-20" @@ -167,11 +169,11 @@ async def test_oauth_token_authenticates_as_a_bearer_with_the_oauth_beta(message @pytest.mark.asyncio async def test_metadata_is_reduced_to_the_fields_anthropic_accepts(messages_server: RecordingServer) -> None: - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments(messages_server, metadata={"user_id": "u-1", "trace_id": "internal"}) ) - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert body["metadata"] == {"user_id": "u-1"} @@ -179,17 +181,17 @@ async def test_metadata_is_reduced_to_the_fields_anthropic_accepts(messages_serv async def test_additional_drop_params_remove_nested_fields_from_the_wire(messages_server: RecordingServer) -> None: tools: Final = [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}] - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments(messages_server, tools=tools, additional_drop_params=["tools[*].input_examples"]) ) - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert body["tools"] == [{"name": "lookup", "input_schema": {"type": "object"}}] @pytest.mark.asyncio async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(messages_server: RecordingServer) -> None: - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments( messages_server, provider_specific_header=[ @@ -199,7 +201,7 @@ async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(mess ) ) - _, headers = sent(messages_server) + _, headers = sent(messages_server, response) assert headers["x-scoped"] == "yes" assert "x-other" not in headers @@ -208,7 +210,7 @@ async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(mess async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers( messages_server: RecordingServer, ) -> None: - await litellm.anthropic.messages.acreate( + response: Final = await litellm.anthropic.messages.acreate( **arguments( messages_server, headers={"x-priority": "forwarded", "x-forwarded-only": "kept"}, @@ -217,7 +219,7 @@ async def test_scoped_headers_override_extra_headers_which_override_forwarded_he ) ) - _, headers = sent(messages_server) + _, headers = sent(messages_server, response) assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == { "x-priority": "scoped", "x-forwarded-only": "kept", @@ -298,20 +300,12 @@ async def test_native_messages_observes_runtime_capabilities_and_separate_caller @pytest.mark.asyncio async def test_native_messages_reads_optional_positional_body_parameters(messages_server: RecordingServer) -> None: - from litellm.messages.dispatch import _MESSAGES, _public_request - from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES - - native: Final = NATIVE_AMESSAGES.load() - assert native is not None metadata: Final = {"user_id": "caller"} args: Final = (16, MESSAGES, "anthropic/claude-test", metadata, None, False, "Be brief", 0.25) kwargs: Final = {"api_key": "test-key", "api_base": messages_server.base_url} - call: Final = _public_request(_MESSAGES, args, kwargs) - assert call is not None + response: Final = await litellm.anthropic.messages.acreate(*args, **kwargs) - await native(call) - - body, _ = sent(messages_server) + body, _ = sent(messages_server, response) assert body["temperature"] == args[7] assert body["system"] == args[6] assert body["metadata"] == metadata diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py index 7b4b9fab579..a29a45376d9 100644 --- a/tests/test_litellm_rust/ocr/test_dispatch.py +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -6,6 +6,7 @@ import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE +from tests.test_litellm_rust.support.served import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -28,8 +29,7 @@ def test_public_ocr_uses_native_route_independently_of_flag(ocr_server: Recordin assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "native OCR response" - assert len(ocr_server.requests) == 1 - assert not ocr_server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + assert_served_natively(ocr_server, response) @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 7c62133202d..8c4f3a403d6 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -19,6 +19,7 @@ from tests.test_litellm_rust.support.requests import ( call_native_aocr, call_native_ocr, ) +from tests.test_litellm_rust.support.served import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -169,7 +170,7 @@ async def test_ocr_contract_native_format_supported( assert response.get_provider_native_response() == payload assert len(ocr_server.requests) == 1 if ocr_backend: - assert_native_request(ocr_server) + assert_served_natively(ocr_server, response) @pytest.mark.asyncio @@ -217,7 +218,7 @@ async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token( assert calls == ["token"] assert response.pages[0].markdown == "native OCR response" - assert_native_request(ocr_server) + assert_served_natively(ocr_server, response) assert ocr_server.requests[0].headers["authorization"] == "Bearer callback-token" @@ -227,11 +228,6 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer: return recording_server -def assert_native_request(server: RecordingServer) -> None: - assert len(server.requests) == 1 - assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") - - def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None: ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) @@ -252,7 +248,7 @@ def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(ocr_se ) assert response.pages[0].markdown == "native OCR response" - assert_native_request(ocr_server) + assert_served_natively(ocr_server, response) assert ocr_server.requests[0].body == { "model": "mistral-ocr-latest", "document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}, @@ -614,7 +610,7 @@ async def test_native_settings_observe_mutation_between_calls( monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.") response: Final = await call_native(ocr_server, asynchronous, num_retries=0) assert response.pages[0].markdown == "native OCR response" - assert_native_request(ocr_server) + assert_served_natively(ocr_server, response) monkeypatch.setattr(litellm, "ssl_certificate", 1) with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"): await call_native(ocr_server, asynchronous, num_retries=0) diff --git a/tests/test_litellm_rust/support/served.py b/tests/test_litellm_rust/support/served.py new file mode 100644 index 00000000000..d3c568937f3 --- /dev/null +++ b/tests/test_litellm_rust/support/served.py @@ -0,0 +1,11 @@ +from typing import Final + +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from tests.test_litellm_rust.support.recording_server import RecordingServer + + +def assert_served_natively(server: RecordingServer, response: object) -> None: + assert len(server.requests) == 1 + additional_headers: Final = get_hidden_params_dict(response).get("additional_headers") + assert isinstance(additional_headers, dict) + assert additional_headers.get("x-litellm-rust") == "true" diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index 02c454e0460..8e1251de8c2 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -66,14 +66,11 @@ _SDK_CONTRACT = textwrap.dedent( import litellm from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict class Handler(BaseHTTPRequestHandler): def do_POST(self): self.rfile.read(int(self.headers["Content-Length"])) - if self.headers.get("User-Agent", "").startswith("python-httpx"): - self.send_response(418) - self.end_headers() - return body = json.dumps({ "pages": [{"index": 0, "markdown": "native", "images": [], "dimensions": None}], "model": "mistral-ocr-latest", @@ -108,7 +105,10 @@ _SDK_CONTRACT = textwrap.dedent( return REFUSED except Exception: return OTHER - return SERVED if response.pages[0].markdown == "native" else OTHER + return SERVED if ( + response.pages[0].markdown == "native" + and get_hidden_params_dict(response).get("additional_headers", {}).get("x-litellm-rust") == "true" + ) else OTHER def forked(asynchronous): pid = os.fork() diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index 46600e0bf60..5ebc3bf7336 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -489,6 +489,9 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) self.call_count = 0 + def _admit_async_call(self, call): + return call() + @_redis_circuit_breaker_guard async def boom(self): self.call_count += 1 diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 0a041694d52..d27645d19fa 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -3,7 +3,7 @@ import time import types from collections.abc import Iterator from datetime import timedelta -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -971,18 +971,17 @@ async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its @pytest.mark.asyncio -async def test_breaker_metrics_track_state_and_failure_class(): +async def test_breaker_metrics_track_state_and_failure_class(monkeypatch): """Breaker accounting must be observable: failure class, transitions, and state.""" from prometheus_client import REGISTRY from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _breaker_metrics, is_redis_timeout_failure + from litellm._service_logger import ServiceLogging + from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure - metrics: Final = _breaker_metrics() - for collector in (metrics._state_gauge, metrics._transitions, metrics._failures): - if collector is not None and collector not in REGISTRY._collector_to_names: - REGISTRY.register(collector) + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + service_logger: Final = ServiceLogging() def sample(name, labels=None): return REGISTRY.get_sample_value(name, labels) or 0.0 @@ -993,7 +992,9 @@ async def test_breaker_metrics_track_state_and_failure_class(): open_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "open"}) closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) - breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0) + breaker = RedisCircuitBreaker( + failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0, service_logger=service_logger + ) breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("t"))) breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) @@ -1011,6 +1012,24 @@ async def test_breaker_metrics_track_state_and_failure_class(): assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1 +def test_prometheus_services_logger_starts_from_the_breakers_already_built(monkeypatch): + """A breaker built before prometheus_system is on shows up in the state gauge once the services logger starts.""" + from prometheus_client import REGISTRY + + from litellm._service_logger import ServiceLogging + from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts + + monkeypatch.setattr(litellm, "service_callback", []) + early: Final = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, service_logger=ServiceLogging()) + assert early.is_open() is False + + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + ServiceLogging() + assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( + redis_breaker_state_counts()["closed"] + ) + + def test_sync_guard_counts_a_timeout_as_a_timeout(): from redis.exceptions import TimeoutError as RedisTimeoutError @@ -1308,6 +1327,32 @@ async def test_pool_wait_timeout_is_a_timeout_failure_not_hard_connectivity(): await client.aclose() +@pytest.mark.asyncio +async def test_cluster_pool_exhaustion_is_a_timeout_failure_not_hard_connectivity(): + """The async cluster client has no blocking pool: a node whose max_connections are all busy raises + MaxConnectionsError at once. That is the cluster shape of a busy pool, so it must sit behind the + timeout duration gate like the standalone pool's wait timeout instead of opening the breaker.""" + from redis import exceptions as redis_exceptions + + max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError + "type[Exception]", + redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs + ) + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + async def pool_exhausted() -> None: + raise max_connections_error() + + for _ in range(breaker.failure_threshold * 2): + with pytest.raises(max_connections_error): + await _run_under_circuit_breaker(breaker, "op", pool_exhausted) + + assert breaker.is_open() is False, "an exhausted client-side pool is a timeout, not a dead Redis" + + def test_timeout_classification_follows_the_explicit_cause_chain_only(): from redis.exceptions import ConnectionError as RedisConnectionError diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 68260bc6870..85422819926 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -1,20 +1,22 @@ import asyncio -from importlib import import_module +import gc import json import ssl +import weakref +from collections.abc import Awaitable, Callable +from importlib import import_module +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi.testclient import TestClient from redis.asyncio import Redis, RedisCluster from redis.asyncio.cluster import ClusterNode from redis.asyncio.connection import SSLConnection - -from litellm.caching.redis_cache import RedisCache -from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cluster_cache import ClusterConnectionGate, RedisClusterCache @patch("litellm._redis.init_redis_cluster") @@ -73,9 +75,7 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): @patch("litellm._redis.get_redis_connection_pool") @patch("litellm._redis.get_redis_client") @patch.object(import_module("litellm.caching.redis_cache").RedisCache, "_setup_health_pings") -def test_cache_init_creates_cluster_cache_from_env_var( - mock_health, mock_get_client, mock_get_pool, monkeypatch -): +def test_cache_init_creates_cluster_cache_from_env_var(mock_health, mock_get_client, mock_get_pool, monkeypatch): """ Test that Cache() creates RedisClusterCache when REDIS_CLUSTER_NODES env var is set. @@ -153,9 +153,7 @@ def test_cache_init_creates_redis_cache_without_cluster_config( ), ], ) -def test_router_create_redis_cache_cluster_detection( - startup_nodes, env_var, expected_cache_type, monkeypatch -): +def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, expected_cache_type, monkeypatch): """ Test that Router._create_redis_cache() creates RedisClusterCache when either startup_nodes is in config or REDIS_CLUSTER_NODES env var is set. @@ -416,3 +414,186 @@ def test_evicted_async_client_is_closed_only_for_cluster(cluster: bool) -> None: client.aclose.assert_not_awaited() asyncio.run(exercise()) + + +class _InFlightCountingClient: + """Async cluster client double that records how many commands were in flight at once.""" + + def __init__(self, command_seconds: float) -> None: + self.command_seconds = command_seconds + self.in_flight = 0 + self.peak_in_flight = 0 + self.completed = 0 + + async def set(self, *args: object, **kwargs: object) -> bool: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + await asyncio.sleep(self.command_seconds) + self.in_flight -= 1 + self.completed += 1 + return True + + +def _cluster_set_cache_with_client( + client: _InFlightCountingClient, + max_connections: int | None = None, + connection_gate: ClusterConnectionGate | None = None, +) -> Callable[[str, str], Awaitable[object]]: + with patch("litellm._redis.init_redis_cluster"): + cache = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 7000}], + max_connections=max_connections, + connection_gate=connection_gate, + ) + setattr(cache, "init_async_client", MagicMock(return_value=client)) + return cast( + "Callable[[str, str], Awaitable[object]]", cache.async_set_cache + ) # cast-ok: wrapped **kwargs are untyped + + +@pytest.mark.asyncio +async def test_cluster_cache_queues_commands_beyond_max_connections_instead_of_failing() -> None: + """A burst larger than max_connections used to hand every command to the node at once, so redis-py raised + MaxConnectionsError on the overflow and the circuit breaker opened. The overflow now waits for a slot.""" + client = _InFlightCountingClient(command_seconds=0.01) + set_cache = _cluster_set_cache_with_client(client, max_connections=3) + + await asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(20))) + + assert client.completed == 20 + assert client.peak_in_flight <= 3 + + +@pytest.mark.asyncio +async def test_cluster_cache_without_max_connections_does_not_throttle() -> None: + client = _InFlightCountingClient(command_seconds=0.01) + set_cache = _cluster_set_cache_with_client(client) + + await asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(20))) + + assert client.completed == 20 + assert client.peak_in_flight == 20 + + +@pytest.mark.asyncio +async def test_cluster_cache_sends_after_the_connection_wait_timeout() -> None: + """A command that cannot get a slot within the pool timeout is still sent, so a saturated pool keeps + surfacing through the client's own errors and the breaker rather than stalling forever.""" + client = _InFlightCountingClient(command_seconds=0.2) + set_cache = _cluster_set_cache_with_client( + client, connection_gate=ClusterConnectionGate(limit=1, wait_timeout=0.02) + ) + + await asyncio.wait_for(asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(3))), timeout=1.0) + + assert client.completed == 3 + assert client.peak_in_flight == 3 + + +@pytest.mark.asyncio +async def test_cluster_caches_sharing_a_client_share_the_connection_cap() -> None: + """init_async_client hands caches with identical settings the same client, so two caches must not each admit + a full max_connections worth of commands against that one pool.""" + client = _InFlightCountingClient(command_seconds=0.01) + first = _cluster_set_cache_with_client(client, max_connections=3) + second = _cluster_set_cache_with_client(client, max_connections=3) + + await asyncio.gather( + *(first(f"first-{i}", "value") for i in range(10)), *(second(f"second-{i}", "value") for i in range(10)) + ) + + assert client.completed == 20 + assert client.peak_in_flight <= 3 + + +@pytest.mark.asyncio +async def test_guarded_command_nested_in_a_guarded_command_reuses_the_slot() -> None: + """batch_cache_write flushes through async_set_cache_pipeline while it is itself admitted. With one slot the + inner command must run on the outer slot instead of waiting for the pool timeout.""" + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + + async def inner() -> str: + return "flushed" + + async def outer() -> str: + return await gate.run(inner, "client-a") + + assert await asyncio.wait_for(gate.run(outer, "client-a"), timeout=1.0) == "flushed" + + +@pytest.mark.asyncio +async def test_nested_guarded_command_on_another_client_waits_for_that_clients_slot() -> None: + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + client_b_busy = asyncio.Event() + release_b = asyncio.Event() + peak_b = 0 + in_flight_b = 0 + + async def hold_b() -> None: + nonlocal peak_b, in_flight_b + in_flight_b += 1 + peak_b = max(peak_b, in_flight_b) + client_b_busy.set() + await release_b.wait() + in_flight_b -= 1 + + async def touch_b() -> None: + nonlocal peak_b, in_flight_b + in_flight_b += 1 + peak_b = max(peak_b, in_flight_b) + await asyncio.sleep(0) + in_flight_b -= 1 + + async def outer_on_a() -> None: + await client_b_busy.wait() + await gate.run(touch_b, "client-b") + + holder = asyncio.create_task(gate.run(hold_b, "client-b")) + nested = asyncio.create_task(gate.run(outer_on_a, "client-a")) + await client_b_busy.wait() + await asyncio.sleep(0.05) + release_b.set() + await asyncio.gather(holder, nested) + + assert peak_b == 1 + + +@pytest.mark.asyncio +async def test_child_task_spawned_inside_a_guarded_command_does_not_inherit_the_slot() -> None: + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + peak = 0 + in_flight = 0 + + async def command() -> None: + nonlocal peak, in_flight + in_flight += 1 + peak = max(peak, in_flight) + await asyncio.sleep(0.02) + in_flight -= 1 + + spawned: list[asyncio.Task[None]] = [] + + async def parent() -> None: + spawned.append(asyncio.create_task(gate.run(command, "client-a"))) + await command() + + await gate.run(parent, "client-a") + await asyncio.gather(*spawned) + + assert peak == 1 + + +def test_closed_event_loop_is_collected_after_contended_gate_use() -> None: + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + + async def command() -> None: + await asyncio.sleep(0.01) + + async def contend() -> "weakref.ref[asyncio.AbstractEventLoop]": + await asyncio.gather(gate.run(command, "client-a"), gate.run(command, "client-a")) + return weakref.ref(asyncio.get_running_loop()) + + loop_ref = asyncio.run(contend()) + gc.collect() + + assert loop_ref() is None diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 5c2b9412d16..17d8e8ad9d1 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -13,6 +13,8 @@ from pydantic import BaseModel, ConfigDict import litellm from litellm.integrations.anthropic_cache_control_hook import ( + CARRY_UNMATCHED_MESSAGE_POINTS, + EXTERNAL_BREAKPOINTS_STAMP, AnthropicCacheControlHook, configured_injection_points, supports_openai_prompt_cache_breakpoint, @@ -1180,7 +1182,9 @@ def test_tool_call_census_matches_the_breakpoints_providers_send(mark: object, t message: Final = cast(AllMessageValues, {"role": "assistant", "content": None, "tool_calls": [tool_call]}) bedrock_cache_points: Final = sum( 1 - for block in _convert_to_bedrock_tool_call_invoke([tool_call], model="anthropic.claude-sonnet-4-5-20250929-v1:0") + for block in _convert_to_bedrock_tool_call_invoke( + [tool_call], model="anthropic.claude-sonnet-4-5-20250929-v1:0" + ) if "cachePoint" in block ) anthropic_marks: Final = sum( @@ -1895,7 +1899,9 @@ class TestEnableAnthropicPromptCaching: ) @pytest.mark.parametrize("infer_provider", [False, True]) @pytest.mark.parametrize("supported", [False, True]) - def test_claude_transport_defaults(self, monkeypatch, local_model_cost_map, family, provider, template, infer_provider, supported): + def test_claude_transport_defaults( + self, monkeypatch, local_model_cost_map, family, provider, template, infer_provider, supported + ): from litellm.utils import supports_prompt_caching model = template.format(f"claude-{family}") @@ -1908,14 +1914,21 @@ class TestEnableAnthropicPromptCaching: resolved_provider = None if infer_provider else provider assert supports_prompt_caching(model=target, custom_llm_provider=resolved_provider) is supported points = AnthropicCacheControlHook.get_default_injection_points( - messages=copy.deepcopy(self.MESSAGES), system=None, model=target, - custom_llm_provider=resolved_provider, enable_prompt_caching=True, + messages=copy.deepcopy(self.MESSAGES), + system=None, + model=target, + custom_llm_provider=resolved_provider, + enable_prompt_caching=True, ) assert [point["index"] for point in points] == ([None, -1] if supported else []) affinity_messages = AnthropicCacheControlHook.messages_with_default_injections( - copy.deepcopy(self.MESSAGES), models=[qualified], enable_prompt_caching=True, + copy.deepcopy(self.MESSAGES), + models=[qualified], + enable_prompt_caching=True, + ) + assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in affinity_messages) == ( + 2 if supported else 0 ) - assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in affinity_messages) == (2 if supported else 0) @pytest.mark.parametrize( "provider, model", @@ -1929,7 +1942,9 @@ class TestEnableAnthropicPromptCaching: ("anthropic", "unknown-model"), ], ) - def test_non_claude_caching_capability_does_not_enable_defaults(self, monkeypatch, local_model_cost_map, provider, model): + def test_non_claude_caching_capability_does_not_enable_defaults( + self, monkeypatch, local_model_cost_map, provider, model + ): from litellm.utils import supports_prompt_caching qualified = f"{provider}/{model}" @@ -1946,21 +1961,56 @@ class TestEnableAnthropicPromptCaching: @pytest.mark.parametrize("client_control", ["none", "message", "system", "tool", "function", "top_level"]) @pytest.mark.parametrize("envelope", ["request", "extra_body"]) @pytest.mark.parametrize("configured", [False, True]) - def test_new_transports_preserve_client_controls(self, monkeypatch, local_model_cost_map, provider, client_control, envelope, configured): + def test_new_transports_preserve_client_controls( + self, monkeypatch, local_model_cost_map, provider, client_control, envelope, configured + ): from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig model = "claude-sonnet-5" monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) - monkeypatch.setitem(litellm.model_cost, f"{provider}/{model}", { - **litellm.model_cost[f"{provider}/{model}"], "supports_prompt_caching": True, - }) + monkeypatch.setitem( + litellm.model_cost, + f"{provider}/{model}", + { + **litellm.model_cost[f"{provider}/{model}"], + "supports_prompt_caching": True, + }, + ) control = {"type": "ephemeral"} - messages = [{"role": "user", "content": [{"type": "text", "text": "question", **({"cache_control": control} if client_control == "message" else {})}]}] - system = [{"type": "text", "text": "stable context", **({"cache_control": control} if client_control == "system" else {})}] - tools = [{"name": "lookup", "description": "Lookup", "input_schema": {"type": "object", "properties": {}}, **({"cache_control": control} if client_control == "tool" else {})}] + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "question", + **({"cache_control": control} if client_control == "message" else {}), + } + ], + } + ] + system = [ + { + "type": "text", + "text": "stable context", + **({"cache_control": control} if client_control == "system" else {}), + } + ] + tools = [ + { + "name": "lookup", + "description": "Lookup", + "input_schema": {"type": "object", "properties": {}}, + **({"cache_control": control} if client_control == "tool" else {}), + } + ] if client_control == "function": tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}, "cache_control": control}}] - kwargs = {"metadata": {}, "model_info": {"id": "selected-deployment"}, **({"cache_control": control} if client_control == "top_level" else {})} + kwargs = { + "metadata": {}, + "model_info": {"id": "selected-deployment"}, + **({"cache_control": control} if client_control == "top_level" else {}), + } if envelope == "extra_body": kwargs["extra_body"] = {"messages": messages, "system": system, "tools": tools} if "cache_control" in kwargs: @@ -1974,7 +2024,12 @@ class TestEnableAnthropicPromptCaching: seeded = copy.deepcopy(kwargs) original = copy.deepcopy((messages, system, tools)) result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( - messages, system, kwargs, model, provider, tools=tools, + messages, + system, + kwargs, + model, + provider, + tools=tools, ) if client_control != "none" and not configured: assert (result_messages, result_system, tools) == original @@ -1989,32 +2044,46 @@ class TestEnableAnthropicPromptCaching: assert (result_system == original[1]) == (envelope == "request" and client_control == "system") if provider == "vertex_ai": wire = VertexAIAnthropicConfig().transform_request( - model=model, messages=[{"role": "system", "content": result_system}, *result_messages], - optional_params={"max_tokens": 8}, litellm_params={}, headers={}, + model=model, + messages=[{"role": "system", "content": result_system}, *result_messages], + optional_params={"max_tokens": 8}, + litellm_params={}, + headers={}, ) assert wire["system"][0]["cache_control"] == control assert wire["messages"][-1]["content"][-1]["cache_control"] == control affinity = AnthropicCacheControlHook.messages_with_default_injections( - [{"role": "system", "content": original[1]}, *original[0]], [f"{provider}/{model}"], - tools=tools, request_kwargs=seeded, + [{"role": "system", "content": original[1]}, *original[0]], + [f"{provider}/{model}"], + tools=tools, + request_kwargs=seeded, ) if client_control != "none": assert affinity == [{"role": "system", "content": original[1]}, *original[0]] AnthropicCacheControlHook.maybe_seed_default_injection_points( - seeded, [{"role": "system", "content": original[1]}, *original[0]], model, provider, tools=tools, + seeded, + [{"role": "system", "content": original[1]}, *original[0]], + model, + provider, + tools=tools, ) assert bool(seeded.get("cache_control_injection_points")) == (client_control == "none" or configured) @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) - @pytest.mark.parametrize("model, target, client_control, expected", [ - ("vertex_ai/claude-sonnet-5", "bedrock/amazon.nova-pro-v1:0", False, 0), - ("azure_ai/gpt-6-astra", "azure_ai/claude-sonnet-5", False, 2), - ("azure_ai/claude-sonnet-5", None, False, 2), - ("azure_ai/claude-sonnet-5", None, True, 1), - ("azure_ai/model_router/claude-replacement", None, False, 2), - ]) - async def test_public_completion_cache_ownership(self, monkeypatch, local_model_cost_map, asynchronous, model, target, client_control, expected): + @pytest.mark.parametrize( + "model, target, client_control, expected", + [ + ("vertex_ai/claude-sonnet-5", "bedrock/amazon.nova-pro-v1:0", False, 0), + ("azure_ai/gpt-6-astra", "azure_ai/claude-sonnet-5", False, 2), + ("azure_ai/claude-sonnet-5", None, False, 2), + ("azure_ai/claude-sonnet-5", None, True, 1), + ("azure_ai/model_router/claude-replacement", None, False, 2), + ], + ) + async def test_public_completion_cache_ownership( + self, monkeypatch, local_model_cost_map, asynchronous, model, target, client_control, expected + ): import httpx from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -2027,22 +2096,46 @@ class TestEnableAnthropicPromptCaching: monkeypatch.setitem(litellm.model_cost, qualified, entry) monkeypatch.setitem(litellm.model_cost, qualified.split("/", 1)[-1], entry) sent = [] + def respond(request): sent.append(json.loads(request.content)) - return httpx.Response(200, request=request, json={ - "id": "msg-test", "type": "message", "role": "assistant", "model": "claude-sonnet-5", - "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn", "stop_sequence": None, - "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, "stopReason": "end_turn", - "usage": {"input_tokens": 10, "output_tokens": 1, "inputTokens": 10, "outputTokens": 1, "totalTokens": 11}, - }) + return httpx.Response( + 200, + request=request, + json={ + "id": "msg-test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": { + "input_tokens": 10, + "output_tokens": 1, + "inputTokens": 10, + "outputTokens": 1, + "totalTokens": 11, + }, + }, + ) + control = {"type": "ephemeral", "ttl": "1h"} messages = [{"role": "system", "content": "stable context"}, {"role": "user", "content": "question"}] metadata = {} kwargs = { - "model": model, "messages": copy.deepcopy(messages), "max_tokens": 32, "num_retries": 0, + "model": model, + "messages": copy.deepcopy(messages), + "max_tokens": 32, + "num_retries": 0, "litellm_metadata": metadata, - "api_base": "https://rig.services.ai.azure.com/anthropic", "api_key": "synthetic-test-key", - "aws_access_key_id": "synthetic", "aws_secret_access_key": "synthetic", "aws_region_name": "us-east-1", + "api_base": "https://rig.services.ai.azure.com/anthropic", + "api_key": "synthetic-test-key", + "aws_access_key_id": "synthetic", + "aws_secret_access_key": "synthetic", + "aws_region_name": "us-east-1", **({"extra_body": {"cache_control": control}} if client_control else {}), } if asynchronous: @@ -2131,7 +2224,9 @@ class TestEnableAnthropicPromptCaching: """OpenAI-shaped tools nest cache_control under ``function``; the Anthropic chat transform honors that location, so the stand-down must see it too.""" monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) - tools = [{"type": "function", "function": {"name": "t", "parameters": {}, "cache_control": {"type": "ephemeral"}}}] + tools = [ + {"type": "function", "function": {"name": "t", "parameters": {}, "cache_control": {"type": "ephemeral"}}} + ] assert self._points(tools=tools) == [] def test_seed_stands_down_when_only_tools_carry_cache_control(self, monkeypatch): @@ -2912,7 +3007,10 @@ class TestConfiguredInjectionPointsSurviveClientMarks: self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES), tools=[self.MARKED_TOOL_TOP_LEVEL]) affinity = AnthropicCacheControlHook.messages_with_default_injections( - copy.deepcopy(self.CLEAN_MESSAGES), ["claude-sonnet-4-5"], tools=[self.MARKED_TOOL_TOP_LEVEL], request_kwargs=params + copy.deepcopy(self.CLEAN_MESSAGES), + ["claude-sonnet-4-5"], + tools=[self.MARKED_TOOL_TOP_LEVEL], + request_kwargs=params, ) assert [p["index"] for p in params["cache_control_injection_points"]] == [None, -1] @@ -2966,9 +3064,7 @@ class TestAnthropicPromptCachingEnvVars: print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl])) """ ) - result = subprocess.run( - [sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300 - ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300) assert result.returncode == 0, result.stderr enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1]) return enabled, ttl @@ -3337,7 +3433,11 @@ class TestOpenAIPromptCacheBreakpointPlacementRules: def test_tool_message_text_is_marked_on_chat_path(self): messages = [ {"role": "user", "content": "weather?"}, - {"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "w", "arguments": "{}"}}]}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "w", "arguments": "{}"}}], + }, {"role": "tool", "tool_call_id": "c1", "content": "sunny"}, ] out, params = self._chat(messages, [{"location": "message", "index": -1}]) @@ -3561,9 +3661,9 @@ class TestChatPathProviderStamp: class TestClientBreakpointsCountedOnce: def test_client_message_breakpoints_are_not_double_counted(self): - messages = [{"role": "user", "content": [{"type": "text", "text": "m0", "cache_control": {"type": "ephemeral"}}]}] + [ - {"role": "user", "content": [{"type": "text", "text": f"m{i}"}]} for i in range(1, 4) - ] + messages = [ + {"role": "user", "content": [{"type": "text", "text": "m0", "cache_control": {"type": "ephemeral"}}]} + ] + [{"role": "user", "content": [{"type": "text", "text": f"m{i}"}]} for i in range(1, 4)] out, system, _ = AnthropicCacheControlHook.apply_to_anthropic_messages_request( messages=messages, system="sys", @@ -3721,7 +3821,6 @@ class TestPromptCacheBreakpointCapability: yield litellm.utils.cached_get_model_info_helper.cache_clear() - def test_listed_model_uses_the_model_map_flag(self, monkeypatch): flagged = {**litellm.model_cost["gpt-4.1"], "supports_prompt_cache_breakpoint": True} monkeypatch.setitem(litellm.model_cost, "gpt-4.1", flagged) @@ -3747,7 +3846,6 @@ class TestPromptCacheBreakpointCapability: ) assert supports_openai_prompt_cache_breakpoint(model) is False - def test_published_map_without_the_flag_still_injects_on_gpt_5_6(self, monkeypatch): unflagged = {k: v for k, v in litellm.model_cost["gpt-5.6"].items() if k != "supports_prompt_cache_breakpoint"} monkeypatch.setitem(litellm.model_cost, "gpt-5.6", unflagged) @@ -3854,9 +3952,7 @@ class TestHostedOpenAIDialectFlag: def test_entry_provider_must_match_the_request_provider(self, monkeypatch): """A flagged entry does not license a different provider serving the same model string.""" self._register(monkeypatch, self.MANTLE_MODEL, "bedrock_mantle") - assert ( - AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "azure") is False - ) + assert AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(self.MANTLE_MODEL, "azure") is False def test_openai_entries_still_go_through_the_api_base_check(self, monkeypatch): """gpt-5.6 is flagged and openai-provided, so it must not bypass the host gate.""" @@ -3968,7 +4064,8 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: def test_every_flagged_mantle_row_is_an_openai_gpt_5_6_or_newer_model(self): flagged = { - key for key, entry in litellm.model_cost.items() + key + for key, entry in litellm.model_cost.items() if key.startswith("bedrock_mantle/") and entry.get("supports_prompt_cache_breakpoint") is True } assert self.MANTLE_MODEL in flagged @@ -3995,7 +4092,9 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: prompt_variables=None, dynamic_callback_params={}, ) - assert messages[0]["content"] == [{"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}}] + assert messages[0]["content"] == [ + {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] assert params["prompt_cache_options"] == {"mode": "implicit"} assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 1 @@ -4009,7 +4108,9 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: prompt_variables=None, dynamic_callback_params={}, ) - assert messages[0]["content"] == [{"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}}] + assert messages[0]["content"] == [ + {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} + ] assert params["prompt_cache_options"] == {"mode": "implicit"} BARE_MODEL = "openai.gpt-5.6-sol" @@ -4052,10 +4153,10 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: assert out["prompt_cache_options"] == {"mode": "implicit"} def test_the_responses_stamp_carries_the_resolved_provider_for_a_bare_deployment_name(self): - from litellm.responses.main import _stamp_injection_points_with_dialect + from litellm.responses.main import _stamp_injection_points_for_the_hook kwargs: dict = {"cache_control_injection_points": copy.deepcopy(self.POINTS)} - _stamp_injection_points_with_dialect(kwargs, self.BARE_MODEL, "bedrock_mantle") + _stamp_injection_points_for_the_hook(kwargs, self.BARE_MODEL, "bedrock_mantle", None, None) _, messages, out = AnthropicCacheControlHook().get_chat_completion_prompt( model=self.BARE_MODEL, messages=copy.deepcopy(self.MESSAGES), @@ -4068,6 +4169,156 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: assert out["prompt_cache_options"] == {"mode": "implicit"} +class TestResponsesLayerReservesMarksOutsideTheProvisionalMessages: + """The Responses layer's hook pass sees only the role-bearing input items, while the bridge + forwards a function_call or function_call_output mark onto the tool call or tool message it + builds, so the pass reserves those marks before placing an ordinal point and the bridge's + pass, which sees them in its own messages, replaces the reservation with its own count.""" + + MODEL = "claude-sonnet-4-6" + MARK = {"type": "ephemeral"} + + @classmethod + def _items(cls, marked_calls: int) -> list[dict]: + calls = [ + { + "type": "function_call", + "call_id": f"call_{index}", + "name": "weather", + "arguments": "{}", + "cache_control": cls.MARK, + } + for index in range(marked_calls) + ] + outputs = [ + {"type": "function_call_output", "call_id": f"call_{index}", "output": "sunny"} + for index in range(marked_calls) + ] + return [ + {"role": "user", "content": "ask", "cache_control": cls.MARK}, + *calls, + *outputs, + {"role": "user", "content": "next"}, + ] + + @classmethod + def _tools(cls, marked_tools: int) -> list[dict[str, object]]: + return [ + {"type": "function", "name": f"tool_{index}", "parameters": {}, "cache_control": cls.MARK} + for index in range(marked_tools) + ] + + def _marks_on_the_last_message_after_the_responses_pass( + self, marked_calls: int, marked_tools: int = 0, other_items: tuple[dict[str, object], ...] = () + ) -> int: + from litellm.responses.main import _stamp_injection_points_for_the_hook + from litellm.responses.utils import ResponsesAPIRequestUtils + + *earlier, last = self._items(marked_calls) + items = [*earlier, *other_items, last] + kwargs: dict = { + "cache_control_injection_points": [{"location": "message", "index": -1}], + CARRY_UNMATCHED_MESSAGE_POINTS: True, + } + _stamp_injection_points_for_the_hook(kwargs, self.MODEL, "anthropic", items, self._tools(marked_tools)) + _, messages, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model=self.MODEL, + messages=ResponsesAPIRequestUtils.responses_input_to_chat_messages(items), + non_default_params=kwargs, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + return AnthropicCacheControlHook._count_cache_control_blocks(messages[-1]) + + def test_function_call_marks_filling_the_cap_leave_the_ordinal_point_unplaced(self): + assert self._marks_on_the_last_message_after_the_responses_pass(marked_calls=3) == 0 + + def test_function_call_marks_under_the_cap_leave_the_ordinal_point_its_slot(self): + assert self._marks_on_the_last_message_after_the_responses_pass(marked_calls=2) == 1 + + def test_function_tool_marks_filling_the_cap_leave_the_ordinal_point_unplaced(self): + assert self._marks_on_the_last_message_after_the_responses_pass(marked_calls=1, marked_tools=2) == 0 + + def test_function_tool_marks_under_the_cap_leave_the_ordinal_point_its_slot(self): + assert self._marks_on_the_last_message_after_the_responses_pass(marked_calls=1, marked_tools=1) == 1 + + def test_marked_reasoning_items_leave_the_ordinal_point_its_slot(self): + reasoning = tuple( + { + "type": "reasoning", + "id": f"rs_{index}", + "summary": [{"type": "summary_text", "text": "checked the weather"}], + "cache_control": self.MARK, + } + for index in range(3) + ) + assert self._marks_on_the_last_message_after_the_responses_pass(marked_calls=0, other_items=reasoning) == 1 + + def test_a_marked_tool_of_another_type_is_not_reserved(self): + from litellm.responses.main import _cache_control_marks_outside_the_provisional_messages + + tools = [{"type": "web_search", "cache_control": self.MARK}, *self._tools(1)] + assert _cache_control_marks_outside_the_provisional_messages(self._items(1), tools) == 2 + + def test_sdk_typed_output_items_are_skipped_by_the_reserve(self): + from openai.types.responses import ResponseFunctionToolCall + + from litellm.responses.main import _cache_control_marks_outside_the_provisional_messages + + items = [ + *self._items(1), + ResponseFunctionToolCall(type="function_call", call_id="call_typed", name="weather", arguments="{}"), + ] + assert _cache_control_marks_outside_the_provisional_messages(items, None) == 1 + + def test_sdk_typed_tools_are_skipped_by_the_reserve(self): + from openai.types.responses import FunctionTool + + from litellm.responses.main import _cache_control_marks_outside_the_provisional_messages + + tools = [*self._tools(1), FunctionTool(type="function", name="typed", parameters={}, strict=False)] + assert _cache_control_marks_outside_the_provisional_messages(self._items(1), tools) == 2 + + def test_the_chat_pass_replaces_the_count_the_responses_layer_stamped(self): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "ask", "cache_control": self.MARK}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_0", + "type": "function", + "function": {"name": "weather", "arguments": "{}"}, + "cache_control": self.MARK, + } + ], + }, + {"role": "tool", "tool_call_id": "call_0", "content": "sunny"}, + {"role": "user", "content": "next"}, + ] + params: dict = { + "cache_control_injection_points": [{"location": "message", "role": "system", EXTERNAL_BREAKPOINTS_STAMP: 3}] + } + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=copy.deepcopy(messages), + model=self.MODEL, + custom_llm_provider="anthropic", + ) + _, out, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model=self.MODEL, + messages=copy.deepcopy(messages), + non_default_params=params, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + assert AnthropicCacheControlHook._count_cache_control_blocks(out[0]) == 1 + + class TestRecordGatewayInjection: """The injection marker spend accounting gates prompt-caching savings on.""" @@ -4274,8 +4525,8 @@ class TestMalformedInjectionPointsAreIgnored: @pytest.mark.parametrize("value", SHAPES) def test_responses_dialect_stamp_leaves_the_request_untouched(self, value): - from litellm.responses.main import _stamp_injection_points_with_dialect + from litellm.responses.main import _stamp_injection_points_for_the_hook kwargs: dict = {"cache_control_injection_points": copy.deepcopy(value)} - _stamp_injection_points_with_dialect(kwargs, "gpt-5.6", "openai") + _stamp_injection_points_for_the_hook(kwargs, "gpt-5.6", "openai", None, None) assert kwargs == {"cache_control_injection_points": value} diff --git a/tests/unit/integrations/test_prometheus_metric_name_consistency.py b/tests/unit/integrations/test_prometheus_metric_name_consistency.py index d648afcd087..4db34274e82 100644 --- a/tests/unit/integrations/test_prometheus_metric_name_consistency.py +++ b/tests/unit/integrations/test_prometheus_metric_name_consistency.py @@ -21,7 +21,6 @@ from prometheus_client import REGISTRY, Gauge from prometheus_client.registry import Collector import litellm -from litellm.caching.redis_cache import _breaker_metrics from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.prometheus_services import PrometheusServicesLogger from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import SpendLogCleanupMetrics @@ -62,10 +61,7 @@ def _lazy_owner_collectors() -> tuple[Collector, ...]: assert SpendLogCleanupMetrics.runs is not None in_flight: Final = InFlightRequestsMiddleware._get_gauge() assert in_flight is not None - breaker: Final = _breaker_metrics() - assert breaker._state_gauge is not None - assert breaker._transitions is not None - assert breaker._failures is not None + services: Final = PrometheusServicesLogger() return ( SpendLogCleanupMetrics.rows_deleted, SpendLogCleanupMetrics.batch_duration, @@ -73,9 +69,9 @@ def _lazy_owner_collectors() -> tuple[Collector, ...]: SpendLogCleanupMetrics.batch_failures, SpendLogCleanupMetrics.runs, in_flight, - breaker._state_gauge, - breaker._transitions, - breaker._failures, + services.redis_breaker_state, + services.redis_breaker_transitions, + services.redis_breaker_failures, ) diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 2c3dae21058..e9d96f057ae 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -11,6 +11,7 @@ content frame, restoring per-token cadence on cache hits. import pytest +import litellm from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _split_assembled_content_for_replay, convert_to_streaming_response, @@ -87,6 +88,71 @@ async def test_async_yields_multiple_content_chunks_with_lossless_join(): assert reassembled == text +@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"]) +@pytest.mark.asyncio +async def test_async_and_sync_replays_carry_the_cached_provider_specific_fields(content): + provider_specific_fields = {"stop_reason": "refusal", "stop_details": {"category": "cyber"}} + + def payload(): + cached = _async_payload(content=content) + cached["choices"][0]["message"]["provider_specific_fields"] = provider_specific_fields + return cached + + async_chunks = await _collect_async(payload()) + sync_chunks = list(convert_to_streaming_response(payload())) + + assert [c.choices[0].delta.provider_specific_fields for c in async_chunks] == [ + c.choices[0].delta.provider_specific_fields for c in sync_chunks + ] + assert async_chunks[0].choices[0].delta.provider_specific_fields == provider_specific_fields + + +@pytest.mark.parametrize("provider_fields", [{"thought_signatures": ["sig"]}, None], ids=["provider_fields", "text"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_the_cached_completion_tokens_details(replay, provider_fields): + cached = _async_payload() + cached["choices"][0]["message"]["provider_specific_fields"] = provider_fields + cached["usage"]["completion_tokens_details"] = {"reasoning_tokens": 4, "text_tokens": 3} + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.get("provider_specific_fields") == provider_fields + assert rebuilt.usage.completion_tokens_details.reasoning_tokens == 4 + assert rebuilt.usage.completion_tokens_details.text_tokens == 3 + + +@pytest.mark.parametrize("cached_fields", [7, [1, 2], "", "x" * 5000], ids=["int", "list", "empty", "long"]) +@pytest.mark.asyncio +async def test_the_async_replay_ignores_cached_provider_fields_that_are_not_an_object(cached_fields): + cached = _async_payload() + cached["choices"][0]["message"]["provider_specific_fields"] = cached_fields + + chunks = await _collect_async(cached) + + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "Hello world! How are you?" + assert all(c.choices[0].delta.provider_specific_fields is None for c in chunks) + + +@pytest.mark.parametrize( + "details", + [{"reasoning_tokens": "abc"}, {"reasoning_tokens": 1.5}, {"reasoning_tokens": {"a": 1}}], + ids=["text", "fraction", "object"], +) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_not_validate(replay, details): + cached = _async_payload() + cached["usage"]["completion_tokens_details"] = details + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.content == "Hello world! How are you?" + assert rebuilt.usage.completion_tokens == 7 + + @pytest.mark.asyncio async def test_async_finish_reason_only_on_last_chunk(): chunks = await _collect_async(_async_payload()) @@ -204,9 +270,7 @@ def test_sync_multi_chunk_and_lossless_join(): def test_sync_delta_and_choice_metadata_only_on_first_chunk(): - thinking_blocks = [ - {"type": "thinking", "thinking": "cached thinking", "signature": "sig"} - ] + thinking_blocks = [{"type": "thinking", "thinking": "cached thinking", "signature": "sig"}] payload = { "id": "chatcmpl-test-sync", "object": "chat.completion", diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 0172f0e9aeb..a604e61b9c2 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -732,6 +732,42 @@ def test_set_response_id_propagation_valid_to_invalid( assert initialized_custom_stream_wrapper.response_id == "first-valid-id" +@pytest.mark.parametrize( + "last_delta", + [ + Delta(content="Hello"), + Delta(provider_specific_fields={"stop_reason": "refusal"}), + ], + ids=["content", "provider_specific_fields"], +) +def test_logged_response_keeps_received_model_id_when_the_finish_chunk_is_synthesized( + logging_obj: Logging, last_delta: Delta +): + upstream_id: Final = "chatcmpl-upstream-id" + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[ + ModelResponseStream( + id=upstream_id, + model="gpt-4o-mini", + choices=[StreamingChoices(index=0, delta=last_delta, finish_reason="content_filter")], + ) + ] + ), + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + emitted: Final = list(wrapper) + logged: Final = litellm.stream_chunk_builder(chunks=wrapper.chunks) + + assert len(emitted) == 2 + assert [chunk._hidden_params.get("received_model_id") for chunk in emitted] == [upstream_id, upstream_id] + assert logged is not None + assert logged._hidden_params.get("received_model_id") == upstream_id + + @pytest.mark.asyncio async def test_streaming_completion_start_time(logging_obj: Logging): """Test that the start time is set correctly""" diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index e60f049a8ed..75ea6e2a698 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -5,6 +5,7 @@ import re import subprocess import sys import threading +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Final, NoReturn, Optional @@ -2083,3 +2084,101 @@ def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch assert ProxyExtrasDBManager.repair_invalid_indexes() is False + +@pytest.mark.parametrize( + "environment", + ( + {"DATABASE_URL": "postgresql://pool/db"}, + {"DATABASE_URL": "postgresql://pool/db?pgbouncer=true", "DIRECT_URL": "postgresql://writer/db"}, + ), + ids=("pooled-only", "pooled-with-direct"), +) +def test_prisma_runs_its_migrations_without_a_session_advisory_lock(environment): + from litellm_proxy_extras.migration_lock import migration_environment + + migrated = migration_environment(environment) + + assert migrated["PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK"] == "1" + assert "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK" not in migration_environment({"PRISMA_OFFLINE_MODE": "true"}) + + +class TestMigrateDeployRunsUnderTheCoordinatorLock: + """`prisma migrate deploy` only runs while the injected lock is held, and the lock is + released again before the resolver handles the outcome, success or failure.""" + + @staticmethod + def _recording_lock(events: list[str]): + from contextlib import contextmanager + + @contextmanager + def lock(): + events.append("lock") + try: + yield + finally: + events.append("unlock") + + return lock + + def test_deploy_runs_between_lock_and_unlock(self, monkeypatch, tmp_path): + import litellm_proxy_extras.utils as utils_module + + events: list[str] = [] + harness = _MigrateDeployHarness(monkeypatch, tmp_path, ["ok"]) + original_run = harness._fake_run + + def run_and_record(cmd, **kwargs): + events.append("deploy") + return original_run(cmd, **kwargs) + + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record) + + assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True + assert events == ["lock", "deploy", "unlock"] + + def test_a_failed_deploy_releases_the_lock_before_the_retry(self, monkeypatch, tmp_path): + import litellm_proxy_extras.utils as utils_module + + events: list[str] = [] + harness = _MigrateDeployHarness( + monkeypatch, tmp_path, ["Error: P1002\nTimed out waiting for the advisory lock", "ok"] + ) + original_run = harness._fake_run + + def run_and_record(cmd, **kwargs): + events.append("deploy") + return original_run(cmd, **kwargs) + + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record) + + assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True + assert events == ["lock", "deploy", "unlock", "lock", "deploy", "unlock"] + + def test_a_deadlock_rollback_runs_under_the_coordinator_lock( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + import litellm_proxy_extras.utils as utils_module + + migration_name: Final = "20260921190000_agent_identity" + stderr: Final = ( + "Error: P3018\n\n" + "A migration failed to apply. New migrations cannot be applied before the error is recovered from.\n\n" + f"Migration name: {migration_name}\n\n" + "Database error code: 40P01\n\n" + f"Database error:\n{_P3009_DEADLOCK_LOGS}" + ) + events: Final[list[str]] = [] + harness: Final = _MigrateDeployHarness(monkeypatch, tmp_path, [stderr, "ok"]) + original_run: Final = harness._fake_run + + def run_and_record(cmd: Sequence[str], *, timeout: float, env: Mapping[str, str]) -> _FakeCompleted: + if cmd[1:] == ["migrate", "deploy"]: + events.append("deploy") + elif cmd[1:4] == ["migrate", "resolve", "--rolled-back"]: + events.append("rollback") + return original_run(cmd, timeout=timeout, env=env) + + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record) + + assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True + assert events == ["lock", "deploy", "unlock", "lock", "rollback", "unlock", "lock", "deploy", "unlock"] diff --git a/tests/unit/llms/anthropic/anthropic_docs_refusal.py b/tests/unit/llms/anthropic/anthropic_docs_refusal.py new file mode 100644 index 00000000000..2bd75b4aab5 --- /dev/null +++ b/tests/unit/llms/anthropic/anthropic_docs_refusal.py @@ -0,0 +1,19 @@ +from typing import Final + +ANTHROPIC_DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response" + +# Copied verbatim from ANTHROPIC_DOCS_REFUSAL_URL on 2026-10-12. Streams carry the same stop_details on message_delta. +ANTHROPIC_DOCS_REFUSAL_MESSAGE: Final = { + "id": "msg_01XFUDYJgAACzvnptvVoYEL", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", + }, + "usage": {"input_tokens": 412, "output_tokens": 0}, +} diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 47375afa90f..b46f1707a26 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -22,6 +22,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import OutputCodeInterpreterCall from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponse, ModelResponseStream +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @pytest.mark.asyncio @@ -3434,3 +3435,50 @@ def test_anthropic_streaming_invalid_key_maps_to_authentication_error() -> None: assert exc_info.value.status_code == 401 assert exc_info.value.llm_provider == "anthropic" assert route.calls.last.request.headers["x-api-key"] == "bad-anthropic-key" + + +def _message_stream_events(stop_reason: str, stop_details: dict | None) -> tuple[dict, ...]: + return ( + { + "type": "message_start", + "message": { + "id": ANTHROPIC_DOCS_REFUSAL_MESSAGE["id"], + "type": "message", + "role": "assistant", + "model": ANTHROPIC_DOCS_REFUSAL_MESSAGE["model"], + "content": [], + "usage": {"input_tokens": 412, "output_tokens": 0}, + }, + }, + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None, "stop_details": stop_details}, + "usage": {"output_tokens": 0}, + }, + {"type": "message_stop"}, + ) + + +def _assemble_stream(events: tuple[dict, ...]) -> litellm.ModelResponse: + iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False) + chunks: Final = [iterator.chunk_parser(event) for event in events] + return litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + +def test_streamed_refusal_stop_details_reach_the_assembled_message(): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + assembled: Final = _assemble_stream(_message_stream_events("refusal", stop_details)) + + assert assembled.choices[0].finish_reason == "content_filter" + provider_fields: Final = assembled.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == stop_details + + +def test_streamed_end_turn_carries_no_stop_details(): + assembled: Final = _assemble_stream(_message_stream_events("end_turn", None)) + + assert assembled.choices[0].finish_reason == "stop" + provider_fields: Final = assembled.choices[0].message.provider_specific_fields or {} + assert "stop_details" not in provider_fields + assert "stop_reason" not in provider_fields diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index d3e31ef0dee..9d40979cd7c 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -42,6 +42,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE def test_response_format_transformation_unit_test(): @@ -9168,3 +9169,93 @@ def test_completion_surfaces_response_citations_in_provider_specific_fields(resp assert response.choices[0].message.provider_specific_fields["citations"] == [ [{**citation, "supported_text": "the grass is green"}] ] + + +def _refusal_message_body(stop_reason: str, stop_details: dict | None) -> dict: + return {**ANTHROPIC_DOCS_REFUSAL_MESSAGE, "stop_reason": stop_reason, "stop_details": stop_details} + + +def _transform_message_body(body: dict) -> litellm.ModelResponse: + return AnthropicConfig().transform_parsed_response( + completion_response=body, + raw_response=httpx.Response(status_code=200, headers={}), + model_response=litellm.ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + + +def test_refusal_stop_details_kept_in_provider_specific_fields(): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + result: Final = _transform_message_body(ANTHROPIC_DOCS_REFUSAL_MESSAGE) + + assert result.choices[0].finish_reason == "content_filter" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == stop_details + assert result._hidden_params["provider_specific_fields"]["stop_details"] == stop_details + + +def test_refusal_without_stop_details_keeps_the_native_stop_reason(): + result: Final = _transform_message_body(_refusal_message_body("refusal", None)) + + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert "stop_details" not in provider_fields + + +def test_end_turn_response_carries_no_stop_details(): + result: Final = _transform_message_body(_refusal_message_body("end_turn", None)) + + assert result.choices[0].finish_reason == "stop" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert "stop_details" not in provider_fields + assert "stop_reason" not in provider_fields + + +def _json_mode_message_body(stop_reason: str, stop_details: dict | None) -> dict: + return { + **_refusal_message_body(stop_reason, stop_details), + "content": [ + {"type": "tool_use", "id": "toolu_json", "name": RESPONSE_FORMAT_TOOL_NAME, "input": {"answer": "Hello"}} + ], + } + + +def _transform_json_mode_body(body: dict) -> litellm.ModelResponse: + return AnthropicConfig().transform_parsed_response( + completion_response=body, + raw_response=httpx.Response(status_code=200, headers={}), + model_response=litellm.ModelResponse(), + json_mode=True, + prefix_prompt=None, + ) + + +def test_json_mode_refusal_keeps_stop_details_on_the_message(): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + result: Final = _transform_json_mode_body(_json_mode_message_body("refusal", stop_details)) + + message: Final = result.choices[0].message + assert json.loads(message.content) == {"answer": "Hello"} + assert message.provider_specific_fields == {"stop_reason": "refusal", "stop_details": stop_details} + + +def test_json_mode_tool_use_turn_adds_no_stop_fields_to_the_message(): + result: Final = _transform_json_mode_body(_json_mode_message_body("tool_use", None)) + + message: Final = result.choices[0].message + assert json.loads(message.content) == {"answer": "Hello"} + assert not message.provider_specific_fields + + +def test_non_dict_upstream_body_still_transforms_without_refusal_fields() -> None: + result = AnthropicConfig().transform_parsed_response( + completion_response=MagicMock(), + raw_response=httpx.Response(200, headers={}), + model_response=litellm.ModelResponse(), + json_mode=False, + prefix_prompt=None, + ) + provider_specific_fields = result.choices[0].message.provider_specific_fields or {} + assert "stop_details" not in provider_specific_fields diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py new file mode 100644 index 00000000000..6472ccd1b2e --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py @@ -0,0 +1,28 @@ +from typing import Final + +import pytest + +from litellm.llms.anthropic.pass_through.messages.utils import AnthropicMessagesRequestUtils + +REQUEST: Final[dict[str, object]] = { + "max_tokens": 64, + "service_tier": "standard_only", + "messages": [{"role": "user", "content": "Hello"}], +} + + +@pytest.mark.parametrize("custom_llm_provider", ["anthropic", "azure_ai", "bedrock_mantle", None]) +def test_messages_forwards_service_tier_on_anthropic_api_routes(custom_llm_provider: str | None) -> None: + optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params=dict(REQUEST), model="claude-sonnet-4-6", drop_params=False, custom_llm_provider=custom_llm_provider + ) + assert dict(optional_params).get("service_tier") == "standard_only" + + +@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "bedrock"]) +def test_messages_drops_service_tier_where_the_provider_rejects_it(custom_llm_provider: str) -> None: + optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params=dict(REQUEST), model="claude-sonnet-4-6", drop_params=False, custom_llm_provider=custom_llm_provider + ) + assert "service_tier" not in optional_params + assert dict(optional_params).get("max_tokens") == 64 diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index d9b01c4e75a..1890c68b30d 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -343,15 +343,18 @@ def test_router_deployment_stream_chunk_size_reaches_iter_bytes( client: Final = HTTPHandler() client.post = MagicMock(return_value=mock_response) router: Final = litellm.Router( - model_list=[{"model_name": "invoke-chunked", "litellm_params": {**INVOKE_DEPLOYMENT, **deployment_extras}}] + model_list=[{"model_name": "invoke-chunked", "litellm_params": {**INVOKE_DEPLOYMENT, **deployment_extras}}], + num_retries=0, ) - router.completion( - model="invoke-chunked", - messages=[{"role": "user", "content": "hi"}], - stream=True, - client=client, - ) + # the empty body is a stream that ends before its first chunk, which the Router now reports from the call + with pytest.raises(litellm.BadGatewayError): + router.completion( + model="invoke-chunked", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size) data: Final = client.post.call_args.kwargs["data"] diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 859b57c7cf3..8cb9ad9b73c 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -4,7 +4,7 @@ import json import uuid from types import SimpleNamespace from typing import Final -from unittest.mock import patch +from unittest.mock import MagicMock, patch import httpx import pytest @@ -17,6 +17,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE ONE_PIXEL_PNG = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" @@ -1689,3 +1690,23 @@ class TestBedrockAnthropicCombinedRegressions: user_msg = result["messages"][0] has_cache_control = any(isinstance(c, dict) and "cache_control" in c for c in user_msg["content"]) assert has_cache_control + + +def test_refusal_stop_details_kept_through_bedrock_invoke_transform(): + body: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE + result: Final = AmazonAnthropicClaudeConfig().transform_response( + model="anthropic.claude-sonnet-5-5-v1:0", + raw_response=httpx.Response(status_code=200, headers={}, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].finish_reason == "content_filter" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] diff --git a/tests/unit/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py index 6279b74bc0d..d0360b07957 100644 --- a/tests/unit/llms/chat/test_converse_handler.py +++ b/tests/unit/llms/chat/test_converse_handler.py @@ -417,15 +417,18 @@ def test_router_deployment_stream_chunk_size_reaches_iter_bytes( client: Final = HTTPHandler() client.post = MagicMock(return_value=mock_response) router: Final = litellm.Router( - model_list=[{"model_name": "converse-chunked", "litellm_params": {**CONVERSE_DEPLOYMENT, **deployment_extras}}] + model_list=[{"model_name": "converse-chunked", "litellm_params": {**CONVERSE_DEPLOYMENT, **deployment_extras}}], + num_retries=0, ) - router.completion( - model="converse-chunked", - messages=[{"role": "user", "content": "hi"}], - stream=True, - client=client, - ) + # the empty body is a stream that ends before its first chunk, which the Router now reports from the call + with pytest.raises(litellm.BadGatewayError): + router.completion( + model="converse-chunked", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size) data: Final = client.post.call_args.kwargs["data"] diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index f9e7628f23e..e47d7fc702c 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -371,6 +371,7 @@ class TestChatGPTResponsesAPITransformation: "type": "error", "error": {"message": "ChatGPT upstream failed"}, }, + {"type": "error", "code": "server_error", "message": "ChatGPT upstream failed"}, ], ) def test_chatgpt_non_stream_sse_response_raises_openai_error(self, error_chunk): @@ -392,5 +393,5 @@ class TestChatGPTResponsesAPITransformation: logging_obj=logging_obj, ) - assert "ChatGPT upstream failed" in str(exc_info.value) + assert exc_info.value.message == "ChatGPT upstream failed" assert exc_info.value.status_code == 502 diff --git a/tests/unit/llms/litellm_proxy/skills/test_skill_search.py b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py index e28e927cf61..c0754f6f1af 100644 --- a/tests/unit/llms/litellm_proxy/skills/test_skill_search.py +++ b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py @@ -89,6 +89,19 @@ class FixedDimensionEmbedder: return tuple((1.0,) * self.dimensions for _ in texts) +class OversizedAwareEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + if any("OVERSIZED" in text for text in texts): + raise litellm.ContextWindowExceededError( + message="input is too long", model="text-embedding-3-large", llm_provider="openai" + ) + return tuple(VECTORS[text] for text in texts) + + class TestSkillSearchText: def test_joins_title_description_and_instructions(self) -> None: assert skill_search_text(TRANSLATOR) == ( @@ -133,6 +146,19 @@ class TestSkillSearchIndex: assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file", "trip-planner"] assert outcome.hits[0].score > outcome.hits[1].score + @pytest.mark.asyncio + async def test_omits_skills_with_oversized_search_text(self) -> None: + oversized: Final = LiteLLM_SkillsTable( + skill_id="oversized", display_title="Oversized", instructions="OVERSIZED instructions" + ) + + outcome: Final = await SkillSearchIndex().search( + "language translation", (*SKILLS, oversized), top_k=5, embed=OversizedAwareEmbedder(), embedding_model="m" + ) + + assert isinstance(outcome, SkillSearchHits) + assert {hit.skill.skill_id for hit in outcome.hits} == {skill.skill_id for skill in SKILLS} + @pytest.mark.asyncio async def test_second_search_only_embeds_the_query(self) -> None: index = SkillSearchIndex() diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 8bd242308f7..57e66032e5a 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -898,10 +898,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert tool_call["id"] == "call_4SjsMeA6DUHwGKaE87ZojgOF" assert tool_call["type"] == "function" assert tool_call["function"]["name"] == "get_current_weather" - assert ( - tool_call["function"]["arguments"] - == '{"location":"Boston, MA","unit":"celsius"}' - ) + assert tool_call["function"]["arguments"] == '{"location":"Boston, MA","unit":"celsius"}' assert tool_call["index"] == 0 def test_extract_tool_call_from_dict_format(self): @@ -942,10 +939,7 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert tool_call["id"] == "call_4SjsMeA6DUHwGKaE87ZojgOF" assert tool_call["type"] == "function" assert tool_call["function"]["name"] == "get_current_weather" - assert ( - tool_call["function"]["arguments"] - == '{"location":"Boston, MA","unit":"celsius"}' - ) + assert tool_call["function"]["arguments"] == '{"location":"Boston, MA","unit":"celsius"}' @pytest.mark.parametrize( "output_item", @@ -1481,9 +1475,7 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: logging_obj: Optional[Any] = None, ) -> GenericGuardrailAPIInputs: texts = inputs.get("texts", []) - inputs["texts"] = [ - t.replace("", "john@example.com") for t in texts - ] + inputs["texts"] = [t.replace("", "john@example.com") for t in texts] return inputs handler = OpenAIResponsesHandler() @@ -1519,15 +1511,11 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: litellm_logging_obj=None, ) - completed_chunk = next( - c - for c in result - if isinstance(c, dict) and c.get("type") == "response.completed" - ) + completed_chunk = next(c for c in result if isinstance(c, dict) and c.get("type") == "response.completed") output_text = completed_chunk["response"]["output"][0]["content"][0]["text"] - assert ( - output_text == "send to john@example.com" - ), f"Expected PII token to be unmasked in response.completed output, got: {output_text!r}" + assert output_text == "send to john@example.com", ( + f"Expected PII token to be unmasked in response.completed output, got: {output_text!r}" + ) @pytest.mark.asyncio async def test_process_output_streaming_response_pass_through_unchanged(self): @@ -1585,7 +1573,11 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: "content_index": 0, "part": {"type": "output_text", "text": "hello world"}, }, - {"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}}, + { + "type": "response.output_item.done", + "output_index": 0, + "item": {**item, "content": [dict(c) for c in content]}, + }, { "type": "response.completed", "response": { @@ -1646,15 +1638,29 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: return [ {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, - {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'}, - {"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_123", + "output_index": 0, + "delta": '{"fruit":', + }, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_123", + "output_index": 0, + "delta": ' "persimmon"}', + }, { "type": "response.function_call_arguments.done", "item_id": "fc_123", "output_index": 0, "arguments": '{"fruit": "persimmon"}', }, - {"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.output_item.done", + "output_index": 0, + "item": item('{"fruit": "persimmon"}', "completed"), + }, { "type": "response.completed", "response": { @@ -1759,8 +1765,18 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: return [ {"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")}, {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "echo "}, - {"type": "response.custom_tool_call_input.delta", "item_id": "ctc_1", "output_index": 0, "delta": "persimmon"}, - {"type": "response.custom_tool_call_input.done", "item_id": "ctc_1", "output_index": 0, "input": "echo persimmon"}, + { + "type": "response.custom_tool_call_input.delta", + "item_id": "ctc_1", + "output_index": 0, + "delta": "persimmon", + }, + { + "type": "response.custom_tool_call_input.done", + "item_id": "ctc_1", + "output_index": 0, + "input": "echo persimmon", + }, {"type": "response.output_item.done", "output_index": 0, "item": item("echo persimmon", "completed")}, { "type": "response.completed", @@ -1893,19 +1909,43 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: return [ {"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)}, {"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)}, - {"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}}, - {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"}, + { + "type": "response.output_item.added", + "output_index": 0, + "item": {**message, "status": "in_progress", "content": []}, + }, + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "Looking that up", + }, {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}}, {"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")}, - {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'}, - {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'}, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_1", + "output_index": 1, + "delta": '{"fruit":', + }, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_1", + "output_index": 1, + "delta": ' "persimmon"}', + }, { "type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 1, "arguments": '{"fruit": "persimmon"}', }, - {"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")}, + { + "type": "response.output_item.done", + "output_index": 1, + "item": function_call('{"fruit": "persimmon"}', "completed"), + }, { "type": "response.completed", "response": { @@ -2303,9 +2343,7 @@ class TestGetStructuredMessages: } result = handler.get_structured_messages(data) assert result is not None - has_system = any( - isinstance(msg, dict) and msg.get("role") == "system" for msg in result - ) + has_system = any(isinstance(msg, dict) and msg.get("role") == "system" for msg in result) assert has_system, f"Expected system message from instructions, got: {result}" def test_should_return_none_when_no_input(self): @@ -2405,9 +2443,7 @@ class StructuredRewriteGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: messages = list(inputs.get("structured_messages") or []) first_user = next(i for i, m in enumerate(messages) if m.get("role") == "user") - rewritten = [ - {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) - ] + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages)] return {**inputs, "structured_messages": rewritten} @@ -2453,9 +2489,7 @@ class ToolOutputRewriteGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: messages = list(inputs.get("structured_messages") or []) first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool") - rewritten = [ - {**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages) - ] + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages)] return {**inputs, "structured_messages": rewritten} @@ -2472,9 +2506,7 @@ class DroppingRewriteGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: messages = list(inputs.get("structured_messages") or []) first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user") - rewritten = [ - {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) - ] + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages)] return {**inputs, "structured_messages": rewritten[:-1]} @@ -2520,6 +2552,21 @@ class TestStructuredMessagesWriteBack: assert [_texts(item) for item in result["input"]] == [[COMPRESSED_MARKER]] assert "instructions" not in result + @pytest.mark.asyncio + async def test_rows_the_guardrail_left_alone_keep_their_cache_control_past_the_provider_cap(self): + handler = OpenAIResponsesHandler() + ephemeral = {"type": "ephemeral"} + marked_turns = [{"role": "user", "content": f"turn {index}", "cache_control": ephemeral} for index in range(4)] + data = { + "model": "bedrock/us.anthropic.claude-sonnet-4-6", + "input": [{"role": "user", "content": "memo " * 400, "cache_control": ephemeral}, *marked_turns], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert _texts(result["input"][0]) == [COMPRESSED_MARKER] + assert result["input"][1:] == marked_turns + @pytest.mark.asyncio async def test_developer_item_preserved_verbatim_by_row_patch(self): handler = OpenAIResponsesHandler() @@ -2816,6 +2863,20 @@ class TestPerMessageRewriteWriteBack: only where every row has a scanned text (instructions plus a string input) and is otherwise rejected by name instead of sent unrewritten.""" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "cache_control", [{"type": "ephemeral", "ttl": "30m"}, {"type": "persistent"}], ids=["ttl-30m", "unknown-type"] + ) + async def test_an_item_mark_no_chat_row_can_hold_still_lets_the_guardrail_run(self, cache_control: dict): + guardrail = _per_message_redactor() + marked_item = {"role": "user", "content": "Look up " + SSN + " for me.", "cache_control": cache_control} + data = {"model": "gemini/gemini-2.5-flash", "input": [marked_item]} + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["input"] == [{**marked_item, "content": "Look up " + REDACTED_SSN + " for me."}] + @pytest.mark.asyncio async def test_structured_rows_land_on_instructions_and_tool_output(self): guardrail = _per_message_redactor() @@ -3034,9 +3095,7 @@ class SystemRewriteGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: messages = list(inputs.get("structured_messages") or []) first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system") - rewritten = [ - {**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages) - ] + rewritten = [{**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages)] return {**inputs, "structured_messages": rewritten} @@ -3411,9 +3470,7 @@ class TestBuildBlockSseChunks: {"type": "message", "id": "msg_live", "status": "in_progress", "role": "assistant", "content": []} ) yielded = [ - OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=open_item - ), + OutputItemAddedEvent(type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=open_item), ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id="msg_live", @@ -3627,7 +3684,11 @@ class TestOpenAIResponsesHandlerStreamingScanKey: "delta": '{"city":', } function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city":'} - terminal = {"type": terminal_type, "sequence_number": 2, "response": {"id": "resp_1", "output": [function_call]}} + terminal = { + "type": terminal_type, + "sequence_number": 2, + "response": {"id": "resp_1", "output": [function_call]}, + } mid_stream_key = handler.get_streaming_scan_key([arguments_delta]) ended_key = handler.get_streaming_scan_key([arguments_delta, terminal]) assert ended_key.stream_ended is True @@ -3661,7 +3722,13 @@ class TestOpenAIResponsesHandlerStreamingScanKey: added = { "type": "response.output_item.added", "sequence_number": 1, - "item": {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather", "arguments": ""}, + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": "", + }, } arguments_delta = { "type": "response.function_call_arguments.delta", @@ -3686,7 +3753,11 @@ class TestOpenAIResponsesHandlerStreamingScanKey: arguments = json.dumps({"city": city}) pending = {"type": "function_call", "id": item_id, "call_id": "call_" + item_id, "name": "get_weather"} return ( - {"type": "response.output_item.added", "sequence_number": sequence_number, "item": {**pending, "arguments": ""}}, + { + "type": "response.output_item.added", + "sequence_number": sequence_number, + "item": {**pending, "arguments": ""}, + }, { "type": "response.function_call_arguments.delta", "sequence_number": sequence_number + 1, @@ -3710,7 +3781,9 @@ class TestOpenAIResponsesHandlerStreamingScanKey: ended_key = handler.get_streaming_scan_key(handler.released_stream_as_ended(released)) assert ended_key.stream_ended is True assert ended_key.texts == ("hi",) - cities = tuple(city for fingerprint in ended_key.tool_calls for city in ("Paris", "Rome") if city in fingerprint) + cities = tuple( + city for fingerprint in ended_key.tool_calls for city in ("Paris", "Rome") if city in fingerprint + ) assert cities == ("Paris", "Rome"), ended_key def test_released_stream_as_ended_keys_a_message_whose_item_already_finished(self): @@ -3720,7 +3793,9 @@ class TestOpenAIResponsesHandlerStreamingScanKey: "sequence_number": 1, "item": {"type": "message", "id": "msg_1", "content": [{"type": "output_text", "text": "hi"}]}, } - ended_key = handler.get_streaming_scan_key(handler.released_stream_as_ended((self._delta(0, "hi"), message_done))) + ended_key = handler.get_streaming_scan_key( + handler.released_stream_as_ended((self._delta(0, "hi"), message_done)) + ) assert ended_key.stream_ended is True assert ended_key.texts == ("hi",) diff --git a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py index 344b8e5573d..8ddbb7c1a6b 100644 --- a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py @@ -18,7 +18,7 @@ import pytest from litellm.llms.snowflake.chat.transformation import ( SnowflakeConfig, - _is_claude_model, + is_snowflake_claude_model, ) from litellm.types.utils import ModelResponse @@ -69,6 +69,7 @@ def _make_anthropic_response(content: str = "Hello!") -> httpx.Response: # ─── SnowflakeConfig (OpenAI-compatible) ─────────────────────────────────── + class TestSnowflakeConfigURL: def setup_method(self): self.cfg = SnowflakeConfig() @@ -266,6 +267,7 @@ class TestSnowflakeConfigResponse: # ─── SnowflakeConfig ──────────────────────────────────────── + class TestAnthropicConfigURL: def setup_method(self): self.cfg = SnowflakeConfig() @@ -533,21 +535,23 @@ class TestAnthropicConfigResponse: # ─── Model detection helper ──────────────────────────────────────────────── + class TestIsClaudeModel: def test_claude_model_detected(self): - assert _is_claude_model("snowflake/claude-sonnet-4-5") is True - assert _is_claude_model("claude-3-haiku") is True - assert _is_claude_model("snowflake/claude-opus-4") is True + assert is_snowflake_claude_model("snowflake/claude-sonnet-4-5") is True + assert is_snowflake_claude_model("claude-3-haiku") is True + assert is_snowflake_claude_model("snowflake/claude-opus-4") is True def test_non_claude_not_detected(self): - assert _is_claude_model("snowflake/llama3.1-70b") is False - assert _is_claude_model("snowflake/mistral-large") is False - assert _is_claude_model("snowflake/deepseek-r1") is False - assert _is_claude_model("snowflake/snowflake-arctic") is False + assert is_snowflake_claude_model("snowflake/llama3.1-70b") is False + assert is_snowflake_claude_model("snowflake/mistral-large") is False + assert is_snowflake_claude_model("snowflake/deepseek-r1") is False + assert is_snowflake_claude_model("snowflake/snowflake-arctic") is False # ─── Anthropic Tool Transformation Tests ────────────────────────────────── + class TestAnthropicToolTransformation: def setup_method(self): self.cfg = SnowflakeConfig() diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 56e76df25a4..ace24c8a3f4 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -3,6 +3,7 @@ import copy import json from collections.abc import Mapping from typing import Final +from unittest.mock import MagicMock import httpx import pytest @@ -19,6 +20,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im ) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VertexPartnerProvider +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @pytest.mark.parametrize( @@ -1125,3 +1127,23 @@ def test_chat_flagged_model_replays_a_byte_identical_prefix_around_a_mid_convers _assert_prefix_stable(requests) assert [m["role"] for m in requests[1]["messages"]] == ["user", "assistant", "user", "system"] assert [m["role"] for m in requests[2]["messages"]] == ["user", "assistant", "user", "system", "assistant", "user"] + + +def test_refusal_stop_details_kept_through_vertex_anthropic_transform(): + body: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE + result: Final = VertexAIAnthropicConfig().transform_response( + model="claude-sonnet-5-5", + raw_response=httpx.Response(status_code=200, headers={}, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].finish_reason == "content_filter" + provider_fields: Final = result.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py index 90403e5553f..1524aad5b66 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py @@ -7,7 +7,7 @@ from fastapi import HTTPException from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler -from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth from litellm.proxy.auth import auth_checks from litellm.types.agents import AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext @@ -114,6 +114,41 @@ async def test_access_groups_cap_agent_servers_without_granting_new_ones( assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == [] +@pytest.mark.asyncio +async def test_managed_agent_tool_ceiling_reads_the_admitted_policy_not_the_registry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.common_utils import registry_read_through + + groups: Final = { + "stale": LiteLLM_AccessGroupTable( + access_group_id="stale", access_group_name="Stale", access_mcp_server_ids=["slack"] + ), + "current": LiteLLM_AccessGroupTable( + access_group_id="current", + access_group_name="Current", + access_mcp_server_ids=["slack"], + mcp_tool_permissions={"slack": ["read"]}, + ), + } + + async def load_group(access_group_id: str, **_: object) -> LiteLLM_AccessGroupTable: + return groups[access_group_id] + + monkeypatch.setattr(auth_checks, "get_access_object", load_group) + auth: Final = actor(("read", "write")) + assert auth.managed_agent_policy is not None + stale_registry_row: Final = auth.managed_agent_policy.model_copy(update={"access_group_ids": ["stale"]}) + monkeypatch.setattr( + registry_read_through, "get_agent_with_read_through", AsyncMock(return_value=stale_registry_row) + ) + auth.managed_agent_policy = auth.managed_agent_policy.model_copy(update={"access_group_ids": ["current"]}) + + assert tuple(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == ("slack",) + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ["read"] + + @pytest.mark.asyncio @pytest.mark.parametrize("change", ["tools", "servers", "disabled", "outage"]) async def test_delegated_mcp_revokes_warm_human_policy_before_tool_execution( @@ -400,9 +435,9 @@ async def test_delegation_never_borrows_another_teams_server_or_tools( team_id=name, models=[], members_with_roles=[{"user_id": "human", "role": "user"}], - object_permission=permission if name == "selected" else LiteLLM_ObjectPermissionTable( - object_permission_id="other-grant", mcp_servers=["slack", "linear"] - ), + object_permission=permission + if name == "selected" + else LiteLLM_ObjectPermissionTable(object_permission_id="other-grant", mcp_servers=["slack", "linear"]), ) for name in ("selected", "other") } @@ -439,7 +474,9 @@ async def test_managed_mcp_rejects_unavailable_authoritative_entitlements( ) auth: Final = actor(None) assert auth.managed_agent_policy is not None - auth.managed_agent_policy = auth.managed_agent_policy.model_copy(update={"object_permission": permission.model_dump()}) + auth.managed_agent_policy = auth.managed_agent_policy.model_copy( + update={"object_permission": permission.model_dump()} + ) auth.requires_fresh_policy = True with pytest.raises(HTTPException) as failure: await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) @@ -484,6 +521,7 @@ async def test_managed_agent_mcp_access_is_capped_at_the_invoking_callers_grants "_apply_user_tool_ceiling", AsyncMock(side_effect=lambda tools, _server_id, _auth: tools), ) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="callers"))) auth: Final = actor(("read", "write")) auth.agent_caller = AgentCaller(user_id="alice", team_id="callers") @@ -496,25 +534,34 @@ async def test_managed_agent_mcp_access_is_capped_at_the_invoking_callers_grants @pytest.mark.parametrize("fresh", [False, True]) @pytest.mark.parametrize("caller_kind", ["team", "user"]) async def test_caller_mcp_revocation_uses_fresh_policy( - monkeypatch: pytest.MonkeyPatch, fresh: bool, caller_kind: str, + monkeypatch: pytest.MonkeyPatch, + fresh: bool, + caller_kind: str, ) -> None: from litellm.proxy._types import LiteLLM_TeamTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key from litellm.types.agents import AgentCaller cached_permission: Final = LiteLLM_ObjectPermissionTable( - object_permission_id="caller-permission", mcp_servers=["slack", "linear"], + object_permission_id="caller-permission", + mcp_servers=["slack", "linear"], mcp_tool_permissions={"slack": ["read", "write"]}, ) current_permission: Final = LiteLLM_ObjectPermissionTable( - object_permission_id="caller-permission", mcp_servers=["slack"], + object_permission_id="caller-permission", + mcp_servers=["slack"], mcp_tool_permissions={"slack": ["read"]}, ) team: Final = LiteLLM_TeamTable( - team_id="caller", object_permission_id="caller-permission", object_permission=current_permission, + team_id="caller", + object_permission_id="caller-permission", + object_permission=current_permission, ) user: Final = LiteLLM_UserTable( - user_id="caller", teams=[], object_permission_id="caller-permission", object_permission=current_permission, + user_id="caller", + teams=[], + object_permission_id="caller-permission", + object_permission=current_permission, ) database: Final = MagicMock() database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team) @@ -531,13 +578,16 @@ async def test_caller_mcp_revocation_uses_fresh_policy( auth.agent_caller = AgentCaller(team_id="caller") if caller_kind == "team" else AgentCaller(user_id="caller") assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == ({"slack"} if fresh else {"slack", "linear"}) - assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == (["read"] if fresh else ["read", "write"]) + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ( + ["read"] if fresh else ["read", "write"] + ) @pytest.mark.asyncio @pytest.mark.parametrize("fresh", [False, True]) async def test_caller_team_outage_cannot_remove_authoritative_server_ceiling( - monkeypatch: pytest.MonkeyPatch, fresh: bool, + monkeypatch: pytest.MonkeyPatch, + fresh: bool, ) -> None: from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.types.agents import AgentCaller diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 078fff0266b..4e2597508fc 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1,6 +1,8 @@ import contextlib import json import os +from collections.abc import Sequence +from collections.abc import Set as AbstractSet from datetime import datetime, timedelta, timezone from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -10,6 +12,8 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers +from litellm.constants import MCP_ALL_TOOLS_WILDCARD +from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, @@ -18,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( ) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, ProxyException, SpecialHeaders, SpecialMCPServerNames, @@ -26,17 +31,21 @@ from litellm.proxy._types import ( from litellm.types.agents import AgentCaller -@pytest.mark.parametrize("headers,expected", [ - ({"Authorization": "Bearer caller.jwt.signature"}, "caller.jwt.signature"), - ({"AUTHORIZATION": "bEaReR caller.jwt.signature"}, "caller.jwt.signature"), - ({"authorization": "Bearer "}, ""), - ({"authorization": "Basic credentials"}, None), - ({"authorization-extra": "Bearer wrong.jwt.signature"}, None), - ({"x-litellm-api-key": "sk-key"}, None), - ({}, None), -]) +@pytest.mark.parametrize( + "headers,expected", + [ + ({"Authorization": "Bearer caller.jwt.signature"}, "caller.jwt.signature"), + ({"AUTHORIZATION": "bEaReR caller.jwt.signature"}, "caller.jwt.signature"), + ({"authorization": "Bearer "}, ""), + ({"authorization": "Basic credentials"}, None), + ({"authorization-extra": "Bearer wrong.jwt.signature"}, None), + ({"x-litellm-api-key": "sk-key"}, None), + ({}, None), + ], +) def test_incoming_guardrail_bearer_requires_exact_authorization_header( - headers: dict[str, str], expected: str | None, + headers: dict[str, str], + expected: str | None, ) -> None: token: Final = MCPRequestHandler.get_incoming_bearer_token(headers) assert (token.get_secret_value() if token is not None else None) == expected @@ -8372,7 +8381,9 @@ class TestUserSubjectTeamUnion: assert servers == ["srv1"] assert tools == ["add"] assert resolve.await_args_list - assert all(call.kwargs == {"toolset_ids": ["ts-1"], "requires_fresh_policy": True} for call in resolve.await_args_list) + assert all( + call.kwargs == {"toolset_ids": ["ts-1"], "requires_fresh_policy": True} for call in resolve.await_args_list + ) async def test_key_based_caller_uses_single_team_only(self): """A key-based caller (api_key set) with a team_id sees ONLY that team, even though the @@ -9564,9 +9575,7 @@ class TestSessionBearerEgressScrub: ) -> None: caller_key: Final = "sk-caller-admission-key-123" upstream_token: Final = "Bearer real-upstream-token" - oauth2_headers: Final = { - "Authorization": caller_value if context == "oauth2_authorization" else upstream_token - } + oauth2_headers: Final = {"Authorization": caller_value if context == "oauth2_authorization" else upstream_token} raw_headers: Final = { "x-litellm-api-key": f"Bearer {caller_key}", "authorization": caller_value if context == "raw_authorization" else upstream_token, @@ -9823,7 +9832,9 @@ class TestSessionBearerEgressScrub: {}, ) - async def test_caller_admission_credential_falls_back_to_standard_headers_when_custom_header_is_absent(self) -> None: + async def test_caller_admission_credential_falls_back_to_standard_headers_when_custom_header_is_absent( + self, + ) -> None: caller_value: Final = "Bearer sk-caller-admission-key-123" headers: Final = Headers( { @@ -10243,7 +10254,11 @@ class TestGetUserObjectPermission: prisma_client = MagicMock() from litellm.proxy._types import LiteLLM_UserTable - row = LiteLLM_UserTable(user_id="human", object_permission_id=user_row.object_permission_id) if user_row is not None else None + row = ( + LiteLLM_UserTable(user_id="human", object_permission_id=user_row.object_permission_id) + if user_row is not None + else None + ) prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row) return prisma_client @@ -10259,7 +10274,10 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())), + ), patch( "litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock, @@ -10286,7 +10304,10 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())), + ), patch("litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock) as mock_get_perm, ): assert await MCPRequestHandler._get_user_object_permission(auth) is None @@ -10308,7 +10329,10 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())), + ), ): assert await MCPRequestHandler._get_user_object_permission(auth) is None @@ -10322,7 +10346,10 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())), + ), ): assert await MCPRequestHandler._get_user_object_permission(auth) is None @@ -10339,7 +10366,10 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())), + ), patch( "litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock, @@ -10759,7 +10789,9 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc from litellm.proxy import proxy_server from litellm.proxy._types import LiteLLM_UserTable - current = LiteLLM_UserTable(user_id="catalog-user", organization_id="current-org", user_role="internal_user", teams=[]) + current = LiteLLM_UserTable( + user_id="catalog-user", organization_id="current-org", user_role="internal_user", teams=[] + ) table = SimpleNamespace(find_unique=AsyncMock(return_value=current)) database = SimpleNamespace(writer_db=SimpleNamespace(litellm_usertable=table)) monkeypatch.setattr(proxy_server, "prisma_client", database) @@ -10779,10 +10811,25 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc @pytest.mark.parametrize("current_groups", [[], ["replacement-group"]]) async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session_scope(monkeypatch, current_groups): permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_servers=["current-server"]) - current = UserAPIKeyAuth(object_permission=permission, object_permission_id="current-policy", team_id="new-team", org_id="new-org", project_id="new-project", user_id="new-owner", access_group_ids=current_groups) + current = UserAPIKeyAuth( + object_permission=permission, + object_permission_id="current-policy", + team_id="new-team", + org_id="new-org", + project_id="new-project", + user_id="new-owner", + access_group_ids=current_groups, + ) reload_key = AsyncMock(return_value=current) monkeypatch.setattr(MCPRequestHandler, "_reload_admitted_key", reload_key) - caller = UserAPIKeyAuth(api_key="owned-key-hash", team_id="old-team", org_id="old-org", project_id="old-project", user_id="old-owner", access_group_ids=["original-group"]) + caller = UserAPIKeyAuth( + api_key="owned-key-hash", + team_id="old-team", + org_id="old-org", + project_id="old-project", + user_id="old-owner", + access_group_ids=["original-group"], + ) caller.via_virtual_key = True caller.mcp_session_resource_server_id = "session-server" caller.mcp_toolset_id = "session-toolset" @@ -10790,7 +10837,12 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True) assert refreshed.object_permission == permission assert refreshed.object_permission_id == "current-policy" - assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ("new-team", "new-org", "new-project", "new-owner") + assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ( + "new-team", + "new-org", + "new-project", + "new-owner", + ) assert refreshed.mcp_session_resource_server_id == "session-server" assert refreshed.mcp_toolset_id == "session-toolset" assert refreshed.via_virtual_key and refreshed.requires_fresh_policy @@ -10820,33 +10872,40 @@ async def test_admission_request_body_serves_stashed_peek_callable(): ) assert await with_peek.body() == admission_body - without_peek: Final = _admission_request( - {"type": "http", "method": "POST", "path": "/mcp", "headers": []} - ) + without_peek: Final = _admission_request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) assert await without_peek.body() == b"{}" -@pytest.mark.parametrize("bearer,api_key,claims,explicit_key,expected", [ - (None, "sk-key", None, None, None), - ("sk-key", "sk-key", None, None, None), - ("sk-key", "sk-key", {"sub": "previous-jwt"}, "sk-key", None), - ("caller.jwt.assertion", "sk-key", {"sub": "caller"}, None, "caller.jwt.assertion"), - ("caller.jwt.assertion", None, {"sub": "caller"}, None, "caller.jwt.assertion"), - ("upstream", "sk-key", None, "sk-key", "upstream"), - ("llm_session_test", None, {"sub": "caller"}, None, None), - ("llm_srefresh_test", None, {"sub": "caller"}, None, None), - ("opaque-master", "litellm_proxy_master_key", None, None, None), -]) +@pytest.mark.parametrize( + "bearer,api_key,claims,explicit_key,expected", + [ + (None, "sk-key", None, None, None), + ("sk-key", "sk-key", None, None, None), + ("sk-key", "sk-key", {"sub": "previous-jwt"}, "sk-key", None), + ("caller.jwt.assertion", "sk-key", {"sub": "caller"}, None, "caller.jwt.assertion"), + ("caller.jwt.assertion", None, {"sub": "caller"}, None, "caller.jwt.assertion"), + ("upstream", "sk-key", None, "sk-key", "upstream"), + ("llm_session_test", None, {"sub": "caller"}, None, None), + ("llm_srefresh_test", None, {"sub": "caller"}, None, None), + ("opaque-master", "litellm_proxy_master_key", None, None, None), + ], +) def test_guardrail_bearer_preserves_identity_without_exposing_gateway_credentials( - bearer: str | None, api_key: str | None, claims: dict[str, str] | None, - explicit_key: str | None, expected: str | None, + bearer: str | None, + api_key: str | None, + claims: dict[str, str] | None, + explicit_key: str | None, + expected: str | None, ) -> None: - headers: Final = Headers({ - **({"Authorization": f"Bearer {bearer}"} if bearer is not None else {}), - **({"x-litellm-api-key": explicit_key} if explicit_key else {}), - }) + headers: Final = Headers( + { + **({"Authorization": f"Bearer {bearer}"} if bearer is not None else {}), + **({"x-litellm-api-key": explicit_key} if explicit_key else {}), + } + ) token: Final = MCPRequestHandler.get_guardrail_bearer_token( - headers, UserAPIKeyAuth(api_key=api_key, jwt_claims=claims), + headers, + UserAPIKeyAuth(api_key=api_key, jwt_claims=claims), ) assert (token.get_secret_value() if token is not None else None) == expected @@ -10860,3 +10919,545 @@ def test_guardrail_bearer_respects_custom_admission_header(monkeypatch: pytest.M UserAPIKeyAuth(api_key="litellm_proxy_master_key"), ) assert token is None + + +def _access_group(group_id, servers, tool_permissions=None, **owners): + from litellm.models.access_group import LiteLLM_AccessGroupTable + + return LiteLLM_AccessGroupTable( + access_group_id=group_id, + access_group_name=group_id, + access_mcp_server_ids=list(servers), + mcp_tool_permissions=tool_permissions, + **owners, + ) + + +class TestAccessGroupToolsForServer: + """Tool grants of unified access groups on one server, before they meet the key/team ceilings.""" + + def test_no_group_grants_the_server_is_no_restriction(self): + groups = [_access_group("g", ["other"], {"other": ["x"], "srv": ["leak"]})] + assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") is None + + def test_granting_group_without_entry_allows_every_tool(self): + groups = [_access_group("g", ["srv"])] + assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") is None + + def test_granting_group_with_explicit_empty_list_grants_no_tools(self): + groups = [_access_group("g", ["srv"], {"srv": []})] + assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") == () + + def test_tools_union_across_granting_groups(self): + groups = [_access_group("a", ["srv"], {"srv": ["read"]}), _access_group("b", ["srv"], {"srv": ["write"]})] + assert sorted(MCPRequestHandler.access_group_tools_for_server(groups, "srv")) == ["read", "write"] + + def test_unrestricted_granting_group_wins_the_union(self): + groups = [_access_group("a", ["srv"], {"srv": ["read"]}), _access_group("b", ["srv"])] + assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") is None + + def test_non_granting_group_entry_does_not_widen(self): + groups = [_access_group("a", ["srv"], {"srv": []}), _access_group("b", ["other"], {"srv": ["write"]})] + assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") == () + + +@pytest.mark.asyncio +class TestAccessGroupToolPermissionsInheritance: + """Group tool allowlists reach the key, team and agent levels of get_allowed_tools_for_server.""" + + @staticmethod + def _team(access_group_ids): + from litellm.proxy._types import LiteLLM_TeamTable + + return LiteLLM_TeamTable(team_id="team-1", access_group_ids=access_group_ids) + + @staticmethod + def _groups_loader(groups): + by_id = {group.access_group_id: group for group in groups} + + async def load(*, access_group_id, **_): + return by_id[access_group_id] + + return AsyncMock(side_effect=load) + + async def test_key_access_group_tool_allowlist_restricts_the_server(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": ["read"]}, assigned_key_ids=[auth.token])] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_key_access_group_explicit_empty_list_denies_every_tool(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": []}, assigned_key_ids=[auth.token])] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_key_server_only_access_group_stays_unrestricted(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], assigned_key_ids=[auth.token])] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None + + async def test_key_object_permission_tools_intersect_with_group_tools(self): + auth = _key_auth_reaching("srv", tools=["read", "delete"], access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": ["read", "write"]}, assigned_key_ids=[auth.token])] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + allowed = await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) + assert sorted(allowed) == ["delete", "read", "write"] + + async def test_team_access_group_tool_allowlist_caps_the_key(self): + auth = _key_auth_reaching("srv", team_id="team-1") + team = self._team(["g"]) + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_team_access_group_explicit_empty_list_denies_even_with_key_tools(self): + auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1") + team = self._team(["g"]) + groups = [_access_group("g", ["srv"], {"srv": []})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_team_and_key_group_tools_intersect(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", team_id="team-1", access_group_ids=["kg"]) + team = self._team(["tg"]) + groups = [ + _access_group("kg", ["srv"], {"srv": ["read", "write"]}, assigned_team_ids=["team-1"]), + _access_group("tg", ["srv"], {"srv": ["write", "delete"]}), + ] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["write"] + + async def test_agent_access_group_tool_allowlist_caps_an_agent_key(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + agent = MagicMock(access_group_ids=["g"]) + with ( + _entitlement_fault_globals(_agent_prisma()), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + +class TestAccessGroupToolOwnershipAndFaults: + """Group tool grants bind only through groups that own the caller, fail closed when a named group + cannot be read, and cap rather than widen an agent's own grants. Every case runs the real + ``get_allowed_tools_for_server`` resolver with ``get_access_object`` as the only patched boundary.""" + + @staticmethod + def _loader( + groups: Sequence[LiteLLM_AccessGroupTable], + unreadable: AbstractSet[str] = frozenset(), + missing: AbstractSet[str] = frozenset(), + ) -> AsyncMock: + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + by_id = {group.access_group_id: group for group in groups} + + async def load(*, access_group_id: str, **_: object) -> LiteLLM_AccessGroupTable: + if access_group_id in unreadable: + raise HTTPException(status_code=404, detail={"error": "Access group lookup failed. Error: db down"}) + if access_group_id in missing: + raise AccessGroupNotFoundError(access_group_id=access_group_id) + return by_id[access_group_id] + + return AsyncMock(side_effect=load) + + async def test_unowned_key_group_contributes_no_tools(self): + auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": ["write"]})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_key_group_owned_through_the_team_binds(self): + auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1", access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": ["write"]}, assigned_team_ids=["team-1"])] + team = LiteLLM_TeamTable(team_id="team-1") + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert sorted(await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)) == ["read", "write"] + + async def test_unreadable_key_group_denies_the_server_tools(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"]) + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_unreadable_key_group_denies_even_when_the_team_grants_the_server(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", team_id="team-1", access_group_ids=["g"]) + team = LiteLLM_TeamTable( + team_id="team-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-team", mcp_servers=["srv"], mcp_tool_permissions={"srv": ["read", "write"]} + ), + ) + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_unreadable_key_group_keeps_the_direct_allowlist(self): + auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g"]) + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_unreadable_team_group_denies_the_server_tools(self): + auth = _key_auth_reaching("srv", team_id="team-1") + team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg"]) + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"tg"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_nonexistent_key_group_leaves_the_team_grant_alone(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", team_id="team-1", access_group_ids=["typo-group"]) + team = LiteLLM_TeamTable( + team_id="team-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-team", mcp_servers=["srv"], mcp_tool_permissions={"srv": ["read", "write"]} + ), + ) + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], missing={"typo-group"})), + ): + assert sorted(await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)) == ["read", "write"] + + @pytest.mark.parametrize("requires_fresh_policy", [False, True]) + async def test_nonexistent_key_group_resolves_like_no_group_at_all(self, requires_fresh_policy): + without_group = _key_auth_reaching("srv", requires_fresh_policy=requires_fresh_policy) + with_typo = _key_auth_reaching( + "srv", requires_fresh_policy=requires_fresh_policy, access_group_ids=["typo-group"] + ) + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], missing={"typo-group"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server( + "srv", with_typo + ) == await MCPRequestHandler.get_allowed_tools_for_server("srv", without_group) + + async def test_nonexistent_key_group_does_not_lift_a_real_attached_bound(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["typo-group", "g"]) + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})), + patch( + "litellm.proxy.auth.auth_checks.get_mcp_server_ids_from_access_groups", AsyncMock(return_value=["srv"]) + ), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_nonexistent_team_group_resolves_like_a_team_without_groups(self): + auth = _key_auth_reaching("srv", team_id="team-1") + with_typo = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["typo-group"]) + without_group = LiteLLM_TeamTable(team_id="team-1") + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], missing={"typo-group"})), + ): + with patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=with_typo)): + typo_tools = await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) + with patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=without_group)): + assert typo_tools == await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) + + @pytest.mark.parametrize("group_tools, expected", [(["read"], ["read"]), ([], [])]) + async def test_unowned_group_only_key_is_bounded_by_the_group_it_is_admitted_through(self, group_tools, expected): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": group_tools})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + patch( + "litellm.proxy.auth.auth_checks.get_mcp_server_ids_from_access_groups", AsyncMock(return_value=["srv"]) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == ["srv"] + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == expected + + @pytest.mark.parametrize("key_owns_group", [True, False], ids=["owned-group", "attached-only-group"]) + async def test_explicit_key_wildcard_is_not_narrowed_by_a_group_bound(self, key_owns_group): + auth = _key_auth_reaching("srv", tools=[MCP_ALL_TOOLS_WILDCARD], access_group_ids=["g"]) + owners = {"assigned_key_ids": [auth.token]} if key_owns_group else {} + groups = [_access_group("g", ["srv"], {"srv": ["read"]}, **owners)] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None + + async def test_owned_group_without_tool_entry_keeps_the_key_unrestricted(self): + auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g", "h"]) + groups = [ + _access_group("g", ["srv"], assigned_key_ids=[auth.token]), + _access_group("h", ["srv"], {"srv": ["read"]}), + ] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None + + async def test_direct_server_grant_without_tool_entry_is_not_narrowed_by_an_attached_group(self): + auth = _key_auth_reaching("srv", access_group_ids=["g"]) + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None + + async def test_legacy_tag_colliding_with_the_server_name_is_not_a_direct_grant(self): + auth = UserAPIKeyAuth( + api_key="sk-hash", + user_id="u1", + access_group_ids=["g"], + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_access_groups=["srv"]), + ) + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_legacy_tag_carried_by_the_server_keeps_the_key_unrestricted(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + auth = UserAPIKeyAuth( + api_key="sk-hash", + user_id="u1", + access_group_ids=["g"], + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_access_groups=["tag"]), + ) + srv = MCPServer( + server_id="srv", + name="srv", + url="https://srv.example.com", + transport=MCPTransport.http, + access_groups=["tag"], + ) + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch.object(global_mcp_server_manager, "config_mcp_servers", {"srv": srv}), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None + + async def test_duplicate_key_group_ids_do_not_deny(self): + auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g", "g"]) + groups = [_access_group("g", ["srv"], {"srv": ["write"]}, assigned_key_ids=[auth.token])] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert sorted(await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)) == ["read", "write"] + + async def test_duplicate_team_group_ids_do_not_deny(self): + auth = _key_auth_reaching("srv", team_id="team-1") + team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg", "tg"]) + groups = [_access_group("tg", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_duplicate_agent_group_ids_do_not_deny(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + agent = MagicMock(access_group_ids=["g", "g"]) + with ( + _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read", "write"])), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + @staticmethod + def _writer_down(*_: object, **__: object) -> None: + raise RuntimeError("writer down") + + async def test_fresh_key_group_fault_denies_the_server_tools(self): + auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g"]).model_copy( + update={"requires_fresh_policy": True} + ) + assert auth.requires_fresh_policy is True + with ( + _entitlement_fault_globals(), + patch.object(MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=None)), + patch("litellm.proxy.auth.auth_checks.get_access_object", AsyncMock(side_effect=self._writer_down)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_fresh_team_group_fault_denies_the_server_tools(self): + auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1").model_copy( + update={"requires_fresh_policy": True} + ) + assert auth.requires_fresh_policy is True + team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg"]) + with ( + _entitlement_fault_globals(), + patch.object(MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=None)), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", AsyncMock(side_effect=self._writer_down)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_team_lookup_fault_in_group_resolution_denies_the_server_tools(self): + auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1") + with ( + _entitlement_fault_globals(), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(side_effect=self._writer_down)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_agent_group_fault_denies_rather_than_shedding_the_agents_own_allowlist(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + agent = MagicMock(access_group_ids=["g"]) + with ( + _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read"])), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", AsyncMock(side_effect=self._writer_down)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_echoed_caller_team_group_tools_cap_the_agent_key(self): + from litellm.types.agents import AgentCaller + + auth = _key_auth_reaching("srv", tools=["read", "write"], agent_id="agent-1") + auth.agent_caller = AgentCaller(team_id="team-1") + team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg"]) + groups = [_access_group("tg", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(), + patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)), + patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + capped = await MCPRequestHandler.apply_agent_caller_tool_ceiling(["read", "write"], "srv", auth) + assert sorted(capped) == ["read"] + + @staticmethod + def _agent_prisma_with_own_tools(tools: Sequence[str]) -> MagicMock: + prisma_client = _agent_prisma(object_permission_id="op-agent") + prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=LiteLLM_ObjectPermissionTable( + object_permission_id="op-agent", mcp_servers=["srv"], mcp_tool_permissions={"srv": tools} + ) + ) + return prisma_client + + async def test_agent_group_ceiling_caps_the_agents_own_tools(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + groups = [_access_group("g", ["srv"], {"srv": ["read"]})] + agent = MagicMock(access_group_ids=["g"]) + with ( + _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read", "write"])), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + async def test_agent_group_tools_outside_the_agents_own_grant_are_denied(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + groups = [_access_group("g", ["srv"], {"srv": ["write"]})] + agent = MagicMock(access_group_ids=["g"]) + with ( + _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read"])), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_unreadable_agent_group_denies_the_server_tools(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + agent = MagicMock(access_group_ids=["g"]) + with ( + _entitlement_fault_globals(_agent_prisma()), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == [] + + async def test_nonexistent_agent_group_leaves_the_real_groups_tools_alone(self): + auth = _key_auth_reaching("srv", agent_id="agent-1") + agent = MagicMock(access_group_ids=["real", "typo-group"]) + groups = [_access_group("real", ["srv"], {"srv": ["read"]})] + with ( + _entitlement_fault_globals(_agent_prisma()), + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})), + ): + assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 0d6b83b3275..fd94fccf8fd 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -37,6 +37,7 @@ from litellm.proxy._experimental.mcp_server.tool_search import ( coerce_top_k, get_virtual_tool_definitions, handle_mcp_tool_search, + rank_mcp_tools, search_mcp_tools, search_tools, ) @@ -193,6 +194,19 @@ class RecordingEmbedder: return tuple(FAKE_VECTORS[text] for text in texts) +class OversizedRecordingEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + if any("OVERSIZED" in text for text in texts): + raise litellm.ContextWindowExceededError( + message="input is too long", model="text-embedding-3-large", llm_provider="openai" + ) + return tuple((1.0, 0.0) if "weather" in text.lower() else (0.0, 1.0) for text in texts) + + def _ranker(embedder: RecordingEmbedder | None = None) -> SemanticToolRanker: return SemanticToolRanker(embed=embedder or RecordingEmbedder(), embedding_model="emb", index=SemanticTextIndex()) @@ -203,6 +217,26 @@ def _names(results: Sequence[ToolSearchResult] | EmbeddingFailed) -> list[str]: class TestSearchMcpTools: + @pytest.mark.asyncio + async def test_rank_mcp_tools_omits_oversized_description(self) -> None: + tools: Final = _make_tools( + [ + ("get_weather", "Get the current weather forecast for a city"), + ("send_email", "Send an email to a recipient"), + ("list_files", "List files in a directory"), + ("huge_tool", "OVERSIZED description"), + ] + ) + embedder = OversizedRecordingEmbedder() + ranker = SemanticToolRanker(embed=embedder, embedding_model="emb", index=SemanticTextIndex()) + + hits: Final = await rank_mcp_tools("weather", tools, 5, MCPToolSearchSettings(embedding_model="emb"), ranker) + + assert not isinstance(hits, EmbeddingFailed) + names: Final = [hit.tool.name for hit in hits] + assert names[0] == "get_weather" + assert "huge_tool" not in names + @pytest.mark.asyncio async def test_semantic_mode_finds_foreign_exchange_tool_for_fx(self) -> None: keyword_only = await search_mcp_tools("FX", CATALOG, 5, MCPToolSearchSettings(), ranker=None) diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py index 08cceb0d967..cb9a63512d4 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -61,13 +61,11 @@ async def test_unknown_agent_has_no_ceiling(): @pytest.mark.asyncio async def test_ceiling_is_the_union_of_every_attached_group(): - load_agent, load_group = _loaders( - _agent(["g1", "g2"]), - { - "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)), - "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)), - }, - ) + groups: Final = { + "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)), + "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)), + } + load_agent, load_group = _loaders(_agent(["g1", "g2"]), groups) ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) @@ -76,12 +74,14 @@ async def test_ceiling_is_the_union_of_every_attached_group(): models=frozenset({"gpt-5", "claude-sonnet"}), mcp_server_ids=frozenset({"mcp-a", "mcp-b"}), agent_ids=frozenset({"agent-b", "agent-c"}), + access_groups=(groups["g1"], groups["g2"]), ) @pytest.mark.asyncio async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applies(): - load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": _group("g1", models=("gpt-5",))}) + g1: Final = _group("g1", models=("gpt-5",)) + load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": g1}) ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) @@ -90,6 +90,7 @@ async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applie models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset(), + access_groups=(g1,), ) @@ -109,7 +110,8 @@ async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): async def test_default_agent_loader_reads_the_attached_groups_from_the_registry(): from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - _, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))}) + g1: Final = _group("g1", models=("gpt-5",)) + _, load_group = _loaders(None, {"g1": g1}) global_agent_registry.register_agent(_agent(["g1"])) try: ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group) @@ -117,7 +119,11 @@ async def test_default_agent_loader_reads_the_attached_groups_from_the_registry( global_agent_registry.deregister_agent("agent") assert ceiling == AgentAccessGroupCeiling( - access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset() + access_group_ids=("g1",), + models=frozenset({"gpt-5"}), + mcp_server_ids=frozenset(), + agent_ids=frozenset(), + access_groups=(g1,), ) @@ -136,6 +142,61 @@ async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: assert await _load_access_group("gone") is None +@pytest.mark.asyncio +async def test_default_loader_lets_a_provably_absent_group_surface(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + async def absent_group(*, access_group_id: str, **_: object) -> LiteLLM_AccessGroupTable: + raise AccessGroupNotFoundError(access_group_id=access_group_id) + + monkeypatch.setattr(proxy_server, "prisma_client", object()) + monkeypatch.setattr(auth_checks, "get_access_object", absent_group) + + with pytest.raises(AccessGroupNotFoundError): + await _load_access_group("gone") + + +@pytest.mark.asyncio +async def test_ceiling_drops_a_group_that_does_not_exist_but_keeps_an_unreadable_one(): + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + real = LiteLLM_AccessGroupTable(access_group_id="real", access_group_name="real", access_mcp_server_ids=["srv"]) + + async def ids(_agent_id: str) -> tuple[str, ...]: + return ("real", "typo-group", "unreadable") + + async def load(group_id: str) -> LiteLLM_AccessGroupTable | None: + if group_id == "typo-group": + raise AccessGroupNotFoundError(access_group_id=group_id) + return real if group_id == "real" else None + + ceiling = await resolve_agent_access_group_ceiling("agent-1", ids, load) + assert ceiling is not None + assert ceiling.access_group_ids == ("real", "unreadable") + assert ceiling.access_groups == (real,) + assert ceiling.mcp_server_ids == frozenset({"srv"}) + + +@pytest.mark.asyncio +async def test_ceiling_stays_and_caps_to_nothing_when_every_named_group_does_not_exist(): + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + async def ids(_agent_id: str) -> tuple[str, ...]: + return ("typo-group",) + + async def load(group_id: str) -> LiteLLM_AccessGroupTable | None: + raise AccessGroupNotFoundError(access_group_id=group_id) + + assert await resolve_agent_access_group_ceiling("agent-1", ids, load) == AgentAccessGroupCeiling( + access_group_ids=(), models=frozenset(), mcp_server_ids=frozenset(), agent_ids=frozenset(), access_groups=() + ) + + @pytest.mark.asyncio async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server @@ -159,7 +220,9 @@ async def test_authoritative_group_ceiling_propagates_policy_outages( database: Final = MagicMock() database.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) - database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock( + side_effect=RuntimeError("database unavailable") + ) monkeypatch.setattr(proxy_server, "prisma_client", database) monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) if strict: diff --git a/tests/unit/proxy/agent_endpoints/test_agent_search.py b/tests/unit/proxy/agent_endpoints/test_agent_search.py index c02ed1f37e5..750fadadab0 100644 --- a/tests/unit/proxy/agent_endpoints/test_agent_search.py +++ b/tests/unit/proxy/agent_endpoints/test_agent_search.py @@ -88,6 +88,19 @@ class FixedDimensionEmbedder: return tuple((1.0,) * self.dimensions for _ in texts) +class OversizedAwareEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + if any("OVERSIZED" in text for text in texts): + raise litellm.ContextWindowExceededError( + message="input is too long", model="text-embedding-3-large", llm_provider="openai" + ) + return tuple(VECTORS[text] for text in texts) + + class TestAgentSearchText: def test_joins_name_description_and_skills_with_tags(self) -> None: assert agent_search_text(TRANSLATOR) == ( @@ -127,6 +140,21 @@ class TestAgentSearchIndex: assert [hit.agent.agent_id for hit in outcome.hits] == ["translator", "trip"] assert outcome.hits[0].score > outcome.hits[1].score + @pytest.mark.asyncio + async def test_omits_agents_with_oversized_search_text(self) -> None: + oversized: Final = AgentResponse( + agent_id="oversized", + agent_name="oversized-agent", + agent_card_params={"description": "OVERSIZED description"}, + ) + + outcome: Final = await AgentSearchIndex().search( + "language translation", (*AGENTS, oversized), top_k=5, embed=OversizedAwareEmbedder(), embedding_model="m" + ) + + assert isinstance(outcome, AgentSearchHits) + assert {hit.agent.agent_id for hit in outcome.hits} == {agent.agent_id for agent in AGENTS} + @pytest.mark.asyncio async def test_second_search_only_embeds_the_query(self) -> None: index = AgentSearchIndex() diff --git a/tests/unit/proxy/anthropic_endpoints/test_claude_code_skill_access.py b/tests/unit/proxy/anthropic_endpoints/test_claude_code_skill_access.py index 3e21ac9c796..a12a90b5a78 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_claude_code_skill_access.py +++ b/tests/unit/proxy/anthropic_endpoints/test_claude_code_skill_access.py @@ -1,12 +1,16 @@ """Unit tests for claude_code_skill_access.py: key/team grant resolution.""" -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( + PUBLIC_ONLY, + AccessGroupSkills, + access_group_skills, granted_skills, + resolve_skill_visibility, skill_visibility, ) @@ -77,3 +81,331 @@ def test_where_clause_bounds_the_plugin_query_to_what_the_caller_may_see(): assert skill_visibility(_key(["b", "a"], None)).where() == {"OR": [{"enabled": True}, {"name": {"in": ["a", "b"]}}]} assert skill_visibility(_key(["a", "b"], ["b"])).where() == {"OR": [{"enabled": True}, {"name": {"in": ["b"]}}]} assert skill_visibility(admin).where() == {} + + +@pytest.mark.parametrize( + ("key_skills", "team_skills", "inherited", "expected"), + [ + (None, None, AccessGroupSkills(key=frozenset({"g"})), frozenset({"g"})), + (None, None, AccessGroupSkills(team=frozenset({"g"})), frozenset({"g"})), + (None, None, AccessGroupSkills(), frozenset()), + ([], [], AccessGroupSkills(key=frozenset({"g1"}), team=frozenset({"g2"})), frozenset()), + (["a"], None, AccessGroupSkills(key=frozenset({"g"})), frozenset({"a", "g"})), + (None, ["a"], AccessGroupSkills(team=frozenset({"g"})), frozenset({"a", "g"})), + (None, ["a"], AccessGroupSkills(key=frozenset({"g"})), frozenset()), + (["a"], ["a", "b"], AccessGroupSkills(key=frozenset({"b"})), frozenset({"a", "b"})), + (["a"], ["a"], AccessGroupSkills(key=frozenset({"g"}), team=frozenset({"g"})), frozenset({"a", "g"})), + ], +) +def test_granted_skills_adds_access_group_skills_to_their_own_scope_before_intersecting( + key_skills, team_skills, inherited, expected +): + assert granted_skills(_key(key_skills, team_skills), inherited) == expected + + +def test_visibility_with_inherited_skills_only_opens_the_inherited_private_plugin(): + inherited = AccessGroupSkills(key=frozenset({"group-skill"})) + visibility = skill_visibility(_key(None, None), inherited) + + assert visibility.allows(_plugin("group-skill", enabled=False)) + assert not visibility.allows(_plugin("other-private", enabled=False)) + assert visibility.allows(_plugin("public", enabled=True)) + assert visibility.where() == {"OR": [{"enabled": True}, {"name": {"in": ["group-skill"]}}]} + + +@pytest.mark.asyncio +async def test_resolve_skill_visibility_folds_in_access_group_skills_for_non_admins(): + loader = AsyncMock(return_value=AccessGroupSkills(team=frozenset({"group-skill"}))) + caller = _key(None, None) + + visibility = await resolve_skill_visibility(caller, loader) + + loader.assert_awaited_once_with(caller) + assert visibility.allows(_plugin("group-skill", enabled=False)) + assert not visibility.allows(_plugin("other-private", enabled=False)) + + +@pytest.mark.asyncio +async def test_resolve_skill_visibility_skips_group_lookup_for_anonymous_and_admin_callers(): + loader = AsyncMock(return_value=AccessGroupSkills(key=frozenset({"group-skill"}))) + admin = UserAPIKeyAuth(api_key="sk-9876", user_role=LitellmUserRoles.PROXY_ADMIN) + + assert not (await resolve_skill_visibility(None, loader)).allows(_plugin("group-skill", enabled=False)) + assert (await resolve_skill_visibility(admin, loader)).sees_private + loader.assert_not_awaited() + + +def _team(team_id: str, access_group_ids: list[str], blocked: bool = False) -> MagicMock: + team = MagicMock() + team.team_id = team_id + team.access_group_ids = access_group_ids + team.blocked = blocked + return team + + +def _group(access_group_id: str, skills: list[str], assigned_key_ids: list[str] | None = None) -> MagicMock: + group = MagicMock() + group.access_group_id = access_group_id + group.access_skill_names = skills + group.assigned_team_ids = [] + group.assigned_key_ids = assigned_key_ids or [] + return group + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_blocked", "expected_team_skills"), + [(False, frozenset({"team-group-skill"})), (True, frozenset())], +) +async def test_access_group_skills_resolves_key_groups_and_team_groups(monkeypatch, team_blocked, expected_team_skills): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + caller = UserAPIKeyAuth( + api_key="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-1", + access_group_ids=["key-group", "stray-group"], + ).model_copy(update={"requires_fresh_policy": True}) + groups = { + "key-group": _group("key-group", ["key-group-skill"], assigned_key_ids=[caller.token or ""]), + "stray-group": _group("stray-group", ["stray-skill"]), + "team-group": _group("team-group", ["team-group-skill"]), + } + + stale_groups = {group_id: _group(group_id, ["stale-skill"]) for group_id in groups} + + async def fake_access_object(access_group_id, check_db_only=False, **_kwargs): + return (groups if check_db_only else stale_groups)[access_group_id] + + async def fake_skill_names(access_group_ids, **_kwargs): + return frozenset().union(*(groups[group_id].access_skill_names for group_id in access_group_ids)) + + monkeypatch.setattr(auth_checks, "get_access_object", fake_access_object) + monkeypatch.setattr(auth_checks, "get_skill_names_from_access_groups", fake_skill_names) + monkeypatch.setattr( + auth_checks, "get_team_object", AsyncMock(return_value=_team("team-1", ["team-group"], blocked=team_blocked)) + ) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + inherited = await access_group_skills(caller) + + assert inherited == AccessGroupSkills(key=frozenset({"key-group-skill"}), team=expected_team_skills) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fresh", [True, False]) +async def test_access_group_skills_surfaces_key_group_lookup_failures(monkeypatch, fresh): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + caller = UserAPIKeyAuth( + api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, access_group_ids=["key-group"] + ).model_copy(update={"requires_fresh_policy": fresh}) + monkeypatch.setattr(auth_checks, "get_access_object", AsyncMock(side_effect=RuntimeError("writer down"))) + monkeypatch.setattr(auth_checks, "get_skill_names_from_access_groups", AsyncMock(return_value=frozenset())) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + + with pytest.raises(RuntimeError, match="writer down"): + await access_group_skills(caller) + + +@pytest.mark.asyncio +async def test_access_group_skills_without_groups_or_team_grants_nothing(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + get_team_object = AsyncMock() + monkeypatch.setattr(auth_checks, "get_team_object", get_team_object) + monkeypatch.setattr(auth_checks, "get_skill_names_from_access_groups", AsyncMock(return_value=frozenset())) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + caller = UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER) + + assert await access_group_skills(caller) == AccessGroupSkills() + get_team_object.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fresh", [True, False]) +async def test_access_group_skills_keeps_team_ceiling_when_team_lookup_fails(monkeypatch, fresh): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + caller = UserAPIKeyAuth( + api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-1", access_group_ids=["key-group"] + ).model_copy(update={"requires_fresh_policy": fresh}) + key_group = _group("key-group", ["secret-skill"], assigned_key_ids=[caller.token or ""]) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(side_effect=RuntimeError("writer down"))) + monkeypatch.setattr(auth_checks, "get_access_object", AsyncMock(return_value=key_group)) + monkeypatch.setattr(auth_checks, "get_skill_names_from_access_groups", AsyncMock(return_value=frozenset())) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + + with pytest.raises(RuntimeError, match="writer down"): + await access_group_skills(caller) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fresh", [True, False]) +async def test_access_group_skills_surfaces_team_group_lookup_failures(monkeypatch, fresh): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + caller = UserAPIKeyAuth( + api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-1", access_group_ids=["key-group"] + ).model_copy(update={"requires_fresh_policy": fresh}) + key_group = _group("key-group", ["secret-skill"], assigned_key_ids=[caller.token or ""]) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=_team("team-1", ["team-group"]))) + + async def lookup(access_group_id: str, **_kwargs: object) -> MagicMock: + if access_group_id == "team-group": + raise RuntimeError("writer down") + return key_group + + monkeypatch.setattr(auth_checks, "get_access_object", lookup) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + + with pytest.raises(RuntimeError, match="writer down"): + await access_group_skills(caller) + + +@pytest.mark.asyncio +async def test_access_group_skills_treats_confirmed_missing_team_like_no_team(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + caller = UserAPIKeyAuth( + api_key="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="removed-team", + access_group_ids=["key-group"], + ) + key_group = _group("key-group", ["secret-skill"], assigned_key_ids=[caller.token or ""]) + monkeypatch.setattr( + auth_checks, "get_team_object", AsyncMock(side_effect=auth_checks.TeamNotFoundError(team_id="removed-team")) + ) + monkeypatch.setattr(auth_checks, "get_access_object", AsyncMock(return_value=key_group)) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + + assert await access_group_skills(caller) == AccessGroupSkills(key=frozenset({"secret-skill"})) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fresh", [True, False]) +@pytest.mark.parametrize("missing_scope", ["key", "team"]) +async def test_access_group_skills_marks_scope_unresolved_when_a_listed_group_is_missing( + monkeypatch, fresh, missing_scope +): + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + key_groups = ["typo-group", "key-group"] if missing_scope == "key" else ["key-group"] + team_groups = ["team-group"] if missing_scope == "key" else ["stale-team-group", "team-group"] + caller = UserAPIKeyAuth( + api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-1", access_group_ids=key_groups + ).model_copy(update={"requires_fresh_policy": fresh}) + groups = { + "key-group": _group("key-group", ["secret-skill"], assigned_key_ids=[caller.token or ""]), + "team-group": _group("team-group", ["team-skill"]), + } + + async def lookup(access_group_id: str, **_kwargs: object) -> MagicMock: + if access_group_id not in groups: + raise auth_checks.AccessGroupNotFoundError(access_group_id=access_group_id) + return groups[access_group_id] + + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=_team("team-1", team_groups))) + monkeypatch.setattr(auth_checks, "get_access_object", lookup) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", MagicMock()) + + assert await access_group_skills(caller) == AccessGroupSkills(unresolved=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing_scope", ["key", "team"]) +async def test_access_group_skills_marks_scope_unresolved_through_the_real_group_lookup(monkeypatch, missing_scope): + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import cache_access_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key_groups = ["typo-group", "key-group"] if missing_scope == "key" else ["key-group"] + team_groups = ["team-group"] if missing_scope == "key" else ["stale-team-group", "team-group"] + caller = UserAPIKeyAuth( + api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-1", access_group_ids=key_groups + ) + cache = UserApiKeyCache() + for group in ( + LiteLLM_AccessGroupTable( + access_group_id="key-group", + access_group_name="key-group", + access_skill_names=["secret-skill"], + assigned_key_ids=[caller.token or ""], + ), + LiteLLM_AccessGroupTable( + access_group_id="team-group", access_group_name="team-group", access_skill_names=["team-skill"] + ), + ): + await cache_access_object( + access_group_id=group.access_group_id, access_group_table=group, user_api_key_cache=cache + ) + prisma = MagicMock() + prisma.db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=_team("team-1", team_groups))) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + + assert await access_group_skills(caller) == AccessGroupSkills(unresolved=True) + prisma.db.litellm_accessgrouptable.find_unique.assert_awaited_once() + + +def test_skill_visibility_is_public_only_while_a_group_scope_is_unresolved(): + caller = UserAPIKeyAuth( + api_key="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=_perm(["private-a"]), + team_object_permission=_perm(["private-a", "private-b"]), + ) + + assert skill_visibility(caller, AccessGroupSkills(unresolved=True)) == PUBLIC_ONLY + assert skill_visibility(caller, AccessGroupSkills()).granted == frozenset({"private-a"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check_db_only", [True, False]) +@pytest.mark.parametrize(("strict_errors", "expect_raise"), [(True, True), (False, False)]) +async def test_shared_group_loader_skips_missing_groups_unless_strict( + monkeypatch, check_db_only, strict_errors, expect_raise +): + """A group with no row grants nothing on cached and fresh reads alike; only ``strict_errors`` callers, + which turn the typed absence into an unresolved policy, see ``AccessGroupNotFoundError``.""" + from litellm.proxy.auth import auth_checks + + groups = {"real-group": _group("real-group", [])} + groups["real-group"].access_mcp_server_ids = ["mcp-1"] + + async def lookup(access_group_id: str, **_kwargs: object) -> MagicMock: + if access_group_id not in groups: + raise auth_checks.AccessGroupNotFoundError(access_group_id=access_group_id) + return groups[access_group_id] + + monkeypatch.setattr(auth_checks, "get_access_object", lookup) + call = auth_checks._get_resources_from_access_groups( + access_group_ids=["typo-group", "real-group"], + resource_field="access_mcp_server_ids", + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + check_db_only=check_db_only, + strict_errors=strict_errors, + ) + if expect_raise: + with pytest.raises(auth_checks.AccessGroupNotFoundError): + await call + return + assert await call == ["mcp-1"] diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 0da1148ac84..a8c11acbb30 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -2169,3 +2169,31 @@ def test_is_mcp_discovery_request(route, request_body, expected): from litellm.proxy.auth.auth_checks import is_mcp_discovery_request assert is_mcp_discovery_request(route=route, request_body=request_body) is expected + + +@pytest.mark.asyncio +async def test_get_skill_names_from_access_groups_reads_cached_group_rows(): + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import cache_access_object, get_skill_names_from_access_groups + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + for group_id, skills in (("group-a", ["skill-a", "shared"]), ("group-b", ["skill-b", "shared"])): + await cache_access_object( + access_group_id=group_id, + access_group_table=LiteLLM_AccessGroupTable( + access_group_id=group_id, access_group_name=group_id, access_skill_names=skills + ), + user_api_key_cache=cache, + ) + untouched_prisma = MagicMock() + untouched_prisma.db.litellm_accessgrouptable.find_unique = AsyncMock( + side_effect=AssertionError("cache hit must not reach the DB") + ) + + assert await get_skill_names_from_access_groups( + access_group_ids=("group-a", "group-b"), + prisma_client=untouched_prisma, + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + ) == frozenset({"skill-a", "skill-b", "shared"}) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 25022621edc..1f8c85ce337 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -8856,6 +8856,30 @@ async def test_get_access_object_db_fetch_returns_validated_access_group(): assert result.access_model_names == ["gpt-4"] +@pytest.mark.asyncio +async def test_get_access_object_distinguishes_absent_group_from_unreadable_row(): + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + absent_client = MagicMock() + absent_client.db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=None) + with pytest.raises(AccessGroupNotFoundError) as absent_info: + await get_access_object("typo-group", absent_client, mock_cache) + assert absent_info.value.status_code == 404 + assert "Access group doesn't exist in db" in str(absent_info.value.detail) + + unreadable_client = MagicMock() + unreadable_client.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=ConnectionError("db unreachable")) + with pytest.raises(HTTPException) as unreadable_info: + await get_access_object("real-group", unreadable_client, mock_cache) + assert unreadable_info.value.status_code == 404 + assert not isinstance(unreadable_info.value, AccessGroupNotFoundError) + + @pytest.mark.asyncio async def test_get_team_object_by_alias_db_fetch_returns_cached_obj(): from litellm.proxy._types import LiteLLM_TeamTableCachedObj diff --git a/tests/unit/proxy/auth/test_fallback_budget.py b/tests/unit/proxy/auth/test_fallback_budget.py index 3d31f1e0c78..1117ca43657 100644 --- a/tests/unit/proxy/auth/test_fallback_budget.py +++ b/tests/unit/proxy/auth/test_fallback_budget.py @@ -429,7 +429,11 @@ async def test_check_fails_closed_when_the_spend_lookup_breaks(monkeypatch): @pytest.mark.asyncio async def test_router_skips_the_paid_fallback_target_when_over_budget(): - from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget + from litellm.router_utils.fallback_event_handlers import ( + AttemptedFallbackTargets, + _is_fallback_target_within_budget, + ) + from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget router = Router( model_list=[FREE_MODEL, PAID_MODEL], @@ -438,19 +442,34 @@ async def test_router_skips_the_paid_fallback_target_when_over_budget(): ) over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} under = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + over_attempted: Final = AttemptedFallbackTargets() + under_attempted: Final = AttemptedFallbackTargets() - assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is False - assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", under) is True + assert ( + await _is_fallback_target_within_budget(router, "paid-model", "free-model", over, attempted=over_attempted) + is False + ) + assert ( + await _is_fallback_target_within_budget( + router, "paid-model", "free-model", under, attempted=under_attempted + ) + is True + ) + assert over_attempted.skipped == (SkippedFallbackTarget("paid-model", FallbackSkipReason.OVER_BUDGET),) @pytest.mark.asyncio async def test_router_without_a_budget_check_attempts_every_fallback(): - from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets, _is_fallback_target_within_budget router = _router() # fallback_budget_check defaults to None over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + attempted: Final = AttemptedFallbackTargets() - assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is True + assert ( + await _is_fallback_target_within_budget(router, "paid-model", "free-model", over, attempted=attempted) + is True + ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/common_utils/test_semantic_text_index.py b/tests/unit/proxy/common_utils/test_semantic_text_index.py new file mode 100644 index 00000000000..21fc33f839e --- /dev/null +++ b/tests/unit/proxy/common_utils/test_semantic_text_index.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Sequence +from typing import Final + +import pytest + +import litellm +from litellm.proxy.common_utils.semantic_text_index import EmbeddingFailed, SemanticTextIndex, Vector + +QUERY: Final = "weather query" +TEXTS: Final = ( + "valid one", + "OVERSIZED first", + "valid two", + "valid three", + "OVERSIZED second", + "valid four", + "valid five", +) + + +class ContextWindowEmbedder: + def __init__(self, dimensions: int = 2) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + self.vector: Final = (1.0,) * dimensions + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + if any("OVERSIZED" in text for text in texts): + raise litellm.ContextWindowExceededError( + message="input is too long", model="text-embedding-3-large", llm_provider="openai" + ) + return tuple(self.vector for _ in texts) + + +class RateLimitedEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + raise litellm.RateLimitError(message="rate limited", model="text-embedding-3-large", llm_provider="openai") + + +class RateLimitedAfterSplitEmbedder: + def __init__(self, rate_limited: bool = True) -> None: + self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs + self.rate_limited = rate_limited + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + if any("OVERSIZED" in text for text in texts): + raise litellm.ContextWindowExceededError( + message="input is too long", model="text-embedding-3-large", llm_provider="openai" + ) + if self.rate_limited and any("RATE LIMITED" in text for text in texts): + raise litellm.RateLimitError(message="rate limited", model="text-embedding-3-large", llm_provider="openai") + return tuple((1.0, 1.0) for _ in texts) + + +class ConcurrencyTrackingEmbedder: + def __init__(self) -> None: + self.in_flight = 0 + self.max_in_flight = 0 + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.in_flight += 1 + self.max_in_flight = max(self.max_in_flight, self.in_flight) + await asyncio.sleep(0) + self.in_flight -= 1 + if any("OVERSIZED" in text for text in texts): + raise litellm.ContextWindowExceededError( + message="input is too long", model="text-embedding-3-large", llm_provider="openai" + ) + return tuple((1.0, 1.0) for _ in texts) + + +class FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +class FixedDimensionEmbedder: + def __init__(self, dimensions: int) -> None: + self.vector: Final = (1.0,) * dimensions + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + return tuple(self.vector for _ in texts) + + +class TestSemanticTextIndex: + @pytest.mark.asyncio + async def test_omits_multiple_oversized_texts_and_preserves_score_order(self) -> None: + embedder = ContextWindowEmbedder() + scores: Final = await SemanticTextIndex().scores(QUERY, TEXTS, embedder, "emb") + + assert not isinstance(scores, EmbeddingFailed) + assert len(scores) == len(TEXTS) + assert scores[0] == pytest.approx(1.0) + assert scores[1] is None + assert scores[2] == pytest.approx(1.0) + assert scores[3] == pytest.approx(1.0) + assert scores[4] is None + assert scores[5] == pytest.approx(1.0) + assert scores[6] == pytest.approx(1.0) + + @pytest.mark.asyncio + async def test_repeat_search_only_embeds_the_query(self) -> None: + index = SemanticTextIndex() + embedder = ContextWindowEmbedder() + first: Final = await index.scores(QUERY, TEXTS, embedder, "emb") + assert not isinstance(first, EmbeddingFailed) + first_call_count: Final = len(embedder.calls) + + second: Final = await index.scores(QUERY, TEXTS, embedder, "emb") + + assert not isinstance(second, EmbeddingFailed) + assert embedder.calls[first_call_count:] == [(QUERY,)] + assert second == first + + @pytest.mark.asyncio + async def test_oversized_query_returns_embedding_failed(self) -> None: + result: Final = await SemanticTextIndex().scores( + "OVERSIZED query", ("valid text",), ContextWindowEmbedder(), "emb" + ) + + assert isinstance(result, EmbeddingFailed) + + @pytest.mark.asyncio + async def test_non_context_provider_error_fails_after_one_embed_call(self) -> None: + embedder = RateLimitedEmbedder() + + result: Final = await SemanticTextIndex().scores(QUERY, ("valid text",), embedder, "emb") + + assert isinstance(result, EmbeddingFailed) + assert len(embedder.calls) == 1 + + @pytest.mark.asyncio + async def test_oversized_text_is_retried_once_the_exclusion_expires(self) -> None: + clock = FakeClock() + index = SemanticTextIndex(too_long_ttl_seconds=60.0, clock=clock) + embedder = ContextWindowEmbedder() + texts: Final = ("valid one", "OVERSIZED first") + _ = await index.scores(QUERY, texts, embedder, "emb") + clock.now = 59.0 + before_expiry: Final = len(embedder.calls) + _ = await index.scores(QUERY, texts, embedder, "emb") + assert embedder.calls[before_expiry:] == [(QUERY,)] + + clock.now = 61.0 + after_expiry: Final = len(embedder.calls) + scores: Final = await index.scores(QUERY, texts, embedder, "emb") + + assert scores == (pytest.approx(1.0), None) + assert embedder.calls[after_expiry] == (QUERY, "OVERSIZED first") + + @pytest.mark.asyncio + async def test_splitting_sends_one_embedding_request_at_a_time(self) -> None: + embedder = ConcurrencyTrackingEmbedder() + + scores: Final = await SemanticTextIndex().scores(QUERY, TEXTS, embedder, "emb") + + assert not isinstance(scores, EmbeddingFailed) + assert embedder.max_in_flight == 1 + + @pytest.mark.asyncio + async def test_non_context_provider_error_while_splitting_fails_the_search(self) -> None: + result: Final = await SemanticTextIndex().scores( + QUERY, ("RATE LIMITED text", "OVERSIZED text", "valid text"), RateLimitedAfterSplitEmbedder(), "emb" + ) + + assert isinstance(result, EmbeddingFailed) + assert "rate limited" in result.reason + + @pytest.mark.asyncio + async def test_failed_split_keeps_vectors_and_too_long_texts_for_the_next_search(self) -> None: + index = SemanticTextIndex() + embedder = RateLimitedAfterSplitEmbedder() + texts: Final = ("valid text", "OVERSIZED text", "RATE LIMITED text") + + first: Final = await index.scores(QUERY, texts, embedder, "emb") + assert isinstance(first, EmbeddingFailed) + assert "rate limited" in first.reason + + embedder.rate_limited = False + before: Final = len(embedder.calls) + second: Final = await index.scores(QUERY, texts, embedder, "emb") + + assert second == (pytest.approx(1.0), None, pytest.approx(1.0)) + assert embedder.calls[before:] == [(QUERY, "RATE LIMITED text")] + + @pytest.mark.asyncio + async def test_mixed_dimension_reembedding_omits_oversized_text(self) -> None: + index = SemanticTextIndex() + seeded: Final = await index.scores(QUERY, ("cached text", "OVERSIZED text"), FixedDimensionEmbedder(3), "emb") + assert not isinstance(seeded, EmbeddingFailed) + embedder = ContextWindowEmbedder(2) + + scores: Final = await index.scores(QUERY, ("cached text", "OVERSIZED text"), embedder, "emb") + + assert not isinstance(scores, EmbeddingFailed) + assert scores[0] == pytest.approx(1.0) + assert scores[1] is None + assert (QUERY, "cached text", "OVERSIZED text") in embedder.calls + assert ("OVERSIZED text",) in embedder.calls diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 7b8d000a8d5..99b500ec315 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -1,8 +1,10 @@ +import asyncio import json import os import signal import sys import urllib.parse +from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -10,7 +12,13 @@ from fastapi.testclient import TestClient -from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema +from litellm.proxy.db.prisma_client import ( + PrismaManager, + PrismaWrapper, + _PrismaDrainTracker, + _TrackedPrismaEngine, + should_update_prisma_schema, +) @pytest.fixture(autouse=True) @@ -470,3 +478,97 @@ def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None) assert run() is outcome + + +class _SlowEngine: + """Engine double whose requests only answer once the test lets them, recording what the engine saw.""" + + def __init__(self) -> None: + self.gates: dict[str, asyncio.Event] = {} # mutable-ok: one gate per request the test lets through + self.queries_finished = 0 + self.requests_cancelled = 0 + self.rolled_back: list[str] = [] # mutable-ok: test records + self.process = object() + + def gate(self, request: str) -> asyncio.Event: + return self.gates.setdefault(request, asyncio.Event()) + + async def _answer_once_released(self, request: str) -> None: + try: + await self.gate(request).wait() + except asyncio.CancelledError: + self.requests_cancelled += 1 + raise + + async def query(self, content: str, *, tx_id: str | None) -> object: + await self._answer_once_released(content) + self.queries_finished += 1 + return {"content": content, "tx_id": tx_id} + + async def start_transaction(self, *, content: str) -> str: + await self._answer_once_released("start") + return "tx-opened-after-caller-left" + + async def commit_transaction(self, tx_id: str) -> None: + raise AssertionError(f"unexpected commit of {tx_id}") + + async def rollback_transaction(self, tx_id: str) -> None: + self.rolled_back.append(tx_id) + + +def _tracked(engine: _SlowEngine) -> tuple[_TrackedPrismaEngine, _PrismaDrainTracker]: + tracker: Final = _PrismaDrainTracker() + return _TrackedPrismaEngine(cast(Any, engine), tracker), tracker + + +@pytest.mark.asyncio +async def test_cancelling_a_query_caller_lets_the_engine_finish_the_statement() -> None: + engine: Final = _SlowEngine() + tracked, tracker = _tracked(engine) + caller: Final = asyncio.create_task(tracked.query("x", tx_id=None)) + await asyncio.sleep(0) + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + drained: Final = asyncio.create_task(tracker.wait_until_drained()) + await asyncio.sleep(0) + assert engine.requests_cancelled == 0 + assert not drained.done() + engine.gate("x").set() + await asyncio.wait_for(drained, timeout=1) + assert engine.queries_finished == 1 + + +@pytest.mark.asyncio +async def test_transaction_opened_after_its_caller_was_cancelled_is_rolled_back() -> None: + engine: Final = _SlowEngine() + tracked, tracker = _tracked(engine) + caller: Final = asyncio.create_task(tracked.start_transaction(content="{}")) + await asyncio.sleep(0) + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + engine.gate("start").set() + await asyncio.wait_for(tracker.wait_until_drained(), timeout=1) + assert engine.requests_cancelled == 0 + assert engine.rolled_back == ["tx-opened-after-caller-left"] + + +@pytest.mark.asyncio +async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain_open() -> None: + engine: Final = _SlowEngine() + tracked, tracker = _tracked(engine) + first: Final = asyncio.create_task(tracked.query("first", tx_id=None)) + await asyncio.sleep(0) + await asyncio.sleep(0) + drained: Final = asyncio.create_task(tracker.wait_until_drained()) + await asyncio.sleep(0) + second: Final = asyncio.create_task(tracked.query("second", tx_id=None)) + engine.gate("first").set() + await first + await asyncio.sleep(0) + assert not drained.done() + engine.gate("second").set() + await second + await asyncio.wait_for(drained, timeout=1) + assert engine.queries_finished == 2 diff --git a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 34a7f78659b..1ef0eee761a 100644 --- a/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -1,11 +1,11 @@ import os +from typing import Final from unittest.mock import MagicMock, patch import pytest from fastapi import FastAPI from fastapi.testclient import TestClient - from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import router from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry @@ -21,7 +21,6 @@ def test_ui_discovery_endpoints_with_defaults(): patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -44,7 +43,6 @@ def test_ui_discovery_endpoints_with_custom_server_root_path(): patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -69,7 +67,6 @@ def test_ui_discovery_endpoints_with_proxy_base_url_when_set(): patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): - response = client.get("/litellm/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -80,6 +77,41 @@ def test_ui_discovery_endpoints_with_proxy_base_url_when_set(): assert data["sso_configured"] is False +@pytest.mark.parametrize( + ("gateway_url", "expected_gateway_url"), + [ + ("https://gw.example.com/", "https://gw.example.com"), + (" ", None), + (None, None), + ], +) +def test_ui_discovery_endpoints_normalizes_gateway_base_url( + monkeypatch: pytest.MonkeyPatch, + gateway_url: str | None, + expected_gateway_url: str | None, +) -> None: + app: Final[FastAPI] = FastAPI() + app.include_router(router) + client: Final[TestClient] = TestClient(app) + + if gateway_url is None: + monkeypatch.delenv("LITELLM_GATEWAY_URL", raising=False) + else: + monkeypatch.setenv("LITELLM_GATEWAY_URL", gateway_url) + + with ( + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + patch("litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com"), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), + patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), + ): + response: Final = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + assert response.json()["proxy_base_url"] == "https://proxy.example.com" + assert response.json()["gateway_base_url"] == expected_gateway_url + + def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_enabled(): app = FastAPI() app.include_router(router) @@ -98,7 +130,6 @@ def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_enabled(): clear=False, ), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -155,7 +186,6 @@ def test_ui_discovery_endpoints_with_sso_configured_but_auto_redirect_disabled() clear=False, ), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -181,7 +211,6 @@ def test_ui_discovery_endpoints_with_sso_not_configured_but_auto_redirect_enable clear=False, ), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -210,7 +239,6 @@ def test_ui_discovery_endpoints_both_routes_return_same_data(): clear=False, ), ): - response1 = client.get("/.well-known/litellm-ui-config") response2 = client.get("/litellm/.well-known/litellm-ui-config") @@ -265,7 +293,6 @@ def test_ui_discovery_endpoints_with_auto_redirect_env_var_overrides_general_set clear=False, ), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -284,7 +311,6 @@ def test_ui_discovery_endpoints_with_admin_ui_disabled(): patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "true"}, clear=False), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -303,9 +329,7 @@ def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured(): mock_config = MagicMock() mock_config.worker_registry = [ - WorkerRegistryEntry( - worker_id="team-a", name="Team A", url="https://worker-1:4001" - ), + WorkerRegistryEntry(worker_id="team-a", name="Team A", url="https://worker-1:4001"), ] with ( @@ -315,7 +339,6 @@ def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured(): patch("litellm.proxy.proxy_server.proxy_config", mock_config), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -405,7 +428,6 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_via_env_var(): clear=False, ), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 @@ -453,7 +475,6 @@ def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers(): patch("litellm.proxy.proxy_server.proxy_config", mock_config), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): - response = client.get("/.well-known/litellm-ui-config") assert response.status_code == 200 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py new file mode 100644 index 00000000000..61f41b0482a --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py @@ -0,0 +1,625 @@ +import json +import uuid +from collections.abc import AsyncGenerator +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +import litellm +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail +from litellm.proxy.guardrails.init_guardrails import ( + init_guardrails_v2, # pyright: ignore[reportUnknownVariableType] # untyped upstream +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel +from litellm.types.utils import GenericGuardrailAPIInputs + +API_BASE = "https://gateway.airia.ai" +API_KEY = "ak-test-key" + + +def _make_guardrail(**kwargs: object) -> AiriaGuardrail: + """Build a guardrail; construction defaults still exercise the real HTTP client.""" + kwargs.setdefault("api_base", API_BASE) + kwargs.setdefault("api_key", API_KEY) + kwargs.setdefault("guardrail_name", "airia-guard") + return AiriaGuardrail(**kwargs) + + +def _make_guardrail_with_response(payload: object, status_code: int = 200, **kwargs: object) -> AiriaGuardrail: + """Build a guardrail with a fake HTTP client injected at construction, per CLAUDE.md's + dependency-injection rule: no test mutates `async_handler` after the fact.""" + handler = MagicMock() + handler.post = AsyncMock(return_value=_mock_response(payload, status_code)) + return _make_guardrail(async_handler=handler, **kwargs) + + +def _mock_response(payload: object, status_code: int = 200) -> MagicMock: + response = MagicMock() + response.status_code = status_code + response.json.return_value = payload + response.raise_for_status = MagicMock() + return response + + +def _as_mock(value: object) -> MagicMock: + """Forget the declared type, which is the real handler's, and recover the mock that was injected.""" + assert isinstance(value, MagicMock) + return value + + +def _post_mock(guardrail: AiriaGuardrail) -> AsyncMock: + """The `post` the test injected, typed as the mock it is so call inspection stays checked.""" + post = _as_mock(guardrail.async_handler).post + assert isinstance(post, AsyncMock) + return post + + +def _timeout(guardrail: AiriaGuardrail) -> httpx.Timeout: + timeout = guardrail.async_handler.timeout # pyright: ignore[reportUnknownMemberType] # the handler's attribute is partially untyped + assert isinstance(timeout, httpx.Timeout) + return timeout + + +def _inputs(**overrides: object) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs(**{"texts": ["hello"], "model": "gpt-4o", **overrides}) + + +def test_airia_guardrail_config(monkeypatch: pytest.MonkeyPatch) -> None: + """The guardrail registers through init_guardrails_v2 under the `airia` key.""" + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setenv("AIRIA_GATEWAY_URL", API_BASE) + monkeypatch.setenv("AIRIA_API_KEY", API_KEY) + + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "airia-guard", + "litellm_params": { + "guardrail": "airia", + "mode": "pre_call", + "default_on": True, + "timeout": 45.0, + }, + } + ], + config_file_path="", + ) + + registered = [c for c in litellm.callbacks if isinstance(c, AiriaGuardrail)] + assert len(registered) == 1 + assert registered[0].guardrail_name == "airia-guard" + assert registered[0].default_on is True + assert registered[0].event_hook == "pre_call" + assert _timeout(registered[0]).read == 45.0 + + +def test_during_call_is_not_a_supported_event_hook() -> None: + """during_call runs concurrently with the upstream call, so a block can land too late.""" + hooks = _make_guardrail().supported_event_hooks + + assert hooks is not None + assert GuardrailEventHooks.during_call not in hooks + assert list(hooks) == [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + +@pytest.mark.parametrize("missing", ["api_base", "api_key"]) +def test_missing_credentials_raise_at_construction(missing: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Fail at startup rather than on the first request.""" + monkeypatch.delenv("AIRIA_GATEWAY_URL", raising=False) + monkeypatch.delenv("AIRIA_API_KEY", raising=False) + + kwargs = {"api_base": API_BASE, "api_key": API_KEY} + kwargs[missing] = None + + with pytest.raises(ValueError, match="AiriaGuardrail requires"): + _make_guardrail(**kwargs) + + +def test_credentials_fall_back_to_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AIRIA_GATEWAY_URL", API_BASE) + monkeypatch.setenv("AIRIA_API_KEY", API_KEY) + + guardrail = _make_guardrail(api_base=None, api_key=None) + + assert guardrail.api_base == API_BASE + assert guardrail.api_key == API_KEY + + +def test_trailing_slash_is_stripped_from_api_base() -> None: + guardrail = _make_guardrail(api_base=f"{API_BASE}/") + + assert guardrail.api_base == API_BASE + + +def test_custom_timeout_from_kwargs() -> None: + guardrail = _make_guardrail(timeout=45.0) + + assert _timeout(guardrail).read == 45.0 + assert _timeout(guardrail).connect == 5.0 + + +def test_timeout_falls_back_to_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AIRIA_TIMEOUT", "30") + + guardrail = _make_guardrail() + + assert _timeout(guardrail).read == 30.0 + + +@pytest.mark.asyncio +async def test_request_payload_and_auth_header() -> None: + guardrail = _make_guardrail_with_response({"action": "NONE"}) + + inputs = _inputs( + images=["data:image/png;base64,AAAA"], + structured_messages=[{"role": "user", "content": "hello"}], + tools=[{"type": "function", "function": {"name": "f"}}], + tool_calls=[{"id": "c1", "function": {"name": "f", "arguments": "{}"}}], + ) + await guardrail.apply_guardrail(inputs=inputs, request_data={"litellm_call_id": "call-123"}, input_type="request") + + call = _post_mock(guardrail).call_args + assert call.args[0] == f"{API_BASE}/v1/guardrails/litellm" + assert call.kwargs["headers"] == {"Authorization": f"Bearer {API_KEY}"} + + payload = call.kwargs["json"] + assert payload["input_type"] == "request" + assert payload["texts"] == ["hello"] + assert payload["structured_messages"] == [{"role": "user", "content": "hello"}] + assert payload["images"] == ["data:image/png;base64,AAAA"] + assert payload["tools"] == [{"type": "function", "function": {"name": "f"}}] + assert payload["tool_calls"] == [{"id": "c1", "function": {"name": "f", "arguments": "{}"}}] + assert payload["model"] == "gpt-4o" + assert payload["litellm_call_id"] == "call-123" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request", "response"]) +async def test_input_type_is_forwarded_verbatim(input_type: Literal["request", "response"]) -> None: + guardrail = _make_guardrail_with_response({"action": "NONE"}) + + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type=input_type) + + assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == input_type + + +@pytest.mark.asyncio +async def test_call_id_falls_back_to_a_generated_uuid() -> None: + guardrail = _make_guardrail_with_response({"action": "NONE"}) + + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + call_id = _post_mock(guardrail).call_args.kwargs["json"]["litellm_call_id"] + assert uuid.UUID(call_id).version == 4 + + +@pytest.mark.asyncio +async def test_action_none_returns_inputs_unchanged() -> None: + guardrail = _make_guardrail_with_response({"action": "NONE"}) + inputs = _inputs(structured_messages=[{"role": "user", "content": "hello"}]) + + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + assert result["structured_messages"] == [{"role": "user", "content": "hello"}] + + +@pytest.mark.asyncio +async def test_action_blocked_raises_with_blocked_content_set() -> None: + guardrail = _make_guardrail_with_response({"action": "BLOCKED", "blocked_reason": "Contains a secret"}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + assert "Contains a secret" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_action_blocked_without_a_reason_uses_the_default_message() -> None: + guardrail = _make_guardrail_with_response({"action": "BLOCKED"}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + assert "Blocked by your organization's content policy." in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_intervened_replaces_both_text_bearing_fields() -> None: + """Substituting only one field would leave the other still carrying the unredacted text.""" + guardrail = _make_guardrail_with_response( + { + "action": "GUARDRAIL_INTERVENED", + "texts": ["my email is [EmailAddress1]"], + "structured_messages": [{"role": "user", "content": "my email is [EmailAddress1]"}], + } + ) + inputs = _inputs( + texts=["my email is a@b.com"], + structured_messages=[{"role": "user", "content": "my email is a@b.com"}], + ) + + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result["texts"] == ["my email is [EmailAddress1]"] + assert result["structured_messages"] == [{"role": "user", "content": "my email is [EmailAddress1]"}] + assert "a@b.com" not in str(result) + + +@pytest.mark.asyncio +async def test_intervened_leaves_a_field_alone_when_airia_omits_it() -> None: + """An omitted field means "not rewritten", not "rewritten to empty".""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["redacted"]}) + inputs = _inputs(texts=["original"], structured_messages=[{"role": "user", "content": "keep me"}]) + + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result["texts"] == ["redacted"] + assert result["structured_messages"] == [{"role": "user", "content": "keep me"}] + + +@pytest.mark.asyncio +async def test_intervened_with_a_non_list_field_is_treated_as_a_block() -> None: + """A rewrite this version cannot apply must not let the original text through.""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": "not a list"}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(texts=["secret"]), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +async def test_intervened_with_a_non_list_structured_messages_is_treated_as_a_block() -> None: + guardrail = _make_guardrail_with_response( + {"action": "GUARDRAIL_INTERVENED", "structured_messages": {"role": "user"}} + ) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +async def test_intervened_without_any_rewritten_field_is_treated_as_a_block() -> None: + """An intervention the proxy cannot apply must not let the original content through.""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "tool_calls": []}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(texts=["secret"]), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [["not", "an", "object"], "just a string", 42, None]) +async def test_a_non_object_response_body_fails_closed_but_is_not_reported_as_a_verdict(body: object) -> None: + """response.json() is not guaranteed to return a dict; body.get(...) on a list or scalar + would otherwise raise AttributeError and surface as an internal error, not a clean refusal.""" + guardrail = _make_guardrail_with_response(body) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is False + assert "could not evaluate" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_intervened_with_a_texts_rewrite_of_the_wrong_length_is_treated_as_a_block() -> None: + """A rewrite with a different item count than what was sent would misalign the positional + write-back downstream; accepting it could drop content or leave other items unredacted.""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["only one redacted item"]}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail( + inputs=_inputs(texts=["first secret", "second secret"]), request_data={}, input_type="request" + ) + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("structured_messages", [[], [{"role": "user", "content": "[Redacted]"}]]) +async def test_intervened_rewrite_for_a_field_that_was_not_sent_is_treated_as_a_block( + structured_messages: list[object], +) -> None: + """post_call sends only `texts`. A `structured_messages` rewrite then has nothing to replace, and an + empty one must not count as "applied" while the `texts` that carried the content pass through.""" + guardrail = _make_guardrail_with_response( + {"action": "GUARDRAIL_INTERVENED", "structured_messages": structured_messages} + ) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(texts=["secret"]), request_data={}, input_type="response") + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +async def test_intervened_returns_a_copy_and_leaves_the_caller_object_untouched() -> None: + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["redacted"]}) + inputs = _inputs(texts=["original"]) + + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result["texts"] == ["redacted"] + assert inputs["texts"] == ["original"] + + +@pytest.mark.asyncio +async def test_streaming_redaction_actually_reaches_the_client() -> None: + """block_only (the framework default) drops guardrail rewrites on the streaming path, so a + REDACT verdict would silently leak the original text. Runs the real UnifiedLLMGuardrails + streaming iterator end to end (not just this guardrail's config attributes) and checks what + the client would actually receive. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail: Final = _make_guardrail_with_response( + {"action": "GUARDRAIL_INTERVENED", "texts": ["Contact me at [EmailAddress1]"]}, + event_hook="post_call", + default_on=True, + ) + + async def source_stream() -> AsyncGenerator[ModelResponseStream, None]: + for content in ("Contact ", "me at ", "ada@example.com"): + yield ModelResponseStream( + model="gpt-4o-mini", choices=[StreamingChoices(index=0, delta=Delta(content=content))] + ) + yield ModelResponseStream( + model="gpt-4o-mini", + choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")], + ) + + assembled = "" + async for chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( # pyright: ignore[reportUnknownMemberType] # annotates request_data as a bare dict + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"), + response=source_stream(), + request_data={ # mutable-ok: one-shot request_data the framework requires as a plain dict + "messages": [{"role": "user", "content": "share your contact"}], + "stream": True, + }, + guardrail_to_apply=guardrail, + ): + for choice in chunk.choices: + if choice.delta.content: + assembled += choice.delta.content + + assert "[EmailAddress1]" in assembled + assert "ada@example.com" not in assembled + _post_mock(guardrail).assert_called_once() + + +def _streamed_request(route: str | None) -> dict[str, object]: + """A proxy request_data as the streaming hook sees it: `stream` set, and the route the proxy + stamped into the metadata bucket (`litellm_metadata` for the non-chat routes under test).""" + request_data: dict[str, object] = {"stream": True} # mutable-ok: built once per test + if route is not None: + request_data["litellm_metadata"] = {"user_api_key_request_route": route} + return request_data + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/messages", "/v1/responses", "/responses", "/unknown/route", None]) +async def test_streamed_redaction_on_a_route_without_stream_rewrites_is_enforced_as_a_block( + route: str | None, +) -> None: + """incremental_diff exists only for chat completions; every other streamed route runs block_only, + which drops rewrites and replays the original text. There a REDACT verdict must block instead.""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail( + inputs=_inputs(texts=["secret"]), request_data=_streamed_request(route), input_type="response" + ) + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/chat/completions", "/v1/chat/completions"]) +async def test_streamed_redaction_on_chat_completions_is_still_delivered(route: str) -> None: + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]}) + request_data: dict[str, object] = {"stream": True, "metadata": {"user_api_key_request_route": route}} + + result = await guardrail.apply_guardrail( + inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response" + ) + + assert result["texts"] == ["[Redacted]"] + + +@pytest.mark.asyncio +async def test_unstreamed_redaction_on_a_non_chat_route_is_delivered() -> None: + """Only the stream cannot carry a rewrite; a non-streamed Responses call is rewritten as usual.""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]}) + request_data: dict[str, object] = {"litellm_metadata": {"user_api_key_request_route": "/v1/responses"}} + + result = await guardrail.apply_guardrail( + inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response" + ) + + assert result["texts"] == ["[Redacted]"] + + +@pytest.mark.asyncio +async def test_streamed_request_side_redaction_is_unaffected() -> None: + """A pre_call rewrite lands in the prompt before the model is called; streaming is irrelevant there.""" + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]}) + + result = await guardrail.apply_guardrail( + inputs=_inputs(texts=["secret"]), request_data=_streamed_request("/v1/messages"), input_type="request" + ) + + assert result["texts"] == ["[Redacted]"] + + +@pytest.mark.asyncio +async def test_stream_flag_is_also_read_from_the_proxied_request_body() -> None: + guardrail = _make_guardrail_with_response({"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]}) + request_data: dict[str, object] = { + "proxy_server_request": {"body": {"stream": True}}, + "litellm_metadata": {"user_api_key_request_route": "/v1/messages"}, + } + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail( + inputs=_inputs(texts=["secret"]), request_data=request_data, input_type="response" + ) + + assert excinfo.value.blocked_content is True + + +def _anthropic_sse(event_type: str, data: dict[str, object]) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + +async def _anthropic_messages_stream() -> AsyncGenerator[bytes, None]: + """A complete Anthropic /v1/messages SSE stream, the shape the proxy relays on that route.""" + yield _anthropic_sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_orig", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ) + yield _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ) + for text in ("Contact ", "me at ", "ada@example.com"): + yield _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + yield _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}) + yield _anthropic_sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + yield _anthropic_sse("message_stop", {"type": "message_stop"}) + + +async def _drain_into(stream: AsyncGenerator[object, None], sink: list[object]) -> None: + """Forward every chunk the client would have received into `sink`, keeping what leaked before a + raise (a comprehension would drop it). Lives outside the test so `pytest.raises` wraps one await.""" + async for chunk in stream: + sink.append(chunk) # noqa: PERF401 # the partial list is the point when the stream is refused + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "verdict", + [ + {"action": "BLOCKED", "blocked_reason": "policy"}, + {"action": "GUARDRAIL_INTERVENED", "texts": ["Contact me at [EmailAddress1]"]}, + ], +) +async def test_non_chat_streamed_route_releases_nothing_before_the_verdict(verdict: dict[str, object]) -> None: + """Drives the real UnifiedLLMGuardrails iterator on /v1/messages, which incremental_diff does not + support, so the framework takes its block_only fallback. streaming_buffer_until_moderated holds + the text until the end-of-stream verdict; a block, and a redaction block_only could not deliver, + both refuse the response without one original chunk reaching the client.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail: Final = _make_guardrail_with_response(verdict, event_hook="post_call", default_on=True) + + released: list[object] = [] # mutable-ok: what the client received before the stream was refused + hook: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( # pyright: ignore[reportUnknownMemberType] # annotates request_data as a bare dict + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/messages"), + response=_anthropic_messages_stream(), + request_data={ # mutable-ok: one-shot request_data the framework requires as a plain dict + "messages": [{"role": "user", "content": "share your contact"}], + "stream": True, + }, + guardrail_to_apply=guardrail, + ) + with pytest.raises(GuardrailRaisedException) as excinfo: + await _drain_into(hook, released) + + assert excinfo.value.blocked_content is True + assert "ada@example.com" not in "".join(c.decode() if isinstance(c, bytes) else str(c) for c in released) + _post_mock(guardrail).assert_called_once() + + +def test_config_model_is_exposed() -> None: + assert AiriaGuardrail.get_config_model() is AiriaGuardrailConfigModel + assert AiriaGuardrailConfigModel.ui_friendly_name() == "Airia Guardrail" + + +@pytest.mark.asyncio +async def test_unrecognized_action_is_treated_as_a_block() -> None: + """An action this version cannot interpret must not become an allow.""" + guardrail = _make_guardrail_with_response({"action": "SOME_FUTURE_ACTION"}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +async def test_missing_action_is_treated_as_a_block() -> None: + """A response with no action at all is not an allow either.""" + guardrail = _make_guardrail_with_response({}) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is True + + +@pytest.mark.asyncio +async def test_transport_failure_fails_closed_but_is_not_reported_as_a_verdict() -> None: + """Still refused, but blocked_content stays False: could not evaluate is not a verdict.""" + handler: Final = MagicMock() + handler.post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + guardrail: Final = _make_guardrail(async_handler=handler) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is False + assert "could not evaluate" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_non_2xx_fails_closed_but_is_not_reported_as_a_verdict() -> None: + response: Final = _mock_response({}, status_code=503) + response.raise_for_status.side_effect = httpx.HTTPStatusError("503", request=MagicMock(), response=MagicMock()) + handler: Final = MagicMock() + handler.post = AsyncMock(return_value=response) + guardrail: Final = _make_guardrail(async_handler=handler) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await guardrail.apply_guardrail(inputs=_inputs(), request_data={}, input_type="request") + + assert excinfo.value.blocked_content is False diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 7d70b96eeff..19eb3c71a5b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -115,14 +115,9 @@ class TestGenericGuardrailAPIConfiguration: headers={"Authorization": "Bearer test-key"}, additional_provider_specific_params={"custom_param": "value"}, ) - assert ( - guardrail.api_base - == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" - ) + assert guardrail.api_base == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" assert guardrail.headers == {"Authorization": "Bearer test-key"} - assert guardrail.additional_provider_specific_params == { - "custom_param": "value" - } + assert guardrail.additional_provider_specific_params == {"custom_param": "value"} def test_init_with_env_vars(self): """Test initialization with environment variables""" @@ -133,10 +128,7 @@ class TestGenericGuardrailAPIConfiguration: }, ): guardrail = GenericGuardrailAPI() - assert ( - guardrail.api_base - == "https://env.api.guardrail.com/beta/litellm_basic_guardrail_api" - ) + assert guardrail.api_base == "https://env.api.guardrail.com/beta/litellm_basic_guardrail_api" def test_init_without_api_base_raises_error(self): """Test that initialization without API base raises ValueError""" @@ -149,20 +141,14 @@ class TestGenericGuardrailAPIConfiguration: guardrail = GenericGuardrailAPI( api_base="https://api.test.guardrail.com/v1", ) - assert ( - guardrail.api_base - == "https://api.test.guardrail.com/v1/beta/litellm_basic_guardrail_api" - ) + assert guardrail.api_base == "https://api.test.guardrail.com/v1/beta/litellm_basic_guardrail_api" def test_api_base_not_duplicated(self): """Test that endpoint path is not duplicated if already present""" guardrail = GenericGuardrailAPI( api_base="https://api.test.guardrail.com/beta/litellm_basic_guardrail_api", ) - assert ( - guardrail.api_base - == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" - ) + assert guardrail.api_base == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" def test_api_key_sets_x_api_key_header(self): """Test that api_key is set as x-api-key header""" @@ -224,9 +210,7 @@ class TestExtraHeadersForwarding: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=mock_response) as mock_post: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=request_data, @@ -264,9 +248,7 @@ class TestExtraHeadersForwarding: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=mock_response) as mock_post: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=request_data, @@ -285,9 +267,7 @@ class TestMetadataExtraction: """Test metadata extraction from request data""" @pytest.mark.asyncio - async def test_extract_metadata_from_input_request( - self, generic_guardrail, mock_request_data_input - ): + async def test_extract_metadata_from_input_request(self, generic_guardrail, mock_request_data_input): """Test extracting metadata from input request (metadata field)""" # Mock API response mock_response = MagicMock() @@ -297,9 +277,7 @@ class TestMetadataExtraction: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["Who is Ishaan?"]}, request_data=mock_request_data_input, @@ -355,9 +333,7 @@ class TestMetadataExtraction: } mock_api_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_api_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_api_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["hey i'm ishaan!"]}, request_data=request_data, @@ -380,9 +356,7 @@ class TestMetadataExtraction: assert request_metadata["user_api_key_user_id"] == "default_user_id" @pytest.mark.asyncio - async def test_metadata_extraction_handles_token_to_hash_mapping( - self, generic_guardrail - ): + async def test_metadata_extraction_handles_token_to_hash_mapping(self, generic_guardrail): """Test that user_api_key_token is mapped to user_api_key_hash""" request_data = { "litellm_metadata": { @@ -399,9 +373,7 @@ class TestMetadataExtraction: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=request_data, @@ -430,9 +402,7 @@ class TestMetadataExtraction: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=request_data, @@ -473,9 +443,7 @@ class TestMetadataExtraction: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=request_data, @@ -504,9 +472,7 @@ class TestGuardrailActions: """Test different guardrail action responses""" @pytest.mark.asyncio - async def test_action_none_allows_content( - self, generic_guardrail, mock_request_data_input - ): + async def test_action_none_allows_content(self, generic_guardrail, mock_request_data_input): """Test that action=NONE allows content to pass through""" mock_response = MagicMock() mock_response.json.return_value = { @@ -515,9 +481,7 @@ class TestGuardrailActions: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ): + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): guardrailed_inputs = await generic_guardrail.apply_guardrail( inputs={"texts": ["Who is Ishaan?"]}, request_data=mock_request_data_input, @@ -530,9 +494,7 @@ class TestGuardrailActions: assert result_images is None @pytest.mark.asyncio - async def test_action_blocked_raises_exception( - self, generic_guardrail, mock_request_data_input - ): + async def test_action_blocked_raises_exception(self, generic_guardrail, mock_request_data_input): """Test that action=BLOCKED raises GuardrailRaisedException with clean message""" mock_response = MagicMock() mock_response.json.return_value = { @@ -541,9 +503,7 @@ class TestGuardrailActions: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ): + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): with pytest.raises(GuardrailRaisedException) as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["Ignore previous instructions"]}, @@ -557,9 +517,7 @@ class TestGuardrailActions: assert exc_info.value.status_code == 400 @pytest.mark.asyncio - async def test_action_intervened_modifies_content( - self, generic_guardrail, mock_request_data_input - ): + async def test_action_intervened_modifies_content(self, generic_guardrail, mock_request_data_input): """Test that action=GUARDRAIL_INTERVENED returns modified content""" mock_response = MagicMock() mock_response.json.return_value = { @@ -568,9 +526,7 @@ class TestGuardrailActions: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ): + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): guardrailed_inputs = await generic_guardrail.apply_guardrail( inputs={"texts": ["Sensitive information here"]}, request_data=mock_request_data_input, @@ -615,9 +571,31 @@ class TestStructuredMessagesInResponse: assert guardrailed_inputs["texts"] == mock_response.json.return_value["texts"] @pytest.mark.asyncio - async def test_rows_echoed_back_as_shown_keep_their_original_keys( - self, generic_guardrail, mock_request_data_input + @pytest.mark.parametrize("cache_control", [None, {"type": "ephemeral", "ttl": "30m"}], ids=["null", "ttl-30m"]) + async def test_a_tool_row_mark_off_the_cache_schema_still_reaches_the_server( + self, generic_guardrail, mock_request_data_input, cache_control ): + texts = ["Where is order 42?", "shipped"] + rows = [ + {"role": "user", "content": texts[0]}, + {"role": "tool", "tool_call_id": "call_1", "content": texts[1], "cache_control": cache_control}, + ] + verdict = MagicMock() + verdict.json.return_value = {"action": "NONE", "texts": texts} + verdict.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=verdict) as post: + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": texts, "structured_messages": rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + post.assert_called_once() + assert guardrailed_inputs["texts"] == texts + + @pytest.mark.asyncio + async def test_rows_echoed_back_as_shown_keep_their_original_keys(self, generic_guardrail, mock_request_data_input): tool_call_row = { "role": "assistant", "content": None, @@ -726,9 +704,7 @@ class TestImageSupport: """Test image handling in guardrail requests""" @pytest.mark.asyncio - async def test_images_passed_in_request( - self, generic_guardrail, mock_request_data_input - ): + async def test_images_passed_in_request(self, generic_guardrail, mock_request_data_input): """Test that images are passed to the API""" mock_response = MagicMock() mock_response.json.return_value = { @@ -738,9 +714,7 @@ class TestImageSupport: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: guardrailed_inputs = await generic_guardrail.apply_guardrail( inputs={ "texts": ["What's in this image?"], @@ -778,9 +752,7 @@ class TestApiKeyHeader: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=mock_response) as mock_post: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -814,9 +786,7 @@ class TestAdditionalParams: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=mock_response) as mock_post: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -826,23 +796,15 @@ class TestAdditionalParams: # Verify API was called with additional params call_args = mock_post.call_args json_payload = call_args.kwargs["json"] - assert ( - json_payload["additional_provider_specific_params"]["custom_threshold"] - == 0.8 - ) - assert ( - json_payload["additional_provider_specific_params"]["enable_feature"] - is True - ) + assert json_payload["additional_provider_specific_params"]["custom_threshold"] == 0.8 + assert json_payload["additional_provider_specific_params"]["enable_feature"] is True class TestModelParameter: """Test model parameter handling in guardrail requests""" @pytest.mark.asyncio - async def test_model_passed_from_inputs( - self, generic_guardrail, mock_request_data_input - ): + async def test_model_passed_from_inputs(self, generic_guardrail, mock_request_data_input): """Test that model is passed to the API when provided in inputs""" mock_response = MagicMock() mock_response.json.return_value = { @@ -851,9 +813,7 @@ class TestModelParameter: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"], "model": "gpt-4"}, request_data=mock_request_data_input, @@ -866,9 +826,7 @@ class TestModelParameter: assert json_payload["model"] == "gpt-4" @pytest.mark.asyncio - async def test_model_none_when_not_provided( - self, generic_guardrail, mock_request_data_input - ): + async def test_model_none_when_not_provided(self, generic_guardrail, mock_request_data_input): """Test that model is None when not provided in inputs""" mock_response = MagicMock() mock_response.json.return_value = { @@ -877,9 +835,7 @@ class TestModelParameter: } mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, # No model in inputs request_data=mock_request_data_input, @@ -896,18 +852,14 @@ class TestErrorHandling: """Test error handling scenarios""" @pytest.mark.asyncio - async def test_api_failure_handling( - self, generic_guardrail, mock_request_data_input - ): + async def test_api_failure_handling(self, generic_guardrail, mock_request_data_input): """Test API failure handling""" with patch.object( generic_guardrail.async_handler, "post", - side_effect=httpx.HTTPStatusError( - "API Error", request=MagicMock(), response=MagicMock(status_code=500) - ), + side_effect=httpx.HTTPStatusError("API Error", request=MagicMock(), response=MagicMock(status_code=500)), ): - with pytest.raises(Exception, match='Generic Guardrail API failed: API Error') as exc_info: + with pytest.raises(Exception, match="Generic Guardrail API failed: API Error") as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -917,16 +869,14 @@ class TestErrorHandling: assert "Generic Guardrail API failed" in str(exc_info.value) @pytest.mark.asyncio - async def test_network_error_handling( - self, generic_guardrail, mock_request_data_input - ): + async def test_network_error_handling(self, generic_guardrail, mock_request_data_input): """Test network error handling""" with patch.object( generic_guardrail.async_handler, "post", side_effect=httpx.RequestError("Connection failed", request=MagicMock()), ): - with pytest.raises(Exception, match='Generic Guardrail API failed: Connection failed') as exc_info: + with pytest.raises(Exception, match="Generic Guardrail API failed: Connection failed") as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -950,7 +900,7 @@ class TestErrorHandling: "post", side_effect=httpx.RequestError("Connection failed", request=MagicMock()), ): - with pytest.raises(Exception, match='Generic Guardrail API failed: Connection failed') as exc_info: + with pytest.raises(Exception, match="Generic Guardrail API failed: Connection failed") as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -1078,9 +1028,7 @@ class TestMultimodalSupport: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=mock_response) as mock_post: # This should not raise SerializationIterator error await guardrail.apply_guardrail( inputs={ @@ -1142,9 +1090,7 @@ class TestMultimodalSupport: } mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=mock_response) as mock_post: await guardrail.apply_guardrail( inputs={ "texts": ["Hello", "World"], @@ -1325,9 +1271,7 @@ class TestGenericGuardrailAPIStreamingConfig: guardrail_config = {"guardrail_name": "test-generic-streaming"} - with patch( - "litellm.logging_callback_manager.add_litellm_callback" - ): + with patch("litellm.logging_callback_manager.add_litellm_callback"): guardrail = initialize_guardrail(litellm_params, guardrail_config) assert guardrail.streaming_end_of_stream_only is False @@ -1398,9 +1342,7 @@ class TestGenericGuardrailAPIStreamingConfig: guardrail_config = {"guardrail_name": "test-generic-streaming-mixed"} - with patch( - "litellm.logging_callback_manager.add_litellm_callback" - ): + with patch("litellm.logging_callback_manager.add_litellm_callback"): guardrail = initialize_guardrail(litellm_params, guardrail_config) assert guardrail.streaming_end_of_stream_only is True @@ -1430,9 +1372,7 @@ class TestGenericGuardrailAPIStreamingConfig: guardrail_config = {"guardrail_name": "test-generic-streaming-nested-wins"} - with patch( - "litellm.logging_callback_manager.add_litellm_callback" - ): + with patch("litellm.logging_callback_manager.add_litellm_callback"): guardrail = initialize_guardrail(litellm_params, guardrail_config) assert guardrail.streaming_end_of_stream_only is True @@ -1461,9 +1401,7 @@ class TestGenericGuardrailAPIStreamingConfig: guardrail_config = {"guardrail_name": "test-generic-streaming-dict-optional"} - with patch( - "litellm.logging_callback_manager.add_litellm_callback" - ): + with patch("litellm.logging_callback_manager.add_litellm_callback"): guardrail = initialize_guardrail(litellm_params, guardrail_config) assert guardrail.streaming_end_of_stream_only is True @@ -1492,9 +1430,7 @@ class TestGenericGuardrailAPIStreamingConfig: guardrail_config = {"guardrail_name": "test-generic-streaming-dict-sibling"} - with patch( - "litellm.logging_callback_manager.add_litellm_callback" - ): + with patch("litellm.logging_callback_manager.add_litellm_callback"): guardrail = initialize_guardrail(litellm_params, guardrail_config) assert guardrail.streaming_end_of_stream_only is True @@ -1605,11 +1541,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: finish_reason="stop" if i == len(chunks_data) - 1 else None, ) - mock_post = AsyncMock( - return_value=_mock_guardrail_post_response( - action="NONE", texts=["Hello world! Goodbye"] - ) - ) + mock_post = AsyncMock(return_value=_mock_guardrail_post_response(action="NONE", texts=["Hello world! Goodbye"])) with ( patch.object(guardrail.async_handler, "post", mock_post), @@ -1618,9 +1550,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: return_value=_make_assembled_model_response("Hello world! Goodbye"), ), ): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, @@ -1662,9 +1592,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: ) mock_post = AsyncMock( - return_value=_mock_guardrail_post_response( - action="BLOCKED", blocked_reason="Ishaan is not allowed" - ) + return_value=_mock_guardrail_post_response(action="BLOCKED", blocked_reason="Ishaan is not allowed") ) with ( @@ -1674,9 +1602,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: return_value=_make_assembled_model_response("Hello ishaan here"), ), ): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, @@ -1717,11 +1643,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: finish_reason="stop" if i == len(chunks_data) - 1 else None, ) - mock_post = AsyncMock( - return_value=_mock_guardrail_post_response( - action="NONE", texts=["ABCDEFGHIJ"] - ) - ) + mock_post = AsyncMock(return_value=_mock_guardrail_post_response(action="NONE", texts=["ABCDEFGHIJ"])) with ( patch.object(guardrail.async_handler, "post", mock_post), @@ -1730,9 +1652,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: return_value=_make_assembled_model_response("ABCDEFGHIJ"), ), ): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, @@ -1776,11 +1696,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: finish_reason="stop" if i == len(chunks_data) - 1 else None, ) - mock_post = AsyncMock( - return_value=_mock_guardrail_post_response( - action="NONE", texts=["ABCDEFGHIJ"] - ) - ) + mock_post = AsyncMock(return_value=_mock_guardrail_post_response(action="NONE", texts=["ABCDEFGHIJ"])) with ( patch.object(guardrail.async_handler, "post", mock_post), @@ -1789,9 +1705,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: return_value=_make_assembled_model_response("ABCDEFGHIJ"), ), ): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, @@ -1806,8 +1720,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: pass assert mock_post.await_count == 1, ( - f"Expected exactly one guardrail call at end of stream, " - f"got {mock_post.await_count}" + f"Expected exactly one guardrail call at end of stream, got {mock_post.await_count}" ) @pytest.mark.asyncio @@ -1836,9 +1749,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: finish_reason="stop" if i == len(chunks_data) - 1 else None, ) - mock_post = AsyncMock( - return_value=_mock_guardrail_post_response(action="NONE", texts=["ABCDEF"]) - ) + mock_post = AsyncMock(return_value=_mock_guardrail_post_response(action="NONE", texts=["ABCDEF"])) with ( patch.object(guardrail.async_handler, "post", mock_post), @@ -1847,9 +1758,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: return_value=_make_assembled_model_response("ABCDEF"), ), ): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, @@ -1886,9 +1795,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: async def mock_stream(): for i, content in enumerate(["A", "B", "C"]): - yield _make_stream_chunk( - content, finish_reason="stop" if i == 2 else None - ) + yield _make_stream_chunk(content, finish_reason="stop" if i == 2 else None) mock_post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) @@ -1899,9 +1806,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: return_value=_make_assembled_model_response("ABC"), ), ): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, @@ -1937,16 +1842,10 @@ class TestGenericGuardrailAPIStreamingViaUnified: for event in _make_responses_stream_events("Hello world"): yield event - mock_post = AsyncMock( - return_value=_mock_guardrail_post_response( - action="NONE", texts=["Hello world"] - ) - ) + mock_post = AsyncMock(return_value=_mock_guardrail_post_response(action="NONE", texts=["Hello world"])) with patch.object(guardrail.async_handler, "post", mock_post): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/v1/responses" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses") request_data = { "input": "hi", "guardrail_to_apply": guardrail, @@ -1963,8 +1862,7 @@ class TestGenericGuardrailAPIStreamingViaUnified: assert events_received == 6 assert mock_post.await_count == 1, ( - f"Expected exactly one guardrail call at end of /v1/responses stream, " - f"got {mock_post.await_count}" + f"Expected exactly one guardrail call at end of /v1/responses stream, got {mock_post.await_count}" ) assert mock_post.await_args.kwargs["json"]["input_type"] == "response" @@ -1989,15 +1887,11 @@ class TestGenericGuardrailAPIStreamingViaUnified: yield event mock_post = AsyncMock( - return_value=_mock_guardrail_post_response( - action="BLOCKED", blocked_reason="Responses content not allowed" - ) + return_value=_mock_guardrail_post_response(action="BLOCKED", blocked_reason="Responses content not allowed") ) with patch.object(guardrail.async_handler, "post", mock_post): - user_api_key_dict = UserAPIKeyAuth( - api_key="test", request_route="/v1/responses" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses") request_data = { "input": "hi", "guardrail_to_apply": guardrail, @@ -2014,13 +1908,12 @@ class TestGenericGuardrailAPIStreamingViaUnified: assert "Responses content not allowed" in str(exc_info.value) + class TestToolSupport: """Test tool handling in guardrail requests""" @pytest.mark.asyncio - async def test_builtin_tools_without_function_block_do_not_crash( - self, generic_guardrail - ): + async def test_builtin_tools_without_function_block_do_not_crash(self, generic_guardrail): """Built-in tools (code_interpreter, file_search) have no `function` block. Regression for a 500 where serializing them raised a Pydantic @@ -2041,9 +1934,7 @@ class TestToolSupport: mock_response.json.return_value = {"action": "NONE", "texts": ["hi"]} mock_response.raise_for_status = MagicMock() - with patch.object( - generic_guardrail.async_handler, "post", return_value=mock_response - ) as mock_post: + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response) as mock_post: await generic_guardrail.apply_guardrail( inputs={"texts": ["hi"], "tools": tools}, request_data={}, @@ -2069,16 +1960,10 @@ class TestFailOnError: ) @pytest.mark.asyncio - async def test_endpoint_error_continues_when_fail_on_error_false( - self, fail_open_guardrail - ): + async def test_endpoint_error_continues_when_fail_on_error_false(self, fail_open_guardrail): """A non-unreachable endpoint error (HTTP 400) is swallowed and the request proceeds unchanged.""" - error = httpx.HTTPStatusError( - "bad request", request=MagicMock(), response=MagicMock(status_code=400) - ) - with patch.object( - fail_open_guardrail.async_handler, "post", side_effect=error - ): + error = httpx.HTTPStatusError("bad request", request=MagicMock(), response=MagicMock(status_code=400)) + with patch.object(fail_open_guardrail.async_handler, "post", side_effect=error): result = await fail_open_guardrail.apply_guardrail( inputs={"texts": ["hi"]}, request_data={}, @@ -2088,9 +1973,7 @@ class TestFailOnError: assert result == {"texts": ["hi"]} @pytest.mark.asyncio - async def test_internal_error_continues_without_calling_endpoint( - self, fail_open_guardrail - ): + async def test_internal_error_continues_without_calling_endpoint(self, fail_open_guardrail): """An error while building the request (here: invalid input_type) fails open too. Proves the request construction runs inside the protected block: the @@ -2107,9 +1990,7 @@ class TestFailOnError: assert result == {"texts": ["hi"]} @pytest.mark.asyncio - async def test_valid_block_still_blocks_when_fail_on_error_false( - self, fail_open_guardrail - ): + async def test_valid_block_still_blocks_when_fail_on_error_false(self, fail_open_guardrail): """Only a valid response acts: a BLOCKED decision still raises even with fail_on_error=False.""" mock_response = MagicMock() mock_response.json.return_value = { @@ -2118,9 +1999,7 @@ class TestFailOnError: } mock_response.raise_for_status = MagicMock() - with patch.object( - fail_open_guardrail.async_handler, "post", return_value=mock_response - ): + with patch.object(fail_open_guardrail.async_handler, "post", return_value=mock_response): with pytest.raises(GuardrailRaisedException): await fail_open_guardrail.apply_guardrail( inputs={"texts": ["hi"]}, @@ -2131,9 +2010,7 @@ class TestFailOnError: @pytest.mark.asyncio async def test_endpoint_error_raises_by_default(self, generic_guardrail): """Default fail_on_error=True keeps blocking on a non-unreachable endpoint error.""" - error = httpx.HTTPStatusError( - "bad request", request=MagicMock(), response=MagicMock(status_code=400) - ) + error = httpx.HTTPStatusError("bad request", request=MagicMock(), response=MagicMock(status_code=400)) with patch.object(generic_guardrail.async_handler, "post", side_effect=error): with pytest.raises(Exception, match="Generic Guardrail API failed"): await generic_guardrail.apply_guardrail( @@ -2143,16 +2020,10 @@ class TestFailOnError: ) @pytest.mark.asyncio - async def test_response_path_continues_when_fail_on_error_false( - self, fail_open_guardrail - ): + async def test_response_path_continues_when_fail_on_error_false(self, fail_open_guardrail): """fail_on_error governs the response path identically to the request path.""" - error = httpx.HTTPStatusError( - "bad request", request=MagicMock(), response=MagicMock(status_code=400) - ) - with patch.object( - fail_open_guardrail.async_handler, "post", side_effect=error - ): + error = httpx.HTTPStatusError("bad request", request=MagicMock(), response=MagicMock(status_code=400)) + with patch.object(fail_open_guardrail.async_handler, "post", side_effect=error): result = await fail_open_guardrail.apply_guardrail( inputs={"texts": ["model output"]}, request_data={}, @@ -2171,9 +2042,7 @@ class TestFailOnError: } mock_response.raise_for_status = MagicMock() - with patch.object( - fail_open_guardrail.async_handler, "post", return_value=mock_response - ): + with patch.object(fail_open_guardrail.async_handler, "post", return_value=mock_response): with pytest.raises(GuardrailRaisedException): await fail_open_guardrail.apply_guardrail( inputs={"texts": ["model output"]}, diff --git a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 88b28fffbd3..85b3baed01d 100644 --- a/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/unit/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -5,13 +5,15 @@ from collections.abc import Callable, Mapping, Sequence from itertools import chain from types import MappingProxyType from typing import Final -from unittest.mock import AsyncMock, MagicMock, call +from unittest.mock import AsyncMock, MagicMock +from unittest.mock import call as mock_call import pytest from fastapi import FastAPI, HTTPException from httpx import ASGITransport, AsyncClient from pytest_mock import MockerFixture +import litellm.proxy.management_endpoints.scim.scim_v2 as scim_v2_module from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_UserTable, @@ -24,6 +26,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( + GroupMemberExtractionResult, SCIMRosterSyncError, UserProvisionerHelpers, _apply_group_patch_updates, @@ -32,8 +35,8 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( _extract_ids_from_path_filter, _handle_group_membership_changes, _handle_team_membership_changes, + _member_ids_to_drop, _parse_member_entries, - premium_user_check, _process_group_patch_operations, _recompute_scim_member_roles, _resolve_group_member_ids, @@ -48,6 +51,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( patch_group, patch_team_membership, patch_user, + premium_user_check, scim_router, update_group, update_user, @@ -1616,6 +1620,7 @@ async def test_update_group_metadata_serialization_issue(mocker): # Mock team operations mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Mock user operations @@ -1624,7 +1629,7 @@ async def test_update_group_metadata_serialization_issue(mocker): mock_user.user_email = "user1@example.com" # Add proper string value for user_email mock_user.teams = [group_id] mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(mock_user,)) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock the _get_prisma_client_or_raise_exception to return our mock @@ -1805,6 +1810,7 @@ async def test_update_group_e2e(mocker): # Mock database operations mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock the updated team that gets returned from database updated_team = LiteLLM_TeamTable( @@ -1825,9 +1831,9 @@ async def test_update_group_e2e(mocker): # Mock user validation (all users exist) mock_user = mocker.MagicMock() - mock_user.user_id = "test-user" + mock_users: Final = tuple(LiteLLM_UserTable(user_id=user_id) for user_id in ("user2", "user3", "user4")) mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=mock_users) # Mock dependencies mocker.patch( @@ -1913,13 +1919,20 @@ async def test_update_group_e2e(mocker): def _rows_by_exact_id( user_row: Callable[[Mapping[str, str]], LiteLLM_UserTable | MagicMock | None], ) -> Callable[..., tuple[LiteLLM_UserTable | MagicMock, ...]]: - """``find_many`` stand-in for the classifier's cross-field read on a table where a - member value only ever matches as an exact ``user_id``.""" + """``find_many`` stand-in for a table where members match only exact user ids.""" def rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable | MagicMock, ...]: clauses: Final = where["OR"] assert isinstance(clauses, list) - found: Final = tuple(user_row(clause) for clause in clauses if "user_id" in clause) + fields: Final = tuple(next(iter(clause)) for clause in clauses) + assert fields == ("user_id", "sso_user_id", "user_email"), fields + user_filter: Final = clauses[0]["user_id"] + assert isinstance(user_filter, dict) + assert tuple(user_filter) == ("in",) + values: Final = user_filter["in"] + assert isinstance(values, list) + assert all(isinstance(value, str) for value in values) + found: Final = tuple(user_row({"user_id": value}) for value in values) return tuple(row for row in found if row is not None) return rows @@ -1970,6 +1983,7 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): # Mock team operations - team doesn't exist yet mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock user lookup - only existing-user exists def mock_user_lookup(where): @@ -2048,6 +2062,7 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_existing_team if where["team_id"] == group_id else None mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=mock_team_lookup) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock updated team response mock_updated_team = mocker.MagicMock() @@ -2125,6 +2140,7 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker # Mock team operations - team doesn't exist yet mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock user lookup - only existing-user exists initially def mock_user_lookup(where): @@ -2214,6 +2230,7 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock user lookup - only existing-user exists initially def mock_user_lookup(where): @@ -2283,6 +2300,7 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock user lookup - only existing-user exists def mock_user_lookup(where): @@ -2347,6 +2365,7 @@ async def test_process_group_patch_operations_with_flag_true_creates_users(mocke mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -2404,6 +2423,7 @@ async def test_process_group_patch_operations_with_flag_false_rejects(mocker, mo mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) # Execute the function - should raise HTTPException with pytest.raises(HTTPException) as exc_info: @@ -2933,10 +2953,13 @@ async def test_update_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock( + return_value=tuple(LiteLLM_UserTable(user_id=user_id) for user_id in ("user2", "user3")) + ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2991,10 +3014,13 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock( + return_value=(LiteLLM_UserTable(user_id="user1"), LiteLLM_UserTable(user_id="user2")) + ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3220,9 +3246,12 @@ async def test_create_group_recomputes_roles_for_members(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock( + return_value=(LiteLLM_UserTable(user_id="user1"), LiteLLM_UserTable(user_id="user2")) + ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3277,10 +3306,11 @@ async def test_update_group_rename_recomputes_retained_members(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(LiteLLM_UserTable(user_id="user1"),)) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3332,10 +3362,13 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock( + return_value=(LiteLLM_UserTable(user_id="user1"), LiteLLM_UserTable(user_id="user2")) + ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -4001,69 +4034,80 @@ def _member_resolution_prisma( else ({email: (user_id,) for email, user_id in email_to_user_id.items()} if email_to_user_id else {}) ) ssos_to_ids: Final[Mapping[str, str]] = dict(sso_user_id_to_user_id) if sso_user_id_to_user_id else {} + user_ids: Final = tuple( + dict.fromkeys(chain(users, chain.from_iterable(emails_to_ids.values()), ssos_to_ids.values())) + ) + user_rows: Final = tuple( + LiteLLM_UserTable( + user_id=user_id, + user_email=next((email for email, ids in emails_to_ids.items() if user_id in ids), None), + sso_user_id=next((subject for subject, mapped_id in ssos_to_ids.items() if mapped_id == user_id), None), + ) + for user_id in user_ids + ) def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]: - """Stand-in for the cross-field lookup, honouring the comparison mode - production actually asks for per field, so a field that stops folding case, or - starts folding it, fails here instead of passing. - - A caller that must know which accounts match rather than merely how many - passes take=None, so an unbounded read returns every match. The row keyed by - the value comes last, the order a bounded read is least prepared for, since - the database promises no order at all. - """ + """Stand-in for the cross-field lookup with the production comparison modes.""" clauses: Final = where["OR"] assert isinstance(clauses, list) fields: Final = tuple(next(iter(clause)) for clause in clauses) assert fields in (("user_id", "sso_user_id", "user_email"), ("sso_user_id", "user_email")), fields - def comparison(clause: Mapping[str, object]) -> tuple[str, bool]: - """The needle and whether production asked for a case-insensitive compare, - read per field so a field that stops folding case fails here.""" - criterion = next(iter(clause.values())) + def values_for(clause: Mapping[str, object], field: str, *, insensitive: bool) -> tuple[str, ...]: + criterion: Final = clause[field] if isinstance(criterion, str): - return criterion, False + assert not insensitive + return (criterion,) assert isinstance(criterion, dict), criterion - return criterion["equals"], criterion.get("mode") == "insensitive" + if "in" in criterion: + expected_keys: Final = ("in", "mode") if insensitive else ("in",) + assert tuple(criterion) == expected_keys + if insensitive: + assert criterion["mode"] == "insensitive" + values: Final = criterion["in"] + assert isinstance(values, list) + assert all(isinstance(value, str) for value in values) + return tuple(values) + assert insensitive + assert tuple(criterion) == ("equals", "mode") + assert criterion["mode"] == "insensitive" + value: Final = criterion["equals"] + assert isinstance(value, str) + return (value,) - by_field: Final = dict(zip(fields, (comparison(clause) for clause in clauses))) - sso_needle, sso_insensitive = by_field["sso_user_id"] - email_needle, email_insensitive = by_field["user_email"] - - def same(stored: str, needle: str, insensitive: bool) -> bool: - return stored.casefold() == needle.casefold() if insensitive else stored == needle - - matched: Final = tuple( - chain( - ( - user_id - for sso_user_id, user_id in ssos_to_ids.items() - if same(sso_user_id, sso_needle, sso_insensitive) - ), - ( - user_id - for email, user_ids in emails_to_ids.items() - if same(email, email_needle, email_insensitive) - for user_id in user_ids - ), - ( - user_id - for user_id in users - if "user_id" in by_field and same(user_id, by_field["user_id"][0], by_field["user_id"][1]) - ), - ) + by_field: Final = dict(zip(fields, clauses)) + user_id_values: Final = ( + values_for(by_field["user_id"], "user_id", insensitive=False) if "user_id" in by_field else () ) - found: Final = tuple(dict.fromkeys(matched)) - return tuple(LiteLLM_UserTable(user_id=user_id) for user_id in (found[:take] if take else found)) + sso_values: Final = values_for(by_field["sso_user_id"], "sso_user_id", insensitive=False) + email_values: Final = values_for(by_field["user_email"], "user_email", insensitive=True) + matched: Final = tuple( + row + for row in user_rows + if row.user_id in user_id_values + or row.sso_user_id in sso_values + or (row.user_email or "").casefold() in {value.casefold() for value in email_values} + ) + return matched[:take] if take else matched def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None: team_id: Final = where["team_id"] return team_row(team_id) + async def team_rows(where: Mapping[str, object]) -> tuple[LiteLLM_TeamTable, ...]: + criterion: Final = where["team_id"] + assert isinstance(criterion, dict) + assert tuple(criterion) == ("in",) + team_ids: Final = criterion["in"] + assert isinstance(team_ids, list) + assert all(isinstance(team_id, str) for team_id in team_ids) + return tuple(team for team_id in team_ids if (team := team_row(team_id)) is not None) + prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows) prisma_client.db.litellm_teamtable = mocker.MagicMock() prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) + prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=team_rows) return prisma_client @@ -4620,8 +4664,8 @@ async def test_process_group_patch_team_match_needs_scim_provenance( members_with_roles=[], ) prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set()) - prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=LiteLLM_TeamTable(team_id="child-team", metadata=team_metadata) + prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=(LiteLLM_TeamTable(team_id="child-team", metadata=team_metadata),) ) create_user_mock = mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", @@ -4895,17 +4939,17 @@ async def test_resolve_group_member_ids_dedupes_repeated_member(mocker, scim_ups assert result.all_member_ids == ["dup-user"] -def _identity_lookup(value: str) -> object: - """The single cross-field lookup the classifier is expected to issue per member.""" - return call( +def _identity_lookup(*values: str) -> object: + unique_values: Final = tuple(dict.fromkeys(values)) + subjects: Final = list(dict.fromkeys(value.strip() for value in unique_values)) + return mock_call( where={ "OR": [ - {"user_id": value}, - {"sso_user_id": value}, - {"user_email": {"equals": value, "mode": "insensitive"}}, + {"user_id": {"in": list(unique_values)}}, + {"sso_user_id": {"in": subjects}}, + {"user_email": {"in": subjects, "mode": "insensitive"}}, ] - }, - take=2, + } ) @@ -4925,7 +4969,7 @@ async def test_resolve_group_member_ids_matches_sso_user_id(mocker, scim_upsert_ ) result = await _resolve_group_member_ids( - members=[SCIMMember(value="member-sub")], + members=[SCIMMember(value=" member-sub ")], created_via="scim_group_membership", prisma_client=prisma_client, ) @@ -4934,7 +4978,7 @@ async def test_resolve_group_member_ids_matches_sso_user_id(mocker, scim_upsert_ assert result.existing_member_ids == ["sso-user"] assert result.created_users == [] assert result.all_member_ids == ["sso-user"] - assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("member-sub")] + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup(" member-sub ")] @pytest.mark.asyncio @@ -5440,13 +5484,10 @@ async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_acc @pytest.mark.asyncio -async def test_resolve_group_member_ids_reads_the_exact_id_when_two_other_accounts_fill_the_lookup( +async def test_resolve_group_member_ids_detects_all_identity_collisions_with_an_exact_id( mocker, scim_upsert_user_enabled ): - """A value that is one account's id and two other accounts' identities fills the - bounded lookup with the other two. The account keyed by the value must still be - found, or the id would lose its precedence and a non-canonical type would skip - a member that names a real user.""" + """A value that names three accounts remains ambiguous, even when one account's id matches exactly.""" prisma_client = _member_resolution_prisma( mocker, users={"shared"}, @@ -5470,15 +5511,11 @@ async def test_resolve_group_member_ids_reads_the_exact_id_when_two_other_accoun assert "shared" in str(exc_info.value.detail) create_user_mock.assert_not_called() assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("shared")] - prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "shared"}) + prisma_client.db.litellm_usertable.find_unique.assert_not_awaited() @pytest.mark.asyncio -async def test_resolve_group_member_ids_reads_the_user_table_once_per_member(mocker, scim_upsert_user_enabled): - """Every member costs one read of the user table, however it resolves: by its exact - id (which still outranks a non-canonical type), by identity, as a SCIM team, or not - at all. Looking the exact id up on its own before the identity read doubled the - reads of a push, and the identity read is a scan.""" +async def test_resolve_group_member_ids_reads_the_user_table_once_per_push(mocker, scim_upsert_user_enabled): prisma_client = _member_resolution_prisma( mocker, users={"by-id"}, @@ -5503,11 +5540,122 @@ async def test_resolve_group_member_ids_reads_the_user_table_once_per_member(moc assert result.all_member_ids == ["by-id", "email-user", "nobody"] prisma_client.db.litellm_usertable.find_unique.assert_not_awaited() + prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited() assert prisma_client.db.litellm_usertable.find_many.await_args_list == [ - _identity_lookup("by-id"), - _identity_lookup("by-email@example.com"), - _identity_lookup("by-team"), - _identity_lookup("nobody"), + _identity_lookup("by-id", "by-email@example.com", "by-team", "nobody") + ] + assert prisma_client.db.litellm_teamtable.find_many.await_args_list == [ + mock_call(where={"team_id": {"in": ["by-email@example.com", "by-team", "nobody"]}}) + ] + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_chunks_member_lookups_by_batch_size( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, scim_upsert_user_enabled: None +): + members: Final = ( + SCIMMember(value="member-id", type="direct"), + SCIMMember(value="member@example.com"), + SCIMMember(value="member-sub"), + SCIMMember(value="unknown-one"), + SCIMMember(value="unknown-two"), + SCIMMember(value="member-id", type="direct"), + ) + + async def resolve() -> tuple[GroupMemberExtractionResult, MagicMock]: + prisma_client: Final = _member_resolution_prisma( + mocker, + users={"member-id"}, + teams=set(), + email_to_user_id={"member@example.com": "email-user"}, + sso_user_id_to_user_id={"member-sub": "sso-user"}, + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock( + side_effect=[ + NewUserResponse(user_id="unknown-one", key="key-one"), + NewUserResponse(user_id="unknown-two", key="key-two"), + ] + ), + ) + result: Final = await _resolve_group_member_ids( + members=members, + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + return result, prisma_client + + default_result, _ = await resolve() + with monkeypatch.context() as context: + context.setattr(scim_v2_module, "SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE", 2) + chunked_result, prisma_client = await resolve() + + assert ( + chunked_result.all_member_ids + == default_result.all_member_ids + == [ + "member-id", + "email-user", + "sso-user", + "unknown-one", + "unknown-two", + ] + ) + assert ( + chunked_result.existing_member_ids + == default_result.existing_member_ids + == [ + "member-id", + "email-user", + "sso-user", + "member-id", + ] + ) + assert ( + chunked_result.created_users + == default_result.created_users + == [ + NewUserResponse(user_id="unknown-one", key="key-one"), + NewUserResponse(user_id="unknown-two", key="key-two"), + ] + ) + user_calls: Final = prisma_client.db.litellm_usertable.find_many.await_args_list + assert len(user_calls) == 3 + user_id_chunks: Final = tuple(tuple(entry.kwargs["where"]["OR"][0]["user_id"]["in"]) for entry in user_calls) + assert all(len(chunk) <= 2 for chunk in user_id_chunks) + assert tuple(chain.from_iterable(user_id_chunks)) == ( + "member-id", + "member@example.com", + "member-sub", + "unknown-one", + "unknown-two", + ) + + +@pytest.mark.asyncio +async def test_member_ids_to_drop_batches_user_id_email_and_sso_resolution(mocker: MockerFixture): + prisma_client: Final = _member_resolution_prisma( + mocker, + users={"user-id"}, + teams=set(), + email_to_user_id={"member@example.com": "email-user"}, + sso_user_id_to_user_id={"member-sub": "sso-user"}, + ) + + dropped: Final = await _member_ids_to_drop( + members=[ + SCIMMember(value="user-id"), + SCIMMember(value="MEMBER@example.com"), + SCIMMember(value="member-sub"), + ], + roster=frozenset({"user-id", "email-user", "sso-user"}), + prisma_client=prisma_client, + ) + + assert dropped == frozenset({"user-id", "email-user", "sso-user"}) + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [ + _identity_lookup("MEMBER@example.com", "member-sub", "user-id") ] @@ -5841,10 +5989,13 @@ async def test_update_group_roster_failure_propagates(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock( + return_value=(LiteLLM_UserTable(user_id="user1"), LiteLLM_UserTable(user_id="user2")) + ) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -6007,18 +6158,44 @@ def _shadow_tenant_prisma( async def find_unique(where: Mapping[str, str]) -> LiteLLM_UserTable | None: return users.get(where["user_id"]) + def values_for(clause: Mapping[str, object], field: str, *, insensitive: bool) -> tuple[str, ...]: + criterion: Final = clause[field] + if isinstance(criterion, str): + assert not insensitive + return (criterion,) + assert isinstance(criterion, dict), criterion + if "in" in criterion: + expected_keys: Final = ("in", "mode") if insensitive else ("in",) + assert tuple(criterion) == expected_keys + if insensitive: + assert criterion["mode"] == "insensitive" + values: Final = criterion["in"] + assert isinstance(values, list) + assert all(isinstance(value, str) for value in values) + return tuple(values) + assert insensitive + assert tuple(criterion) == ("equals", "mode") + assert criterion["mode"] == "insensitive" + value: Final = criterion["equals"] + assert isinstance(value, str) + return (value,) + def clause_matches(row: LiteLLM_UserTable, clause: Mapping[str, object]) -> bool: - if "user_id" in clause: - return row.user_id == clause["user_id"] - if "sso_user_id" in clause: - return row.sso_user_id == clause["sso_user_id"] - email_filter: Final = clause["user_email"] - assert isinstance(email_filter, dict) - return (row.user_email or "").casefold() == str(email_filter["equals"]).casefold() + field: Final = next(iter(clause)) + if field == "user_id": + return row.user_id in values_for(clause, field, insensitive=False) + if field == "sso_user_id": + return row.sso_user_id in values_for(clause, field, insensitive=False) + email_values: Final = frozenset( + value.casefold() for value in values_for(clause, "user_email", insensitive=True) + ) + return (row.user_email or "").casefold() in email_values async def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]: clauses: Final = where["OR"] assert isinstance(clauses, list) + fields: Final = tuple(next(iter(clause)) for clause in clauses) + assert fields in (("user_id", "sso_user_id", "user_email"), ("sso_user_id", "user_email")), fields matched: Final = tuple(row for row in users.values() if any(clause_matches(row, clause) for clause in clauses)) return matched[:take] if take else matched @@ -6031,6 +6208,15 @@ def _shadow_tenant_prisma( async def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None: return team if where["team_id"] == team.team_id else None + async def team_rows(where: Mapping[str, object]) -> tuple[LiteLLM_TeamTable, ...]: + criterion: Final = where["team_id"] + assert isinstance(criterion, dict) + assert tuple(criterion) == ("in",) + team_ids: Final = criterion["in"] + assert isinstance(team_ids, list) + assert all(isinstance(team_id, str) for team_id in team_ids) + return (team,) if team.team_id in team_ids else () + prisma_client = mocker.MagicMock() prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() @@ -6039,6 +6225,7 @@ def _shadow_tenant_prisma( prisma_client.db.litellm_usertable.delete = AsyncMock(side_effect=delete) prisma_client.db.litellm_teamtable = mocker.MagicMock() prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) + prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=team_rows) prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team) prisma_client.db.litellm_verificationtoken = mocker.MagicMock() prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=keys_for) @@ -6179,7 +6366,9 @@ class _PatchedTeamRow: async def update(self, *, where: dict[str, object], data: dict[str, object]) -> LiteLLM_TeamTable: self.written = data - self.team = LiteLLM_TeamTable(**{**self.team.model_dump(), **data, "metadata": json.loads(str(data["metadata"]))}) + self.team = LiteLLM_TeamTable( + **{**self.team.model_dump(), **data, "metadata": json.loads(str(data["metadata"]))} + ) return self.team diff --git a/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py b/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py index 13e57408bd0..f06008bb0ef 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py @@ -29,6 +29,7 @@ def _make_access_group_record( access_model_names: list | None = None, access_mcp_server_ids: list | None = None, access_agent_ids: list | None = None, + access_skill_names: list | None = None, assigned_team_ids: list | None = None, assigned_key_ids: list | None = None, created_by: str | None = "admin-user", @@ -44,6 +45,7 @@ def _make_access_group_record( "access_model_names": access_model_names or [], "access_mcp_server_ids": access_mcp_server_ids or [], "access_agent_ids": access_agent_ids or [], + "access_skill_names": access_skill_names or [], "assigned_team_ids": assigned_team_ids or [], "assigned_key_ids": assigned_key_ids or [], "created_at": created_at_val, @@ -89,6 +91,7 @@ def client_and_mocks(monkeypatch): access_model_names=data.get("access_model_names", []), access_mcp_server_ids=data.get("access_mcp_server_ids", []), access_agent_ids=data.get("access_agent_ids", []), + access_skill_names=data.get("access_skill_names", []), assigned_team_ids=data.get("assigned_team_ids", []), assigned_key_ids=data.get("assigned_key_ids", []), created_by=data.get("created_by"), @@ -106,6 +109,7 @@ def client_and_mocks(monkeypatch): access_model_names=data.get("access_model_names", []), access_mcp_server_ids=data.get("access_mcp_server_ids", []), access_agent_ids=data.get("access_agent_ids", []), + access_skill_names=data.get("access_skill_names", []), assigned_team_ids=data.get("assigned_team_ids", []), assigned_key_ids=data.get("assigned_key_ids", []), updated_by=data.get("updated_by"), @@ -1602,3 +1606,131 @@ def test_update_access_group_response_carries_resolved_names(client_and_mocks): assert body["access_mcp_servers"] == [{"id": "mcp-new", "name": "Linear"}] assert body["access_agents"] == [{"id": "agent-a", "name": "support-bot"}] assert body["access_mcp_server_ids"] == ["mcp-new"] + + +def test_create_access_group_persists_mcp_tool_permissions_as_json(client_and_mocks): + """Tool allowlists are stored on the group as JSON, an explicit [] survives as a deny-all entry.""" + client, _, mock_table, *_ = client_and_mocks + tool_permissions = {"mcp-1": ["read_file"], "mcp-2": []} + + resp = client.post( + "/v1/access_group", + json={ + "access_group_name": "tools", + "access_mcp_server_ids": ["mcp-1", "mcp-2"], + "mcp_tool_permissions": tool_permissions, + }, + ) + + assert resp.status_code == 201, resp.text + written = mock_table.create.await_args.kwargs["data"] + assert written["mcp_tool_permissions"].data == tool_permissions + + +def test_create_access_group_without_tool_permissions_writes_null(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + + resp = client.post("/v1/access_group", json={"access_group_name": "servers-only", "access_mcp_server_ids": ["m"]}) + + assert resp.status_code == 201, resp.text + assert mock_table.create.await_args.kwargs["data"]["mcp_tool_permissions"].data is None + + +@pytest.mark.parametrize("tool_permissions", [{"mcp-1": ["a", "b"]}, {"mcp-1": []}, None]) +def test_update_access_group_writes_mcp_tool_permissions(client_and_mocks, tool_permissions): + client, _, mock_table, *_ = client_and_mocks + mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-update")) + + resp = client.put("/v1/access_group/ag-update", json={"mcp_tool_permissions": tool_permissions}) + + assert resp.status_code == 200, resp.text + written = mock_table.update.await_args.kwargs["data"]["mcp_tool_permissions"] + assert written.data == tool_permissions + + +def test_update_access_group_leaves_tool_permissions_alone_when_not_sent(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-update")) + + resp = client.put("/v1/access_group/ag-update", json={"description": "x"}) + + assert resp.status_code == 200, resp.text + assert "mcp_tool_permissions" not in mock_table.update.await_args.kwargs["data"] + + +def test_get_access_group_returns_mcp_tool_permissions(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + record = _make_access_group_record(access_group_id="ag-tools", access_mcp_server_ids=["mcp-1"]) + record.dict = lambda: {**record.model_dump(), "mcp_tool_permissions": {"mcp-1": []}} + mock_table.find_unique = AsyncMock(return_value=record) + + resp = client.get("/v1/access_group/ag-tools") + + assert resp.status_code == 200, resp.text + assert resp.json()["mcp_tool_permissions"] == {"mcp-1": []} + + +def test_create_access_group_persists_and_returns_skill_names(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "skills-group", "access_skill_names": ["private-plugin", "other-plugin"]}, + ) + + assert resp.status_code == 201 + assert resp.json()["access_skill_names"] == ["private-plugin", "other-plugin"] + assert mock_table.create.await_args.kwargs["data"]["access_skill_names"] == ["private-plugin", "other-plugin"] + + +def test_create_access_group_without_skills_stores_empty_list(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + + resp = client.post("/v1/access_group", json={"access_group_name": "plain-group"}) + + assert resp.status_code == 201 + assert resp.json()["access_skill_names"] == [] + assert mock_table.create.await_args.kwargs["data"]["access_skill_names"] == [] + + +@pytest.mark.parametrize( + ("update_payload", "expected_skills"), + [ + ({"access_skill_names": ["private-plugin"]}, ["private-plugin"]), + ({"access_skill_names": []}, []), + ({"access_skill_names": None}, []), + ], +) +def test_update_access_group_writes_skill_names(client_and_mocks, update_payload, expected_skills): + client, _, mock_table, *_ = client_and_mocks + mock_table.find_unique = AsyncMock( + return_value=_make_access_group_record(access_group_id="ag-skills", access_skill_names=["old-plugin"]) + ) + + resp = client.put("/v1/access_group/ag-skills", json=update_payload) + + assert resp.status_code == 200 + assert resp.json()["access_skill_names"] == expected_skills + assert mock_table.update.await_args.kwargs["data"]["access_skill_names"] == expected_skills + + +def test_update_access_group_without_skill_field_leaves_skills_untouched(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + mock_table.find_unique = AsyncMock( + return_value=_make_access_group_record(access_group_id="ag-skills", access_skill_names=["old-plugin"]) + ) + + resp = client.put("/v1/access_group/ag-skills", json={"description": "renamed"}) + + assert resp.status_code == 200 + assert "access_skill_names" not in mock_table.update.await_args.kwargs["data"] + + +def test_get_and_list_access_group_return_skill_names(client_and_mocks): + client, _, mock_table, *_ = client_and_mocks + record = _make_access_group_record(access_group_id="ag-skills", access_skill_names=["private-plugin"]) + mock_table.find_unique = AsyncMock(return_value=record) + mock_table.find_many = AsyncMock(return_value=[record]) + + assert client.get("/v1/access_group/ag-skills").json()["access_skill_names"] == ["private-plugin"] + assert client.get("/v1/access_group").json()[0]["access_skill_names"] == ["private-plugin"] diff --git a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py index 0c9ed9356f8..0ffb3694942 100644 --- a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -19,6 +19,7 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( TeamAdminFieldNotPermitted, TeamAdminKeyEditAllowed, TeamAdminMemberKeyEditingDisabled, + TeamMemberDefaults, changed_key_fields, changed_team_fields, resolve_team_admin_editable_fields, @@ -160,8 +161,8 @@ class TestChangedTeamFields: assert changed_team_fields(UpdateTeamRequest(team_id="team-1", model_aliases={}), _team()) == frozenset() def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self): - data = UpdateTeamRequest(team_id="team-1", team_member_budget=10.0) - assert changed_team_fields(data, _team()) == frozenset({"team_member_budget"}) + data = UpdateTeamRequest(team_id="team-1", team_member_key_duration="30d") + assert changed_team_fields(data, _team()) == frozenset({"team_member_key_duration"}) class TestTeamAdminEditVerdict: @@ -192,27 +193,105 @@ class TestTeamAdminEditVerdict: assert verdict == TeamAdminFieldNotPermitted(field="blocked") -def _enabled(*fields: str) -> frozenset[str]: +def _permitted(*fields: str) -> frozenset[str]: return resolve_team_admin_editable_fields( {"team_admin_editable_team_fields": list(fields)}, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS ) +_STORED_MEMBER_DEFAULTS: Final = TeamMemberDefaults( + team_member_budget=10.0, team_member_budget_duration="30d", team_member_tpm_limit=100, team_member_rpm_limit=5 +) + +_MEMBER_DEFAULT_EDITS: Final = ( + pytest.param("team_member_budget", 25.0, 10.0, id="team_member_budget"), + pytest.param("team_member_budget_duration", "7d", "30d", id="team_member_budget_duration"), + pytest.param("team_member_tpm_limit", 200, 100, id="team_member_tpm_limit"), + pytest.param("team_member_rpm_limit", 9, 5, id="team_member_rpm_limit"), +) + + +class TestTeamMemberDefaultEdits: + @pytest.mark.parametrize("field,new,_stored", _MEMBER_DEFAULT_EDITS) + def test_enabled_field_change_is_allowed(self, field: str, new: object, _stored: object): + data: Final = UpdateTeamRequest.model_validate({"team_id": "team-1", field: new}) + verdict: Final = team_admin_edit_verdict(data, _team(), _permitted(field), _STORED_MEMBER_DEFAULTS) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", field: new} + + @pytest.mark.parametrize("field,new,_stored", _MEMBER_DEFAULT_EDITS) + def test_change_without_the_field_enabled_names_it(self, field: str, new: object, _stored: object): + data: Final = UpdateTeamRequest.model_validate({"team_id": "team-1", field: new}) + verdict: Final = team_admin_edit_verdict(data, _team(), _permitted("tpm_limit"), _STORED_MEMBER_DEFAULTS) + assert verdict == TeamAdminFieldNotPermitted(field=field) + + @pytest.mark.parametrize("field,_new,stored", _MEMBER_DEFAULT_EDITS) + def test_resending_the_stored_value_is_dropped_without_the_field_enabled( + self, field: str, _new: object, stored: object + ): + data: Final = UpdateTeamRequest.model_validate({"team_id": "team-1", field: stored, "tpm_limit": 6}) + verdict: Final = team_admin_edit_verdict( + data, _team(tpm_limit=5), _permitted("tpm_limit"), _STORED_MEMBER_DEFAULTS + ) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6} + + def test_nulls_resent_for_a_team_without_member_defaults_are_dropped(self): + data: Final = UpdateTeamRequest( + team_id="team-1", + team_member_budget=None, + team_member_budget_duration=None, + team_member_tpm_limit=None, + team_member_rpm_limit=None, + tpm_limit=6, + ) + verdict: Final = team_admin_edit_verdict(data, _team(tpm_limit=5), _permitted("tpm_limit"), None) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6} + + def test_first_member_budget_keeps_an_explicit_never_resets_duration(self): + data: Final = UpdateTeamRequest(team_id="team-1", team_member_budget=5.0, team_member_budget_duration=None) + permitted: Final = _permitted("team_member_budget", "team_member_budget_duration") + verdict: Final = team_admin_edit_verdict(data, _team(budget_duration="7d"), permitted, None) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == { + "team_id": "team-1", + "team_member_budget": 5.0, + "team_member_budget_duration": None, + } + + def test_first_member_budget_drops_a_duration_the_team_admin_may_not_set(self): + data: Final = UpdateTeamRequest(team_id="team-1", team_member_budget=5.0, team_member_budget_duration=None) + verdict: Final = team_admin_edit_verdict( + data, _team(budget_duration="7d"), _permitted("team_member_budget"), None + ) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "team_member_budget": 5.0} + + def test_existing_member_budget_drops_a_resent_duration_next_to_a_changed_budget(self): + data: Final = UpdateTeamRequest(team_id="team-1", team_member_budget=5.0, team_member_budget_duration="7d") + stored: Final = TeamMemberDefaults(team_member_budget=1.0, team_member_budget_duration="7d") + permitted: Final = _permitted("team_member_budget", "team_member_budget_duration") + verdict: Final = team_admin_edit_verdict(data, _team(), permitted, stored) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "team_member_budget": 5.0} + + class TestSoftBudgetEdits: def test_enabled_soft_budget_change_is_allowed(self): data: Final = UpdateTeamRequest(team_id="team-1", soft_budget=20.0) - verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("soft_budget")) + verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _permitted("soft_budget")) assert isinstance(verdict, TeamAdminEditAllowed) assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "soft_budget": 20.0} def test_soft_budget_change_without_it_enabled_names_it(self): data: Final = UpdateTeamRequest(team_id="team-1", soft_budget=20.0) - verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("max_budget")) + verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _permitted("max_budget")) assert verdict == TeamAdminFieldNotPermitted(field="soft_budget") - def test_unchanged_soft_budget_resend_is_dropped_when_not_enabled(self): + def test_unchanged_soft_budget_resend_is_dropped_when_not_permitted(self): data: Final = UpdateTeamRequest(team_id="team-1", tpm_limit=6, soft_budget=10) - verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("tpm_limit")) + verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _permitted("tpm_limit")) assert isinstance(verdict, TeamAdminEditAllowed) assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6} diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index bbb83ec1248..2da9e8bd77c 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -11386,6 +11386,7 @@ class TestResolveTeamAccessGroupResources: row1.access_model_names = ["gpt-4", "claude-3"] row1.access_mcp_server_ids = ["mcp-1"] row1.access_agent_ids = [] + row1.access_skill_names = ["private-plugin", "shared-plugin"] row2 = MagicMock() row2.access_group_id = "ag-2" @@ -11393,6 +11394,7 @@ class TestResolveTeamAccessGroupResources: row2.access_model_names = ["claude-3", "gemini"] row2.access_mcp_server_ids = [] row2.access_agent_ids = ["agent-1"] + row2.access_skill_names = ["shared-plugin"] fake_prisma = MagicMock() fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( @@ -11413,14 +11415,44 @@ class TestResolveTeamAccessGroupResources: ] assert resolved.access_group_mcp_server_ids == ["mcp-1"] assert resolved.access_group_agent_ids == ["agent-1"] + assert resolved.access_group_skill_names == () + assert resolved.unresolved_access_group_ids == ("ag-missing",) assert [ - (d.access_group_id, d.access_group_name, d.models, d.mcp_server_ids, d.agent_ids) + (d.access_group_id, d.access_group_name, d.models, d.mcp_server_ids, d.agent_ids, d.skill_names) for d in (resolved.access_group_details or []) ] == [ - ("ag-1", "shared-models", ("gpt-4", "claude-3"), ("mcp-1",), ()), - ("ag-2", "extra-models", ("claude-3", "gemini"), (), ("agent-1",)), + ("ag-1", "shared-models", ("gpt-4", "claude-3"), ("mcp-1",), (), ("private-plugin", "shared-plugin")), + ("ag-2", "extra-models", ("claude-3", "gemini"), (), ("agent-1",), ("shared-plugin",)), ] + @pytest.mark.asyncio + async def test_skill_names_resolve_when_every_listed_group_exists(self): + """The flat skill list is the effective grant: names from every group when all rows + exist, nothing while any listed id has no row, since the skill check then clamps + the team's keys to public plugins.""" + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + row = MagicMock() + row.access_group_id = "ag-1" + row.access_group_name = "skills-group" + row.access_model_names = [] + row.access_mcp_server_ids = [] + row.access_agent_ids = [] + row.access_skill_names = ["private-plugin"] + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[row]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resolved = await _resolve_team_access_group_resources( + TeamInfoResponseObjectTeamTable(team_id="team-1", access_group_ids=["ag-1"]) + ) + + assert resolved.access_group_skill_names == ("private-plugin",) + assert resolved.unresolved_access_group_ids == () + @pytest.mark.asyncio async def test_no_access_groups_leaves_details_unset(self): from litellm.proxy._types import TeamInfoResponseObjectTeamTable @@ -16950,6 +16982,41 @@ async def test_update_team_team_admin_echoing_unpermitted_fields_unchanged_is_al assert prisma.db.litellm_teamtable.update.called +@pytest.mark.asyncio +async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an_edit( + disable_audit_logging_for_mocked_team, +): + import contextlib + from litellm.models.budget import LiteLLM_BudgetTable + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {"team_member_budget_id": "member-budget-1"}) + prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=LiteLLM_BudgetTable( + budget_id="member-budget-1", max_budget=10.0, budget_duration="30d", tpm_limit=100, rpm_limit=5 + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["tpm_limit"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global + ) + result = await update_team( + data=UpdateTeamRequest( + team_id="test_team_id", + team_member_budget=10.0, + team_member_budget_duration="30d", + team_member_tpm_limit=100, + team_member_rpm_limit=5, + tpm_limit=6, + ), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert result["data"].team_id == "test_team_id" + assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 6 + assert not prisma.db.litellm_budgettable.update.called + + @pytest.mark.asyncio async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enables_it( disable_audit_logging_for_mocked_team, diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 973d6bc9f7b..88c5e2d2c1b 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -16,6 +16,7 @@ from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) +from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE async def _drain_tasks(): @@ -2951,3 +2952,58 @@ def test_get_user_from_metadata(end_user_id): ) assert response == "test" + + +def _anthropic_sse_stream(stop_reason: str, stop_details: dict | None) -> list[str]: + events: Final = ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": ANTHROPIC_DOCS_REFUSAL_MESSAGE["id"], + "type": "message", + "role": "assistant", + "model": ANTHROPIC_DOCS_REFUSAL_MESSAGE["model"], + "content": [], + "usage": {"input_tokens": 412, "output_tokens": 0}, + }, + }, + ), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None, "stop_details": stop_details}, + "usage": {"output_tokens": 0}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + return [f"event: {name}\ndata: {json.dumps(data)}\n\n" for name, data in events] + + +def test_streamed_messages_refusal_stop_details_on_the_logged_response(mock_logging_obj): + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + logged: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response( + all_chunks=_anthropic_sse_stream("refusal", stop_details), + litellm_logging_obj=mock_logging_obj, + model="claude-sonnet-5-5", + ) + + assert logged.choices[0].finish_reason == "content_filter" + provider_fields: Final = logged.choices[0].message.provider_specific_fields + assert provider_fields["stop_reason"] == "refusal" + assert provider_fields["stop_details"] == stop_details + + +def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj): + logged: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response( + all_chunks=_anthropic_sse_stream("end_turn", None), + litellm_logging_obj=mock_logging_obj, + model="claude-sonnet-5-5", + ) + + assert logged.choices[0].finish_reason == "stop" + provider_fields: Final = logged.choices[0].message.provider_specific_fields or {} + assert "stop_details" not in provider_fields diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index ad86c3c5267..01400384842 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -23,7 +23,7 @@ from __future__ import annotations import asyncio from datetime import datetime -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -830,31 +830,25 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 -@pytest.mark.asyncio -async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( - monkeypatch, -): - """A failing pipeline must invalidate every pending counter so the next - request reseeds from the DB (which already holds this request's cost) - instead of trusting a value the write may have partially applied.""" - from redis.exceptions import MaxConnectionsError +_PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", +) + +def _warm_counters_behind_a_failing_pipeline(monkeypatch, pipeline_error: Exception) -> AsyncMock: counter_cache = ps.DualCache() - pending_keys = ( - "spend:key:hashed", - "spend:team:team-1", - "spend:team_member:user-1:team-1", - "spend:user:user-1", - "spend:end_user:eu-1", - "spend:tag:tag-a", - "spend:tag:tag-b", - "spend:org:org-1", - ) - for key in pending_keys: + for key in _PENDING_SPEND_COUNTER_KEYS: counter_cache.in_memory_cache.set_cache(key=key, value=1.0) fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm - fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError()) + fake_redis.async_increment_pipeline = AsyncMock(side_effect=pipeline_error) fake_redis.async_increment = AsyncMock() fake_redis.async_delete_cache = AsyncMock() fake_redis.get_ttl = MagicMock(return_value=None) @@ -862,23 +856,63 @@ async def test_increment_spend_counters_pipeline_failure_invalidates_all_counter monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) monkeypatch.setattr(ps, "prisma_client", None) + return fake_redis - with pytest.raises(MaxConnectionsError): - await ps.increment_spend_counters( - token="hashed", - team_id="team-1", - user_id="user-1", - response_cost=0.5, - org_id="org-1", - end_user_id="eu-1", - tags=["tag-a", "tag-b"], - ) +async def _increment_every_spend_counter() -> None: + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + +def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None: assert fake_redis.async_increment.await_count == 0 deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} - assert deleted_keys == set(pending_keys) - for key in pending_keys: - assert counter_cache.in_memory_cache.get_cache(key=key) is None + assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS) + for key in _PENDING_SPEND_COUNTER_KEYS: + assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( + monkeypatch, +): + """A failing pipeline must invalidate every pending counter so the next + request reseeds from the DB (which already holds this request's cost) + instead of trusting a value the write may have partially applied.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, RedisConnectionError("Redis is down")) + + with pytest.raises(RedisConnectionError): + await _increment_every_spend_counter() + + _assert_every_pending_counter_invalidated(fake_redis) + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pool_exhaustion_invalidates_all_counters_without_failing( + monkeypatch, +): + """A cluster node with every connection busy raises MaxConnectionsError at once. That is a busy pool, the + same as the standalone pool's wait timeout, so the counters are invalidated and the request still succeeds.""" + from redis import exceptions as redis_exceptions + + max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError + "type[Exception]", + redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs + ) + fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, max_connections_error()) + + await _increment_every_spend_counter() + + _assert_every_pending_counter_invalidated(fake_redis) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5dac1d25572..c1ff8205a04 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -43,27 +43,57 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( model: Final = "gpt-6-astra" message: Final = "Upstream cannot complete this response" code: Final = { - "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "429", - "server_error": "server_error", "response_failed": "server_error", "cyber_policy": "cyber_policy", + "rate_limit": "rate_limit_exceeded", + "numeric_rate_limit": "429", + "server_error": "server_error", + "response_failed": "server_error", + "cyber_policy": "cyber_policy", }[error_kind] error: Final = {"message": message, "code": code, "type": None, "param": "input"} - response: Final = {"id": "resp_upstream", "object": "response", "created_at": 1, - "status": "in_progress", "model": model, "output": [], - "parallel_tool_calls": True, "tool_choice": "auto", "tools": []} + response: Final = { + "id": "resp_upstream", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": model, + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } created: Final = {"type": "response.created", "sequence_number": 0, "response": response} - tool_added: Final = {"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, - "item": {"type": "function_call", "id": "fc_partial", "call_id": "call_partial", - "name": "read_file", "arguments": "", "status": "in_progress"}} - tool_delta: Final = {"type": "response.function_call_arguments.delta", "sequence_number": 2, - "item_id": "fc_partial", "output_index": 0, "delta": '{"path":"partial'} + tool_added: Final = { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": { + "type": "function_call", + "id": "fc_partial", + "call_id": "call_partial", + "name": "read_file", + "arguments": "", + "status": "in_progress", + }, + } + tool_delta: Final = { + "type": "response.function_call_arguments.delta", + "sequence_number": 2, + "item_id": "fc_partial", + "output_index": 0, + "delta": '{"path":"partial', + } failed: Final = ( - {"type": "response.failed", "sequence_number": 9, - "response": {**response, "status": "failed", "error": error}} - if error_kind in ("response_failed", "cyber_policy") else {"type": "error", "error": error} + {"type": "response.failed", "sequence_number": 9, "response": {**response, "status": "failed", "error": error}} + if error_kind in ("response_failed", "cyber_policy") + else {"type": "error", "error": error} ) - chat: Final = {"id": "chatcmpl_partial", "object": "chat.completion.chunk", "created": 1, - "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, - "finish_reason": None}]} + chat: Final = { + "id": "chatcmpl_partial", + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {"content": "partial"}, "finish_reason": None}], + } is_chat: Final = path == "/v1/chat/completions" partial: Final = path != "/v1/responses" or error_kind in ("numeric_rate_limit", "response_failed", "cyber_policy") response_events: Final = (created, tool_added, tool_delta, failed) if partial else (failed,) @@ -71,8 +101,12 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) upstream_url: Final = "https://streaming.example/v1" router: Final = litellm.Router( - model_list=[{"model_name": model, "litellm_params": { - "model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}}], + model_list=[ + { + "model_name": model, + "litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}, + } + ], num_retries=0, ) monkeypatch.setattr(ps, "llm_router", router) @@ -84,15 +118,18 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: result: Final = await client.post( - path, json={ - "model": model, "stream": True, + path, + json={ + "model": model, + "stream": True, **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), }, ) frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) events: Final = tuple( json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) - for frame in frames if "data: [DONE]" not in frame + for frame in frames + if "data: [DONE]" not in frame ) assert result.status_code == 200, result.text @@ -102,8 +139,10 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert frames[-2].startswith("event: response.failed\n"), result.text if partial: assert [event["type"] for event in events] == [ - "response.created", "response.output_item.added", - "response.function_call_arguments.delta", "response.failed", + "response.created", + "response.output_item.added", + "response.function_call_arguments.delta", + "response.failed", ] assert events[2]["delta"] == tool_delta["delta"] assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 @@ -113,10 +152,16 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert events[0]["sequence_number"] == 0 assert events[0]["response"]["id"].startswith("resp_") assert events[-1]["response"]["status"] == "failed" - assert events[-1]["response"]["error"]["code"] == { - "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "rate_limit_exceeded", - "server_error": "server_error", "response_failed": "server_error", "cyber_policy": "cyber_policy", - }[error_kind] + assert ( + events[-1]["response"]["error"]["code"] + == { + "rate_limit": "rate_limit_exceeded", + "numeric_rate_limit": "rate_limit_exceeded", + "server_error": "server_error", + "response_failed": "server_error", + "cyber_policy": "cyber_policy", + }[error_kind] + ) assert events[-1]["response"]["error"]["message"] == message else: assert events[0]["object"] == "chat.completion.chunk", result.text @@ -769,7 +814,9 @@ class TestResponsesWSFirstFrameModelAuth: replayed_input = [{"type": "reasoning", "id": "encitem_abc", "encrypted_content": "litellm_enc:abc;blob"}] payload = {"model": "gpt-4o-mini", "input": replayed_input, "previous_response_id": "resp_prev"} - first_frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload} + first_frame = ( + {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload} + ) raw_first_frame = json.dumps(first_frame) ws = MagicMock() @@ -1115,7 +1162,6 @@ class TestReadWSModelFromFirstFrameErrors: ws.close.assert_not_awaited() - @pytest.mark.parametrize( "configured,expected", [ @@ -2421,6 +2467,27 @@ class TestResponsesInputTokens: ] assert token_request.tools == tools + def test_item_cache_control_is_left_out_of_the_counted_messages(self): + response, counter = self._post_input_tokens( + { + "model": "anthropic/claude-sonnet-4-6", + "input": [ + {"role": "system", "content": "Long handbook", "cache_control": {"type": "ephemeral"}}, + { + "role": "user", + "content": "What does a reviewer do first?", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + ) + + assert response.status_code == 200, response.text + assert counter.call_args.kwargs["request"].messages == [ + {"role": "system", "content": "Long handbook"}, + {"role": "user", "content": "What does a reviewer do first?"}, + ] + def test_missing_model_returns_openai_400(self): response, counter = self._post_input_tokens({"input": "Hello"}) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 5563c2826bb..0280a6fcc0c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -45,7 +45,10 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): handler = LiteLLMCompletionTransformationHandler() async def fake_session_handler( - previous_response_id: str, litellm_completion_request: dict[str, object], instructions: str | None = None + previous_response_id: str, + litellm_completion_request: dict[str, object], + instructions: str | None = None, + request_kwargs: Mapping[str, object] | None = None, ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request @@ -305,9 +308,7 @@ async def test_previous_response_id_without_instructions_carries_the_latest_ones def _anthropic_sse(*events: Mapping[str, object]) -> bytes: - return b"".join( - f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events - ) + return b"".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) def test_streamed_responses_over_anthropic_emit_the_responses_event_sequence(respx_mock: respx.MockRouter) -> None: diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 7006feb22b2..6cdccf3958b 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1,9 +1,12 @@ +import asyncio import json +from collections.abc import Mapping, Sequence from copy import deepcopy -from typing import Final, Literal +from typing import Final, Literal, cast import pytest from openai.types.responses import ResponseFunctionToolCall +from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.response_function_web_search import ( ActionFind, ActionOpenPage, @@ -11,15 +14,22 @@ from openai.types.responses.response_function_web_search import ( ActionSearchSource, ResponseFunctionWebSearch, ) +from pydantic import TypeAdapter import litellm +from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR -from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt +from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + anthropic_messages_pt, +) from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, + ChatCompletionSession, LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.base import HiddenParams +from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -5162,6 +5172,766 @@ async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) +def test_cache_control_on_user_message(): + """cache_control on a user message should be passed through to GenericChatCompletionMessage""" + input_item = { + "role": "user", + "content": "Hello, world!", + "cache_control": {"type": "ephemeral"}, + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "user" + assert msg.get("cache_control") == {"type": "ephemeral"} + + +def test_cache_control_on_assistant_message(): + """cache_control on an assistant message should be passed through""" + input_item = { + "role": "assistant", + "content": "I can help with that.", + "cache_control": {"type": "ephemeral"}, + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "assistant" + assert msg.get("cache_control") == {"type": "ephemeral"} + + +def test_no_cache_control_when_absent(): + """Messages without cache_control should not have the field set""" + input_item = { + "role": "user", + "content": "Hello", + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(result) == 1 + msg = result[0] + assert "cache_control" not in msg + + +def test_cache_control_on_function_call_output(): + """cache_control on a function_call_output should be passed through to tool message""" + TOOL_CALLS_CACHE.set_cache( + key="call_123", + value={ + "id": "call_123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "NYC"}'}, + }, + ) + + input_item = { + "type": "function_call_output", + "call_id": "call_123", + "output": '{"temp": 72}', + "cache_control": {"type": "ephemeral"}, + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + # Should have assistant message + tool message + assert len(result) == 2 + tool_msg = result[1] + assert tool_msg["role"] == "tool" + assert tool_msg.get("cache_control") == {"type": "ephemeral"} + + +def test_cache_control_on_function_call(): + """cache_control on a function_call input should land on the tool call it builds""" + input_item = { + "type": "function_call", + "call_id": "call_456", + "name": "get_weather", + "arguments": '{"city": "NYC"}', + "cache_control": {"type": "ephemeral"}, + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(result) == 1 + msg = result[0] + assert msg["role"] == "assistant" + assert "cache_control" not in msg + assert msg["tool_calls"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_cache_control_preserved_in_full_input_transform(): + """cache_control should survive the full transform_responses_api_input_to_messages path""" + input_items = [ + { + "role": "user", + "content": "What is the weather?", + }, + { + "role": "assistant", + "content": "Let me check.", + "cache_control": {"type": "ephemeral"}, + }, + { + "role": "user", + "content": "Thanks!", + "cache_control": {"type": "ephemeral"}, + }, + ] + + result = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input_items, + responses_api_request={}, + item_cache_control=True, + ) + + # 3 messages, no system + assert len(result) == 3 + assert "cache_control" not in result[0] + assert result[1].get("cache_control") == {"type": "ephemeral"} + assert result[2].get("cache_control") == {"type": "ephemeral"} + + +def test_cache_control_propagated_to_last_content_block_for_list_content(): + """ + For Anthropic/Bedrock, an item-level cache_control over list content activates + prompt caching only as a mark on the last content block, so the bridge moves + it there instead of keeping a message-level copy the hook's census would also count. + """ + input_item = { + "role": "user", + "content": [ + {"type": "input_text", "text": "first block"}, + {"type": "input_text", "text": "last block"}, + ], + "cache_control": {"type": "ephemeral"}, + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(result) == 1 + msg = result[0] + assert "cache_control" not in msg + assert isinstance(msg["content"], list) + assert "cache_control" not in msg["content"][0] + assert msg["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result) == 1 + + +def test_cache_control_does_not_overwrite_existing_block_level_cache_control(): + """ + If the last content block already carries its own cache_control + (preserved by _transform_responses_api_content_to_chat_completion_content), + the message-level cache_control must not silently overwrite it. + """ + input_item = { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "last block", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + }, + ], + "cache_control": {"type": "ephemeral"}, + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=input_item + ) + + assert len(result) == 1 + msg = result[0] + assert "cache_control" not in msg + assert msg["content"][-1]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result) == 1 + + +def test_function_call_cache_control_reaches_the_provider_tool_use_block(): + cache_control: Final = {"type": "ephemeral"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + {"role": "user", "content": "Summarize the gateway latency report."}, + { + "type": "function_call", + "call_id": "call_cached_tool_use", + "name": "get_latency_report", + "arguments": '{"service": "gateway"}', + "cache_control": cache_control, + }, + {"type": "function_call_output", "call_id": "call_cached_tool_use", "output": "p50 120ms"}, + ], + responses_api_request={}, + item_cache_control=True, + ) + + anthropic_turns: Final = anthropic_messages_pt( + messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic" + ) + tool_use_block: Final = anthropic_turns[1]["content"][0] + assert tool_use_block["type"] == "tool_use" + assert tool_use_block["cache_control"] == cache_control + + bedrock_blocks: Final = _bedrock_converse_messages_pt( + messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock" + )[1]["content"] + tool_use_index: Final = next(index for index, block in enumerate(bedrock_blocks) if "toolUse" in block) + assert "cachePoint" in bedrock_blocks[tool_use_index + 1] + + +def test_assistant_text_merged_into_a_function_call_turn_keeps_its_cache_control(): + cache_control: Final = {"type": "ephemeral"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + {"role": "user", "content": "Summarize the gateway latency report."}, + { + "type": "function_call", + "call_id": "call_merged_text", + "name": "get_latency_report", + "arguments": '{"service": "gateway"}', + }, + {"role": "assistant", "content": "Looking that up now.", "cache_control": cache_control}, + {"type": "function_call_output", "call_id": "call_merged_text", "output": "p50 120ms"}, + ], + responses_api_request={}, + item_cache_control=True, + ) + + assert len(messages) == 3 + assert messages[1]["cache_control"] == cache_control + + anthropic_turns: Final = anthropic_messages_pt( + messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic" + ) + text_block: Final = anthropic_turns[1]["content"][0] + assert text_block["type"] == "text" + assert text_block["cache_control"] == cache_control + + bedrock_blocks: Final = _bedrock_converse_messages_pt( + messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock" + )[1]["content"] + assert any("cachePoint" in block for block in bedrock_blocks) + + +def test_assistant_text_merged_behind_earlier_text_puts_cache_control_on_the_last_block(): + cache_control: Final = {"type": "ephemeral"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + {"role": "user", "content": "Summarize the gateway latency report."}, + { + "type": "function_call", + "call_id": "call_two_texts", + "name": "get_latency_report", + "arguments": '{"service": "gateway"}', + }, + {"role": "assistant", "content": "Looking that up now."}, + {"role": "assistant", "content": "One moment.", "cache_control": cache_control}, + {"type": "function_call_output", "call_id": "call_two_texts", "output": "p50 120ms"}, + ], + responses_api_request={}, + item_cache_control=True, + ) + + merged_turn: Final = messages[1] + assert "cache_control" not in merged_turn + assert "cache_control" not in merged_turn["content"][0] + assert merged_turn["content"][-1]["cache_control"] == cache_control + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 1 + + +def test_marked_assistant_text_merged_ahead_of_later_text_keeps_its_cache_control_on_its_block(): + cache_control: Final = {"type": "ephemeral"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + {"role": "user", "content": "Summarize the gateway latency report."}, + { + "type": "function_call", + "call_id": "call_marked_first_text", + "name": "get_latency_report", + "arguments": '{"service": "gateway"}', + }, + {"role": "assistant", "content": "Looking that up now.", "cache_control": cache_control}, + {"role": "assistant", "content": "One moment."}, + {"type": "function_call_output", "call_id": "call_marked_first_text", "output": "p50 120ms"}, + ], + responses_api_request={}, + item_cache_control=True, + ) + + merged_turn: Final = messages[1] + assert "cache_control" not in merged_turn + assert merged_turn["content"][0]["cache_control"] == cache_control + assert "cache_control" not in merged_turn["content"][-1] + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 1 + + anthropic_turns: Final = anthropic_messages_pt( + messages=messages, model="claude-sonnet-4-6", llm_provider="anthropic" + ) + blocks: Final = TypeAdapter(list[dict[str, object]]).validate_python(anthropic_turns[1]["content"]) + marked_blocks: Final = [block for block in blocks if block.get("cache_control")] + assert len(marked_blocks) == 1 + assert marked_blocks[0]["text"] == "Looking that up now." + + +_CLAUDE_ON_BEDROCK: Final = "us.anthropic.claude-sonnet-4-6" +_EPHEMERAL: Final = {"type": "ephemeral"} + + +def _replayed_messages_from( + turns: Sequence[Mapping[str, object]], + current: Mapping[str, object], + *, + model: str = _CLAUDE_ON_BEDROCK, + tools: Sequence[Mapping[str, object]] = (), + **request_kwargs: object, +) -> Sequence[Mapping[str, object]]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + assert previous_response_id == "resp_marked_history" + return ChatCompletionSession(messages=list(turns), litellm_session_id=None) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_marked_history", + litellm_completion_request={ + "model": model, + "custom_llm_provider": "bedrock", + "messages": [current], + **({"tools": list(tools)} if tools else {}), + }, + instructions=None, + session_lookup=session_lookup, + request_kwargs=request_kwargs, + ) + ) + return TypeAdapter(list[dict[str, object]]).validate_python(request["messages"]) + + +def test_a_replayed_turn_over_the_remaining_budget_yields_its_marks_while_older_turns_keep_theirs(): + three_marked_blocks: Final = [ + {"type": "text", "text": f"block {index}", "cache_control": _EPHEMERAL} for index in range(3) + ] + replayed_turns: Final = [ + {"role": "user", "content": "turn 0", "cache_control": _EPHEMERAL}, + {"role": "user", "content": "turn 1", "cache_control": _EPHEMERAL}, + {"role": "user", "content": three_marked_blocks}, + {"role": "user", "content": "turn 3", "cache_control": _EPHEMERAL}, + ] + + messages: Final = _replayed_messages_from( + replayed_turns, {"role": "user", "content": "newest", "cache_control": _EPHEMERAL} + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == MAX_CACHE_CONTROL_BLOCKS + assert [AnthropicCacheControlHook._count_cache_control_blocks(message) for message in messages] == [1, 1, 0, 1, 1] + assert messages[2] == { + "role": "user", + "content": [{"type": "text", "text": f"block {index}"} for index in range(3)], + } + + +def test_replayed_session_turns_yield_their_cache_control_to_the_newest_markers(): + replayed_turns: Final = [ + {"role": "user", "content": f"turn {index}", "cache_control": _EPHEMERAL} for index in range(5) + ] + + messages: Final = _replayed_messages_from( + replayed_turns, {"role": "user", "content": "newest", "cache_control": _EPHEMERAL} + ) + assert ["cache_control" in message for message in messages] == [False, False, True, True, True, True] + assert messages[0] == {"role": "user", "content": "turn 0"} + assert messages[-1] == {"role": "user", "content": "newest", "cache_control": _EPHEMERAL} + + +def test_a_marked_tool_takes_a_slot_from_the_replayed_turns(): + replayed_turns: Final = [ + {"role": "user", "content": f"turn {index}", "cache_control": _EPHEMERAL} for index in range(3) + ] + marked_tool: Final = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + "cache_control": _EPHEMERAL, + } + + messages: Final = _replayed_messages_from( + replayed_turns, {"role": "user", "content": "newest", "cache_control": _EPHEMERAL}, tools=[marked_tool] + ) + + assert ["cache_control" in message for message in messages] == [False, True, True, True] + + +@pytest.mark.parametrize( + "request_kwargs", + [{"cache_control": _EPHEMERAL}, {"extra_body": {"cache_control": _EPHEMERAL}}], + ids=["top_level", "extra_body"], +) +def test_a_top_level_cache_control_takes_a_slot_from_the_replayed_turns(request_kwargs: Mapping[str, object]): + replayed_turns: Final = [ + {"role": "user", "content": f"turn {index}", "cache_control": _EPHEMERAL} for index in range(3) + ] + + messages: Final = _replayed_messages_from( + replayed_turns, {"role": "user", "content": "newest", "cache_control": _EPHEMERAL}, **request_kwargs + ) + + assert ["cache_control" in message for message in messages] == [False, True, True, True] + + +def test_a_null_cache_control_on_the_current_turn_takes_no_slot_from_the_replayed_turns(): + current: Final = {"role": "user", "content": [{"type": "text", "text": "newest", "cache_control": None}]} + + messages: Final = _replayed_messages_from(list(_marked_turns(4)), current) + + assert ["cache_control" in message for message in messages[:4]] == [True, True, True, True] + + +@pytest.mark.parametrize("drop_params", [False, True]) +def test_replayed_marks_a_provider_would_refuse_drop_with_or_without_drop_params(drop_params: bool): + replayed_turns: Final = [{"role": "user", "content": "turn 0", "cache_control": {"type": "forever"}}] + current: Final = {"role": "user", "content": "newest"} + + assert _replayed_messages_from(replayed_turns, current, drop_params=drop_params)[0] == { + "role": "user", + "content": "turn 0", + } + + +def test_replayed_tool_marks_drop_for_nova_without_drop_params(): + replayed_turns: Final = [ + {"role": "user", "content": "Summarize the gateway latency report."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_latency_report", "arguments": "{}"}, + "cache_control": _EPHEMERAL, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "p50 12 ms", "cache_control": _EPHEMERAL}, + ] + + messages: Final = _replayed_messages_from( + replayed_turns, {"role": "user", "content": "newest"}, model="us.amazon.nova-2-lite-v1:0" + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 0 + assert messages[2] == {"role": "tool", "tool_call_id": "call_1", "content": "p50 12 ms"} + + +def test_replayed_marks_stay_untouched_for_a_target_without_cache_points(): + replayed_turns: Final = [ + {"role": "user", "content": f"turn {index}", "cache_control": _EPHEMERAL} for index in range(5) + ] + + messages: Final = _replayed_messages_from( + replayed_turns, {"role": "user", "content": "newest"}, model="meta.llama3-70b-instruct-v1:0" + ) + + assert messages[:5] == replayed_turns + + +def _bridged_request( + *items: Mapping[str, object], + model: str = f"bedrock/{_CLAUDE_ON_BEDROCK}", + custom_llm_provider: str = "bedrock", + tools: Sequence[Mapping[str, object]] = (), + **kwargs: object, +) -> Sequence[Mapping[str, object]]: + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=cast(ResponseInputParam, list(items)), # cast-ok: cache_control is not a key of the input TypedDicts + responses_api_request=cast( # cast-ok: cache_control is not a key of the tool TypedDicts + ResponsesAPIOptionalRequestParams, {"instructions": "Be brief.", "tools": list(tools)} + ), + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + return TypeAdapter(list[dict[str, object]]).validate_python(request["messages"]) + + +def _marked_turns(count: int) -> tuple[Mapping[str, object], ...]: + return tuple({"role": "user", "content": f"turn {index}", "cache_control": _EPHEMERAL} for index in range(count)) + + +def test_marks_past_the_provider_cap_reach_it_as_sent_without_drop_params(): + messages: Final = _bridged_request(*_marked_turns(5)) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == MAX_CACHE_CONTROL_BLOCKS + 1 + + +@pytest.mark.parametrize("drop_params", [{"drop_params": True}, {"drop_params": "true"}], ids=["bool", "string"]) +def test_marks_past_the_provider_cap_keep_the_newest_under_drop_params(drop_params: Mapping[str, object]): + messages: Final = _bridged_request(*_marked_turns(5), **drop_params) + + assert ["cache_control" in message for message in messages] == [False, False, True, True, True, True] + assert messages[1] == {"role": "user", "content": "turn 0"} + + +def test_global_drop_params_keeps_the_newest_marks(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", True) + + messages: Final = _bridged_request(*_marked_turns(5)) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == MAX_CACHE_CONTROL_BLOCKS + + +def test_a_marked_tool_takes_a_slot_from_the_request_marks_under_drop_params(): + marked_tool: Final = { + "type": "function", + "name": "lookup", + "parameters": {"type": "object"}, + "cache_control": _EPHEMERAL, + } + + messages: Final = _bridged_request(*_marked_turns(4), tools=[marked_tool], drop_params=True) + + assert ["cache_control" in message for message in messages] == [False, False, True, True, True] + + +@pytest.mark.parametrize( + "request_kwargs", + [{"cache_control": _EPHEMERAL}, {"extra_body": {"cache_control": _EPHEMERAL}}], + ids=["top_level", "extra_body"], +) +def test_a_top_level_cache_control_takes_a_slot_from_the_request_marks_under_drop_params( + request_kwargs: Mapping[str, object], +): + messages: Final = _bridged_request(*_marked_turns(4), drop_params=True, **request_kwargs) + + assert ["cache_control" in message for message in messages] == [False, False, True, True, True] + + +def test_an_item_mark_replaces_a_null_mark_on_its_last_block(): + item: Final = { + "role": "user", + "content": [{"type": "input_text", "text": "x", "cache_control": None}], + "cache_control": _EPHEMERAL, + } + + assert _bridged_request(item)[1]["content"] == [{"type": "text", "text": "x", "cache_control": _EPHEMERAL}] + + +_MARKED_TOOL_ROUND_TRIP: Final = ( + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}", "cache_control": _EPHEMERAL}, + {"type": "function_call_output", "call_id": "call_1", "output": "42", "cache_control": _EPHEMERAL}, + {"role": "user", "content": "thanks", "cache_control": _EPHEMERAL}, +) + + +def _marked_tool_round_trip_marks(messages: Sequence[Mapping[str, object]]) -> tuple[bool, bool, bool]: + tool_calls: Final = TypeAdapter(list[dict[str, object]]).validate_python(messages[1]["tool_calls"]) + return ("cache_control" in tool_calls[0], "cache_control" in messages[2], "cache_control" in messages[3]) + + +def test_nova_gets_no_marks_on_tool_calls_and_results_under_drop_params(): + nova: Final = "bedrock/us.amazon.nova-2-lite-v1:0" + + assert _marked_tool_round_trip_marks(_bridged_request(*_MARKED_TOOL_ROUND_TRIP, model=nova)) == (True, True, True) + assert _marked_tool_round_trip_marks(_bridged_request(*_MARKED_TOOL_ROUND_TRIP, model=nova, drop_params=True)) == ( + False, + False, + True, + ) + + +def test_claude_keeps_marks_on_tool_calls_and_results_under_drop_params(): + assert _marked_tool_round_trip_marks(_bridged_request(*_MARKED_TOOL_ROUND_TRIP, drop_params=True)) == ( + True, + True, + True, + ) + + +@pytest.mark.parametrize("value", [1, "ephemeral", ["ephemeral"], {"type": "forever"}]) +def test_a_cache_control_value_the_provider_would_refuse_reaches_it_as_sent_without_drop_params(value: object): + assert _bridged_request({"role": "user", "content": "hi", "cache_control": value})[1] == { + "role": "user", + "content": "hi", + "cache_control": value, + } + + +@pytest.mark.parametrize("value", [1, "ephemeral", ["ephemeral"], {"type": "forever"}]) +def test_a_cache_control_value_the_provider_would_refuse_is_dropped_under_drop_params(value: object): + assert _bridged_request({"role": "user", "content": "hi", "cache_control": value}, drop_params=True)[1] == { + "role": "user", + "content": "hi", + } + + +def test_drop_params_keeps_a_valid_mark_with_its_ttl_and_without_unknown_keys(): + marked: Final = {"role": "user", "content": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h", "extra": 1}} + + assert _bridged_request(marked, drop_params=True)[1] == { + "role": "user", + "content": "hi", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + + +def test_a_marked_developer_item_reaches_a_claude_target_as_a_marked_system_message(): + developer: Final = { + "role": "developer", + "content": "Answer in French.", + "name": "policy", + "cache_control": _EPHEMERAL, + } + + assert _bridged_request(developer)[1] == { + "role": "system", + "content": "Answer in French.", + "cache_control": _EPHEMERAL, + } + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("gemini/gemini-2.5-flash", "gemini"), + ("mistral/mistral-small-latest", "mistral"), + ("bedrock/meta.llama3-70b-instruct-v1:0", "bedrock"), + ("openrouter/openai/gpt-5", "openrouter"), + ("databricks/databricks-meta-llama-3-3-70b-instruct", "databricks"), + ("snowflake/llama-claude-distill", "snowflake"), + ], +) +def test_item_marks_stay_off_a_target_without_cache_points(model: str, custom_llm_provider: str): + messages: Final = _bridged_request( + {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL}, + {"role": "developer", "content": "Answer in French.", "cache_control": _EPHEMERAL}, + {"role": "user", "content": "hi"}, + model=model, + custom_llm_provider=custom_llm_provider, + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 0 + assert messages[2] == {"role": "developer", "content": "Answer in French."} + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("databricks/databricks-claude-sonnet-4", "databricks"), + ("dashscope/qwen-plus", "dashscope"), + ("minimax/MiniMax-M2", "minimax"), + ("zai/glm-4.6", "zai"), + ("sap/anthropic--claude-4-sonnet", "sap"), + ("openrouter/anthropic/claude-sonnet-4.5", "openrouter"), + ("openrouter/google/gemini-2.5-flash", "openrouter"), + ("snowflake/claude-sonnet-4-5", "snowflake"), + ("anthropic/my-gateway-model", "anthropic"), + ], +) +def test_item_marks_reach_a_target_whose_chat_route_keeps_them(model: str, custom_llm_provider: str): + messages: Final = _bridged_request( + {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL}, + model=model, + custom_llm_provider=custom_llm_provider, + ) + + assert messages[1] == {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL} + + +def test_a_block_mark_reaches_a_target_without_cache_points_as_before(): + block: Final = {"type": "input_text", "text": "House rules", "cache_control": _EPHEMERAL} + + messages: Final = _bridged_request( + {"role": "system", "content": [block]}, model="gemini/gemini-2.5-flash", custom_llm_provider="gemini" + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 1 + + +_APPLICATION_INFERENCE_PROFILE: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123" + + +def test_an_application_inference_profile_keeps_its_item_marks_without_drop_params(): + messages: Final = _bridged_request( + {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL}, + model=f"bedrock/{_APPLICATION_INFERENCE_PROFILE}", + ) + + assert messages[1] == {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL} + + +def test_an_application_inference_profile_gets_no_item_marks_under_drop_params(): + messages: Final = _bridged_request( + {"role": "system", "content": "House rules", "cache_control": _EPHEMERAL}, + {"role": "user", "content": "Summarize the gateway latency report."}, + { + "type": "function_call", + "call_id": "call_1", + "name": "get_latency_report", + "arguments": "{}", + "cache_control": _EPHEMERAL, + }, + {"type": "function_call_output", "call_id": "call_1", "output": "p50 12 ms", "cache_control": _EPHEMERAL}, + model=f"bedrock/{_APPLICATION_INFERENCE_PROFILE}", + drop_params=True, + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 0 + + +def test_the_shared_input_transform_keeps_every_item_mark_for_a_request_it_sends(): + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast(ResponseInputParam, list(_marked_turns(5))), # cast-ok: cache_control is not an input TypedDict key + responses_api_request={}, + item_cache_control=True, + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == MAX_CACHE_CONTROL_BLOCKS + 1 + + +def test_guardrails_and_token_counting_read_the_input_as_if_it_carried_no_item_marks(): + marked_turns: Final = list(_marked_turns(5)) + + inspected: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast(ResponseInputParam, marked_turns), # cast-ok: cache_control is not an input TypedDict key + responses_api_request={}, + ) + unmarked: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast(ResponseInputParam, [{"role": "user", "content": turn["content"]} for turn in marked_turns]), + responses_api_request={}, + ) + + assert inspected == unmarked + + +def test_a_single_item_input_reads_the_same_whether_or_not_item_marks_are_kept(): + single_item: Final = cast( # cast-ok: a bare item where the input list belongs, as some clients send it + ResponseInputParam, {"type": "message", "role": "user", "content": "hi"} + ) + + inspected: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=single_item, responses_api_request={} + ) + sent: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=single_item, responses_api_request={}, item_cache_control=True + ) + + assert inspected == sent + + @pytest.mark.parametrize("index", [0, 1, 2]) def test_convert_response_function_tool_call_keeps_the_index_it_is_given(index: int) -> None: item: Final = ResponseFunctionToolCall( diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 60376408326..1e8880424ea 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -620,6 +620,7 @@ def test_sync_prompt_swap_cross_provider_with_credentials_raises(): with pytest.raises(litellm.BadRequestError, match="Refusing to send"): _apply_prompt_management_to_responses_call( input="hi", + tools=None, model="anthropic/claude-haiku-4-5", custom_llm_provider="anthropic", litellm_logging_obj=logging_obj, diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index b29e9a6b7db..9f9deb2ca0f 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -3119,7 +3119,7 @@ class TestNativeWebSocketEncryptedContentAffinity: @pytest.mark.asyncio @pytest.mark.parametrize( - "failure_frame, expected_status", + "failure_frame, expected_status, expected_message", [ ( { @@ -3131,6 +3131,7 @@ class TestNativeWebSocketEncryptedContentAffinity: }, }, 400, + "The encrypted content for item rs_1 could not be verified.", ), ( { @@ -3142,11 +3143,17 @@ class TestNativeWebSocketEncryptedContentAffinity: }, }, 500, + "upstream blew up", + ), + ( + {"type": "error", "code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.6"}, + 429, + "Rate limit reached for gpt-5.6", ), ], ) async def test_backend_to_client_books_failure_frames_as_failures( - self, failure_frame: dict[str, object], expected_status: int + self, failure_frame: dict[str, object], expected_status: int, expected_message: str ): import asyncio from unittest.mock import AsyncMock @@ -3183,7 +3190,7 @@ class TestNativeWebSocketEncryptedContentAffinity: logging_obj.dispatch_failure_handlers.assert_awaited_once() exception = logging_obj.dispatch_failure_handlers.await_args[0][0] assert exception.status_code == expected_status - assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception) + assert expected_message in str(exception) @pytest.mark.asyncio async def test_backend_to_client_bills_completed_turns_before_a_failure(self): diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index a29b9877c0d..c37e57a43ca 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -587,3 +587,24 @@ def test_maybe_raise_for_error_event_maps_vector_store_timeout_to_retriable_504( with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 504 + + +@pytest.mark.parametrize( + ("code", "status_code", "exception_type"), + [ + ("rate_limit_exceeded", 429, litellm.RateLimitError), + ("server_error", 500, litellm.InternalServerError), + (429, 429, litellm.RateLimitError), + ], +) +def test_openai_top_level_error_event_keeps_the_provider_code_and_message( + code: str | int, status_code: int, exception_type: type[Exception] +): + chunk = ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, sequence_number=3, code=code, message=f"upstream said {code}" + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + _make_iterator()._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == status_code + assert type(exc_info.value.original_exception) is exception_type + assert f"upstream said {code}" in str(exc_info.value.original_exception) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index febb33aab7d..56184fa299a 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -29,6 +29,7 @@ from litellm.router_utils.fallback_event_handlers import ( fallback_attempt_key, get_fallback_model_group, get_pre_routing_selection, + is_bad_request_without_fallback, log_failure_fallback_event, log_success_fallback_event, mid_stream_fallback_snapshot_kwargs, @@ -38,6 +39,7 @@ from litellm.router_utils.fallback_event_handlers import ( routed_deployment_id, run_async_fallback, ) +from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE from typing import Dict @@ -1378,9 +1380,7 @@ class TestOrderedFallbackLookupGroups: "smart-router", "requested-model", ) - assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == ( - "requested-model", - ) + assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == ("requested-model",) def test_fallback_hop_resumes_the_original_groups_chain_last(self): from litellm.router_utils.fallback_event_handlers import fallback_lookup_groups @@ -1419,6 +1419,19 @@ class TestHasUnattemptedFallbackTarget: assert has_unattempted_fallback_target(None, {}) is False +def test_attempted_fallback_targets_deduplicates_skips_in_order_and_keeps_first_reason() -> None: + attempted: Final = AttemptedFallbackTargets() + attempted.record_skip("a", FallbackSkipReason.NOT_AUTHORIZED) + attempted.record_skip("a", FallbackSkipReason.NOT_AUTHORIZED) + attempted.record_skip("b", FallbackSkipReason.OVER_BUDGET) + attempted.record_skip("a", FallbackSkipReason.OVER_BUDGET) + + assert attempted.skipped == ( + SkippedFallbackTarget("a", FallbackSkipReason.NOT_AUTHORIZED), + SkippedFallbackTarget("b", FallbackSkipReason.OVER_BUDGET), + ) + + def test_get_fallback_model_group_matches_provider_prefixed_key(): """A bare model group routed via a wildcard (e.g. "gpt-4o" through "openai/*") must match a fallback keyed on the provider-prefixed name, @@ -1539,7 +1552,11 @@ def test_mid_stream_retry_kwargs_strips_what_the_retry_wrapper_pops_and_keeps_th @pytest.mark.parametrize( "kwargs,expected", [ - pytest.param({"litellm_metadata": {"attempted_retries": 2}, "metadata": {"attempted_retries": 5}}, 2, id="litellm_metadata-wins"), + pytest.param( + {"litellm_metadata": {"attempted_retries": 2}, "metadata": {"attempted_retries": 5}}, + 2, + id="litellm_metadata-wins", + ), pytest.param({"metadata": {"attempted_retries": 1}}, 1, id="metadata-bucket"), pytest.param({"litellm_metadata": {"attempted_retries": "2"}}, 0, id="string-is-not-a-count"), pytest.param({"litellm_metadata": {"attempted_retries": -1}}, 0, id="negative-is-not-a-count"), @@ -1566,8 +1583,12 @@ def test_record_retry_attempt_stamps_the_bucket_the_retry_wrapper_reads(): "kwargs,expected", [ pytest.param({"litellm_metadata": {"attempted_retries": 1, "max_retries": 3}}, 3, id="committed-by-a-retry"), - pytest.param({"litellm_metadata": {"attempted_retries": 0, "max_retries": 3}}, None, id="stamped-before-any-retry"), - pytest.param({"litellm_metadata": {"attempted_retries": 1, "max_retries": "3"}}, None, id="string-is-not-a-budget"), + pytest.param( + {"litellm_metadata": {"attempted_retries": 0, "max_retries": 3}}, None, id="stamped-before-any-retry" + ), + pytest.param( + {"litellm_metadata": {"attempted_retries": 1, "max_retries": "3"}}, None, id="string-is-not-a-budget" + ), pytest.param({"litellm_metadata": {"attempted_retries": 1}}, None, id="no-budget"), pytest.param({}, None, id="no-bucket"), ], @@ -1602,6 +1623,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown(): """ @@ -1622,6 +1644,7 @@ def setup_and_teardown(): loop.close() asyncio.set_event_loop(None) + REFUSAL_RESPONSE: dict[str, Any] = { "id": "msg_refusal", "type": "message", @@ -1647,9 +1670,11 @@ OK_RESPONSE: dict[str, Any] = { "usage": {"input_tokens": 25, "output_tokens": 2}, } + def _sse(event: str, data: dict[str, Any]) -> bytes: return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = ( _sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}), _sse( @@ -1672,10 +1697,12 @@ OK_STREAM_FRAMES: tuple[bytes, ...] = ( _sse("message_stop", {"type": "message_stop"}), ) + def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]: """Split each frame's data line in half, modeling a transport chunk boundary.""" return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :])) + class _FrameStream(httpx.AsyncByteStream): def __init__(self, frames: tuple[bytes, ...]) -> None: self._frames = frames @@ -1687,6 +1714,7 @@ class _FrameStream(httpx.AsyncByteStream): async def aclose(self) -> None: return None + class FakeAnthropicUpstream: """Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable models, answers on others. The router deliberately does not forward caller-injected @@ -1724,6 +1752,7 @@ class FakeAnthropicUpstream: return patch("httpx.AsyncClient.send", new=_send) + FABLE_TIER = { "model_name": "fable-tier", "litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"}, @@ -1734,12 +1763,15 @@ OPUS_TARGET = { "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"}, } + def _router(content_policy_fallbacks: list | None) -> Router: return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks) + async def _collect(stream: AsyncIterator[bytes]) -> bytes: return b"".join([chunk async for chunk in stream]) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response(): @@ -1756,6 +1788,7 @@ async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response assert len(fake.calls) == 2 assert "claude-opus-5" in fake.calls[1] + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1780,6 +1813,7 @@ async def test_non_streaming_refusal_passes_through_untouched(content_policy_fal assert response.get("stop_details") == upstream_body.get("stop_details") assert len(fake.calls) == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_with_fallback_row_streams_fallback_frames(): @@ -1796,6 +1830,7 @@ async def test_streaming_refusal_with_fallback_row_streams_fallback_frames(): assert b"text_delta" in body assert len(fake.calls) == 2 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_split_across_chunks_still_falls_back(): @@ -1812,6 +1847,7 @@ async def test_streaming_refusal_split_across_chunks_still_falls_back(): assert b"text_delta" in body assert len(fake.calls) == 2 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_without_fallback_row_passes_frames_through(): @@ -1828,6 +1864,7 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through(): assert b"stop_details" in body assert len(fake.calls) == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): @@ -1861,6 +1898,7 @@ async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_i assert b"text_delta" in body assert len(fake.calls) == 2 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain(): @@ -1880,6 +1918,7 @@ async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain assert b'"stop_reason": "refusal"' in body assert len(fake.calls) == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_tier_stamp_never_reaches_provider_bound_metadata(): @@ -1916,6 +1955,7 @@ async def test_tier_stamp_never_reaches_provider_bound_metadata(): for body in fake.bodies: assert body.get("metadata") == {"user_id": "u1"} + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_record_pre_routing_selection_writes_only_the_internal_bucket(): """The Anthropic request's own metadata field must never carry the tier stamp.""" @@ -1926,6 +1966,7 @@ def test_record_pre_routing_selection_writes_only_the_internal_bucket(): assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"} assert kwargs["metadata"] == {"user_id": "u1"} + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) @@ -1949,6 +1990,7 @@ async def test_generic_only_row_recovers_safeguard_refusal(stream): assert len(fake.calls) == 2 assert "claude-opus-5" in fake.calls[1] + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_configured_content_policy_list_stays_authoritative_over_generic_rows(): @@ -1967,6 +2009,7 @@ async def test_configured_content_policy_list_stays_authoritative_over_generic_r assert response["stop_reason"] == "refusal" assert len(fake.calls) == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy(): router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}]) @@ -1976,6 +2019,7 @@ def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_po assert router._refusal_fallback_available("router-group", {}) is False assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_chat_content_filter_gate_unchanged_by_generic_rows(): """The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's @@ -1987,6 +2031,7 @@ def test_chat_content_filter_gate_unchanged_by_generic_rows(): assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) @@ -2012,6 +2057,7 @@ async def test_disable_fallbacks_returns_the_refusal_instead_of_raising(stream): assert body["stop_reason"] == "refusal" assert len(fake.calls) == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_disable_fallbacks_beats_a_content_policy_row_too(): @@ -2032,6 +2078,7 @@ async def test_disable_fallbacks_beats_a_content_policy_row_too(): assert response["stop_reason"] == "refusal" assert len(fake.calls) == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_refusal_gate_keys_on_pre_routing_tier_stamp(): router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}]) @@ -2059,6 +2106,7 @@ def test_refusal_gate_keys_on_pre_routing_tier_stamp(): is False ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_has_content_policy_fallback_default_fallbacks_arm(): router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}]) @@ -2066,6 +2114,7 @@ def test_has_content_policy_fallback_default_fallbacks_arm(): assert router._has_content_policy_fallback("any-group", {}) is True assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested(): router = _router(content_policy_fallbacks=None) @@ -2079,6 +2128,7 @@ def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested ) == ["backup-b"] assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_refusal_gate_ignores_other_generic_call_types(): router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) @@ -2269,9 +2319,7 @@ async def test_failed_fallbacks_raise_most_recent_exception(function_name): model="gpt-3.5-turbo", ) - request_kwargs: Dict[str, Any] = { - "metadata": {"previous_models": ["gpt-3.5-turbo"]} - } + request_kwargs: Dict[str, Any] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}} if function_name == "_aembedding": request_kwargs["input"] = "hello this is a test for run_async_fallback" @@ -2323,27 +2371,19 @@ class CustomTestLogger(CustomLogger): self.success_fallback_events = [] self.failure_fallback_events = [] - async def log_success_fallback_event( - self, original_model_group, kwargs, original_exception - ): + async def log_success_fallback_event(self, original_model_group, kwargs, original_exception): print( "in log_success_fallback_event for original_model_group: ", original_model_group, ) - self.success_fallback_events.append( - (original_model_group, kwargs, original_exception) - ) + self.success_fallback_events.append((original_model_group, kwargs, original_exception)) - async def log_failure_fallback_event( - self, original_model_group, kwargs, original_exception - ): + async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception): print( "in log_failure_fallback_event for original_model_group: ", original_model_group, ) - self.failure_fallback_events.append( - (original_model_group, kwargs, original_exception) - ) + self.failure_fallback_events.append((original_model_group, kwargs, original_exception)) def create_test_router_2(): @@ -2458,3 +2498,73 @@ async def test_multiple_fallbacks(function_name, respx_mock: respx.MockRouter, m print(result._hidden_params) assert result._hidden_params["api_base"] == FAKE_OPENAI_API_BASE + + +def _provider_error(status_code: int) -> httpx.Response: + return httpx.Response(status_code=status_code, request=httpx.Request(method="POST", url="https://provider.test")) + + +@pytest.mark.parametrize( + ("exception", "stops_fallbacks"), + [ + pytest.param( + litellm.BadRequestError(message="Unsupported value: temperature", model="m", llm_provider="openai"), + True, + id="provider 400", + ), + pytest.param( + litellm.UnsupportedParamsError( + message="reasoning_effort is not supported", model="m", llm_provider="openai" + ), + True, + id="client-side rejection", + ), + pytest.param( + litellm.UnsupportedParamsError(message="stop is not supported"), + True, + id="client-side rejection without a provider", + ), + pytest.param( + litellm.BadRequestError( + message="Deployment ptu is reserved for another team", model="ptu", llm_provider="" + ), + False, + id="router-raised 400 keeps falling back", + ), + pytest.param( + litellm.BadRequestError( + message="VertexAIException BadRequestError - 403", + model="m", + llm_provider="vertex_ai", + response=_provider_error(403), + ), + False, + id="vertex 403 wrapped as BadRequestError", + ), + pytest.param( + litellm.PaymentRequiredError(message="insufficient credits", model="m", llm_provider="openrouter"), + False, + id="402", + ), + pytest.param( + litellm.ContextWindowExceededError(message="maximum context length", model="m", llm_provider="openai"), + False, + id="context window keeps its own fallbacks", + ), + pytest.param( + litellm.ContentPolicyViolationError(message="content_policy_violation", model="m", llm_provider="openai"), + False, + id="content policy keeps its own fallbacks", + ), + pytest.param( + litellm.InternalServerError(message="overloaded", model="m", llm_provider="openai"), + False, + id="500 still falls back", + ), + ], +) +def test_is_bad_request_without_fallback_only_stops_on_a_real_400_when_the_flag_is_off( + exception: Exception, stops_fallbacks: bool +) -> None: + assert is_bad_request_without_fallback(exception, fallback_on_unsupported_params=False) is stops_fallbacks + assert is_bad_request_without_fallback(exception, fallback_on_unsupported_params=True) is False diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 3694ecea4ed..d532ce186dd 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -1,5 +1,5 @@ import logging -from typing import Dict, List, Optional, Union +from typing import Dict, Final, List, Optional, Union from unittest.mock import Mock import pytest @@ -12,6 +12,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + format_fallback_outcome_message, provider_for_generic_call, resolve_model_group_alias, resolve_served_model, @@ -19,6 +20,7 @@ from litellm.router_utils.common_utils import ( PROVIDER_SCOPED_CREDENTIAL_PARAMS, warn_on_provider_credential_mismatch, ) +from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget class TestFilterTeamBasedModels: @@ -610,6 +612,56 @@ class TestTruncateFallbackErrorDetail: assert len(truncate_fallback_error_detail("x" * 1_000_000)) < 3_000 +def test_format_fallback_outcome_reports_unauthorized_skips_without_failure_detail() -> None: + message: Final = format_fallback_outcome_message( + "primary", + ("restricted-model",), + "", + skipped=(SkippedFallbackTarget("restricted-model", FallbackSkipReason.NOT_AUTHORIZED),), + ) + + assert "Fallback to restricted-model was skipped because the caller is not authorized to call it." in message + assert "also failed" not in message + + +def test_format_fallback_outcome_separates_attempted_failures_from_skipped_targets() -> None: + message: Final = format_fallback_outcome_message( + "primary", + ("restricted-model", "dead-allowed"), + "boom", + skipped=(SkippedFallbackTarget("restricted-model", FallbackSkipReason.NOT_AUTHORIZED),), + ) + failed_sentence: Final = message.split("Fallback to dead-allowed also failed:", 1)[1].split( + " Fallback to restricted-model was skipped", 1 + )[0] + + assert "Fallback to dead-allowed also failed: boom" in message + assert "Fallback to restricted-model was skipped because the caller is not authorized" in message + assert "restricted-model" not in failed_sentence + + +def test_format_fallback_outcome_uses_plural_for_multiple_over_budget_skips() -> None: + message: Final = format_fallback_outcome_message( + "primary", + ("paid-one", "paid-two"), + "", + skipped=( + SkippedFallbackTarget("paid-one", FallbackSkipReason.OVER_BUDGET), + SkippedFallbackTarget("paid-two", FallbackSkipReason.OVER_BUDGET), + ), + ) + + assert "Fallback to paid-one, paid-two was skipped because the caller is over budget for them." in message + + +def test_format_fallback_outcome_preserves_attempted_failure_wording_without_skips() -> None: + assert ( + format_fallback_outcome_message("primary", ("fallback",), "detail") + == "\n\nLiteLLM: model group 'primary' failed with the error above. " + "Fallback to fallback also failed: detail" + ) + + class TestWarnOnProviderCredentialMismatch: """A deployment that carries one provider's credentials while resolving to another is silently broken: litellm ignores the credentials and sends the diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 64f3f88cd89..0a795ff2a0f 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1638,7 +1638,8 @@ def test_arouter_responses_api_bridge(): ) ## CONFIRM BRIDGE IS CALLED - with patch.object(litellm, "responses", return_value=AsyncMock()) as mock_responses: + # an empty stream: the Router reads the first chunk before returning, and a stream of mocks never produces one + with patch.object(litellm, "responses", return_value=iter([])) as mock_responses: result = router.completion( model="[IP-approved] o3-pro", messages=[{"role": "user", "content": "Hello, world!"}], @@ -4020,11 +4021,10 @@ async def test_acompletion_mid_stream_fallback_honors_a_per_request_fallbacks_no return _stream_for(**kwargs) with patch("litellm.acompletion", side_effect=fake_acompletion) as provider_calls: - response: Final = await router.acompletion( - model="primary", messages=[{"role": "user", "content": "hi"}], stream=True, fallbacks=None - ) with pytest.raises(litellm.InternalServerError, match="provider 500 from openai/primary-model"): - [chunk async for chunk in response] + await router.acompletion( + model="primary", messages=[{"role": "user", "content": "hi"}], stream=True, fallbacks=None + ) assert _groups_called(provider_calls) == ["primary"] @@ -16837,11 +16837,10 @@ async def test_acompletion_disable_fallbacks_reaches_the_mid_stream_hop(): return _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error), model=kwargs["model"]) with patch("litellm.acompletion", side_effect=primary_stream) as provider_calls: - response = await router.acompletion( - model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True - ) with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in response] + await router.acompletion( + model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True + ) assert raised.value is primary_error assert [call.kwargs["metadata"]["model_group"] for call in provider_calls.call_args_list] == ["primary"] @@ -16856,11 +16855,10 @@ def test_completion_disable_fallbacks_reaches_the_mid_stream_hop(): return _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error), model=kwargs["model"]) with patch("litellm.completion", side_effect=primary_stream) as provider_calls: - response = router.completion( - model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True - ) with pytest.raises(litellm.InternalServerError) as raised: - list(response) + router.completion( + model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True + ) assert raised.value is primary_error assert [call.kwargs["metadata"]["model_group"] for call in provider_calls.call_args_list] == ["primary"] @@ -17803,6 +17801,48 @@ def _router_with_failing_primary(fallback_access_check) -> Router: ) +class _DenyFallbackBudgetCheck: + async def __call__( + self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router + ) -> bool: + return False + + +def _router_with_failing_fallback_chain( + fallback_targets: tuple[str, ...], + *, + fallback_access_check: _AllowlistFallbackAccessCheck | None = None, + fallback_budget_check: _DenyFallbackBudgetCheck | None = None, +) -> Router: + return Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/primary", + "api_key": "k", + "mock_response": Exception("primary is down"), + }, + }, + *( + { + "model_name": target, + "litellm_params": { + "model": f"openai/{target}", + "api_key": "k", + "mock_response": Exception(f"{target} is down"), + }, + } + for target in fallback_targets + ), + ], + fallbacks=[{"primary": list(fallback_targets)}], + num_retries=0, + fallback_access_check=fallback_access_check, + fallback_budget_check=fallback_budget_check, + ) + + @pytest.mark.asyncio async def test_fallback_access_check_blocks_config_fallback_the_caller_cannot_use(): access_check = _AllowlistFallbackAccessCheck(allowed_models=frozenset()) @@ -17814,6 +17854,59 @@ async def test_fallback_access_check_blocks_config_fallback_the_caller_cannot_us assert access_check.checked_models == ["secret-fallback"] +@pytest.mark.asyncio +async def test_fallback_access_check_reports_denied_fallback_as_skipped(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "expose_router_debug_in_errors", True) + access_check: Final = _AllowlistFallbackAccessCheck(allowed_models=frozenset()) + router: Final = _router_with_failing_primary(access_check) + primary_only_router: Final = Router(model_list=router.model_list, fallbacks=[], num_retries=0) + + with pytest.raises(litellm.InternalServerError) as primary_exception: + await primary_only_router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + with pytest.raises(litellm.InternalServerError) as fallback_exception: + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + + message: Final = str(fallback_exception.value) + assert type(fallback_exception.value) is type(primary_exception.value) + assert fallback_exception.value.status_code == primary_exception.value.status_code + assert message.count("primary is down") == 1 + assert "Fallback to secret-fallback was skipped because the caller is not authorized to call it." in message + assert "also failed" not in message + + +@pytest.mark.asyncio +async def test_fallback_access_check_separates_denied_and_failed_targets(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "expose_router_debug_in_errors", True) + access_check: Final = _AllowlistFallbackAccessCheck(allowed_models=frozenset({"allowed-fallback"})) + router: Final = _router_with_failing_fallback_chain( + ("secret-fallback", "allowed-fallback"), + fallback_access_check=access_check, + ) + + with pytest.raises(litellm.InternalServerError) as exception: + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + + message: Final = str(exception.value) + assert "Fallback to allowed-fallback also failed:" in message + assert "Fallback to secret-fallback was skipped because the caller is not authorized" in message + + +@pytest.mark.asyncio +async def test_fallback_budget_check_reports_skipped_targets(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "expose_router_debug_in_errors", True) + router: Final = _router_with_failing_fallback_chain( + ("paid-fallback",), + fallback_budget_check=_DenyFallbackBudgetCheck(), + ) + + with pytest.raises(litellm.InternalServerError) as exception: + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + + message: Final = str(exception.value) + assert "Fallback to paid-fallback was skipped because the caller is over budget for it." in message + assert "also failed" not in message + + @pytest.mark.asyncio async def test_fallback_access_check_lets_an_authorized_config_fallback_through(): router = _router_with_failing_primary(_AllowlistFallbackAccessCheck(allowed_models=frozenset({"secret-fallback"}))) @@ -26135,3 +26228,98 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None assert text == "stream answer" assert chunks[-1].choices[0].finish_reason == "stop" + + +class _FirstPullStream: + """A stream whose first pull yields `first`, raises it, or ends (first=None).""" + + def __init__(self, first: object) -> None: + self._first = first + + def _pull(self, stop: type[Exception]) -> object: + if self._first is None: + raise stop + if isinstance(self._first, Exception): + raise self._first + return self._first + + def __next__(self) -> object: + return self._pull(StopIteration) + + async def __anext__(self) -> object: + return self._pull(StopAsyncIteration) + + +def _first_chunk_cases(): + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _FirstChunk + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))]) + provider_error = litellm.InternalServerError(message="primary died at stream start", llm_provider="openai", model="m") + wrapped = MidStreamFallbackError( + message=str(provider_error), + model="m", + llm_provider="openai", + original_exception=provider_error, + is_pre_first_chunk=True, + ) + return [ + pytest.param(chunk, _FirstChunk(chunk=chunk), None, id="first-chunk"), + pytest.param(None, _FirstChunk(chunk=None), None, id="empty-stream"), + pytest.param(wrapped, None, provider_error, id="pre-content-failure"), + ] + + +@pytest.mark.parametrize("first, expected, raised", _first_chunk_cases()) +def test_read_first_chunk_direct_call(first, expected, raised): + from litellm.router import _read_first_chunk + + if raised is None: + assert _read_first_chunk(_FirstPullStream(first)) == expected + return + with pytest.raises(type(raised)) as caught: + _read_first_chunk(_FirstPullStream(first)) + assert caught.value is raised + + +@pytest.mark.asyncio +@pytest.mark.parametrize("first, expected, raised", _first_chunk_cases()) +async def test_aread_first_chunk_direct_call(first, expected, raised): + from litellm.router import _aread_first_chunk + + if raised is None: + assert await _aread_first_chunk(_FirstPullStream(first)) == expected + return + with pytest.raises(type(raised)) as caught: + await _aread_first_chunk(_FirstPullStream(first)) + assert caught.value is raised + + +def _chunks_with_first_cases(): + from litellm.router import _FirstChunk + + return [ + pytest.param(None, ["a", "b"], ["a", "b"], id="nothing-read-stream-as-is"), + pytest.param(_FirstChunk(chunk=None), ["a", "b"], [], id="empty-stream-not-touched"), + pytest.param(_FirstChunk(chunk="first"), ["a", "b"], ["first", "a", "b"], id="first-chunk-put-back"), + ] + + +@pytest.mark.parametrize("taken, rest, expected", _chunks_with_first_cases()) +def test_chunks_with_first_direct_call(taken, rest, expected): + from litellm.router import _chunks_with_first + + assert list(_chunks_with_first(iter(rest), taken)) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("taken, rest, expected", _chunks_with_first_cases()) +async def test_achunks_with_first_direct_call(taken, rest, expected): + from litellm.router import _achunks_with_first + + async def stream(): + for item in rest: + yield item + + assert [item async for item in _achunks_with_first(stream(), taken)] == expected diff --git a/tests/unit/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py index 48b5194a66e..918d215c89f 100644 --- a/tests/unit/test_router_retry_policy_update.py +++ b/tests/unit/test_router_retry_policy_update.py @@ -300,6 +300,21 @@ def test_update_settings_persists_retry_policy_dict(): assert router.retry_policy.TimeoutErrorRetries == 3 +def test_update_settings_fallback_on_unsupported_params_round_trips_through_get_settings(): + router = _build_router() + assert router.get_settings()["fallback_on_unsupported_params"] is True + + router.update_settings(fallback_on_unsupported_params=False) + + assert router.fallback_on_unsupported_params is False + assert router.get_settings()["fallback_on_unsupported_params"] is False + + +def test_update_router_config_keeps_fallback_on_unsupported_params_payload(): + payload = UpdateRouterConfig(fallback_on_unsupported_params=False).model_dump(exclude_none=True) + assert payload["fallback_on_unsupported_params"] is False + + def test_update_settings_accepts_retry_policy_object_unchanged(): """A pre-built ``RetryPolicy`` instance must pass through verbatim so callers that already constructed one (e.g. tests or programmatic diff --git a/ui/Dockerfile b/ui/Dockerfile index 073dadb2da4..4b6b0c1581a 100644 --- a/ui/Dockerfile +++ b/ui/Dockerfile @@ -33,6 +33,9 @@ RUN rm -f /etc/nginx/conf.d/default.conf # Static export → web root. COPY --from=builder /app/out /usr/share/nginx/html +# Runtime backend discovery config. +COPY ui/docker-entrypoint.d/40-litellm-runtime-config.sh /docker-entrypoint.d/ + # Routing rules — see ui/nginx.conf for the full description. COPY ui/nginx.conf /etc/nginx/nginx.conf diff --git a/ui/docker-entrypoint.d/40-litellm-runtime-config.sh b/ui/docker-entrypoint.d/40-litellm-runtime-config.sh new file mode 100755 index 00000000000..4b53ab0b19f --- /dev/null +++ b/ui/docker-entrypoint.d/40-litellm-runtime-config.sh @@ -0,0 +1,32 @@ +#!/bin/sh +set -eu + +runtime_dir=/tmp/litellm-ui +runtime_file=$runtime_dir/runtime-config.json +rm -f "$runtime_file" 2>/dev/null || true + +if [ "${LITELLM_BACKEND_URL+x}" = "" ]; then + exit 0 +fi + +newline=' +' +case "$LITELLM_BACKEND_URL" in + *"$newline"*) + printf '%s\n' 'WARN: ignoring invalid LITELLM_BACKEND_URL' >&2 + exit 0 + ;; +esac + +if ! printf '%s\n' "$LITELLM_BACKEND_URL" | grep -Eq '^https?://[^"\\ ]+$'; then + printf '%s\n' 'WARN: ignoring invalid LITELLM_BACKEND_URL' >&2 + exit 0 +fi + +backend_url=$LITELLM_BACKEND_URL +while [ "${backend_url%/}" != "$backend_url" ]; do + backend_url=${backend_url%/} +done + +mkdir -p "$runtime_dir" +printf '{"backend_url":"%s"}\n' "$backend_url" > "$runtime_file" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx index aad63e979ac..a4a615fe798 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx @@ -1,6 +1,6 @@ import { useAccessGroupDetails } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails"; import { AccessGroupResponse } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../../../tests/test-utils"; @@ -58,6 +58,7 @@ const createMockAccessGroup = (overrides: Partial = {}): Ac assigned_key_ids: ["key-1", "key-2"], access_mcp_servers: [{ id: "mcp-1", name: "GitHub MCP" }], access_agents: [{ id: "agent-1", name: "Support Agent" }], + access_skill_names: ["private-plugin"], assigned_teams: [{ id: "team-1", name: "Platform Team" }], assigned_keys: [ { id: "key-1", name: "ci-key" }, @@ -336,6 +337,34 @@ describe("AccessGroupDetail", () => { expect(screen.getByText("No MCP servers assigned to this group")).toBeInTheDocument(); }); + + it("should list the allowed tools per server, with deny-all and unrestricted servers spelled out", async () => { + const user = userEvent.setup(); + renderWith({ + access_mcp_servers: [ + { id: "mcp-1", name: "GitHub MCP" }, + { id: "mcp-2", name: "Jira MCP" }, + { id: "mcp-3", name: "Slack MCP" }, + ], + mcp_tool_permissions: { "mcp-1": ["list_issues", "create_issue"], "mcp-2": [] }, + }); + + await user.click(screen.getByRole("tab", { name: /MCP Servers/i })); + + const list = screen.getByRole("region", { name: "MCP tool permissions" }); + expect(within(list).getByText("list_issues, create_issue")).toBeInTheDocument(); + expect(within(list).getByText("No tools")).toBeInTheDocument(); + expect(within(list).getByText("All tools")).toBeInTheDocument(); + }); + + it("should not render a tool permission list when the group has none", async () => { + const user = userEvent.setup(); + renderWith({ mcp_tool_permissions: null }); + + await user.click(screen.getByRole("tab", { name: /MCP Servers/i })); + + expect(screen.queryByRole("region", { name: "MCP tool permissions" })).not.toBeInTheDocument(); + }); }); describe("Agents tab", () => { @@ -368,6 +397,26 @@ describe("AccessGroupDetail", () => { }); }); + describe("Skills tab", () => { + it("should list the plugin names granted by the group", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("tab", { name: /Skills/i })); + + expect(screen.getByText("private-plugin")).toBeInTheDocument(); + }); + + it("should show empty state when none assigned", async () => { + const user = userEvent.setup(); + renderWith({ access_skill_names: [] }); + + await user.click(screen.getByRole("tab", { name: /Skills/i })); + + expect(screen.getByText("No skills assigned to this group")).toBeInTheDocument(); + }); + }); + it("should show empty state in Models tab when no models assigned", () => { renderWith({ access_model_names: [] }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx index 118a8655d5a..11d75c0f750 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx @@ -1,5 +1,14 @@ import { useAccessGroupDetails } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails"; -import { ArrowLeftIcon, BotIcon, EditIcon, KeyIcon, LayersIcon, ServerIcon, UsersIcon } from "lucide-react"; +import { + ArrowLeftIcon, + BotIcon, + EditIcon, + KeyIcon, + LayersIcon, + ServerIcon, + SparklesIcon, + UsersIcon, +} from "lucide-react"; import { useState } from "react"; import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag"; import { BadgeLink } from "@/components/shared/BadgeLink"; @@ -25,6 +34,37 @@ const MAX_PREVIEW = 5; const shortId = (id: string) => (id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id); +const describeAllowedTools = (tools: readonly string[] | undefined): string => { + if (tools === undefined) { + return "All tools"; + } + return tools.length === 0 ? "No tools" : tools.join(", "); +}; + +function MCPToolPermissionList({ + servers, + toolPermissions, +}: { + servers: readonly AccessGroupResource[]; + toolPermissions: Readonly> | null; +}) { + if (toolPermissions === null || servers.length === 0) { + return null; + } + return ( +
+ {servers.map(({ id, name }) => { + return ( +
+ {name ?? id} + {describeAllowedTools(toolPermissions[id])} +
+ ); + })} +
+ ); +} + function ResourceList({ items, emptyMessage }: { items: readonly AccessGroupResource[]; emptyMessage: string }) { if (items.length === 0) { return

{emptyMessage}

; @@ -95,6 +135,7 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr const models = accessGroup.access_model_names.map((id) => ({ id, name: null })); const mcpServers = accessGroup.access_mcp_servers; const agents = accessGroup.access_agents; + const skills = (accessGroup.access_skill_names ?? []).map((name) => ({ id: name, name: null })); const keys = accessGroup.assigned_keys; const teams = accessGroup.assigned_teams; @@ -231,16 +272,25 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr Agents {agents.length} + + + Skills + {skills.length} + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx index 91e008de402..43eaf11b60d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx @@ -1,15 +1,18 @@ "use client"; -import { BotIcon, InfoIcon, LayersIcon, ServerIcon } from "lucide-react"; +import { BotIcon, InfoIcon, LayersIcon, ServerIcon, SparklesIcon } from "lucide-react"; import type { UseFormReturn } from "react-hook-form"; import { z } from "zod"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { MultiSelect } from "@/components/shared/MultiSelect"; +import SkillSelector from "@/components/skills/SkillSelector"; import { Input } from "@/components/ui/input"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Textarea } from "@/components/ui/textarea"; @@ -19,7 +22,9 @@ export const accessGroupFormSchema = z.object({ description: z.string(), modelIds: z.array(z.string()), mcpServerIds: z.array(z.string()), + mcpToolPermissions: z.record(z.string(), z.array(z.string())), agentIds: z.array(z.string()), + skillNames: z.array(z.string()), }); export type AccessGroupFormValues = z.output; @@ -28,12 +33,44 @@ export const GENERAL_TAB = "general"; export const MODELS_TAB = "models"; export const MCP_SERVERS_TAB = "mcp-servers"; export const AGENTS_TAB = "agents"; +export const SKILLS_TAB = "skills"; interface AccessGroupBaseFormProps { form: UseFormReturn; isNameDisabled?: boolean; activeTab: string; onTabChange: (tab: string) => void; + initialServerIds?: readonly string[]; +} + +export function AccessGroupMCPToolPermissions({ + selectedServers, + toolPermissions, + onChange, + preserveUnrestrictedServers, +}: { + selectedServers: string[]; + toolPermissions: Record; + onChange: (next: Record) => void; + preserveUnrestrictedServers?: readonly string[]; +}) { + const { accessToken } = useAuthorized(); + + if (selectedServers.length === 0) { + return null; + } + + return ( +
+ +
+ ); } export function AccessGroupBaseForm({ @@ -41,7 +78,9 @@ export function AccessGroupBaseForm({ isNameDisabled = false, activeTab, onTabChange, + initialServerIds, }: AccessGroupBaseFormProps) { + const { accessToken } = useAuthorized(); const { data: agentsData } = useAgents(); const { data: mcpServersData } = useMCPServers(); @@ -73,6 +112,10 @@ export function AccessGroupBaseForm({ Agents + + + Skills + @@ -108,6 +151,12 @@ export function AccessGroupBaseForm({ /> )} + form.setValue("mcpToolPermissions", next, { shouldDirty: true })} + preserveUnrestrictedServers={initialServerIds} + /> @@ -123,6 +172,24 @@ export function AccessGroupBaseForm({ )} + + + + {({ value, onChange }) => ( + + )} + + ); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx index 7c4e218261f..ca1e09d1d46 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx @@ -3,6 +3,7 @@ import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event" import { fireEvent, renderWithProviders, screen, waitFor, within } from "../../../../../../tests/test-utils"; import { AccessGroupEditModal } from "./AccessGroupEditModal"; import { AccessGroupResponse } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import * as networking from "@/components/networking"; const mutate = vi.fn(); @@ -31,6 +32,19 @@ vi.mock("@/components/ModelSelect/ModelSelect", () => ({ ), })); +vi.mock("@/components/networking"); + +vi.mock("@/components/skills/SkillSelector", () => ({ + default: ({ value, onChange }: { value?: string[]; onChange: (next: string[]) => void }) => ( +
+ +
+ ), +})); + vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn(), error: vi.fn() }, })); @@ -44,7 +58,9 @@ const accessGroup: AccessGroupResponse = { description: "Engineers", access_model_names: ["gpt-4"], access_mcp_server_ids: ["srv-1"], + mcp_tool_permissions: { "srv-1": ["read_file"] }, access_agent_ids: ["agent-1"], + access_skill_names: ["old-plugin"], assigned_team_ids: [], assigned_key_ids: [], access_mcp_servers: [{ id: "srv-1", name: "Server One" }], @@ -67,6 +83,18 @@ const variables = () => mutate.mock.calls.at(-1)?.[0] as { accessGroupId: string describe("AccessGroupEditModal submit payload", () => { beforeEach(() => { vi.clearAllMocks(); + vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: "srv-1", server_name: "Files", alias: "Files" }, + ]); + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: [ + { name: "read_file", description: "Read a file" }, + { name: "write_file", description: "Write a file" }, + ], + error: false, + }); }); it("sends exactly the antd payload for an untouched save", async () => { @@ -86,9 +114,37 @@ describe("AccessGroupEditModal submit payload", () => { access_model_names: undefined, access_mcp_server_ids: undefined, access_agent_ids: undefined, + access_skill_names: undefined, }); }); + it("sends the skills picked on the Skills tab", async () => { + const user = setup(); + renderModal(); + await screen.findByDisplayValue("Engineering"); + + await user.click(screen.getByRole("tab", { name: /Skills/ })); + await user.click(screen.getByLabelText("skill-select")); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.access_skill_names).toStrictEqual(["old-plugin", "private-plugin"]); + expect(variables().params.access_agent_ids).toBeUndefined(); + }); + + it("sends an empty list when every skill is removed on the Skills tab", async () => { + const user = setup(); + renderModal(); + await screen.findByDisplayValue("Engineering"); + + await user.click(screen.getByRole("tab", { name: /Skills/ })); + await user.click(screen.getByLabelText("skill-clear")); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.access_skill_names).toStrictEqual([]); + }); + it("sends a tab's field only once that tab has been visited", async () => { const user = setup(); renderModal(); @@ -106,6 +162,7 @@ describe("AccessGroupEditModal submit payload", () => { access_model_names: undefined, access_mcp_server_ids: ["srv-1"], access_agent_ids: ["agent-1"], + access_skill_names: undefined, }); }); @@ -169,6 +226,20 @@ describe("AccessGroupEditModal submit payload", () => { expect(mutate).not.toHaveBeenCalled(); }); + it("sends the edited tool checklist as mcp_tool_permissions, an emptied one as an explicit []", async () => { + const user = setup(); + renderModal(); + await screen.findByDisplayValue("Engineering"); + + await user.click(screen.getByRole("tab", { name: /MCP Servers/ })); + await user.click(await screen.findByRole("radio", { name: "Flat List" })); + await user.click(await screen.findByRole("checkbox", { name: "read_file" })); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.mcp_tool_permissions).toStrictEqual({ "srv-1": [] }); + }); + it("renders each selected MCP server as its own removable chip and drops one on remove", async () => { const user = setup(); renderModal(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx index 7370b3bc195..9ee2583474b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx @@ -1,6 +1,7 @@ "use client"; import React, { useState } from "react"; +import { useFormState } from "react-hook-form"; import { toast } from "@/lib/toast"; import { useZodForm } from "@/lib/forms/useZodForm"; @@ -12,6 +13,7 @@ import { AccessGroupBaseForm, accessGroupFormSchema, AGENTS_TAB, + SKILLS_TAB, GENERAL_TAB, MCP_SERVERS_TAB, MODELS_TAB, @@ -31,11 +33,14 @@ const toFormValues = (accessGroup: AccessGroupResponse): AccessGroupFormValues = description: accessGroup.description ?? "", modelIds: accessGroup.access_model_names ?? [], mcpServerIds: accessGroup.access_mcp_server_ids ?? [], + mcpToolPermissions: accessGroup.mcp_tool_permissions ?? {}, agentIds: accessGroup.access_agent_ids ?? [], + skillNames: accessGroup.access_skill_names ?? [], }); function AccessGroupEditForm({ accessGroup, onCancel, onSuccess }: Omit) { const form = useZodForm(accessGroupFormSchema, { defaultValues: toFormValues(accessGroup) }); + const { dirtyFields } = useFormState({ control: form.control }); const editMutation = useEditAccessGroup(); const [activeTab, setActiveTab] = useState(GENERAL_TAB); const [visitedTabs, setVisitedTabs] = useState>(new Set([GENERAL_TAB])); @@ -52,7 +57,9 @@ function AccessGroupEditForm({ accessGroup, onCancel, onSuccess }: Omit event.preventDefault()}> - +