Merge remote-tracking branch 'origin/main' into beluga_decrease_anys_20261102

.github/workflows/test-unit.yml

@@ -977,10 +977,10 @@ jobs:
             dist: loadscope
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
-        timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
+        timeout-minutes: 3
         with:
           persist-credentials: false
-          fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
+          fetch-depth: 1
 
       - name: Detect relevant changes
         id: changes
@@ -1013,13 +1013,11 @@ jobs:
         with:
           artifact: ${{ needs.rust-bridge.outputs.artifact }}
 
-      - name: Install dependencies and check schema
+      - name: Install dependencies
         if: steps.changes.outputs.decision != 'skip'
-        timeout-minutes: 8
+        timeout-minutes: 5
         env:
           RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
-          UNIT_SHARD: ${{ matrix.shard }}
-          SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }}
         run: |
           diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
           if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
@@ -1031,9 +1029,21 @@ jobs:
             uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
           fi
           uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
-          if [ "$UNIT_SHARD" = "db-and-spend" ]; then
-            uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF"
-          fi
+
+      - name: Check schema.prisma against the merge base
+        if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request'
+        timeout-minutes: 3
+        env:
+          GH_TOKEN: ${{ github.token }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
+        run: |
+          base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref)
+          merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
+          test -n "$merge_base"
+          git fetch --no-tags --depth=1 origin "$merge_base"
+          uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base"
 
       - name: Cache Prisma binaries
         if: steps.changes.outputs.decision != 'skip'

helm/litellm-helm/Chart.yaml

@@ -18,7 +18,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 1.1.3
+version: 1.1.4
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to

helm/litellm-helm/README.md

@@ -56,8 +56,8 @@ If `db.useStackgresOperator` is used (not yet implemented):
 
 | `billingMetrics.enabled`    | Enable enterprise billable-request metering. Requires an enterprise license.                                                                                                                                                                  | `false`                   |
 | `billingMetrics.endpoint`   | Collector that the billable-request counter is pushed to.                                                                                                                                                                                    | `https://telemetry.litellm.ai` |
-| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`.                                                                                                                                       | `litellm-billing-metrics-mtls` |
-| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI.                                                                    | `""`                      |
+| `billingMetrics.tokenUrl`   | Cognito OAuth2 client-credentials token endpoint.                                                                                                                                                                                               | `https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token` |
+| `billingMetrics.secretName` | Name of an existing Secret with `client-id` and `client-secret` keys.                                                                                                                                                                           | `litellm-billing-metrics` |
 | `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset.                                                                                                                                             | `""`                      |
 
 #### Example `proxy_config` ConfigMap from values (default):
@@ -109,10 +109,10 @@ type: Opaque
 
 #### Enterprise billable-request metering
 
-Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
+Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector with Cognito bearer authentication. Create a Secret with `client-id` and `client-secret` keys, then enable the chart values:
 
 ```
-kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
+kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
 ```
 
 ```
@@ -120,7 +120,8 @@ billingMetrics:
   enabled: true
 ```
 
-Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
+The endpoint and token URL must use HTTPS. The chart requires both URLs and
+the Secret name when billing metrics are enabled
 
 ### Database Settings
 

helm/litellm-helm/templates/_helpers.tpl

@@ -50,53 +50,27 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
 app.kubernetes.io/instance: {{ .Release.Name }}
 {{- end }}
 
-{{/*
-Enterprise billable-request metering. The client certificate identifies the
-deployment to LiteLLM's collector, so it is mounted read-only from an existing
-Secret rather than passed through the environment.
-*/}}
-{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
-{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
-
 {{- define "litellm.billingMetricsEnv" -}}
 - name: LITELLM_BILLING_METRICS_ENDPOINT
   value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_CERT
-  value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_KEY
-  value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: LITELLM_BILLING_METRICS_CA_CERT
-  value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
-{{- end }}
+- name: LITELLM_BILLING_METRICS_TOKEN_URL
+  value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
+- name: LITELLM_BILLING_METRICS_CLIENT_ID
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-id
+- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-secret
 {{- with .Values.billingMetrics.exportIntervalMs }}
 - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
   value: {{ . | quote }}
 {{- end }}
 {{- end -}}
 
-{{- define "litellm.billingMetricsVolumes" -}}
-- name: billing-metrics-mtls
-  secret:
-    secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  secret:
-    secretName: {{ .Values.billingMetrics.caSecretName }}
-{{- end }}
-{{- end -}}
-
-{{- define "litellm.billingMetricsVolumeMounts" -}}
-- name: billing-metrics-mtls
-  mountPath: {{ include "litellm.billingMetrics.certDir" . }}
-  readOnly: true
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  mountPath: {{ include "litellm.billingMetrics.caDir" . }}
-  readOnly: true
-{{- end }}
-{{- end -}}
-
 {{/*
 Create the name of the service account to use
 */}}

helm/litellm-helm/templates/deployment.yaml

@@ -148,9 +148,6 @@ spec:
             - name: npm
               mountPath: /.npm
           {{- end }}
-          {{- if .Values.billingMetrics.enabled }}
-            {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
-          {{- end }}
           {{- if include "litellm.collector.socketDir" . }}
             - name: collector-socket
               mountPath: {{ include "litellm.collector.socketDir" . }}
@@ -234,9 +231,6 @@ spec:
             items:
               - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
                 path: "config.yaml"
-        {{- if .Values.billingMetrics.enabled }}
-        {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
-        {{- end }}
         {{- if include "litellm.collector.socketDir" . }}
         - name: collector-socket
           emptyDir:

helm/litellm-helm/tests/billing_metrics_tests.yaml

@@ -1,31 +1,21 @@
-suite: test billingMetrics wiring on the proxy deployment
+suite: test bearer-only billingMetrics wiring
 templates:
   - deployment.yaml
   - configmap-litellm.yaml
-  - migrations-job.yaml
 tests:
-  - it: is off by default, adding no env, volume, or mount
+  - it: is disabled by default
     template: deployment.yaml
     asserts:
-      - notContains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-      - notContains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
       - notContains:
           path: spec.template.spec.containers[0].env
           content:
             name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
+      - notContains:
+          path: spec.template.spec.volumes
+          content:
+            name: billing-metrics
 
-  - it: renders the endpoint and the mounted cert paths when enabled
+  - it: renders required URLs and Cognito secret references when enabled
     template: deployment.yaml
     set:
       billingMetrics:
@@ -39,119 +29,34 @@ tests:
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_CERT
-            value: /etc/litellm/billing-mtls/tls.crt
+            name: LITELLM_BILLING_METRICS_TOKEN_URL
+            value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_KEY
-            value: /etc/litellm/billing-mtls/tls.key
-
-  # The conventional Secret name is the default, so enabling the block is enough.
-  - it: mounts the default cert secret read-only alongside the config volume
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-
-  - it: honours a secretName override
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: my-billing-mtls
-    asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: my-billing-mtls
-      - notContains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-
-  - it: honours an endpoint override
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        endpoint: https://collector.internal:4318
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_ID
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-id
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://collector.internal:4318
-
-  # The production collector presents a public web-PKI certificate, so the CA
-  # override must stay absent unless a private collector is configured.
-  - it: omits the CA env, volume, and mount when no caSecretName is set
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-secret
       - notContains:
           path: spec.template.spec.volumes
           content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
+            name: billing-metrics
       - notContains:
           path: spec.template.spec.containers[0].volumeMounts
           content:
-            name: billing-metrics-mtls-ca
-            mountPath: /etc/litellm/billing-mtls-ca
-            readOnly: true
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
+            name: billing-metrics
 
-  - it: mounts the CA secret when caSecretName is set
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        caSecretName: billing-ca
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls-ca
-            mountPath: /etc/litellm/billing-mtls-ca
-            readOnly: true
-
-  - it: passes the export interval through only when set
+  - it: renders the export interval only when set
     template: deployment.yaml
     set:
       billingMetrics:
@@ -164,129 +69,7 @@ tests:
             name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
             value: "5000"
 
-  - it: omits the export interval when unset
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
-            value: "60000"
-
-  # Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
-  # entries must render after .Values.envVars or a user could silently redirect
-  # the metering export. The three billing entries are the last ones emitted here
-  # (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
-  - it: renders the billing endpoint after envVars so it cannot be shadowed
-    template: deployment.yaml
-    set:
-      migrationJob:
-        enabled: false
-      billingMetrics:
-        enabled: true
-      envVars:
-        LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://shadowed.example
-      - equal:
-          path: spec.template.spec.containers[0].env[-3]
-          value:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - equal:
-          path: spec.template.spec.containers[0].env[-2].name
-          value: LITELLM_BILLING_METRICS_CLIENT_CERT
-      - equal:
-          path: spec.template.spec.containers[0].env[-1].name
-          value: LITELLM_BILLING_METRICS_CLIENT_KEY
-
-  - it: keeps user-supplied volumes and mounts alongside the billing secret
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-      volumes:
-        - name: custom-callbacks
-          configMap:
-            name: my-callbacks
-      volumeMounts:
-        - name: custom-callbacks
-          mountPath: /app/callbacks
-    asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: custom-callbacks
-            configMap:
-              name: my-callbacks
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: custom-callbacks
-            mountPath: /app/callbacks
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-
-  - it: still mounts the proxy config when enabled
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: litellm-config
-            mountPath: /etc/litellm/config.yaml
-            subPath: config.yaml
-
-  # Only the proxy serves billable traffic. The migrations Job must never mount
-  # the client certificate, and it renders its own env and volumes, so nothing
-  # stops a future edit from wiring the billing include into it by mistake.
-  - it: does not touch the migrations job when enabled
-    template: migrations-job.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - notExists:
-          path: spec.template.spec.containers[0].volumeMounts
-      - notExists:
-          path: spec.template.spec.volumes
-
-  - it: fails loudly when enabled with an emptied secretName
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: ""
-    asserts:
-      - failedTemplate:
-          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
-
-  - it: fails loudly when enabled without an endpoint
+  - it: requires the endpoint when enabled
     template: deployment.yaml
     set:
       billingMetrics:
@@ -295,3 +78,23 @@ tests:
     asserts:
       - failedTemplate:
           errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
+
+  - it: requires the token URL when enabled
+    template: deployment.yaml
+    set:
+      billingMetrics:
+        enabled: true
+        tokenUrl: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
+
+  - it: requires the secret name when enabled
+    template: deployment.yaml
+    set:
+      billingMetrics:
+        enabled: true
+        secretName: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true

helm/litellm-helm/tests/collector_tests.yaml

@@ -191,7 +191,7 @@ tests:
       metricsServer.port: 9090
       billingMetrics.enabled: true
       billingMetrics.endpoint: https://metering.example.com
-      billingMetrics.secretName: billing-mtls
+      billingMetrics.secretName: billing-metrics
     asserts:
       - contains:
           path: spec.template.spec.containers[0].env
@@ -214,9 +214,9 @@ tests:
             name: LITELLM_BILLING_METRICS_ENDPOINT
           any: true
       - notContains:
-          path: spec.template.spec.containers[1].volumeMounts
+          path: spec.template.spec.containers[1].env
           content:
-            name: billing-metrics-mtls
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
           any: true
 
   - it: should give the sidecar the same scratch mounts as the proxy on a read-only root

helm/litellm-helm/values.yaml

@@ -154,19 +154,13 @@ masterkeySecretName: ""
 # if set, use this secret key for the master key; otherwise, use the default key
 masterkeySecretKey: ""
 
-# Optional: enterprise billable-request metering. When enabled, the proxy counts
-# successful requests to inference, MCP, and A2A endpoints and pushes them to
-# LiteLLM's collector over mutual TLS. Requires an enterprise license.
-# The client certificate identifies the deployment, so it is mounted read-only
-# from an existing Secret and never passed through the environment.
+# Optional: enterprise billable-request metering. Requires an enterprise license.
 billingMetrics:
   enabled: false
-  endpoint: https://telemetry.litellm.ai   # collector to push the counter to
-  secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
-  # Only for private or test collectors whose server certificate is not on the
-  # public web PKI. The production collector needs no CA override.
-  caSecretName: ""                         # existing Secret holding ca.crt
-  exportIntervalMs: ""                     # push cadence; the proxy defaults to 60000
+  endpoint: https://telemetry.litellm.ai
+  tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
+  secretName: litellm-billing-metrics
+  exportIntervalMs: ""
 
 proxyConfigMap:
   # when true, creates a new configmap

helm/litellm/Chart.yaml

@@ -2,7 +2,7 @@ apiVersion: v2
 name: litellm
 description: LiteLLM componentized — gateway, UI backend, and UI as separate services
 type: application
-version: 0.1.0
+version: 0.1.1
 appVersion: "0.1.0"
 annotations:
   org.opencontainers.image.source: "https://github.com/BerriAI/litellm"

helm/litellm/README.md

@@ -22,6 +22,23 @@ backend:
   gatewayUrl: https://gateway.example.com
 ```
 
+## Enterprise billing metrics
+
+Billing metrics use Cognito bearer authentication. Create the Secret with
+`client-id` and `client-secret` keys, then enable billing metrics:
+
+```bash
+kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
+```
+
+```yaml
+billingMetrics:
+  enabled: true
+```
+
+The chart requires an HTTPS endpoint, an HTTPS token URL, and the Secret name
+when billing metrics are enabled
+
 ## Routes to another Service
 
 For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields

helm/litellm/templates/_helpers.tpl

@@ -58,56 +58,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
 helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
 {{- end -}}
 
-{{/*
-Enterprise billable-request metering. Wired into gateway and backend, not the
-migrations job. The gateway serves nearly all billable traffic, but the backend
-keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
-SpendLogs row, so metering only the gateway would silently drop that traffic.
-The client certificate identifies the deployment to LiteLLM's collector, so it is
-mounted read-only from an existing Secret rather than passed through the
-environment.
-*/}}
-{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
-{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
-
 {{- define "litellm.billingMetricsEnv" -}}
 - name: LITELLM_BILLING_METRICS_ENDPOINT
   value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_CERT
-  value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_KEY
-  value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: LITELLM_BILLING_METRICS_CA_CERT
-  value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
-{{- end }}
+- name: LITELLM_BILLING_METRICS_TOKEN_URL
+  value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
+- name: LITELLM_BILLING_METRICS_CLIENT_ID
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-id
+- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-secret
 {{- with .Values.billingMetrics.exportIntervalMs }}
 - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
   value: {{ . | quote }}
 {{- end }}
 {{- end -}}
 
-{{- define "litellm.billingMetricsVolumes" -}}
-- name: billing-metrics-mtls
-  secret:
-    secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  secret:
-    secretName: {{ .Values.billingMetrics.caSecretName }}
-{{- end }}
-{{- end -}}
-
-{{- define "litellm.billingMetricsVolumeMounts" -}}
-- name: billing-metrics-mtls
-  mountPath: {{ include "litellm.billingMetrics.certDir" . }}
-  readOnly: true
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  mountPath: {{ include "litellm.billingMetrics.caDir" . }}
-  readOnly: true
-{{- end }}
-{{- end -}}
 
 {{/*
 Per-component selector labels — used in both Service selectors and Deployment matchLabels.

helm/litellm/templates/backend/deployment.yaml

@@ -79,16 +79,13 @@ spec:
             {{- include "litellm.billingMetricsEnv" . | nindent 12 }}
             {{- end }}
           {{- include "litellm.envFrom" .Values.backend | nindent 10 }}
-          {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
+          {{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
           volumeMounts:
             {{- if .Values.gateway.config.create }}
             - name: gateway-config
               mountPath: /app/config/config.yaml
               subPath: config.yaml
             {{- end }}
-            {{- if .Values.billingMetrics.enabled }}
-            {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
-            {{- end }}
             {{- with .Values.backend.volumeMounts }}
             {{- toYaml . | nindent 12 }}
             {{- end }}
@@ -114,16 +111,13 @@ spec:
       {{- with .Values.backend.extraContainers }}
         {{- tpl (toYaml .) $ | nindent 8 }}
       {{- end }}
-      {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
+      {{- if or .Values.gateway.config.create .Values.backend.volumes }}
       volumes:
         {{- if .Values.gateway.config.create }}
         - name: gateway-config
           configMap:
             name: {{ include "litellm.gateway.fullname" . }}-config
         {{- end }}
-        {{- if .Values.billingMetrics.enabled }}
-        {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
-        {{- end }}
         {{- with .Values.backend.volumes }}
         {{- toYaml . | nindent 8 }}
         {{- end }}

helm/litellm/templates/gateway/deployment.yaml

@@ -86,7 +86,7 @@ spec:
             {{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
             {{- end }}
           {{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
-          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
+          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
           volumeMounts:
             {{- if .Values.gateway.config.create }}
             - name: gateway-config
@@ -101,9 +101,6 @@ spec:
             - name: collector-socket
               mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
             {{- end }}
-            {{- if .Values.billingMetrics.enabled }}
-            {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
-            {{- end }}
             {{- with .Values.gateway.volumeMounts }}
             {{- toYaml . | nindent 12 }}
             {{- end }}
@@ -206,7 +203,7 @@ spec:
       {{- with .Values.gateway.extraContainers }}
         {{- tpl (toYaml .) $ | nindent 8 }}
       {{- end }}
-      {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
+      {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
       volumes:
         {{- if .Values.gateway.config.create }}
         - name: gateway-config
@@ -222,9 +219,6 @@ spec:
           emptyDir:
             sizeLimit: 1Mi
         {{- end }}
-        {{- if .Values.billingMetrics.enabled }}
-        {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
-        {{- end }}
         {{- with .Values.gateway.volumes }}
         {{- toYaml . | nindent 8 }}
         {{- end }}

helm/litellm/tests/billing_metrics_tests.yaml

@@ -1,39 +1,30 @@
-suite: test billingMetrics wiring on gateway and backend
+suite: test bearer-only billingMetrics wiring
 templates:
   - gateway/deployment.yaml
   - gateway/configmap.yaml
   - backend/deployment.yaml
+  - backend/configmap.yaml
   - migrations-job.yaml
 values:
   - ./values/required.yaml
 tests:
-  - it: is off by default, adding no env, volume, or mount
+  - it: is disabled by default
     template: gateway/deployment.yaml
     asserts:
-      - notContains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
-      - equal:
-          path: spec.template.spec.containers[0].volumeMounts
-          value:
-            - name: gateway-config
-              mountPath: /app/config/config.yaml
-              subPath: config.yaml
       - notContains:
           path: spec.template.spec.containers[0].env
           content:
             name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
+      - notContains:
+          path: spec.template.spec.volumes
+          content:
+            name: billing-metrics
 
-  - it: renders the endpoint and the mounted cert paths when enabled
+  - it: wires gateway URLs and bearer credentials from Secret keys
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
     asserts:
       - contains:
           path: spec.template.spec.containers[0].env
@@ -43,87 +34,68 @@ tests:
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_CERT
-            value: /etc/litellm/billing-mtls/tls.crt
+            name: LITELLM_BILLING_METRICS_TOKEN_URL
+            value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_KEY
-            value: /etc/litellm/billing-mtls/tls.key
-
-  - it: mounts the cert secret read-only alongside the config volume
-    template: gateway/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_ID
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-id
       - contains:
-          path: spec.template.spec.volumes
+          path: spec.template.spec.containers[0].env
           content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-
-  # The production collector presents a public web-PKI certificate, so the CA
-  # override must stay absent unless a private collector is configured.
-  - it: omits the CA env, volume, and mount when no caSecretName is set
-    template: gateway/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-secret
       - notContains:
           path: spec.template.spec.volumes
           content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
+            name: billing-metrics
       - notContains:
-          path: spec.template.spec.containers[0].env
+          path: spec.template.spec.containers[0].volumeMounts
           content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
+            name: billing-metrics
 
-  - it: mounts the CA secret when caSecretName is set
-    template: gateway/deployment.yaml
+  - it: wires billing metrics on the backend without volumes
+    template: backend/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
-        caSecretName: billing-ca
     asserts:
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
-      - contains:
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-secret
+      - notContains:
           path: spec.template.spec.volumes
           content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls-ca
-            mountPath: /etc/litellm/billing-mtls-ca
-            readOnly: true
+            name: billing-metrics
 
-  - it: passes the export interval through only when set
+  - it: does not add billing metrics to migrations
+    template: migrations-job.yaml
+    set:
+      billingMetrics:
+        enabled: true
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_BILLING_METRICS_ENDPOINT
+
+  - it: renders an optional export interval
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
         exportIntervalMs: 5000
     asserts:
       - contains:
@@ -132,102 +104,27 @@ tests:
             name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
             value: "5000"
 
-  - it: keeps user-supplied gateway volumes alongside the billing secret
+  - it: requires the endpoint when enabled
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
-      gateway.volumes:
-        - name: custom-callbacks
-          configMap:
-            name: my-callbacks
-      gateway.volumeMounts:
-        - name: custom-callbacks
-          mountPath: /app/callbacks
+        endpoint: ""
     asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: custom-callbacks
-            configMap:
-              name: my-callbacks
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
+      - failedTemplate:
+          errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
 
-  # The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
-  # which writes a SpendLogs row, so it must meter too or that traffic is lost.
-  - it: meters the backend as well, since it serves the MCP transport
-    template: backend/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
-
-  - it: leaves the backend alone when metering is off
-    template: backend/deployment.yaml
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-
-  # The migrations job runs prisma and serves no traffic; it must never receive
-  # the client key.
-  - it: never mounts the billing cert on the migrations job
-    template: migrations-job.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - isNull:
-          path: spec.template.spec.volumes
-
-  # The conventional Secret name is the default, so enabling metering needs no
-  # secretName at all; the guard below only fires on an explicitly blanked one.
-  - it: uses the conventional secret name by default
+  - it: requires the token URL when enabled
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
+        tokenUrl: ""
     asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
+      - failedTemplate:
+          errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
 
-  - it: fails loudly when the secretName is explicitly blanked
+  - it: requires the secret name when enabled
     template: gateway/deployment.yaml
     set:
       billingMetrics:
@@ -235,15 +132,4 @@ tests:
         secretName: ""
     asserts:
       - failedTemplate:
-          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
-
-  - it: fails loudly when enabled without an endpoint
-    template: gateway/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        endpoint: ""
-        secretName: billing-mtls
-    asserts:
-      - failedTemplate:
-          errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
+          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true

helm/litellm/tests/migration_job_tests.yaml

@@ -91,7 +91,7 @@ tests:
             app.kubernetes.io/name: litellm
             app.kubernetes.io/instance: RELEASE-NAME
             app.kubernetes.io/managed-by: Helm
-            helm.sh/chart: litellm-0.1.0
+            helm.sh/chart: litellm-0.1.1
             app.kubernetes.io/component: migrations
 
   - it: renders pod-level and container-level securityContext in their own scopes

helm/litellm/values.yaml

@@ -180,23 +180,14 @@ masterKey:
   secretKey: master-key
 
 # Optional: enterprise billable-request metering. When enabled, the gateway and
-# backend count successful requests to inference, MCP, and A2A endpoints and push
-# them to LiteLLM's collector over mutual TLS. Both components serve billable
-# routes: the backend keeps the named-server MCP transport. Requires an
-# enterprise license. The client certificate identifies the deployment, so it is
-# mounted read-only from an existing Secret and never passed through the env.
+# backend count successful requests to inference, MCP, and A2A endpoints. Requires
+# an enterprise license.
 billingMetrics:
   enabled: false
-  endpoint: https://telemetry.litellm.ai              # collector to push the counter to
-  # An existing Secret holding the client certificate under tls.crt and its key
-  # under tls.key, usually created from the onboarding artifact. The default is
-  # the conventional name, so the common path is to create that Secret and set
-  # enabled: true. Override only if yours is named differently.
-  secretName: litellm-billing-metrics-mtls
-  # Only for private or test collectors whose server certificate is not on the
-  # public web PKI. The production collector needs no CA override.
-  caSecretName: ""                                    # existing Secret holding ca.crt
-  exportIntervalMs: ""                                # push cadence; the proxy defaults to 60000
+  endpoint: https://telemetry.litellm.ai
+  tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
+  secretName: litellm-billing-metrics
+  exportIntervalMs: ""
 
 # External Postgres connection. 
 database:

litellm-rust/Cargo.lock

@@ -3389,15 +3389,19 @@ version = "0.1.0"
 dependencies = [
  "aws-config",
  "aws-credential-types",
+ "aws-runtime",
  "aws-sdk-sts",
  "aws-sigv4",
  "aws-smithy-runtime-api",
+ "aws-smithy-types",
  "aws-types",
  "bytes",
  "futures-util",
  "litellm-auth-types",
  "litellm-error",
  "litellm-http",
+ "litellm-secret-value",
+ "litellm-secrets",
  "moka",
  "reqwest",
  "rstest",
@@ -3406,6 +3410,7 @@ dependencies = [
  "sha2 0.10.9",
  "thiserror 2.0.19",
  "tokio",
+ "url",
  "veil",
 ]
 
@@ -3418,6 +3423,7 @@ dependencies = [
  "litellm-auth-types",
  "litellm-error",
  "litellm-secret-value",
+ "litellm-secrets",
  "moka",
  "rstest",
  "serde_json",
@@ -3431,6 +3437,7 @@ dependencies = [
 name = "litellm-auth-gcp"
 version = "0.1.0"
 dependencies = [
+ "axum",
  "base64 0.22.1",
  "gcp_auth",
  "google-cloud-auth",
@@ -3440,10 +3447,12 @@ dependencies = [
  "litellm-secret-value",
  "moka",
  "rstest",
+ "rustls 0.23.42",
  "serde",
  "serde_json",
  "sha2 0.10.9",
  "tempfile",
+ "time",
  "tokio",
  "veil",
 ]

litellm-rust/crates/auth-aws/Cargo.toml

@@ -11,14 +11,16 @@ litellm-error.workspace = true
 litellm-http.workspace = true
 reqwest.workspace = true
 
-moka = { workspace = true, features = ["sync"] }
+moka.workspace = true
+litellm-secret-value.workspace = true
 serde.workspace = true
 serde_json = { workspace = true, features = ["preserve_order"] }
 sha2.workspace = true
 thiserror.workspace = true
+url.workspace = true
 veil.workspace = true
 
-aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
+aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
 aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
 aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
 aws-sigv4 = "1.5.1"
@@ -26,6 +28,9 @@ aws-types = "1.4.0"
 aws-smithy-runtime-api = "1.13.0"
 
 [dev-dependencies]
+litellm-secrets.workspace = true
+aws-smithy-types = "1.6.4"
+aws-runtime = { version = "1.9.4", default-features = false }
 rstest.workspace = true
 litellm-http = { workspace = true, features = ["test-support"] }
 bytes.workspace = true

litellm-rust/crates/auth-aws/src/constants.rs

@@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0];
 pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0];
 pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0];
 pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0];
-pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
-pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
+pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV;
+pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV;
 pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0];
 pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0];
 pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";

litellm-rust/crates/auth-aws/src/error.rs

@@ -2,6 +2,8 @@ use litellm_error::Source;
 
 #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
 pub enum Error {
+    #[error("AWS credential provider returned expired credentials")]
+    ExpiredCredentials,
     #[error("could not load AWS {flow} credentials")]
     Credentials {
         flow: CredentialFlow,

litellm-rust/crates/auth-aws/src/signer.rs

@@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer {
         )
         .map_err(|error| match error {
             Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
-            Error::Credentials { .. }
+            Error::ExpiredCredentials
+            | Error::Credentials { .. }
             | Error::MissingWebIdentityCredentials
             | Error::WebIdentityExpiration(_)
             | Error::Signing { .. } => litellm_http::Error::signature(error),

litellm-rust/crates/auth-azure/Cargo.toml

@@ -8,17 +8,19 @@ repository.workspace = true
 [dependencies]
 litellm-secret-value.workspace = true
 litellm-auth-types.workspace = true
+litellm-error.workspace = true
 
 moka.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
 strum.workspace = true
+tokio.workspace = true
 url.workspace = true
 
 azure_core = "1.0.0"
 azure_identity = { version = "1.0.0", features = ["tokio"] }
 
 [dev-dependencies]
-litellm-error.workspace = true
+litellm-secrets.workspace = true
 rstest.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["test-util"] }

litellm-rust/crates/auth-azure/src/lib.rs

@@ -3,5 +3,5 @@ mod native;
 mod resolve;
 mod types;
 
-pub use resolve::{AzureAuthService, SECRET_NAMES};
+pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES};
 pub use types::{AzureAuthInputs, ConfigValue};

litellm-rust/crates/auth-azure/src/native.rs

@@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest {
 pub(crate) struct ValidatedAzureRequest {
     request: NativeAzureRequest,
     credential_source: InputSource,
+    probe_imds: bool,
 }
 
 impl ValidatedAzureRequest {
@@ -68,9 +69,18 @@ impl ValidatedAzureRequest {
         Ok(Self {
             request,
             credential_source,
+            probe_imds: false,
         })
     }
 
+    pub(crate) fn with_imds_probe(self, enabled: bool) -> Self {
+        Self {
+            probe_imds: enabled
+                && matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }),
+            ..self
+        }
+    }
+
     pub(crate) fn credential_source(&self) -> InputSource {
         self.credential_source
     }
@@ -87,9 +97,34 @@ impl ValidatedAzureRequest {
     }
 }
 
+pub(crate) enum CredentialAttempt {
+    Resolved(ResolvedCredential),
+    Unavailable(Error),
+    Failed(Error),
+}
+
+impl CredentialAttempt {
+    pub(crate) fn into_result(self) -> Result<ResolvedCredential, Error> {
+        match self {
+            Self::Resolved(credential) => Ok(credential),
+            Self::Unavailable(error) | Self::Failed(error) => Err(error),
+        }
+    }
+}
+
 pub(crate) struct NativeAzureTokenAcquirer {
     cache: AzureCredentialProviderCache,
     transport: Option<azure_core::http::Transport>,
+    imds_available: tokio::sync::Mutex<ImdsAvailability>,
+}
+
+enum ImdsAvailability {
+    Unknown,
+    Available,
+    Unavailable {
+        retry_at: tokio::time::Instant,
+        error: Error,
+    },
 }
 
 impl Default for NativeAzureTokenAcquirer {
@@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer {
         Self {
             cache: AzureCredentialProviderCache::new(cache_capacity),
             transport: None,
+            imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
         }
     }
 
@@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer {
         Self {
             cache: AzureCredentialProviderCache::new(cache_capacity),
             transport: Some(transport),
+            imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
         }
     }
 
-    pub(crate) async fn acquire(
+    pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt {
+        let discover_imds = request.probe_imds;
+        if discover_imds && let Err(error) = self.discover_imds().await {
+            return CredentialAttempt::Unavailable(error);
+        }
+        match self.acquire_token(request).await {
+            Ok(credential) => CredentialAttempt::Resolved(credential),
+            Err(error) => {
+                let missing_identity = discover_imds && unassigned_identity(&error);
+                if missing_identity {
+                    CredentialAttempt::Unavailable(error)
+                } else {
+                    CredentialAttempt::Failed(error)
+                }
+            }
+        }
+    }
+
+    async fn discover_imds(&self) -> Result<(), Error> {
+        let mut availability = self.imds_available.lock().await;
+        match &*availability {
+            ImdsAvailability::Available => return Ok(()),
+            ImdsAvailability::Unavailable { retry_at, error }
+                if tokio::time::Instant::now() < *retry_at =>
+            {
+                return Err(error.clone());
+            }
+            ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {}
+        }
+        let result = probe_imds(self.transport.clone().unwrap_or_default()).await;
+        *availability = match &result {
+            Ok(()) => ImdsAvailability::Available,
+            Err(error) => ImdsAvailability::Unavailable {
+                retry_at: tokio::time::Instant::now() + Duration::from_secs(30),
+                error: error.clone(),
+            },
+        };
+        result
+    }
+
+    async fn acquire_token(
         &self,
         request: ValidatedAzureRequest,
     ) -> Result<ResolvedCredential, Error> {
@@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer {
     }
 }
 
+fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool {
+    std::iter::successors(Some(error), |cause| cause.source()).any(|cause| {
+        cause
+            .downcast_ref::<azure_core::Error>()
+            .or_else(|| {
+                cause
+                    .downcast_ref::<litellm_error::Source>()?
+                    .downcast_ref::<azure_core::Error>()
+            })
+            .is_some_and(|error| {
+                error.http_status() == Some(azure_core::http::StatusCode::BadRequest)
+            })
+    })
+}
+
+async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> {
+    use azure_core::http::{Context, Method, Request, Url};
+
+    let endpoint =
+        Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01")
+            .map_err(|error| {
+                Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error))
+            })?;
+    let mut request = Request::new(endpoint, Method::Get);
+    let response = tokio::time::timeout(
+        Duration::from_secs(1),
+        transport.send(&Context::new(), &mut request),
+    )
+    .await
+    .map_err(|error| {
+        Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
+    })?
+    .map_err(|error| {
+        Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
+    })?;
+    // Headerless IMDS discovery expects 400, checked 2026-11-01
+    // https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging
+    if response.status() != azure_core::http::StatusCode::BadRequest {
+        return Err(Error::CredentialAcquisition(
+            ErrorDetail::UnexpectedStatus {
+                operation: "Azure IMDS availability probe",
+                status: u16::from(response.status()),
+            },
+        ));
+    }
+    Ok(())
+}
+
 impl NativeAzureRequest {
     fn scope(&self) -> &str {
         match self {
@@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion {
 #[cfg(test)]
 mod tests {
     use std::sync::{Arc, Mutex};
+    use std::time::Duration;
 
     use azure_core::http::headers::Headers;
     use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport};
@@ -569,6 +695,296 @@ mod tests {
         }
     }
 
+    #[derive(Debug)]
+    struct AvailableImdsClient {
+        probe_status: StatusCode,
+        probe_delay: Duration,
+        recover_after: Option<usize>,
+        response_status: Option<StatusCode>,
+        probes: std::sync::atomic::AtomicUsize,
+    }
+
+    impl Default for AvailableImdsClient {
+        fn default() -> Self {
+            Self {
+                probe_status: StatusCode::BadRequest,
+                probe_delay: Duration::ZERO,
+                recover_after: None,
+                response_status: None,
+                probes: std::sync::atomic::AtomicUsize::new(0),
+            }
+        }
+    }
+
+    impl HttpClient for AvailableImdsClient {
+        fn execute_request<'life0, 'life1, 'async_trait>(
+            &'life0 self,
+            request: &'life1 Request,
+        ) -> std::pin::Pin<
+            Box<dyn std::future::Future<Output = Result<AsyncRawResponse>> + Send + 'async_trait>,
+        >
+        where
+            'life0: 'async_trait,
+            'life1: 'async_trait,
+            Self: 'async_trait,
+        {
+            Box::pin(async move {
+                if request
+                    .headers()
+                    .get_optional_str(&azure_core::http::headers::HeaderName::from_static(
+                        "metadata",
+                    ))
+                    .is_none()
+                {
+                    let probe = self
+                        .probes
+                        .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                    if self.recover_after.is_some_and(|failures| probe >= failures) {
+                        return Ok(AsyncRawResponse::from_bytes(
+                            StatusCode::BadRequest,
+                            Headers::new(),
+                            "metadata header required",
+                        ));
+                    }
+                    tokio::time::sleep(self.probe_delay).await;
+                    return Ok(AsyncRawResponse::from_bytes(
+                        self.probe_status,
+                        Headers::new(),
+                        "metadata header required",
+                    ));
+                }
+                if let Some(status) = self.response_status {
+                    return Ok(AsyncRawResponse::from_bytes(
+                        status,
+                        Headers::new(),
+                        "identity endpoint failure",
+                    ));
+                }
+                tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+                let expires = std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .unwrap()
+                    .as_secs()
+                    + 3600;
+                Ok(AsyncRawResponse::from_bytes(
+                    StatusCode::Ok,
+                    Headers::new(),
+                    serde_json::json!({
+                        "access_token": "managed-token", "expires_on": expires.to_string(),
+                        "token_type": "Bearer", "resource": "https://service.test",
+                    })
+                    .to_string(),
+                ))
+            })
+        }
+    }
+
+    #[rstest::rstest]
+    #[tokio::test(start_paused = true)]
+    async fn available_imds_retains_full_sdk_token_acquisition() {
+        let acquirer = NativeAzureTokenAcquirer::with_transport(
+            4,
+            Transport::new(Arc::new(AvailableImdsClient {
+                probe_status: StatusCode::BadRequest,
+                response_status: None,
+                ..Default::default()
+            })),
+        );
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true);
+        let credential = acquirer.acquire(request).await.into_result().unwrap();
+        assert_eq!(credential.secret().expose(), "managed-token");
+    }
+
+    #[rstest::rstest]
+    #[case::another_metadata_service(StatusCode::NotFound)]
+    #[case::token_required(StatusCode::Unauthorized)]
+    #[case::proxy_rejection(StatusCode::Forbidden)]
+    #[case::unavailable(StatusCode::ServiceUnavailable)]
+    #[tokio::test(start_paused = true)]
+    async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition(
+        #[case] status: StatusCode,
+    ) {
+        let client = Arc::new(AvailableImdsClient {
+            probe_status: status,
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true);
+        assert!(matches!(
+            acquirer.acquire(request.clone()).await,
+            super::CredentialAttempt::Unavailable(_)
+        ));
+        assert!(matches!(
+            acquirer.acquire(request).await,
+            super::CredentialAttempt::Unavailable(_)
+        ));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+
+    #[rstest::rstest]
+    #[tokio::test(start_paused = true)]
+    async fn an_imds_timeout_is_not_repeated_on_the_next_request() {
+        let client = Arc::new(AvailableImdsClient {
+            probe_delay: Duration::from_secs(10),
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true);
+        assert!(matches!(
+            acquirer.acquire(request.clone()).await,
+            super::CredentialAttempt::Unavailable(_)
+        ));
+        let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request))
+            .await
+            .unwrap();
+        assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+
+    #[rstest::fixture]
+    fn imds_request() -> ValidatedAzureRequest {
+        ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true)
+    }
+
+    #[rstest::rstest]
+    #[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))]
+    #[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)]
+    #[tokio::test(start_paused = true)]
+    async fn failed_imds_discovery_recovers_without_restarting_the_acquirer(
+        imds_request: ValidatedAzureRequest,
+        #[case] status: StatusCode,
+        #[case] delay: Duration,
+    ) {
+        let client = Arc::new(AvailableImdsClient {
+            probe_status: status,
+            probe_delay: delay,
+            recover_after: Some(1),
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        let (first, second) = tokio::join!(
+            acquirer.acquire(imds_request.clone()),
+            acquirer.acquire(imds_request.clone()),
+        );
+        assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
+        assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
+
+        tokio::time::advance(Duration::from_secs(30)).await;
+        let (first, second) = tokio::join!(
+            acquirer.acquire(imds_request.clone()),
+            acquirer.acquire(imds_request.clone()),
+        );
+        assert_eq!(
+            first.into_result().unwrap().secret().expose(),
+            "managed-token"
+        );
+        assert_eq!(
+            second.into_result().unwrap().secret().expose(),
+            "managed-token"
+        );
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
+
+        tokio::time::advance(Duration::from_secs(60)).await;
+        assert_eq!(
+            acquirer
+                .acquire(imds_request)
+                .await
+                .into_result()
+                .unwrap()
+                .secret()
+                .expose(),
+            "managed-token",
+        );
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
+    }
+
+    #[rstest::rstest]
+    #[tokio::test(start_paused = true)]
+    async fn a_persistent_imds_outage_reuses_each_failed_probe(
+        imds_request: ValidatedAzureRequest,
+    ) {
+        let client = Arc::new(AvailableImdsClient {
+            probe_delay: Duration::from_secs(10),
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        assert!(matches!(
+            acquirer.acquire(imds_request.clone()).await,
+            super::CredentialAttempt::Unavailable(_),
+        ));
+        tokio::time::advance(Duration::from_secs(30)).await;
+        let (first, second) = tokio::join!(
+            acquirer.acquire(imds_request.clone()),
+            acquirer.acquire(imds_request),
+        );
+        assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
+        assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
+    }
+
+    #[rstest::rstest]
+    #[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)]
+    #[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)]
+    #[case::default_identity_forbidden(true, StatusCode::Forbidden, false)]
+    #[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)]
+    #[tokio::test(start_paused = true)]
+    async fn only_unavailable_default_identities_allow_fallback(
+        #[case] discovery: bool,
+        #[case] status: StatusCode,
+        #[case] unavailable: bool,
+    ) {
+        let acquirer = NativeAzureTokenAcquirer::with_transport(
+            4,
+            Transport::new(Arc::new(AvailableImdsClient {
+                probe_status: StatusCode::BadRequest,
+                response_status: Some(status),
+                ..Default::default()
+            })),
+        );
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(discovery);
+        let outcome = acquirer.acquire(request).await;
+        let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_));
+        let error = outcome.into_result().unwrap_err();
+        assert_eq!(
+            is_unavailable,
+            unavailable,
+            "{}",
+            litellm_error::Report::new(&error)
+        );
+    }
+
+    #[rstest::rstest]
     #[tokio::test]
     async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() {
         let transport = Arc::new(RecordingTokenClient::default());
@@ -582,8 +998,12 @@ mod tests {
             "https://login.test",
         );
 
-        let first = acquirer.acquire(request.clone()).await.unwrap();
-        let second = acquirer.acquire(request).await.unwrap();
+        let first = acquirer
+            .acquire(request.clone())
+            .await
+            .into_result()
+            .unwrap();
+        let second = acquirer.acquire(request).await.into_result().unwrap();
 
         assert_eq!(first.secret().expose(), "native-token");
         assert_eq!(second.secret().expose(), "native-token");
@@ -599,6 +1019,7 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
     #[tokio::test]
     async fn credential_provider_cache_isolates_every_client_secret_identity_field() {
         let transport = Arc::new(RecordingTokenClient::default());
@@ -645,10 +1066,10 @@ mod tests {
             ),
         ];
 
-        acquirer.acquire(base.clone()).await.unwrap();
-        acquirer.acquire(base).await.unwrap();
+        acquirer.acquire(base.clone()).await.into_result().unwrap();
+        acquirer.acquire(base).await.into_result().unwrap();
         for request in variants.into_iter().skip(1) {
-            acquirer.acquire(request).await.unwrap();
+            acquirer.acquire(request).await.into_result().unwrap();
         }
 
         assert_eq!(transport.requests.lock().unwrap().len(), 6);

litellm-rust/crates/auth-azure/src/resolve.rs

@@ -1,3 +1,4 @@
+use crate::native::CredentialAttempt;
 use litellm_auth_types::{
     CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced,
     TokenProviderHandle,
@@ -12,25 +13,31 @@ use std::sync::Arc;
 use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
 use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE};
 
-const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
-const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID";
-const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID";
-const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET";
-const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE";
-const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST";
-const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL";
-const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE";
+#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)]
+enum AzureSecretName {
+    #[strum(serialize = "AZURE_AD_TOKEN")]
+    AdToken,
+    #[strum(serialize = "AZURE_TENANT_ID")]
+    TenantId,
+    #[strum(serialize = "AZURE_CLIENT_ID")]
+    ClientId,
+    #[strum(serialize = "AZURE_CLIENT_SECRET")]
+    ClientSecret,
+    #[strum(serialize = "AZURE_SCOPE")]
+    Scope,
+    #[strum(serialize = "AZURE_AUTHORITY_HOST")]
+    AuthorityHost,
+    #[strum(serialize = "AZURE_CREDENTIAL")]
+    Credential,
+    #[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")]
+    FederatedTokenFile,
+    #[strum(serialize = "IDENTITY_ENDPOINT")]
+    IdentityEndpoint,
+    #[strum(serialize = "MSI_ENDPOINT")]
+    MsiEndpoint,
+}
 
-pub const SECRET_NAMES: &[&str] = &[
-    AZURE_AD_TOKEN_ENV,
-    AZURE_TENANT_ID_ENV,
-    AZURE_CLIENT_ID_ENV,
-    AZURE_CLIENT_SECRET_ENV,
-    AZURE_SCOPE_ENV,
-    AZURE_AUTHORITY_HOST_ENV,
-    AZURE_CREDENTIAL_ENV,
-    AZURE_FEDERATED_TOKEN_FILE_ENV,
-];
+pub const SECRET_NAMES: &[&str] = <AzureSecretName as strum::VariantNames>::VARIANTS;
 
 #[derive(Clone, Debug)]
 pub(crate) enum AzureCredentialPlan {
@@ -57,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync {
     fn acquire(
         &self,
         request: ValidatedAzureRequest,
-    ) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
+    ) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>>;
 }
 
 impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
     fn acquire(
         &self,
         request: ValidatedAzureRequest,
-    ) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
+    ) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
         Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
     }
 }
@@ -83,12 +90,36 @@ impl AzureAuthService {
         Self { native }
     }
 
+    pub async fn identity_token(
+        &self,
+        identity: &AzureIdentity,
+        scope: &str,
+        env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
+    ) -> Result<Option<ResolvedCredential>, Error> {
+        let plan = identity_plan(identity, scope, env_lookup)?;
+        let inputs = AzureAuthInputs::default_credential_for_scope(scope);
+        Ok(self
+            .run_plan(plan, &inputs, env_lookup)
+            .await?
+            .map(Sourced::into_value))
+    }
+
     pub async fn get_azure_ad_token(
         &self,
         inputs: &AzureAuthInputs,
         env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
     ) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
-        match select_auth_plan(inputs, env_lookup)? {
+        self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup)
+            .await
+    }
+
+    async fn run_plan(
+        &self,
+        plan: AzureCredentialPlan,
+        inputs: &AzureAuthInputs,
+        env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
+    ) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
+        match plan {
             AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
             AzureCredentialPlan::Caller(caller) => {
                 let credential = caller.acquire().await?;
@@ -123,6 +154,7 @@ impl AzureAuthService {
                 self.native
                     .acquire(request)
                     .await
+                    .into_result()
                     .map(|credential| Sourced::new(credential, source))
                     .map(Some)
             }
@@ -131,6 +163,7 @@ impl AzureAuthService {
                 self.native
                     .acquire(request)
                     .await
+                    .into_result()
                     .map(|credential| Some(Sourced::new(credential, source)))
             }
             AzureCredentialPlan::Chain(requests) => {
@@ -138,8 +171,16 @@ impl AzureAuthService {
                 for request in requests {
                     let source = request.credential_source();
                     match self.native.acquire(request).await {
-                        Ok(credential) => return Ok(Some(Sourced::new(credential, source))),
-                        Err(error) => failures.push(error),
+                        CredentialAttempt::Resolved(credential) => {
+                            return Ok(Some(Sourced::new(credential, source)));
+                        }
+                        CredentialAttempt::Unavailable(error) => failures.push(error),
+                        CredentialAttempt::Failed(error) => {
+                            return Err(Error::CredentialChain {
+                                last: Box::new(error),
+                                earlier: failures,
+                            });
+                        }
                     }
                 }
                 match failures.pop() {
@@ -155,39 +196,161 @@ impl AzureAuthService {
     }
 }
 
+/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings
+/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads.
+#[derive(Clone, Debug)]
+pub enum AzureIdentity {
+    ClientSecret {
+        tenant_id: String,
+        client_id: String,
+        client_secret: SecretValue,
+    },
+    ManagedIdentity {
+        client_id: String,
+    },
+    Default,
+}
+
+fn identity_plan(
+    identity: &AzureIdentity,
+    scope: &str,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> Result<AzureCredentialPlan, Error> {
+    let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment);
+    let environment = |name: &str| {
+        env_lookup(name)
+            .filter(|value| !value.trim().is_empty())
+            .map(|value| Sourced::new(value, InputSource::Environment))
+    };
+    let scope = deployment(scope);
+    let authority = environment(AzureSecretName::AuthorityHost.into());
+    match identity {
+        AzureIdentity::ClientSecret {
+            tenant_id,
+            client_id,
+            client_secret,
+        } => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
+            NativeAzureRequest::ClientSecret {
+                tenant_id: deployment(tenant_id),
+                client_id: deployment(client_id),
+                client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment),
+                scope,
+                authority,
+            },
+        )?)),
+        AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native(
+            ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+                client_id: Some(deployment(client_id)),
+                scope,
+                selection_source: InputSource::Deployment,
+            })?,
+        )),
+        AzureIdentity::Default => {
+            let tenant_id = environment(AzureSecretName::TenantId.into());
+            let client_id = environment(AzureSecretName::ClientId.into());
+            let from_environment = match (
+                tenant_id.clone(),
+                client_id.clone(),
+                environment(AzureSecretName::ClientSecret.into()),
+            ) {
+                (Some(tenant_id), Some(client_id), Some(client_secret)) => {
+                    Some(NativeAzureRequest::ClientSecret {
+                        tenant_id,
+                        client_id,
+                        client_secret: client_secret.map(SecretValue::new),
+                        scope: scope.clone(),
+                        authority: authority.clone(),
+                    })
+                }
+                _ => None,
+            };
+            let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
+                && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
+            let AzureCredentialPlan::Chain(chain) = select_native_plan(
+                Some(Sourced::new(
+                    AzureCredentialType::DefaultAzureCredential,
+                    InputSource::Deployment,
+                )),
+                tenant_id,
+                client_id,
+                environment(AzureSecretName::FederatedTokenFile.into()),
+                scope,
+                authority,
+                InputSource::Environment,
+            )?
+            else {
+                unreachable!("DefaultAzureCredential always plans a chain")
+            };
+            let environment_request = from_environment
+                .map(ValidatedAzureRequest::new)
+                .transpose()?;
+            Ok(AzureCredentialPlan::Chain(
+                environment_request
+                    .into_iter()
+                    .chain(chain)
+                    .map(|request| request.with_imds_probe(probe_imds))
+                    .collect(),
+            ))
+        }
+    }
+}
+
 pub(crate) fn select_auth_plan(
     inputs: &AzureAuthInputs,
     env_lookup: &dyn Fn(&str) -> Option<String>,
 ) -> Result<AzureCredentialPlan, Error> {
-    let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
-    let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
-    let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
-    let client_secret =
-        configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup);
-    let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup)
-        .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
-    let authority = configured_string(
-        &inputs.azure_authority_host,
-        AZURE_AUTHORITY_HOST_ENV,
+    let token = configured_secret(
+        &inputs.azure_ad_token,
+        AzureSecretName::AdToken.into(),
         env_lookup,
     );
-    let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup)
-        .map(|value| {
-            value
-                .value()
-                .parse::<AzureCredentialType>()
-                .map(|selector| Sourced::new(selector, value.source()))
-        })
-        .transpose()
-        .map_err(|error| {
-            Error::InvalidConfiguration(ErrorDetail::failed(
-                "Azure credential selector parsing",
-                error,
-            ))
-        })?;
+    let tenant_id = configured_string(
+        &inputs.tenant_id,
+        AzureSecretName::TenantId.into(),
+        env_lookup,
+    );
+    let client_id = configured_string(
+        &inputs.client_id,
+        AzureSecretName::ClientId.into(),
+        env_lookup,
+    );
+    let client_secret = configured_secret(
+        &inputs.client_secret,
+        AzureSecretName::ClientSecret.into(),
+        env_lookup,
+    );
+    let scope = configured_string(
+        &inputs.azure_scope,
+        AzureSecretName::Scope.into(),
+        env_lookup,
+    )
+    .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
+    let authority = configured_string(
+        &inputs.azure_authority_host,
+        AzureSecretName::AuthorityHost.into(),
+        env_lookup,
+    );
+    let selector = configured_string(
+        &inputs.azure_credential,
+        AzureSecretName::Credential.into(),
+        env_lookup,
+    )
+    .map(|value| {
+        value
+            .value()
+            .parse::<AzureCredentialType>()
+            .map(|selector| Sourced::new(selector, value.source()))
+    })
+    .transpose()
+    .map_err(|error| {
+        Error::InvalidConfiguration(ErrorDetail::failed(
+            "Azure credential selector parsing",
+            error,
+        ))
+    })?;
     let federated_token_file = configured_string(
         &inputs.federated_token_file,
-        AZURE_FEDERATED_TOKEN_FILE_ENV,
+        AzureSecretName::FederatedTokenFile.into(),
         env_lookup,
     );
 
@@ -237,7 +400,12 @@ pub(crate) fn select_auth_plan(
         return Ok(AzureCredentialPlan::Missing);
     }
 
-    select_native_plan(
+    let probe_imds = selector
+        .as_ref()
+        .is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential)
+        && env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
+        && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
+    match select_native_plan(
         selector,
         tenant_id,
         client_id,
@@ -245,7 +413,15 @@ pub(crate) fn select_auth_plan(
         scope,
         authority,
         inputs.enable_azure_ad_token_refresh.source(),
-    )
+    )? {
+        AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain(
+            requests
+                .into_iter()
+                .map(|request| request.with_imds_probe(probe_imds))
+                .collect(),
+        )),
+        plan => Ok(plan),
+    }
 }
 
 fn select_native_plan(
@@ -509,7 +685,10 @@ mod tests {
     #[derive(Debug)]
     struct FileResolver;
 
+    use crate::native::CredentialAttempt;
+
     struct ChainAcquirer {
+        terminal_on: Option<&'static str>,
         requests: Mutex<Vec<&'static str>>,
         succeed_on: Option<&'static str>,
     }
@@ -518,20 +697,25 @@ mod tests {
         fn acquire(
             &self,
             request: ValidatedAzureRequest,
-        ) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
-        {
+        ) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
             let kind = request.kind();
             self.requests.lock().unwrap().push(kind);
             Box::pin(async move {
                 if self.succeed_on == Some(kind) {
-                    Ok(ResolvedCredential::AccessToken {
+                    CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
                         token: SecretValue::new("chain-token"),
                         expires_on: None,
                     })
                 } else {
-                    Err(Error::CredentialAcquisition(
-                        litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
-                    ))
+                    let error =
+                        Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved {
+                            subject: kind,
+                        });
+                    if self.terminal_on == Some(kind) {
+                        CredentialAttempt::Failed(error)
+                    } else {
+                        CredentialAttempt::Unavailable(error)
+                    }
                 }
             })
         }
@@ -690,9 +874,210 @@ mod tests {
         assert_eq!(resolved, Some(SecretValue::new("rotated-assertion")));
     }
 
+    #[derive(Debug)]
+    struct UnavailableImdsClient;
+
+    impl azure_core::http::HttpClient for UnavailableImdsClient {
+        fn execute_request<'life0, 'life1, 'async_trait>(
+            &'life0 self,
+            _request: &'life1 azure_core::http::Request,
+        ) -> std::pin::Pin<
+            Box<
+                dyn Future<Output = azure_core::Result<azure_core::http::AsyncRawResponse>>
+                    + Send
+                    + 'async_trait,
+            >,
+        >
+        where
+            'life0: 'async_trait,
+            'life1: 'async_trait,
+            Self: 'async_trait,
+        {
+            Box::pin(std::future::pending())
+        }
+    }
+
+    struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer);
+
+    impl AzureTokenAcquirer for OfflineChainAcquirer {
+        fn acquire(
+            &self,
+            request: ValidatedAzureRequest,
+        ) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
+            Box::pin(async move {
+                if request.kind() == "developer-tools" {
+                    return CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
+                        token: SecretValue::new("cli-token"),
+                        expires_on: None,
+                    });
+                }
+                self.0.acquire(request).await
+            })
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::default_off_azure("DefaultAzureCredential", None, true)]
+    #[case::explicit_managed_identity("ManagedIdentityCredential", None, false)]
+    #[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)]
+    #[tokio::test(start_paused = true)]
+    async fn only_default_imds_discovery_is_bounded(
+        #[case] selector: &str,
+        #[case] identity_endpoint: Option<&str>,
+        #[case] expect_cli: bool,
+    ) {
+        let native = super::NativeAzureTokenAcquirer::with_transport(
+            4,
+            azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
+        );
+        let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
+        let inputs = AzureAuthInputs {
+            azure_credential: crate::types::ConfigValue::Value(Sourced::new(
+                selector.into(),
+                InputSource::Deployment,
+            )),
+            enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
+            ..Default::default()
+        };
+        let environment = |name: &str| {
+            (name == "IDENTITY_ENDPOINT")
+                .then_some(identity_endpoint)
+                .flatten()
+                .map(str::to_string)
+        };
+        let outcome = tokio::time::timeout(
+            std::time::Duration::from_secs(2),
+            service.get_azure_ad_token(&inputs, &environment),
+        )
+        .await;
+
+        assert_eq!(outcome.is_ok(), expect_cli);
+        if expect_cli {
+            assert_eq!(
+                outcome.unwrap().unwrap().unwrap().value().secret().expose(),
+                "cli-token"
+            );
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::app_service(Some("IDENTITY_ENDPOINT"), false)]
+    #[case::legacy_app_service(Some("MSI_ENDPOINT"), false)]
+    #[case::off_azure(None, true)]
+    #[tokio::test(start_paused = true)]
+    async fn default_chain_uses_endpoints_from_the_secret_snapshot(
+        #[case] endpoint: Option<&str>,
+        #[case] expect_cli: bool,
+        #[values(false, true)] consumer_identity: bool,
+    ) {
+        use litellm_secrets::source::SecretSource;
+
+        struct EndpointSecrets<'a>(Option<&'a str>);
+
+        impl SecretSource for EndpointSecrets<'_> {
+            fn get_secret_str<'a>(
+                &'a self,
+                name: &'a str,
+            ) -> std::pin::Pin<
+                Box<
+                    dyn Future<Output = Result<Option<SecretValue>, litellm_secrets::Error>>
+                        + Send
+                        + 'a,
+                >,
+            > {
+                Box::pin(async move {
+                    Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test")))
+                })
+            }
+        }
+
+        let snapshot = EndpointSecrets(endpoint)
+            .resolve(SECRET_NAMES)
+            .await
+            .unwrap();
+        let native = super::NativeAzureTokenAcquirer::with_transport(
+            4,
+            azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
+        );
+        let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
+        let lookup = |name: &str| snapshot.get(name);
+        let scope = "https://service.test/.default";
+        let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async {
+            if consumer_identity {
+                service
+                    .identity_token(&super::AzureIdentity::Default, scope, &lookup)
+                    .await
+            } else {
+                service
+                    .get_azure_ad_token(
+                        &AzureAuthInputs::default_credential_for_scope(scope),
+                        &lookup,
+                    )
+                    .await
+                    .map(|token| token.map(Sourced::into_value))
+            }
+        })
+        .await;
+
+        assert_eq!(outcome.is_ok(), expect_cli);
+        if expect_cli {
+            assert_eq!(
+                outcome.unwrap().unwrap().unwrap().secret().expose(),
+                "cli-token"
+            );
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::litellm_provider_settings_are_ignored(
+        &[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")],
+        &["managed-identity", "developer-tools"]
+    )]
+    #[case::environment_client_secret_comes_first(
+        &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")],
+        &["client-secret", "managed-identity", "developer-tools"]
+    )]
+    #[case::aks_workload_identity(
+        &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")],
+        &["workload-identity", "managed-identity", "developer-tools"]
+    )]
+    #[tokio::test]
+    async fn the_default_identity_follows_azure_identity_defaults(
+        #[case] environment: &[(&str, &str)],
+        #[case] expected: &[&str],
+    ) {
+        let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: None,
+            requests: Mutex::new(Vec::new()),
+            succeed_on: Some("developer-tools"),
+        });
+        let service = AzureAuthService::with_acquirer(acquirer.clone());
+        let lookup = |name: &str| {
+            environment
+                .iter()
+                .find(|(key, _)| *key == name)
+                .map(|(_, value)| (*value).to_owned())
+        };
+
+        let credential = service
+            .identity_token(
+                &super::AzureIdentity::Default,
+                "https://redis.azure.com/.default",
+                &lookup,
+            )
+            .await
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(credential.secret().expose(), "chain-token");
+        assert_eq!(*acquirer.requests.lock().unwrap(), expected);
+    }
+
+    #[rstest::rstest]
     #[tokio::test]
     async fn default_chain_uses_declared_order_and_stops_after_success() {
         let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: None,
             requests: Mutex::new(Vec::new()),
             succeed_on: Some("developer-tools"),
         });
@@ -715,9 +1100,13 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
+    #[case::all_unavailable(false)]
+    #[case::terminal_failure_after_unavailable(true)]
     #[tokio::test]
-    async fn chain_reports_each_acquisition_failure() {
+    async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) {
         let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: terminal.then_some("developer-tools"),
             requests: Mutex::new(Vec::new()),
             succeed_on: None,
         });
@@ -744,6 +1133,27 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() {
+        let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: Some("managed-identity"),
+            requests: Mutex::new(Vec::new()),
+            succeed_on: Some("developer-tools"),
+        });
+        let service = AzureAuthService::with_acquirer(acquirer.clone());
+        let inputs = AzureAuthInputs {
+            enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
+            ..Default::default()
+        };
+        let error = service
+            .get_azure_ad_token(&inputs, &|_| None)
+            .await
+            .unwrap_err();
+        assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]);
+        assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty()));
+    }
+
     #[derive(Debug)]
     struct CallerToken(&'static str);
 

litellm-rust/crates/auth-gcp/Cargo.toml

@@ -6,7 +6,7 @@ license.workspace = true
 repository.workspace = true
 
 [features]
-google-sdk = ["dep:google-cloud-auth", "dep:http"]
+google-sdk = ["dep:http"]
 
 [dependencies]
 litellm-secret-value.workspace = true
@@ -21,10 +21,14 @@ tokio.workspace = true
 veil.workspace = true
 
 gcp_auth = "0.12.7"
-google-cloud-auth = { workspace = true, optional = true }
+google-cloud-auth.workspace = true
 http = { workspace = true, optional = true }
 
 [dev-dependencies]
+time = { workspace = true, features = ["formatting"] }
+axum = { workspace = true, features = ["json"] }
 litellm-error.workspace = true
 rstest.workspace = true
+rustls.workspace = true
 tempfile.workspace = true
+tokio = { workspace = true, features = ["test-util"] }

litellm-rust/crates/auth-gcp/src/auth.rs

@@ -1,20 +1,11 @@
-use std::{future::Future, pin::Pin, sync::Arc};
+use std::sync::Arc;
 
-use gcp_auth::{CustomServiceAccount, TokenProvider};
-use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
-use moka::future::Cache;
-use serde_json::Value;
-use sha2::{Digest, Sha256};
+use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential};
 
 use crate::{
-    config::{
-        CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
-        is_base64_encoded_json, names_a_credentials_file, non_empty_env,
-    },
-    constants::{
-        CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
-        VERTEXAI_API_KEY_ENV,
-    },
+    config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env},
+    constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV},
+    identity::{GcpAuth, GcpProviderLoader, GcpTokenSource},
 };
 
 pub struct VertexEnvironment {
@@ -22,28 +13,15 @@ pub struct VertexEnvironment {
     pub project_id: String,
 }
 
-struct VertexAccessToken {
-    token: String,
-    project_id: String,
-}
-
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct VertexAuth {
-    providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
-    loader: Arc<dyn VertexProviderLoader>,
-}
-
-impl Default for VertexAuth {
-    fn default() -> Self {
-        Self::new(Arc::new(GcpProviderLoader))
-    }
+    identity: GcpAuth,
 }
 
 impl VertexAuth {
-    pub fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
+    pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
         Self {
-            providers: Cache::builder().max_capacity(64).build(),
-            loader,
+            identity: GcpAuth::new(loader),
         }
     }
 
@@ -52,7 +30,10 @@ impl VertexAuth {
         config: &VertexConfig,
         env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
     ) -> Result<String, Error> {
-        self.load_provider(config, env_lookup).await?.token().await
+        self.identity
+            .access_token(credential_source(config, env_lookup))
+            .await
+            .map(|token| token.expose().to_owned())
     }
 
     pub async fn validate_environment(
@@ -77,10 +58,18 @@ impl VertexAuth {
             (true, _) => None,
             (false, Some(token)) => Some(token),
             (false, None) => {
-                let access = self.get_access_token(config, env_lookup).await?;
+                let provider = self.load_provider(config, env_lookup).await?;
+                let (token, project_id) = match project_id {
+                    Some(project_id) => (provider.token().await?, project_id),
+                    None => tokio::try_join!(provider.token(), provider.project_id())?,
+                };
                 return Ok(VertexEnvironment {
-                    headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
-                    project_id: project_id.unwrap_or(access.project_id),
+                    headers: apply_credential(
+                        headers,
+                        token.expose(),
+                        CredentialPlacement::Bearer,
+                    )?,
+                    project_id,
                 });
             }
         };
@@ -104,248 +93,171 @@ impl VertexAuth {
         })
     }
 
-    async fn get_access_token(
-        &self,
-        config: &VertexConfig,
-        env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
-    ) -> Result<VertexAccessToken, Error> {
-        let provider = self.load_provider(config, env_lookup).await?;
-        let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
-        Ok(VertexAccessToken { token, project_id })
-    }
-
     async fn load_provider(
         &self,
         config: &VertexConfig,
         env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
-    ) -> Result<Arc<dyn VertexTokenSource>, Error> {
-        let source = credential_source(config, env_lookup);
-        let key = source.cache_key();
-        self.providers
-            .try_get_with(key, self.loader.load(source))
+    ) -> Result<Arc<dyn GcpTokenSource>, Error> {
+        self.identity
+            .load_provider(credential_source(config, env_lookup))
             .await
-            .map_err(|error| (*error).clone())
     }
 }
 
-pub trait VertexTokenSource: Send + Sync {
-    fn project_id(&self) -> VertexAuthFuture<'_, String>;
-    fn token(&self) -> VertexAuthFuture<'_, String>;
-}
-
-pub trait VertexProviderLoader: Send + Sync {
-    fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
-}
-
-pub type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-
-struct GcpTokenSource(Arc<dyn TokenProvider>);
-
-impl VertexTokenSource for GcpTokenSource {
-    fn project_id(&self) -> VertexAuthFuture<'_, String> {
-        Box::pin(async move {
-            self.0
-                .project_id()
-                .await
-                .map(|project| project.to_string())
-                .map_err(auth_acquisition_error)
-        })
-    }
-
-    fn token(&self) -> VertexAuthFuture<'_, String> {
-        Box::pin(async move {
-            self.0
-                .token(&[CLOUD_PLATFORM_SCOPE])
-                .await
-                .map(|token| token.as_str().to_string())
-                .map_err(auth_acquisition_error)
-        })
-    }
-}
-
-struct GcpProviderLoader;
-
-impl VertexProviderLoader for GcpProviderLoader {
-    fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
-        Box::pin(async move {
-            let provider: Arc<dyn TokenProvider> = match source {
-                CredentialSource::Inline(configured) => Arc::new(
-                    CustomServiceAccount::from_json(validate_request_credentials(
-                        configured.expose(),
-                    )?)
-                    .map_err(auth_acquisition_error)?,
-                ),
-                CredentialSource::Trusted(configured) => {
-                    Arc::new(load_trusted_service_account(configured.expose())?)
-                }
-                CredentialSource::ApplicationCredentials(path) => {
-                    Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
-                }
-                CredentialSource::Adc => {
-                    gcp_auth::provider().await.map_err(auth_acquisition_error)?
-                }
-            };
-            Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
-        })
-    }
-}
-
-fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
-    let token_uri = serde_json::from_str::<Value>(configured)
-        .ok()
-        .and_then(|credentials| {
-            credentials
-                .get("token_uri")
-                .and_then(Value::as_str)
-                .map(str::to_string)
-        });
-    if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
-        return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
-            field: "request-controlled vertex_credentials token_uri".into(),
-            expected: "the canonical Google OAuth token endpoint",
-        }));
-    }
-    Ok(configured)
-}
-
-impl CredentialSource {
-    fn cache_key(&self) -> CredentialCacheKey {
-        match self {
-            Self::Inline(configured) => {
-                CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
-            }
-            Self::Trusted(configured) => {
-                CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
-            }
-            Self::ApplicationCredentials(path) => {
-                CredentialCacheKey::ApplicationCredentials(path.clone())
-            }
-            Self::Adc => CredentialCacheKey::Adc,
-        }
-    }
-}
-
-#[derive(Clone, Debug, Hash, PartialEq, Eq)]
-enum CredentialCacheKey {
-    Inline([u8; 32]),
-    Trusted([u8; 32]),
-    ApplicationCredentials(String),
-    Adc,
-}
-
-fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
-    if is_base64_encoded_json(configured) {
-        return Err(Error::InvalidConfiguration(
-            ErrorDetail::Base64EncodedJson {
-                field: "vertex_credentials",
-            },
-        ));
-    }
-    if names_a_credentials_file(configured) {
-        CustomServiceAccount::from_file(configured)
-    } else {
-        CustomServiceAccount::from_json(configured)
-    }
-    .map_err(auth_acquisition_error)
-}
-
-fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
-    Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
-        "Vertex AI credentials",
-        error,
-    ))
-}
-
 #[cfg(test)]
 mod tests {
+    use std::{
+        collections::{BTreeMap, BTreeSet},
+        sync::{
+            Arc,
+            atomic::{AtomicUsize, Ordering},
+        },
+    };
+
+    use crate::{
+        CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth,
+        VertexConfig, get_vertex_ai_location, secret_names,
+    };
     use litellm_secret_value::SecretValue;
+    use rstest::rstest;
+    use serde_json::{Value, json};
 
-    use super::*;
-
-    #[rstest::rstest]
-    #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
-    #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
-    #[case::missing_endpoint("{}", false)]
-    fn request_credentials_require_canonical_token_endpoint(
-        #[case] credentials: &str,
-        #[case] accepted: bool,
-    ) {
-        assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
+    struct FakeProvider {
+        project: Option<&'static str>,
+        calls: Arc<AtomicUsize>,
     }
 
-    #[rstest::rstest]
-    #[case::missing_file(
-        "/nonexistent/vertex-key.json",
-        "failed to open application credentials file"
-    )]
-    #[case::malformed_json(
-        " {\"type\": \"service_account\",",
-        "failed to deserialize ApplicationCredentials"
-    )]
-    #[case::quoted_json(
-        "'{\"type\": \"service_account\"}'",
-        "failed to deserialize ApplicationCredentials"
-    )]
-    #[tokio::test]
-    async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
-        #[case] configured: &str,
-        #[case] expected: &str,
-    ) {
-        let Err(error) = GcpProviderLoader
-            .load(CredentialSource::Trusted(SecretValue::new(configured)))
-            .await
-        else {
-            panic!("{configured:?} should not load");
-        };
-        let report = litellm_error::Report::new(&error).to_string();
-        assert!(report.contains(expected), "{report}");
-    }
+    impl GcpTokenSource for FakeProvider {
+        fn project_id(&self) -> GcpAuthFuture<'_, String> {
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            Box::pin(async move {
+                self.project.map(str::to_owned).ok_or(
+                    litellm_auth_types::Error::CredentialAcquisition(
+                        litellm_auth_types::ErrorDetail::Unresolved {
+                            subject: "credential project",
+                        },
+                    ),
+                )
+            })
+        }
 
-    #[rstest::rstest]
-    #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
-    #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
-    #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
-    #[tokio::test]
-    async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
-        assert!(
-            trusted_load_error(configured)
-                .await
-                .contains("is base64-encoded JSON, not the JSON itself")
-        );
-    }
-
-    async fn trusted_load_error(configured: &str) -> String {
-        match GcpProviderLoader
-            .load(CredentialSource::Trusted(SecretValue::new(configured)))
-            .await
-        {
-            Ok(_) => panic!("{configured:?} should not load"),
-            Err(error) => litellm_error::Report::new(&error).to_string(),
+        fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            Box::pin(async { Ok(SecretValue::new("adc-token")) })
         }
     }
 
-    #[rstest::rstest]
-    #[tokio::test]
-    async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
-        let key = tempfile::Builder::new()
-            .prefix("{vertex")
-            .suffix(".json")
-            .tempfile_in(".")
-            .unwrap();
-        std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
-        let relative_name = key.path().file_name().unwrap().to_str().unwrap();
-        let absolute_path = std::fs::canonicalize(key.path()).unwrap();
-        assert_eq!(
-            trusted_load_error(relative_name).await,
-            trusted_load_error(absolute_path.to_str().unwrap()).await
-        );
+    struct FakeLoader {
+        loads: Arc<AtomicUsize>,
+        provider: Arc<dyn GcpTokenSource>,
     }
 
-    #[rstest::rstest]
-    fn inline_and_trusted_credentials_never_share_a_cache_entry() {
-        assert_ne!(
-            CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
-            CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
-        );
+    impl GcpProviderLoader for FakeLoader {
+        fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
+            let loads = self.loads.clone();
+            let provider = self.provider.clone();
+            Box::pin(async move {
+                loads.fetch_add(1, Ordering::SeqCst);
+                Ok(provider)
+            })
+        }
+    }
+
+    fn config(value: Value) -> VertexConfig {
+        VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
+            .unwrap()
+    }
+
+    fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
+        let provider: Arc<dyn GcpTokenSource> = Arc::new(FakeProvider {
+            calls,
+            project: Some("adc-project"),
+        });
+        VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn secret_names_cover_environment_reads() {
+        let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
+        let recorded = seen.clone();
+        let env = |name: &str| {
+            recorded.lock().unwrap().insert(name.to_string());
+            None
+        };
+        let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
+        auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
+            .await
+            .unwrap();
+        get_vertex_ai_location(&VertexConfig::default(), &env);
+        let allowed = secret_names()
+            .iter()
+            .map(|name| (*name).to_owned())
+            .collect::<BTreeSet<_>>();
+        assert!(seen.lock().unwrap().is_subset(&allowed));
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn explicit_token_and_header_do_not_acquire_adc() {
+        let loads = Arc::new(AtomicUsize::new(0));
+        let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
+        let configured = config(json!({"vertex_project":"project-1"}));
+        let explicit = auth
+            .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
+            .await
+            .unwrap();
+        assert_eq!(explicit.headers[0].1, "Bearer access-token");
+        let existing = auth
+            .validate_environment(
+                vec![("authorization".into(), "Bearer existing".into())],
+                None,
+                &configured,
+                &|_| None,
+            )
+            .await
+            .unwrap();
+        assert_eq!(existing.headers[0].1, "Bearer existing");
+        assert_eq!(loads.load(Ordering::SeqCst), 0);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn provider_is_reused_across_authentication_calls() {
+        let calls = Arc::new(AtomicUsize::new(0));
+        let loads = Arc::new(AtomicUsize::new(0));
+        let auth = auth(calls.clone(), loads.clone());
+        let config = VertexConfig::default();
+        let environment = |_: &str| None;
+        let (first, second) = tokio::try_join!(
+            auth.validate_environment(Vec::new(), None, &config, &environment),
+            auth.validate_environment(Vec::new(), None, &config, &environment),
+        )
+        .unwrap();
+        assert_eq!(first.project_id, "adc-project");
+        assert_eq!(first.headers[0].1, "Bearer adc-token");
+        assert_eq!(first.project_id, second.project_id);
+        assert_eq!(first.headers, second.headers);
+        assert_eq!(loads.load(Ordering::SeqCst), 1);
+        assert_eq!(calls.load(Ordering::SeqCst), 4);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() {
+        let auth = VertexAuth::new(Arc::new(FakeLoader {
+            loads: Arc::new(AtomicUsize::new(0)),
+            provider: Arc::new(FakeProvider {
+                calls: Arc::new(AtomicUsize::new(0)),
+                project: None,
+            }),
+        }));
+        let configured = VertexConfig::new(None, Some("configured-project".into()), None);
+        let environment = auth
+            .validate_environment(Vec::new(), None, &configured, &|_| None)
+            .await
+            .unwrap();
+        assert_eq!(environment.project_id, "configured-project");
+        assert_eq!(environment.headers[0].1, "Bearer adc-token");
     }
 }

litellm-rust/crates/auth-gcp/src/config.rs

@@ -146,6 +146,20 @@ pub enum CredentialSource {
     Adc,
 }
 
+impl CredentialSource {
+    pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
+        VertexParams::CREDENTIALS
+            .resolve(&|_| None, env_lookup)
+            .map(|configured| Self::Trusted(SecretValue::new(configured)))
+            .unwrap_or_else(|| Self::application_default(env_lookup))
+    }
+
+    pub fn application_default(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
+        non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
+            .map_or(Self::Adc, Self::ApplicationCredentials)
+    }
+}
+
 pub(crate) fn credential_source(
     config: &VertexConfig,
     env_lookup: &dyn Fn(&str) -> Option<String>,
@@ -158,12 +172,7 @@ pub(crate) fn credential_source(
             }
         };
     }
-    if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) {
-        return CredentialSource::Trusted(SecretValue::new(configured));
-    }
-    non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
-        .map(CredentialSource::ApplicationCredentials)
-        .unwrap_or(CredentialSource::Adc)
+    CredentialSource::from_vertex_environment(env_lookup)
 }
 
 fn optional_credentials(
@@ -319,6 +328,24 @@ mod tests {
         ));
     }
 
+    #[rstest::rstest]
+    fn application_default_ignores_vertex_credentials() {
+        let env = |name: &str| match name {
+            VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()),
+            GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()),
+            _ => None,
+        };
+        assert!(
+            matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
+        );
+        assert!(matches!(
+            CredentialSource::application_default(
+                &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())
+            ),
+            CredentialSource::Adc
+        ));
+    }
+
     #[rstest::rstest]
     fn params_become_a_deployment_sourced_config() {
         let params = VertexParams {

litellm-rust/crates/auth-gcp/src/identity.rs

@@ -0,0 +1,445 @@
+use crate::{
+    config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file},
+    constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT},
+};
+use gcp_auth::{CustomServiceAccount, TokenProvider};
+use litellm_auth_types::{Error, ErrorDetail};
+use litellm_secret_value::SecretValue;
+use moka::future::Cache;
+use serde_json::Value;
+use sha2::{Digest, Sha256};
+use std::{future::Future, pin::Pin, sync::Arc};
+
+#[derive(Clone)]
+pub struct GcpAuth {
+    providers: Cache<CredentialCacheKey, Arc<dyn GcpTokenSource>>,
+    loader: Arc<dyn GcpProviderLoader>,
+}
+
+impl Default for GcpAuth {
+    fn default() -> Self {
+        Self::new(Arc::new(NativeProviderLoader))
+    }
+}
+
+impl GcpAuth {
+    pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
+        Self {
+            providers: Cache::builder().max_capacity(64).build(),
+            loader,
+        }
+    }
+
+    pub async fn access_token(&self, source: CredentialSource) -> Result<SecretValue, Error> {
+        self.load_provider(source).await?.token().await
+    }
+
+    pub(crate) async fn load_provider(
+        &self,
+        source: CredentialSource,
+    ) -> Result<Arc<dyn GcpTokenSource>, Error> {
+        self.providers
+            .try_get_with(source.cache_key(), self.loader.load(source))
+            .await
+            .map_err(|error| (*error).clone())
+    }
+}
+
+pub trait GcpTokenSource: Send + Sync {
+    fn project_id(&self) -> GcpAuthFuture<'_, String>;
+    fn token(&self) -> GcpAuthFuture<'_, SecretValue>;
+}
+
+pub trait GcpProviderLoader: Send + Sync {
+    fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>>;
+}
+
+pub type GcpAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
+
+struct ServiceAccountTokenSource(Arc<dyn TokenProvider>);
+
+impl GcpTokenSource for ServiceAccountTokenSource {
+    fn project_id(&self) -> GcpAuthFuture<'_, String> {
+        Box::pin(async move {
+            self.0
+                .project_id()
+                .await
+                .map(|project| project.to_string())
+                .map_err(auth_acquisition_error)
+        })
+    }
+
+    fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
+        Box::pin(async move {
+            self.0
+                .token(&[CLOUD_PLATFORM_SCOPE])
+                .await
+                .map(|token| SecretValue::new(token.as_str()))
+                .map_err(auth_acquisition_error)
+        })
+    }
+}
+
+struct SdkTokenSource {
+    credentials: google_cloud_auth::credentials::AccessTokenCredentials,
+    project_id: Option<String>,
+}
+
+impl GcpTokenSource for SdkTokenSource {
+    fn project_id(&self) -> GcpAuthFuture<'_, String> {
+        Box::pin(async move {
+            self.project_id
+                .clone()
+                .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
+                    subject: "GCP resource project",
+                }))
+        })
+    }
+
+    fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
+        Box::pin(async move {
+            self.credentials
+                .access_token()
+                .await
+                .map(|token| SecretValue::new(token.token))
+                .map_err(auth_acquisition_error)
+        })
+    }
+}
+
+async fn load_application_credentials(path: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
+    let contents = tokio::fs::read_to_string(path)
+        .await
+        .map_err(auth_acquisition_error)?;
+    load_credentials_json(&contents)
+}
+
+fn load_credentials_json(contents: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
+    let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?;
+    let kind = configured.get("type").and_then(Value::as_str);
+    let project_id = configured
+        .get("project_id")
+        .and_then(Value::as_str)
+        .map(str::to_string);
+    let credentials = match kind {
+        Some("authorized_user") => {
+            google_cloud_auth::credentials::user_account::Builder::new(configured)
+                .with_scopes([CLOUD_PLATFORM_SCOPE])
+                .build_access_token_credentials()
+        }
+        Some("impersonated_service_account") => {
+            google_cloud_auth::credentials::impersonated::Builder::new(configured)
+                .with_scopes([CLOUD_PLATFORM_SCOPE])
+                .build_access_token_credentials()
+        }
+        Some("external_account") => {
+            google_cloud_auth::credentials::external_account::Builder::new(configured)
+                .with_scopes([CLOUD_PLATFORM_SCOPE])
+                .build_access_token_credentials()
+        }
+        Some("service_account") | None => {
+            let provider =
+                CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?;
+            return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider))));
+        }
+        Some(_) => {
+            return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
+                field: "Google application credentials type".into(),
+                expected: "a supported Google credential type",
+            }));
+        }
+    }
+    .map_err(auth_acquisition_error)?;
+    Ok(Arc::new(SdkTokenSource {
+        credentials,
+        project_id,
+    }))
+}
+
+struct NativeProviderLoader;
+
+impl GcpProviderLoader for NativeProviderLoader {
+    fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
+        Box::pin(async move {
+            let provider: Arc<dyn TokenProvider> = match source {
+                CredentialSource::Inline(configured) => Arc::new(
+                    CustomServiceAccount::from_json(validate_request_credentials(
+                        configured.expose(),
+                    )?)
+                    .map_err(auth_acquisition_error)?,
+                ),
+                CredentialSource::Trusted(configured) => {
+                    return load_trusted_credentials(configured.expose()).await;
+                }
+                CredentialSource::ApplicationCredentials(path) => {
+                    return load_application_credentials(&path).await;
+                }
+                CredentialSource::Adc => {
+                    gcp_auth::provider().await.map_err(auth_acquisition_error)?
+                }
+            };
+            Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc<dyn GcpTokenSource>)
+        })
+    }
+}
+
+fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
+    let token_uri = serde_json::from_str::<Value>(configured)
+        .ok()
+        .and_then(|credentials| {
+            credentials
+                .get("token_uri")
+                .and_then(Value::as_str)
+                .map(str::to_string)
+        });
+    if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
+        return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
+            field: "request-controlled vertex_credentials token_uri".into(),
+            expected: "the canonical Google OAuth token endpoint",
+        }));
+    }
+    Ok(configured)
+}
+
+impl CredentialSource {
+    fn cache_key(&self) -> CredentialCacheKey {
+        match self {
+            Self::Inline(configured) => {
+                CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
+            }
+            Self::Trusted(configured) => {
+                CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
+            }
+            Self::ApplicationCredentials(path) => {
+                CredentialCacheKey::ApplicationCredentials(path.clone())
+            }
+            Self::Adc => CredentialCacheKey::Adc,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+enum CredentialCacheKey {
+    Inline([u8; 32]),
+    Trusted([u8; 32]),
+    ApplicationCredentials(String),
+    Adc,
+}
+
+async fn load_trusted_credentials(configured: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
+    if is_base64_encoded_json(configured) {
+        return Err(Error::InvalidConfiguration(
+            ErrorDetail::Base64EncodedJson {
+                field: "vertex_credentials",
+            },
+        ));
+    }
+    if names_a_credentials_file(configured) {
+        load_application_credentials(configured).await
+    } else {
+        load_credentials_json(configured)
+    }
+}
+
+fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error {
+    Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
+        "Google credentials",
+        error,
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use litellm_secret_value::SecretValue;
+
+    use super::*;
+
+    #[rstest::rstest]
+    #[case::authorized_user("authorized_user", 2)]
+    #[case::external_account("external_account", 2)]
+    #[case::impersonated_service_account("impersonated_service_account", 4)]
+    #[tokio::test]
+    async fn explicit_application_credentials_refresh_expired_tokens(
+        #[case] kind: &str,
+        #[case] expected_requests: usize,
+        #[values(false, true)] inline: bool,
+    ) {
+        use axum::{Json, Router, routing::post};
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let _ = rustls::crypto::ring::default_provider().install_default();
+        let requests = Arc::new(AtomicUsize::new(0));
+        let recorded = requests.clone();
+        let impersonations = requests.clone();
+        let app = Router::new().route("/token", post(move || {
+            let recorded = recorded.clone();
+            async move {
+                Json(serde_json::json!({
+                    "access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)),
+                    "expires_in": 3600,
+                    "token_type": "Bearer",
+                    "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
+                }))
+            }
+        }));
+        let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || {
+            let recorded = impersonations.clone();
+            async move {
+                let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap();
+                Json(serde_json::json!({
+                    "accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)),
+                    "expireTime": expiration,
+                }))
+            }
+        }));
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+        let file = tempfile::NamedTempFile::new().unwrap();
+        let subject_token = tempfile::NamedTempFile::new().unwrap();
+        std::fs::write(subject_token.path(), "subject-token").unwrap();
+        let configured = if kind == "external_account" {
+            serde_json::json!({
+                "type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test",
+                "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"),
+                "credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project"
+            })
+        } else {
+            let source = serde_json::json!({
+                "type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret",
+                "refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token")
+            });
+            if kind == "impersonated_service_account" {
+                serde_json::json!({
+                    "type": kind, "source_credentials": source, "quota_project_id": "quota-project",
+                    "service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken")
+                })
+            } else {
+                source
+            }
+        };
+        std::fs::write(file.path(), configured.to_string()).unwrap();
+        let source = if inline {
+            CredentialSource::Trusted(SecretValue::new(configured.to_string()))
+        } else {
+            CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into())
+        };
+        let provider = NativeProviderLoader.load(source).await.unwrap();
+
+        assert!(matches!(
+            provider.project_id().await,
+            Err(Error::CredentialAcquisition(ErrorDetail::Unresolved {
+                subject: "GCP resource project"
+            }))
+        ));
+        let original = provider.token().await.unwrap();
+        tokio::time::pause();
+        tokio::time::advance(std::time::Duration::from_secs(3601)).await;
+        tokio::time::resume();
+        let refreshed = provider.token().await.unwrap();
+        assert_ne!(original, refreshed);
+        assert_eq!(requests.load(Ordering::SeqCst), expected_requests);
+        server.abort();
+    }
+
+    #[rstest::rstest]
+    #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
+    #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
+    #[case::missing_endpoint("{}", false)]
+    fn request_credentials_require_canonical_token_endpoint(
+        #[case] credentials: &str,
+        #[case] accepted: bool,
+    ) {
+        assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
+    }
+
+    #[rstest::rstest]
+    #[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")]
+    #[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")]
+    #[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")]
+    #[tokio::test]
+    async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
+        #[case] configured: &str,
+        #[case] expected: &str,
+    ) {
+        let Err(error) = NativeProviderLoader
+            .load(CredentialSource::Trusted(SecretValue::new(configured)))
+            .await
+        else {
+            panic!("{configured:?} should not load");
+        };
+        let report = litellm_error::Report::new(&error).to_string();
+        assert!(report.contains(expected), "{report}");
+    }
+
+    #[rstest::rstest]
+    #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
+    #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
+    #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
+    #[tokio::test]
+    async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
+        assert!(
+            trusted_load_error(configured)
+                .await
+                .contains("is base64-encoded JSON, not the JSON itself")
+        );
+    }
+
+    async fn trusted_load_error(configured: &str) -> String {
+        match NativeProviderLoader
+            .load(CredentialSource::Trusted(SecretValue::new(configured)))
+            .await
+        {
+            Ok(_) => panic!("{configured:?} should not load"),
+            Err(error) => litellm_error::Report::new(&error).to_string(),
+        }
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
+        let key = tempfile::Builder::new()
+            .prefix("{vertex")
+            .suffix(".json")
+            .tempfile_in(".")
+            .unwrap();
+        std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
+        let relative_name = key.path().file_name().unwrap().to_str().unwrap();
+        let absolute_path = std::fs::canonicalize(key.path()).unwrap();
+        assert_eq!(
+            trusted_load_error(relative_name).await,
+            trusted_load_error(absolute_path.to_str().unwrap()).await
+        );
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() {
+        let file = tempfile::NamedTempFile::new().unwrap();
+        std::fs::write(
+            file.path(),
+            r#"{"type":"unsupported","private_key":"credential-value"}"#,
+        )
+        .unwrap();
+        let result = GcpAuth::default()
+            .access_token(CredentialSource::ApplicationCredentials(
+                file.path().to_str().unwrap().into(),
+            ))
+            .await;
+        let error = result.unwrap_err();
+        assert!(matches!(error, Error::InvalidConfiguration(_)));
+        assert!(
+            !litellm_error::Report::new(&error)
+                .to_string()
+                .contains("credential-value")
+        );
+    }
+
+    #[rstest::rstest]
+    fn inline_and_trusted_credentials_never_share_a_cache_entry() {
+        assert_ne!(
+            CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
+            CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
+        );
+    }
+}

litellm-rust/crates/auth-gcp/src/lib.rs

@@ -1,16 +1,16 @@
 mod auth;
 mod config;
 pub mod constants;
+mod identity;
 #[cfg(feature = "google-sdk")]
 mod sdk;
 
-pub use auth::{
-    VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource,
-};
+pub use auth::{VertexAuth, VertexEnvironment};
 pub use config::{
     CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project,
     get_vertex_ai_project_from_credentials,
 };
 pub use constants::secret_names;
+pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource};
 #[cfg(feature = "google-sdk")]
 pub use sdk::GoogleCredentials;

litellm-rust/crates/auth-gcp/tests/vertex_auth.rs

@@ -1,120 +0,0 @@
-use std::{
-    collections::{BTreeMap, BTreeSet},
-    sync::{
-        Arc,
-        atomic::{AtomicUsize, Ordering},
-    },
-};
-
-use litellm_auth_gcp::{
-    CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader,
-    VertexTokenSource, get_vertex_ai_location, secret_names,
-};
-use rstest::rstest;
-use serde_json::{Value, json};
-
-struct FakeProvider {
-    calls: Arc<AtomicUsize>,
-}
-
-impl VertexTokenSource for FakeProvider {
-    fn project_id(&self) -> VertexAuthFuture<'_, String> {
-        self.calls.fetch_add(1, Ordering::SeqCst);
-        Box::pin(async { Ok("adc-project".into()) })
-    }
-
-    fn token(&self) -> VertexAuthFuture<'_, String> {
-        self.calls.fetch_add(1, Ordering::SeqCst);
-        Box::pin(async { Ok("adc-token".into()) })
-    }
-}
-
-struct FakeLoader {
-    loads: Arc<AtomicUsize>,
-    provider: Arc<dyn VertexTokenSource>,
-}
-
-impl VertexProviderLoader for FakeLoader {
-    fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
-        let loads = self.loads.clone();
-        let provider = self.provider.clone();
-        Box::pin(async move {
-            loads.fetch_add(1, Ordering::SeqCst);
-            Ok(provider)
-        })
-    }
-}
-
-fn config(value: Value) -> VertexConfig {
-    VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
-        .unwrap()
-}
-
-fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
-    let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
-    VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
-}
-
-#[rstest]
-#[tokio::test]
-async fn secret_names_cover_environment_reads() {
-    let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
-    let recorded = seen.clone();
-    let env = |name: &str| {
-        recorded.lock().unwrap().insert(name.to_string());
-        None
-    };
-    let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
-    auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
-        .await
-        .unwrap();
-    get_vertex_ai_location(&VertexConfig::default(), &env);
-    assert!(
-        seen.lock()
-            .unwrap()
-            .iter()
-            .all(|name| secret_names().contains(&name.as_str()))
-    );
-}
-
-#[rstest]
-#[tokio::test]
-async fn explicit_token_and_header_do_not_acquire_adc() {
-    let loads = Arc::new(AtomicUsize::new(0));
-    let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
-    let configured = config(json!({"vertex_project":"project-1"}));
-    let explicit = auth
-        .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
-        .await
-        .unwrap();
-    assert_eq!(explicit.headers[0].1, "Bearer access-token");
-    let existing = auth
-        .validate_environment(
-            vec![("authorization".into(), "Bearer existing".into())],
-            None,
-            &configured,
-            &|_| None,
-        )
-        .await
-        .unwrap();
-    assert_eq!(existing.headers[0].1, "Bearer existing");
-    assert_eq!(loads.load(Ordering::SeqCst), 0);
-}
-
-#[rstest]
-#[tokio::test]
-async fn provider_is_reused_across_authentication_calls() {
-    let calls = Arc::new(AtomicUsize::new(0));
-    let loads = Arc::new(AtomicUsize::new(0));
-    let auth = auth(calls.clone(), loads.clone());
-    for _ in 0..2 {
-        let environment = auth
-            .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
-            .await
-            .unwrap();
-        assert_eq!(environment.project_id, "adc-project");
-        assert_eq!(environment.headers[0].1, "Bearer adc-token");
-    }
-    assert_eq!(loads.load(Ordering::SeqCst), 1);
-    assert_eq!(calls.load(Ordering::SeqCst), 4);
-}

litellm-rust/crates/auth-types/src/error.rs

@@ -92,6 +92,11 @@ impl Error {
 
 #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
 pub enum ErrorDetail {
+    #[error("{operation} returned HTTP {status}")]
+    UnexpectedStatus {
+        operation: &'static str,
+        status: u16,
+    },
     #[error("{field} must be {expected}")]
     InvalidType {
         field: String,

litellm-rust/crates/auth-types/src/params/aws.rs

@@ -53,6 +53,9 @@ pub struct AwsParams {
 }
 
 impl AwsParams {
+    pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN";
+    pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
+
     pub const ACCESS_KEY_ID: ParamSpec = spec(
         "access key id",
         ParamKind::Credential,
@@ -149,6 +152,10 @@ impl AwsParams {
             AwsParams::SPECS
                 .iter()
                 .flat_map(|spec| spec.env.iter().copied())
+                .chain([
+                    AwsParams::ROLE_ARN_ENV,
+                    AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV,
+                ])
                 .fold(Vec::new(), |mut names, name| {
                     if !names.contains(&name) {
                         names.push(name);

litellm-rust/crates/auth-types/tests/aws_params.rs

@@ -3,21 +3,13 @@ use rstest::rstest;
 use serde_json::{Map, Value};
 
 #[rstest]
-fn fields_name_every_param_once_and_in_declaration_order() {
+fn every_named_param_roundtrips_without_loss() {
     let filled: Map<String, Value> = AwsParams::fields()
         .map(|name| (name.to_string(), Value::from(format!("value-of-{name}"))))
         .collect();
     let typed = AwsParams::from_optional_params(&filled);
     let serialized = serde_json::to_value(&typed).unwrap();
-    assert_eq!(
-        serialized
-            .as_object()
-            .unwrap()
-            .keys()
-            .map(String::as_str)
-            .collect::<Vec<_>>(),
-        AwsParams::fields().collect::<Vec<_>>()
-    );
+    assert_eq!(filled.len(), AwsParams::fields().count());
     assert_eq!(serialized, Value::Object(filled.clone()));
     assert_eq!(
         serde_json::from_value::<AwsParams>(Value::Object(filled)).unwrap(),

litellm-rust/crates/cache-gcs/src/cache.rs

@@ -8,7 +8,7 @@ use litellm_cache::{
 use litellm_http::Client;
 use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
 
-use crate::{GcpTokenSource, TokenSource};
+use crate::{GcsTokenSource, TokenSource};
 
 pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
 
@@ -62,7 +62,7 @@ pub struct GcsCache<S: CacheCodec> {
 
 impl<S: CacheCodec> GcsCache<S> {
     pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
-        let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
+        let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone()));
         Self::with_token_source(config, client, codec, token)
     }
 

litellm-rust/crates/cache-gcs/src/lib.rs

@@ -2,4 +2,4 @@ mod cache;
 mod token;
 
 pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
-pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};
+pub use token::{GcsTokenSource, StaticTokenSource, TokenSource};

litellm-rust/crates/cache-gcs/src/token.rs

@@ -1,7 +1,6 @@
 use std::{future::Future, pin::Pin};
 
-use litellm_auth_gcp::{VertexAuth, VertexConfig};
-use litellm_auth_types::{InputSource, Sourced};
+use litellm_auth_gcp::{CredentialSource, GcpAuth};
 use litellm_secret_value::SecretValue;
 
 use litellm_cache::Error;
@@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static {
     fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
 }
 
-pub struct GcpTokenSource {
-    auth: VertexAuth,
-    config: VertexConfig,
+pub struct GcsTokenSource {
+    auth: GcpAuth,
+    credentials: Option<String>,
 }
 
-impl GcpTokenSource {
+impl GcsTokenSource {
     pub fn new(path_service_account: Option<String>) -> Self {
-        let credentials = path_service_account
-            .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
         Self {
-            auth: VertexAuth::default(),
-            config: VertexConfig::new(credentials, None, None),
+            auth: GcpAuth::default(),
+            credentials: path_service_account.filter(|path| !path.trim().is_empty()),
         }
     }
 }
 
-impl TokenSource for GcpTokenSource {
+impl TokenSource for GcsTokenSource {
     fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
         Box::pin(async move {
             self.auth
-                .access_token(&self.config, &|name| std::env::var(name).ok())
+                .access_token(match &self.credentials {
+                    Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())),
+                    None => CredentialSource::application_default(&|name| std::env::var(name).ok()),
+                })
                 .await
+                .map(|token| token.expose().to_owned())
                 .map_err(Error::unavailable)
         })
     }

litellm-rust/crates/cache-s3/src/auth.rs

@@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials {
                 self.config.secret_access_key.clone(),
             ) {
                 return Ok(AwsCredentials::new(
-                    access_key_id,
-                    secret_access_key,
-                    self.config.session_token.clone(),
+                    access_key_id.expose(),
+                    secret_access_key.expose(),
+                    self.config
+                        .session_token
+                        .as_ref()
+                        .map(|value| value.expose().to_owned()),
                     None,
                     "litellm-s3-cache",
                 ));

litellm-rust/crates/cache-s3/tests/auth.rs

@@ -5,9 +5,9 @@ use rstest::rstest;
 
 fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
     AwsAuthConfig {
-        access_key_id: Some("key".to_string()),
-        secret_access_key: Some("secret".to_string()),
-        session_token: session_token.map(str::to_string),
+        access_key_id: Some("key".into()),
+        secret_access_key: Some("secret".into()),
+        session_token: session_token.map(Into::into),
         region_name: Some("us-east-1".to_string()),
         ..Default::default()
     }

litellm-rust/crates/cache-s3/tests/support/mod.rs

@@ -23,8 +23,8 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
             url: endpoint.to_string(),
         }),
         auth: AwsAuthConfig {
-            access_key_id: Some("key".to_string()),
-            secret_access_key: Some("secret".to_string()),
+            access_key_id: Some("key".into()),
+            secret_access_key: Some("secret".into()),
             region_name: Some("us-east-1".to_string()),
             ..Default::default()
         },

litellm-rust/crates/inference-ocr/tests/resources.rs

@@ -7,7 +7,7 @@ use std::sync::{
 
 use litellm_auth::AuthServices;
 use litellm_auth_gcp::{
-    CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource,
+    CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth,
 };
 use litellm_http::{HttpSettings, Resolution};
 use litellm_inference::resources::CoreResources;
@@ -20,28 +20,28 @@ use support::{ReceivedRequest, json_response, upstream};
 
 struct TokenSource(String);
 
-impl VertexTokenSource for TokenSource {
-    fn project_id(&self) -> VertexAuthFuture<'_, String> {
+impl GcpTokenSource for TokenSource {
+    fn project_id(&self) -> GcpAuthFuture<'_, String> {
         Box::pin(async { Ok(self.0.clone()) })
     }
 
-    fn token(&self) -> VertexAuthFuture<'_, String> {
-        Box::pin(async { Ok(self.0.clone()) })
+    fn token(&self) -> GcpAuthFuture<'_, litellm_secret_value::SecretValue> {
+        Box::pin(async { Ok(litellm_secret_value::SecretValue::new(self.0.clone())) })
     }
 }
 
 #[derive(Default)]
 struct Loader(AtomicUsize);
 
-impl VertexProviderLoader for Loader {
-    fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
+impl GcpProviderLoader for Loader {
+    fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
         Box::pin(async move {
             self.0.fetch_add(1, Ordering::SeqCst);
             let identity = match source {
                 CredentialSource::Trusted(secret) => secret.expose().to_string(),
                 other => panic!("unexpected credential source: {other:?}"),
             };
-            Ok(Arc::new(TokenSource(identity)) as Arc<dyn VertexTokenSource>)
+            Ok(Arc::new(TokenSource(identity)) as Arc<dyn GcpTokenSource>)
         })
     }
 }

litellm-rust/crates/secret-value/src/secret_value.rs

@@ -16,6 +16,18 @@ impl SecretValue {
     }
 }
 
+impl From<String> for SecretValue {
+    fn from(value: String) -> Self {
+        Self::new(value)
+    }
+}
+
+impl From<&str> for SecretValue {
+    fn from(value: &str) -> Self {
+        Self::new(value)
+    }
+}
+
 impl PartialEq for SecretValue {
     fn eq(&self, other: &Self) -> bool {
         subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into()

litellm-rust/crates/secrets-aws/src/auth.rs

@@ -39,30 +39,15 @@ impl Credentials {
         Self {
             auth,
             config: AwsAuthConfig {
-                access_key_id: context
-                    .access_key_id
-                    .as_ref()
-                    .map(|value| value.expose().to_owned()),
-                secret_access_key: context
-                    .secret_access_key
-                    .as_ref()
-                    .map(|value| value.expose().to_owned()),
-                session_token: context
-                    .session_token
-                    .as_ref()
-                    .map(|value| value.expose().to_owned()),
+                access_key_id: context.access_key_id.as_ref().cloned(),
+                secret_access_key: context.secret_access_key.as_ref().cloned(),
+                session_token: context.session_token.as_ref().cloned(),
                 region_name: region(settings, environment.as_ref()).ok(),
                 role_name: settings.aws_role_name.clone(),
                 session_name: settings.aws_session_name.clone(),
-                external_id: settings
-                    .aws_external_id
-                    .as_ref()
-                    .map(|v| v.expose().to_owned()),
+                external_id: settings.aws_external_id.as_ref().cloned(),
                 profile_name: settings.aws_profile_name.clone(),
-                web_identity_token: settings
-                    .aws_web_identity_token
-                    .as_ref()
-                    .map(|v| v.expose().to_owned()),
+                web_identity_token: settings.aws_web_identity_token.as_ref().cloned(),
                 sts_endpoint: settings.aws_sts_endpoint.clone(),
             },
             environment,

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -1183,11 +1183,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             include_namespace=True,
             replay_openai_history=provider == "openai",
         )
-        # OpenAI's Responses API rejects an empty input. For a system-only
-        # request, carry the system message as a system-role input item instead
-        # of instructions, mirroring how non-string system content is already
-        # handled in convert_chat_completion_messages_to_responses_api.
-        is_system_only_request: Final = not converted_input_items and converted_instructions is not None
+        # Keep the system text in input instead of instructions when the Responses API needs it there.
+        # It rejects an empty input, and json_object requires the word "json" in input, ignoring instructions.
+        typed_optional_params: Final = cast("Mapping[str, object]", optional_params)  # cast-ok: untyped dict
+        system_goes_in_input: Final = converted_instructions is not None and (
+            not converted_input_items or _payload_type(typed_optional_params.get("response_format")) == "json_object"
+        )
         target_input_items: Final = (
             _without_audio_input_parts(converted_input_items)
             if _drops_audio_input(model, litellm_params)
@@ -1199,12 +1200,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     "type": "message",
                     "role": "system",
                     "content": [{"type": "input_text", "text": converted_instructions}],
-                }
+                },
+                *target_input_items,
             ]
-            if is_system_only_request
+            if system_goes_in_input
             else target_input_items
         )
-        instructions: Final = None if is_system_only_request else converted_instructions
+        instructions: Final = None if system_goes_in_input else converted_instructions
 
         # Build responses API request using the reverse transformation logic
         responses_api_request: Final = ResponsesAPIOptionalRequestParams()

litellm/cost_calculator.py

@@ -621,11 +621,17 @@ def cost_per_token(
         )
         if per_second is not None:
             return per_second
+        if custom_llm_provider == "anthropic":
+            return anthropic_cost_per_token(
+                model=model, usage=usage_block, service_tier=service_tier, model_info=requested_model_info_override
+            )
         return generic_cost_per_token(
             model=model,
             usage=usage_block,
             custom_llm_provider=custom_llm_provider or "",
             service_tier=service_tier,
+            data_residency=data_residency,
+            vertex_location=vertex_location,
             model_info=requested_model_info_override,
         )
 
@@ -3342,15 +3348,35 @@ def batch_cost_calculator(
     return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift
 
 
+def _standard_rate(model_info: ModelInfo, key: str, prompt_tokens: int) -> float | None:
+    return get_cost_per_unit(model_info, get_modality_token_cost_key(model_info, key, prompt_tokens), None)
+
+
+def _cached_audio_batch_fallback(
+    model_info: ModelInfo, batch_cache_read_rate: float | None, cache_read_rate: float | None, prompt_tokens: int
+) -> float:
+    if batch_cache_read_rate is not None:
+        return batch_cache_read_rate
+    standard_audio_rate: Final = _standard_rate(model_info, "cache_read_input_audio_token_cost", prompt_tokens)
+    if standard_audio_rate is not None:
+        return standard_audio_rate
+    return 0.0 if cache_read_rate is None else cache_read_rate
+
+
 def _batch_prompt_cost(
     usage: Usage,
     model_info: ModelInfo,
     input_rate: float,
-    cache_read_rate: float | None,
+    batch_cache_read_rate: float | None,
     cache_creation_rate: float | None,
     cache_creation_rates_by_ttl: Mapping[str, float],
 ) -> float:
     details: Final = parse_prompt_tokens_details(usage)
+    cache_read_rate: Final = (
+        _standard_rate(model_info, "cache_read_input_token_cost", usage.prompt_tokens)
+        if batch_cache_read_rate is None
+        else batch_cache_read_rate
+    )
     cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0
     cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0
     bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl)
@@ -3371,7 +3397,11 @@ def _batch_prompt_cost(
         + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate)
         + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0)
         + cached_audio_tokens
-        * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0)
+        * _batch_rate(
+            model_info,
+            "cache_read_input_audio_token_cost_batches",
+            _cached_audio_batch_fallback(model_info, batch_cache_read_rate, cache_read_rate, usage.prompt_tokens),
+        )
         + (
             calculate_cache_writing_cost(
                 written_tokens,

litellm/integrations/prometheus.py

@@ -54,6 +54,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import is_registered_cust
 from litellm.litellm_core_utils.service_tier_utils import (
     get_service_tier_from_standard_logging_payload,
 )
+from litellm.llms.azure.common_utils import caused_by_entra_rejection
 from litellm.llms.github_copilot.per_user_auth import github_copilot_user_session_from
 from litellm.llms.openai_like.json_loader import JSONProviderRegistry
 from litellm.models.end_user import LiteLLM_EndUserTable
@@ -448,7 +449,7 @@ def _is_deployment_failure(exception: openai.APIError, logged_call: Mapping[str,
     return not (
         isinstance(exception, CallerCredentialAuthenticationError)
         or bool(logged_call.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL))
-        or logged_call.get("api_call_start_time") is None
+        or (logged_call.get("api_call_start_time") is None and not caused_by_entra_rejection(exception))
         or _used_a_caller_credential(logged_call)
     )
 

litellm/litellm_core_utils/litellm_logging.py

@@ -7178,6 +7178,7 @@ class StandardLoggingPayloadSetup:
         init_response_obj: object,
         api_base: str | None = None,
         router_model_id: str | None = None,
+        custom_model_info: ModelInfo | None = None,
     ) -> StandardLoggingModelInformation:
         model_cost_name: Final = select_model_name_for_cost_calc(
             model=base_model if custom_pricing else None,
@@ -7196,6 +7197,8 @@ class StandardLoggingPayloadSetup:
             if router_model_id is not None
             else model_cost_name
         )
+        if model_cost_name is not None and custom_model_info is not None:
+            return StandardLoggingModelInformation(model_map_key=model_cost_name, model_map_value=custom_model_info)
         if model_cost_name is None:
             model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None)
         else:
@@ -7952,6 +7955,7 @@ def get_standard_logging_object_payload(
             init_response_obj=init_response_obj,
             api_base=litellm_params.get("api_base"),
             router_model_id=logging_obj.get_router_model_id(),
+            custom_model_info=logging_obj.custom_model_info if custom_pricing else None,
         )
 
         error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(

litellm/litellm_core_utils/url_utils.py

@@ -102,6 +102,14 @@ class SSRFError(ValueError):
     """Raised when a URL targets a blocked network."""
 
 
+def is_https_url_with_host(url: str) -> bool:
+    try:
+        parsed_url: Final = urlsplit(url)
+        return parsed_url.scheme == "https" and bool(parsed_url.hostname)
+    except ValueError:
+        return False
+
+
 def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str:
     """Percent-encode one user-controlled URL path segment.
 

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -781,12 +781,17 @@ class AnthropicMessagesHandler(BaseTranslation):
             for event in AnthropicMessagesHandler._parse_sse_data_line(line)
         )
 
-    def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
+    def _translate_to_openai(
+        self,
+        data: dict,  # mutable-ok: the request body the adapter reads
+        *,
+        keep_cache_control: bool = False,
+    ) -> ChatCompletionRequest:
         """Translate Anthropic request to OpenAI chat completion format."""
         (
             chat_completion_compatible_request,
             _tool_name_mapping,
-        ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
+        ) = LiteLLMAnthropicMessagesAdapter(keep_cache_control=keep_cache_control).translate_anthropic_to_openai(
             anthropic_message_request=cast(AnthropicMessagesRequest, scannable_tool_result_request(data)),
             preserve_midturn_system=True,
         )
@@ -813,9 +818,13 @@ class AnthropicMessagesHandler(BaseTranslation):
         data: dict,
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
+        *,
+        request_data: dict[str, object] | None = None,  # mutable-ok: the request dict guardrails receive
+        keep_cache_control: bool = False,
     ) -> Mapping[str, object]:
         """
-        Process input messages by applying guardrails to text content.
+        Process input messages by applying guardrails to text content. ``request_data`` is what the guardrail sees as
+        the request when ``data`` is only its body, e.g. a provider pass-through route that nests the body
         """
         messages: Final = data.get("messages")
         if messages is None:
@@ -827,8 +836,17 @@ class AnthropicMessagesHandler(BaseTranslation):
 
         # The top-level prompt is translated on its own below so it can be hoisted in front of
         # any mid-turn system entries and scanned first, aligned with that structured position.
-        translation_source: Final = {key: value for key, value in data.items() if key != "system"}
-        chat_completion_compatible_request: Final = self._translate_to_openai(translation_source)
+        body_model: Final = data.get("model")
+        request_model: Final = (request_data or {}).get("model")
+        model: Final[str] = (
+            body_model
+            if isinstance(body_model, str) and body_model
+            else (request_model if isinstance(request_model, str) else "")
+        )
+        translation_source: Final = {**{key: value for key, value in data.items() if key != "system"}, "model": model}
+        chat_completion_compatible_request: Final = self._translate_to_openai(
+            translation_source, keep_cache_control=keep_cache_control
+        )
 
         full_structured_messages: Final = cast(
             list[AllMessageValues],
@@ -837,7 +855,9 @@ class AnthropicMessagesHandler(BaseTranslation):
         has_midturn_system_message: Final = any(
             str(message.get("role") or "").lower() == "system" for message in full_structured_messages
         )
-        hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data)
+        hoisted_system_message: Final = (
+            None if skip_system else self._hoisted_top_level_system_message(data, model, keep_cache_control)
+        )
         if hoisted_system_message is not None:
             full_structured_messages.insert(0, hoisted_system_message)
         # skip_system already excluded the trusted top-level prompt (it is simply not hoisted);
@@ -899,7 +919,6 @@ class AnthropicMessagesHandler(BaseTranslation):
             if structured_messages:
                 inputs["structured_messages"] = structured_messages
             # Include model information if available
-            model: Final = data.get("model")
             if model:
                 inputs["model"] = model
             original_tool_names: Final = _original_tool_names(
@@ -907,7 +926,7 @@ class AnthropicMessagesHandler(BaseTranslation):
             )
             guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
                 inputs=inputs,
-                request_data=data,
+                request_data=data if request_data is None else request_data,
                 input_type="request",
                 logging_obj=litellm_logging_obj,
             )
@@ -955,6 +974,7 @@ class AnthropicMessagesHandler(BaseTranslation):
                     hoisted_system_message=hoisted_system_message,
                     preserve_system_messages=has_midturn_system_message,
                     original_tool_names=original_tool_names.used,
+                    model=model,
                 )
             else:
                 if guardrailed_texts and len(guardrailed_texts) != len(scanned):
@@ -977,17 +997,20 @@ class AnthropicMessagesHandler(BaseTranslation):
 
         return data
 
-    def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None:
+    def _hoisted_top_level_system_message(
+        self, data: Mapping[str, object], model: str, keep_cache_control: bool
+    ) -> AllMessageValues | None:
         """Return the system message produced by translating the top-level prompt."""
         system: Final = data.get("system")
         if not system:
             return None
         probe: Final = self._translate_to_openai(
             {
-                "model": data.get("model") or "",
+                "model": model,
                 "messages": [],
                 "system": system,
-            }
+            },
+            keep_cache_control=keep_cache_control,
         )
         hoisted: Final = probe.get("messages") or []
         return hoisted[0] if hoisted else None
@@ -1121,6 +1144,7 @@ class AnthropicMessagesHandler(BaseTranslation):
         preserve_system_messages: bool = False,
         *,
         original_tool_names: Mapping[str, str],
+        model: str,
     ) -> None:
         """Write a guardrail's structured-message rewrite back without losing corrections."""
         from litellm.litellm_core_utils.prompt_templates.factory import (
@@ -1129,7 +1153,6 @@ class AnthropicMessagesHandler(BaseTranslation):
         )
 
         _is_system: Final = AnthropicMessagesHandler._is_system
-        model: Final = str(data.get("model") or "")
         converted: Final[list] = []  # mutable-ok: API message payload
 
         def _convert_run(run: list) -> None:  # mutable-ok: API message payload

litellm/llms/anthropic/messages/adapters/transformation.py

@@ -380,8 +380,8 @@ _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.va
 
 
 class LiteLLMAnthropicMessagesAdapter:
-    def __init__(self):
-        pass
+    def __init__(self, *, keep_cache_control: bool = False) -> None:
+        self.keep_cache_control: Final = keep_cache_control
 
     ### FOR [BETA] `/v1/messages` endpoint support
 
@@ -437,7 +437,7 @@ class LiteLLMAnthropicMessagesAdapter:
         cache_control: Final = (
             source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None)
         )
-        if cache_control and model and self.target_consumes_cache_control(model):
+        if cache_control and (self.keep_cache_control or (model and self.target_consumes_cache_control(model))):
             # TypedDict objects support dict operations at runtime
             # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432)
             if isinstance(target, dict):

litellm/llms/azure/common_utils.py

@@ -1,15 +1,16 @@
 import asyncio
 import hashlib
+import itertools
 import json
 import os
 import re
-from collections.abc import Awaitable, Callable, Mapping
+from collections.abc import Awaitable, Callable, Iterator, Mapping
 from contextlib import closing
 from types import MappingProxyType
-from typing import Final, Literal, NamedTuple, cast
+from typing import Final, Literal, NamedTuple, NoReturn, cast
 
 import httpx
-from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
+from openai import APIError, AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 from pydantic import ConfigDict, TypeAdapter, with_config
 from typing_extensions import ReadOnly, TypedDict
 
@@ -20,7 +21,7 @@ from litellm.constants import DEFAULT_MAX_RETRIES
 from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache
 from litellm.litellm_core_utils.url_utils import ensure_api_route
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
-from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.llms.custom_httpx.http_handler import HTTPHandler, MaskedHTTPStatusError
 from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient
 from litellm.secret_managers.get_azure_ad_token_provider import (
     get_azure_ad_token_provider,
@@ -52,6 +53,13 @@ class _AsyncAzureClientParams(TypedDict, total=False, extra_items=object):
 
 
 _ASYNC_AZURE_CLIENT_PARAMS_ADAPTER: Final = TypeAdapter(_AsyncAzureClientParams)
+_AZURE_AD_TOKEN_PROVIDER_ADAPTER: Final = TypeAdapter[Callable[[], str | Awaitable[str]] | None](
+    Callable[[], str | Awaitable[str]] | None
+)
+
+
+def _callable_azure_ad_token_provider(provider: object) -> Callable[[], str | Awaitable[str]] | None:
+    return _AZURE_AD_TOKEN_PROVIDER_ADAPTER.validate_python(provider) if callable(provider) else None
 
 
 class _AzureGatewayClientParams(TypedDict, total=False, extra_items=object):
@@ -81,6 +89,160 @@ class AzureOpenAIError(BaseLLMException):
         )
 
 
+class EntraRejectionError(AzureOpenAIError):
+    pass
+
+
+_GUID: Final = re.compile(r"[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}")
+_AADSTS_REASON: Final = re.compile(r"\bAADSTS(\d+)\b.*?(?:\.(?=[\s'\"\\]|$)|$)", re.MULTILINE)
+_ENTRA_CREDENTIAL_REJECTION_CODES: Final = frozenset(
+    {
+        50012,
+        50013,
+        50027,
+        50034,
+        50053,
+        50055,
+        50057,
+        50059,
+        50076,
+        50079,
+        50126,
+        50128,
+        53003,
+        65001,
+        70001,
+        70002,
+        70011,
+        70021,
+        70025,
+        70052,
+        90002,
+        90094,
+        500011,
+        500014,
+        700016,
+        700023,
+        700024,
+        700025,
+        700027,
+        700211,
+        700213,
+        700223,
+        700238,
+        7000112,
+        7000215,
+        7000218,
+        7000222,
+    }
+)
+
+
+def _entra_rejection_reason(message: str) -> str | None:
+    reason: Final = _AADSTS_REASON.search(message)
+    if reason is None or int(reason.group(1)) not in _ENTRA_CREDENTIAL_REJECTION_CODES:
+        return None
+    return _GUID.sub("<redacted>", reason.group(0))
+
+
+def _entra_rejection(reason: str) -> EntraRejectionError:
+    return EntraRejectionError(status_code=401, message=f"Failed to get Azure AD token: {reason}")
+
+
+def entra_rejection_as_auth_error(error: Exception) -> EntraRejectionError | None:
+    try:
+        from azure.core.exceptions import ClientAuthenticationError
+    except ImportError:
+        return None
+    if not isinstance(error, ClientAuthenticationError):
+        return None
+    status_code: Final = error.status_code
+    if status_code is not None and (status_code == 429 or status_code >= 500):
+        return None
+    reason: Final = _entra_rejection_reason(str(error))
+    return None if reason is None else _entra_rejection(reason)
+
+
+def _raise_entra_rejection_as_auth_error(error: Exception) -> NoReturn:
+    auth_error: Final = entra_rejection_as_auth_error(error)
+    if auth_error is None:
+        raise error
+    raise auth_error from error
+
+
+def entra_token(provider: Callable[[], str]) -> str:
+    try:
+        return provider()
+    except entra_rejection_types() as error:
+        _raise_entra_rejection_as_auth_error(error)
+
+
+def entra_rejection_types() -> tuple[type[Exception], ...]:
+    try:
+        from azure.core.exceptions import ClientAuthenticationError
+    except ImportError:
+        return ()
+    return (ClientAuthenticationError,)
+
+
+_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16
+
+
+def _exception_chain(error: BaseException) -> Iterator[BaseException]:
+    link: BaseException | None = error  # rebind-ok: advances one link of the cause chain per iteration
+    for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH):
+        if link is None:
+            return
+        yield link
+        link = link.__cause__ or link.__context__
+
+
+_REQUEST_FAILURE_TYPES: Final = (APIError, httpx.HTTPError)
+
+
+def _ends_the_failed_attempt(link: BaseException, failed_attempt: BaseException) -> bool:
+    return not isinstance(link, EntraRejectionError) and (
+        isinstance(link, _REQUEST_FAILURE_TYPES) or type(link) is type(failed_attempt)
+    )
+
+
+def _links_within_the_failed_attempt(error: BaseException) -> Iterator[BaseException]:
+    yield error
+    yield from itertools.takewhile(
+        lambda link: not _ends_the_failed_attempt(link, error), itertools.islice(_exception_chain(error), 1, None)
+    )
+
+
+def entra_rejection_in(error: BaseException) -> EntraRejectionError | None:
+    return next(
+        (link for link in _links_within_the_failed_attempt(error) if isinstance(link, EntraRejectionError)), None
+    )
+
+
+def caused_by_entra_rejection(error: BaseException) -> bool:
+    return entra_rejection_in(error) is not None
+
+
+async def _awaited_token(token: Awaitable[str]) -> str:
+    try:
+        return await token
+    except entra_rejection_types() as error:
+        _raise_entra_rejection_as_auth_error(error)
+
+
+def _entra_rejections_as_auth_errors(
+    provider: Callable[[], str | Awaitable[str]],
+) -> Callable[[], str | Awaitable[str]]:
+    def call_provider() -> str | Awaitable[str]:
+        try:
+            token: Final = provider()
+        except entra_rejection_types() as error:
+            _raise_entra_rejection_as_auth_error(error)
+        return token if isinstance(token, str) else _awaited_token(token)
+
+    return call_provider
+
+
 def process_azure_headers(headers: httpx.Headers | dict) -> dict:
     openai_headers: Final = {}
     if "x-ratelimit-limit-requests" in headers:
@@ -91,7 +253,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict:
         openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"]
     if "x-ratelimit-remaining-tokens" in headers:
         openai_headers["x-ratelimit-remaining-tokens"] = headers["x-ratelimit-remaining-tokens"]
-    llm_response_headers: Final = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()}
+    llm_response_headers: Final = {f"llm_provider-{k}": v for k, v in headers.items()}
 
     return {**llm_response_headers, **openai_headers}
 
@@ -257,22 +419,34 @@ def _cached_oidc_token(cache_key: str) -> str | None:
     return cached if isinstance(cached, str) else None
 
 
+def _oidc_exchange_rejection_reason(error: MaskedHTTPStatusError) -> str | None:
+    if error.response.status_code not in (400, 401):
+        return None
+    return _entra_rejection_reason(error.text or error.response.text)
+
+
 def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str:
     cached_access_token: Final = _cached_oidc_token(cache_key)
     if cached_access_token is not None:
         return cached_access_token
 
     with closing(HTTPHandler(timeout=litellm.request_timeout)) as client:
-        req_token: Final = client.post(
-            token_url,
-            data={
-                "client_id": client_id,
-                "grant_type": "client_credentials",
-                "scope": scope,
-                "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
-                "client_assertion": oidc_token,
-            },
-        )
+        try:
+            req_token: Final = client.post(
+                token_url,
+                data={
+                    "client_id": client_id,
+                    "grant_type": "client_credentials",
+                    "scope": scope,
+                    "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+                    "client_assertion": oidc_token,
+                },
+            )
+        except MaskedHTTPStatusError as error:
+            reason: Final = _oidc_exchange_rejection_reason(error)
+            if reason is None:
+                raise
+            raise _entra_rejection(reason) from error
 
     if req_token.status_code != 200:
         raise AzureOpenAIError(
@@ -483,6 +657,9 @@ def get_azure_ad_token(
             raise
         except Exception as e:
             verbose_logger.error("Error calling Azure AD token provider: %s", e)
+            auth_error: Final = entra_rejection_as_auth_error(e)
+            if auth_error is not None:
+                raise auth_error from e
             raise RuntimeError(f"Failed to get Azure AD token: {e}") from e
 
     return azure_ad_token
@@ -696,7 +873,7 @@ class BaseAzureLLM(BaseOpenAILLM):
         api_version: str | None,
         is_async: bool,
     ) -> dict:
-        azure_ad_token_provider = litellm_params.get("azure_ad_token_provider")
+        azure_ad_token_provider = _callable_azure_ad_token_provider(litellm_params.get("azure_ad_token_provider"))
         # If we have api_key, then we have higher priority
         azure_ad_token = litellm_params.get("azure_ad_token")
 
@@ -775,7 +952,7 @@ class BaseAzureLLM(BaseOpenAILLM):
             azure_client_params["timeout"] = timeout
 
         if azure_ad_token_provider is not None:
-            azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider
+            azure_client_params["azure_ad_token_provider"] = _entra_rejections_as_auth_errors(azure_ad_token_provider)
         # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client
         # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client
 

litellm/llms/azure/decisions/transformation.py

@@ -8,6 +8,7 @@ import litellm
 from litellm.llms.azure.common_utils import (
     BaseAzureLLM,
     cached_azure_ad_token_refresh_provider,
+    entra_token,
 )
 from litellm.llms.openai.decisions.transformation import OpenAIDecisionsConfig
 from litellm.secret_managers.main import get_secret_str
@@ -61,7 +62,7 @@ class AzureDecisionsConfig(OpenAIDecisionsConfig):
         scope: Final = (
             params.azure_scope or get_secret_str("AZURE_SCOPE") or "https://cognitiveservices.azure.com/.default"
         )
-        azure_headers["Authorization"] = f"Bearer {cached_azure_ad_token_refresh_provider(scope)()}"
+        azure_headers["Authorization"] = f"Bearer {entra_token(cached_azure_ad_token_refresh_provider(scope))}"
         return azure_headers
 
     async def avalidate_environment(

litellm/llms/azure/search/transformation.py

@@ -36,7 +36,7 @@ import httpx
 from pydantic import ConfigDict, ValidationError
 
 from litellm.litellm_core_utils.url_utils import ensure_api_route
-from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider
+from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider, entra_rejection_as_auth_error
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.search.responses import (
     ResponsesSearchAnnotation,
@@ -270,6 +270,9 @@ class BingGroundingSearchConfig(BaseSearchConfig):
         try:
             return minter()
         except Exception as e:
+            rejection: Final = entra_rejection_as_auth_error(e)
+            if rejection is not None:
+                raise rejection from e
             raise ValueError(
                 f"Grounding with Bing Search: no credential available. Pass api_key, set {TOKEN_ENV} "
                 f"to an Entra bearer token, or configure azure-identity (AZURE_CLIENT_ID / "

litellm/llms/azure_ai/anthropic/count_tokens/handler.py

@@ -16,6 +16,7 @@ from litellm.anthropic_beta_headers_manager import aload_beta_headers_config
 from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
 from litellm.llms.anthropic.common_utils import AnthropicError
+from litellm.llms.azure.common_utils import EntraRejectionError
 from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
     AzureAIAnthropicCountTokensConfig,
 )
@@ -137,6 +138,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
         except AnthropicError:
             # Re-raise Anthropic exceptions as-is
             raise
+        except EntraRejectionError as rejection:
+            raise AnthropicError(status_code=rejection.status_code, message=rejection.message) from rejection
         except httpx.HTTPStatusError as e:
             # HTTP errors - preserve the actual status code
             verbose_logger.error("HTTP error in CountTokens handler: %s", e)

litellm/llms/azure_ai/common_utils.py

@@ -1,5 +1,7 @@
 import asyncio
-from collections.abc import Mapping
+import os
+from collections.abc import Callable, Mapping
+from functools import partial
 from types import MappingProxyType
 from typing import Final, Literal
 from urllib.parse import urlparse
@@ -52,6 +54,28 @@ AZURE_ENTRA_CREDENTIAL_PARAM_KEYS: Final = frozenset({"azure_ad_token", "client_
 AZURE_ENTRA_LITELLM_PARAM_KEYS: Final = AZURE_ENTRA_CREDENTIAL_PARAM_KEYS | frozenset(
     {"tenant_id", "client_id", "azure_username", "azure_scope"}
 )
+AZURE_DEFAULT_SCOPE: Final = "https://cognitiveservices.azure.com/.default"
+_AZURE_ENTRA_ENVIRONMENT_VARIABLES: Final = MappingProxyType(
+    {
+        "tenant_id": "AZURE_TENANT_ID",
+        "client_id": "AZURE_CLIENT_ID",
+        "client_secret": "AZURE_CLIENT_SECRET",
+        "azure_username": "AZURE_USERNAME",
+        "azure_password": "AZURE_PASSWORD",
+        "azure_ad_token": "AZURE_AD_TOKEN",
+    }
+)
+
+
+def azure_entra_environment_defaults() -> Mapping[str, object]:
+    return MappingProxyType(
+        {
+            **{key: os.getenv(variable) for key, variable in _AZURE_ENTRA_ENVIRONMENT_VARIABLES.items()},
+            "azure_scope": os.getenv("AZURE_SCOPE", AZURE_DEFAULT_SCOPE),
+        }
+    )
+
+
 AZURE_ENTRA_CREDENTIAL_HELP: Final = (
     "Set `tenant_id` + `client_id` + `client_secret`, `azure_ad_token` (an `oidc/` token also needs "
     "`tenant_id` + `client_id`), or `client_id` + `azure_username` + `azure_password` in the agent's `litellm_params`"
@@ -70,6 +94,17 @@ def _resolve_config_secret(value: object) -> str | None:
     return get_secret_str(value) if value.startswith("os.environ/") else value
 
 
+def _agent_entra_token(provider: Callable[[], str], model: str) -> str:
+    from litellm.llms.azure.common_utils import EntraRejectionError, entra_token
+
+    try:
+        return entra_token(provider)
+    except EntraRejectionError as rejection:
+        raise litellm.AuthenticationError(
+            message=rejection.message, llm_provider="azure_ai", model=model
+        ) from rejection
+
+
 def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str:
     """Mints the Entra bearer from the agent's own litellm_params, never from process-wide AZURE_* env vars."""
     from litellm.llms.azure.common_utils import (
@@ -88,18 +123,32 @@ def get_azure_ai_agent_entra_token(litellm_params: Mapping[str, object]) -> str:
     azure_username: Final = resolved["azure_username"]
     azure_password: Final = resolved["azure_password"]
     azure_ad_token: Final = resolved["azure_ad_token"]
+    model: Final = str(litellm_params.get("model") or "")
     if tenant_id and client_id and client_secret:
-        return get_azure_ad_token_from_entra_id(
-            tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope
-        )()
+        return _agent_entra_token(
+            get_azure_ad_token_from_entra_id(
+                tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope
+            ),
+            model,
+        )
     if client_id and azure_username and azure_password:
-        return get_azure_ad_token_from_username_password(
-            client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope
-        )()
+        return _agent_entra_token(
+            get_azure_ad_token_from_username_password(
+                client_id=client_id, azure_username=azure_username, azure_password=azure_password, scope=scope
+            ),
+            model,
+        )
     federated: Final = azure_ad_token is not None and azure_ad_token.startswith("oidc/")
     if azure_ad_token and federated and tenant_id and client_id:
-        return get_azure_ad_token_from_oidc(
-            azure_ad_token=azure_ad_token, azure_client_id=client_id, azure_tenant_id=tenant_id, scope=scope
+        return _agent_entra_token(
+            partial(
+                get_azure_ad_token_from_oidc,
+                azure_ad_token=azure_ad_token,
+                azure_client_id=client_id,
+                azure_tenant_id=tenant_id,
+                scope=scope,
+            ),
+            model,
         )
     if azure_ad_token and not federated:
         return azure_ad_token

litellm/llms/bedrock/passthrough/guardrail_translation/handler.py

@@ -14,7 +14,7 @@ from pydantic import InstanceOf, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
-from litellm.litellm_core_utils.core_helpers import decoded_json_object
+from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object
 from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
 from litellm.llms.base_llm.guardrail_translation.base_translation import (
     BaseTranslation,
@@ -61,6 +61,8 @@ if TYPE_CHECKING:
 
 _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
 _INVOKE_STREAM_ACTION: Final = "invoke-with-response-stream"
+_INVOKE_ACTIONS: Final = frozenset({"invoke", _INVOKE_STREAM_ACTION})
+_BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN: Final = frozenset({"document", "search_result"})
 _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream"
 _EVENT_STREAM_MEDIA_TYPE: Final = "application/vnd.amazon.eventstream"
 
@@ -74,6 +76,83 @@ def _is_invoke_stream_endpoint(endpoint: str) -> bool:
     return endpoint.rstrip("/").split("/")[-1] == _INVOKE_STREAM_ACTION
 
 
+def _anthropic_invoke_body(
+    data: Mapping[str, object],
+) -> dict[str, object] | None:  # mutable-ok: the body guardrails write back into
+    """The body of an InvokeModel call in Anthropic Messages format, which Bedrock requires to carry anthropic_version"""
+    endpoint: Final = data.get("endpoint")
+    body: Final = data.get("data")
+    if not isinstance(endpoint, str) or endpoint.rstrip("/").split("/")[-1] not in _INVOKE_ACTIONS:
+        return None
+    if not isinstance(body, dict):
+        return None
+    invoke_body: Final = _BODY_BY_REFERENCE.validate_python(body)
+    if not isinstance(invoke_body.get("anthropic_version"), str) or not isinstance(invoke_body.get("messages"), list):
+        return None
+    messages: Final = _ARRAY_BY_REFERENCE.validate_python(invoke_body["messages"])
+    if not all(_is_well_formed_message(message) for message in messages):
+        return None
+    blocks: Final = _anthropic_message_blocks(messages)
+    if any(block.get("type") in _BLOCK_TYPES_THE_ANTHROPIC_HANDLER_DOES_NOT_SCAN for block in blocks):
+        return None
+    return invoke_body
+
+
+def _body_level_request_data(
+    passthrough_request: Mapping[str, object],
+) -> dict[str, object]:  # mutable-ok: the request dict guardrails receive
+    """What `/v1/messages` hands a guardrail as `request_data`: the Anthropic body at the top level, with the proxy's own keys (model, metadata, headers) on top"""
+    body: Final = as_str_mapping(passthrough_request.get("data")) or {}
+    proxy_keys: Final = {key: value for key, value in passthrough_request.items() if key != "data"}
+    return {**body, **proxy_keys}
+
+
+def _is_well_formed_message(message: object) -> bool:
+    """Whether the Messages translation can read `message`, so a malformed body still reaches Bedrock to be rejected there"""
+    fields: Final = as_str_mapping(message)
+    if fields is None or not isinstance(fields.get("role"), str):
+        return False
+    content: Final = fields.get("content")
+    if isinstance(content, str):
+        return True
+    return isinstance(content, list) and all(
+        as_str_mapping(block) is not None for block in _ARRAY_BY_REFERENCE.validate_python(content)
+    )
+
+
+def _anthropic_message_blocks(messages: Sequence[object]) -> tuple[Mapping[str, object], ...]:
+    fields: Final = (as_str_mapping(message) for message in messages)
+    return tuple(chain.from_iterable(_anthropic_content_blocks(field.get("content")) for field in fields if field))
+
+
+def _anthropic_content_blocks(content: object) -> tuple[Mapping[str, object], ...]:
+    blocks: Final[list[Mapping[str, object]]] = []  # mutable-ok: collected while walking the content
+    pending: Final[list[object]] = [content]  # mutable-ok: content lists still to walk, nested tool results included
+    while pending:
+        current = pending.pop()
+        if not isinstance(current, list):
+            continue
+        for item in _ARRAY_BY_REFERENCE.validate_python(current):
+            block = as_str_mapping(item)
+            if block is None:
+                continue
+            blocks.append(block)
+            if block.get("type") == "tool_result":
+                pending.append(block.get("content"))
+    return tuple(blocks)
+
+
+def _anthropic_invoke_reply(
+    response: object,
+    request_data: object,
+) -> dict[str, object] | None:  # mutable-ok: the reply guardrails write back into
+    if not isinstance(response, dict) or not isinstance(request_data, dict):
+        return None
+    if _anthropic_invoke_body(_OBJECT_BY_REFERENCE.validate_python(request_data)) is None:
+        return None
+    return _BODY_BY_REFERENCE.validate_python(response)
+
+
 def _generic_passthrough_handler() -> "PassThroughEndpointHandler":
     """
     Fallback for non-Converse Bedrock routes (e.g. invoke). The generic
@@ -1287,6 +1366,17 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
         body: Final = data.get("data")
 
         if not _is_converse_endpoint(endpoint):
+            passthrough_request: Final = _BODY_BY_REFERENCE.validate_python(data)
+            anthropic_body: Final = _anthropic_invoke_body(passthrough_request)
+            if anthropic_body is not None:
+                await AnthropicMessagesHandler().process_input_messages(  # pyright: ignore[reportUnknownMemberType]  # takes the untyped request dict
+                    data=anthropic_body,
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=litellm_logging_obj,
+                    request_data=_body_level_request_data(passthrough_request),
+                    keep_cache_control=True,
+                )
+                return passthrough_request
             return await _generic_passthrough_handler().process_input_messages(
                 data=data,
                 guardrail_to_apply=guardrail_to_apply,
@@ -1412,6 +1502,15 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
             and not _is_converse_endpoint(endpoint)
             and not _is_invoke_stream_endpoint(endpoint)
         ):
+            anthropic_reply: Final = _anthropic_invoke_reply(response, request_data)  # pyright: ignore[reportUnknownArgumentType]  # the base signature types request_data as a bare dict
+            if anthropic_reply is not None:
+                return await AnthropicMessagesHandler().process_output_response(  # pyright: ignore[reportUnknownMemberType]  # takes the untyped request dict
+                    response=anthropic_reply,  # pyright: ignore[reportArgumentType]  # the Anthropic reply dict, written back in place
+                    guardrail_to_apply=guardrail_to_apply,
+                    litellm_logging_obj=litellm_logging_obj,
+                    user_api_key_dict=user_api_key_dict,
+                    request_data=_body_level_request_data(_OBJECT_BY_REFERENCE.validate_python(request_data or {})),
+                )
             return await _generic_passthrough_handler().process_output_response(
                 response=response,
                 guardrail_to_apply=guardrail_to_apply,

litellm/llms/github_copilot/github_host.py

@@ -21,8 +21,10 @@ class GithubHostEndpoints:
     device_code_url: str
     access_token_url: str
     api_key_url: str
+    copilot_user_url: str
     user_url: str
     client_id: str
+    exchanges_copilot_token: bool
     copilot_api_hosts: frozenset[str]
     copilot_api_host_suffix: str | None
     default_copilot_api_base: str
@@ -49,14 +51,17 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints:
     if normalized is None:
         raise ValueError("host must be github.com or a single-subdomain GHE.com hostname")
     client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
+    exchanges_copilot_token: Final = client_id == DEFAULT_GITHUB_CLIENT_ID
     if normalized == GITHUB_COPILOT_DEFAULT_HOST:
         return GithubHostEndpoints(
             host=normalized,
             device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL),
             access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL),
             api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL),
+            copilot_user_url="https://api.github.com/copilot_internal/user",
             user_url="https://api.github.com/user",
             client_id=client_id,
+            exchanges_copilot_token=exchanges_copilot_token,
             copilot_api_hosts=frozenset({"githubcopilot.com"}),
             copilot_api_host_suffix=".githubcopilot.com",
             default_copilot_api_base="https://api.githubcopilot.com",
@@ -70,8 +75,10 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints:
         device_code_url=f"https://{normalized}/login/device/code",
         access_token_url=f"https://{normalized}/login/oauth/access_token",
         api_key_url=f"https://{api_host}/copilot_internal/v2/token",
+        copilot_user_url=f"https://{api_host}/copilot_internal/user",
         user_url=f"https://{api_host}/user",
         client_id=client_id,
+        exchanges_copilot_token=exchanges_copilot_token,
         copilot_api_hosts=frozenset({copilot_api_host}),
         copilot_api_host_suffix=None,
         default_copilot_api_base=f"https://{copilot_api_host}",

litellm/llms/github_copilot/per_user_auth.py

@@ -46,6 +46,7 @@ _LLM_PROVIDER: Final = "github_copilot"
 _DEVICE_FLOW_SCOPE: Final = "read:user"
 _SESSION_CACHE_MAX_SIZE: Final = 1000
 _SYNC_LOCK_COUNT: Final = 64
+_ACCOUNT_LOOKUP_TTL_SECONDS: Final = 1800
 
 
 class _SessionCache(Protocol):
@@ -125,12 +126,20 @@ class _CopilotEndpointsPayload(TypedDict):
     api: ReadOnly[NotRequired[object]]
 
 
+@with_config(ConfigDict(extra="allow", strict=True))
+class _CopilotAccountPayload(TypedDict):
+    access_type_sku: ReadOnly[NotRequired[object]]
+    chat_enabled: ReadOnly[NotRequired[object]]
+    endpoints: ReadOnly[NotRequired[object]]
+
+
 @with_config(ConfigDict(extra="allow", strict=True))
 class _GithubUserPayload(TypedDict):
     login: ReadOnly[str]
 
 
 _COPILOT_TOKEN: Final = TypeAdapter(_CopilotTokenPayload)
+_COPILOT_ACCOUNT: Final = TypeAdapter(_CopilotAccountPayload)
 _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload)
 _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload)
 _GITHUB_USER: Final = TypeAdapter(_GithubUserPayload)
@@ -273,6 +282,8 @@ def _session_from_response(
             llm_provider=_LLM_PROVIDER,
             model="",
         )
+    if not endpoints.exchanges_copilot_token:
+        return _session_from_account(response, github_token, credential_name, cache_key, endpoints)
     try:
         payload: Final = _COPILOT_TOKEN.validate_python(response.json())
     except Exception as e:
@@ -282,13 +293,7 @@ def _session_from_response(
             model="",
         ) from e
 
-    payload_endpoints: Final[object] = payload.get("endpoints")
-    endpoints_api: Final[object] = (
-        _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
-        if isinstance(payload_endpoints, Mapping)
-        else None
-    )
-    api_base: Final = validated_copilot_api_base(endpoints_api, endpoints)
+    api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints)
     session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base)
     expires_at: Final = _expires_at_seconds(payload.get("expires_at"))
     if expires_at > 0:
@@ -296,6 +301,41 @@ def _session_from_response(
     return session
 
 
+def _session_from_account(
+    response: httpx.Response,
+    github_token: str,
+    credential_name: str,
+    cache_key: str,
+    endpoints: GithubHostEndpoints,
+) -> GithubCopilotUserSession:
+    try:
+        payload: Final = _COPILOT_ACCOUNT.validate_python(response.json())
+    except Exception as e:
+        raise ServiceUnavailableError(
+            message="GitHub Copilot account lookup returned an invalid response",
+            llm_provider=_LLM_PROVIDER,
+            model="",
+        ) from e
+    if payload.get("chat_enabled") is False or payload.get("access_type_sku") == "no_access":
+        _SESSION_CACHE.delete_cache(cache_key)
+        raise _reject_connection(credential_name)
+
+    api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints)
+    session: Final = GithubCopilotUserSession(token=github_token, api_base=api_base)
+    _cache_session(cache_key, github_token, session, time.time() + _ACCOUNT_LOOKUP_TTL_SECONDS)
+    return session
+
+
+def _endpoints_api(payload_endpoints: object) -> object:
+    if not isinstance(payload_endpoints, Mapping):
+        return None
+    return _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api")
+
+
+def _session_lookup_url(endpoints: GithubHostEndpoints) -> str:
+    return endpoints.api_key_url if endpoints.exchanges_copilot_token else endpoints.copilot_user_url
+
+
 def _to_int(raw: object) -> int:
     if isinstance(raw, bool):
         return 0
@@ -337,7 +377,7 @@ def _do_exchange(
             _SyncGetClient, litellm.module_level_client
         )
         response: Final = client.get(
-            endpoints.api_key_url,
+            _session_lookup_url(endpoints),
             headers=github_api_headers(github_token),
         )
     except Exception as e:
@@ -367,7 +407,7 @@ async def _do_aexchange(
     client: Final = _new_async_github_client()
     try:
         response: Final = await client.get(
-            endpoints.api_key_url,
+            _session_lookup_url(endpoints),
             headers=github_api_headers(github_token),
         )
     except Exception as e:

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
 from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
+from litellm.llms.azure.common_utils import caused_by_entra_rejection
 from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
 from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 from litellm.proxy.a2a.agent_card import point_agent_card_at_proxy
@@ -1202,6 +1203,8 @@ async def invoke_agent_a2a(
             pass
         if isinstance(e, litellm.BadRequestError):
             return _jsonrpc_error(body.get("id"), -32602, e.message, 400)
+        if isinstance(e, litellm.AuthenticationError) and caused_by_entra_rejection(e):
+            return _jsonrpc_error(body.get("id"), -32000, e.message, 401)
         return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500)
 
 

litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py

@@ -26,7 +26,6 @@ from collections.abc import Mapping, Sequence
 from datetime import datetime, timezone
 from types import MappingProxyType
 from typing import Annotated, Final, Protocol, TypedDict, get_args
-from urllib.parse import urlsplit
 
 from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import JSONResponse, Response
@@ -34,6 +33,7 @@ from pydantic import Json, TypeAdapter
 from typing_extensions import ReadOnly
 
 from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.url_utils import is_https_url_with_host
 from litellm.models.plugin_upload import PluginUpload
 from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
 from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import (
@@ -377,14 +377,6 @@ _VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-
 _VALID_SHA256_RE: Final = re.compile(r"^[0-9a-fA-F]{64}$")
 
 
-def _is_https_url_with_host(url: str) -> bool:
-    try:
-        parts: Final = urlsplit(url)
-    except ValueError:
-        return False
-    return parts.scheme == "https" and bool(parts.hostname)
-
-
 def _validate_plugin_source(source: Mapping[str, str]) -> None:
     """Validate plugin source format, raising HTTPException on invalid input."""
     source_type: Final = source.get("source")
@@ -421,7 +413,7 @@ def _validate_plugin_source(source: Mapping[str, str]) -> None:
                 },
             )
     elif source_type == "archive":
-        if not _is_https_url_with_host(source.get("url", "")):
+        if not is_https_url_with_host(source.get("url", "")):
             raise HTTPException(
                 status_code=400,
                 detail={

litellm/proxy/enterprise_billing/billing_metrics.py

@@ -2,29 +2,33 @@
 Push-based OTLP metering for enterprise litellm deployments.
 
 Owns a dedicated OpenTelemetry meter provider and an OTLP/HTTP exporter
-authenticated to our global collector with a TLS client certificate. The
-collector front end terminates mutual TLS: the client certificate presented
-here is validated against our CA at the edge, and the verified subject is
-what identifies the deployment. It is intentionally isolated from the global
-meter provider so the customer's own OTEL metrics are untouched and ours
-never leak into their backend.
+authenticated to our global collector with a Cognito bearer token. It is
+intentionally isolated from the global meter provider so the customer's own
+OTEL metrics are untouched and ours never leak into their backend.
 
-The deployment's identity rides on the TLS client certificate, not on the
-payload; the secret license key is never sent as an attribute or header.
+The secret license key is never sent as an attribute or header.
 """
 
 import os
-import tempfile
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Final, Optional
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, TypedDict
 
+import requests
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
 from opentelemetry.metrics import Counter
 from opentelemetry.sdk.metrics import MeterProvider
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 from opentelemetry.sdk.resources import Resource
+from pydantic import BaseModel, ConfigDict
+from requests.auth import AuthBase
+from typing_extensions import ReadOnly, Unpack
 
 from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.url_utils import is_https_url_with_host
 from litellm.proxy.middleware.billable_request_metrics_middleware import (
     BillableCategory,
 )
@@ -33,23 +37,15 @@ if TYPE_CHECKING:
     from litellm.proxy._types import EnterpriseLicenseData
 
 ENDPOINT_ENV: Final = "LITELLM_BILLING_METRICS_ENDPOINT"
-CLIENT_CERT_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_CERT"
-CLIENT_KEY_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_KEY"
-CA_CERT_ENV: Final = "LITELLM_BILLING_METRICS_CA_CERT"
+CLIENT_ID_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_ID"
+CLIENT_SECRET_ENV: Final = "LITELLM_BILLING_METRICS_CLIENT_SECRET"
+TOKEN_URL_ENV: Final = "LITELLM_BILLING_METRICS_TOKEN_URL"
 EXPORT_INTERVAL_ENV: Final = "LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS"
 DEFAULT_EXPORT_INTERVAL_MS: Final = 60_000
 SHUTDOWN_FLUSH_TIMEOUT_MS: Final = 5_000
+_TOKEN_REFRESH_SKEW_SECONDS: Final = 60
 _METRICS_PATH: Final = "/v1/metrics"
 
-# The cert env vars take a path or the PEM itself. Secret stores that inject
-# values as env content cannot mount them as files, so inline PEM is written out.
-_PEM_PREFIX: Final = "-----BEGIN"
-_PEM_DIR_PREFIX: Final = "litellm-billing-mtls-"
-_PEM_FILE_MODE: Final = 0o600
-_CLIENT_CERT_FILENAME: Final = "client.crt"
-_CLIENT_KEY_FILENAME: Final = "client.key"
-_CA_CERT_FILENAME: Final = "ca.crt"
-
 METRIC_NAME: Final = "litellm.enterprise.billable_requests"
 METER_NAME: Final = "litellm.enterprise.billing"
 
@@ -59,21 +55,28 @@ AttributeValue = str | int
 @dataclass(frozen=True, slots=True)
 class BillingMetricsConfig:
     endpoint: str
-    client_cert_path: str
-    client_key_path: str
-    ca_cert_path: str | None
+    token_url: str
+    client_id: str
+    client_secret: str = field(repr=False)
     export_interval_ms: int
     litellm_version: str
     license_id: str | None
 
 
+class _BillingMetricsOAuthToken(BaseModel):
+    model_config = ConfigDict(frozen=True, strict=True)
+
+    access_token: str
+    expires_in: int
+
+
 def _metrics_endpoint(endpoint: str) -> str:
     """The OTLP/HTTP metric exporter wants the full URL including the signal path."""
     trimmed: Final = endpoint.rstrip("/")
     return trimmed if trimmed.endswith(_METRICS_PATH) else f"{trimmed}{_METRICS_PATH}"
 
 
-def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]:
+def resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]:
     base: Final[dict[str, AttributeValue]] = {
         "service.name": "litellm-proxy",
         "litellm.version": config.litellm_version,
@@ -82,7 +85,7 @@ def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeVal
     return {**base, **license_attr}
 
 
-def _billable_attributes(
+def billable_attributes(
     category: BillableCategory, route: str, status_code: int, model_id: str | None
 ) -> dict[str, AttributeValue]:
     base: Final[dict[str, AttributeValue]] = {
@@ -94,23 +97,95 @@ def _billable_attributes(
     return {**base, **model_attr}
 
 
-def build_mtls_meter_provider(config: BillingMetricsConfig) -> MeterProvider:
-    """OTLP/HTTP exporter presenting a TLS client certificate.
+class BillingMetricsBearerAuth(AuthBase):
+    def __init__(
+        self,
+        *,
+        client_id: str,
+        client_secret: str,
+        token_url: str,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self._client_id: Final = client_id
+        self._client_secret: Final = client_secret
+        self._token_url: Final = token_url
+        self._clock: Final = clock
+        self._lock: Final = threading.Lock()
+        self._cached_token: str | None = None
+        self._refresh_at: float = 0.0
 
-    The collector's load balancer terminates mutual TLS and validates the client
-    certificate against our CA. Server verification uses the system trust store
-    (the collector presents a public web-PKI certificate); ca_cert_path overrides
-    it only for private/test collectors.
-    """
-    exporter: Final = OTLPMetricExporter(
-        endpoint=_metrics_endpoint(config.endpoint),
-        # None -> exporter falls back to the system trust store.
-        certificate_file=config.ca_cert_path,
-        client_certificate_file=config.client_cert_path,
-        client_key_file=config.client_key_path,
+    def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
+        token: Final = self._get_token()
+        authenticated_request: Final = request.copy()
+        authenticated_request.headers["Authorization"] = f"Bearer {token}"
+        return authenticated_request
+
+    def invalidate(self) -> None:
+        with self._lock:
+            self._cached_token = None
+            self._refresh_at = 0.0
+
+    def _get_token(self) -> str:
+        with self._lock:
+            now: Final = self._clock()
+            if self._cached_token is not None and now < self._refresh_at:
+                return self._cached_token
+            token, cache_seconds = self._fetch_token()
+            self._cached_token = token
+            self._refresh_at = now + cache_seconds
+            return token
+
+    def _fetch_token(self) -> tuple[str, float]:
+        try:
+            # The OTel periodic exporter's background thread makes blocking token I/O safe here.
+            response: Final = requests.post(
+                self._token_url,
+                auth=(self._client_id, self._client_secret),
+                data={"grant_type": "client_credentials", "scope": "telemetry/write"},
+                timeout=10,
+            )
+            response.raise_for_status()
+            token_response: Final = _BillingMetricsOAuthToken.model_validate(response.json())
+            return token_response.access_token, max(0.0, float(token_response.expires_in) - _TOKEN_REFRESH_SKEW_SECONDS)
+        except (requests.RequestException, ValueError):
+            raise requests.RequestException("Billing metrics token request failed") from None
+
+
+class _SessionHookKwargs(TypedDict, total=False):
+    stream: ReadOnly[object]
+    timeout: ReadOnly[object]
+    verify: ReadOnly[object]
+    cert: ReadOnly[object]
+    proxies: ReadOnly[object]
+
+
+class _BillingMetricsBearerSession(requests.Session):
+    def __init__(self, auth: BillingMetricsBearerAuth) -> None:
+        super().__init__()
+        self._bearer_auth: Final = auth
+        self.auth = auth
+        self.hooks = {"response": [self._invalidate_token_on_unauthorized]}
+
+    def _invalidate_token_on_unauthorized(
+        self,
+        response: requests.Response,
+        **kwargs: Unpack[_SessionHookKwargs],  # kwargs-ok: requests hooks forward send options
+    ) -> requests.Response:
+        if response.status_code == 401:
+            self._bearer_auth.invalidate()
+        return response
+
+
+def build_bearer_meter_provider(config: BillingMetricsConfig) -> MeterProvider:
+    auth: Final = BillingMetricsBearerAuth(
+        client_id=config.client_id,
+        client_secret=config.client_secret,
+        token_url=config.token_url,
     )
+    session: Final = _BillingMetricsBearerSession(auth)
+    exporter: Final = OTLPMetricExporter(endpoint=_metrics_endpoint(config.endpoint), session=session)
     reader: Final = PeriodicExportingMetricReader(exporter, export_interval_millis=config.export_interval_ms)
-    return MeterProvider(metric_readers=[reader], resource=Resource.create(_resource_attributes(config)))
+    return MeterProvider(metric_readers=[reader], resource=Resource.create(resource_attributes(config)))
 
 
 class BillingMetricsRecorder:
@@ -125,7 +200,7 @@ class BillingMetricsRecorder:
         )
 
     def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: str | None) -> None:
-        self._counter.add(1, _billable_attributes(category, route, status_code, model_id))
+        self._counter.add(1, billable_attributes(category, route, status_code, model_id))
 
     def shutdown(self) -> None:
         """Final flush + exporter-thread stop. Without this, up to one export
@@ -146,115 +221,34 @@ def _export_interval_ms() -> int:
         return DEFAULT_EXPORT_INTERVAL_MS
 
 
-@dataclass(frozen=True, slots=True)
-class _CredentialPaths:
-    client_cert_path: str
-    client_key_path: str
-    ca_cert_path: str | None
-
-
-def _is_pem_content(value: str) -> bool:
-    return value.lstrip().startswith(_PEM_PREFIX)
-
-
-def _write_pem(directory: str, filename: str, pem: str) -> str:
-    path: Final = os.path.join(directory, filename)
-    with open(path, "w", encoding="utf-8") as handle:
-        handle.write(pem if pem.endswith("\n") else f"{pem}\n")
-    os.chmod(path, _PEM_FILE_MODE)
-    return path
-
-
-def _resolve_credential_paths(*, client_cert: str, client_key: str, ca_cert: str | None) -> _CredentialPaths:
-    """
-    Accept either a filesystem path or inline PEM content for each credential.
-
-    Secret stores that inject values as environment content rather than mounted
-    files (ECS tasks reading AWS Secrets Manager, Cloud Run reading Secret
-    Manager) can only deliver the certificate as a string. The OTLP exporter
-    takes paths, so inline PEM is written to a private directory once, when the
-    recorder is built. Raises OSError if that write fails; the caller disables
-    metering rather than propagating.
-    """
-    inline: Final = tuple(value for value in (client_cert, client_key, ca_cert) if value and _is_pem_content(value))
-    if not inline:
-        return _CredentialPaths(client_cert, client_key, ca_cert)
-
-    # mkdtemp is 0o700, so the 0o600 key file it holds is unreachable by other users.
-    directory: Final = tempfile.mkdtemp(prefix=_PEM_DIR_PREFIX)
-    return _CredentialPaths(
-        client_cert_path=(
-            _write_pem(directory, _CLIENT_CERT_FILENAME, client_cert) if _is_pem_content(client_cert) else client_cert
-        ),
-        client_key_path=(
-            _write_pem(directory, _CLIENT_KEY_FILENAME, client_key) if _is_pem_content(client_key) else client_key
-        ),
-        ca_cert_path=(
-            _write_pem(directory, _CA_CERT_FILENAME, ca_cert) if ca_cert and _is_pem_content(ca_cert) else ca_cert
-        ),
-    )
-
-
 def load_billing_metrics_config(
-    *, license_data: Optional["EnterpriseLicenseData"], litellm_version: str
+    *, license_data: "EnterpriseLicenseData | None", litellm_version: str
 ) -> BillingMetricsConfig | None:
-    endpoint: Final = os.getenv(ENDPOINT_ENV)
-    client_cert: Final = os.getenv(CLIENT_CERT_ENV)
-    client_key: Final = os.getenv(CLIENT_KEY_ENV)
-    # Optional: only for private/test collectors whose server cert is not on the
-    # public web PKI. The production collector needs no CA override.
-    ca_cert: Final = os.getenv(CA_CERT_ENV)
-
-    missing: Final = [
-        name
-        for name, value in (
-            (ENDPOINT_ENV, endpoint),
-            (CLIENT_CERT_ENV, client_cert),
-            (CLIENT_KEY_ENV, client_key),
-        )
-        if not value
-    ]
-    if not endpoint or not client_cert or not client_key:
+    env: Final = MappingProxyType(
+        {name: os.getenv(name) or "" for name in (ENDPOINT_ENV, CLIENT_ID_ENV, CLIENT_SECRET_ENV, TOKEN_URL_ENV)}
+    )
+    missing: Final = tuple(name for name, value in env.items() if not value)
+    if missing:
         verbose_proxy_logger.warning(
             "Enterprise billing metrics disabled: licensed deployment missing config (%s)",
             ", ".join(missing),
         )
         return None
 
-    try:
-        paths: Final = _resolve_credential_paths(client_cert=client_cert, client_key=client_key, ca_cert=ca_cert)
-    except OSError as exc:
-        verbose_proxy_logger.warning(
-            "Enterprise billing metrics disabled: could not write inline certificate content to disk: %s", exc
-        )
-        return None
-
-    # Report the variable names, never their values. A value that is neither a
-    # readable path nor recognizable PEM is still secret material, and this
-    # warning would otherwise copy a client key straight into the proxy logs.
-    unreadable: Final = [
+    insecure_urls: Final = tuple(
         env_name
-        for env_name, path in (
-            (CLIENT_CERT_ENV, paths.client_cert_path),
-            (CLIENT_KEY_ENV, paths.client_key_path),
-            (CA_CERT_ENV, paths.ca_cert_path),
-        )
-        if path and not os.path.isfile(path)
-    ]
-    if unreadable:
-        verbose_proxy_logger.warning(
-            "Enterprise billing metrics disabled: %s did not resolve to a readable certificate file. "
-            "Set each to a file path, or to inline PEM content beginning with '%s'.",
-            ", ".join(unreadable),
-            _PEM_PREFIX,
-        )
+        for env_name, url in ((ENDPOINT_ENV, env[ENDPOINT_ENV]), (TOKEN_URL_ENV, env[TOKEN_URL_ENV]))
+        if not is_https_url_with_host(url)
+    )
+    if insecure_urls:
+        verbose_proxy_logger.warning("Enterprise billing metrics disabled: %s must use HTTPS", insecure_urls[0])
         return None
 
     return BillingMetricsConfig(
-        endpoint=endpoint,
-        client_cert_path=paths.client_cert_path,
-        client_key_path=paths.client_key_path,
-        ca_cert_path=paths.ca_cert_path,
+        endpoint=env[ENDPOINT_ENV],
+        client_id=env[CLIENT_ID_ENV],
+        client_secret=env[CLIENT_SECRET_ENV],
+        token_url=env[TOKEN_URL_ENV],
         export_interval_ms=_export_interval_ms(),
         litellm_version=litellm_version,
         license_id=(license_data or {}).get("user_id"),
@@ -282,7 +276,7 @@ _ACTIVE_RECORDER: Final = _ActiveRecorderRegistry()
 
 
 def build_billing_metrics_recorder(
-    *, premium: bool, license_data: Optional["EnterpriseLicenseData"], litellm_version: str
+    *, premium: bool, license_data: "EnterpriseLicenseData | None", litellm_version: str
 ) -> BillingMetricsRecorder | None:
     """Build the recorder, or None when the deployment is not licensed or metering is unconfigured."""
     if not premium:
@@ -296,14 +290,12 @@ def build_billing_metrics_recorder(
         return None
 
     try:
-        recorder: Final = BillingMetricsRecorder(build_mtls_meter_provider(config))
+        provider: Final = build_bearer_meter_provider(config)
+        recorder: Final = BillingMetricsRecorder(provider)
     except Exception as exc:  # noqa: BLE001 -- metering must never break proxy startup
         verbose_proxy_logger.warning("Enterprise billing metrics disabled: failed to initialize exporter: %s", exc)
         return None
     _ACTIVE_RECORDER.set(recorder)
-    # The only positive signal that this component meters. Without it, a silent
-    # return above is indistinguishable from a working exporter in the logs, and
-    # a component that carries the cert but no license would look healthy.
     verbose_proxy_logger.info(
         "Enterprise billing metrics enabled: exporting to %s every %d ms",
         config.endpoint,

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -31,7 +31,7 @@ from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm import get_llm_provider
-from litellm._logging import verbose_proxy_logger
+from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.constants import (
     ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
@@ -58,6 +58,7 @@ from litellm.llms.anthropic.common_utils import (
     anthropic_pass_through_api_base,
     merge_anthropic_beta_headers,
 )
+from litellm.llms.azure.common_utils import entra_rejection_in
 from litellm.llms.azure.passthrough.transformation import (
     foreign_azure_deployment,
     is_azure_body_model_inference_endpoint,
@@ -2381,6 +2382,16 @@ async def _relay_router_model(
         )
     except httpx.HTTPStatusError as upstream_error:
         return await _relay_upstream_response(upstream_error.response)
+    except Exception as error:
+        rejection: Final = entra_rejection_in(error)
+        if rejection is None:
+            raise
+        raise ProxyException(
+            message=redact_internal_details_from_client_message(rejection.message),
+            type=ProxyErrorTypes.auth_error,
+            param=None,
+            code=rejection.status_code,
+        ) from error
 
     if not is_streaming_request:
         return await _relay_upstream_response(cast(httpx.Response, result))

litellm/router_api.py

@@ -1402,6 +1402,9 @@ class RoutingContext:
     def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]:
         return self._router._params_sent_to_deployment(request, deployment)
 
+    def default_litellm_params(self) -> Mapping[str, object]:
+        return self._router._active_default_litellm_params()
+
     def settings(self) -> RouterSettings:
         return self._router.settings()
 

litellm/router_deprecated.py

@@ -27,13 +27,14 @@ import inspect
 import sys
 import warnings
 from collections import Counter
-from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload
+from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload
 
 from pydantic import TypeAdapter
 
 from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet
+from litellm.router_strategy.least_busy import picks_reserve
 
 if TYPE_CHECKING:
     from typing import Any  # noqa: TID251  # only to repeat two old signatures that had Any, unchanged
@@ -60,6 +61,7 @@ if TYPE_CHECKING:
     from litellm.types.utils import ModelInfo as ModelMapInfo
 
 _T: Final = TypeVar("_T")
+_P: Final = ParamSpec("_P")
 
 DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter()  # mutable-ok: the live tally _note increments
 
@@ -168,6 +170,18 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib
     return _DeprecatedAttribute(replacement, read)
 
 
+def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coroutine[object, object, _T]]:
+    @functools.wraps(lookup)
+    async def looked_up(
+        *args: _P.args,
+        **kwargs: _P.kwargs,  # kwargs-ok: forwards the wrapped lookup's ParamSpec unchanged
+    ) -> _T:
+        with picks_reserve(False):
+            return await lookup(*args, **kwargs)
+
+    return looked_up
+
+
 # (id(router), member name) pairs whose subclass override is running in this context.
 _OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar(
     "router_overrides_running", default=frozenset()
@@ -334,7 +348,7 @@ class DeprecatedRouterMembers:
     deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map)
     async_get_available_deployment = _attr(
         "select(model, request_kwargs)",
-        lambda r: r._core_async_get_available_deployment,
+        lambda r: _without_reserving(r._core_async_get_available_deployment),
     )
 
     # -- observers ------------------------------------------------------------------------------

litellm/router_python_core.py

@@ -297,7 +297,15 @@ from litellm.router_strategy.complexity_router.context_compaction import (
     reject_recursive_compactor,
     surface_for_call,
 )
-from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
+from litellm.router_strategy.least_busy import (
+    CALL_IDENTITY,
+    SELECTION_RESERVATION_MAX_KEYS,
+    LeastBusyLoggingHandler,
+    call_id_of,
+    owning_picks,
+    pick_reserves,
+    picks_reserve,
+)
 from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
 from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
 from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler
@@ -321,6 +329,7 @@ from litellm.router_utils.cooldown_handlers import (
     get_cooldown_deployment_info,
     get_cooldown_deployments,
     is_background_response_poll_not_found,
+    is_caller_entra_credential_failure,
     is_caller_scoped_auth_failure,
     is_caller_timeout_408,
     is_failure_outside_selected_deployment,
@@ -1497,13 +1506,14 @@ class PythonCore(DeprecatedRouterMembers):
         (``pass_through``).
         """
         if kind == "pick":
-            return await self._core_async_get_available_deployment(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the selection cores return untyped deployment dicts
-                model=model,
-                request_kwargs=request_kwargs if request_kwargs is not None else {},
-                messages=messages,
-                input=input,
-                specific_deployment=specific_deployment,
-            )
+            with picks_reserve(False):
+                return await self._core_async_get_available_deployment(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the selection cores return untyped deployment dicts
+                    model=model,
+                    request_kwargs=request_kwargs if request_kwargs is not None else {},
+                    messages=messages,
+                    input=input,
+                    specific_deployment=specific_deployment,
+                )
         if kind == "healthy":
             return await self._core_async_get_healthy_deployments(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the selection cores return untyped deployment dicts
                 model=model,
@@ -1515,13 +1525,14 @@ class PythonCore(DeprecatedRouterMembers):
                 health_check_probe=health_check_probe,
             )
         if kind == "pass_through":
-            return self._core_get_available_deployment_for_pass_through(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the selection cores return untyped deployment dicts
-                model=model,
-                messages=messages,
-                input=input,
-                specific_deployment=specific_deployment,
-                request_kwargs=request_kwargs,
-            )
+            with picks_reserve(False):
+                return self._core_get_available_deployment_for_pass_through(  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the selection cores return untyped deployment dicts
+                    model=model,
+                    messages=messages,
+                    input=input,
+                    specific_deployment=specific_deployment,
+                    request_kwargs=request_kwargs,
+                )
         raise ValueError(f"unknown select kind {kind!r}")
 
     def apply(self, changes: ChangeSet) -> ApplyResult:
@@ -1803,6 +1814,9 @@ class PythonCore(DeprecatedRouterMembers):
         self._core_cache = DualCache(
             redis_cache=redis_cache, in_memory_cache=InMemoryCache()
         )  # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc.
+        self._core_least_busy_reservations: InMemoryCache = InMemoryCache(
+            max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS
+        )
         self._claude_code_session_router_cache: DualCache = DualCache(
             redis_cache=redis_cache,
             in_memory_cache=InMemoryCache(),
@@ -2306,7 +2320,9 @@ class PythonCore(DeprecatedRouterMembers):
         selector: RouterStrategySelector | None = None
         match self._normalize_strategy(strategy):
             case RoutingStrategy.LEAST_BUSY.value:
-                selector = LeastBusyLoggingHandler(router_cache=self._core_cache)
+                selector = LeastBusyLoggingHandler(
+                    router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations
+                )
             case RoutingStrategy.USAGE_BASED_ROUTING.value:
                 selector = LowestTPMLoggingHandler(
                     router_cache=self._core_cache,
@@ -2771,6 +2787,7 @@ class PythonCore(DeprecatedRouterMembers):
         messages: list[dict[str, str]] | None,
         input: str | list | None,
         request_kwargs: dict | None,
+        reserve: bool,
     ) -> Any | None:
         """
         Asks the strategy selector for a deployment. Caller handles
@@ -2786,6 +2803,8 @@ class PythonCore(DeprecatedRouterMembers):
                 ).async_get_available_deployments(
                     model_group=model,
                     healthy_deployments=healthy_deployments,
+                    reserve=reserve,
+                    call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})),
                 )
             case "usage-based-routing":
                 # `LowestTPMLoggingHandler` (v1) only exposes the sync
@@ -2832,6 +2851,7 @@ class PythonCore(DeprecatedRouterMembers):
         messages: list[dict[str, str]] | None,
         input: str | list | None,
         request_kwargs: dict | None,
+        reserve: bool,
     ) -> Any | None:
         """
         Sync sibling of `_select_deployment_async`. Caller handles
@@ -2850,6 +2870,8 @@ class PythonCore(DeprecatedRouterMembers):
                 ).get_available_deployments(
                     model_group=model,
                     healthy_deployments=healthy_deployments,
+                    reserve=reserve,
+                    call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})),
                 )
             case "usage-based-routing" | "usage-based-routing-v2":
                 return require_attributes(
@@ -3573,12 +3595,13 @@ class PythonCore(DeprecatedRouterMembers):
             )
 
             # pick the one that is available (lowest TPM/RPM)
-            deployment = self.get_available_deployment(
-                model=model,
-                messages=messages,
-                specific_deployment=kwargs.pop("specific_deployment", None),
-                request_kwargs=kwargs,
-            )
+            with picks_reserve(True):
+                deployment = self.get_available_deployment(
+                    model=model,
+                    messages=messages,
+                    specific_deployment=kwargs.pop("specific_deployment", None),
+                    request_kwargs=kwargs,
+                )
             self._drop_unsupported_classifier_reasoning_effort(
                 deployment=cast(DeploymentTypedDict, deployment),  # cast-ok: selection returns a router deployment
                 model=model,
@@ -6154,12 +6177,13 @@ class PythonCore(DeprecatedRouterMembers):
         deployment = None  # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
         try:
             verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
-            deployment = self.get_available_deployment(  # rebind-ok: set on success, see pre-init
-                model=model,
-                messages=[{"role": "user", "content": "prompt"}],
-                specific_deployment=kwargs.pop("specific_deployment", None),
-                request_kwargs=kwargs,
-            )
+            with picks_reserve(True):
+                deployment = self.get_available_deployment(  # rebind-ok: set on success, see pre-init
+                    model=model,
+                    messages=[{"role": "user", "content": "prompt"}],
+                    specific_deployment=kwargs.pop("specific_deployment", None),
+                    request_kwargs=kwargs,
+                )
             self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
             data: Final = deployment["litellm_params"].copy()
 
@@ -6546,37 +6570,38 @@ class PythonCore(DeprecatedRouterMembers):
         **kwargs,
     ):
         messages: Final = [{"role": "user", "content": prompt}]
-        try:
-            kwargs["model"] = model
-            kwargs["prompt"] = prompt
-            kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
-            kwargs.setdefault("metadata", {}).update({"model_group": model})
+        with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
+            try:
+                kwargs["model"] = model
+                kwargs["prompt"] = prompt
+                kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
+                kwargs.setdefault("metadata", {}).update({"model_group": model})
 
-            # pick the one that is available (lowest TPM/RPM)
-            with endpoint_selection("text_completion", model=model):
-                deployment: Final = self.get_available_deployment(
-                    model=model,
-                    messages=messages,
-                    specific_deployment=kwargs.pop("specific_deployment", None),
-                    request_kwargs=kwargs,
+                # pick the one that is available (lowest TPM/RPM)
+                with endpoint_selection("text_completion", model=model), picks_reserve(True):
+                    deployment: Final = self.get_available_deployment(
+                        model=model,
+                        messages=messages,
+                        specific_deployment=kwargs.pop("specific_deployment", None),
+                        request_kwargs=kwargs,
+                    )
+
+                data: Final = deployment["litellm_params"].copy()
+                for k, v in self._core_default_litellm_params.items():
+                    if k not in kwargs:  # prioritize model-specific params > default router params
+                        kwargs[k] = v
+                    elif k == "metadata":
+                        kwargs[k].update(v)
+                pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data))
+                if pinned_model_id is not None:
+                    kwargs["model_id"] = pinned_model_id  # rebind-ok: this call's own kwargs, filled like the defaults
+
+                # call via litellm.completion()
+                return self._host.invoke("text_completion")(
+                    **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs}
                 )
-
-            data: Final = deployment["litellm_params"].copy()
-            for k, v in self._core_default_litellm_params.items():
-                if k not in kwargs:  # prioritize model-specific params > default router params
-                    kwargs[k] = v
-                elif k == "metadata":
-                    kwargs[k].update(v)
-            pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data))
-            if pinned_model_id is not None:
-                kwargs["model_id"] = pinned_model_id  # rebind-ok: this call's own kwargs, filled like the defaults
-
-            # call via litellm.completion()
-            return self._host.invoke("text_completion")(
-                **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs}
-            )
-        except Exception as e:
-            raise e
+            except Exception as e:
+                raise e
 
     async def atext_completion(
         self,
@@ -7719,64 +7744,69 @@ class PythonCore(DeprecatedRouterMembers):
         handler_name: Final = original_function.__name__
         metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call")
         deployment = None  # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
-        try:
-            verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model)
-            self._update_kwargs_before_fallbacks(
-                model=model,
-                kwargs=kwargs,
-                metadata_variable_name=metadata_variable_name,
-            )
-            with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name):
-                deployment = self.get_available_deployment(  # rebind-ok: set on success, see pre-init
+        with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
+            try:
+                verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model)
+                self._update_kwargs_before_fallbacks(
                     model=model,
-                    messages=kwargs.get("messages", None),
-                    input=kwargs.get("input", None),
-                    specific_deployment=kwargs.pop("specific_deployment", None),
-                    request_kwargs=kwargs,
+                    kwargs=kwargs,
+                    metadata_variable_name=metadata_variable_name,
+                )
+                with (
+                    self._video_call_pin(handler_name, kwargs),
+                    self._route_capability_pin(handler_name),
+                    picks_reserve(True),
+                ):
+                    deployment = self.get_available_deployment(  # rebind-ok: set on success, see pre-init
+                        model=model,
+                        messages=kwargs.get("messages", None),
+                        input=kwargs.get("input", None),
+                        specific_deployment=kwargs.pop("specific_deployment", None),
+                        request_kwargs=kwargs,
+                    )
+                self._update_kwargs_with_deployment(
+                    deployment=deployment,
+                    kwargs=kwargs,
+                    function_name="generic_api_call",
+                    splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
                 )
-            self._update_kwargs_with_deployment(
-                deployment=deployment,
-                kwargs=kwargs,
-                function_name="generic_api_call",
-                splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
-            )
 
-            data: Final = deployment["litellm_params"].copy()
-            model_name: Final = data["model"]
+                data: Final = deployment["litellm_params"].copy()
+                model_name: Final = data["model"]
 
-            self.total_calls[model_name] += 1
+                self.total_calls[model_name] += 1
 
-            # For passthrough routes, use the actual model from deployment
-            # and swap model name in endpoint if present
-            if "endpoint" in kwargs and kwargs["endpoint"]:
-                kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name)
-            kwargs["model"] = model_name
+                # For passthrough routes, use the actual model from deployment
+                # and swap model name in endpoint if present
+                if "endpoint" in kwargs and kwargs["endpoint"]:
+                    kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name)
+                kwargs["model"] = model_name
 
-            # Perform pre-call checks for routing strategy
-            self.routing_strategy_pre_call_checks(deployment=deployment)
+                # Perform pre-call checks for routing strategy
+                self.routing_strategy_pre_call_checks(deployment=deployment)
 
-            custom_llm_provider: Final = provider_for_generic_call(data)
+                custom_llm_provider: Final = provider_for_generic_call(data)
 
-            response: Final = original_function(
-                **{
-                    **data,
-                    "custom_llm_provider": custom_llm_provider,
-                    "caching": self._core_cache_responses,
-                    **kwargs,
-                }
-            )
+                response: Final = original_function(
+                    **{
+                        **data,
+                        "custom_llm_provider": custom_llm_provider,
+                        "caching": self._core_cache_responses,
+                        **kwargs,
+                    }
+                )
 
-            self.success_calls[model_name] += 1
-            verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name)
-            return response
-        except Exception as e:
-            verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e)
-            if model is not None:
-                self.fail_calls[model] += 1
-            self._stamp_failed_attempt(
-                e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
-            )
-            raise e
+                self.success_calls[model_name] += 1
+                verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name)
+                return response
+            except Exception as e:
+                verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e)
+                if model is not None:
+                    self.fail_calls[model] += 1
+                self._stamp_failed_attempt(
+                    e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
+                )
+                raise e
 
     def embedding(
         self,
@@ -7800,7 +7830,7 @@ class PythonCore(DeprecatedRouterMembers):
         deployment = None  # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
         try:
             verbose_router_logger.debug("Inside embedding()- model: %s", model)
-            with embedding_deployment_selection():
+            with embedding_deployment_selection(), picks_reserve(True):
                 deployment = self.get_available_deployment(  # rebind-ok: set on success, see pre-init
                     model=model,
                     input=input,
@@ -8109,68 +8139,71 @@ class PythonCore(DeprecatedRouterMembers):
         Returns:
             VectorStoreCreateResponse
         """
-        try:
-            # If model is None, use the factory function approach (direct SDK call)
-            if model is None:
-                from litellm.vector_stores.main import acreate
+        with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
+            try:
+                # If model is None, use the factory function approach (direct SDK call)
+                if model is None:
+                    from litellm.vector_stores.main import acreate
 
-                # Use the factory function to handle the call
-                factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create")
-                return await factory_fn(**kwargs)
+                    # Use the factory function to handle the call
+                    factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create")
+                    return await factory_fn(**kwargs)
 
-            from litellm.vector_stores import acreate as avector_store_create_sdk
+                from litellm.vector_stores import acreate as avector_store_create_sdk
 
-            parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
-            deployment: Final = await self._core_async_get_available_deployment(
-                model=model,
-                messages=[{"role": "user", "content": "vector-store-api-fake-text"}],
-                specific_deployment=kwargs.pop("specific_deployment", None),
-                request_kwargs=kwargs,
-            )
-            data: Final = deployment["litellm_params"].copy()
-            model_name: Final = data["model"]
-            self._update_kwargs_with_deployment(
-                deployment=deployment,
-                kwargs=kwargs,
-                function_name="avector_store_create",
-            )
-
-            model_client: Final = self._get_async_openai_model_client(
-                deployment=deployment,
-                kwargs=kwargs,
-            )
-            self.total_calls[model_name] += 1
-
-            # Get custom provider from deployment params
-            custom_llm_provider = data.get("custom_llm_provider")
-            _, inferred_custom_llm_provider, _, _ = get_llm_provider(
-                model=data["model"],
-                custom_llm_provider=custom_llm_provider,
-            )
-            custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
-
-            async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
-                response = await avector_store_create_sdk(
-                    **{
-                        **data,
-                        "custom_llm_provider": custom_llm_provider,
-                        "caching": self._core_cache_responses,
-                        "client": model_client,
-                        **kwargs,
-                    }
+                parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
+                deployment: Final = await self._core_async_get_available_deployment(
+                    model=model,
+                    messages=[{"role": "user", "content": "vector-store-api-fake-text"}],
+                    specific_deployment=kwargs.pop("specific_deployment", None),
+                    request_kwargs=kwargs,
+                )
+                data: Final = deployment["litellm_params"].copy()
+                model_name: Final = data["model"]
+                self._update_kwargs_with_deployment(
+                    deployment=deployment,
+                    kwargs=kwargs,
+                    function_name="avector_store_create",
                 )
 
-            self.success_calls[model_name] += 1
-            verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
+                model_client: Final = self._get_async_openai_model_client(
+                    deployment=deployment,
+                    kwargs=kwargs,
+                )
+                self.total_calls[model_name] += 1
 
-            return response
-        except Exception as e:
-            verbose_router_logger.exception(
-                "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e
-            )
-            if model is not None:
-                self.fail_calls[model] += 1
-            raise e
+                # Get custom provider from deployment params
+                custom_llm_provider = data.get("custom_llm_provider")
+                _, inferred_custom_llm_provider, _, _ = get_llm_provider(
+                    model=data["model"],
+                    custom_llm_provider=custom_llm_provider,
+                )
+                custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
+
+                async with self._deployment_slot(
+                    deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span
+                ):
+                    response = await avector_store_create_sdk(
+                        **{
+                            **data,
+                            "custom_llm_provider": custom_llm_provider,
+                            "caching": self._core_cache_responses,
+                            "client": model_client,
+                            **kwargs,
+                        }
+                    )
+
+                self.success_calls[model_name] += 1
+                verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
+
+                return response
+            except Exception as e:
+                verbose_router_logger.exception(
+                    "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e
+                )
+                if model is not None:
+                    self.fail_calls[model] += 1
+                raise e
 
     def _override_vector_store_methods_for_router(self):
         """
@@ -8418,12 +8451,13 @@ class PythonCore(DeprecatedRouterMembers):
         try:
             verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model)
             parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
-            deployment: Final = await self._core_async_get_available_deployment(
-                model=model,
-                messages=[{"role": "user", "content": "batch-api-fake-text"}],
-                specific_deployment=kwargs.pop("specific_deployment", None),
-                request_kwargs=kwargs,
-            )
+            with picks_reserve(False):
+                deployment: Final = await self._core_async_get_available_deployment(
+                    model=model,
+                    messages=[{"role": "user", "content": "batch-api-fake-text"}],
+                    specific_deployment=kwargs.pop("specific_deployment", None),
+                    request_kwargs=kwargs,
+                )
 
             selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id")
             data: Final = deployment["litellm_params"].copy()
@@ -9578,38 +9612,39 @@ class PythonCore(DeprecatedRouterMembers):
 
         mock_timeout: Final = kwargs.pop("mock_timeout", None)
 
-        try:
-            self._handle_mock_testing_fallbacks(
-                kwargs=kwargs,
-                model_group=model_group,
-                fallbacks=fallbacks,
-                context_window_fallbacks=context_window_fallbacks,
-                content_policy_fallbacks=content_policy_fallbacks,
-            )
+        with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
+            try:
+                self._handle_mock_testing_fallbacks(
+                    kwargs=kwargs,
+                    model_group=model_group,
+                    fallbacks=fallbacks,
+                    context_window_fallbacks=context_window_fallbacks,
+                    content_policy_fallbacks=content_policy_fallbacks,
+                )
 
-            if mock_timeout is not None:
-                response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout)
-            else:
-                response = await self.async_function_with_retries(*args, **kwargs)
-            if verbose_router_logger.isEnabledFor(logging.DEBUG):
-                verbose_router_logger.debug("Async Response: %s", response)
-            response = add_fallback_headers_to_response(
-                response=response,
-                attempted_fallbacks=0,
-            )
-            return response
-        except Exception as e:
-            return await self.async_function_with_fallbacks_common_utils(
-                e,
-                disable_fallbacks,
-                fallbacks,
-                context_window_fallbacks,
-                content_policy_fallbacks,
-                model_group,
-                args,
-                kwargs,
-                include_fallback_errors=include_fallback_errors,
-            )
+                if mock_timeout is not None:
+                    response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout)
+                else:
+                    response = await self.async_function_with_retries(*args, **kwargs)
+                if verbose_router_logger.isEnabledFor(logging.DEBUG):
+                    verbose_router_logger.debug("Async Response: %s", response)
+                response = add_fallback_headers_to_response(
+                    response=response,
+                    attempted_fallbacks=0,
+                )
+                return response
+            except Exception as e:
+                return await self.async_function_with_fallbacks_common_utils(
+                    e,
+                    disable_fallbacks,
+                    fallbacks,
+                    context_window_fallbacks,
+                    content_policy_fallbacks,
+                    model_group,
+                    args,
+                    kwargs,
+                    include_fallback_errors=include_fallback_errors,
+                )
 
     def _handle_mock_testing_fallbacks(
         self,
@@ -10373,8 +10408,9 @@ class PythonCore(DeprecatedRouterMembers):
                 return False
 
             exception_status: Final = getattr(exception, "status_code", "")
+            failure_call_details: Final = _CALL_KWARGS_ADAPTER.validate_python(kwargs)
 
-            if is_caller_timeout_408(kwargs, exception_status):
+            if is_caller_timeout_408(failure_call_details, exception_status):
                 verbose_router_logger.debug(
                     "Router: Exiting 'deployment_callback_on_failure' without cooldown. "
                     "A timeout the caller set caused this 408, not the deployment's health."
@@ -10398,6 +10434,18 @@ class PythonCore(DeprecatedRouterMembers):
                     "Caller-scoped OAuth authentication failed, not the deployment."
                 )
                 return False
+            if is_caller_entra_credential_failure(
+                caller_failure_deployment,
+                self._active_default_litellm_params(),
+                failure_call_details,
+                exception_status,
+                exception,
+            ):
+                verbose_router_logger.debug(
+                    "Router: Exiting 'deployment_callback_on_failure' without cooldown. "
+                    "Entra ID rejected a credential the caller sent, not the deployment's own."
+                )
+                return False
 
             exception_headers: Final = exception_mapping_utils._get_response_headers(original_exception=exception)
 
@@ -10906,6 +10954,7 @@ class PythonCore(DeprecatedRouterMembers):
             messages=None,
             input=None,
             request_kwargs=routing_kwargs,
+            reserve=False,
         )
         if selected is None:
             no_deployment: Final = await async_raise_no_deployment_exception(
@@ -16264,6 +16313,7 @@ class PythonCore(DeprecatedRouterMembers):
                 messages=messages,
                 input=input,
                 specific_deployment=specific_deployment,
+                reserve=pick_reserves(default=True),
             )
 
     async def _async_get_available_deployment(
@@ -16273,6 +16323,7 @@ class PythonCore(DeprecatedRouterMembers):
         messages: list[dict[str, str]] | None,
         input: str | list | None,
         specific_deployment: bool | None,
+        reserve: bool,
     ):
         if (
             self._core_routing_strategy != "usage-based-routing-v2"
@@ -16369,6 +16420,7 @@ class PythonCore(DeprecatedRouterMembers):
                     messages=messages,
                     input=input,
                     request_kwargs=request_kwargs,
+                    reserve=reserve and request_kwargs not in (None, {}),
                 )
             if deployment is None:
                 exception: Final = await async_raise_no_deployment_exception(
@@ -16538,6 +16590,7 @@ class PythonCore(DeprecatedRouterMembers):
                 messages=messages,
                 input=input,
                 request_kwargs=request_kwargs,
+                reserve=pick_reserves(default=False) and request_kwargs not in (None, {}),
             )
 
             if deployment is None:
@@ -17415,6 +17468,7 @@ class PythonCore(DeprecatedRouterMembers):
             messages=messages,
             input=input,
             request_kwargs=request_kwargs,
+            reserve=pick_reserves(default=False) and request_kwargs not in (None, {}),
         )
 
         if deployment is None:
@@ -17595,6 +17649,7 @@ class PythonCore(DeprecatedRouterMembers):
             messages=messages,
             input=input,
             request_kwargs=request_kwargs,
+            reserve=pick_reserves(default=False) and request_kwargs not in (None, {}),
         )
 
         if deployment is None:

litellm/router_strategy/least_busy.py

@@ -1,6 +1,11 @@
+import contextlib
 import logging
-from collections.abc import Mapping, Sequence
-from typing import Final
+import threading
+import time
+import uuid
+from collections.abc import Callable, Generator, Mapping, Sequence
+from contextvars import ContextVar
+from typing import Final, Protocol, TypeAlias
 
 from pydantic import TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict
@@ -8,11 +13,31 @@ from typing_extensions import ReadOnly, TypedDict
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.caching.redis_cache import log_redis_failure
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
+from litellm.types.utils import CallTypes
 
 IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
+SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0
+SELECTION_RESERVATION_MAX_KEYS: Final = 10_000
+_RESERVATIONS_LOCK: Final = threading.Lock()
+_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves")
+_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None)
+
+
+@contextlib.contextmanager
+def picks_reserve(reserve: bool) -> Generator[None, None, None]:
+    token: Final = _PICK_RESERVES.set(reserve)
+    try:
+        yield
+    finally:
+        _PICK_RESERVES.reset(token)
+
+
+def pick_reserves(default: bool) -> bool:
+    return _PICK_RESERVES.get(default)
 
 
 class _ModelInfo(TypedDict, total=False):
@@ -21,10 +46,12 @@ class _ModelInfo(TypedDict, total=False):
 
 class _Metadata(TypedDict, total=False):
     model_group: ReadOnly[str | None]
+    model_info: ReadOnly[_ModelInfo | None]
 
 
 class _LitellmParams(TypedDict, total=False):
     metadata: ReadOnly[_Metadata | None]
+    litellm_metadata: ReadOnly[_Metadata | None]
     model_info: ReadOnly[_ModelInfo | None]
 
 
@@ -32,6 +59,10 @@ class _CallKwargs(TypedDict, total=False):
     litellm_params: ReadOnly[_LitellmParams | None]
 
 
+class CallIdentity(TypedDict, total=False):
+    litellm_call_id: ReadOnly[object]
+
+
 class _DeploymentModelInfo(TypedDict):
     id: ReadOnly[str | int]
 
@@ -41,38 +72,85 @@ class _Deployment(TypedDict):
 
 
 _CALL_KWARGS: Final = TypeAdapter(_CallKwargs)
+CALL_IDENTITY: Final = TypeAdapter(CallIdentity)
 _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
-_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None)
+_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None)
+_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None)
+_Pick: TypeAlias = tuple[float, str | None]
+_PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None)
+_PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None)
 
 
 def _request_count_key(model_group: str, deployment_id: str) -> str:
     return f"{model_group}_request_count:{deployment_id}"
 
 
+def _model_group(metadata: _Metadata | None) -> str | None:
+    return metadata.get("model_group") if metadata else None
+
+
+def _deployment_id(model_info: _ModelInfo | None) -> str | int | None:
+    return model_info.get("id") if model_info else None
+
+
+def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None:
+    return _deployment_id(metadata.get("model_info")) if metadata else None
+
+
 def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
-    if is_batch_retrieve_call_type(kwargs.get("call_type")):
+    call_type: Final = kwargs.get("call_type")
+    if is_batch_retrieve_call_type(call_type) or call_type == CallTypes.arealtime.value:
         return None
     try:
         call: Final = _CALL_KWARGS.validate_python(kwargs)
     except ValidationError:
         return None
-    litellm_params: Final = call.get("litellm_params")
-    metadata: Final = litellm_params.get("metadata") if litellm_params else None
-    model_info: Final = litellm_params.get("model_info") if litellm_params else None
-    model_group: Final = metadata.get("model_group") if metadata else None
-    deployment_id: Final = model_info.get("id") if model_info else None
+    litellm_params: Final = call.get("litellm_params") or _LitellmParams()
+    routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata"))
+    model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None)
+    deployment_id: Final = next(
+        (
+            deployment_id
+            for deployment_id in (
+                _deployment_id(litellm_params.get("model_info")),
+                *map(_routed_deployment_id, routing_metadata),
+            )
+            if deployment_id is not None
+        ),
+        None,
+    )
     if model_group is None or deployment_id is None:
         return None
     return model_group, str(deployment_id)
 
 
-def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
+def _reservation_key(deployment_id: str) -> str:
+    return f"least_busy_reservations:{deployment_id}"
+
+
+def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str:
+    return f"least_busy_started:{call_id}:{_request_count_key(*ref)}"
+
+
+def _picked_deployments_key(call_id: str) -> str:
+    return f"least_busy_picked:{call_id}"
+
+
+def call_id_of(kwargs: Mapping[str, object]) -> str | None:
+    call_id: Final = kwargs.get("litellm_call_id")
+    return call_id if isinstance(call_id, str) else None
+
+
+def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
     return tuple(
-        _request_count_key(model_group, str(deployment["model_info"]["id"]))
-        for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
+        str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
     )
 
 
+def _request_count_keys(model_group: str, deployment_ids: tuple[str, ...]) -> tuple[str, ...]:
+    return tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids)
+
+
 def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]:
     return tuple(0 if value is None else int(value) for value in values)
 
@@ -84,12 +162,21 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]:
     return _as_counts(values)
 
 
-def _least_busy(
-    healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...]
-) -> Mapping[str, object] | None:
-    if not healthy_deployments:
-        return None
-    return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])]
+def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]:
+    return tuple(pick for pick in picks if now - pick[0] < SELECTION_RESERVATION_TTL_SECONDS)
+
+
+def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]:
+    tags: Final = tuple(tag for _, tag in picks)
+    owner: Final = next((tag for tag in (call_id, _PICK_OWNER.get()) if tag is not None and tag in tags), None)
+    if owner not in tags:
+        return picks
+    index: Final = tags.index(owner)
+    return (*picks[:index], *picks[index + 1 :])
+
+
+def _least_busy_index(counts: tuple[int, ...]) -> int:
+    return min(range(len(counts)), key=lambda index: counts[index])
 
 
 def _warn_unreadable(model_group: str, error: Exception) -> None:
@@ -111,24 +198,156 @@ def _warn_unwritable(key: str, error: Exception) -> None:
     )
 
 
+class _ReservationStore(Protocol):
+    def get_cache(self, key: str) -> object: ...
+
+    def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: ...
+
+    def delete_cache(self, key: str) -> None: ...
+
+
+def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]:
+    return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or ()
+
+
+def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]:
+    return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or ()
+
+
+def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]:
+    return _PICKED_DEPLOYMENTS.validate_python(store.get_cache(_picked_deployments_key(call_id))) or ()
+
+
+def _put(store: _ReservationStore, key: str, values: tuple[object, ...], ttl: float) -> None:
+    # InMemoryCache keeps a live key's first TTL, so drop the key to restart it
+    store.delete_cache(key)
+    if values:
+        store.set_cache(key, values, ttl=ttl)
+
+
+def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str) -> None:
+    kept: Final = tuple(pick for pick in _picks(store, deployment_id) if pick[1] != call_id)
+    _put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS)
+
+
+def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None:
+    if call_id is None:
+        return
+    if not _picked_deployments(store, call_id):
+        return
+    with _RESERVATIONS_LOCK:
+        for deployment_id in frozenset(_picked_deployments(store, call_id)):
+            _release_call_on(store, deployment_id, call_id)
+        store.delete_cache(_picked_deployments_key(call_id))
+
+
+@contextlib.contextmanager
+def owning_picks(store: _ReservationStore, call_id: str | None) -> Generator[None, None, None]:
+    owner: Final = call_id or _PICK_OWNER.get() or uuid.uuid4().hex
+    token: Final = _PICK_OWNER.set(owner)
+    try:
+        yield
+    finally:
+        _PICK_OWNER.reset(token)
+        release_unstarted_picks(store, owner)
+
+
 class LeastBusyLoggingHandler(CustomLogger):
     test_flag: bool = False
     logged_success: int = 0
     logged_failure: int = 0
 
-    def __init__(self, router_cache: DualCache):
+    def __init__(
+        self,
+        router_cache: DualCache,
+        clock: Callable[[], float] = time.monotonic,
+        reservation_store: _ReservationStore | None = None,
+    ):
         self.router_cache = router_cache
         self.router_cache_id = str(id(router_cache))
+        self._clock = clock
+        self._reservations: Final[_ReservationStore] = (
+            reservation_store
+            if reservation_store is not None
+            else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock)
+        )
+
+    def _live_reservations(self, deployment_id: str, now: float) -> tuple[_Pick, ...]:
+        return _unexpired(_picks(self._reservations, deployment_id), now)
+
+    def _store_reservations(self, deployment_id: str, picks: tuple[_Pick, ...]) -> None:
+        _put(self._reservations, _reservation_key(deployment_id), picks, SELECTION_RESERVATION_TTL_SECONDS)
+
+    def _take_reservation(self, deployment_id: str, call_id: str | None, now: float) -> None:
+        self._store_reservations(deployment_id, _without_one_pick(self._live_reservations(deployment_id, now), call_id))
+
+    def _record_pick(self, call_id: str, deployment_id: str) -> None:
+        picked: Final = (*_picked_deployments(self._reservations, call_id), deployment_id)
+        _put(self._reservations, _picked_deployments_key(call_id), picked, SELECTION_RESERVATION_TTL_SECONDS)
+
+    def _select(
+        self,
+        healthy_deployments: Sequence[Mapping[str, object]],
+        deployment_ids: tuple[str, ...],
+        keys: tuple[str, ...],
+        shared: tuple[int, ...] | None,
+        reserve: bool,
+        call_id: str | None,
+    ) -> Mapping[str, object] | None:
+        if not healthy_deployments:
+            return None
+        with _RESERVATIONS_LOCK:
+            now: Final = self._clock()
+            counts: Final = (
+                self._this_workers_counts(keys) if shared is None else self._at_least_this_workers(keys, shared)
+            )
+            live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids)
+            index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True)))
+            if reserve:
+                tag: Final = call_id if call_id is not None else _PICK_OWNER.get()
+                self._store_reservations(deployment_ids[index], (*live[index], (now, tag)))
+                if tag is not None:
+                    self._record_pick(tag, deployment_ids[index])
+        return healthy_deployments[index]
+
+    def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None:
+        with _RESERVATIONS_LOCK:
+            now: Final = self._clock()
+            self._take_reservation(ref[1], call_id, now)
+            self._count_locally(_request_count_key(*ref), 1)
+            if call_id is not None:
+                key: Final = _started_attempts_key(call_id, ref)
+                started: Final = (*_attempt_times(self._reservations, key), now)
+                _put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS)
+
+    def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool:
+        if call_id is None:
+            with _RESERVATIONS_LOCK:
+                self._count_locally(_request_count_key(*ref), -1)
+            return True
+        key: Final = _started_attempts_key(call_id, ref)
+        with _RESERVATIONS_LOCK:
+            started: Final = _attempt_times(self._reservations, key)
+            if started:
+                _put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS)
+                self._count_locally(_request_count_key(*ref), -1)
+            else:
+                self._take_reservation(ref[1], call_id, self._clock())
+        return bool(started)
 
     @with_service_target("router_usage")
     def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
-        self._increment(kwargs, 1)
+        ref: Final = _deployment_ref(kwargs)
+        if ref is None:
+            return
+        self._start_attempt(ref, call_id_of(kwargs))
+        self._count_shared(_request_count_key(*ref), 1)
 
     @with_service_target("router_usage")
     def log_success_event(
         self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
     ) -> None:
-        self._increment(kwargs, -1)
+        self._settle(kwargs)
         if self.test_flag:
             self.logged_success += 1
 
@@ -136,7 +355,7 @@ class LeastBusyLoggingHandler(CustomLogger):
     def log_failure_event(
         self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
     ) -> None:
-        self._increment(kwargs, -1)
+        self._settle(kwargs)
         if self.test_flag:
             self.logged_failure += 1
 
@@ -144,7 +363,7 @@ class LeastBusyLoggingHandler(CustomLogger):
     async def async_log_success_event(
         self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
     ) -> None:
-        await self._async_increment(kwargs, -1)
+        await self._async_settle(kwargs)
         if self.test_flag:
             self.logged_success += 1
 
@@ -152,15 +371,20 @@ class LeastBusyLoggingHandler(CustomLogger):
     async def async_log_failure_event(
         self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
     ) -> None:
-        await self._async_increment(kwargs, -1)
+        await self._async_settle(kwargs)
         if self.test_flag:
             self.logged_failure += 1
 
     @with_service_target("router_usage")
     def get_available_deployments(
-        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
+        self,
+        model_group: str,
+        healthy_deployments: Sequence[Mapping[str, object]],
+        reserve: bool = False,
+        call_id: str | None = None,
     ) -> Mapping[str, object] | None:
-        keys: Final = _request_count_keys(model_group, healthy_deployments)
+        deployment_ids: Final = _deployment_ids(healthy_deployments)
+        keys: Final = _request_count_keys(model_group, deployment_ids)
         redis_cache: Final = self.router_cache.redis_cache
         if redis_cache is not None:
             try:
@@ -168,15 +392,19 @@ class LeastBusyLoggingHandler(CustomLogger):
             except Exception as e:
                 _warn_unreadable(model_group, e)
             else:
-                return _least_busy(healthy_deployments, shared)
-        local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
-        return _least_busy(healthy_deployments, local)
+                return self._select(healthy_deployments, deployment_ids, keys, shared, reserve, call_id)
+        return self._select(healthy_deployments, deployment_ids, keys, None, reserve, call_id)
 
     @with_service_target("router_usage")
     async def async_get_available_deployments(
-        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
+        self,
+        model_group: str,
+        healthy_deployments: Sequence[Mapping[str, object]],
+        reserve: bool = False,
+        call_id: str | None = None,
     ) -> Mapping[str, object] | None:
-        keys: Final = _request_count_keys(model_group, healthy_deployments)
+        deployment_ids: Final = _deployment_ids(healthy_deployments)
+        keys: Final = _request_count_keys(model_group, deployment_ids)
         redis_cache: Final = self.router_cache.redis_cache
         if redis_cache is not None:
             try:
@@ -184,42 +412,51 @@ class LeastBusyLoggingHandler(CustomLogger):
             except Exception as e:
                 _warn_unreadable(model_group, e)
             else:
-                return _least_busy(healthy_deployments, shared)
-        local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys)
-        return _least_busy(healthy_deployments, local)
+                return self._select(healthy_deployments, deployment_ids, keys, shared, reserve, call_id)
+        return self._select(healthy_deployments, deployment_ids, keys, None, reserve, call_id)
 
-    def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
+    def _settle(self, kwargs: Mapping[str, object]) -> None:
         ref: Final = _deployment_ref(kwargs)
-        if ref is None:
+        if ref is None or not self._end_attempt(ref, call_id_of(kwargs)):
             return
-        key: Final = _request_count_key(*ref)
-        redis_cache: Final = self.router_cache.redis_cache
+        self._count_shared(_request_count_key(*ref), -1)
+
+    async def _async_settle(self, kwargs: Mapping[str, object]) -> None:
+        ref: Final = _deployment_ref(kwargs)
+        if ref is None or not self._end_attempt(ref, call_id_of(kwargs)):
+            return
+        await self._async_count_shared(_request_count_key(*ref), -1)
+
+    def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]:
+        return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
+
+    def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]:
+        return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True))
+
+    def _count_locally(self, key: str, delta: int) -> None:
         try:
             local: Final = self.router_cache.increment_cache(
                 key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
             )
             if local < 0:
                 self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
-            if redis_cache is None:
-                return
+        except TypeError as e:
+            _warn_unwritable(key, e)
+
+    def _count_shared(self, key: str, delta: int) -> None:
+        redis_cache: Final = self.router_cache.redis_cache
+        if redis_cache is None:
+            return
+        try:
             redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
         except Exception as e:
             _warn_unwritable(key, e)
 
-    async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None:
-        ref: Final = _deployment_ref(kwargs)
-        if ref is None:
-            return
-        key: Final = _request_count_key(*ref)
+    async def _async_count_shared(self, key: str, delta: int) -> None:
         redis_cache: Final = self.router_cache.redis_cache
+        if redis_cache is None:
+            return
         try:
-            local: Final = await self.router_cache.async_increment_cache(
-                key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
-            )
-            if local is not None and local < 0:
-                await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
-            if redis_cache is None:
-                return
             await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
         except Exception as e:
             _warn_unwritable(key, e)

litellm/router_utils/cooldown_handlers.py

@@ -847,6 +847,45 @@ def is_stream_stall(exception: BaseException) -> bool:
     return isinstance(exception, litellm.Timeout) and isinstance(exception.__cause__, MidStreamFallbackError)
 
 
+def is_caller_entra_credential_failure(
+    deployment: Deployment | None,
+    router_defaults: Mapping[str, object],
+    model_call_details: Mapping[str, object],
+    exception_status: str | int,
+    exception: BaseException,
+) -> bool:
+    if cast_exception_status_to_int(exception_status) != 401 or deployment is None:
+        return False
+
+    from litellm.llms.azure.common_utils import caused_by_entra_rejection
+
+    if not caused_by_entra_rejection(exception):
+        return False
+    from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, azure_entra_environment_defaults
+
+    call_params: Final = _CREDENTIAL_VALUES_ADAPTER.validate_python(model_call_details.get("litellm_params") or {})
+
+    credential_name: Final = deployment.litellm_params.litellm_credential_name
+    configured_by_precedence: Final[tuple[Mapping[str, object], ...]] = (
+        router_defaults,
+        _CREDENTIAL_VALUES_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)),
+        _CREDENTIAL_VALUES_ADAPTER.validate_python(
+            CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {}
+        ),
+    )
+    environment: Final = azure_entra_environment_defaults()
+    return any(
+        call_params.get(key) is not None
+        and (call_params.get(key) or environment.get(key))
+        != (_first_configured(configured_by_precedence, key) or environment.get(key))
+        for key in AZURE_ENTRA_LITELLM_PARAM_KEYS
+    )
+
+
+def _first_configured(sources: tuple[Mapping[str, object], ...], key: str) -> object:
+    return next((value for source in sources if (value := source.get(key)) is not None), None)
+
+
 def is_caller_scoped_auth_failure(deployment: Deployment | None, exception_status: str | int) -> bool:
     if cast_exception_status_to_int(exception_status) not in (401, 403) or deployment is None:
         return False

litellm/router_utils/fallback_event_handlers.py

@@ -28,6 +28,7 @@ from litellm.router_utils.cooldown_handlers import (
     _first_present,  # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils
     cast_exception_status_to_int,
     is_advisor_orchestration_failure,
+    is_caller_entra_credential_failure,
     is_caller_scoped_auth_failure,
     is_caller_timeout_408,
     is_stream_stall,
@@ -132,6 +133,14 @@ def _trigger_cooldown_for_failed_deployment(
                 deployment_id,
             )
             return
+        if is_caller_entra_credential_failure(
+            deployment, ctx.default_litellm_params(), model_call_details, exception_status, exception
+        ):
+            verbose_router_logger.debug(
+                "Not triggering cooldown for fallback deployment %s: Entra ID rejected a credential the caller sent.",
+                deployment_id,
+            )
+            return
 
         # Priority: deployment config > response header > router default, matching
         # Router.deployment_callback_on_failure's precedence for the primary path.

scripts/check_schema_compatibility.py

@@ -52,14 +52,20 @@ def check_breaking_changes(old_schema: Schema, new_schema: Schema) -> tuple[str,
     return tuple(changes())
 
 
-def main() -> int:
-    parser: Final = argparse.ArgumentParser()
-    parser.add_argument("--base-ref", required=True)
-    options: Final = parser.parse_args()
-    base_ref: Final = TypeAdapter(str).validate_python(options.base_ref)
-    base: Final = subprocess.run(
+def _merge_base(base_ref: str) -> str:
+    return subprocess.run(
         ("git", "merge-base", "HEAD", base_ref), capture_output=True, text=True, check=True
     ).stdout.strip()
+
+
+def main() -> int:
+    parser: Final = argparse.ArgumentParser()
+    source: Final = parser.add_mutually_exclusive_group(required=True)
+    source.add_argument("--base-ref", help="compare against the merge base of HEAD and this ref (needs history)")
+    source.add_argument("--base-commit", help="compare against this commit as is (works in a shallow clone)")
+    options: Final = TypeAdapter(dict[str, str | None]).validate_python(vars(parser.parse_args()))
+    base_commit: Final = options["base_commit"]
+    base: Final = base_commit if base_commit else _merge_base(options["base_ref"] or "")
     previous: Final = subprocess.run(
         ("git", "show", f"{base}:schema.prisma"), capture_output=True, text=True, check=True
     ).stdout

terraform/litellm/aws/README.md

@@ -212,33 +212,23 @@ orthogonal to the OTLP variables above; their credentials still go in
 
 ### Enterprise billing metrics
 
-License-gated request metering is opt-in and gated entirely on
-`billing_metrics_endpoint`. Empty (default) and no billing env is added to
-the container, so existing deployments are unchanged. Set it and both
-gateway and backend export billable-request counts over OTLP/HTTP,
-authenticating to the collector with the mTLS client certificate issued for
-your deployment.
-
-The proxy accepts the certificate, key, and CA bundle as either a file path
-or literal PEM content. This stack takes the PEM, writes each one to its own
-Secrets Manager entry, grants the task-execution role
-`secretsmanager:GetSecretValue` on them, and injects them as
-`LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and `_CA_CERT` when
-set), so no volume mount is needed on Fargate.
+License-gated request metering uses Cognito bearer authentication. Set the
+endpoint, client ID, and token URL as plain values, and provide the client
+secret through the sensitive Terraform variable. The client secret is stored
+in Secrets Manager and injected into both gateway and backend tasks
 
 ```hcl
 billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics"
+billing_metrics_client_id = "cognito-client-id"
+billing_metrics_token_url = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token"
 ```
 
 ```bash
-export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)"
-export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)"
+export TF_VAR_billing_metrics_client_secret="cognito-client-secret"
 ```
 
-`billing_metrics_ca_cert_pem` is only for private or test collectors whose
-CA is not in the system trust store; leave it empty against
-`telemetry.litellm.ai`. Metering requires an enterprise license, so pair
-this with `litellm_license`. To tune the export cadence, set
+The client ID and client secret are required when the endpoint is set
+Metering also requires an enterprise license. To tune the export cadence, set
 `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` /
 `backend_extra_env`
 

terraform/litellm/aws/ecs.tf

@@ -79,32 +79,19 @@ locals {
     { name = "OTEL_HEADERS", valueFrom = var.otel_headers_secret_arn },
   ] : []
 
-  # Enterprise request metering, gated on billing_metrics_endpoint. The
-  # endpoint rides in as a plain env var; the mTLS material is stored in
-  # Secrets Manager (secrets.tf) and injected as PEM-valued env vars, which
-  # the proxy accepts in place of file paths. Each PEM is wired only when the
-  # operator supplied it, so an empty ca_cert_pem falls back to the system
-  # trust store.
-  billing_metrics_enabled             = var.billing_metrics_endpoint != ""
-  billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != ""
-  billing_metrics_client_key_enabled  = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != ""
-  billing_metrics_ca_cert_enabled     = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != ""
+  # Enterprise request metering, gated on billing_metrics_endpoint.
+  billing_metrics_enabled               = var.billing_metrics_endpoint != ""
+  billing_metrics_client_secret_enabled = local.billing_metrics_enabled && var.billing_metrics_client_secret != ""
 
   billing_metrics_env = local.billing_metrics_enabled ? [
     { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint },
+    { name = "LITELLM_BILLING_METRICS_CLIENT_ID", value = var.billing_metrics_client_id },
+    { name = "LITELLM_BILLING_METRICS_TOKEN_URL", value = var.billing_metrics_token_url },
   ] : []
 
-  billing_metrics_secrets = concat(
-    local.billing_metrics_client_cert_enabled ? [
-      { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_client_cert[0].arn },
-    ] : [],
-    local.billing_metrics_client_key_enabled ? [
-      { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", valueFrom = aws_secretsmanager_secret.billing_metrics_client_key[0].arn },
-    ] : [],
-    local.billing_metrics_ca_cert_enabled ? [
-      { name = "LITELLM_BILLING_METRICS_CA_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_ca_cert[0].arn },
-    ] : [],
-  )
+  billing_metrics_secrets = local.billing_metrics_client_secret_enabled ? [
+    { name = "LITELLM_BILLING_METRICS_CLIENT_SECRET", valueFrom = aws_secretsmanager_secret.billing_metrics_client_secret[0].arn },
+  ] : []
 
   managed_db_env = var.create_database ? [
     { name = "IAM_TOKEN_DB_AUTH", value = "true" },
@@ -338,27 +325,12 @@ locals {
 
 # ---------- Gateway ----------
 resource "aws_ecs_task_definition" "gateway" {
-  # Metering needs a client certificate AND its key. Each secret is created only
-  # when its own PEM is supplied, so an endpoint set with a missing key would
-  # otherwise apply cleanly and leave the proxy logging "missing config" and
-  # never exporting. ca_cert_pem stays optional: empty means fall back to the
-  # system trust store.
-  #
-  # The guard lives here, on an unconditional resource, rather than on the cert
-  # secret: that secret is count-gated on the cert itself, so it has zero
-  # instances in exactly the case this must catch. Adding count or for_each to
-  # this resource would silently stop the guard from evaluating.
-  #
-  #   endpoint  cert  key  -> result
-  #   ""        any   any  -> metering off, no secrets created
-  #   set       set   set  -> metering on
-  #   set       any-missing -> plan fails here
   lifecycle {
     precondition {
       condition = var.billing_metrics_endpoint == "" || (
-        var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
+        var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != ""
       )
-      error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
+      error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set."
     }
 
     precondition {
@@ -465,9 +437,7 @@ resource "aws_ecs_service" "gateway" {
     aws_secretsmanager_secret_version.license,
     aws_secretsmanager_secret_version.database_url,
     aws_secretsmanager_secret_version.redis_url,
-    aws_secretsmanager_secret_version.billing_metrics_client_cert,
-    aws_secretsmanager_secret_version.billing_metrics_client_key,
-    aws_secretsmanager_secret_version.billing_metrics_ca_cert,
+    aws_secretsmanager_secret_version.billing_metrics_client_secret,
   ]
 
   tags = local.tags
@@ -481,9 +451,9 @@ resource "aws_ecs_task_definition" "backend" {
   lifecycle {
     precondition {
       condition = var.billing_metrics_endpoint == "" || (
-        var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
+        var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != ""
       )
-      error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
+      error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set."
     }
   }
 
@@ -566,9 +536,7 @@ resource "aws_ecs_service" "backend" {
     aws_secretsmanager_secret_version.ui_password,
     aws_secretsmanager_secret_version.database_url,
     aws_secretsmanager_secret_version.redis_url,
-    aws_secretsmanager_secret_version.billing_metrics_client_cert,
-    aws_secretsmanager_secret_version.billing_metrics_client_key,
-    aws_secretsmanager_secret_version.billing_metrics_ca_cert,
+    aws_secretsmanager_secret_version.billing_metrics_client_secret,
   ]
 
   tags = local.tags

terraform/litellm/aws/iam.tf

@@ -53,9 +53,7 @@ data "aws_iam_policy_document" "secrets_access" {
       [aws_secretsmanager_secret.master_key.arn],
       aws_secretsmanager_secret.license[*].arn,
       aws_secretsmanager_secret.ui_password[*].arn,
-      aws_secretsmanager_secret.billing_metrics_client_cert[*].arn,
-      aws_secretsmanager_secret.billing_metrics_client_key[*].arn,
-      aws_secretsmanager_secret.billing_metrics_ca_cert[*].arn,
+      aws_secretsmanager_secret.billing_metrics_client_secret[*].arn,
       aws_secretsmanager_secret.database_url[*].arn,
       aws_secretsmanager_secret.redis_url[*].arn,
       local.extra_secret_arns,

terraform/litellm/aws/secrets.tf

@@ -75,59 +75,21 @@ resource "aws_secretsmanager_secret_version" "ui_password" {
   secret_string = var.ui_password
 }
 
-# Billing-metrics mTLS material — only created when metering is enabled
-# (billing_metrics_endpoint non-empty) and the operator supplied the PEM.
-# The task-execution role gets GetSecretValue via iam.tf, and gateway +
-# backend pick the env vars up through shared_secrets in ecs.tf.
-resource "aws_secretsmanager_secret" "billing_metrics_client_cert" {
-  count = local.billing_metrics_client_cert_enabled ? 1 : 0
+resource "aws_secretsmanager_secret" "billing_metrics_client_secret" {
+  count = local.billing_metrics_client_secret_enabled ? 1 : 0
 
-  name                    = "${local.name}-billing-metrics-client-cert"
-  description             = "LITELLM_BILLING_METRICS_CLIENT_CERT for gateway + backend."
+  name                    = "${local.name}-billing-metrics-client-secret"
+  description             = "LITELLM_BILLING_METRICS_CLIENT_SECRET for gateway + backend."
   recovery_window_in_days = 0
 
   tags = local.tags
 }
 
-resource "aws_secretsmanager_secret_version" "billing_metrics_client_cert" {
-  count = local.billing_metrics_client_cert_enabled ? 1 : 0
+resource "aws_secretsmanager_secret_version" "billing_metrics_client_secret" {
+  count = local.billing_metrics_client_secret_enabled ? 1 : 0
 
-  secret_id     = aws_secretsmanager_secret.billing_metrics_client_cert[0].id
-  secret_string = var.billing_metrics_client_cert_pem
-}
-
-resource "aws_secretsmanager_secret" "billing_metrics_client_key" {
-  count = local.billing_metrics_client_key_enabled ? 1 : 0
-
-  name                    = "${local.name}-billing-metrics-client-key"
-  description             = "LITELLM_BILLING_METRICS_CLIENT_KEY for gateway + backend."
-  recovery_window_in_days = 0
-
-  tags = local.tags
-}
-
-resource "aws_secretsmanager_secret_version" "billing_metrics_client_key" {
-  count = local.billing_metrics_client_key_enabled ? 1 : 0
-
-  secret_id     = aws_secretsmanager_secret.billing_metrics_client_key[0].id
-  secret_string = var.billing_metrics_client_key_pem
-}
-
-resource "aws_secretsmanager_secret" "billing_metrics_ca_cert" {
-  count = local.billing_metrics_ca_cert_enabled ? 1 : 0
-
-  name                    = "${local.name}-billing-metrics-ca-cert"
-  description             = "LITELLM_BILLING_METRICS_CA_CERT for gateway + backend."
-  recovery_window_in_days = 0
-
-  tags = local.tags
-}
-
-resource "aws_secretsmanager_secret_version" "billing_metrics_ca_cert" {
-  count = local.billing_metrics_ca_cert_enabled ? 1 : 0
-
-  secret_id     = aws_secretsmanager_secret.billing_metrics_ca_cert[0].id
-  secret_string = var.billing_metrics_ca_cert_pem
+  secret_id     = aws_secretsmanager_secret.billing_metrics_client_secret[0].id
+  secret_string = var.billing_metrics_client_secret
 }
 
 resource "aws_secretsmanager_secret" "db_master_password" {

terraform/litellm/aws/variables.tf

@@ -764,13 +764,7 @@ variable "otel_headers_secret_arn" {
 
 # ---------- Enterprise billing metrics ----------
 #
-# License-gated request metering. Opt-in and gated entirely on
-# billing_metrics_endpoint: leave it empty (the default) and nothing
-# metering-related lands in the container env. Set it and gateway + backend
-# export billable-request counts over OTLP/HTTP, authenticating to the
-# collector with an mTLS client cert. The proxy accepts the cert, key, and CA
-# as either a file path or literal PEM content, so on Fargate they are
-# injected straight from Secrets Manager as env vars and no volume is needed.
+# License-gated request metering. The client secret is stored in Secrets Manager.
 
 variable "billing_metrics_endpoint" {
   description = <<-EOT
@@ -784,44 +778,34 @@ variable "billing_metrics_endpoint" {
   default     = ""
 }
 
-variable "billing_metrics_client_cert_pem" {
+variable "billing_metrics_client_id" {
   description = <<-EOT
-    PEM content of the mTLS client certificate issued for this deployment.
-    When billing_metrics_endpoint is set, the stack stores this in a
-    `<tenant>-litellm-<env>-billing-metrics-client-cert` Secrets Manager
-    entry, grants the task-execution role GetSecretValue on it, and exposes
-    it to gateway + backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required
-    whenever metering is enabled.
+    Cognito OAuth2 client ID exposed to the gateway and backend as
+    LITELLM_BILLING_METRICS_CLIENT_ID. Required when billing_metrics_endpoint
+    is set.
+  EOT
+  type        = string
+  default     = ""
+}
+
+variable "billing_metrics_client_secret" {
+  description = <<-EOT
+    Cognito OAuth2 client secret stored in Secrets Manager and exposed to
+    the gateway and backend as LITELLM_BILLING_METRICS_CLIENT_SECRET.
+    Required when billing_metrics_endpoint is set.
   EOT
   type        = string
   default     = ""
   sensitive   = true
 }
 
-variable "billing_metrics_client_key_pem" {
+variable "billing_metrics_token_url" {
   description = <<-EOT
-    PEM content of the private key matching
-    billing_metrics_client_cert_pem. Stored in a
-    `<tenant>-litellm-<env>-billing-metrics-client-key` Secrets Manager
-    entry and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required
-    whenever metering is enabled.
+    Cognito OAuth2 client-credentials token endpoint exposed to gateway and
+    backend as LITELLM_BILLING_METRICS_TOKEN_URL.
   EOT
   type        = string
-  default     = ""
-  sensitive   = true
-}
-
-variable "billing_metrics_ca_cert_pem" {
-  description = <<-EOT
-    PEM content of the CA bundle used to verify the metering collector.
-    Only needed for private or test collectors whose CA is not in the
-    system trust store; telemetry.litellm.ai is publicly trusted, so leave
-    this empty for production. When set, it is exposed as
-    LITELLM_BILLING_METRICS_CA_CERT.
-  EOT
-  type        = string
-  default     = ""
-  sensitive   = true
+  default     = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token"
 }
 
 # ---------- Collector sidecar ----------

terraform/litellm/gcp/README.md

@@ -210,33 +210,23 @@ Behavior matches the AWS stack 1:1; the only naming differences are
 
 ### Enterprise billing metrics
 
-License-gated request metering is opt-in and gated entirely on
-`billing_metrics_endpoint`. Empty (default) and no billing env is added to
-the container, so existing deployments are unchanged. Set it and both
-gateway and backend export billable-request counts over OTLP/HTTP,
-authenticating to the collector with the mTLS client certificate issued for
-your deployment.
-
-The proxy accepts the certificate, key, and CA bundle as either a file path
-or literal PEM content. This stack takes the PEM, writes each one to its own
-Secret Manager entry, grants the runtime service account
-`roles/secretmanager.secretAccessor` on them, and injects them as Cloud Run
-secret env vars `LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and
-`_CA_CERT` when set), so no volume mount is needed.
+License-gated request metering uses Cognito bearer authentication. Set the
+endpoint, client ID, and token URL as plain values, and provide the client
+secret through the sensitive Terraform variable. The client secret is stored
+in Secret Manager and injected into both gateway and backend services
 
 ```hcl
 billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics"
+billing_metrics_client_id = "cognito-client-id"
+billing_metrics_token_url = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token"
 ```
 
 ```bash
-export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)"
-export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)"
+export TF_VAR_billing_metrics_client_secret="cognito-client-secret"
 ```
 
-`billing_metrics_ca_cert_pem` is only for private or test collectors whose
-CA is not in the system trust store; leave it empty against
-`telemetry.litellm.ai`. Metering requires an enterprise license, so pair
-this with `litellm_license`. To tune the export cadence, set
+The client ID and client secret are required when the endpoint is set
+Metering also requires an enterprise license. To tune the export cadence, set
 `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` /
 `backend_extra_env`
 

terraform/litellm/gcp/cloudrun.tf

@@ -62,32 +62,19 @@ locals {
     { name = "OTEL_HEADERS", secret = var.otel_headers_secret, version = "latest" },
   ] : []
 
-  # Enterprise request metering, gated on billing_metrics_endpoint. The
-  # endpoint rides in as a plain env var; the mTLS material lives in Secret
-  # Manager (secrets.tf) and is injected as PEM-valued env vars, which the
-  # proxy accepts in place of file paths. Each PEM is wired only when the
-  # operator supplied it, so an empty ca_cert_pem falls back to the system
-  # trust store.
-  billing_metrics_enabled             = var.billing_metrics_endpoint != ""
-  billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != ""
-  billing_metrics_client_key_enabled  = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != ""
-  billing_metrics_ca_cert_enabled     = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != ""
+  # Enterprise request metering, gated on billing_metrics_endpoint.
+  billing_metrics_enabled               = var.billing_metrics_endpoint != ""
+  billing_metrics_client_secret_enabled = local.billing_metrics_enabled && var.billing_metrics_client_secret != ""
 
   billing_metrics_env_kv = local.billing_metrics_enabled ? [
     { name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint },
+    { name = "LITELLM_BILLING_METRICS_CLIENT_ID", value = var.billing_metrics_client_id },
+    { name = "LITELLM_BILLING_METRICS_TOKEN_URL", value = var.billing_metrics_token_url },
   ] : []
 
-  billing_metrics_env_secrets = concat(
-    local.billing_metrics_client_cert_enabled ? [
-      { name = "LITELLM_BILLING_METRICS_CLIENT_CERT", secret = google_secret_manager_secret.billing_metrics_client_cert[0].id, version = "latest" },
-    ] : [],
-    local.billing_metrics_client_key_enabled ? [
-      { name = "LITELLM_BILLING_METRICS_CLIENT_KEY", secret = google_secret_manager_secret.billing_metrics_client_key[0].id, version = "latest" },
-    ] : [],
-    local.billing_metrics_ca_cert_enabled ? [
-      { name = "LITELLM_BILLING_METRICS_CA_CERT", secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id, version = "latest" },
-    ] : [],
-  )
+  billing_metrics_env_secrets = local.billing_metrics_client_secret_enabled ? [
+    { name = "LITELLM_BILLING_METRICS_CLIENT_SECRET", secret = google_secret_manager_secret.billing_metrics_client_secret[0].id, version = "latest" },
+  ] : []
 
   # Cloud Run v2 secret env vars use value_source.secret_key_ref pointing at a
   # secret resource ID. Shared between gateway and backend (the migrations
@@ -481,9 +468,7 @@ resource "google_cloud_run_v2_service" "gateway" {
     google_secret_manager_secret_iam_member.license,
     google_secret_manager_secret_iam_member.extras,
     google_secret_manager_secret_iam_member.otel_headers,
-    google_secret_manager_secret_iam_member.billing_metrics_client_cert,
-    google_secret_manager_secret_iam_member.billing_metrics_client_key,
-    google_secret_manager_secret_iam_member.billing_metrics_ca_cert,
+    google_secret_manager_secret_iam_member.billing_metrics_client_secret,
     google_secret_manager_secret_iam_member.metrics_run_monitoring,
     google_project_iam_member.runtime_metric_writer,
     google_storage_bucket_iam_member.proxy_config_runtime,
@@ -603,9 +588,7 @@ resource "google_cloud_run_v2_service" "backend" {
     google_secret_manager_secret_iam_member.ui_password,
     google_secret_manager_secret_iam_member.extras,
     google_secret_manager_secret_iam_member.otel_headers,
-    google_secret_manager_secret_iam_member.billing_metrics_client_cert,
-    google_secret_manager_secret_iam_member.billing_metrics_client_key,
-    google_secret_manager_secret_iam_member.billing_metrics_ca_cert,
+    google_secret_manager_secret_iam_member.billing_metrics_client_secret,
     google_storage_bucket_iam_member.proxy_config_runtime,
     google_sql_user.app,
     terraform_data.migration,

terraform/litellm/gcp/iam.tf

@@ -10,9 +10,9 @@ resource "google_service_account" "runtime" {
   lifecycle {
     precondition {
       condition = !var.create_runtime || var.billing_metrics_endpoint == "" || (
-        var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
+        var.billing_metrics_client_id != "" && var.billing_metrics_client_secret != ""
       )
-      error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set and create_runtime is true."
+      error_message = "billing_metrics_client_id and billing_metrics_client_secret are required when billing_metrics_endpoint is set and create_runtime is true."
     }
   }
 }
@@ -91,28 +91,10 @@ resource "google_secret_manager_secret_iam_member" "otel_headers" {
   member    = "serviceAccount:${google_service_account.runtime.email}"
 }
 
-# Billing-metrics mTLS accessors — only created when request metering is
-# enabled and the matching PEM was supplied.
-resource "google_secret_manager_secret_iam_member" "billing_metrics_client_cert" {
-  count = local.billing_metrics_client_cert_enabled ? 1 : 0
+resource "google_secret_manager_secret_iam_member" "billing_metrics_client_secret" {
+  count = local.billing_metrics_client_secret_enabled ? 1 : 0
 
-  secret_id = google_secret_manager_secret.billing_metrics_client_cert[0].id
-  role      = "roles/secretmanager.secretAccessor"
-  member    = "serviceAccount:${google_service_account.runtime.email}"
-}
-
-resource "google_secret_manager_secret_iam_member" "billing_metrics_client_key" {
-  count = local.billing_metrics_client_key_enabled ? 1 : 0
-
-  secret_id = google_secret_manager_secret.billing_metrics_client_key[0].id
-  role      = "roles/secretmanager.secretAccessor"
-  member    = "serviceAccount:${google_service_account.runtime.email}"
-}
-
-resource "google_secret_manager_secret_iam_member" "billing_metrics_ca_cert" {
-  count = local.billing_metrics_ca_cert_enabled ? 1 : 0
-
-  secret_id = google_secret_manager_secret.billing_metrics_ca_cert[0].id
+  secret_id = google_secret_manager_secret.billing_metrics_client_secret[0].id
   role      = "roles/secretmanager.secretAccessor"
   member    = "serviceAccount:${google_service_account.runtime.email}"
 }

terraform/litellm/gcp/secrets.tf

@@ -64,59 +64,21 @@ resource "google_secret_manager_secret_version" "ui_password" {
   secret_data = var.ui_password
 }
 
-# Billing-metrics mTLS material — only created when metering is enabled
-# (billing_metrics_endpoint non-empty) and the operator supplied the PEM.
-# The runtime SA gets accessor permission via iam.tf, and gateway + backend
-# pick the env vars up through billing_metrics_env_secrets in cloudrun.tf.
-resource "google_secret_manager_secret" "billing_metrics_client_cert" {
-  count = local.billing_metrics_client_cert_enabled ? 1 : 0
+resource "google_secret_manager_secret" "billing_metrics_client_secret" {
+  count = local.billing_metrics_client_secret_enabled ? 1 : 0
 
-  secret_id = "${local.name}-billing-metrics-client-cert"
+  secret_id = "${local.name}-billing-metrics-client-secret"
   labels    = local.labels
   replication {
     auto {}
   }
 }
 
-resource "google_secret_manager_secret_version" "billing_metrics_client_cert" {
-  count = local.billing_metrics_client_cert_enabled ? 1 : 0
+resource "google_secret_manager_secret_version" "billing_metrics_client_secret" {
+  count = local.billing_metrics_client_secret_enabled ? 1 : 0
 
-  secret      = google_secret_manager_secret.billing_metrics_client_cert[0].id
-  secret_data = var.billing_metrics_client_cert_pem
-}
-
-resource "google_secret_manager_secret" "billing_metrics_client_key" {
-  count = local.billing_metrics_client_key_enabled ? 1 : 0
-
-  secret_id = "${local.name}-billing-metrics-client-key"
-  labels    = local.labels
-  replication {
-    auto {}
-  }
-}
-
-resource "google_secret_manager_secret_version" "billing_metrics_client_key" {
-  count = local.billing_metrics_client_key_enabled ? 1 : 0
-
-  secret      = google_secret_manager_secret.billing_metrics_client_key[0].id
-  secret_data = var.billing_metrics_client_key_pem
-}
-
-resource "google_secret_manager_secret" "billing_metrics_ca_cert" {
-  count = local.billing_metrics_ca_cert_enabled ? 1 : 0
-
-  secret_id = "${local.name}-billing-metrics-ca-cert"
-  labels    = local.labels
-  replication {
-    auto {}
-  }
-}
-
-resource "google_secret_manager_secret_version" "billing_metrics_ca_cert" {
-  count = local.billing_metrics_ca_cert_enabled ? 1 : 0
-
-  secret      = google_secret_manager_secret.billing_metrics_ca_cert[0].id
-  secret_data = var.billing_metrics_ca_cert_pem
+  secret      = google_secret_manager_secret.billing_metrics_client_secret[0].id
+  secret_data = var.billing_metrics_client_secret
 }
 
 resource "google_secret_manager_secret" "metrics_run_monitoring" {

terraform/litellm/gcp/variables.tf

@@ -602,14 +602,7 @@ variable "gateway_metrics_collector_image" {
 
 # ---------- Enterprise billing metrics ----------
 #
-# License-gated request metering. Opt-in and gated entirely on
-# billing_metrics_endpoint: leave it empty (the default) and nothing
-# metering-related is added to the container env. Set it and gateway +
-# backend export billable-request counts over OTLP/HTTP, authenticating to
-# the collector with an mTLS client cert. The proxy accepts the cert, key,
-# and CA as either a file path or literal PEM content, so on Cloud Run they
-# are injected straight from Secret Manager as env vars and no volume is
-# needed.
+# License-gated request metering. The client secret is stored in Secret Manager.
 
 variable "billing_metrics_endpoint" {
   description = <<-EOT
@@ -623,44 +616,34 @@ variable "billing_metrics_endpoint" {
   default     = ""
 }
 
-variable "billing_metrics_client_cert_pem" {
+variable "billing_metrics_client_id" {
   description = <<-EOT
-    PEM content of the mTLS client certificate issued for this deployment.
-    When billing_metrics_endpoint is set, the stack stores this in a
-    `<tenant>-litellm-<env>-billing-metrics-client-cert` Secret Manager
-    entry, grants the runtime SA accessor on it, and exposes it to gateway +
-    backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required whenever
-    metering is enabled.
+    Cognito OAuth2 client ID exposed to the gateway and backend as
+    LITELLM_BILLING_METRICS_CLIENT_ID. Required when billing_metrics_endpoint
+    is set.
+  EOT
+  type        = string
+  default     = ""
+}
+
+variable "billing_metrics_client_secret" {
+  description = <<-EOT
+    Cognito OAuth2 client secret stored in Secret Manager and exposed to
+    the gateway and backend as LITELLM_BILLING_METRICS_CLIENT_SECRET.
+    Required when billing_metrics_endpoint is set.
   EOT
   type        = string
   default     = ""
   sensitive   = true
 }
 
-variable "billing_metrics_client_key_pem" {
+variable "billing_metrics_token_url" {
   description = <<-EOT
-    PEM content of the private key matching
-    billing_metrics_client_cert_pem. Stored in a
-    `<tenant>-litellm-<env>-billing-metrics-client-key` Secret Manager entry
-    and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required whenever
-    metering is enabled.
+    Cognito OAuth2 client-credentials token endpoint exposed to gateway and
+    backend as LITELLM_BILLING_METRICS_TOKEN_URL.
   EOT
   type        = string
-  default     = ""
-  sensitive   = true
-}
-
-variable "billing_metrics_ca_cert_pem" {
-  description = <<-EOT
-    PEM content of the CA bundle used to verify the metering collector.
-    Only needed for private or test collectors whose CA is not in the
-    system trust store; telemetry.litellm.ai is publicly trusted, so leave
-    this empty for production. When set, it is exposed as
-    LITELLM_BILLING_METRICS_CA_CERT.
-  EOT
-  type        = string
-  default     = ""
-  sensitive   = true
+  default     = "https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token"
 }
 
 # ---------- Collector sidecar ----------

tests/integration/observability/test_arize_otel_v2_openinference_chaos.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
+import socket
 import threading
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
 from pathlib import Path
+from queue import SimpleQueue
 from typing import Final
-from urllib.parse import urlsplit
 
 import httpx
 import pytest
@@ -44,7 +45,7 @@ from _openinference_support import (
 )
 from integration._support.client import Gateway, eventually
 from integration._support.process import graceful_stop_seconds
-from integration._support.wire import Reply, Request, wire_server
+from integration._support.wire import Reply, Request, Wire, wire_server
 
 
 def _call(
@@ -249,11 +250,12 @@ def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: P
             _assert_response(response, marker, surface, stream, model_name)
 
     with ExitStack() as servers:
-        initial_stack: Final = servers.enter_context(ExitStack())
-        stopped_destination: Final = initial_stack.enter_context(wire_server(_owned_sink_handler(sink)))
-        sink_port: Final = urlsplit(stopped_destination.url).port
-        assert sink_port is not None, stopped_destination.url
-        initial_stack.close()
+        outage_socket: Final = servers.enter_context(socket.socket(socket.AF_INET, socket.SOCK_STREAM))
+        outage_socket.bind(("127.0.0.1", 0))  # test-quality-ok: socket stays bound throughout the outage
+        sink_port: Final = outage_socket.getsockname()[1]
+        stopped_destination: Final = Wire(
+            f"http://127.0.0.1:{sink_port}", SimpleQueue(), SimpleQueue(), SimpleQueue(), SimpleQueue()
+        )
         with _rig(
             gateway,
             tmp_path,
@@ -272,6 +274,7 @@ def test_arize_otel_v2_f1_sink_outage_and_recovery(gateway: Gateway, tmp_path: P
                 with httpx.Client(trust_env=False) as client, pytest.raises(httpx.ConnectError):
                     client.get(stopped_destination.url + "/health", timeout=2)
                 recovered_stack: Final = servers.enter_context(ExitStack())
+                outage_socket.close()
                 recovered_destination: Final = recovered_stack.enter_context(
                     wire_server(_owned_sink_handler(sink), port=sink_port)
                 )

tests/integration/routing/test_least_busy_concurrent_decisions.py

@@ -0,0 +1,90 @@
+import json
+import threading
+from collections import Counter
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import yaml
+from integration._support.client import Gateway, gateway_from_environment
+from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, wire_server
+
+pytestmark = pytest.mark.timeout(240)
+
+_GROUP: Final = "least-busy-decider"
+_DEPLOYMENTS: Final = ("least-busy-decider-a", "least-busy-decider-b")
+_BURST: Final = 4
+_REPLY: Final = json.dumps(
+    {
+        "model": "Qwen/Qwen3-0.6B",
+        "answers": {"billing": {"type": "noul", "noul": 0.91}},
+        "usage": {"input_tokens": 367, "output_tokens": 3},
+    }
+).encode()
+_DECISION: Final = {
+    "model": _GROUP,
+    "input": "I was charged twice for order 4417.",
+    "questions": [{"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}],
+}
+
+
+def _held_until_the_burst_arrives(arrived: threading.Barrier) -> Callable[[Request], Reply]:
+    def respond(request: Request) -> Reply:
+        arrived.wait()
+        return Reply(body=_REPLY)
+
+    return respond
+
+
+def _write_config(path: Path, provider_url: str) -> Path:
+    path.write_text(
+        yaml.safe_dump(
+            {
+                "model_list": [
+                    {
+                        "model_name": _GROUP,
+                        "litellm_params": {"model": "hosted_vllm/Qwen/Qwen3-0.6B", "api_base": provider_url},
+                        "model_info": {"id": deployment_id},
+                    }
+                    for deployment_id in _DEPLOYMENTS
+                ],
+                "general_settings": {
+                    "master_key": "os.environ/LITELLM_MASTER_KEY",
+                    "database_url": "os.environ/DATABASE_URL",
+                    "disable_model_info_refresh": True,
+                },
+                "router_settings": {"routing_strategy": "least-busy", "num_retries": 0, "disable_cooldowns": True},
+            }
+        )
+    )
+    return path
+
+
+def _decide(proxy: Gateway) -> httpx.Response:
+    return proxy.client.post(
+        "/v1/decisions", json=_DECISION, headers={"Authorization": f"Bearer {proxy.key}"}, timeout=120
+    )
+
+
+def test_least_busy_spreads_a_burst_of_decisions_calls_across_the_group(tmp_path: Path) -> None:
+    arrived: Final = threading.Barrier(_BURST, timeout=60)
+    with (
+        wire_server(_held_until_the_burst_arrives(arrived)) as upstream,
+        gateway_from_environment() as rig,
+        owned_proxy(
+            rig,
+            tmp_path,
+            {"STORE_MODEL_IN_DB": "False"},
+            config=_write_config(tmp_path / "least-busy.yaml", upstream.url),
+        ) as proxy,
+        ThreadPoolExecutor(max_workers=_BURST) as pool,
+    ):
+        responses: Final = tuple(pool.map(lambda _: _decide(proxy), range(_BURST)))
+    assert [response.status_code for response in responses] == [200] * _BURST, [r.text for r in responses]
+    assert Counter(response.headers["x-litellm-model-id"] for response in responses) == Counter(
+        {deployment_id: _BURST // len(_DEPLOYMENTS) for deployment_id in _DEPLOYMENTS}
+    )

tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py

@@ -1856,6 +1856,51 @@ def test_transform_request_system_only_message_maps_to_system_input_item():
     assert not result.get("instructions")
 
 
+_JSON_SYSTEM_PROMPT: Final = "You output JSON objects."
+_PLAIN_USER_ITEM: Final = {
+    "type": "message",
+    "role": "user",
+    "content": [{"type": "input_text", "text": "List three colors."}],
+}
+_COLORS_SCHEMA: Final = {
+    "type": "json_schema",
+    "json_schema": {"name": "colors", "schema": {"type": "object", "properties": {}}},
+}
+
+
+@pytest.mark.parametrize(
+    ("response_format", "expected_input", "expected_instructions"),
+    [
+        (
+            {"type": "json_object"},
+            [
+                {"type": "message", "role": "system", "content": [{"type": "input_text", "text": _JSON_SYSTEM_PROMPT}]},
+                _PLAIN_USER_ITEM,
+            ],
+            None,
+        ),
+        (_COLORS_SCHEMA, [_PLAIN_USER_ITEM], _JSON_SYSTEM_PROMPT),
+    ],
+    ids=["json_object_keeps_system_in_input", "json_schema_keeps_system_in_instructions"],
+)
+def test_transform_request_json_object_keeps_system_prompt_where_the_json_word_check_reads_it(
+    response_format, expected_input, expected_instructions
+):
+    result = LiteLLMResponsesTransformationHandler().transform_request(
+        model="gpt-5.4-mini",
+        messages=[
+            {"role": "system", "content": _JSON_SYSTEM_PROMPT},
+            {"role": "user", "content": "List three colors."},
+        ],
+        optional_params={"response_format": response_format},
+        litellm_params={},
+        headers={},
+        litellm_logging_obj=Mock(),
+    )
+
+    assert (result["input"], result.get("instructions")) == (expected_input, expected_instructions), result
+
+
 def test_transform_request_single_char_keys_not_matched():
     """Test that single-character keys are not incorrectly matched to 'metadata' or 'previous_response_id'
 

tests/unit/decisions/test_main.py

@@ -6,8 +6,10 @@ import copy
 import json
 import re
 import threading
+import uuid
 from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
 from concurrent.futures import wait
+from datetime import datetime
 from functools import partial
 from types import MappingProxyType
 from typing import Final, TypeAlias
@@ -300,6 +302,43 @@ async def test_router_dispatches_typesafe_decisions_without_api_base(
     assert response.answers["sentiment"].choice == "positive"
 
 
+@pytest.mark.asyncio
+async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_finishes(
+    respx_mock: respx.MockRouter,
+) -> None:
+    respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE)
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "jev",
+                "litellm_params": {"model": "typesafe/jev-latest", "api_key": "k"},
+                "model_info": {"id": deployment_id},
+            }
+            for deployment_id in ("jev-a", "jev-b")
+        ],
+        routing_strategy="least-busy",
+    )
+
+    async def call_through_proxy_logging() -> object:
+        logging_obj, _ = litellm.utils.function_setup(
+            original_function="adecisions",
+            rules_obj=litellm.utils.Rules(),
+            start_time=datetime(2026, 1, 1),
+            model="jev",
+            litellm_call_id=str(uuid.uuid4()),
+        )
+        response: Final = await router.adecisions(
+            model="jev",
+            input="Is this a defect?",
+            questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
+            litellm_logging_obj=logging_obj,
+        )
+        await _drain_logging_worker()
+        return response._hidden_params["model_id"]
+
+    assert (await call_through_proxy_logging(), await call_through_proxy_logging()) == ("jev-a", "jev-a")
+
+
 def test_decisions_uses_the_same_wire_contract_for_sync_calls(respx_mock: respx.MockRouter) -> None:
     route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE)
 

tests/unit/integrations/test_prometheus_invalid_key_filtering.py

@@ -512,6 +512,20 @@ def _provider_auth_failure_kwargs() -> dict[str, object]:
     }
 
 
+def _deployment_entra_rejection_kwargs() -> dict[str, object]:
+    from litellm.llms.azure.common_utils import EntraRejectionError
+
+    reason: Final = "Failed to get Azure AD token: AADSTS7000215: Invalid client secret provided."
+    try:
+        raise litellm.AuthenticationError(message=reason, llm_provider="azure", model="gpt-4.1") from (
+            EntraRejectionError(status_code=401, message=reason)
+        )
+    except litellm.AuthenticationError as raised:
+        return {
+            key: value for key, value in _provider_auth_failure_kwargs().items() if key not in CALLED_THE_PROVIDER
+        } | {"exception": raised}
+
+
 def _metric_values(name: str, model_id: str) -> list[float]:
     return [
         sample.value
@@ -541,6 +555,20 @@ class TestProviderAuthFailures:
         assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0]
         assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0]
 
+    @pytest.mark.asyncio
+    async def test_log_failure_event_records_entra_rejecting_the_deployment_credential_before_the_call(
+        self, prometheus_logger
+    ):
+        await prometheus_logger.async_log_failure_event(
+            kwargs=_deployment_entra_rejection_kwargs(),
+            response_obj=None,
+            start_time=None,
+            end_time=None,
+        )
+
+        assert _metric_values("litellm_llm_api_failed_requests_metric_total", "deployment-1") == [1.0]
+        assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0]
+
     def test_deployment_failure_metrics_record_a_provider_401(self, prometheus_logger):
         prometheus_logger.set_llm_deployment_failure_metrics(_provider_auth_failure_kwargs())
 

tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py

@@ -4332,6 +4332,37 @@ def test_azure_gpt_6_foundry_price_sheet(_local_model_cost_map, model_base):
         assert azure_eu_info[field] == pytest.approx(1.2 * base)
 
 
+@pytest.mark.parametrize("model", ["azure/gpt-6-luna", "azure/gpt-6-luna-2026-09-22"])
+@pytest.mark.parametrize(
+    ("prompt_tokens", "input_rate", "cache_read_rate", "output_rate"),
+    [
+        (100_000, 5e-08, 5e-09, 2.5e-07),
+        (300_000, 1e-07, 1e-08, 3.75e-07),
+    ],
+)
+def test_generic_cost_per_token_azure_gpt_6_luna_flex_tiers(
+    _local_model_cost_map, model, prompt_tokens, input_rate, cache_read_rate, output_rate
+):
+    """azure/gpt-6-luna bills Azure's Flex Global rates, half of Standard, with the long context step past 272K."""
+    cached_tokens = 20_000
+    completion_tokens = 1_000
+    usage = Usage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=prompt_tokens + completion_tokens,
+        prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
+    )
+    prompt_cost, completion_cost = generic_cost_per_token(
+        model=model,
+        usage=usage,
+        custom_llm_provider="azure",
+        service_tier="flex",
+    )
+    expected_prompt = (prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate
+    assert prompt_cost == pytest.approx(expected_prompt)
+    assert completion_cost == pytest.approx(completion_tokens * output_rate)
+
+
 @pytest.mark.parametrize("region_prefix", ["azure/", "azure/us/", "azure/eu/"])
 def test_azure_gpt_5_6_alias_matches_sol_pricing(_local_model_cost_map, region_prefix):
     """The bare gpt-5.6 alias routes to GPT-5.6 Sol, so every Azure region must bill the

tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py

@@ -327,7 +327,7 @@ EMPTY_ASSISTANT = pytest.mark.parametrize(
 )
 
 
-def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_linear_time():
+def test_flagged_placement_turns_reminders_before_a_long_run_of_unanswered_tool_search_uses_into_user_turns():
     search_uses: Final = tuple(
         {
             "role": "assistant",
@@ -349,12 +349,9 @@ def test_flagged_placement_handles_a_long_run_of_unanswered_tool_search_uses_in_
         {"role": "user", "content": "end"},
     ]
 
-    started: Final = time.perf_counter()
     placed: Final = place_mid_conversation_system(messages, supports_mid_conversation_system=True)
-    elapsed: Final = time.perf_counter() - started
 
     assert set(_roles(placed)[:21]) == {"user"}
-    assert elapsed < 1.0
 
 
 def test_flagged_placement_merges_every_system_run_that_shares_a_user_turn():

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -14643,3 +14643,50 @@ async def test_a_recovered_attempt_bills_a_failed_agentic_follow_up_but_not_a_fa
     assert own_cost > 0
     assert after_a_dropped_stream == pytest.approx(own_cost)
     assert after_a_failed_follow_up == pytest.approx(own_cost + _FAILED_ATTEMPT[1])
+
+
+def test_standard_logging_payload_reports_the_per_call_prices_the_request_was_billed_from() -> None:
+    """A direct call that sets its own prices is logged at those prices, not at the catalog row's rates."""
+    from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload
+    from litellm.litellm_core_utils.per_call_pricing import per_call_model_info
+
+    rates: Final = {"input_cost_per_token": 3e-06, "output_cost_per_token": 7e-06}
+    now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0)
+    logging_obj: Final = LitellmLogging(
+        model="gpt-5.4-nano",
+        messages=[{"role": "user", "content": "Hi"}],
+        stream=False,
+        call_type="completion",
+        start_time=now,
+        litellm_call_id="per-call-prices",
+        function_id="fn",
+    )
+    logging_obj.update_environment_variables(
+        model="gpt-5.4-nano", user="", optional_params={}, litellm_params=rates, custom_llm_provider="openai"
+    )
+    logging_obj.set_per_call_pricing(per_call_model_info("gpt-5.4-nano", "openai", rates, "chat"))
+    response: Final = ModelResponse(
+        model="gpt-5.4-nano",
+        choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))],
+        usage=litellm.Usage(prompt_tokens=10, completion_tokens=6, total_tokens=16),
+    )
+
+    payload: Final = get_standard_logging_object_payload(
+        kwargs={
+            "model": "gpt-5.4-nano",
+            "custom_llm_provider": "openai",
+            "litellm_params": rates,
+            "response_cost": logging_obj.response_cost_calculator(result=response),
+        },
+        init_response_obj=response,
+        start_time=now,
+        end_time=now,
+        logging_obj=logging_obj,
+        status="success",
+    )
+
+    assert payload is not None
+    assert payload["response_cost"] == pytest.approx(10 * 3e-06 + 6 * 7e-06)
+    model_map_value: Final = payload["model_map_information"]["model_map_value"]
+    assert model_map_value is not None
+    assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (3e-06, 7e-06)

tests/unit/litellm_core_utils/test_url_utils.py

@@ -16,11 +16,25 @@ from litellm.litellm_core_utils.url_utils import (
     encode_url_path_segment,
     encode_url_path_segments,
     ensure_api_route,
+    is_https_url_with_host,
     join_api_base,
     validate_url,
 )
 
 
+@pytest.mark.parametrize(
+    ("url", "expected"),
+    (
+        ("https://example.com/path", True),
+        ("http://example.com/path", False),
+        ("https:///path", False),
+        ("https://[", False),
+    ),
+)
+def test_is_https_url_with_host(url: str, expected: bool) -> None:
+    assert is_https_url_with_host(url) is expected
+
+
 @pytest.fixture
 def mock_dns_public(monkeypatch):
     """Resolve any hostname to 93.184.216.34 (public)."""

tests/unit/llms/azure/decisions/test_transformation.py

@@ -125,6 +125,29 @@ async def test_default_credential_refreshes_without_opt_in(
     )
 
 
+@pytest.mark.parametrize("is_async", (False, True))
+@pytest.mark.parametrize("systemone", (False, True))
+async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason(
+    monkeypatch: pytest.MonkeyPatch, is_async: bool, systemone: bool
+) -> None:
+    from azure.core.exceptions import ClientAuthenticationError
+
+    app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2"
+    credential = Mock(spec=["get_token"])
+    credential.get_token.side_effect = ClientAuthenticationError(
+        message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being "
+        f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'."
+    )
+    monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential))
+
+    with pytest.raises(litellm.AuthenticationError) as excinfo:
+        await call_decisions(is_async, systemone, {})
+
+    assert excinfo.value.status_code == 401
+    assert "AADSTS7000215" in str(excinfo.value)
+    assert app_client_id not in str(excinfo.value)
+
+
 @pytest.mark.parametrize("kind", ("service-principal", "oidc", "provider"))
 async def test_azure_decisions_reuses_existing_azure_credential_resolution(
     monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, kind: str

tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py

@@ -466,6 +466,12 @@ def _mock_response():
     response.content = json.dumps(MOCK_BING_GROUNDING_RESPONSE).encode()
     return response
 
+async def _search_what_is_litellm(is_async: bool) -> object:
+    if is_async:
+        return await litellm.asearch(query="what is litellm", search_provider="bing_grounding")
+    return litellm.search(query="what is litellm", search_provider="bing_grounding")
+
+
 @pytest.mark.usefixtures("_vcr_outcome_gate")
 class TestBingGroundingSearchTransformation:
     """
@@ -546,6 +552,36 @@ class TestBingGroundingSearchTransformation:
             }
         ]
 
+    @pytest.mark.parametrize("is_async", (False, True))
+    async def test_entra_rejecting_the_default_credential_answers_401_with_only_the_aadsts_reason(
+        self, monkeypatch: pytest.MonkeyPatch, is_async: bool
+    ):
+        from azure.core.exceptions import ClientAuthenticationError
+
+        from litellm.llms.azure import common_utils
+
+        app_client_id: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2"
+        credential = Mock(spec=["get_token"])
+        credential.get_token.side_effect = ClientAuthenticationError(
+            message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being "
+            f"sent in the request is the client secret value, for a secret added to app '{app_client_id}'."
+        )
+        monkeypatch.delenv("BING_GROUNDING_TOKEN")
+        for name in ("AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET", "AZURE_USERNAME", "AZURE_PASSWORD"):
+            monkeypatch.delenv(name, raising=False)
+        monkeypatch.setattr("azure.identity.DefaultAzureCredential", Mock(return_value=credential))
+        common_utils._cached_azure_ad_token_refresh_provider.cache_clear()
+
+        try:
+            with pytest.raises(litellm.AuthenticationError) as excinfo:
+                await _search_what_is_litellm(is_async)
+        finally:
+            common_utils._cached_azure_ad_token_refresh_provider.cache_clear()
+
+        assert excinfo.value.status_code == 401
+        assert "AADSTS7000215" in str(excinfo.value)
+        assert app_client_id not in str(excinfo.value)
+
     @pytest.mark.asyncio
     async def test_bing_grounding_asearch(self):
         with patch(  # test-quality-ok: litellm.asearch has no client injection seam

tests/unit/llms/azure/test_azure_common_utils.py

@@ -1,16 +1,23 @@
+import asyncio
+import io
 import json
 import os
+import sys
 import threading
 import traceback
 import uuid
+from collections.abc import Mapping
 from concurrent.futures import ThreadPoolExecutor
 from queue import SimpleQueue
-from typing import Callable, Final, Optional
+from typing import Awaitable, Callable, Final, Optional
 from urllib.parse import parse_qs
 from unittest.mock import ANY, MagicMock, patch
 
 import httpx
+import openai
 import pytest
+import requests
+from azure.core.pipeline.transport import RequestsTransport
 
 import litellm
 from litellm.llms.azure.common_utils import (
@@ -2500,3 +2507,487 @@ def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(m
     closer.reap()
 
     assert wrapper.is_closed() is True
+
+
+_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2"
+_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500"
+_ENTRA_REJECTIONS: Final = {
+    "wrong_secret": (
+        "AADSTS7000215",
+        "Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent in the "
+        "request is the client secret value, not the client secret ID, for a secret added to app "
+        f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID} Correlation ID: ecfd2c73-ec0a-474c-bba7-205be903159d "
+        "Timestamp: 2026-10-29 07:19:40Z",
+    ),
+    "expired_secret": (
+        "AADSTS7000222",
+        f"Authentication failed: AADSTS7000222: The provided client secret keys for app '{_ENTRA_APP_CLIENT_ID}' "
+        f"are expired. Visit the Azure portal to create new keys for your app. Trace ID: {_ENTRA_TRACE_ID} "
+        "Correlation ID: 354f9b57-5fed-47e4-a67d-89a7c56f6ddc Timestamp: 2026-10-29 07:19:50Z",
+    ),
+}
+
+
+def _raising_token_provider(error: Exception) -> Callable[[], str]:
+    def token_provider() -> str:
+        raise error
+
+    return token_provider
+
+
+def _azure_chat(token_provider: Callable[[], str]) -> object:
+    return litellm.completion(
+        model="azure/gpt-4.1",
+        api_base="https://litellm-qa.openai.azure.com",
+        api_version="2025-04-01-preview",
+        messages=[{"role": "user", "content": "hi"}],
+        azure_ad_token_provider=token_provider,
+        max_retries=0,
+    )
+
+
+def _azure_async_chat(token_provider: Callable[[], str]) -> object:
+    return asyncio.run(
+        litellm.acompletion(
+            model="azure/gpt-4.1",
+            api_base="https://litellm-qa.openai.azure.com",
+            api_version="2025-04-01-preview",
+            messages=[{"role": "user", "content": "hi"}],
+            azure_ad_token_provider=token_provider,
+        )
+    )
+
+
+def _azure_responses(token_provider: Callable[[], str]) -> object:
+    return litellm.responses(
+        model="azure/gpt-4.1",
+        api_base="https://litellm-qa.openai.azure.com",
+        api_version="2025-04-01-preview",
+        input="hi",
+        azure_ad_token_provider=token_provider,
+    )
+
+
+def _azure_ai_decisions(token_provider: Callable[[], str]) -> object:
+    return litellm.systemone(
+        model="azure_ai/decision-1-qa",
+        api_base="https://litellm-qa.services.ai.azure.com",
+        state={"ticket": "The login page returns a 500 error after deploy"},
+        questions={"is_bug": {"type": "noul", "instructions": "Is this a software bug report?"}},
+        azure_ad_token_provider=token_provider,
+    )
+
+
+_AZURE_CALLS: Final = {
+    "chat": _azure_chat,
+    "async_chat": _azure_async_chat,
+    "responses": _azure_responses,
+    "decisions": _azure_ai_decisions,
+}
+
+
+@pytest.fixture
+def no_azure_api_key(monkeypatch):
+    for env_var in (
+        "AZURE_API_KEY",
+        "AZURE_OPENAI_API_KEY",
+        "AZURE_AI_API_KEY",
+        "AZURE_AD_TOKEN",
+        "AZURE_TENANT_ID",
+        "AZURE_CLIENT_ID",
+        "AZURE_CLIENT_SECRET",
+        "AZURE_USERNAME",
+        "AZURE_PASSWORD",
+    ):
+        monkeypatch.delenv(env_var, raising=False)
+    monkeypatch.setattr(litellm, "api_key", None)
+    monkeypatch.setattr(litellm, "azure_key", None)
+
+
+@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys())
+@pytest.mark.parametrize(("aadsts_code", "entra_message"), _ENTRA_REJECTIONS.values(), ids=_ENTRA_REJECTIONS.keys())
+def test_entra_rejecting_the_credential_answers_401_with_only_the_aadsts_reason(
+    no_azure_api_key, call, aadsts_code, entra_message
+):
+    from azure.core.exceptions import ClientAuthenticationError
+
+    with pytest.raises(litellm.AuthenticationError) as excinfo:
+        call(_raising_token_provider(ClientAuthenticationError(message=entra_message)))
+
+    assert excinfo.value.status_code == 401
+    assert aadsts_code in str(excinfo.value)
+    assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value)
+    assert _ENTRA_TRACE_ID not in str(excinfo.value)
+
+
+@pytest.mark.parametrize("api_version", ["2025-04-01-preview", "v1"])
+def test_entra_rejecting_an_async_token_provider_answers_401(no_azure_api_key, api_version):
+    from azure.core.exceptions import ClientAuthenticationError
+
+    aadsts_code, entra_message = _ENTRA_REJECTIONS["wrong_secret"]
+
+    async def token_provider() -> str:
+        raise ClientAuthenticationError(message=entra_message)
+
+    async_token_provider: Final[Callable[[], Awaitable[str]]] = token_provider
+
+    with pytest.raises(litellm.AuthenticationError) as excinfo:
+        asyncio.run(
+            litellm.acompletion(
+                model="azure/gpt-4.1",
+                api_base="https://litellm-qa.openai.azure.com",
+                api_version=api_version,
+                messages=[{"role": "user", "content": "hi"}],
+                azure_ad_token_provider=async_token_provider,
+            )
+        )
+
+    assert excinfo.value.status_code == 401
+    assert aadsts_code in str(excinfo.value)
+    assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value)
+    assert _ENTRA_TRACE_ID not in str(excinfo.value)
+
+
+def _azure_identity_wrapped_token_provider(inner_error: Exception) -> Callable[[], str]:
+    from azure.core.exceptions import ClientAuthenticationError
+
+    def token_provider() -> str:
+        try:
+            raise inner_error
+        except Exception as error:
+            raise ClientAuthenticationError(message=f"Authentication failed: {error}") from error
+
+    return token_provider
+
+
+@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys())
+@pytest.mark.parametrize("failed_leg", ["request", "response"])
+def test_unreachable_entra_endpoint_stays_a_retryable_server_error(no_azure_api_key, call, failed_leg):
+    from azure.core.exceptions import ServiceRequestError, ServiceResponseError
+
+    transport_error: Final = {"request": ServiceRequestError, "response": ServiceResponseError}[failed_leg](
+        "('Unable to connect to proxy', NewConnectionError(\"Failed to establish a new connection: "
+        "[Errno 111] Connection refused\"))"
+    )
+
+    with pytest.raises(openai.APIError) as excinfo:
+        call(_azure_identity_wrapped_token_provider(transport_error))
+
+    assert not isinstance(excinfo.value, litellm.AuthenticationError)
+    assert excinfo.value.status_code != 401
+
+
+_ENTRA_TENANT_ID: Final = "5b3e1c7a-2f4d-4e8b-9a61-0c2d7f8e9b14"
+_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default"
+_ENTRA_OPENID_CONFIGURATION: Final = {
+    "token_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token",
+    "authorization_endpoint": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/authorize",
+    "issuer": f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/v2.0",
+}
+
+
+class _FakeEntraAdapter(requests.adapters.BaseAdapter):
+    def __init__(self, status_code: int | None, body: Mapping[str, object] | str) -> None:
+        super().__init__()
+        self.status_code: Final = status_code
+        self.body: Final = body
+
+    def send(self, request: requests.PreparedRequest, *args: object, **kwargs: object) -> requests.Response:
+        if "openid-configuration" in str(request.url):
+            return self._answer(request, 200, _ENTRA_OPENID_CONFIGURATION)
+        if self.status_code is None:
+            raise requests.exceptions.ConnectionError("Failed to establish a new connection: Connection refused")
+        return self._answer(request, self.status_code, self.body)
+
+    def close(self) -> None:
+        pass
+
+    @staticmethod
+    def _answer(
+        request: requests.PreparedRequest, status_code: int, body: Mapping[str, object] | str
+    ) -> requests.Response:
+        response: Final = requests.Response()
+        response.status_code = status_code
+        response.raw = io.BytesIO((body if isinstance(body, str) else json.dumps(body)).encode())
+        response.headers = requests.structures.CaseInsensitiveDict(
+            {"Content-Type": "application/json", "Retry-After": "0"}
+        )
+        response.url = str(request.url)
+        response.request = request
+        return response
+
+
+def _entra_answering(status_code: int | None, body: Mapping[str, object] | str = "") -> RequestsTransport:
+    session: Final = requests.Session()
+    session.mount("https://", _FakeEntraAdapter(status_code, body))
+    return RequestsTransport(session=session, session_owner=False)
+
+
+def _entra_error(aadsts_code: str, description: str) -> Mapping[str, object]:
+    return {
+        "error": "invalid_client",
+        "error_description": f"{aadsts_code}: {description} Trace ID: {_ENTRA_TRACE_ID}",
+        "error_codes": [int(aadsts_code.removeprefix("AADSTS"))],
+    }
+
+
+def _real_credential_token_provider(credential_kind: str, entra: RequestsTransport) -> Callable[[], str]:
+    from azure.identity import (
+        ChainedTokenCredential,
+        ClientAssertionCredential,
+        ClientSecretCredential,
+        get_bearer_token_provider,
+    )
+
+    client_secret: Final = ClientSecretCredential(
+        _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, "wrong-secret", transport=entra, disable_instance_discovery=True
+    )
+    credentials: Final = {
+        "client_secret": client_secret,
+        "default_azure_credential_chain": ChainedTokenCredential(client_secret),
+        "workload_identity": ClientAssertionCredential(
+            _ENTRA_TENANT_ID, _ENTRA_APP_CLIENT_ID, lambda: "federated-token", transport=entra, retry_total=0
+        ),
+    }
+    return get_bearer_token_provider(credentials[credential_kind], _ENTRA_SCOPE)
+
+
+_REAL_CREDENTIAL_REJECTIONS: Final = {
+    "wrong_secret": ("client_secret", "AADSTS7000215", "Invalid client secret provided."),
+    "wrong_secret_through_a_credential_chain": (
+        "default_azure_credential_chain",
+        "AADSTS7000215",
+        "Invalid client secret provided.",
+    ),
+    "unknown_app_through_workload_identity": (
+        "workload_identity",
+        "AADSTS700016",
+        f"Application with identifier '{_ENTRA_APP_CLIENT_ID}' was not found in the directory.",
+    ),
+    "unknown_federated_subject_through_workload_identity": (
+        "workload_identity",
+        "AADSTS700213",
+        "No matching federated identity record found for presented assertion subject 'system:serviceaccount:qa:proxy'.",
+    ),
+}
+
+
+@pytest.mark.parametrize(
+    ("credential_kind", "aadsts_code", "description"),
+    _REAL_CREDENTIAL_REJECTIONS.values(),
+    ids=_REAL_CREDENTIAL_REJECTIONS.keys(),
+)
+def test_entra_rejection_through_a_real_azure_identity_credential_answers_401(
+    no_azure_api_key, credential_kind, aadsts_code, description
+):
+    entra: Final = _entra_answering(401, _entra_error(aadsts_code, description))
+
+    with pytest.raises(litellm.AuthenticationError) as excinfo:
+        _azure_chat(_real_credential_token_provider(credential_kind, entra))
+
+    assert excinfo.value.status_code == 401
+    assert aadsts_code in str(excinfo.value)
+    assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value)
+    assert _ENTRA_TRACE_ID not in str(excinfo.value)
+
+
+_REAL_CREDENTIAL_NON_REJECTIONS: Final = {
+    "unreachable_through_a_credential_chain": ("default_azure_credential_chain", None, ""),
+    "html_503_through_a_credential_chain": (
+        "default_azure_credential_chain",
+        503,
+        "<html><body>Service Unavailable</body></html>",
+    ),
+    "json_503_through_a_credential_chain": (
+        "default_azure_credential_chain",
+        503,
+        _entra_error("AADSTS90033", "A transient error has occurred."),
+    ),
+    "throttled_client_secret": (
+        "client_secret",
+        429,
+        _entra_error("AADSTS50196", "The server terminated an operation because it encountered a client request loop."),
+    ),
+    "tenant_throttled_client_secret": (
+        "client_secret",
+        429,
+        _entra_error("AADSTS90055", "There are too many incoming requests."),
+    ),
+    "request_budget_exceeded_client_secret": (
+        "client_secret",
+        400,
+        _entra_error("AADSTS90024", "A transient error has occurred. Try again."),
+    ),
+    "retryable_error_client_secret": (
+        "client_secret",
+        400,
+        _entra_error("AADSTS50033", "A transient error has occurred."),
+    ),
+    "external_server_retryable_client_secret": (
+        "client_secret",
+        400,
+        _entra_error("AADSTS90006", "The service is temporarily unavailable."),
+    ),
+    "graph_retryable_client_secret": (
+        "client_secret",
+        400,
+        _entra_error("AADSTS90090", "The service is temporarily unavailable."),
+    ),
+    "token_issuance_error_client_secret": (
+        "client_secret",
+        400,
+        _entra_error("AADSTS50000", "There was an error issuing a token or an issue with our sign-in service."),
+    ),
+    "throttled_without_a_code_through_a_credential_chain": (
+        "default_azure_credential_chain",
+        429,
+        {"error": "temporarily_unavailable", "error_description": "Too many requests."},
+    ),
+    "json_503_through_workload_identity": (
+        "workload_identity",
+        503,
+        _entra_error("AADSTS90033", "A transient error has occurred."),
+    ),
+    "throttled_workload_identity_naming_a_rejection_code": (
+        "workload_identity",
+        429,
+        _entra_error("AADSTS7000215", "Invalid client secret provided."),
+    ),
+}
+
+
+@pytest.mark.parametrize(
+    ("credential_kind", "status_code", "body"),
+    _REAL_CREDENTIAL_NON_REJECTIONS.values(),
+    ids=_REAL_CREDENTIAL_NON_REJECTIONS.keys(),
+)
+def test_entra_outage_or_throttle_through_a_real_azure_identity_credential_stays_retryable(
+    no_azure_api_key, credential_kind, status_code, body
+):
+    with pytest.raises(openai.APIError) as excinfo:
+        _azure_chat(_real_credential_token_provider(credential_kind, _entra_answering(status_code, body)))
+
+    assert not isinstance(excinfo.value, litellm.AuthenticationError)
+    assert excinfo.value.status_code != 401
+
+
+def _azure_chat_with_oidc_token(
+    respx_mock,
+    monkeypatch,
+    entra_status_code: int,
+    entra_body: Mapping[str, object],
+    entra_headers: Mapping[str, str] | None = None,
+):
+    monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token")
+    monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False)
+    respx_mock.post(f"https://login.microsoftonline.com/{_ENTRA_TENANT_ID}/oauth2/v2.0/token").mock(
+        return_value=httpx.Response(entra_status_code, json=entra_body, headers=entra_headers)
+    )
+    return litellm.completion(
+        model="azure/gpt-4.1",
+        api_base="https://litellm-qa.openai.azure.com",
+        api_version="2025-04-01-preview",
+        messages=[{"role": "user", "content": "hi"}],
+        azure_ad_token="oidc/env/LITELLM_QA_FEDERATED_TOKEN",
+        tenant_id=_ENTRA_TENANT_ID,
+        client_id=_ENTRA_APP_CLIENT_ID,
+        max_retries=0,
+    )
+
+
+def test_entra_rejecting_an_oidc_token_answers_401_with_only_the_aadsts_reason(no_azure_api_key, respx_mock, monkeypatch):
+    rejection: Final = _entra_error(
+        "AADSTS700213", f"No matching federated identity record found for app '{_ENTRA_APP_CLIENT_ID}'."
+    )
+
+    with pytest.raises(litellm.AuthenticationError) as excinfo:
+        _azure_chat_with_oidc_token(respx_mock, monkeypatch, 400, rejection)
+
+    assert excinfo.value.status_code == 401
+    assert "AADSTS700213" in str(excinfo.value)
+    assert _ENTRA_APP_CLIENT_ID not in str(excinfo.value)
+    assert _ENTRA_TRACE_ID not in str(excinfo.value)
+
+
+def test_entra_answering_an_oidc_exchange_with_a_transient_code_keeps_its_status(
+    no_azure_api_key, respx_mock, monkeypatch
+):
+    with pytest.raises(openai.APIError) as excinfo:
+        _azure_chat_with_oidc_token(
+            respx_mock, monkeypatch, 400, _entra_error("AADSTS90024", "A transient error has occurred. Try again.")
+        )
+
+    assert excinfo.value.status_code == 400
+    assert not isinstance(excinfo.value, litellm.AuthenticationError)
+
+
+def test_entra_throttling_an_oidc_exchange_keeps_its_retry_after(no_azure_api_key, respx_mock, monkeypatch):
+    with pytest.raises(litellm.RateLimitError) as excinfo:
+        _azure_chat_with_oidc_token(
+            respx_mock,
+            monkeypatch,
+            429,
+            _entra_error("AADSTS90055", "TenantThrottlingError: There are too many pending requests."),
+            {"Retry-After": "7"},
+        )
+
+    assert excinfo.value.response.headers.get("retry-after") == "7"
+
+
+class _EntraServiceUnavailableAnswer:
+    status_code: Final = 503
+    reason: Final = "Service Unavailable"
+
+    def text(self, encoding: str | None = None) -> str:
+        return "<html><body>Service Unavailable</body></html>"
+
+
+@pytest.mark.parametrize("call", _AZURE_CALLS.values(), ids=_AZURE_CALLS.keys())
+@pytest.mark.parametrize("credential_kind", ["client_secret", "managed_identity"])
+def test_entra_answering_a_server_error_stays_a_retryable_server_error(no_azure_api_key, call, credential_kind):
+    from azure.core.exceptions import ClientAuthenticationError
+    from msal.exceptions import MsalServiceError
+
+    token_providers: Final = {
+        "client_secret": _azure_identity_wrapped_token_provider(
+            MsalServiceError("HTTP Error: 503", error=None, error_description=None)
+        ),
+        "managed_identity": _raising_token_provider(
+            ClientAuthenticationError(message="No token received.", response=_EntraServiceUnavailableAnswer())
+        ),
+    }
+
+    with pytest.raises(openai.APIError) as excinfo:
+        call(token_providers[credential_kind])
+
+    assert not isinstance(excinfo.value, litellm.AuthenticationError)
+    assert excinfo.value.status_code != 401
+
+
+def test_token_provider_error_without_azure_core_installed_keeps_the_generic_token_error(monkeypatch):
+    monkeypatch.setitem(sys.modules, "azure.core.exceptions", None)
+
+    with pytest.raises(RuntimeError, match="Failed to get Azure AD token: token service down"):
+        get_azure_ad_token(
+            GenericLiteLLMParams(azure_ad_token_provider=_raising_token_provider(ValueError("token service down")))
+        )
+
+
+def test_image_generation_with_an_api_key_ignores_a_non_callable_azure_ad_token_provider():
+    import respx
+
+    with respx.mock:
+        image_route: Final = respx.post(url__startswith="https://litellm-qa.openai.azure.com").mock(
+            return_value=httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGk="}]})
+        )
+        response: Final = litellm.image_generation(
+            model="azure/gpt-image-1",
+            prompt="a red dot",
+            api_base="https://litellm-qa.openai.azure.com",
+            api_version="2025-04-01-preview",
+            api_key="configured-key",
+            azure_ad_token_provider="not-a-callable",
+        )
+
+    assert image_route.calls.last.request.headers["api-key"] == "configured-key"
+    assert response.data[0].b64_json == "aGk="

tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py

@@ -171,6 +171,79 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_
     )
 
 
+@pytest.mark.asyncio
+async def test_entra_rejecting_the_deployment_oidc_token_becomes_a_401_error_response_without_counting(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token")
+    monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False)
+    with respx.mock:
+        respx.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock(
+            return_value=httpx.Response(
+                401,
+                json={
+                    "error": "invalid_client",
+                    "error_description": "AADSTS700211: No matching federated identity record found for presented "
+                    "assertion issuer 'https://oidc.litellm-qa.invalid'. Trace ID: 0f5c2b1e",
+                    "error_codes": [700211],
+                },
+            )
+        )
+        count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6}))
+        result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
+            model_to_use="claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={
+                "litellm_params": {
+                    "api_base": API_BASE,
+                    "tenant_id": "qa-tenant",
+                    "client_id": "qa-client",
+                    "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN",
+                }
+            },
+            request_model="azure_ai/claude-sonnet-4-5",
+        )
+
+    assert result is not None
+    assert (result.error, result.status_code) == (True, 401)
+    assert "AADSTS700211" in (result.error_message or "")
+    assert not count_route.called
+
+
+@pytest.mark.asyncio
+async def test_entra_throttling_the_deployment_oidc_token_exchange_keeps_its_429(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token")
+    monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False)
+    with respx.mock:
+        respx.post("https://login.microsoftonline.com/throttled-tenant/oauth2/v2.0/token").mock(
+            return_value=httpx.Response(429, json={"error": "temporarily_unavailable", "error_codes": [90055]})
+        )
+        count_route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6}))
+        result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
+            model_to_use="claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hi"}],
+            contents=None,
+            deployment={
+                "litellm_params": {
+                    "api_base": API_BASE,
+                    "tenant_id": "throttled-tenant",
+                    "client_id": "qa-client",
+                    "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN",
+                }
+            },
+            request_model="azure_ai/claude-sonnet-4-5",
+        )
+
+    assert result is not None
+    assert (result.error, result.status_code) == (True, 429)
+    assert not count_route.called
+
+
 @pytest.mark.asyncio
 async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
     monkeypatch: pytest.MonkeyPatch,

tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py

@@ -254,6 +254,57 @@ def test_agent_entra_values_resolve_os_environ_references(monkeypatch):
     assert get_azure_ai_agent_entra_token({"azure_ad_token": "os.environ/FOUNDRY_AGENT_AD_TOKEN"}) == "token-from-env"
 
 
+_AGENT_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2"
+_AGENT_CREDENTIALS: Final = {
+    "service_principal": (
+        "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id",
+        {"tenant_id": "tenant", "client_id": _AGENT_APP_CLIENT_ID, "client_secret": "wrong-secret"},
+    ),
+    "username_password": (
+        "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password",
+        {"client_id": _AGENT_APP_CLIENT_ID, "azure_username": "agent@example.com", "azure_password": "wrong"},
+    ),
+}
+
+
+def _raising(error: Exception):
+    def token() -> str:
+        raise error
+
+    return token
+
+
+@pytest.mark.parametrize(("token_fetch", "agent_params"), _AGENT_CREDENTIALS.values(), ids=_AGENT_CREDENTIALS.keys())
+def test_agent_entra_rejection_raises_401_with_only_the_aadsts_reason(token_fetch, agent_params):
+    from azure.core.exceptions import ClientAuthenticationError
+
+    rejection: Final = ClientAuthenticationError(
+        message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being sent "
+        f"in the request is the client secret value, for a secret added to app '{_AGENT_APP_CLIENT_ID}'."
+    )
+
+    with patch(token_fetch, return_value=_raising(rejection)):  # test-quality-ok: stands in for the Entra token endpoint answering a wrong credential, the shape azure-identity raises
+        with pytest.raises(litellm.AuthenticationError) as excinfo:
+            get_azure_ai_agent_entra_token(agent_params)
+
+    assert excinfo.value.status_code == 401
+    assert "AADSTS7000215" in str(excinfo.value)
+    assert _AGENT_APP_CLIENT_ID not in str(excinfo.value)
+
+
+def test_agent_entra_endpoint_unreachable_is_not_an_auth_error():
+    from azure.core.exceptions import ClientAuthenticationError
+
+    unreachable: Final = ClientAuthenticationError(message="Authentication failed: Connection refused")
+    token_fetch, agent_params = _AGENT_CREDENTIALS["service_principal"]
+
+    with patch(token_fetch, return_value=_raising(unreachable)):  # test-quality-ok: stands in for an unreachable Entra token endpoint, the shape azure-identity raises
+        with pytest.raises(ClientAuthenticationError) as excinfo:
+            get_azure_ai_agent_entra_token(agent_params)
+
+    assert not isinstance(excinfo.value, litellm.AuthenticationError)
+
+
 def test_agent_entra_token_failure_names_the_credential_fields():
     with pytest.raises(ValueError, match="client_secret"):
         get_azure_ai_agent_entra_token({"azure_scope": "https://ai.azure.com/.default"})

tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py

@@ -1394,6 +1394,7 @@ class _ScriptedGuardrail(CustomGuardrail):
         self.skip_system_message_in_guardrail = skip_system
         self.scan_only_tool_results = scan_only_tool_results
         self.received: tuple[GenericGuardrailAPIInputs, ...] = ()
+        self.request_data: tuple[Mapping[str, object], ...] = ()
 
     async def apply_guardrail(
         self,
@@ -1403,6 +1404,7 @@ class _ScriptedGuardrail(CustomGuardrail):
         logging_obj: LiteLLMLoggingObj | None = None,
     ) -> Mapping[str, object]:
         self.received = (*self.received, copy.deepcopy(inputs))
+        self.request_data = (*self.request_data, request_data)
         return self.reply(inputs)
 
 
@@ -1429,16 +1431,30 @@ def _converse_request(body: dict) -> dict:
 
 
 async def _sent_to_bedrock(data: dict) -> dict:
-    """The Converse body Bedrock receives when the passthrough route sends ``data``."""
+    """The body Bedrock receives when the passthrough route sends ``data``."""
     received: list[httpx.Request] = []
 
+    action: Final = data["endpoint"].rsplit("/", 1)[-1]
+
     def respond(request: httpx.Request) -> httpx.Response:
         received.append(request)
-        reply = {
-            "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
-            "stopReason": "end_turn",
-            "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
-        }
+        reply: Final = (
+            {
+                "id": "msg_bdrk_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-haiku-4-5",
+                "content": [{"type": "text", "text": "ok"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            }
+            if action == "invoke"
+            else {
+                "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
+                "stopReason": "end_turn",
+                "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
+            }
+        )
         return httpx.Response(200, json=reply, request=request)
 
     router = litellm.Router(
@@ -1457,7 +1473,7 @@ async def _sent_to_bedrock(data: dict) -> dict:
     await router.allm_passthrough_route(
         model="bedrock-model",
         method="POST",
-        endpoint="/model/bedrock-model/converse",
+        endpoint=f"/model/bedrock-model/{action}",
         data=data["data"],
         client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)),
     )
@@ -3843,6 +3859,260 @@ class TestInvokeStreamGuardrailRelay:
         assert rewritten_key == received_key
 
 
+def _anthropic_invoke_data(
+    action: str = "invoke",
+) -> dict[str, object]:  # mutable-ok: tests write into the request the handler gets
+    return {
+        "endpoint": f"model/anthropic.claude-sonnet/{action}",
+        "custom_llm_provider": "bedrock",
+        "model": "anthropic.claude-sonnet",
+        "litellm_call_id": "call-1",
+        "data": {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 512,
+            "system": "You are a travel assistant. Card on file 4111 1111 1111 1111",
+            "tools": [
+                {
+                    "name": "find_trips",
+                    "description": "Search trips",
+                    "input_schema": {"type": "object", "properties": {"port": {"type": "string"}}},
+                }
+            ],
+            "messages": [
+                {"role": "user", "content": "Pay with 4111 1111 1111 1111"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {"type": "tool_use", "id": "toolu_1", "name": "find_trips", "input": {"port": "Miami"}}
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card 4111 1111 1111 1111 charged"}
+                    ],
+                },
+            ],
+        },
+    }
+
+
+def _masked_anthropic_invoke_body() -> Mapping[str, object]:
+    body: Final = _anthropic_invoke_data()["data"]
+    return {
+        **body,
+        "system": "You are a travel assistant. Card on file [CARD] 1111 1111 1111",
+        "messages": [
+            {"role": "user", "content": "Pay with [CARD] 1111 1111 1111"},
+            body["messages"][1],
+            {
+                "role": "user",
+                "content": [
+                    {"type": "tool_result", "tool_use_id": "toolu_1", "content": "Card [CARD] 1111 1111 1111 charged"}
+                ],
+            },
+        ],
+    }
+
+
+def _rows_masked(inputs: GenericGuardrailAPIInputs) -> Mapping[str, object]:
+    rows = [
+        {**row, "content": row["content"].replace("4111", "[CARD]")} if isinstance(row.get("content"), str) else row
+        for row in inputs["structured_messages"]
+    ]
+    return {"structured_messages": rows}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
+async def test_anthropic_invoke_masks_are_written_back_into_the_body(action: str) -> None:
+    guardrail = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)})
+
+    result = await BedrockPassthroughGuardrailHandler().process_input_messages(
+        _anthropic_invoke_data(action), guardrail
+    )
+
+    assert result["data"] == _masked_anthropic_invoke_body()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_masks_reach_bedrock() -> None:
+    data: Final = _anthropic_invoke_data()
+
+    await BedrockPassthroughGuardrailHandler().process_input_messages(
+        data, _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)})
+    )
+
+    assert await _sent_to_bedrock(data) == _masked_anthropic_invoke_body()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_guardrail_gets_roles_tools_and_the_body_at_the_top_level_of_request_data() -> None:
+    data: Final = _anthropic_invoke_data()
+    guardrail = _replying({})
+
+    await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
+
+    inputs = guardrail.received[-1]
+    assert [row["role"] for row in inputs["structured_messages"]] == ["system", "user", "assistant", "tool"]
+    assert [tool["function"]["name"] for tool in inputs["tools"]] == ["find_trips"]
+    assert inputs["model"] == data["model"]
+    request_data: Final = guardrail.request_data[-1]
+    assert request_data["messages"] is data["data"]["messages"]
+    assert request_data["endpoint"] == data["endpoint"]
+    assert request_data["litellm_call_id"] == data["litellm_call_id"]
+    assert "data" not in request_data
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_structured_message_rewrites_are_written_back() -> None:
+    result = await BedrockPassthroughGuardrailHandler().process_input_messages(
+        _anthropic_invoke_data(), _ScriptedGuardrail(_rows_masked)
+    )
+
+    assert result["data"]["messages"][0] == {
+        "role": "user",
+        "content": [{"type": "text", "text": "Pay with [CARD] 1111 1111 1111"}],
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model", ["anthropic.claude-sonnet", "prod-model", ""])
+async def test_anthropic_invoke_structured_rewrite_keeps_cache_control_whatever_the_model_is_called(
+    model: str,
+) -> None:
+    cached: Final = {"type": "ephemeral"}
+    data: Final = _anthropic_invoke_data()
+    data["model"] = model
+    data["data"]["system"] = [{"type": "text", "text": "Card on file 4111 1111 1111 1111", "cache_control": cached}]
+    data["data"]["tools"] = [
+        {
+            "name": "find_trips",
+            "description": "Search trips",
+            "input_schema": {"type": "object", "properties": {}},
+            "cache_control": cached,
+        }
+    ]
+    data["data"]["messages"] = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Pay with 4111 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}}
+            ],
+        }
+    ]
+
+    result = await BedrockPassthroughGuardrailHandler().process_input_messages(
+        data,
+        _ScriptedGuardrail(
+            lambda inputs: {
+                "structured_messages": json.loads(json.dumps(inputs["structured_messages"]).replace("4111", "[CARD]")),
+                "tools": json.loads(json.dumps(inputs["tools"])),
+            }
+        ),
+    )
+
+    assert result["data"]["system"] == [
+        {"type": "text", "text": "Card on file [CARD] 1111 1111 1111", "cache_control": cached}
+    ]
+    assert result["data"]["messages"][0]["content"] == [
+        {"type": "text", "text": "Pay with [CARD] 1111 1111 1111", "cache_control": {**cached, "ttl": "1h"}}
+    ]
+    assert result["data"]["tools"][0]["cache_control"] == cached
+
+
+@pytest.mark.asyncio
+async def test_invoke_body_without_anthropic_version_is_scanned_whole() -> None:
+    data: Final = {**_anthropic_invoke_data(), "data": {"messages": [{"role": "user", "content": "hello"}]}}
+    guardrail = _replying({})
+
+    await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
+
+    assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"]
+    assert "structured_messages" not in guardrail.received[-1]
+
+
+_PLAIN_TEXT_DOCUMENT: Final = {
+    "type": "document",
+    "source": {"type": "text", "media_type": "text/plain", "data": "Card 4111 1111 1111 1111"},
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "content",
+    [
+        pytest.param([{"type": "text", "text": "Summarize this"}, _PLAIN_TEXT_DOCUMENT], id="document"),
+        pytest.param(
+            [
+                {
+                    "type": "search_result",
+                    "source": "https://example.com/cards",
+                    "title": "Cards",
+                    "content": [{"type": "text", "text": "Card 4111 1111 1111 1111"}],
+                }
+            ],
+            id="search_result",
+        ),
+        pytest.param(
+            [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [_PLAIN_TEXT_DOCUMENT]}],
+            id="document_in_tool_result",
+        ),
+    ],
+)
+async def test_anthropic_invoke_body_with_blocks_the_messages_handler_does_not_scan_is_scanned_whole(content: Sequence[object]) -> None:
+    invoke: Final = _anthropic_invoke_data()
+    data: Final = {**invoke, "data": {**invoke["data"], "messages": [{"role": "user", "content": content}]}}
+    guardrail = _replying({})
+
+    await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
+
+    assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"]
+    assert "structured_messages" not in guardrail.received[-1]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "messages",
+    [
+        pytest.param([{"content": "hi"}], id="message_without_role"),
+        pytest.param([{"role": "user", "content": ["hi"]}], id="bare_string_in_content"),
+        pytest.param([{"role": "user", "content": None}], id="null_content"),
+        pytest.param(["hi"], id="message_is_a_string"),
+    ],
+)
+async def test_anthropic_invoke_malformed_messages_are_scanned_whole_and_left_for_bedrock_to_reject(
+    messages: Sequence[object],
+) -> None:
+    invoke: Final = _anthropic_invoke_data()
+    data: Final = {**invoke, "data": {**invoke["data"], "messages": messages}}
+    guardrail = _replying({})
+
+    result: Final = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
+
+    assert json.loads(guardrail.received[-1]["texts"][0])["data"] == data["data"]
+    assert result["data"]["messages"] == messages
+
+
+@pytest.mark.asyncio
+async def test_anthropic_invoke_response_masks_are_written_back() -> None:
+    response = {
+        "id": "msg_bdrk_1",
+        "type": "message",
+        "role": "assistant",
+        "content": [{"type": "text", "text": "Charged 4111 1111 1111 1111"}],
+        "stop_reason": "end_turn",
+    }
+
+    result: Final = await BedrockPassthroughGuardrailHandler().process_output_response(
+        response=response,
+        guardrail_to_apply=_ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}),
+        request_data=_anthropic_invoke_data(),
+    )
+
+    assert result == {**response, "content": [{"type": "text", "text": "Charged [CARD] 1111 1111 1111"}]}
+
+
 def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None:
     event: Final = {
         "type": "content_block_delta",

tests/unit/llms/github_copilot/test_github_host.py

@@ -54,6 +54,7 @@ def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.Monke
     assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code"
     assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token"
     assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token"
+    assert endpoints.copilot_user_url == "https://api.octo.ghe.com/copilot_internal/user"
     assert endpoints.user_url == "https://api.octo.ghe.com/user"
     assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com"
     assert endpoints.client_id == "custom-client"

tests/unit/llms/github_copilot/test_per_user_auth.py

@@ -307,6 +307,71 @@ async def test_aexchange_401_raises_and_evicts():
         )
 
 
+ACCOUNT_URL = "https://api.github.com/copilot_internal/user"
+_CUSTOM_CLIENT_ID: Final = "Ov23liCustomOrgApp"
+_ACCOUNT_API_BASE: Final = "https://api.business.githubcopilot.com"
+
+
+def _account_payload(access_type_sku="copilot_enterprise_seat", chat_enabled=True, api_base=_ACCOUNT_API_BASE):
+    return {"access_type_sku": access_type_sku, "chat_enabled": chat_enabled, "endpoints": {"api": api_base}}
+
+
+def _account_responder(payload=None, status=200, model_payload=None):
+    def respond(request: httpx.Request) -> httpx.Response:
+        if str(request.url) == ACCOUNT_URL:
+            return httpx.Response(status, json=payload or _account_payload(), request=request)
+        if model_payload is not None and request.url.host.endswith("githubcopilot.com"):
+            return httpx.Response(200, json=model_payload, request=request)
+        pytest.fail(f"unexpected request to {request.url}")
+
+    return respond
+
+
+@pytest.fixture
+def custom_client_id(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", _CUSTOM_CLIENT_ID)
+
+
+@pytest.mark.asyncio
+async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_swap(custom_client_id):
+    endpoints: Final = github_host_endpoints("github.com")
+    with github_http(_account_responder()) as requests:
+        first = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints)
+        second = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints)
+
+    assert [str(r.url) for r in requests] == [ACCOUNT_URL]
+    assert requests[0].headers["Authorization"] == "token gho_org_app"
+    assert first.token == second.token == "gho_org_app"
+    assert first.api_base == _ACCOUNT_API_BASE
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        _account_payload(access_type_sku="no_access", chat_enabled=False),
+        _account_payload(chat_enabled=False),
+        _account_payload(access_type_sku="no_access"),
+    ],
+)
+def test_custom_client_id_without_a_copilot_seat_raises_reconnect(custom_client_id, payload):
+    with github_http(_account_responder(payload)):
+        with pytest.raises(CallerCredentialAuthenticationError):
+            exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
+
+
+@pytest.mark.parametrize("status", [401, 403, 404])
+def test_custom_client_id_rejected_account_lookup_raises_reconnect(custom_client_id, status):
+    with github_http(_account_responder(status=status)):
+        with pytest.raises(CallerCredentialAuthenticationError):
+            exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
+
+
+def test_custom_client_id_account_api_base_must_stay_on_copilot_hosts(custom_client_id):
+    with github_http(_account_responder(_account_payload(api_base="https://evil.example"))):
+        session = exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
+    assert session.api_base == "https://api.githubcopilot.com"
+
+
 def _kwargs_with_connection(
     token: str = "gho_user_token",
     github_host: str | None = "github.com",
@@ -740,6 +805,25 @@ def test_e2e_embedding_per_user_caller_authorization_cannot_override_session(per
     _assert_session_token_on_the_wire(requests)
 
 
+def test_e2e_custom_client_id_sends_the_github_token_to_the_accounts_copilot_host(
+    per_user_credential, custom_client_id
+):
+    with (
+        github_http(_account_responder(model_payload=_embedding_response_payload())) as requests,
+        _no_authenticator(),
+    ):
+        litellm.embedding(
+            model="github_copilot/text-embedding-3-small",
+            input=["hi"],
+            litellm_credential_name="copilot-cred",
+            secret_fields=_secret_fields("user-a", "gho_org_app"),
+        )
+    assert EXCHANGE_URL not in [str(r.url) for r in requests]
+    upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")]
+    assert upstream[0].url.host == httpx.URL(_ACCOUNT_API_BASE).host
+    assert upstream[0].headers["authorization"] == "Bearer gho_org_app"
+
+
 def test_e2e_responses_per_user_caller_authorization_cannot_override_session(per_user_credential):
     with (
         github_http(_copilot_surface_responder(_responses_payload())) as requests,

tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py

@@ -1268,6 +1268,153 @@ async def test_task_method_failure_hook_uses_enriched_request_data():
     assert failure_data.get("agent_id") == "test-agent"
 
 
+_ENTRA_APP_CLIENT_ID: Final = "0d08e61f-6184-4980-9cef-a2e9128f4cb2"
+_ENTRA_TRACE_ID: Final = "f0c149d3-ef20-4194-8b81-c81a5e002500"
+
+
+class _JsonRpcError(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    code: int
+    message: str
+
+
+class _JsonRpcErrorBody(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    error: _JsonRpcError
+
+
+async def _send_message_to_entra_agent(
+    litellm_params: Mapping[str, str], *patches: AbstractContextManager[object]
+) -> tuple[int, _JsonRpcError]:
+    agent = _make_agent_mock()
+    agent.litellm_params = dict(litellm_params)
+    mock_request = _make_request_mock(
+        "message/send",
+        {"message": {"role": "user", "parts": [{"kind": "text", "text": "Hello"}], "messageId": "msg-1"}},
+    )
+    mock_proxy_logging = MagicMock()
+    mock_proxy_logging.pre_call_hook = AsyncMock(
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+    )
+    mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
+
+    with ExitStack() as stack:
+        for p in (*_base_patches(agent), *patches):
+            stack.enter_context(p)
+        stack.enter_context(
+            patch(  # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use
+                "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging
+            )
+        )
+
+        from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
+
+        response = await invoke_agent_a2a(
+            agent_id="test-agent",
+            request=mock_request,
+            fastapi_response=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"),
+        )
+
+    return response.status_code, _JsonRpcErrorBody.model_validate_json(response.body).error
+
+
+@pytest.mark.asyncio
+async def test_entra_rejecting_the_agent_credential_returns_401_with_only_the_aadsts_reason():
+    from azure.core.exceptions import ClientAuthenticationError
+
+    def rejected_token() -> str:
+        raise ClientAuthenticationError(
+            message="Authentication failed: AADSTS7000215: Invalid client secret provided. Ensure the secret being "
+            f"sent in the request is the client secret value, for a secret added to app '{_ENTRA_APP_CLIENT_ID}'."
+        )
+
+    status_code, error = await _send_message_to_entra_agent(
+        {"tenant_id": "tenant", "client_id": _ENTRA_APP_CLIENT_ID, "client_secret": "wrong-secret"},
+        patch(  # test-quality-ok: stands in for the Entra token endpoint answering a wrong secret, the shape azure-identity raises
+            "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=rejected_token
+        ),
+    )
+
+    assert status_code == 401
+    assert "AADSTS7000215" in error.message
+    assert _ENTRA_APP_CLIENT_ID not in error.message
+
+
+@pytest.mark.asyncio
+async def test_entra_rejecting_the_agent_oidc_token_returns_401_with_only_the_aadsts_reason(monkeypatch, respx_mock):
+    monkeypatch.setenv("LITELLM_QA_FEDERATED_TOKEN", "federated-token")
+    monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False)
+    respx_mock.post("https://login.microsoftonline.com/qa-tenant/oauth2/v2.0/token").mock(
+        return_value=httpx.Response(
+            400,
+            json={
+                "error": "invalid_client",
+                "error_description": "AADSTS700213: No matching federated identity record found for app "
+                f"'{_ENTRA_APP_CLIENT_ID}'. Trace ID: {_ENTRA_TRACE_ID}",
+            },
+        )
+    )
+
+    status_code, error = await _send_message_to_entra_agent(
+        {
+            "tenant_id": "qa-tenant",
+            "client_id": _ENTRA_APP_CLIENT_ID,
+            "azure_ad_token": "oidc/env/LITELLM_QA_FEDERATED_TOKEN",
+        }
+    )
+
+    assert status_code == 401
+    assert error.code == -32000
+    assert "AADSTS700213" in error.message
+    assert _ENTRA_APP_CLIENT_ID not in error.message
+    assert _ENTRA_TRACE_ID not in error.message
+
+
+@pytest.mark.asyncio
+async def test_upstream_authentication_error_without_an_entra_rejection_stays_an_internal_error():
+    agent = _make_agent_mock()
+    mock_handler = MagicMock()
+    mock_handler.post = AsyncMock(
+        side_effect=litellm.AuthenticationError(
+            message="Incorrect API key provided.", llm_provider="openai", model="gpt-4.1"
+        )
+    )
+    mock_proxy_logging = MagicMock()
+    mock_proxy_logging.pre_call_hook = AsyncMock(
+        side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
+    )
+    mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
+
+    with ExitStack() as stack:
+        for p in _base_patches(agent):
+            stack.enter_context(p)
+        stack.enter_context(
+            patch(  # test-quality-ok: stands in for the upstream agent rejecting the provider key it was sent
+                "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=mock_handler
+            )
+        )
+        stack.enter_context(
+            patch(  # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook tests use
+                "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging
+            )
+        )
+
+        from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
+
+        response = await invoke_agent_a2a(
+            agent_id="test-agent",
+            request=_make_request_mock("tasks/get", {"id": "task-1"}),
+            fastapi_response=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"),
+        )
+
+    assert response.status_code == 500
+    assert json.loads(response.body.decode())["error"]["code"] == -32603
+
+
 @pytest.mark.asyncio
 async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400():
     from litellm.proxy._types import UserAPIKeyAuth