Merge remote-tracking branch 'upstream/main' into litellm_anthropic-precall-schema-count

.circleci/scripts/run_integration.sh

@@ -201,7 +201,7 @@ if [ "$isolated_workers" = false ]; then
   curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
     -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
     -d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
-  if [ "$suite" = management ] || [ "$suite" = mcp ]; then
+  if [ "$suite" = management ] || [ "$suite" = mcp ] || [ "$suite" = accounting ]; then
     export INTEGRATION_PEER_URL=http://127.0.0.1:4001
     start_proxy 4001 peer.log
     peer_pid="$launched_pid"

.github/CODEOWNERS

@@ -1,2 +1,2 @@
-/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
-/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
+/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
+/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri

.github/workflows/test-unit.yml

@@ -346,10 +346,12 @@ jobs:
               tests/unit/llms/tinyfish
               tests/unit/llms/together_ai
               tests/unit/llms/tool_loop
+              tests/unit/llms/topaz
               tests/unit/llms/triton
               tests/unit/llms/v0
               tests/unit/llms/valkey
               tests/unit/llms/vercel_ai_gateway
+              tests/unit/llms/vllm
               tests/unit/llms/volcengine
               tests/unit/llms/voyage
               tests/unit/llms/wandb
@@ -416,6 +418,7 @@ jobs:
               tests/unit/rag
               tests/unit/realtime_api
               tests/unit/rerank_api
+              tests/unit/search
               tests/unit/vector_stores
               tests/unit/videos
             python-version: "3.12"

Makefile

@@ -290,7 +290,7 @@ test-rust-extension:
 	"$$temporary/venv/bin/python" -I -m mypy.stubtest \
 		--mypy-config-file tests/unit/rust_bridge/stubtest.ini \
 		litellm.rust_bridge._native && \
-	LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
+	LITELLM_LOCAL_MODEL_COST_MAP=True \
 	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
 
 rust-sqlx-prepare:

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -24,7 +24,7 @@ from uuid import NAMESPACE_URL, uuid5
 
 import httpx
 from fastapi import HTTPException
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 from typing_extensions import ReadOnly, Unpack
 
 import litellm
@@ -54,6 +54,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
     can_access_resource,
     resolve_resource_owner_id,
 )
+from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
 from litellm.proxy._types import (
     CallTypes,
     LiteLLM_ManagedFileTable,
@@ -102,6 +103,7 @@ from litellm.types.utils import (
     CallTypesLiteral,
     LiteLLMBatch,
     LiteLLMFineTuningJob,
+    LlmProviders,
     LLMResponseTypes,
     SpecialEnums,
 )
@@ -194,6 +196,32 @@ def _is_transient_file_retrieve_error(error: Exception) -> bool:
     return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
 
 
+_MEASURED_SIZE_SAVE_ATTEMPTS: Final = 3
+
+
+def _stored_json_text(raw_file_object: object) -> str:
+    """JSON text of a stored file object exactly as read, for a write that must not race a newer one."""
+    if isinstance(raw_file_object, BaseModel):
+        return raw_file_object.model_dump_json()
+    return json.dumps(raw_file_object)
+
+
+def _with_measured_size(file_object: OpenAIFileObject, size_bytes: int | None) -> OpenAIFileObject:
+    """Fill a file object whose provider reported no size (0 bytes) with a size measured elsewhere."""
+    if not size_bytes or file_object.bytes:
+        return file_object
+    return file_object.model_copy(update={"bytes": size_bytes})
+
+
+def _over_stored_object(file_object: OpenAIFileObject, stored_object: OpenAIFileObject | None) -> OpenAIFileObject:
+    """Refreshed details for a stored row, keeping its size when none is reported and its stored created_at."""
+    if stored_object is None:
+        return file_object
+    return _with_measured_size(file_object, stored_object.bytes).model_copy(
+        update={"created_at": stored_object.created_at}
+    )
+
+
 def _proxy_llm_router() -> Router | None:
     import litellm.proxy.proxy_server as proxy_server_module
 
@@ -252,6 +280,23 @@ def _has_provider_file_retrieve_route(
     return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None
 
 
+def _is_xai_batch_output(
+    *,
+    llm_router: Router | None,
+    model_id: str | None,
+    model_name: str | None,
+    provider_file_id: str,
+) -> bool:
+    """Whether a batch output row holds an xAI batch's results, from its batch id and its deployment's provider."""
+    if not is_xai_batch_results_id(provider_file_id):
+        return False
+    router_credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
+    provider: Final = (router_credentials or {}).get("custom_llm_provider") or _model_name_file_retrieve_provider(
+        model_name
+    )
+    return provider == LlmProviders.XAI.value
+
+
 class _ManagedFileRow(Protocol):
     unified_file_id: str
     file_object: OpenAIFileObject
@@ -845,6 +890,57 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             litellm_parent_otel_span=None,
         )
 
+    async def _save_xai_refreshed_file_object(
+        self,
+        stored: LiteLLM_ManagedFileTable,
+        file_object: OpenAIFileObject,
+    ) -> OpenAIFileObject | None:
+        """Save an xAI batch output's refreshed details, keeping a known size, only over the row as read; drop the cached copy if a newer write won."""
+        stored_object: Final = _parse_managed_file_object(stored.file_object, stored.unified_file_id)
+        sized_file_object: Final = _over_stored_object(file_object, stored_object)
+        if not await ManagedFileRepository(self.prisma_client).update_file_object(
+            stored.unified_file_id,
+            sized_file_object,
+            if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None,
+        ):
+            await self.internal_usage_cache.async_set_cache(
+                key=stored.unified_file_id, value=None, litellm_parent_otel_span=None
+            )
+            return None
+        await self.internal_usage_cache.async_set_cache(
+            key=stored.unified_file_id,
+            value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(),
+            litellm_parent_otel_span=None,
+        )
+        return sized_file_object
+
+    async def _save_batch_output_file_object(
+        self,
+        stored: LiteLLM_ManagedFileTable,
+        file_object: OpenAIFileObject,
+        size_bytes: int | None,
+    ) -> None:
+        """Save an xAI batch output's details; after losing to a refresh, size the refreshed details instead of replacing them."""
+        current_row: LiteLLM_ManagedFileTable | None = stored  # rebind-ok: the row as last read
+        details: OpenAIFileObject = file_object  # rebind-ok: a concurrent refresh's details replace the job's copy
+        for _ in range(_MEASURED_SIZE_SAVE_ATTEMPTS):
+            if current_row is None:
+                return
+            saved = await self._save_xai_refreshed_file_object(  # rebind-ok: per attempt
+                current_row, _with_measured_size(details, size_bytes)
+            )
+            if saved is not None or not size_bytes:
+                return
+            latest_row = await ManagedFileRepository(  # rebind-ok: per attempt
+                self.prisma_client, use_writer=True
+            ).table.find_first(where={"unified_file_id": stored.unified_file_id})
+            current_row = (
+                None if latest_row is None else LiteLLM_ManagedFileTable.model_validate(latest_row.model_dump())
+            )
+            latest_object = None if current_row is None else current_row.file_object  # rebind-ok: per attempt
+            if latest_object is not None and not latest_object.litellm_details_fallback:
+                details = latest_object
+
     async def store_batch_output_file(
         self,
         *,
@@ -885,6 +981,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             )
             return
         if stored_object is not None and not stored_object.litellm_details_fallback:
+            if (
+                stored_file is not None
+                and size_bytes
+                and not stored_object.bytes
+                and _is_xai_batch_output(
+                    llm_router=_proxy_llm_router(),
+                    model_id=model_id,
+                    model_name=model_name,
+                    provider_file_id=provider_file_id,
+                )
+            ):
+                await self._save_batch_output_file_object(stored_file, stored_object, size_bytes)
             return
 
         fallback_written_recently: Final = (
@@ -924,13 +1032,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             )
         )
 
+        xai_batch_output: Final = _is_xai_batch_output(
+            llm_router=_proxy_llm_router(),
+            model_id=model_id,
+            model_name=model_name,
+            provider_file_id=provider_file_id,
+        )
+        if stored_file is not None and xai_batch_output:
+            await self._save_batch_output_file_object(stored_file, file_object, size_bytes)
+            return
         if stored_file is not None:
             await self._save_refreshed_file_object(stored_file, file_object)
             return
 
         await self.store_unified_file_id(
             file_id=unified_file_id,
-            file_object=file_object,
+            file_object=_with_measured_size(file_object, size_bytes) if xai_batch_output else file_object,
             litellm_parent_otel_span=litellm_parent_otel_span,
             model_mappings={model_id: provider_file_id} if model_id else {},
             user_api_key_dict=owner,
@@ -1943,8 +2060,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     )
                     if refreshed_file_object is None:
                         return _public_file_object(file_object, file_id)
-                    await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
-                    return _public_file_object(refreshed_file_object, file_id)
+                    if not _is_xai_batch_output(
+                        llm_router=llm_router if llm_router is not None else _proxy_llm_router(),
+                        model_id=model_id,
+                        model_name=model_id,
+                        provider_file_id=provider_file_id,
+                    ):
+                        await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
+                        return _public_file_object(refreshed_file_object, file_id)
+                    saved_file_object: Final = await self._save_xai_refreshed_file_object(
+                        stored_file_object, refreshed_file_object
+                    )
+                    return _public_file_object(
+                        saved_file_object or _over_stored_object(refreshed_file_object, file_object), file_id
+                    )
                 except Exception as error:
                     verbose_logger.warning(
                         "Failed to refresh batch file object for "

helm/litellm/README.md

@@ -5,3 +5,21 @@ This chart deploys separate gateway, backend and UI services. Configure Lens thr
 Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
 
 Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for the complete values and verification flow. It also links the GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
+
+For an ingress with separate component hostnames, set all three `ingress.componentHosts` values, leave `ingress.host` and `ingress.extraPaths` empty, configure `ui.backendUrl` as the backend's public origin and `backend.gatewayUrl` as the gateway's public origin, and set `ui.publicUrl` to the browser-facing UI origin. The UI Deployment receives `LITELLM_BACKEND_URL` from `ui.backendUrl` only when `ingress.componentHosts` is configured.
+
+```yaml
+ingress:
+  enabled: true
+  host: ""
+  componentHosts:
+    ui: ui.example.com
+    backend: backend.example.com
+    gateway: gateway.example.com
+  extraPaths: []
+ui:
+  publicUrl: https://ui.example.com
+  backendUrl: https://backend.example.com
+backend:
+  gatewayUrl: https://gateway.example.com
+```

helm/litellm/templates/NOTES.txt

@@ -22,12 +22,14 @@ Reminders:
       - database.writer.{host,port,dbname}                     (writer connection pieces)
       - database.writer.passwordSecret.{name,usernameKey,passwordKey}
                                                                (Secret holding the writer DB username + password)
+      - database.writer.passwordSecret.{hostKey,urlKey}        (optional, read writer host or full URL from that Secret)
       - database.writer.useIAMAuth: true                       (optional — chart sets IAM_TOKEN_DB_AUTH=true and
                                                                 omits DATABASE_PASSWORD / DATABASE_URL so the proxy
                                                                 mints the URL from an IAM token at startup)
-      - database.reader.host                                   (optional — enables read-replica routing; reader
-                                                                .passwordSecret.name is required when set, unless
-                                                                .useIAMAuth is true)
+      - database.reader.host / database.reader.passwordSecret.{hostKey,urlKey}
+                                                                (optional, enables read-replica routing. The Secret
+                                                                name is required for hostKey, urlKey, and for host
+                                                                with password auth)
       - database.reader.useIAMAuth: true                       (optional, requires database.writer.useIAMAuth: true —
                                                                 chart emits DATABASE_*_READ_REPLICA env vars and
                                                                 omits DATABASE_PASSWORD_READ_REPLICA /
@@ -35,6 +37,8 @@ Reminders:
                                                                 reader URL from an IAM token at startup)
       - redis.passwordSecret.name                              (optional — set when redis.host is provided and the
                                                                 cache requires auth)
+      - redis.passwordSecret.{hostKey,urlKey}                  (optional, read Redis host or full URL from that Secret)
+                                                                (the Secret name is required when either key is set)
       - redis.cluster: true                                    (optional — chart sets REDIS_CLUSTER_NODES from
                                                                 redis.host / redis.port so the proxy's Cache()
                                                                 constructs a RedisClusterCache; the cluster client

helm/litellm/templates/_helpers.tpl

@@ -27,6 +27,30 @@ Common naming + label helpers shared by gateway, backend, and ui templates.
 {{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}}
 {{- end -}}
 
+{{- define "litellm.ui.publicOrigin" -}}
+{{- $origin := .Values.ui.publicUrl | trim | trimSuffix "/" -}}
+{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
+{{- fail "ui.publicUrl must be an origin like https://ui.example.com, without a path" -}}
+{{- end -}}
+{{- $origin -}}
+{{- end -}}
+
+{{- define "litellm.ui.backendOrigin" -}}
+{{- $origin := .Values.ui.backendUrl | trim | trimSuffix "/" -}}
+{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
+{{- fail "ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set" -}}
+{{- end -}}
+{{- $origin -}}
+{{- end -}}
+
+{{- define "litellm.backend.gatewayOrigin" -}}
+{{- $origin := .Values.backend.gatewayUrl | trim | trimSuffix "/" -}}
+{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
+{{- fail "backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set" -}}
+{{- end -}}
+{{- $origin -}}
+{{- end -}}
+
 {{- define "litellm.commonLabels" -}}
 app.kubernetes.io/name: {{ include "litellm.name" . }}
 app.kubernetes.io/instance: {{ .Release.Name }}
@@ -202,7 +226,9 @@ the chart context (needed for .Values), `component` selects which component's
 Sensitive values (master key, DB username + password, Redis password) come
 only from referenced Secrets; the chart never accepts inline values for them.
 
-The chart never assembles DATABASE_URL itself. It emits only the discrete
+The chart never assembles DATABASE_URL itself. When
+database.writer.passwordSecret.urlKey is set, the chart passes the operator's
+URL through from the Secret. Otherwise it emits only the discrete
 DATABASE_HOST/PORT/USER/NAME/SCHEMA (+ DATABASE_PASSWORD for password auth)
 vars; the proxy's entrypoint (DatabaseURLSettings in
 litellm/proxy/db/db_url_settings.py) builds the URL from them and
@@ -218,11 +244,11 @@ token instead of a static password. `database.writer.useAzureEntraAuth: true`
 does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
 for Azure Database for PostgreSQL. The two are mutually exclusive.
 
-The read replica is opt-in via `database.reader.host`. The chart emits
-DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
-DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
-modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
-password auth. When `database.reader.useIAMAuth: true` (or
+The read replica is opt-in via `database.reader.host`, `hostKey`, or `urlKey`.
+The chart emits DATABASE_URL_READ_REPLICA alone for urlKey, or discrete
+DATABASE_*_READ_REPLICA vars for host and hostKey configurations, plus
+DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for password auth.
+When `database.reader.useIAMAuth: true` (or
 `database.reader.useAzureEntraAuth: true`) it omits
 DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
 same way. Reader token auth only takes effect when the writer uses the same
@@ -242,8 +268,36 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
   value: {{ $component.logLevel | quote }}
 {{- end }}
 {{- with $root.Values.database.writer }}
+{{- $hostKey := .passwordSecret.hostKey }}
+{{- $urlKey := .passwordSecret.urlKey }}
+{{- if and .host $hostKey }}
+{{- fail "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" }}
+{{- end }}
+{{- if and $urlKey (or .host $hostKey .dbname .schema .sslMode .sslRootCert .useIAMAuth .useAzureEntraAuth) }}
+{{- fail "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" }}
+{{- end }}
+{{- if and (not .host) (not $hostKey) (not $urlKey) }}
+{{- fail "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" }}
+{{- end }}
+{{- if and (or $hostKey $urlKey) (not .passwordSecret.name) }}
+{{- fail "database.writer.passwordSecret.name is required when hostKey or urlKey is set" }}
+{{- end }}
+{{- if $urlKey }}
+- name: DATABASE_URL
+  valueFrom:
+    secretKeyRef:
+      name: {{ .passwordSecret.name | quote }}
+      key: {{ $urlKey | quote }}
+{{- else }}
 - name: DATABASE_HOST
-  value: {{ required "database.writer.host is required" .host | quote }}
+  {{- if $hostKey }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ .passwordSecret.name | quote }}
+      key: {{ $hostKey | quote }}
+  {{- else }}
+  value: {{ .host | quote }}
+  {{- end }}
 - name: DATABASE_PORT
   value: {{ .port | default 5432 | quote }}
 - name: DATABASE_USER
@@ -282,8 +336,29 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
       key: {{ .passwordSecret.passwordKey | default "password" }}
 {{- end }}
 {{- end }}
+{{- end }}
 {{- with $root.Values.database.reader }}
-{{- if .host }}
+{{- $readerSecret := .passwordSecret | default dict }}
+{{- $hostKey := $readerSecret.hostKey }}
+{{- $urlKey := $readerSecret.urlKey }}
+{{- $readerEnabled := or .host $hostKey $urlKey }}
+{{- if and .host $hostKey }}
+{{- fail "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" }}
+{{- end }}
+{{- if and $urlKey (or .host $hostKey .dbname .schema .useIAMAuth .useAzureEntraAuth) }}
+{{- fail "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" }}
+{{- end }}
+{{- if and (or $hostKey $urlKey) (not $readerSecret.name) }}
+{{- fail "database.reader.passwordSecret.name is required when hostKey or urlKey is set" }}
+{{- end }}
+{{- if $readerEnabled }}
+{{- if $urlKey }}
+- name: DATABASE_URL_READ_REPLICA
+  valueFrom:
+    secretKeyRef:
+      name: {{ $readerSecret.name | quote }}
+      key: {{ $urlKey | quote }}
+{{- else }}
 {{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
 {{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
 {{- end }}
@@ -291,11 +366,18 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
 {{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
 {{- end }}
 - name: DATABASE_HOST_READ_REPLICA
+  {{- if $hostKey }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ $readerSecret.name | quote }}
+      key: {{ $hostKey | quote }}
+  {{- else }}
   value: {{ .host | quote }}
+  {{- end }}
 - name: DATABASE_PORT_READ_REPLICA
   value: {{ .port | default 5432 | quote }}
 - name: DATABASE_NAME_READ_REPLICA
-  value: {{ required "database.reader.dbname is required when database.reader.host is set" .dbname | quote }}
+  value: {{ required "database.reader.dbname is required when database.reader.host, hostKey, or urlKey is set" .dbname | quote }}
 {{- if .schema }}
 - name: DATABASE_SCHEMA_READ_REPLICA
   value: {{ .schema | quote }}
@@ -325,6 +407,7 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
 {{- end }}
 {{- end }}
 {{- end }}
+{{- end }}
 {{/*
 The migrations Job (helm.sh/hook: pre-upgrade) is the single owner of
 `prisma migrate deploy`. Without this, every gateway/backend pod also runs
@@ -341,9 +424,34 @@ harmless no-op for the Job and authoritative for the app pods.
      tracking, pod lock manager) via its REDIS_* env fallback. An explicit
      `general_settings.coordination_redis` block in proxy_config takes
      precedence over anything emitted here. */}}
-{{- if $root.Values.redis.host }}
+{{- $redisSecret := $root.Values.redis.passwordSecret | default dict }}
+{{- $hostKey := $redisSecret.hostKey }}
+{{- $urlKey := $redisSecret.urlKey }}
+{{- if and $root.Values.redis.host $hostKey }}
+{{- fail "redis.host and redis.passwordSecret.hostKey cannot both be set" }}
+{{- end }}
+{{- if and $urlKey (or $root.Values.redis.host $hostKey $root.Values.redis.cluster) }}
+{{- fail "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" }}
+{{- end }}
+{{- if and (or $hostKey $urlKey) (not $redisSecret.name) }}
+{{- fail "redis.passwordSecret.name is required when hostKey or urlKey is set" }}
+{{- end }}
+{{- if $urlKey }}
+- name: REDIS_URL
+  valueFrom:
+    secretKeyRef:
+      name: {{ $redisSecret.name | quote }}
+      key: {{ $urlKey | quote }}
+{{- else if or $root.Values.redis.host $hostKey }}
 - name: REDIS_HOST
+  {{- if $hostKey }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ $redisSecret.name | quote }}
+      key: {{ $hostKey | quote }}
+  {{- else }}
   value: {{ $root.Values.redis.host | quote }}
+  {{- end }}
 - name: REDIS_PORT
   value: {{ $root.Values.redis.port | quote }}
 {{- if $root.Values.redis.passwordSecret.name }}
@@ -360,7 +468,11 @@ harmless no-op for the Job and authoritative for the app pods.
      configured endpoint; the cluster client discovers the remaining nodes from
      CLUSTER SLOTS at startup. */}}
 - name: REDIS_CLUSTER_NODES
+  {{- if $hostKey }}
+  value: {{ printf "[{\"host\":\"$(REDIS_HOST)\",\"port\":%v}]" (int $root.Values.redis.port) | quote }}
+  {{- else }}
   value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }}
+  {{- end }}
 {{- end }}
 {{- end }}
 {{- with $component.extraEnv }}

helm/litellm/templates/backend/deployment.yaml

@@ -1,4 +1,6 @@
 {{- if .Values.backend.enabled }}
+{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
+{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
 apiVersion: apps/v1
 kind: Deployment
 metadata:
@@ -58,7 +60,17 @@ spec:
               protocol: TCP
           env:
             {{- include "litellm.lensConnectionEnv" . | nindent 12 }}
+            {{- if .Values.ui.publicUrl }}
+            - name: LITELLM_CORS_ORIGINS
+              value: {{ include "litellm.ui.publicOrigin" . | quote }}
+            - name: LITELLM_CORS_ALLOW_CREDENTIALS
+              value: "true"
+            {{- end }}
             {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
+            {{- if .Values.backend.gatewayUrl }}
+            - name: LITELLM_GATEWAY_URL
+              value: {{ if $hasComponentHosts }}{{ include "litellm.backend.gatewayOrigin" . | quote }}{{ else }}{{ .Values.backend.gatewayUrl | quote }}{{ end }}
+            {{- end }}
             {{- if .Values.gateway.config.create }}
             - name: CONFIG_FILE_PATH
               value: /app/config/config.yaml

helm/litellm/templates/gateway/deployment.yaml

@@ -56,6 +56,10 @@ spec:
               protocol: TCP
           env:
             {{- include "litellm.lensConnectionEnv" . | nindent 12 }}
+            {{- if .Values.ui.publicUrl }}
+            - name: LITELLM_CORS_ORIGINS
+              value: {{ include "litellm.ui.publicOrigin" . | quote }}
+            {{- end }}
             {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
             {{- if .Values.gateway.config.create }}
             - name: CONFIG_FILE_PATH

helm/litellm/templates/ingress.yaml

@@ -1,3 +1,36 @@
+{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
+{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
+{{- if $hasComponentHosts }}
+{{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }}
+{{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }}
+{{- end }}
+{{- if not .Values.gateway.enabled }}
+{{- fail "ingress.componentHosts requires gateway.enabled=true" }}
+{{- end }}
+{{- if not .Values.backend.enabled }}
+{{- fail "ingress.componentHosts requires backend.enabled=true" }}
+{{- end }}
+{{- if not .Values.ui.enabled }}
+{{- fail "ingress.componentHosts requires ui.enabled=true" }}
+{{- end }}
+{{- if .Values.ingress.host }}
+{{- fail "ingress.host must be empty when ingress.componentHosts is configured" }}
+{{- end }}
+{{- if .Values.ingress.extraPaths }}
+{{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }}
+{{- end }}
+{{- if not (trim .Values.ui.backendUrl) }}
+{{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }}
+{{- end }}
+{{- $_ := include "litellm.ui.backendOrigin" . }}
+{{- if not (trim .Values.backend.gatewayUrl) }}
+{{- fail "backend.gatewayUrl is required when ingress.componentHosts is configured" }}
+{{- end }}
+{{- $_ := include "litellm.backend.gatewayOrigin" . }}
+{{- if not (trim .Values.ui.publicUrl) }}
+{{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }}
+{{- end }}
+{{- end }}
 {{- if .Values.ingress.enabled -}}
 {{- $gatewayName := include "litellm.gateway.fullname" . -}}
 {{- $backendName := include "litellm.backend.fullname" . -}}
@@ -90,6 +123,66 @@
   adding to it.
 */}}
 {{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
+{{- if $hasComponentHosts }}
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  name: {{ include "litellm.fullname" . }}
+  labels:
+    {{- include "litellm.commonLabels" . | nindent 4 }}
+  {{- with .Values.ingress.annotations }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+spec:
+  {{- with .Values.ingress.className }}
+  ingressClassName: {{ . | quote }}
+  {{- end }}
+  {{- with .Values.ingress.tls }}
+  tls:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+  rules:
+    - host: {{ $componentHosts.ui | quote }}
+      http:
+        paths:
+          - path: /
+            pathType: Prefix
+            backend:
+              service:
+                name: {{ $uiName }}
+                port:
+                  number: {{ $uiPort }}
+    - host: {{ $componentHosts.backend | quote }}
+      http:
+        paths:
+          - path: /
+            pathType: Prefix
+            backend:
+              service:
+                name: {{ $backendName }}
+                port:
+                  number: {{ $backendPort }}
+    - host: {{ $componentHosts.gateway | quote }}
+      http:
+        paths:
+          {{- if or (eq (include "litellm.lens.mode" .) "bundled") (and (eq (include "litellm.lens.mode" .) "external") .Values.lensWorker.externalServiceName) }}
+          - path: /lens-ingest
+            pathType: Prefix
+            backend:
+              service:
+                name: {{ .Values.lensWorker.externalServiceName | default (printf "%s-lens-worker" (include "litellm.fullname" .)) }}
+                port:
+                  number: {{ .Values.lensWorker.service.port }}
+          {{- end }}
+          - path: /
+            pathType: Prefix
+            backend:
+              service:
+                name: {{ $gatewayName }}
+                port:
+                  number: {{ $gatewayPort }}
+{{- else }}
 apiVersion: networking.k8s.io/v1
 kind: Ingress
 metadata:
@@ -217,3 +310,4 @@ spec:
                 port:
                   number: {{ $backendPort }}
 {{- end }}
+{{- end }}

helm/litellm/templates/ui/deployment.yaml

@@ -1,3 +1,5 @@
+{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
+{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
 {{- if .Values.ui.enabled }}
 apiVersion: apps/v1
 kind: Deployment
@@ -56,9 +58,9 @@ spec:
             - name: LITELLM_LOG
               value: {{ .Values.ui.logLevel | quote }}
             {{- end }}
-            {{- if .Values.ui.backendUrl }}
+            {{- if and .Values.ui.backendUrl $hasComponentHosts }}
             - name: LITELLM_BACKEND_URL
-              value: {{ .Values.ui.backendUrl | quote }}
+              value: {{ include "litellm.ui.backendOrigin" . | quote }}
             {{- end }}
             {{- with .Values.ui.extraEnv }}
             {{- toYaml . | nindent 12 }}

helm/litellm/tests/component_hosts_tests.yaml

@@ -0,0 +1,441 @@
+suite: ingress component hosts
+templates:
+  - ingress.yaml
+  - backend/deployment.yaml
+  - gateway/deployment.yaml
+  - gateway/configmap.yaml
+  - ui/deployment.yaml
+values:
+  - ./values/required.yaml
+tests:
+  - it: routes each hostname to its component and injects the external URLs
+    template: ingress.yaml
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      lensWorker.mode: disabled
+      ingress.tls:
+        - hosts:
+            - ui.example.com
+            - backend.example.com
+            - gateway.example.com
+          secretName: component-hosts
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+    asserts:
+      - equal:
+          path: spec.rules[0].host
+          value: ui.example.com
+      - equal:
+          path: spec.tls[0]
+          value:
+            hosts:
+              - ui.example.com
+              - backend.example.com
+              - gateway.example.com
+            secretName: component-hosts
+      - equal:
+          path: spec.rules[0].http.paths
+          value:
+            - path: /
+              pathType: Prefix
+              backend:
+                service:
+                  name: RELEASE-NAME-litellm-ui
+                  port:
+                    number: 3000
+      - equal:
+          path: spec.rules[1].host
+          value: backend.example.com
+      - equal:
+          path: spec.rules[1].http.paths
+          value:
+            - path: /
+              pathType: Prefix
+              backend:
+                service:
+                  name: RELEASE-NAME-litellm-backend
+                  port:
+                    number: 4001
+      - equal:
+          path: spec.rules[2].host
+          value: gateway.example.com
+      - equal:
+          path: spec.rules[2].http.paths
+          value:
+            - path: /
+              pathType: Prefix
+              backend:
+                service:
+                  name: RELEASE-NAME-litellm-gateway
+                  port:
+                    number: 4000
+
+  - it: keeps Lens ingestion ahead of the gateway catch-all
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+      lensWorker.mode: external
+      lensWorker.externalServiceName: existing-lens
+      lensWorker.externalUrl: https://lens.example.com
+      lensWorker.publicUrl: https://lens.example.com
+      lensWorker.serviceTokenSecret.name: existing-lens-signing
+      lensWorker.gateway.secretName: existing-lens-signing
+    asserts:
+      - equal:
+          path: spec.rules[2].http.paths[0].path
+          value: /lens-ingest
+      - equal:
+          path: spec.rules[2].http.paths[0].backend.service.name
+          value: existing-lens
+      - equal:
+          path: spec.rules[2].http.paths[1].path
+          value: /
+    template: ingress.yaml
+
+  - it: rejects a partial component host configuration
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: ingress.componentHosts.ui, backend, and gateway must all be set together
+    template: ingress.yaml
+
+  - it: rejects component hosts when the gateway is disabled
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+      gateway.enabled: false
+    asserts:
+      - failedTemplate:
+          errorMessage: ingress.componentHosts requires gateway.enabled=true
+    template: ingress.yaml
+
+  - it: rejects component hosts when the backend is disabled
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+      backend.enabled: false
+    asserts:
+      - failedTemplate:
+          errorMessage: ingress.componentHosts requires backend.enabled=true
+    template: ingress.yaml
+
+  - it: rejects component hosts when the UI is disabled
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+      ui.enabled: false
+    asserts:
+      - failedTemplate:
+          errorMessage: ingress.componentHosts requires ui.enabled=true
+    template: ingress.yaml
+
+  - it: rejects a legacy host alongside component hosts
+    set:
+      ingress.enabled: true
+      ingress.host: legacy.example.com
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: ingress.host must be empty when ingress.componentHosts is configured
+    template: ingress.yaml
+
+  - it: rejects extra paths alongside component hosts
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ingress.extraPaths:
+        - path: /custom
+    asserts:
+      - failedTemplate:
+          errorMessage: ingress.extraPaths must be empty when ingress.componentHosts is configured
+    template: ingress.yaml
+
+  - it: requires the UI backend URL in component host mode
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: ui.backendUrl is required when ingress.componentHosts is configured
+    template: ingress.yaml
+
+  - it: requires the backend gateway URL in component host mode
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: backend.gatewayUrl is required when ingress.componentHosts is configured
+    template: ingress.yaml
+
+  - it: requires the public UI origin in component host mode
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      backend.gatewayUrl: https://gateway.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: ui.publicUrl is required when ingress.componentHosts is configured
+    template: ingress.yaml
+
+  - it: exposes the gateway URL in the backend Deployment
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: "  https://gateway.example.com/  "
+    template: backend/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_GATEWAY_URL
+            value: https://gateway.example.com
+
+  - it: preserves the backend gateway URL verbatim in single-host mode
+    set:
+      ingress.enabled: true
+      ingress.host: litellm.example.com
+      backend.gatewayUrl: "  https://gateway.example.com/api/  "
+    template: backend/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_GATEWAY_URL
+            value: "  https://gateway.example.com/api/  "
+
+  - it: exposes the backend URL in the UI Deployment
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: "  https://backend.example.com/  "
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+    template: ui/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_BACKEND_URL
+            value: https://backend.example.com
+
+  - it: omits the backend URL from the UI Deployment in single-host mode
+    set:
+      ingress.enabled: true
+      ingress.host: litellm.example.com
+      ui.backendUrl: https://litellm.example.com/api
+    template: ui/deployment.yaml
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_BACKEND_URL
+          any: true
+
+  - it: enables credentialed CORS from the public UI origin on the backend
+    set:
+      ui.publicUrl: https://ui.example.com
+    template: backend/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_CORS_ORIGINS
+            value: https://ui.example.com
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_CORS_ALLOW_CREDENTIALS
+            value: "true"
+
+  - it: lets backend extraEnv override generated CORS values
+    set:
+      ui.publicUrl: https://ui.example.com
+      backend.extraEnv:
+        - name: LITELLM_CORS_ORIGINS
+          value: https://backend-override.example.com
+        - name: LITELLM_CORS_ALLOW_CREDENTIALS
+          value: "false"
+    template: backend/deployment.yaml
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].env[0].name
+          value: LITELLM_CORS_ORIGINS
+      - equal:
+          path: spec.template.spec.containers[0].env[0].value
+          value: https://ui.example.com
+      - equal:
+          path: spec.template.spec.containers[0].env[1].name
+          value: LITELLM_CORS_ALLOW_CREDENTIALS
+      - equal:
+          path: spec.template.spec.containers[0].env[1].value
+          value: "true"
+      - equal:
+          path: spec.template.spec.containers[0].env[10].name
+          value: LITELLM_CORS_ORIGINS
+      - equal:
+          path: spec.template.spec.containers[0].env[10].value
+          value: https://backend-override.example.com
+      - equal:
+          path: spec.template.spec.containers[0].env[11].name
+          value: LITELLM_CORS_ALLOW_CREDENTIALS
+      - equal:
+          path: spec.template.spec.containers[0].env[11].value
+          value: "false"
+
+  - it: exposes the public UI origin to the gateway without enabling credentials
+    set:
+      ui.publicUrl: https://ui.example.com
+    template: gateway/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_CORS_ORIGINS
+            value: https://ui.example.com
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_CORS_ALLOW_CREDENTIALS
+            value: "true"
+
+  - it: trims the public UI origin before setting backend CORS
+    set:
+      ui.publicUrl: " https://ui.example.com/ "
+    template: backend/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_CORS_ORIGINS
+            value: https://ui.example.com
+
+  - it: trims the public UI origin before setting gateway CORS
+    set:
+      ui.publicUrl: " https://ui.example.com/ "
+    template: gateway/deployment.yaml
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_CORS_ORIGINS
+            value: https://ui.example.com
+
+  - it: rejects a UI public URL with a path
+    set:
+      ui.publicUrl: https://ui.example.com/app
+    template: gateway/deployment.yaml
+    asserts:
+      - failedTemplate:
+          errorMessage: ui.publicUrl must be an origin like https://ui.example.com, without a path
+
+  - it: lets gateway extraEnv override the generated CORS origin
+    set:
+      ui.publicUrl: https://ui.example.com
+      gateway.extraEnv:
+        - name: LITELLM_CORS_ORIGINS
+          value: https://gateway-override.example.com
+    template: gateway/deployment.yaml
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].env[0].name
+          value: LITELLM_CORS_ORIGINS
+      - equal:
+          path: spec.template.spec.containers[0].env[0].value
+          value: https://ui.example.com
+      - equal:
+          path: spec.template.spec.containers[0].env[9].name
+          value: LITELLM_CORS_ORIGINS
+      - equal:
+          path: spec.template.spec.containers[0].env[9].value
+          value: https://gateway-override.example.com
+
+---
+suite: component host UI backend origin validation
+templates:
+  - ingress.yaml
+values:
+  - ./values/required.yaml
+tests:
+  - it: rejects a pathful UI backend URL in component host mode
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com/api
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set
+  - it: rejects a pathful backend gateway URL in component host mode
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: https://gateway.example.com/v1
+    asserts:
+      - failedTemplate:
+          errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set
+  - it: rejects a backend gateway URL without a scheme in component host mode
+    set:
+      ingress.enabled: true
+      ingress.componentHosts.ui: ui.example.com
+      ingress.componentHosts.backend: backend.example.com
+      ingress.componentHosts.gateway: gateway.example.com
+      ui.backendUrl: https://backend.example.com
+      ui.publicUrl: https://ui.example.com
+      backend.gatewayUrl: gateway.example.com
+    asserts:
+      - failedTemplate:
+          errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set

helm/litellm/tests/database_auth_tests.yaml

@@ -191,3 +191,400 @@ tests:
           content:
             name: DATABASE_PASSWORD_READ_REPLICA
           any: true
+
+  - it: writer hostKey reads the host from the Secret and supports IAM auth
+    template: gateway/deployment.yaml
+    set:
+      database.writer.host: ""
+      database.writer.passwordSecret.hostKey: hostname
+      database.writer.useIAMAuth: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST
+            valueFrom:
+              secretKeyRef:
+                name: litellm-writer-secret
+                key: hostname
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+
+  - it: writer urlKey supplies only DATABASE_URL to each database client container and keeps pool env
+    set:
+      database.writer.host: ""
+      database.writer.dbname: ""
+      database.writer.schema: ""
+      database.writer.sslMode: ""
+      database.writer.sslRootCert: ""
+      database.writer.passwordSecret.urlKey: connection-url
+      database.connectionPool.enabled: true
+      gateway.collector.enabled: true
+      gateway.metricsServer.enabled: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            valueFrom:
+              secretKeyRef:
+                name: litellm-writer-secret
+                key: connection-url
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USER
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_NAME
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLMODE
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLROOTCERT
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+        template: gateway/deployment.yaml
+      - notExists:
+          path: >-
+            spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
+            @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
+            @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
+            @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
+            @.name == "DATABASE_PASSWORD")]
+        template: gateway/deployment.yaml
+      - notExists:
+          path: >-
+            spec.template.spec.containers[2].env[?(@.name == "DATABASE_HOST" ||
+            @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
+            @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
+            @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
+            @.name == "DATABASE_PASSWORD")]
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: DATABASE_URL
+            valueFrom:
+              secretKeyRef:
+                name: litellm-writer-secret
+                key: connection-url
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: DATABASE_USER
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: DATABASE_URL
+            valueFrom:
+              secretKeyRef:
+                name: litellm-writer-secret
+                key: connection-url
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            valueFrom:
+              secretKeyRef:
+                name: litellm-writer-secret
+                key: connection-url
+        template: backend/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USER
+          any: true
+        template: backend/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+        template: backend/deployment.yaml
+      - notExists:
+          path: >-
+            spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
+            @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
+            @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
+            @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
+            @.name == "DATABASE_PASSWORD")]
+        template: backend/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            valueFrom:
+              secretKeyRef:
+                name: litellm-writer-secret
+                key: connection-url
+        template: migrations-job.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USER
+          any: true
+        template: migrations-job.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+        template: migrations-job.yaml
+      - notExists:
+          path: >-
+            spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
+            @.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
+            @.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
+            @.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
+            @.name == "DATABASE_PASSWORD")]
+        template: migrations-job.yaml
+
+  - it: writer rejects host together with hostKey
+    template: gateway/deployment.yaml
+    set:
+      database.writer.passwordSecret.hostKey: hostname
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set"
+
+  - it: writer rejects urlKey with discrete connection settings
+    template: gateway/deployment.yaml
+    set:
+      database.writer.host: ""
+      database.writer.passwordSecret.urlKey: connection-url
+      database.writer.schema: public
+      database.writer.sslMode: require
+      database.writer.sslRootCert: /etc/ssl/ca.crt
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
+
+  - it: writer rejects urlKey with IAM auth
+    template: gateway/deployment.yaml
+    set:
+      database.writer.host: ""
+      database.writer.dbname: ""
+      database.writer.passwordSecret.urlKey: connection-url
+      database.writer.useIAMAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
+
+  - it: writer rejects urlKey with Entra auth
+    template: gateway/deployment.yaml
+    set:
+      database.writer.host: ""
+      database.writer.dbname: ""
+      database.writer.passwordSecret.urlKey: connection-url
+      database.writer.useAzureEntraAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
+
+  - it: writer requires an endpoint
+    template: gateway/deployment.yaml
+    set:
+      database.writer.host: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required"
+
+  - it: writer hostKey requires a Secret name
+    template: gateway/deployment.yaml
+    set:
+      database.writer.host: ""
+      database.writer.passwordSecret.hostKey: hostname
+      database.writer.passwordSecret.name: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.writer.passwordSecret.name is required when hostKey or urlKey is set"
+
+  - it: reader renders without replica env vars when passwordSecret is null
+    template: gateway/deployment.yaml
+    set:
+      database.reader.passwordSecret: null
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST_READ_REPLICA
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL_READ_REPLICA
+          any: true
+
+  - it: reader hostKey enables replica routing and reads its host from the Secret
+    template: gateway/deployment.yaml
+    set:
+      database.writer.useIAMAuth: true
+      database.reader.passwordSecret.hostKey: hostname
+      database.reader.dbname: litellm
+      database.reader.useIAMAuth: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST_READ_REPLICA
+            valueFrom:
+              secretKeyRef:
+                name: litellm-reader-secret
+                key: hostname
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_NAME_READ_REPLICA
+            value: litellm
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+
+  - it: reader urlKey emits only the replica URL
+    template: gateway/deployment.yaml
+    set:
+      database.reader.passwordSecret.urlKey: connection-url
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL_READ_REPLICA
+            valueFrom:
+              secretKeyRef:
+                name: litellm-reader-secret
+                key: connection-url
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST_READ_REPLICA
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT_READ_REPLICA
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USER_READ_REPLICA
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD_READ_REPLICA
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_NAME_READ_REPLICA
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA_READ_REPLICA
+          any: true
+
+  - it: reader rejects host together with hostKey
+    template: gateway/deployment.yaml
+    set:
+      database.reader.host: reader.example.com
+      database.reader.passwordSecret.hostKey: hostname
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set"
+
+  - it: reader rejects urlKey with discrete connection settings
+    template: gateway/deployment.yaml
+    set:
+      database.reader.passwordSecret.urlKey: connection-url
+      database.reader.dbname: litellm
+      database.reader.schema: public
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
+
+  - it: reader rejects urlKey with token auth
+    template: gateway/deployment.yaml
+    set:
+      database.writer.useIAMAuth: true
+      database.reader.passwordSecret.urlKey: connection-url
+      database.reader.useIAMAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
+
+  - it: reader urlKey requires a Secret name
+    template: gateway/deployment.yaml
+    set:
+      database.reader.passwordSecret.urlKey: connection-url
+      database.reader.passwordSecret.name: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.reader.passwordSecret.name is required when hostKey or urlKey is set"

helm/litellm/tests/redis_env_tests.yaml

@@ -21,6 +21,27 @@ tests:
             name: REDIS_CLUSTER_NODES
           any: true
 
+  - it: gateway renders without redis env vars when passwordSecret is null
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret: null
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_HOST
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_URL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PASSWORD
+          any: true
+
   - it: gateway emits host, port and password when redis is configured
     template: gateway/deployment.yaml
     set:
@@ -68,6 +89,27 @@ tests:
                 name: redis-secret
                 key: redis-password
 
+  - it: backend renders without redis env vars when passwordSecret is null
+    template: backend/deployment.yaml
+    set:
+      redis.passwordSecret: null
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_HOST
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_URL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PASSWORD
+          any: true
+
   - it: gateway omits REDIS_PASSWORD for an auth-less redis
     template: gateway/deployment.yaml
     set:
@@ -107,3 +149,174 @@ tests:
           content:
             name: REDIS_CLUSTER_NODES
           any: true
+
+  - it: gateway gets redis host from hostKey and keeps port and password
+    template: gateway/deployment.yaml
+    set:
+      redis.host: ""
+      redis.port: 6380
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.hostKey: hostname
+      redis.passwordSecret.passwordKey: redis-password
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_HOST
+            valueFrom:
+              secretKeyRef:
+                name: redis-secret
+                key: hostname
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PORT
+            value: "6380"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PASSWORD
+            valueFrom:
+              secretKeyRef:
+                name: redis-secret
+                key: redis-password
+
+  - it: backend gets redis host from hostKey with the same coordination env
+    template: backend/deployment.yaml
+    set:
+      redis.host: ""
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.hostKey: hostname
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_HOST
+            valueFrom:
+              secretKeyRef:
+                name: redis-secret
+                key: hostname
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PORT
+            value: "6379"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PASSWORD
+            valueFrom:
+              secretKeyRef:
+                name: redis-secret
+                key: password
+
+  - it: gateway expands the redis cluster seed from hostKey at runtime
+    template: gateway/deployment.yaml
+    set:
+      redis.host: ""
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.hostKey: hostname
+      redis.cluster: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_CLUSTER_NODES
+            value: '[{"host":"$(REDIS_HOST)","port":6379}]'
+      - equal:
+          path: spec.template.spec.containers[0].env[8].name
+          value: REDIS_HOST
+      - equal:
+          path: spec.template.spec.containers[0].env[11].name
+          value: REDIS_CLUSTER_NODES
+
+  - it: gateway gets REDIS_URL from urlKey without discrete Redis env
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.urlKey: connection-url
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_URL
+            valueFrom:
+              secretKeyRef:
+                name: redis-secret
+                key: connection-url
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_HOST
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PORT
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_CLUSTER_NODES
+          any: true
+
+  - it: redis rejects host together with hostKey
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.passwordSecret.hostKey: hostname
+      redis.passwordSecret.name: redis-secret
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.host and redis.passwordSecret.hostKey cannot both be set"
+
+  - it: redis urlKey rejects host
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.passwordSecret.urlKey: connection-url
+      redis.passwordSecret.name: redis-secret
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
+
+  - it: redis urlKey rejects hostKey
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.hostKey: hostname
+      redis.passwordSecret.urlKey: connection-url
+      redis.passwordSecret.name: redis-secret
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
+
+  - it: redis urlKey rejects cluster mode
+    template: gateway/deployment.yaml
+    set:
+      redis.cluster: true
+      redis.passwordSecret.urlKey: connection-url
+      redis.passwordSecret.name: redis-secret
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
+
+  - it: redis hostKey requires a Secret name
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.hostKey: hostname
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
+
+  - it: redis urlKey requires a Secret name
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.urlKey: connection-url
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"

helm/litellm/values.yaml

@@ -23,7 +23,11 @@ ingress:
   #          wildcard pathType, so that rule could never match there.
   controller: alb
   annotations: {}
-  host: ""        # optional; if set, becomes the rule's host
+  host: ""        # optional; if set, becomes the single ingress rule's host
+  componentHosts:
+    ui: ""
+    backend: ""
+    gateway: ""
   tls: []
   # Extra HTTP paths appended to the ingress rule. Additive: every built-in
   # UI / gateway / backend path is still rendered, these entries are placed
@@ -208,6 +212,8 @@ database:
       name: litellm-writer-secret
       usernameKey: username
       passwordKey: password
+      hostKey: "" # Secret key containing the writer database host
+      urlKey: "" # Secret key containing the complete writer database URL
     # libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
     # in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
     # sslRootCert on its own implies sslMode verify-full
@@ -229,6 +235,8 @@ database:
       name: litellm-reader-secret
       usernameKey: username
       passwordKey: password
+      hostKey: "" # Secret key containing the reader database host
+      urlKey: "" # Secret key containing the complete reader database URL
 
   # In-container connection pool (PgBouncer, transaction mode) shared by every
   # gateway worker in the pod. Without it each of the `gateway.numWorkers`
@@ -274,6 +282,8 @@ redis:
   passwordSecret:
     name: "" # Leave empty for auth-less Redis
     passwordKey: password
+    hostKey: "" # Secret key containing the Redis host
+    urlKey: "" # Secret key containing the complete Redis URL
 
 # ---------- gateway (LLM data plane) ----------
 gateway:
@@ -493,6 +503,7 @@ gateway:
 backend:
   enabled: true
   logLevel: INFO
+  gatewayUrl: ""
   extraEnv: []
   envConfigMaps: []
   envSecrets: []
@@ -575,9 +586,10 @@ ui:
   service:
     type: ClusterIP
     port: 3000
-  # The dashboard expects to know where to reach the backend API. Set this to
-  # the externally-routable URL (typically the ingress host + /api or similar).
+  # The backend's public origin, passed to the UI only when ingress.componentHosts is configured.
   backendUrl: ""
+  # Public origin of the UI, used for backend and gateway CORS.
+  publicUrl: ""
   resources:
     requests:
       cpu: 500m

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -13,6 +13,7 @@ from litellm_proxy_extras.prisma_toolchain import (
 )
 
 MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
+PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
 
 if TYPE_CHECKING:
     import psycopg
@@ -32,7 +33,7 @@ def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
     if not database_url:
         return environment
     migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
-    return {**environment, "DATABASE_URL": migration_url}
+    return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
 
 
 @dataclass(frozen=True, slots=True)
@@ -55,10 +56,15 @@ class MigrationCoordinator:
     def check_connection(self) -> None:
         self.connection.execute("SELECT 1")
 
+    def release_statement_snapshot(self) -> None:
+        # a parameterized statement leaves its portal snapshot registered on the backend until the
+        # next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
+        self.connection.execute("SELECT 1")
+
     def acquire_prisma_lock(self) -> None:
         deadline: Final = time.monotonic() + migration_lock_timeout()
         while time.monotonic() < deadline:
-            if _try_lock(self.connection, 72707369):
+            if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
                 return
             time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
         raise RuntimeError(
@@ -95,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
     )
 
 
+@contextmanager
+def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
+    """Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
+    pinned connection while a Prisma command runs with its session advisory lock disabled, so a
+    transaction pooler can never strand that lock on an idle backend after the command exits."""
+    with migration_lock(database_url) as coordinator:
+        coordinator.acquire_prisma_lock()
+        coordinator.release_statement_snapshot()
+        yield
+
+
 @contextmanager
 def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
     """A session-level, non-blocking hold of the migration coordinator lock on an autocommit

litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql

@@ -0,0 +1,4 @@
+SET LOCAL lock_timeout = '10s';
+
+-- AlterTable
+ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;

litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_add_access_group_skill_names/migration.sql

@@ -0,0 +1,4 @@
+SET LOCAL lock_timeout = '10s';
+
+-- AlterTable
+ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_skill_names" TEXT[] DEFAULT ARRAY[]::TEXT[];

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable {
   // Resource memberships - explicit arrays per type
   access_model_names        String[] @default([])
   access_mcp_server_ids   String[] @default([])
+  mcp_tool_permissions    Json?
   access_agent_ids        String[] @default([])
+  access_skill_names      String[] @default([])
 
   assigned_team_ids         String[] @default([])
   assigned_key_ids          String[] @default([])

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -9,6 +9,7 @@ import tempfile
 import threading
 import time
 from collections.abc import Callable
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
@@ -1021,7 +1022,11 @@ class ProxyExtrasDBManager:
     def _setup_database_v2(use_migrate: bool) -> bool:
         if not use_migrate:
             return ProxyExtrasDBManager._run_database_v2(False)
-        from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
+        from litellm_proxy_extras.migration_lock import (
+            migration_environment,
+            migration_lock,
+            prisma_migration_lock,
+        )
         from litellm_proxy_extras.migration_recovery import (
             baseline_current_schema,
             recover_completed_migration,
@@ -1055,7 +1060,9 @@ class ProxyExtrasDBManager:
                     migration_environment(_get_prisma_env()),
                 )
 
-        while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
+        while not ProxyExtrasDBManager._run_database_v2(
+            True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url)
+        ):
             continue
         return True
 
@@ -1064,6 +1071,7 @@ class ProxyExtrasDBManager:
         use_migrate: bool,
         recover_completed: Callable[[str], bool] = lambda name: False,
         baseline_existing: "Callable[[str], None] | None" = None,
+        deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
     ) -> bool:
         """
         v2 migration resolver (opt-in via --use_v2_migration_resolver).
@@ -1126,11 +1134,12 @@ class ProxyExtrasDBManager:
             while not budget.exhausted:
                 attempt_started = time.monotonic()
                 try:
-                    result = prisma_toolchain.run_prisma(
-                        [_get_prisma_command(), "migrate", "deploy"],
-                        timeout=deploy_timeout,
-                        env=migration_env,
-                    )
+                    with deploy_lock():
+                        result = prisma_toolchain.run_prisma(
+                            [_get_prisma_command(), "migrate", "deploy"],
+                            timeout=deploy_timeout,
+                            env=migration_env,
+                        )
                     logger.info(f"prisma migrate deploy stdout: {result.stdout}")
                     return True
 
@@ -1155,6 +1164,7 @@ class ProxyExtrasDBManager:
                         e,
                         budget,
                         time.monotonic() - attempt_started,
+                        deploy_lock=deploy_lock,
                     )
 
                 if next_budget.attempts_left < budget.attempts_left:
@@ -1187,15 +1197,19 @@ class ProxyExtrasDBManager:
         return match.group(1) if match else None
 
     @staticmethod
-    def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
+    def _v2_roll_back_migration_best_effort(
+        migration_name: str,
+        deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
+    ) -> None:
         from litellm_proxy_extras.migration_lock import migration_environment
 
         try:
-            prisma_toolchain.run_prisma(
-                [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
-                timeout=prisma_command_timeout(),
-                env=migration_environment(_get_prisma_env()),
-            )
+            with deploy_lock():
+                prisma_toolchain.run_prisma(
+                    [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
+                    timeout=prisma_command_timeout(),
+                    env=migration_environment(_get_prisma_env()),
+                )
         except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
             pass
 
@@ -1204,6 +1218,7 @@ class ProxyExtrasDBManager:
         error: subprocess.CalledProcessError,
         budget: "_MigrateAttemptBudget",
         attempt_seconds: float = 0.0,
+        deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
     ) -> "_MigrateAttemptBudget":
         """Recover from one failed `prisma migrate deploy`, and price the pass.
 
@@ -1225,14 +1240,14 @@ class ProxyExtrasDBManager:
                         "deadlock race, rolling its ledger row back and retrying",
                         migration_name,
                     )
-                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                     return budget.spend()
                 if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
                     logger.info(
                         "Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
                         migration_name,
                     )
-                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                     return budget.spend()
                 if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
                     logger.info(
@@ -1276,7 +1291,7 @@ class ProxyExtrasDBManager:
                     "rolling its ledger row back and retrying",
                     migration_name,
                 )
-                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                 return budget.spend()
 
             if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
@@ -1286,7 +1301,7 @@ class ProxyExtrasDBManager:
                     "its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
                     migration_name,
                 )
-                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
+                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
                 return budget.spend()
 
             raise RuntimeError(

litellm-rust/Cargo.lock

@@ -3406,6 +3406,7 @@ dependencies = [
 name = "litellm-auth-gcp"
 version = "0.1.0"
 dependencies = [
+ "base64 0.22.1",
  "gcp_auth",
  "google-cloud-auth",
  "http 1.4.2",
@@ -3922,13 +3923,16 @@ version = "0.1.0"
 dependencies = [
  "bytes",
  "futures-util",
+ "litellm-coroutine",
  "litellm-host",
+ "parking_lot",
  "pyo3",
  "pyo3-async-runtimes",
  "pythonize",
  "rstest",
  "serde",
  "serde_json",
+ "thiserror 2.0.19",
  "tokio",
 ]
 

litellm-rust/Cargo.toml

@@ -80,6 +80,7 @@ http = "1"
 google-cloud-auth = { version = "1.16.0", default-features = false }
 jsonwebtoken = { version = "11.1.0", default-features = false }
 hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
+parking_lot = "0.12.5"
 proptest = "1.7.0"
 pyo3 = "0.29.2"
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }

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

@@ -11,6 +11,7 @@ google-sdk = ["dep:google-cloud-auth", "dep:http"]
 [dependencies]
 litellm-auth-types.workspace = true
 
+base64.workspace = true
 moka.workspace = true
 serde.workspace = true
 serde_json.workspace = true

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

@@ -1,4 +1,4 @@
-use std::{future::Future, path::Path, pin::Pin, sync::Arc};
+use std::{future::Future, pin::Pin, sync::Arc};
 
 use gcp_auth::{CustomServiceAccount, TokenProvider};
 use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
@@ -8,7 +8,8 @@ use sha2::{Digest, Sha256};
 
 use crate::{
     config::{
-        CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, non_empty_env,
+        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,
@@ -176,14 +177,7 @@ impl VertexProviderLoader for GcpProviderLoader {
                     .map_err(auth_acquisition_error)?,
                 ),
                 CredentialSource::Trusted(configured) => {
-                    let configured = configured.expose();
-                    let service_account = if Path::new(configured).is_file() {
-                        CustomServiceAccount::from_file(configured)
-                    } else {
-                        CustomServiceAccount::from_json(configured)
-                    }
-                    .map_err(auth_acquisition_error)?;
-                    Arc::new(service_account)
+                    Arc::new(load_trusted_service_account(configured.expose())?)
                 }
                 CredentialSource::ApplicationCredentials(path) => {
                     Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
@@ -240,6 +234,23 @@ enum CredentialCacheKey {
     Adc,
 }
 
+const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
+    Set it to a readable file path, or to the credentials JSON itself.";
+
+fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
+    if is_base64_encoded_json(configured) {
+        return Err(Error::InvalidConfiguration(ErrorDetail::Message(
+            BASE64_CREDENTIALS.into(),
+        )));
+    }
+    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",
@@ -264,6 +275,73 @@ mod tests {
         assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
     }
 
+    #[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");
+        };
+        assert!(error.to_string().contains(expected), "{error}");
+    }
+
+    #[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) => 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]
     fn inline_and_trusted_credentials_never_share_a_cache_entry() {
         assert_ne!(

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

@@ -1,5 +1,9 @@
 use std::{collections::BTreeMap, path::Path};
 
+use base64::{
+    Engine, alphabet,
+    engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
+};
 use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced};
 use serde_json::{Map, Value};
 
@@ -80,9 +84,10 @@ pub fn get_vertex_ai_project_from_credentials(
         CredentialSource::Inline(configured) => configured.expose().to_string(),
         CredentialSource::Trusted(configured) => {
             let configured = configured.expose();
-            match std::fs::read_to_string(configured) {
-                Ok(contents) if Path::new(configured).is_file() => contents,
-                _ => configured.to_string(),
+            if names_a_credentials_file(configured) {
+                std::fs::read_to_string(configured).ok()?
+            } else {
+                configured.to_string()
             }
         }
         CredentialSource::ApplicationCredentials(path) => std::fs::read_to_string(path).ok()?,
@@ -105,6 +110,32 @@ pub fn get_vertex_ai_location(
         .or_else(|| VertexParams::LOCATION.resolve(&|_| None, env_lookup))
 }
 
+const INLINE_JSON_LEADING_NOISE: [char; 7] = [' ', '\t', '\r', '\n', '\u{feff}', '"', '\''];
+
+fn is_inline_json(configured: &str) -> bool {
+    configured
+        .trim_start_matches(INLINE_JSON_LEADING_NOISE)
+        .starts_with('{')
+}
+
+pub(crate) fn names_a_credentials_file(configured: &str) -> bool {
+    !is_inline_json(configured) || Path::new(configured).is_file()
+}
+
+const BASE64_ANY_PADDING: GeneralPurpose = GeneralPurpose::new(
+    &alphabet::STANDARD,
+    GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
+);
+
+pub(crate) fn is_base64_encoded_json(configured: &str) -> bool {
+    let compact: String = configured.split_whitespace().collect();
+    let quoted = compact.trim_matches(['"', '\'']);
+    let encoded = quoted.strip_prefix("base64:").unwrap_or(quoted);
+    BASE64_ANY_PADDING
+        .decode(encoded.replace('-', "+").replace('_', "/"))
+        .is_ok_and(|decoded| serde_json::from_slice::<Map<String, Value>>(&decoded).is_ok())
+}
+
 #[derive(Clone, Debug)]
 pub enum CredentialSource {
     Inline(SecretValue),

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

@@ -93,3 +93,22 @@ fn project_is_read_from_the_credentials_that_would_authenticate() {
         None
     );
 }
+
+#[rstest]
+fn project_is_read_from_a_relative_key_file_named_like_json() {
+    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 from_file = VertexConfig::from_params(&VertexParams {
+        vertex_credentials: Some(relative_name.into()),
+        ..VertexParams::default()
+    });
+    assert_eq!(
+        get_vertex_ai_project_from_credentials(&from_file, &|_| None).as_deref(),
+        Some("file-project")
+    );
+}

litellm-rust/crates/callbacks-legacy-python/src/adapter.rs

@@ -9,7 +9,7 @@ use litellm_host::{
     interceptors::{RawResponse, RequestContext, WireRequest},
     lifecycle::{FailureOrigin, Timing, epoch_seconds},
 };
-use litellm_host_python::{HookStep, from_py, missing_state, to_py};
+use litellm_host_python::{Error as HostError, HookStep, from_py, to_py};
 use pyo3::{
     exceptions::{PyBaseException, PyException},
     gc::{PyTraverseError, PyVisit},
@@ -160,7 +160,7 @@ impl LegacyLogging {
         self.response
             .as_ref()
             .map(|response| HookStep::Ready(response.clone_ref(py)))
-            .ok_or_else(missing_state)
+            .ok_or_else(|| HostError::MissingState.into())
     }
 
     fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
@@ -209,7 +209,7 @@ impl LegacyLogging {
 
     fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
         let logger = self.logger()?;
-        let billing = self.stream_billing().ok_or_else(missing_state)?;
+        let billing = self.stream_billing().ok_or(HostError::MissingState)?;
         let billed = Streaming::Success.call(
             py,
             (
@@ -525,7 +525,7 @@ impl LegacyLogging {
             return self.stream_failure(py);
         }
         if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() {
-            let error = self.error.as_ref().ok_or_else(missing_state)?;
+            let error = self.error.as_ref().ok_or(HostError::MissingState)?;
             return Ok(HookStep::Await(
                 DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?,
                 Self::resume_deployment_failure,
@@ -536,7 +536,7 @@ impl LegacyLogging {
 
     pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
         if self.stream_billing().is_none() {
-            return Err(missing_state());
+            return Err(HostError::MissingState.into());
         }
         if let Some(key) = &self.cache_key {
             head.bind(py).set_item("cache_key", key)?;
@@ -551,7 +551,7 @@ impl LegacyLogging {
     }
 
     pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
-        let stream = self.stream.as_mut().ok_or_else(missing_state)?;
+        let stream = self.stream.as_mut().ok_or(HostError::MissingState)?;
         if stream.first_chunk.is_none() {
             stream.first_chunk = Some(datetime(py, epoch_seconds())?);
         }

litellm-rust/crates/coroutine/AGENTS.md

@@ -3,7 +3,7 @@
 Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below. The alternatives section says which one each rejected design breaks
 
 - R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime
-- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
+- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/execution.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
 - R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime
 - R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states
 - R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for

litellm-rust/crates/host-python/AGENTS.md

@@ -6,7 +6,7 @@ This crate owns CPython execution mechanics for generic `litellm-host` machines
 
 `src/native.rs` belongs here: it runs a generic machine through the Python runtime and owns its pending execution and abort handle. Keep provider selection, request projection and public exception policy out of it. A rename to `machine_runner.rs` is optional and must not change behavior
 
-The execution handle receives its Python lifecycle binding from its consumer through `PythonLifecycle` rather than import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
+`run_call` returns the unstarted `Execution` and never names a Python stream class or driver. The consumer drives it from Python and supplies the step classes through `StepTypes` rather than have this crate import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
 
 Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer. `python-bridge::callable::wrap_failure` owns that policy
 
@@ -28,6 +28,7 @@ Terminal notifications share the selected response or exception. An ordinary not
   - `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites
   - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is
   - A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed
+- Consumers import this crate's error as `use litellm_host_python::Error as HostError`, so it never shadows their own `Error`
 - Use standard PyO3 ownership and conversion APIs
   - Prefer `Bound<'py, T>` for attached operations/results, `Py<T>` for retention. Binding/unbinding does not copy payloads
   - Use `pythonize` for selected Serde data, never a JSON-text round trip. Share conversion with `Pythonized<T>`
@@ -37,7 +38,7 @@ Terminal notifications share the selected response or exception. An ordinary not
   - Keep diagnostic counters in the consumer. Wrapper invocations do not measure every interpreter release
   - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers. Expose retained Python edges to GC without calling Python during traversal
 - Keep coroutine driving in the shared Python driver and the native handle
-  - Shared driver implementation: `litellm/rust_bridge/lifecycle.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
+  - Shared driver implementation: `litellm/rust_bridge/execution.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
   - Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract
 - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html)
   - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html)

litellm-rust/crates/host-python/Cargo.toml

@@ -6,11 +6,14 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+litellm-coroutine.workspace = true
 litellm-host.workspace = true
 
 bytes.workspace = true
 futures-util.workspace = true
+parking_lot.workspace = true
 serde.workspace = true
+thiserror.workspace = true
 tokio = { workspace = true, features = ["rt", "sync"] }
 
 pyo3.workspace = true

litellm-rust/crates/host-python/src/caller.rs

@@ -0,0 +1,58 @@
+use std::ops::ControlFlow;
+
+use litellm_coroutine::{Abandoned, Co, Reply};
+use pyo3::prelude::*;
+
+use crate::Error;
+use crate::handle::Suspension;
+
+pub(crate) type Answer = PyResult<Py<PyAny>>;
+
+pub(crate) struct Yielded {
+    pub(crate) suspension: Suspension,
+    pub(crate) reply: Reply<Answer>,
+}
+
+/// Hands suspensions to the Python caller and receives its answers.
+#[derive(Clone)]
+pub(crate) struct PythonCaller {
+    co: Co<Yielded>,
+}
+
+impl PythonCaller {
+    /// Connects the call body to its coroutine's yield handle.
+    pub(crate) fn new(co: Co<Yielded>) -> Self {
+        Self { co }
+    }
+
+    /// Asks the caller to await Python work and return its value or exception.
+    pub(crate) async fn await_python(&self, awaitable: Py<PyAny>) -> Answer {
+        self.suspend(Suspension::Await(awaitable)).await
+    }
+
+    /// Hands the caller a stream head and waits for demand or closure.
+    pub(crate) async fn open_stream(&self, head: Py<PyAny>) -> ControlFlow<()> {
+        self.stream_boundary(Suspension::StreamOpen(head)).await
+    }
+
+    /// Delivers a chunk and waits for the caller to request another or close.
+    pub(crate) async fn yield_chunk(&self, chunk: Py<PyAny>) -> ControlFlow<()> {
+        self.stream_boundary(Suspension::StreamChunk(chunk)).await
+    }
+
+    /// Maps stream demand to continuation and any error to detachment.
+    async fn stream_boundary(&self, suspension: Suspension) -> ControlFlow<()> {
+        match self.suspend(suspension).await {
+            Ok(_) => ControlFlow::Continue(()),
+            Err(_) => ControlFlow::Break(()),
+        }
+    }
+
+    /// Yields a suspension with its reply handle and waits for the answer.
+    async fn suspend(&self, suspension: Suspension) -> Answer {
+        match self.co.yield_(|reply| Yielded { suspension, reply }).await {
+            Ok(answer) => answer,
+            Err(Abandoned) => Err(Error::MissingState.into()),
+        }
+    }
+}

litellm-rust/crates/host-python/src/error.rs

@@ -1,20 +1,58 @@
-use pyo3::prelude::*;
+use pyo3::{
+    PyErr,
+    exceptions::{PyRuntimeError, PyTypeError},
+};
+
+/// A failure host-python raises itself, surfacing in Python as the exception `PyErr::from` picks.
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("missing native call state")]
+    MissingState,
+    #[error("execution is already running")]
+    ExecutionRunning,
+    #[error("execution is closed")]
+    ExecutionClosed,
+    #[error("can't send non-None value to a just-started execution")]
+    SendBeforeStart,
+    #[error("execution requires start before resume and can only start once")]
+    ExecutionOutOfOrder,
+    #[error("execution was closed while running")]
+    ExecutionClosedWhileRunning,
+    #[error("native execution closed")]
+    NativeExecutionClosed,
+    #[error("synchronous native routes cannot run from a Tokio context; use the async route")]
+    SyncInsideTokio,
+    #[error("the context body produced no value")]
+    ContextBodyProducedNothing,
+    #[error("file read must return bytes or str, got {got}")]
+    FileContentType { got: String },
+}
+
+impl From<Error> for PyErr {
+    fn from(error: Error) -> Self {
+        match error {
+            Error::SendBeforeStart | Error::FileContentType { .. } => {
+                PyTypeError::new_err(error.to_string())
+            }
+            Error::MissingState
+            | Error::ExecutionRunning
+            | Error::ExecutionClosed
+            | Error::ExecutionOutOfOrder
+            | Error::ExecutionClosedWhileRunning
+            | Error::NativeExecutionClosed
+            | Error::SyncInsideTokio
+            | Error::ContextBodyProducedNothing => PyRuntimeError::new_err(error.to_string()),
+        }
+    }
+}
 
 /// Why a custom operation the host answered did not produce a result: the route's own code
 /// rejected it, which the route classifies like any other native failure, or Python code
 /// raised, which reaches the caller as it was raised.
-#[derive(Debug)]
+#[derive(Debug, thiserror::Error)]
 pub enum InvokeError<E> {
+    #[error("{0}")]
     Native(E),
-    Python(PyErr),
-}
-
-impl<E> From<PyErr> for InvokeError<E> {
-    fn from(error: PyErr) -> Self {
-        Self::Python(error)
-    }
-}
-
-pub fn missing_state() -> PyErr {
-    pyo3::exceptions::PyRuntimeError::new_err("missing native call state")
+    #[error(transparent)]
+    Python(#[from] PyErr),
 }

litellm-rust/crates/host-python/src/file_reader.rs

@@ -3,13 +3,14 @@
 
 use bytes::Bytes;
 use pyo3::{
-    exceptions::PyTypeError,
     gc::{PyTraverseError, PyVisit},
     prelude::*,
     pybacked::PyBackedBytes,
     types::{PyBytes, PyString},
 };
 
+use crate::Error;
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct FileContent {
     pub bytes: Bytes,
@@ -50,10 +51,10 @@ impl PythonFileReader {
         } else if value.is_instance_of::<PyBytes>() {
             py_bytes(&value)?
         } else {
-            return Err(PyTypeError::new_err(format!(
-                "file read must return bytes or str, got {}",
-                value.get_type(),
-            )));
+            return Err(Error::FileContentType {
+                got: value.get_type().to_string(),
+            }
+            .into());
         };
         Ok(FileContent {
             bytes,

litellm-rust/crates/host-python/src/gil.rs

@@ -3,7 +3,9 @@ use std::sync::{
     atomic::{AtomicU64, Ordering},
 };
 
-use pyo3::{exceptions::PyRuntimeError, prelude::*};
+use pyo3::prelude::*;
+
+use crate::Error;
 
 /// The caller's `contextvars` context, captured at the Python entry point so blocking Python
 /// work started from Rust sees the same request-local values as the Python caller.
@@ -49,7 +51,7 @@ impl PythonContext {
             .lock()
             .expect("context value slot poisoned")
             .take()
-            .ok_or_else(|| PyRuntimeError::new_err("the context body produced no value"))
+            .ok_or_else(|| Error::ContextBodyProducedNothing.into())
     }
 }
 

litellm-rust/crates/host-python/src/handle.rs

@@ -1,19 +1,30 @@
 use std::panic::{AssertUnwindSafe, catch_unwind};
 
-use crate::panic_to_pyerr;
-use pyo3::exceptions::{PyBaseException, PyRuntimeError};
+use crate::{Error, panic_to_pyerr};
+use pyo3::exceptions::PyBaseException;
 use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 
-pub type PythonLifecycle = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
+pub type StepTypes = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
 
-pub enum ExecutionStep {
-    Return(Py<PyAny>),
+/// What a suspended execution waits on: an awaitable the caller awaits, or a stream
+/// boundary the caller's stream reads past.
+pub enum Suspension {
     Await(Py<PyAny>),
     /// The call streams: the caller gets a stream over this execution carrying this head,
     /// and the execution stays suspended until the stream asks for a chunk.
-    Open(Py<PyAny>),
-    Yield(Py<PyAny>),
+    StreamOpen(Py<PyAny>),
+    StreamChunk(Py<PyAny>),
+}
+
+enum Resume {
+    Send(Py<PyAny>),
+    Throw(PyErr),
+}
+
+pub enum ExecutionStep {
+    Suspended(Suspension),
+    Return(Py<PyAny>),
 }
 
 pub trait ExecutionBody: Send + Sync {
@@ -31,77 +42,59 @@ enum ExecutionState {
 #[pyclass]
 pub struct Execution {
     state: ExecutionState,
-    lifecycle: PythonLifecycle,
+    steps: StepTypes,
 }
 
 impl Execution {
-    pub fn new(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
+    pub fn new(body: impl ExecutionBody + 'static, steps: StepTypes) -> Self {
         Self {
             state: ExecutionState::Created(Box::new(body)),
-            lifecycle,
+            steps,
         }
     }
 
-    pub fn into_coroutine(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        let binding = (self.lifecycle)(py)?;
-        let execution = Py::new(py, self)?;
-        binding.getattr("drive")?.call1((execution,))
-    }
-
-    pub(crate) fn into_sync_stream(
-        self,
-        py: Python<'_>,
-        head: Py<PyAny>,
-    ) -> PyResult<Bound<'_, PyAny>> {
-        (self.lifecycle)(py)?
-            .getattr("SyncStream")?
-            .call1((Py::new(py, self)?, head))
-    }
-
-    /// An execution already started elsewhere and now waiting for its next input.
-    pub fn suspended(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
-        Self {
-            state: ExecutionState::Suspended(Box::new(body)),
-            lifecycle,
-        }
-    }
-
-    fn advance(
-        slf: &Bound<'_, Self>,
-        py: Python<'_>,
-        result: Option<PyResult<Py<PyAny>>>,
-    ) -> PyResult<Py<PyAny>> {
-        let mut body = {
+    fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult<Py<PyAny>> {
+        let (mut body, started) = {
             let mut execution = slf.borrow_mut();
-            match (&execution.state, result.is_some()) {
-                (ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {}
+            let started = match (&execution.state, &input) {
+                (ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false,
+                (ExecutionState::Created(_), Resume::Send(_)) => {
+                    return Err(Error::SendBeforeStart.into());
+                }
+                (ExecutionState::Suspended(_), _) => true,
                 (ExecutionState::Running, _) => {
-                    return Err(PyRuntimeError::new_err("execution is already running"));
+                    return Err(Error::ExecutionRunning.into());
                 }
                 (ExecutionState::Closed, _) => {
-                    return Err(PyRuntimeError::new_err("execution is closed"));
+                    return Err(Error::ExecutionClosed.into());
                 }
-                _ => {
-                    return Err(PyRuntimeError::new_err(
-                        "execution requires start before resume and can only start once",
-                    ));
+                (ExecutionState::Created(_), Resume::Throw(_)) => {
+                    return Err(Error::ExecutionOutOfOrder.into());
                 }
-            }
+            };
             match std::mem::replace(&mut execution.state, ExecutionState::Running) {
-                ExecutionState::Created(body) | ExecutionState::Suspended(body) => body,
+                ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started),
                 _ => unreachable!(),
             }
         };
-        let lifecycle = slf.borrow().lifecycle;
+        let result = match input {
+            Resume::Send(value) => started.then_some(Ok(value)),
+            Resume::Throw(error) => Some(Err(error)),
+        };
+        let steps = slf.borrow().steps;
         let outcome = catch_unwind(AssertUnwindSafe(|| {
             let step = body.resume(result)?;
             let (tag, value, suspended) = match step {
-                ExecutionStep::Await(value) => ("Await", value, true),
-                ExecutionStep::Open(head) => ("Open", head, true),
-                ExecutionStep::Yield(value) => ("Yield", value, true),
+                ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
+                ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
+                    ("StreamOpen", head, true)
+                }
+                ExecutionStep::Suspended(Suspension::StreamChunk(chunk)) => {
+                    ("StreamChunk", chunk, true)
+                }
                 ExecutionStep::Return(value) => ("Complete", value, false),
             };
-            let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind();
+            let step = steps(py)?.getattr(tag)?.call1((value,))?.unbind();
             Ok((step, suspended))
         }))
         .map_err(panic_to_pyerr)
@@ -116,9 +109,7 @@ impl Execution {
                 drop(body);
                 outcome.and_then(|(step, suspended)| {
                     if suspended {
-                        Err(PyRuntimeError::new_err(
-                            "execution was closed while running",
-                        ))
+                        Err(Error::ExecutionClosedWhileRunning.into())
                     } else {
                         Ok(step)
                     }
@@ -130,24 +121,16 @@ impl Execution {
 
 #[pymethods]
 impl Execution {
-    fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        Self::advance(slf, py, None)
+    fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
+        Self::advance(slf, py, Resume::Send(value))
     }
 
-    fn resume_value(
-        slf: &Bound<'_, Self>,
-        py: Python<'_>,
-        value: Py<PyAny>,
-    ) -> PyResult<Py<PyAny>> {
-        Self::advance(slf, py, Some(Ok(value)))
-    }
-
-    fn resume_error(
+    fn throw(
         slf: &Bound<'_, Self>,
         py: Python<'_>,
         error: Bound<'_, PyBaseException>,
     ) -> PyResult<Py<PyAny>> {
-        Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any()))))
+        Self::advance(slf, py, Resume::Throw(PyErr::from_value(error.into_any())))
     }
 
     fn close(slf: &Bound<'_, Self>) {

litellm-rust/crates/host-python/src/hooks/chain/adapter.rs

@@ -6,7 +6,7 @@ use pyo3::{
     types::PyDict,
 };
 
-use crate::{HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, missing_state};
+use crate::{Error, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned};
 
 pub(super) enum ChainStep<T> {
     Ready(T),
@@ -110,7 +110,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<ChainStep<Py<PyDict>>> {
         let Some(Continuation::Arguments(resume)) = self.continuation.take() else {
-            return Err(missing_state());
+            return Err(Error::MissingState.into());
         };
         let step = resume(&mut self.hooks, py, result)?;
         Ok(self.step(step, Continuation::Arguments))
@@ -132,7 +132,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<ChainStep<Box<WireRequest>>> {
         let Some(Continuation::Wire(resume)) = self.continuation.take() else {
-            return Err(missing_state());
+            return Err(Error::MissingState.into());
         };
         let step = resume(&mut self.hooks, py, result)?;
         Ok(self.step(step, Continuation::Wire))
@@ -154,7 +154,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<ChainStep<Py<PyAny>>> {
         let Some(Continuation::Response(resume)) = self.continuation.take() else {
-            return Err(missing_state());
+            return Err(Error::MissingState.into());
         };
         let step = resume(&mut self.hooks, py, result)?;
         Ok(self.step(step, Continuation::Response))
@@ -171,7 +171,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<ChainStep<()>> {
         let Some(Continuation::Event(resume)) = self.continuation.take() else {
-            return Err(missing_state());
+            return Err(Error::MissingState.into());
         };
         let step = resume(&mut self.hooks, py, result)?;
         Ok(self.step(step, Continuation::Event))

litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs

@@ -10,9 +10,7 @@ use pyo3::{
     types::PyDict,
 };
 
-use crate::{
-    HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime, missing_state,
-};
+use crate::{Error, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime};
 
 use super::adapter::{ChainHooks, ChainStep, HookAdapter};
 
@@ -72,7 +70,7 @@ impl HookChain {
         py: Python<'_>,
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<HookStep<Self, Py<PyDict>>> {
-        let (index, context) = self.arguments.take().ok_or_else(missing_state)?;
+        let (index, context) = self.arguments.take().ok_or(Error::MissingState)?;
         let result = resume_unless_cancelled(py, result)?;
         let step = self.hooks[index].resume_arguments(py, result)?;
         self.arguments_step(py, index, step, context)
@@ -113,7 +111,7 @@ impl HookChain {
         py: Python<'_>,
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
-        let (index, context) = self.wire.take().ok_or_else(missing_state)?;
+        let (index, context) = self.wire.take().ok_or(Error::MissingState)?;
         let result = resume_unless_cancelled(py, result)?;
         let step = self.hooks[index].resume_wire(py, result)?;
         self.wire_step(py, index, step, &context)
@@ -154,7 +152,7 @@ impl HookChain {
         py: Python<'_>,
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<HookStep<Self, Py<PyAny>>> {
-        let (index, context) = self.response.take().ok_or_else(missing_state)?;
+        let (index, context) = self.response.take().ok_or(Error::MissingState)?;
         let result = resume_unless_cancelled(py, result)?;
         let step = self.hooks[index].resume_response(py, result)?;
         self.response_step(py, index, step, context)
@@ -195,7 +193,7 @@ impl HookChain {
         py: Python<'_>,
         result: PyResult<Py<PyAny>>,
     ) -> PyResult<HookStep<Self, ()>> {
-        let (index, event) = self.event.take().ok_or_else(missing_state)?;
+        let (index, event) = self.event.take().ok_or(Error::MissingState)?;
         let result = resume_unless_cancelled(py, result)?;
         let result = self.hooks[index].resume_event(py, result);
         let step = notification_result(py, is_terminal(&event), result)?;

litellm-rust/crates/host-python/src/lib.rs

@@ -5,6 +5,7 @@
 
 mod argument;
 mod binding;
+mod caller;
 mod conversion_cache;
 mod driver;
 mod error;
@@ -23,11 +24,11 @@ pub use argument::{effective_py_args, lookup, present};
 pub use binding::PythonBinding;
 pub use conversion_cache::{FromPythonCache, ToPythonCache};
 pub use driver::{CallOptions, run_call};
-pub use error::{InvokeError, missing_state};
+pub use error::{Error, InvokeError};
 pub use file_reader::{FileContent, PythonFileReader, py_bytes};
 pub use fork_gate::RuntimeAlreadyStarted;
 pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
-pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle};
+pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension};
 pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
 pub use marshal::{
     Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,

litellm-rust/crates/host-python/src/native.rs

@@ -5,11 +5,10 @@ use futures_util::future::{AbortHandle, Abortable};
 use litellm_host::call::HostedCompletion;
 use litellm_host::machine::{HostFailure, Machine, MachineStep};
 use litellm_host::protocol::Protocol;
-use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use tokio::sync::Mutex;
 
-use crate::missing_state;
+use crate::Error;
 use crate::runtime::{poll_async_value, run_async_value, run_sync_value};
 
 type NativeResult<M> = Result<
@@ -65,7 +64,7 @@ where
         py: Python<'_>,
         interruption: Option<HostFailure<<M::Protocol as Protocol>::Error>>,
     ) -> PyResult<NativePoll<NativeResult<M>>> {
-        let state = Arc::clone(self.state.as_ref().ok_or_else(missing_state)?);
+        let state = Arc::clone(self.state.as_ref().ok_or(Error::MissingState)?);
         let future = async move {
             let mut state = state.lock().await;
             let result = match interruption {
@@ -90,7 +89,7 @@ where
                 run_async_value(py, async move {
                     Abortable::new(future, registration)
                         .await
-                        .map_err(|_| PyRuntimeError::new_err("native execution closed"))?
+                        .map_err(|_| Error::NativeExecutionClosed)?
                 })?
                 .unbind(),
             ))
@@ -103,12 +102,12 @@ where
     pub(super) fn take_result(&self) -> PyResult<NativeResult<M>> {
         self.state
             .as_ref()
-            .ok_or_else(missing_state)?
+            .ok_or(Error::MissingState)?
             .try_lock()
-            .map_err(|_| missing_state())?
+            .map_err(|_| Error::MissingState)?
             .result
             .take()
-            .ok_or_else(missing_state)
+            .ok_or_else(|| Error::MissingState.into())
             .map(|result| {
                 result.map(|step| match step {
                     MachineStep::Suspended(op) => MachineStep::Suspended(op),

litellm-rust/crates/host-python/src/runtime.rs

@@ -5,7 +5,7 @@ use std::task::{Context, Poll, Waker};
 use std::time::Duration;
 
 use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
-use crate::{Pythonized, panic_to_pyerr, release_gil};
+use crate::{Error, Pythonized, panic_to_pyerr, release_gil};
 use futures_util::FutureExt;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -112,9 +112,7 @@ where
     F: Future<Output = PyResult<T>> + Send + 'static,
 {
     if Handle::try_current().is_ok() {
-        return Err(PyRuntimeError::new_err(
-            "synchronous native routes cannot run from a Tokio context; use the async route",
-        ));
+        return Err(Error::SyncInsideTokio.into());
     }
     release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?
 }
@@ -131,9 +129,7 @@ where
     F: Future<Output = Result<T, E>> + Send + 'static,
 {
     if Handle::try_current().is_ok() {
-        return Err(PyRuntimeError::new_err(
-            "synchronous native routes cannot run from a Tokio context; use the async route",
-        ));
+        return Err(Error::SyncInsideTokio.into());
     }
 
     let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;

litellm-rust/crates/host-python/tests/driver.rs

@@ -0,0 +1,770 @@
+mod support;
+
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use litellm_host::{
+    hooks::CallHooks,
+    lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
+    machine::CallMachine,
+    observation::observation_channel,
+};
+use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
+use pyo3::{
+    exceptions::{
+        PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
+        PyValueError, asyncio::CancelledError,
+    },
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::PyDict,
+};
+use rstest::rstest;
+use support::{
+    HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
+    failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
+    run_composed, run_scripted, settled, success_machine,
+};
+
+fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
+    let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
+    let options = CallOptions {
+        asynchronous: false,
+        steps: lifecycle_module,
+        observers: Some(sender),
+    };
+    (options, move || {
+        std::iter::from_fn(|| receiver.try_recv().ok()).collect()
+    })
+}
+
+#[rstest]
+#[case::success(OpScript::AwaitAnswer)]
+#[case::failure(OpScript::AwaitFailure)]
+fn asynchronous_host_operations_resume_the_same_machine(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] op: OpScript,
+) {
+    Python::attach(|py| {
+        let machine: MachineFactory = Box::new(|_| {
+            CallMachine::<Synthetic>::new(None, |host| {
+                Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
+            })
+        });
+        let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
+        match op {
+            OpScript::AwaitAnswer => {
+                assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
+            }
+            OpScript::AwaitFailure => {
+                assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
+                assert_eq!(
+                    log.iter()
+                        .filter(|entry| entry.starts_with("failed:"))
+                        .count(),
+                    1
+                );
+                assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
+            }
+            _ => unreachable!(),
+        }
+    });
+}
+
+#[rstest]
+fn composed_hooks_share_prepared_arguments_and_one_execution(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let chain = HookChain::new()
+            .with(hooks(&log, HookScript::ObserveArguments))
+            .with(hooks(&log, HookScript::RewriteArguments))
+            .with(hooks(&log, HookScript::ReplaceResponse));
+        let result = handed(
+            py,
+            run_call(
+                py,
+                move |_, _, request| Ok(success_machine()(request)),
+                binding(&log, OpScript::Answer),
+                chain,
+                PyDict::new(py).unbind(),
+                call_options(asynchronous),
+            ),
+            asynchronous,
+        );
+        if asynchronous {
+            assert!(log.entries().is_empty());
+        }
+        let response = settled(py, result, asynchronous).unwrap();
+        assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
+        let entries = log.entries();
+        assert!(entries.iter().any(|entry| entry == "adopted:hook"));
+        assert_eq!(
+            entries.iter().filter(|entry| *entry == "project").count(),
+            1
+        );
+        assert_eq!(
+            entries.iter().filter(|entry| *entry == "op:sign").count(),
+            1
+        );
+        assert_eq!(
+            entries
+                .iter()
+                .filter(|entry| *entry == "succeeded:replaced")
+                .count(),
+            3
+        );
+        assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
+    });
+}
+
+#[rstest]
+#[case::draining(4, false)]
+#[case::full(1, false)]
+#[case::closed(4, true)]
+fn observation_delivery_preserves_callback_order_and_response(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] capacity: usize,
+    #[case] closed: bool,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (expected, expected_log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::ReplaceResponse,
+            asynchronous,
+        );
+        let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
+        if closed {
+            receiver.close();
+        }
+        let log = Log::default();
+        let (actual, actual_log) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::ReplaceResponse),
+            std::convert::identity,
+            CallOptions {
+                asynchronous,
+                steps: lifecycle_module,
+                observers: Some(sender.clone()),
+            },
+        );
+        assert_eq!(
+            actual.unwrap().extract::<String>(py).unwrap(),
+            expected.unwrap().extract::<String>(py).unwrap()
+        );
+        assert_eq!(actual_log, expected_log);
+        let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
+        assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
+        if !closed && capacity >= 3 {
+            let [
+                CallEvent::Started { start_time },
+                CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
+                CallEvent::Succeeded { timing, .. },
+            ] = events.as_slice()
+            else {
+                panic!("expected start, provider response and success: {events:?}");
+            };
+            assert_eq!(raw.body, "raw");
+            assert_eq!(*start_time, timing.start_time);
+            assert!(timing.end_time >= timing.start_time);
+        }
+    });
+}
+
+#[rstest]
+#[case::preparation(HookScript::FailBegin, false, 1)]
+#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
+#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
+#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
+fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] script: HookScript,
+    #[case] cancelled: bool,
+    #[case] failure_callbacks: usize,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (options, mut events) = observed(4);
+        let log = Log::default();
+        let (result, log) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, script),
+            std::convert::identity,
+            CallOptions {
+                asynchronous,
+                ..options
+            },
+        );
+        let error = result.unwrap_err();
+        assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
+        let events = events();
+        assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
+        assert_eq!(
+            matches!(events.last(), Some(CallEvent::Cancelled { .. })),
+            cancelled
+        );
+        assert_eq!(
+            matches!(
+                events.last(),
+                Some(CallEvent::Failed {
+                    origin: FailureOrigin::Host,
+                    ..
+                })
+            ),
+            !cancelled
+        );
+        let terminal = events.iter().filter(|event| {
+            matches!(
+                event,
+                CallEvent::Failed { .. }
+                    | CallEvent::Succeeded { .. }
+                    | CallEvent::Cancelled { .. }
+            )
+        });
+        assert_eq!(terminal.count(), 1);
+        assert_eq!(
+            log.iter()
+                .filter(|entry| entry.starts_with("failed:"))
+                .count(),
+            failure_callbacks
+        );
+        assert!(
+            !log.iter()
+                .any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
+        );
+    });
+}
+
+#[rstest]
+#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
+#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
+#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
+fn startup_failures_release_the_factory_without_constructing_a_machine(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] script: HookScript,
+    #[case] op: OpScript,
+    #[values(false, true)] asynchronous: bool,
+) {
+    struct Release(Arc<AtomicBool>);
+    impl Drop for Release {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+    Python::attach(|py| {
+        let constructed = Arc::new(AtomicBool::new(false));
+        let did_construct = constructed.clone();
+        let released = Arc::new(AtomicBool::new(false));
+        let release = Release(released.clone());
+        let machine: MachineFactory = Box::new(move |request| {
+            let _release = release;
+            did_construct.store(true, Ordering::SeqCst);
+            success_machine()(request)
+        });
+        let (result, log) = run_scripted(py, machine, op, script, asynchronous);
+        assert!(result.is_err());
+        assert!(!constructed.load(Ordering::SeqCst));
+        assert!(released.load(Ordering::SeqCst));
+        assert_eq!(
+            log.iter()
+                .filter(|event| event.starts_with("failed:"))
+                .count(),
+            1
+        );
+        assert_eq!(
+            log.iter().filter(|event| *event == "adapter.close").count(),
+            1
+        );
+        assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
+    });
+}
+
+#[rstest]
+fn success_runs_every_step_in_order_and_returns_the_public_response(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::Plain,
+            asynchronous,
+        );
+        assert_eq!(
+            result.unwrap().extract::<String>(py).unwrap(),
+            "project:1|sign|rewritten"
+        );
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "project",
+                "op:sign",
+                "before_provider_request",
+                "response:raw",
+                "complete",
+                "after_success",
+                "succeeded:project:1|sign|rewritten",
+                "adapter.close",
+                "host.close",
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let (options, mut events) = observed(3);
+        let stream = open_stream(
+            py,
+            hooks(&log, HookScript::Plain),
+            CallOptions {
+                asynchronous,
+                ..options
+            },
+        );
+        if asynchronous {
+            for method in ["__anext__", "aclose", "aclose"] {
+                let stop = stream
+                    .call_method0(method)
+                    .unwrap()
+                    .call_method1("send", (py.None(),))
+                    .unwrap_err();
+                assert!(stop.is_instance_of::<PyStopIteration>(py));
+            }
+        } else {
+            assert_eq!(
+                stream
+                    .call_method0("__next__")
+                    .unwrap()
+                    .extract::<String>()
+                    .unwrap(),
+                "first"
+            );
+            stream.call_method0("close").unwrap();
+            stream.call_method0("close").unwrap();
+        }
+        assert!(matches!(
+            events().as_slice(),
+            [CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
+        ));
+        assert_eq!(
+            log.entries(),
+            [
+                "started",
+                "begin",
+                "opened",
+                "delivered",
+                "succeeded:None",
+                "adapter.close"
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn a_stream_carries_its_head_before_the_first_chunk(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let stream = open_stream(
+            py,
+            hooks(&log, HookScript::Plain),
+            call_options(asynchronous),
+        );
+        let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
+            stream.getattr("_hidden_params").unwrap().extract().unwrap();
+        assert_eq!(
+            hidden["additional_headers"],
+            std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
+        );
+        assert_eq!(log.entries(), ["started", "begin", "opened"]);
+        assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
+    });
+}
+
+struct Classification {
+    provider_fails: bool,
+    op: OpScript,
+    classifier_fails: bool,
+    message: &'static str,
+    native_context: Option<&'static str>,
+    steps: &'static [&'static str],
+}
+
+#[rstest]
+#[case::provider_failure(Classification {
+    provider_fails: true,
+    op: OpScript::Answer,
+    classifier_fails: false,
+    message: "classified: provider exploded",
+    native_context: None,
+    steps: &["project", "classify:provider exploded"],
+})]
+#[case::host_op_rejection(Classification {
+    provider_fails: false,
+    op: OpScript::RejectNatively,
+    classifier_fails: false,
+    message: "classified: op rejected",
+    native_context: None,
+    steps: &["project", "op:sign", "classify:op rejected"],
+})]
+#[case::host_op_python_exception(Classification {
+    provider_fails: false,
+    op: OpScript::RaisePython,
+    classifier_fails: false,
+    message: "op failed",
+    native_context: None,
+    steps: &["project", "op:sign"],
+})]
+#[case::failing_classifier(Classification {
+    provider_fails: true,
+    op: OpScript::Answer,
+    classifier_fails: true,
+    message: "classifier failed",
+    native_context: Some("provider exploded"),
+    steps: &["project", "classify:provider exploded"],
+})]
+fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] case: Classification,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let Classification {
+            provider_fails,
+            op,
+            classifier_fails,
+            message,
+            native_context,
+            steps,
+        } = case;
+        let log = Log::default();
+        let binding = if classifier_fails {
+            binding(&log, op).with_failing_classifier()
+        } else {
+            binding(&log, op)
+        };
+        let machine = if provider_fails {
+            failing_machine()
+        } else {
+            success_machine()
+        };
+        let (result, log) = run_composed(
+            py,
+            machine,
+            binding,
+            hooks(&log, HookScript::Plain),
+            std::convert::identity,
+            call_options(asynchronous),
+        );
+        let error = result.unwrap_err();
+        assert_eq!(error.value(py).to_string(), message);
+        match native_context {
+            Some(context) => {
+                assert!(error.is_instance_of::<PyTypeError>(py));
+                let chained = error.context(py).unwrap();
+                assert!(chained.is_instance_of::<PyRuntimeError>(py));
+                assert_eq!(chained.value(py).to_string(), context);
+            }
+            None => assert!(error.is_instance_of::<PyValueError>(py)),
+        }
+        let expected: Vec<String> = ["started", "begin"]
+            .iter()
+            .chain(steps)
+            .map(|step| step.to_string())
+            .chain([
+                format!("failed:Call:{message}"),
+                "adapter.close".into(),
+                "host.close".into(),
+            ])
+            .collect();
+        assert_eq!(log, expected);
+    });
+}
+
+#[rstest]
+fn begin_failures_are_host_failures_without_provider_mapping(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::FailBegin,
+            false,
+        );
+        assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "failed:Host:begin failed",
+                "adapter.close",
+                "host.close"
+            ]
+        );
+    });
+}
+
+enum ArgumentPolicy {
+    Inherit,
+    Reject(Py<PyBaseException>),
+}
+
+impl CallHooks<PythonRuntime> for ArgumentPolicy {
+    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
+        match self {
+            Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
+            Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
+        }
+    }
+}
+
+impl PythonOwned for ArgumentPolicy {
+    fn close(&mut self, _: Python<'_>) {}
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        match self {
+            Self::Inherit => Ok(()),
+            Self::Reject(error) => visit.call(error),
+        }
+    }
+}
+
+#[rstest]
+fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+    #[values(false, true)] fail_setup: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let setup_log = log.clone();
+        let error = PyValueError::new_err("setup failed").into_value(py);
+        let setup_error = error.clone_ref(py);
+        let arguments = PyDict::new(py);
+        arguments.set_item("api_key", "original").unwrap();
+        let result = handed(
+            py,
+            run_call(
+                py,
+                move |py, prepared, request| {
+                    let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
+                    let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
+                    setup_log.push(format!("setup:{source}:{key}"));
+                    if fail_setup {
+                        return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
+                    }
+                    Ok(success_machine()(request))
+                },
+                binding(&log, OpScript::Answer),
+                HookChain::new()
+                    .with(hooks(&log, HookScript::RewriteArguments))
+                    .with(ArgumentPolicy::Inherit),
+                arguments.clone().unbind(),
+                call_options(asynchronous),
+            ),
+            asynchronous,
+        );
+        if asynchronous {
+            assert!(log.entries().is_empty());
+        }
+        let settled = settled(py, result, asynchronous);
+        assert_eq!(
+            arguments
+                .get_item("api_key")
+                .unwrap()
+                .unwrap()
+                .extract::<String>()
+                .unwrap(),
+            "original"
+        );
+        assert_eq!(
+            &log.entries()[..4],
+            ["started", "begin", "project", "setup:hook:inherited"]
+        );
+        if fail_setup {
+            assert!(settled.unwrap_err().value(py).is(error.bind(py)));
+            assert_eq!(
+                log.entries(),
+                [
+                    "started",
+                    "begin",
+                    "project",
+                    "setup:hook:inherited",
+                    "failed:Host:setup failed",
+                    "adapter.close",
+                    "host.close",
+                ]
+            );
+        } else {
+            assert!(settled.is_ok());
+            assert!(
+                log.entries()
+                    .iter()
+                    .any(|entry| entry.starts_with("succeeded:"))
+            );
+        }
+    });
+}
+
+#[rstest]
+fn closing_an_unstarted_call_never_initializes_resources(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let setup_log = log.clone();
+        let pending = handed(
+            py,
+            run_call(
+                py,
+                move |_, _, request| {
+                    setup_log.push("setup");
+                    Ok(success_machine()(request))
+                },
+                binding(&log, OpScript::Answer),
+                hooks(&log, HookScript::Plain),
+                PyDict::new(py).unbind(),
+                call_options(true),
+            ),
+            true,
+        )
+        .unwrap();
+        assert!(log.entries().is_empty());
+        pending.call_method0(py, "close").unwrap();
+        drop(pending);
+        assert_eq!(log.entries(), ["adapter.close", "host.close"]);
+    });
+}
+
+#[rstest]
+fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let raised = PyValueError::new_err("over budget").into_value(py);
+        let log = Log::default();
+        let (result, log) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::Plain),
+            |hooks| {
+                HookChain::new()
+                    .with(hooks)
+                    .with(ArgumentPolicy::Reject(raised.clone_ref(py)))
+            },
+            call_options(asynchronous),
+        );
+        assert!(result.unwrap_err().value(py).is(raised.bind(py)));
+        assert_eq!(
+            log,
+            [
+                "started",
+                "begin",
+                "failed:Host:over budget",
+                "adapter.close",
+                "host.close"
+            ]
+        );
+    });
+}
+
+#[rstest]
+fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let log = Log::default();
+        let (result, _) = run_composed(
+            py,
+            success_machine(),
+            binding(&log, OpScript::Answer),
+            hooks(&log, HookScript::Plain),
+            |hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
+            call_options(asynchronous),
+        );
+        assert_eq!(
+            result.unwrap().extract::<String>(py).unwrap(),
+            "project:2|sign|rewritten"
+        );
+    });
+}
+
+#[rstest]
+fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::Answer,
+            HookScript::FailAfterSuccess,
+            asynchronous,
+        );
+        assert_eq!(
+            result.unwrap_err().value(py).to_string(),
+            "after_success failed"
+        );
+        assert_eq!(
+            &log[log.len() - 4..],
+            [
+                "after_success",
+                "failed:Host:after_success failed",
+                "adapter.close",
+                "host.close"
+            ]
+        );
+        assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
+    });
+}
+
+#[rstest]
+fn cancellation_ends_the_call_without_terminal_dispatch(
+    #[from(interpreter)] _python: Interpreter,
+    #[values(false, true)] asynchronous: bool,
+) {
+    Python::attach(|py| {
+        let (result, log) = run_scripted(
+            py,
+            success_machine(),
+            OpScript::CancelRequest,
+            HookScript::Plain,
+            asynchronous,
+        );
+        let error = result.unwrap_err();
+        assert!(!error.is_instance_of::<PyException>(py));
+        assert_eq!(
+            log,
+            ["started", "begin", "project", "adapter.close", "host.close"]
+        );
+    });
+}

litellm-rust/crates/host-python/tests/handle.rs

@@ -0,0 +1,343 @@
+mod support;
+
+use litellm_host::{hooks::CallHooks, lifecycle::Timing};
+use litellm_host_python::{
+    Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
+    Suspension, run_call,
+};
+use pyo3::{
+    exceptions::PyBaseException,
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::{PyCFunction, PyDict},
+};
+use rstest::rstest;
+use support::{
+    Interpreter, Log, OpScript, binding, call_options, handed, install_lifecycle_module,
+    interpreter, lifecycle_module, success_machine,
+};
+
+struct RetainingHost {
+    retained: Option<Py<PyAny>>,
+}
+
+impl ExecutionBody for RetainingHost {
+    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.retained)
+    }
+}
+
+#[pyfunction]
+fn retaining_coroutine(py: Python<'_>, retained: Py<PyAny>) -> PyResult<Py<Execution>> {
+    Py::new(
+        py,
+        Execution::new(
+            RetainingHost {
+                retained: Some(retained),
+            },
+            lifecycle_module,
+        ),
+    )
+}
+
+struct AwaitBody(Option<Py<PyAny>>);
+
+impl ExecutionBody for AwaitBody {
+    fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        match self.0.take() {
+            Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))),
+            None => result
+                .expect("selected await completed")
+                .map(ExecutionStep::Return),
+        }
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+}
+
+#[pyfunction]
+fn await_execution(awaitable: Py<PyAny>) -> Execution {
+    Execution::new(AwaitBody(Some(awaitable)), lifecycle_module)
+}
+
+struct CallingBody(Py<PyAny>);
+
+impl ExecutionBody for CallingBody {
+    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return))
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+}
+
+#[pyfunction]
+fn calling_execution(callback: Py<PyAny>) -> Execution {
+    Execution::new(CallingBody(callback), lifecycle_module)
+}
+
+struct ErrorBody(Option<Py<PyBaseException>>);
+
+impl ExecutionBody for ErrorBody {
+    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        Python::attach(|py| {
+            Err(PyErr::from_value(
+                self.0.take().unwrap().into_bound(py).into_any(),
+            ))
+        })
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.0)
+    }
+}
+
+#[pyfunction]
+fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
+    Execution::new(ErrorBody(Some(error.unbind())), lifecycle_module)
+}
+
+/// Where a paused call keeps the test's object: in its hooks, or selected as the call's
+/// response or exception, with the call paused in the hook that is told about it.
+#[derive(Clone, Copy, Debug)]
+enum Holder {
+    Hooks,
+    Response,
+    Exception,
+}
+
+struct PausingHooks {
+    retained: Option<Py<PyAny>>,
+    pause: Option<Py<PyAny>>,
+    holder: Holder,
+}
+
+impl CallHooks<PythonRuntime> for PausingHooks {
+    fn transform_response(
+        &mut self,
+        py: Python<'_>,
+        response: Py<PyAny>,
+        _: Timing,
+    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
+        match self.holder {
+            Holder::Hooks => Ok(HookStep::Ready(response)),
+            Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
+            Holder::Exception => Err(PyErr::from_value(
+                self.retained.take().unwrap().into_bound(py),
+            )),
+        }
+    }
+
+    fn on_event(
+        &mut self,
+        _: Python<'_>,
+        event: PythonCallEvent<'_>,
+    ) -> PyResult<HookStep<Self, ()>> {
+        let pauses_here = match self.holder {
+            Holder::Hooks => true,
+            Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
+            Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
+        };
+        Ok(match self.pause.take_if(|_| pauses_here) {
+            Some(pause) => {
+                HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
+            }
+            None => HookStep::Ready(()),
+        })
+    }
+}
+
+impl PythonOwned for PausingHooks {
+    fn close(&mut self, _: Python<'_>) {
+        self.retained = None;
+        self.pause = None;
+    }
+
+    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.retained)?;
+        visit.call(&self.pause)
+    }
+}
+
+fn paused_call(
+    py: Python<'_>,
+    retained: Py<PyAny>,
+    pause: Py<PyAny>,
+    holder: Holder,
+) -> PyResult<Py<PyAny>> {
+    let machine = success_machine();
+    handed(
+        py,
+        run_call(
+            py,
+            move |_, _, request| Ok(machine(request)),
+            binding(&Log::default(), OpScript::Answer),
+            PausingHooks {
+                retained: Some(retained),
+                pause: Some(pause),
+                holder,
+            },
+            PyDict::new(py).unbind(),
+            call_options(true),
+        ),
+        true,
+    )
+}
+
+#[rstest]
+fn every_suspension_is_awaited_inline_on_the_callers_task(
+    #[from(interpreter)] _python: Interpreter,
+) {
+    Python::attach(|py| {
+        let module = install_lifecycle_module(py);
+        let locals = PyDict::new(py);
+        locals
+            .set_item("drive", module.getattr("drive").unwrap())
+            .unwrap();
+        locals
+            .set_item(
+                "await_execution",
+                wrap_pyfunction!(await_execution, py).unwrap(),
+            )
+            .unwrap();
+        locals
+            .set_item(
+                "calling_execution",
+                wrap_pyfunction!(calling_execution, py).unwrap(),
+            )
+            .unwrap();
+        let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
+        py.run(&probe, Some(&locals), Some(&locals)).unwrap();
+    });
+}
+
+#[rstest]
+fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
+    Python::attach(|py| {
+        let locals = PyDict::new(py);
+        locals
+            .set_item(
+                "error_execution",
+                wrap_pyfunction!(error_execution, py).unwrap(),
+            )
+            .unwrap();
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import gc
+import weakref
+
+class Retained:
+    pass
+
+def cycle():
+    retained = Retained()
+    try:
+        raise ValueError('retained traceback')
+    except ValueError as error:
+        retained.owner = error_execution(error)
+    return weakref.ref(retained)
+
+reference = cycle()
+gc.collect()
+assert reference() is None
+"#
+            ),
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+    });
+}
+
+#[rstest]
+fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) {
+    Python::attach(|py| {
+        let locals = PyDict::new(py);
+        locals
+            .set_item(
+                "retaining_coroutine",
+                wrap_pyfunction!(retaining_coroutine, py).unwrap(),
+            )
+            .unwrap();
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import gc
+import weakref
+
+class Retained:
+    pass
+
+def cycle():
+    retained = Retained()
+    coroutine = retaining_coroutine(retained)
+    retained.coroutine = coroutine
+    return weakref.ref(retained)
+
+retained_ref = cycle()
+gc.collect()
+assert retained_ref() is None
+"#
+            ),
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+    });
+}
+
+#[rstest]
+#[case::held_by_the_hooks(Holder::Hooks)]
+#[case::selected_as_the_response(Holder::Response)]
+#[case::selected_as_the_exception(Holder::Exception)]
+fn cycles_through_a_call_suspended_in_a_hook_are_collectable(
+    #[from(interpreter)] _python: Interpreter,
+    #[case] holder: Holder,
+) {
+    Python::attach(|py| {
+        let locals = PyDict::new(py);
+        let call = PyCFunction::new_closure(py, None, None, move |args, _| {
+            let (retained, pause) = args.extract::<(Py<PyAny>, Py<PyAny>)>()?;
+            paused_call(args.py(), retained, pause, holder)
+        })
+        .unwrap();
+        locals.set_item("paused_call", call).unwrap();
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import gc
+import weakref
+
+class Retained(Exception):
+    pass
+
+class Pause:
+    def __await__(self):
+        yield
+
+def cycle():
+    retained = Retained()
+    coroutine = paused_call(retained, Pause())
+    coroutine.send(None)
+    retained.coroutine = coroutine
+    return weakref.ref(retained)
+
+reference = cycle()
+gc.collect()
+assert reference() is None
+"#
+            ),
+            Some(&locals),
+            Some(&locals),
+        )
+        .unwrap();
+    });
+}

litellm-rust/crates/host-python/tests/hook_chain.rs

@@ -504,7 +504,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
                 }
                 HookStep::Await(awaitable, resume) => {
                     self.resume = Some(resume);
-                    Ok(litellm_host_python::ExecutionStep::Await(awaitable))
+                    Ok(litellm_host_python::ExecutionStep::Suspended(
+                        litellm_host_python::Suspension::Await(awaitable),
+                    ))
                 }
             }
         })
@@ -518,9 +520,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
 
 fn lifecycle(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
     let source =
-        std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/lifecycle.py"))
+        std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/execution.py"))
             .unwrap();
-    PyModule::from_code(py, &source, c"lifecycle.py", c"hook_chain_test_lifecycle")
+    PyModule::from_code(py, &source, c"execution.py", c"hook_chain_test_lifecycle")
 }
 
 #[rstest]
@@ -556,9 +558,16 @@ second = Hooks('b', invoke)
             },
             lifecycle,
         );
-        locals
-            .set_item("call", execution.into_coroutine(py).unwrap())
+        let module = lifecycle(py).unwrap();
+        let call = module
+            .getattr("drive")
+            .unwrap()
+            .call1((
+                Py::new(py, execution).unwrap(),
+                module.getattr("AsyncStream").unwrap(),
+            ))
             .unwrap();
+        locals.set_item("call", call).unwrap();
         py.run(
             c"
 async def exercise():

litellm-rust/crates/host-python/tests/lifecycle.py

@@ -36,11 +36,17 @@ async def exercise():
     releaser = asyncio.create_task(release())
     execution = await_execution(CustomAwaitable())
     try:
-        execution.resume_value(None)
+        execution.throw(ValueError("early"))
     except RuntimeError:
         pass
     else:
         raise AssertionError("resumed an unstarted execution")
+    try:
+        execution.send(1)
+    except TypeError:
+        pass
+    else:
+        raise AssertionError("sent a value into an unstarted execution")
     wrapped = drive(execution)
     try:
         wrapped.send(1)
@@ -129,7 +135,7 @@ async def exercise():
 
     def reenter():
         try:
-            active.start()
+            active.send(None)
         except RuntimeError as error:
             observed.append(str(error))
         return result
@@ -151,7 +157,7 @@ async def exercise():
         execution = calling_execution(callback)
         callback.owner = execution
         if started:
-            assert execution.start().value is result
+            assert execution.send(None).value is result
         return weakref.ref(callback)
 
     for started in (False, True):

litellm-rust/crates/host-python/tests/support/mod.rs

@@ -0,0 +1,684 @@
+#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
+
+use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
+
+use litellm_host::{
+    call::{CallOutput, HostedMachine, hosted_call},
+    hooks::CallHooks,
+    interceptors::{Interceptors, RawResponse, RequestContext, WireRequest},
+    lifecycle::{ExecutionEvent, Timing},
+    machine::{CallMachine, MachineFault},
+    protocol::{Protocol, Reply},
+};
+use litellm_host_python::{
+    CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
+    PythonHostCalls, PythonOwned, PythonRuntime, run_call,
+};
+use pyo3::{
+    exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
+    gc::{PyTraverseError, PyVisit},
+    prelude::*,
+    types::{PyDict, PyString},
+};
+use rstest::fixture;
+
+static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
+
+/// One initialized interpreter with the real lifecycle protocol installed, held for the
+/// length of a test so tests sharing the interpreter run one at a time.
+pub struct Interpreter(MutexGuard<'static, ()>);
+
+#[fixture]
+pub fn interpreter() -> Interpreter {
+    let guard = PYTHON_GLOBALS
+        .lock()
+        .unwrap_or_else(PoisonError::into_inner);
+    Python::initialize();
+    Python::attach(|py| {
+        py.import("asyncio").unwrap();
+        install_lifecycle_module(py);
+    });
+    Interpreter(guard)
+}
+
+pub fn lifecycle_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
+    py.import("host_test_lifecycle")
+}
+
+/// Hands a call's execution to its caller the way the product entrypoints do: an async call
+/// gets the coroutine from `drive`, a sync call the value or stream `run` settles to.
+pub fn handed(
+    py: Python<'_>,
+    execution: PyResult<Py<PyAny>>,
+    asynchronous: bool,
+) -> PyResult<Py<PyAny>> {
+    let execution = execution?;
+    let module = lifecycle_module(py)?;
+    let handed = if asynchronous {
+        module.getattr("drive")?.call1((execution,))?
+    } else {
+        module
+            .getattr("run")?
+            .call1((execution, module.getattr("SyncStream")?))?
+    };
+    Ok(handed.unbind())
+}
+
+pub fn call_options(asynchronous: bool) -> CallOptions {
+    CallOptions::new(asynchronous, lifecycle_module)
+}
+
+pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
+    let source = std::ffi::CString::new(concat!(
+        include_str!("../../../../../litellm/rust_bridge/execution.py"),
+        "\n_drive = drive\n\n\nasync def drive(execution):\n    return await _drive(execution, AsyncStream)\n",
+    ))
+    .unwrap();
+    PyModule::from_code(
+        py,
+        &source,
+        pyo3::ffi::c_str!("execution.py"),
+        pyo3::ffi::c_str!("host_test_lifecycle"),
+    )
+    .unwrap()
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Error(pub String);
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter.write_str(&self.0)
+    }
+}
+
+impl From<MachineFault> for Error {
+    fn from(fault: MachineFault) -> Self {
+        Self(format!("{fault:?}"))
+    }
+}
+
+pub struct Synthetic;
+
+impl Protocol for Synthetic {
+    type Response = String;
+    type Error = Error;
+    type Request = String;
+    type HostCall = (&'static str, Reply<String>);
+    type Chunk = std::convert::Infallible;
+    type StreamHead = std::convert::Infallible;
+}
+
+pub fn wire() -> WireRequest {
+    WireRequest {
+        url: "https://example.invalid".into(),
+        headers: Vec::new(),
+        body: serde_json::json!({}),
+    }
+}
+
+pub fn context() -> RequestContext {
+    RequestContext {
+        model: "model".into(),
+        custom_llm_provider: "provider".into(),
+        optional_params: serde_json::json!({}),
+        secret_fields: Vec::new(),
+        api_key: None,
+    }
+}
+
+#[derive(Clone, Default)]
+pub struct Log(Arc<Mutex<Vec<String>>>);
+
+impl Log {
+    pub fn push(&self, entry: impl Into<String>) {
+        self.0.lock().unwrap().push(entry.into());
+    }
+
+    pub fn entries(&self) -> Vec<String> {
+        self.0.lock().unwrap().clone()
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum OpScript {
+    Answer,
+    RaisePython,
+    RejectNatively,
+    RejectRequestNatively,
+    RaiseRequestPython,
+    CancelRequest,
+    AwaitAnswer,
+    AwaitFailure,
+}
+
+/// The fake route's public exception, kept as a value so a test sees what `classify`
+/// produced before the driver raises it.
+#[derive(Debug, PartialEq, Eq)]
+pub struct Classified(pub String);
+
+impl From<Classified> for PyErr {
+    fn from(classified: Classified) -> Self {
+        PyValueError::new_err(format!("classified: {}", classified.0))
+    }
+}
+
+pub struct SyntheticBinding {
+    log: Log,
+    op: OpScript,
+    classifier_fails: bool,
+    pending_reply: Option<Reply<String>>,
+}
+
+pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
+    SyntheticBinding {
+        log: log.clone(),
+        op,
+        classifier_fails: false,
+        pending_reply: None,
+    }
+}
+
+impl SyntheticBinding {
+    pub fn with_failing_classifier(self) -> Self {
+        Self {
+            classifier_fails: true,
+            ..self
+        }
+    }
+
+    fn answer(&self, value: impl FnOnce() -> String) -> Result<String, InvokeError<Error>> {
+        match self.op {
+            OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()),
+            OpScript::RaisePython | OpScript::RaiseRequestPython => {
+                Err(PyValueError::new_err("op failed").into())
+            }
+            OpScript::RejectNatively | OpScript::RejectRequestNatively => {
+                Err(InvokeError::Native(Error("op rejected".into())))
+            }
+            OpScript::CancelRequest => Err(CancelledError::new_err(()).into()),
+        }
+    }
+}
+
+impl PythonBinding for SyntheticBinding {
+    type Protocol = Synthetic;
+    type Failure = Classified;
+
+    fn decode_request(
+        &mut self,
+        _: Python<'_>,
+        arguments: &Bound<'_, PyDict>,
+    ) -> Result<String, InvokeError<Error>> {
+        self.log.push("project");
+        match self.op {
+            OpScript::RejectRequestNatively
+            | OpScript::RaiseRequestPython
+            | OpScript::CancelRequest => self.answer(String::new),
+            _ => Ok(format!("project:{}", arguments.len())),
+        }
+    }
+
+    fn encode_stream_head(
+        &mut self,
+        _: Python<'_>,
+        head: std::convert::Infallible,
+    ) -> PyResult<Py<PyAny>> {
+        match head {}
+    }
+
+    fn encode_chunk(
+        &mut self,
+        _: Python<'_>,
+        chunk: std::convert::Infallible,
+    ) -> PyResult<Py<PyAny>> {
+        match chunk {}
+    }
+
+    fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
+        self.log.push("complete");
+        Ok(PyString::new(py, &response).into_any().unbind())
+    }
+
+    fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
+        self.log.push(format!("classify:{error}"));
+        if self.classifier_fails {
+            return Err(PyTypeError::new_err("classifier failed"));
+        }
+        Ok(Classified(error.0))
+    }
+
+    fn host_error(error: &PyErr) -> Error {
+        Error(error.to_string())
+    }
+}
+
+impl PythonHostCalls<Synthetic> for SyntheticBinding {
+    fn handle_host_call(
+        &mut self,
+        _: Python<'_>,
+        (op, reply): (&'static str, Reply<String>),
+    ) -> Result<(), InvokeError<Error>> {
+        self.log.push(format!("op:{op}"));
+        self.answer(|| op.to_string())
+            .map(|answer| reply.send(answer))
+    }
+
+    fn begin_host_call(
+        &mut self,
+        py: Python<'_>,
+        (op, reply): (&'static str, Reply<String>),
+    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
+        if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
+            return self.handle_host_call(py, (op, reply)).map(|()| None);
+        }
+        self.pending_reply = Some(reply);
+        let module = PyModule::from_code(
+            py,
+            pyo3::ffi::c_str!(
+                "async def answer(fail):\n    if fail:\n        raise LookupError('async host failed')\n    return 'awaited'\n"
+            ),
+            pyo3::ffi::c_str!("host_op.py"),
+            pyo3::ffi::c_str!("host_op"),
+        )?;
+        Ok(Some(
+            module
+                .getattr("answer")?
+                .call1((matches!(self.op, OpScript::AwaitFailure),))?
+                .unbind(),
+        ))
+    }
+
+    fn resume_host_call(
+        &mut self,
+        py: Python<'_>,
+        result: PyResult<Py<PyAny>>,
+    ) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
+        let answer = result?.extract::<String>(py)?;
+        self.pending_reply.take().unwrap().send(answer);
+        Ok(None)
+    }
+}
+
+impl PythonOwned for SyntheticBinding {
+    fn close(&mut self, _: Python<'_>) {
+        self.log.push("host.close");
+    }
+
+    fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum HookScript {
+    Plain,
+    RewriteArguments,
+    ObserveArguments,
+    FailBegin,
+    ReplaceResponse,
+    FailAfterSuccess,
+    CancelTerminal,
+    FailTerminal,
+}
+
+pub struct SyntheticHooks {
+    log: Log,
+    script: HookScript,
+}
+
+pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
+    SyntheticHooks {
+        log: log.clone(),
+        script,
+    }
+}
+
+impl CallHooks<PythonRuntime> for SyntheticHooks {
+    fn prepare_arguments(
+        &mut self,
+        _: Python<'_>,
+        arguments: Py<PyDict>,
+        _: f64,
+    ) -> PyResult<HookStep<Self, Py<PyDict>>> {
+        self.log.push("begin");
+        if matches!(self.script, HookScript::FailBegin) {
+            return Err(PyValueError::new_err("begin failed"));
+        }
+        if matches!(self.script, HookScript::RewriteArguments) {
+            let prepared = Python::attach(|py| -> PyResult<Py<PyDict>> {
+                let prepared = arguments.bind(py).copy()?;
+                prepared.set_item("prepared", "hook")?;
+                prepared.set_item("api_key", "hook-key")?;
+                Ok(prepared.unbind())
+            })?;
+            return Ok(HookStep::Ready(prepared));
+        }
+        Ok(HookStep::Ready(arguments))
+    }
+
+    fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
+        if matches!(self.script, HookScript::ObserveArguments) {
+            let value: String = arguments
+                .bind(py)
+                .get_item("prepared")?
+                .unwrap()
+                .extract()?;
+            self.log.push(format!("adopted:{value}"));
+        }
+        Ok(())
+    }
+
+    fn before_provider_request(
+        &mut self,
+        _: Python<'_>,
+        wire: Box<WireRequest>,
+        _: &RequestContext,
+    ) -> PyResult<HookStep<Self, Box<WireRequest>>> {
+        self.log.push("before_provider_request");
+        Ok(HookStep::Ready(Box::new(WireRequest {
+            url: "rewritten".into(),
+            ..*wire
+        })))
+    }
+
+    fn transform_response(
+        &mut self,
+        py: Python<'_>,
+        response: Py<PyAny>,
+        _: Timing,
+    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
+        self.log.push("after_success");
+        match self.script {
+            HookScript::ReplaceResponse => Ok(HookStep::Ready(
+                "replaced".into_pyobject(py)?.into_any().unbind(),
+            )),
+            HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
+            _ => Ok(HookStep::Ready(response)),
+        }
+    }
+
+    fn on_event(
+        &mut self,
+        py: Python<'_>,
+        event: PythonCallEvent<'_>,
+    ) -> PyResult<HookStep<Self, ()>> {
+        let terminal = matches!(
+            event,
+            PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
+        );
+        if terminal && matches!(self.script, HookScript::CancelTerminal) {
+            return Err(CancelledError::new_err("callback cancelled"));
+        }
+        if terminal && matches!(self.script, HookScript::FailTerminal) {
+            return Err(PyValueError::new_err("callback failed"));
+        }
+        self.log.push(match event {
+            PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
+            PythonCallEvent::Started { .. } => "started".into(),
+            PythonCallEvent::Cancelled { .. } => "cancelled".into(),
+            PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
+                format!("response:{}", raw.body)
+            }
+            PythonCallEvent::Succeeded { response, .. } => {
+                format!("succeeded:{}", response.bind(py))
+            }
+            PythonCallEvent::Failed { origin, error, .. } => {
+                format!("failed:{origin:?}:{}", error.value(py))
+            }
+        });
+        Ok(HookStep::Ready(()))
+    }
+
+    fn on_stream_open(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
+        self.log.push("opened");
+        Ok(())
+    }
+
+    fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
+        self.log.push("delivered");
+        Ok(())
+    }
+}
+
+impl PythonOwned for SyntheticHooks {
+    fn close(&mut self, _: Python<'_>) {
+        self.log.push("adapter.close");
+    }
+
+    fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+}
+
+pub type MachineFactory = Box<dyn FnOnce(String) -> CallMachine<Synthetic> + Send + Sync>;
+
+/// Answers to projection, to the route op and to `before_provider_request` all reach the
+/// response, so a driver that misroutes a reply changes what the call returns.
+pub fn success_machine() -> MachineFactory {
+    Box::new(move |projected| {
+        CallMachine::new(None, move |host| {
+            Box::pin(async move {
+                let signed = host.services.call(|reply| ("sign", reply)).await?;
+                let wire = host
+                    .interceptors
+                    .before_provider_request(wire(), context())
+                    .await?;
+                host.interceptors
+                    .after_provider_response(RawResponse { body: "raw".into() })
+                    .await?;
+                Ok(format!("{projected}|{signed}|{}", wire.url))
+            })
+        })
+    })
+}
+
+pub fn failing_machine() -> MachineFactory {
+    Box::new(move |_| {
+        CallMachine::new(None, |_| {
+            Box::pin(async move { Err(Error("provider exploded".into())) })
+        })
+    })
+}
+
+/// Drives an asynchronous call's coroutine to its settled value from a sync test. A
+/// synchronous call is already settled.
+pub fn settled(
+    py: Python<'_>,
+    result: PyResult<Py<PyAny>>,
+    asynchronous: bool,
+) -> PyResult<Py<PyAny>> {
+    if !asynchronous {
+        return result;
+    }
+    let coroutine = result?;
+    let completed = coroutine
+        .call_method1(py, "send", (py.None(),))
+        .unwrap_err();
+    if !completed.is_instance_of::<pyo3::exceptions::PyStopIteration>(py) {
+        return Err(completed);
+    }
+    completed.value(py).getattr("value").map(Bound::unbind)
+}
+
+pub fn run_scripted(
+    py: Python<'_>,
+    machine: MachineFactory,
+    op: OpScript,
+    script: HookScript,
+    asynchronous: bool,
+) -> (PyResult<Py<PyAny>>, Vec<String>) {
+    let log = Log::default();
+    run_composed(
+        py,
+        machine,
+        binding(&log, op),
+        hooks(&log, script),
+        std::convert::identity,
+        call_options(asynchronous),
+    )
+}
+
+pub fn run_composed<L: PythonCallHooks + 'static>(
+    py: Python<'_>,
+    machine: MachineFactory,
+    binding: SyntheticBinding,
+    hooks: SyntheticHooks,
+    compose: impl FnOnce(SyntheticHooks) -> L,
+    options: CallOptions,
+) -> (PyResult<Py<PyAny>>, Vec<String>) {
+    let asynchronous = options.asynchronous;
+    let log = binding.log.clone();
+    let arguments = PyDict::new(py);
+    arguments.set_item("model", "m").unwrap();
+    let result = handed(
+        py,
+        run_call(
+            py,
+            move |_, _, request| Ok(machine(request)),
+            binding,
+            compose(hooks),
+            arguments.unbind(),
+            options,
+        ),
+        asynchronous,
+    );
+    (settled(py, result, asynchronous), log.entries())
+}
+
+pub struct Streaming;
+
+impl Protocol for Streaming {
+    type Response = ();
+    type Error = Error;
+    type Request = ();
+    type HostCall = std::convert::Infallible;
+    type Chunk = &'static str;
+    type StreamHead = Vec<(&'static str, &'static str)>;
+}
+
+pub struct StreamingBinding;
+
+impl PythonBinding for StreamingBinding {
+    type Protocol = Streaming;
+    type Failure = Classified;
+
+    fn decode_request(
+        &mut self,
+        _: Python<'_>,
+        _: &Bound<'_, PyDict>,
+    ) -> Result<(), InvokeError<Error>> {
+        Ok(())
+    }
+
+    fn encode_stream_head(
+        &mut self,
+        py: Python<'_>,
+        head: Vec<(&'static str, &'static str)>,
+    ) -> PyResult<Py<PyAny>> {
+        let headers = PyDict::new(py);
+        for (name, value) in head {
+            headers.set_item(name, value)?;
+        }
+        let hidden = PyDict::new(py);
+        hidden.set_item("additional_headers", headers)?;
+        Ok(hidden.into_any().unbind())
+    }
+
+    fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
+        Ok(PyString::new(py, chunk).into_any().unbind())
+    }
+
+    fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
+        Ok(py.None())
+    }
+
+    fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
+        Ok(Classified(error.0))
+    }
+
+    fn host_error(error: &PyErr) -> Error {
+        Error(error.to_string())
+    }
+}
+
+impl PythonHostCalls<Streaming> for StreamingBinding {
+    fn handle_host_call(
+        &mut self,
+        _: Python<'_>,
+        op: std::convert::Infallible,
+    ) -> Result<(), InvokeError<Error>> {
+        match op {}
+    }
+}
+
+impl PythonOwned for StreamingBinding {
+    fn close(&mut self, _: Python<'_>) {}
+
+    fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        Ok(())
+    }
+}
+
+pub fn streaming_machine() -> impl FnOnce(()) -> HostedMachine<Streaming> + Send + Sync {
+    move |()| {
+        hosted_call((), None, |(), _, _, _observations| async {
+            Ok(CallOutput::Stream {
+                head: vec![("request-id", "req_1")],
+                chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])),
+            })
+        })
+    }
+}
+
+/// Starts a streaming call and returns the caller's stream object.
+pub fn open_stream(
+    py: Python<'_>,
+    hooks: SyntheticHooks,
+    options: CallOptions,
+) -> Bound<'_, PyAny> {
+    let asynchronous = options.asynchronous;
+    let handed = handed(
+        py,
+        run_call(
+            py,
+            |_, _, request| Ok(streaming_machine()(request)),
+            StreamingBinding,
+            hooks,
+            PyDict::new(py).unbind(),
+            options,
+        ),
+        asynchronous,
+    )
+    .unwrap();
+    if asynchronous {
+        let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
+        stop.value(py).getattr("value").unwrap()
+    } else {
+        handed.into_bound(py)
+    }
+}
+
+/// Drives an `AsyncStream` or `SyncStream` to completion from a sync test.
+pub fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
+    if !asynchronous {
+        return stream
+            .try_iter()
+            .unwrap()
+            .map(|chunk| chunk.unwrap().extract().unwrap())
+            .collect();
+    }
+    std::iter::from_fn(|| {
+        let stop = stream
+            .call_method0("__anext__")
+            .unwrap()
+            .call_method1("send", (py.None(),))
+            .unwrap_err();
+        if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
+            return None;
+        }
+        assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
+        Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
+    })
+    .collect()
+}

litellm-rust/crates/http/src/settings.rs

@@ -140,7 +140,7 @@ impl Default for HttpSettings {
             proxies: EnvironmentProxies::default(),
             connect_timeout: Duration::from_secs(10),
             tcp_keepalive: None,
-            pool_idle_timeout: Duration::from_secs(120),
+            pool_idle_timeout: Duration::from_secs(15),
         }
     }
 }
@@ -330,7 +330,7 @@ mod tests {
             &[("AIOHTTP_TCP_KEEPIDLE", "45")],
         ))]);
         assert_eq!(settings.tcp_keepalive, None);
-        assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
+        assert_eq!(settings.pool_idle_timeout, Duration::from_secs(15));
     }
 
     fn proxy_flags(

litellm-rust/crates/python-bridge/AGENTS.md

@@ -48,8 +48,8 @@ Each step needs focused regression tests in the owning crate and Python integrat
   - Preserve provider-specific upload/submission/poll observation and encoding boundaries. Signed/build-captured bytes must not be silently reserialized
 - Conversion errors and every failure after the call starts are terminal
   - Disabled/unavailable native execution may select legacy once. Callback exceptions never authorize fallback or replay
-- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle and call driver in `litellm-host-python`
-  - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values
+- Use one ordinary inline `async def` driver in `litellm/rust_bridge/execution.py`, with the native handle and call driver in `litellm-host-python`
+  - Contract: generator-style `send` (`send(None)` starts), `throw`, idempotent `close`. Explicitly tagged `Await`/`Complete` preserve awaitable final values
   - Validate Created/Running/Suspended/Closed protocol states. The machine yields ops, the driver emits one terminal event, the adapter chooses dispatch policy
   - Defer effectful setup/context reads/timestamps until start. Unstarted-handle destruction releases inputs independently of Python `finally`
   - Catch only the selected await's errors. Start/resume errors propagate, `GeneratorExit` closes without further awaits

litellm-rust/crates/python-bridge/src/lib.rs

@@ -6,7 +6,6 @@ mod diagnostics;
 mod errors;
 mod execution;
 mod http;
-mod lifecycle;
 mod logger;
 mod marshal;
 mod preflight;

litellm-rust/crates/python-bridge/src/lifecycle.rs

@@ -1,10 +0,0 @@
-use litellm_host_python::CallOptions;
-use pyo3::prelude::*;
-
-pub(crate) fn binding(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
-    py.import("litellm.rust_bridge.streams")
-}
-
-pub(crate) fn call_options(asynchronous: bool) -> CallOptions {
-    CallOptions::new(asynchronous, binding)
-}

litellm-rust/crates/python-bridge/src/routes/mod.rs

@@ -12,9 +12,9 @@ pub(crate) mod traces;
 use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall};
 use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
 use litellm_host_python::{
-    HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
+    CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
 };
-use litellm_host_python::{missing_state, present};
+use litellm_host_python::{Error as HostError, present};
 use pyo3::{
     gc::{PyTraverseError, PyVisit},
     prelude::*,
@@ -64,7 +64,9 @@ impl RequestView {
     }
 
     fn dicts(&self) -> PyResult<&RequestDicts> {
-        self.0.as_ref().ok_or_else(missing_state)
+        self.0
+            .as_ref()
+            .ok_or_else(|| HostError::MissingState.into())
     }
 
     pub(crate) fn resolve(
@@ -72,7 +74,7 @@ impl RequestView {
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> PyResult<()> {
-        let dicts = self.0.as_mut().ok_or_else(missing_state)?;
+        let dicts = self.0.as_mut().ok_or(HostError::MissingState)?;
         dicts.request = effective_py_args(dicts.base.bind(py), arguments)?.unbind();
         Ok(())
     }
@@ -127,6 +129,10 @@ fn call_hooks(
     ))
 }
 
+fn execution_steps(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
+    py.import("litellm.rust_bridge.execution")
+}
+
 fn run_public_call<H, M>(
     py: Python<'_>,
     arguments: Py<PyDict>,
@@ -157,7 +163,7 @@ where
             .with(hooks)
             .with(crate::preflight::SdkPolicy),
         arguments,
-        crate::lifecycle::call_options(asynchronous),
+        CallOptions::new(asynchronous, execution_steps),
     )
 }
 

litellm-rust/crates/python-bridge/src/routes/ocr/host.rs

@@ -1,5 +1,5 @@
 use litellm_auth::ResolvedCredential;
-use litellm_host_python::{InvokeError, PythonBinding, missing_state};
+use litellm_host_python::{Error as HostError, InvokeError, PythonBinding};
 use litellm_host_python::{PythonHostCalls, PythonOwned};
 use litellm_inference_ocr::route::{Ocr, OcrCall, OcrOp};
 use litellm_llms::base_llm::ocr::error::Error;
@@ -42,7 +42,7 @@ impl OcrPythonHost {
     fn handles(&self) -> PyResult<&OcrHostHandles> {
         match &self.data {
             OcrHostData::Projected(handles) => Ok(handles),
-            _ => Err(missing_state()),
+            _ => Err(HostError::MissingState.into()),
         }
     }
 
@@ -50,13 +50,13 @@ impl OcrPythonHost {
         self.handles()?
             .azure_ad_token_provider
             .as_ref()
-            .ok_or_else(missing_state)?
+            .ok_or(HostError::MissingState)?
             .acquire(py)
     }
 
     fn projection(&mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<OcrCall> {
         let OcrHostData::Unprojected = self.data else {
-            return Err(missing_state());
+            return Err(HostError::MissingState.into());
         };
         self.view.resolve(py, arguments)?;
         let (request, handles) = project_request(self.view.request(py)?, arguments)?;

litellm/__init__.py

@@ -1710,6 +1710,12 @@ if TYPE_CHECKING:
     from .llms.hosted_vllm.decisions.transformation import (
         HostedVLLMDecisionsConfig as HostedVLLMDecisionsConfig,
     )
+    from .llms.laya.decisions.transformation import (
+        LayaDecisionsConfig as LayaDecisionsConfig,
+    )
+    from .llms.bespoke.decisions.transformation import (
+        BespokeDecisionsConfig as BespokeDecisionsConfig,
+    )
     from .llms.openai.decisions.transformation import (
         OpenAIDecisionsConfig as OpenAIDecisionsConfig,
     )
@@ -2281,12 +2287,14 @@ if TYPE_CHECKING:
     get_response_string: Callable[..., str]
     supports_function_calling: Callable[..., bool]
     supports_web_search: Callable[..., bool]
+    supports_tool_search: Callable[..., bool]
     supports_url_context: Callable[..., bool]
     supports_response_schema: Callable[..., bool]
     supports_parallel_function_calling: Callable[..., bool]
     supports_vision: Callable[..., bool]
     supports_audio_input: Callable[..., bool]
     supports_video_input: Callable[..., bool]
+    supports_box_questions: Callable[..., bool]
     supports_audio_output: Callable[..., bool]
     supports_system_messages: Callable[..., bool]
     supports_reasoning: Callable[..., bool]

litellm/_lazy_imports_registry.py

@@ -31,12 +31,14 @@ UTILS_NAMES: Final = (
     "create_tokenizer",
     "supports_function_calling",
     "supports_web_search",
+    "supports_tool_search",
     "supports_url_context",
     "supports_response_schema",
     "supports_parallel_function_calling",
     "supports_vision",
     "supports_audio_input",
     "supports_video_input",
+    "supports_box_questions",
     "supports_audio_output",
     "supports_system_messages",
     "supports_reasoning",
@@ -166,6 +168,8 @@ LLM_CONFIG_NAMES: Final = (
     "StrandsDeciderDecisionsConfig",
     "DatabricksDecisionsConfig",
     "HostedVLLMDecisionsConfig",
+    "LayaDecisionsConfig",
+    "BespokeDecisionsConfig",
     "OpenAIDecisionsConfig",
     "AzureAIDecisionsConfig",
     "NvidiaNimRerankConfig",
@@ -475,6 +479,7 @@ _UTILS_IMPORT_MAP: Final = {
     "create_tokenizer": (".utils", "create_tokenizer"),
     "supports_function_calling": (".utils", "supports_function_calling"),
     "supports_web_search": (".utils", "supports_web_search"),
+    "supports_tool_search": (".utils", "supports_tool_search"),
     "supports_url_context": (".utils", "supports_url_context"),
     "supports_response_schema": (".utils", "supports_response_schema"),
     "supports_parallel_function_calling": (
@@ -484,6 +489,7 @@ _UTILS_IMPORT_MAP: Final = {
     "supports_vision": (".utils", "supports_vision"),
     "supports_audio_input": (".utils", "supports_audio_input"),
     "supports_video_input": (".utils", "supports_video_input"),
+    "supports_box_questions": (".utils", "supports_box_questions"),
     "supports_audio_output": (".utils", "supports_audio_output"),
     "supports_system_messages": (".utils", "supports_system_messages"),
     "supports_reasoning": (".utils", "supports_reasoning"),
@@ -736,6 +742,8 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
         ".llms.hosted_vllm.decisions.transformation",
         "HostedVLLMDecisionsConfig",
     ),
+    "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"),
+    "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"),
     "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"),
     "AzureAIDecisionsConfig": (".llms.azure_ai.decisions.transformation", "AzureAIDecisionsConfig"),
     "NvidiaNimRerankConfig": (

litellm/_redis.py

@@ -675,7 +675,10 @@ def _get_redis_client_logic(**env_overrides):
         if redis_kwargs.get("startup_nodes")
         else redis.Redis
     )
-    return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
+    coerced_kwargs: Final = _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
+    if coerced_kwargs.get("socket_timeout") is None:
+        coerced_kwargs["socket_timeout"] = REDIS_SOCKET_TIMEOUT
+    return coerced_kwargs
 
 
 def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
@@ -722,7 +725,6 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
     sentinel_password: Final = redis_kwargs.get("sentinel_password")
     service_name: Final = redis_kwargs.get("service_name")
     connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
-    connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
     sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
 
     if not sentinel_nodes or not service_name:
@@ -755,7 +757,6 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
     sentinel_password: Final = redis_kwargs.get("sentinel_password")
     service_name: Final = redis_kwargs.get("service_name")
     connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
-    connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
     sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
 
     if not sentinel_nodes or not service_name:

litellm/_service_logger.py

@@ -152,6 +152,11 @@ class ServiceLogging(CustomLogger):
         except Exception as e:
             verbose_logger.exception("Error dispatching service event - %s", e)
 
+    def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
+        prometheus_logger: Final[PrometheusServicesLogger | None] = getattr(self, "prometheusServicesLogger", None)
+        if prometheus_logger is not None:
+            prometheus_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
+
     def service_success_hook(
         self,
         service: ServiceTypes,

litellm/caching/azure_blob_cache.py

@@ -65,7 +65,7 @@ class AzureBlobCache(BaseCache):
             print_verbose(f"Get Azure Blob Cache: key: {key}")
             as_bytes: Final = self.container_client.download_blob(key).readall()
             as_str: Final = as_bytes.decode("utf-8")
-            cached_response: Final = json.loads(as_str)
+            cached_response: Final[object] = json.loads(as_str)
 
             verbose_logger.debug(
                 "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
@@ -86,7 +86,7 @@ class AzureBlobCache(BaseCache):
             blob: Final = await self.async_container_client.download_blob(key)
             as_bytes: Final = await blob.readall()
             as_str: Final = as_bytes.decode("utf-8")
-            cached_response: Final = json.loads(as_str)
+            cached_response: Final[object] = json.loads(as_str)
             verbose_logger.debug(
                 "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
                 key,

litellm/caching/caching.py

@@ -709,7 +709,7 @@ class Cache:
                 cache_key = self.get_cache_key(**kwargs)
             if cache_key is not None:
                 if isinstance(result, BaseModel):
-                    result = result.model_dump_json()
+                    result = result.model_dump_json(exclude_unset=True)
 
                 ## DEFAULT TTL ##
                 if self.ttl is not None:

litellm/caching/caching_handler.py

@@ -1088,6 +1088,7 @@ class LLMCachingHandler:
             kwargs=kwargs,
             start_time=self.start_time,
             end_time=datetime.datetime.now(),
+            cache_hit=True,
         )
         return cached_result
 
@@ -1180,7 +1181,7 @@ class LLMCachingHandler:
                         )
                     )
                 else:
-                    result_json: Final = result.model_dump_json()
+                    result_json: Final = result.model_dump_json(exclude_unset=True)
                     create_cache_write_task(
                         lambda: cache.async_add_cache(
                             result_json,

litellm/caching/in_memory_cache.py

@@ -51,6 +51,7 @@ class InMemoryCache(BaseCache):
         self.ttl_dict: dict = {}
         self.expiration_heap: list[tuple[float, str]] = []
         self._increment_lock = threading.Lock()
+        self._write_lock = threading.Lock()
         self._clock = clock if clock is not None else lambda: time.time()
 
     def check_value_size(self, value: Any):
@@ -96,6 +97,9 @@ class InMemoryCache(BaseCache):
         """
         return key in self.ttl_dict and self._clock() > self.ttl_dict[key]
 
+    def _cache_keys(self) -> tuple[str, ...]:
+        return tuple(self.cache_dict)  # pyright: ignore[reportUnknownVariableType]  # cache_dict is untyped legacy state
+
     def _remove_key(self, key: str) -> None:
         """
         Remove a key from both cache_dict and ttl_dict
@@ -133,13 +137,23 @@ class InMemoryCache(BaseCache):
                 # Case 3: Entry is valid and not expired
                 break
 
+        if len(self.expiration_heap) > 2 * self.max_size_in_memory:
+            self.expiration_heap = [entry for entry in self.expiration_heap if self.ttl_dict.get(entry[1]) == entry[0]]
+            heapq.heapify(self.expiration_heap)
+
         # Step 2: Evict if cache is still full
-        while len(self.cache_dict) >= self.max_size_in_memory:
+        while self.expiration_heap and len(self.cache_dict) >= self.max_size_in_memory:
             expiration_time, key = heapq.heappop(self.expiration_heap)
             # Skip if key was removed or updated
             if self.ttl_dict.get(key) == expiration_time:
                 self._remove_key(key)
 
+        if not self.expiration_heap:
+            untracked_keys: Final = self._cache_keys()
+            overflow: Final = len(untracked_keys) - self.max_size_in_memory + 1
+            for untracked_key in untracked_keys[: max(overflow, 0)]:
+                self._remove_key(untracked_key)
+
         # de-reference the removed item
         # https://www.geeksforgeeks.org/diagnosing-and-fixing-memory-leaks-in-python/
         # One of the most common causes of memory leaks in Python is the retention of objects that are no longer being used.
@@ -150,7 +164,7 @@ class InMemoryCache(BaseCache):
         Check if ttl is set for a key
         """
         ttl_time: Final = self.ttl_dict.get(key)
-        if ttl_time is None or float(ttl_time) < self._clock():  # if ttl is not set, allow override
+        if ttl_time is None or ttl_time < self._clock():  # if ttl is not set, allow override
             return True
         else:
             return False
@@ -160,21 +174,22 @@ class InMemoryCache(BaseCache):
         if self.max_size_in_memory == 0:
             return  # Don't cache anything if max size is 0
 
-        # Always prune expired/outdated heap roots before inserting.
-        # This keeps expiration_heap bounded even when the live cache stays
-        # below max_size_in_memory and keys are reinserted after TTL expiry.
-        self.evict_cache()
-        if not self.check_value_size(value):
-            return
+        with self._write_lock:
+            # Always prune expired/outdated heap roots before inserting.
+            # This keeps expiration_heap bounded even when the live cache stays
+            # below max_size_in_memory and keys are reinserted after TTL expiry.
+            self.evict_cache()
+            if not self.check_value_size(value):
+                return
 
-        self.cache_dict[key] = value
-        if self.allow_ttl_override(key):  # if ttl is not set, set it to default ttl
-            if "ttl" in kwargs and kwargs["ttl"] is not None:
-                self.ttl_dict[key] = self._clock() + float(kwargs["ttl"])
-                heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
-            else:
-                self.ttl_dict[key] = self._clock() + self.default_ttl
-                heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
+            self.cache_dict[key] = value
+            if self.allow_ttl_override(key):  # if ttl is not set, set it to default ttl
+                if "ttl" in kwargs and kwargs["ttl"] is not None:
+                    self.ttl_dict[key] = self._clock() + float(kwargs["ttl"])
+                    heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
+                else:
+                    self.ttl_dict[key] = self._clock() + self.default_ttl
+                    heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
 
     async def async_set_cache(self, key, value, **kwargs):
         self.set_cache(key=key, value=value, **kwargs)
@@ -258,14 +273,15 @@ class InMemoryCache(BaseCache):
         return results
 
     def flush_cache(self):
-        self.cache_dict.clear()
-        self.ttl_dict.clear()
-        self.expiration_heap.clear()
+        with self._write_lock:
+            self.cache_dict.clear()
+            self.ttl_dict.clear()
+            self.expiration_heap.clear()
 
     async def disconnect(self):
         pass
 
-    def delete_cache(self, key):
+    def delete_cache(self, key: str) -> None:
         self._remove_key(key)
 
     async def async_get_ttl(self, key: str) -> int | None:

litellm/caching/redis_cache.py

@@ -18,7 +18,9 @@ import json
 import logging
 import threading
 import time
-from collections.abc import Awaitable, Callable, Iterator, Sequence
+import weakref
+from collections import Counter
+from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
 from contextvars import ContextVar
 from dataclasses import dataclass
 from datetime import timedelta
@@ -51,12 +53,12 @@ from .base_cache import BaseCache
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
-    from prometheus_client import Counter as _PromCounter
-    from prometheus_client import Gauge as _PromGauge
     from redis.asyncio import Redis, RedisCluster
     from redis.asyncio.client import Pipeline
     from redis.asyncio.cluster import ClusterPipeline
 
+    from litellm._service_logger import ServiceLogging
+
     pipeline = Pipeline
     cluster_pipeline = ClusterPipeline
     async_redis_client = Redis
@@ -100,8 +102,7 @@ _GENERIC_CALLER_MODULES: Final = frozenset(
         "litellm.caching.redis_batch",
         "litellm.caching.dual_cache",
         "litellm.caching.caching",
-        "litellm.rust_bridge.lifecycle",
-        "litellm.rust_bridge.streams",
+        "litellm.rust_bridge.execution",
         "contextlib",
     }
 )
@@ -243,7 +244,9 @@ class RedisCircuitBreaker:
         recovery_timeout: int,
         enabled: bool = True,
         timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
+        service_logger: "ServiceLogging | None" = None,
     ) -> None:
+        self._service_logger = service_logger
         self.failure_threshold = failure_threshold
         self.recovery_timeout = recovery_timeout
         self.enabled = enabled
@@ -255,7 +258,12 @@ class RedisCircuitBreaker:
         self._opened_at: float | None = None
         self._state = self.CLOSED
         self._generation = 0
-        _breaker_metrics().record_state_change(None, self._state)
+        _LIVE_BREAKERS.add(self)
+        self._report(transitions_to=None, failure_class=None)
+
+    @property
+    def state(self) -> str:
+        return self._state
 
     @property
     def generation(self) -> int:
@@ -301,7 +309,7 @@ class RedisCircuitBreaker:
             self._timeout_count = 0
             self._timeout_streak_started_at = None
         self._opened_at = now
-        _breaker_metrics().record_failure("timeout" if is_timeout else "connectivity")
+        self._report(transitions_to=None, failure_class="timeout" if is_timeout else "connectivity")
         if self._should_open(now):
             if self._state != self.OPEN:
                 verbose_logger.warning(
@@ -327,10 +335,13 @@ class RedisCircuitBreaker:
     def _set_state(self, state: str) -> None:
         if state == self._state:
             return
-        _breaker_metrics().record_transition(state)
-        _breaker_metrics().record_state_change(self._state, state)
         self._state = state
         self._generation += 1
+        self._report(transitions_to=state, failure_class=None)
+
+    def _report(self, transitions_to: str | None, failure_class: str | None) -> None:
+        if self._service_logger is not None:
+            self._service_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
 
 
 _RedisCallResult = TypeVar("_RedisCallResult")
@@ -400,12 +411,19 @@ def _redis_timeout_error_types() -> tuple[type, ...]:
     ``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from
     py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own
     class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither.
+    ``MaxConnectionsError`` is the async cluster client's pool-exhausted signal (its nodes have
+    no blocking pool, so it is raised instantly instead of after a wait): a busy pool, not a
+    dead Redis, like the standalone pool's chained wait timeout.
     """
     try:
-        from redis.exceptions import TimeoutError as RedisTimeoutError
+        from redis import exceptions as redis_exceptions
     except ImportError:
         return (TimeoutError, asyncio.TimeoutError)
-    return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError)
+    max_connections_error: Final = cast(  # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
+        "type[Exception]",
+        redis_exceptions.MaxConnectionsError,  # pyright: ignore[reportAttributeAccessIssue]  # types-redis 4.6 stubs predate it
+    )
+    return (redis_exceptions.TimeoutError, TimeoutError, asyncio.TimeoutError, max_connections_error)
 
 
 _MAX_EXCEPTION_CAUSE_DEPTH: Final = 20
@@ -430,59 +448,6 @@ def is_redis_timeout_failure(exc: BaseException) -> bool:
     return any(isinstance(link, timeout_types) for link in _explicit_causes(exc))
 
 
-class _BreakerMetrics:
-    """Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent.
-
-    Registered lazily on the default registry (which /metrics serves) via the module-level
-    ``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers.
-    """
-
-    def __init__(self) -> None:
-        self._state_gauge: _PromGauge | None = None
-        self._transitions: _PromCounter | None = None
-        self._failures: _PromCounter | None = None
-        try:
-            from prometheus_client import Counter as PromCounter
-            from prometheus_client import Gauge
-        except ImportError:
-            return
-        self._state_gauge = Gauge(
-            "litellm_redis_circuit_breaker_state",
-            "Number of Redis circuit breakers currently in each state",
-            labelnames=("state",),
-        )
-        self._transitions = PromCounter(
-            "litellm_redis_circuit_breaker_transitions",
-            "Redis circuit breaker state transitions",
-            labelnames=("state",),
-        )
-        self._failures = PromCounter(
-            "litellm_redis_circuit_breaker_failures",
-            "Redis health failures counted by the circuit breaker",
-            labelnames=("failure_class",),
-        )
-
-    def record_state_change(self, old_state: str | None, new_state: str) -> None:
-        if self._state_gauge is None:
-            return
-        if old_state is not None:
-            self._state_gauge.labels(old_state).dec()
-        self._state_gauge.labels(new_state).inc()
-
-    def record_transition(self, state: str) -> None:
-        if self._transitions is not None:
-            self._transitions.labels(state).inc()
-
-    def record_failure(self, failure_class: str) -> None:
-        if self._failures is not None:
-            self._failures.labels(failure_class).inc()
-
-
-@functools.lru_cache(maxsize=1)
-def _breaker_metrics() -> _BreakerMetrics:
-    return _BreakerMetrics()
-
-
 def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
     """Record a Redis failure that the calling method is about to swallow.
 
@@ -498,6 +463,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
     _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
 
 
+_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet()
+
+
+def redis_breaker_state_counts() -> Mapping[str, int]:
+    """How many live Redis circuit breakers sit in each state, for the Prometheus services logger."""
+    return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS)))
+
+
 class RedisCircuitBreakerOpenError(Exception):
     pass
 
@@ -637,9 +610,14 @@ def _redis_circuit_breaker_guard(method):
     """
 
     @functools.wraps(method)
-    async def wrapper(self, *args, **kwargs):
+    async def wrapper(self: "RedisCache", *args, **kwargs):
+        async def guarded() -> object:
+            return await method(self, *args, **kwargs)
+
         return await _run_under_circuit_breaker(
-            self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
+            self._circuit_breaker,  # pyright: ignore[reportPrivateUsage]  # the guard is RedisCache's own decorator
+            wrapper.__name__,
+            lambda: self._admit_async_call(guarded),  # pyright: ignore[reportPrivateUsage]  # same decorator
         )
 
     return wrapper
@@ -684,12 +662,9 @@ class RedisCache(BaseCache):
             redis_kwargs["socket_timeout"] = socket_timeout
 
         ### HEALTH MONITORING OBJECT ###
-        if kwargs.get("service_logger_obj", None) is not None and isinstance(
-            kwargs["service_logger_obj"], ServiceLogging
-        ):
+        self.service_logger_obj: ServiceLogging = ServiceLogging()
+        if isinstance(kwargs.get("service_logger_obj"), ServiceLogging):
             self.service_logger_obj = kwargs.pop("service_logger_obj")
-        else:
-            self.service_logger_obj = ServiceLogging()
 
         redis_kwargs.update(kwargs)
         self.redis_client = get_redis_client(**redis_kwargs)
@@ -716,6 +691,7 @@ class RedisCache(BaseCache):
             failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
             recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
             enabled=REDIS_CIRCUIT_BREAKER_ENABLED,
+            service_logger=self.service_logger_obj,
         )
 
         self._setup_health_pings()
@@ -725,6 +701,9 @@ class RedisCache(BaseCache):
         else:
             super().__init__()  # defaults to 60s
 
+    def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
+        return call()
+
     def _setup_health_pings(self):
         """Setup async and sync health pings for Redis."""
         # ASYNC HEALTH PING
@@ -1105,7 +1084,9 @@ class RedisCache(BaseCache):
                     litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor)
                 return await executor(keys=keys, args=args, client=client)
 
-            return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute)
+            return await _run_under_circuit_breaker(
+                self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute)
+            )
 
         return run_script
 

litellm/caching/redis_cluster_cache.py

@@ -5,9 +5,17 @@ Key differences:
 - RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
 """
 
-from typing import TYPE_CHECKING, Any, Final
+import asyncio
+import contextvars
+import os
+import weakref
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar
 
+from litellm._logging import verbose_logger
 from litellm.caching.redis_cache import RedisCache
+from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
@@ -23,9 +31,100 @@ else:
     Span = Any
 
 
+_T: Final = TypeVar("_T")
+
+
+def _connection_limit(max_connections: object) -> int | None:
+    try:
+        limit: Final = int(max_connections)  # pyright: ignore[reportArgumentType]  # config value, str or int
+    except (TypeError, ValueError):
+        return None
+    return limit if limit > 0 else None
+
+
+_ClientSemaphores: TypeAlias = weakref.WeakValueDictionary[str, asyncio.Semaphore]
+
+
+@dataclass(slots=True)
+class _SemaphoreRegistry:
+    """Weak on both layers: a semaphore binds to its loop once a waiter queues on it, so a strong value would keep a
+    closed loop alive. In-flight ``ClusterConnectionGate.run`` calls hold the semaphore they acquired."""
+
+    _per_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _ClientSemaphores]" = field(
+        default_factory=weakref.WeakKeyDictionary, repr=False
+    )
+
+    def get(self, client_key: str, limit: int) -> asyncio.Semaphore:
+        per_loop: Final = self._per_loop.setdefault(asyncio.get_running_loop(), weakref.WeakValueDictionary())
+        existing: Final = per_loop.get(client_key)
+        if existing is not None:
+            return existing
+        created: Final = asyncio.Semaphore(limit)
+        per_loop[client_key] = created
+        return created
+
+
+_Admission: TypeAlias = tuple[str, "asyncio.Task[object] | None"]
+_SEMAPHORES: Final = _SemaphoreRegistry()
+_ADMITTED: Final[contextvars.ContextVar[frozenset[_Admission]]] = contextvars.ContextVar(
+    "redis_cluster_gate_admitted", default=frozenset()
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ClusterConnectionGate:
+    """The async ``RedisCluster`` has no blocking pool, so past ``max_connections`` a node raises
+    ``MaxConnectionsError`` instead of waiting and the circuit breaker opens on a healthy Redis. This gate makes
+    guarded commands wait for a slot the way the standalone ``BlockingConnectionPool`` does. A command that finds no
+    slot within ``wait_timeout`` runs anyway, so a genuinely saturated pool still reaches the client's own error
+    handling and the breaker."""
+
+    limit: int
+    wait_timeout: float = REDIS_CONNECTION_POOL_TIMEOUT
+
+    @classmethod
+    def from_max_connections(cls, max_connections: object) -> "ClusterConnectionGate | None":
+        limit: Final = _connection_limit(max_connections)
+        return None if limit is None else cls(limit=limit)
+
+    async def run(self, call: Callable[[], Awaitable[_T]], client_key: str) -> _T:
+        admission: Final[_Admission] = (client_key, asyncio.current_task())
+        if admission in _ADMITTED.get():
+            return await call()
+        semaphore: Final = _SEMAPHORES.get(client_key, self.limit)
+        try:
+            await asyncio.wait_for(semaphore.acquire(), timeout=self.wait_timeout)
+        except asyncio.TimeoutError:
+            verbose_logger.debug(
+                "Redis Cluster: no free connection slot within %ss (max_connections=%s), sending anyway",
+                self.wait_timeout,
+                self.limit,
+            )
+            return await call()
+        token: Final = _ADMITTED.set(_ADMITTED.get() | {admission})
+        try:
+            return await call()
+        finally:
+            _ADMITTED.reset(token)
+            semaphore.release()
+
+
 class RedisClusterCache(RedisCache):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, *args, max_connections: object = None, connection_gate: object = None, **kwargs):
+        pool_kwargs: Final = {} if max_connections is None else {"max_connections": max_connections}
+        super().__init__(*args, **pool_kwargs, **kwargs)
+        self._connection_gate: Final = (
+            connection_gate
+            if isinstance(connection_gate, ClusterConnectionGate)
+            else ClusterConnectionGate.from_max_connections(
+                os.environ.get("REDIS_MAX_CONNECTIONS") if max_connections is None else max_connections
+            )
+        )
+
+    def _admit_async_call(self, call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]:
+        if self._connection_gate is None:
+            return call()
+        return self._connection_gate.run(call, self._get_async_client_cache_key())
 
     def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
         """

litellm/completion_extras/litellm_responses_transformation/handler.py

@@ -2,8 +2,14 @@
 Handler for transforming /chat/completions api requests to litellm.responses requests
 """
 
-from collections.abc import AsyncIterable, Coroutine, Iterable
-from typing import TYPE_CHECKING, Any, Final, Union
+from collections.abc import AsyncIterable, Coroutine, Iterable, Mapping
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Union,
+    cast,  # noqa: TID251  # optional_params reaches the bridge as an untyped dict
+)
 
 from typing_extensions import TypedDict
 
@@ -245,10 +251,15 @@ class ResponsesToCompletionBridgeHandler:
         else:
             if self._is_preformatted_cached_chat_stream(result):
                 return self._apply_post_stream_processing(result, model, custom_llm_provider)
+            from .transformation import is_structured_output_request
+
             completion_stream: Final = self.transformation_handler.get_model_response_iterator(
                 streaming_response=result,
                 sync_stream=True,
                 json_mode=kwargs.get("json_mode"),
+                structured_output=is_structured_output_request(
+                    cast("Mapping[str, object]", optional_params)  # cast-ok: optional_params is an untyped dict
+                ),
             )
             streamwrapper: Final = CustomStreamWrapper(
                 completion_stream=completion_stream,
@@ -343,10 +354,15 @@ class ResponsesToCompletionBridgeHandler:
         else:
             if self._is_preformatted_cached_chat_stream(result):
                 return self._apply_post_stream_processing(result, model, custom_llm_provider)
+            from .transformation import is_structured_output_request
+
             completion_stream: Final = self.transformation_handler.get_model_response_iterator(
                 streaming_response=result,
                 sync_stream=False,
                 json_mode=kwargs.get("json_mode"),
+                structured_output=is_structured_output_request(
+                    cast("Mapping[str, object]", optional_params)  # cast-ok: optional_params is an untyped dict
+                ),
             )
             streamwrapper: Final = CustomStreamWrapper(
                 completion_stream=completion_stream,

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -20,7 +20,7 @@ from openai.types.responses.response_input_param import (
 from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
 from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
 from openai.types.responses.tool_param import FunctionToolParam
-from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 
 import litellm
 from litellm import ModelResponse
@@ -35,6 +35,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
 from litellm.llms.base_llm.bridges.completion_transformation import (
     CompletionTransformationBridge,
 )
+from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool
 from litellm.responses.sse_output_recovery import (
     parse_sse_json_chunk,
     record_output_item_chunk,
@@ -304,6 +305,136 @@ def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[
     )
 
 
+_TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"})
+_TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_ITEM_TYPES})
+_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items"
+
+
+def _json_field(value: object, key: str) -> object:
+    return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None  # cast-ok: untyped JSON
+
+
+def _tool_search_replay_item(item: object) -> Mapping[str, object] | None:
+    raw: Final[object] = item.model_dump(exclude_unset=True) if isinstance(item, BaseModel) else item
+    if not isinstance(raw, dict):
+        return None
+    fields: Final = cast(dict[str, object], raw)  # cast-ok: provider JSON item, keys forwarded verbatim
+    if fields.get("type") not in _TOOL_SEARCH_REPLAY_ITEM_TYPES:
+        return None
+    return {key: value for key, value in fields.items() if value is not None}
+
+
+def _tool_search_items_from_output_items(output_items: object) -> tuple[Mapping[str, object], ...]:
+    if not isinstance(output_items, list):
+        return ()
+    entries: Final = cast(Sequence[object], output_items)  # cast-ok: each entry is re-validated
+    items: Final = tuple(item for item in (_tool_search_replay_item(entry) for entry in entries) if item is not None)
+    last_search: Final = max(
+        (index for index, item in enumerate(items) if item.get("type") in _TOOL_SEARCH_ITEM_TYPES), default=-1
+    )
+    return items[: last_search + 1]
+
+
+def _tool_search_provider_fields(
+    items: Sequence[Mapping[str, object]],
+) -> dict[str, object] | None:  # mutable-ok: Message.provider_specific_fields is typed as a dict
+    return {_TOOL_SEARCH_ITEMS_KEY: [dict(item) for item in items]} if items else None
+
+
+def _tool_search_input_items(msg: object) -> tuple[Mapping[str, object], ...]:
+    return _tool_search_items_from_output_items(
+        _json_field(_json_field(msg, "provider_specific_fields"), _TOOL_SEARCH_ITEMS_KEY)
+    )
+
+
+def carries_tool_search_items(provider_specific_fields: object) -> bool:
+    return bool(_json_field(provider_specific_fields, _TOOL_SEARCH_ITEMS_KEY))
+
+
+def _tool_definition_name(tool: object) -> tuple[str, ...]:
+    name: Final = _json_field(tool, "server_label" if _json_field(tool, "type") == "mcp" else "name")
+    return (name,) if isinstance(name, str) else ()
+
+
+def _namespace_members(tool: object) -> Sequence[object]:
+    members: Final = _json_field(tool, "tools")
+    if _json_field(tool, "type") != "namespace" or not isinstance(members, list):
+        return (tool,)
+    return cast(Sequence[object], members)  # cast-ok: untyped JSON, each member re-checked
+
+
+def _loaded_tools(item: Mapping[str, object]) -> Sequence[object]:
+    tools: Final = item.get("tools")
+    if item.get("type") != "tool_search_output" or not isinstance(tools, list):
+        return ()
+    return cast(Sequence[object], tools)  # cast-ok: untyped JSON, each tool re-checked
+
+
+def _json_list(value: object) -> Sequence[object] | None:
+    return cast(Sequence[object], value) if isinstance(value, list) else None  # cast-ok: untyped JSON
+
+
+def _message_loaded_tools(message: object) -> Iterator[object]:
+    loaded: Final = chain.from_iterable(_loaded_tools(item) for item in _tool_search_input_items(message))
+    return chain.from_iterable(_namespace_members(tool) for tool in loaded)
+
+
+def replayed_tool_search_tools(request: object) -> tuple[object, ...]:
+    messages: Final = _json_list(_json_field(request, "messages")) or ()
+    return tuple(chain.from_iterable(_message_loaded_tools(message) for message in messages))
+
+
+def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]:
+    return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request)))
+
+
+def _loads_any_tool(message: object, names: frozenset[str]) -> bool:
+    loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message))
+    return any(name in names for name in loaded_names)
+
+
+def _without_tool_search_items(message: object) -> object:
+    fields: Final = _json_field(message, "provider_specific_fields")
+    if not isinstance(fields, dict):
+        return message
+    kept: Final = {
+        key: value
+        for key, value in cast(Mapping[str, object], fields).items()  # cast-ok: untyped JSON
+        if key != _TOOL_SEARCH_ITEMS_KEY
+    }
+    return {**cast(Mapping[str, object], message), "provider_specific_fields": kept}  # cast-ok: dict checked above
+
+
+def without_tool_search_items_loading(
+    request: object, names: frozenset[str]
+) -> list[object] | None:  # mutable-ok: request messages stay a list for every downstream handler
+    messages: Final = _json_list(_json_field(request, "messages"))
+    if messages is None:
+        return None
+    return [_without_tool_search_items(message) if _loads_any_tool(message, names) else message for message in messages]
+
+
+def _effective_tools(request: object) -> Sequence[object]:
+    override: Final = _json_list(_json_field(_json_field(request, "extra_body"), "tools"))
+    return override if override is not None else (_json_list(_json_field(request, "tools")) or ())
+
+
+def _requests_tool_search(request: object) -> bool:
+    return any(_json_field(tool, "type") == "tool_search" for tool in _effective_tools(request))
+
+
+def _reasoning_ids(items: Sequence[Mapping[str, object]]) -> frozenset[str]:
+    ids: Final = (item.get("id") for item in items if item.get("type") == "reasoning")
+    return frozenset(item_id for item_id in ids if isinstance(item_id, str))
+
+
+def _assistant_replay_items(msg: "AllMessageValues", replay_tool_search: bool) -> tuple[Mapping[str, object], ...]:
+    searched: Final = _tool_search_input_items(msg) if replay_tool_search else ()
+    searched_reasoning_ids: Final = _reasoning_ids(searched)
+    reasoning: Final = (item for item in _reasoning_input_items(msg) if item.get("id") not in searched_reasoning_ids)
+    return (*searched, *reasoning)
+
+
 def _as_chat_reasoning_items(
     reasoning_items: Sequence[_BuiltReasoningItem | ChatCompletionReasoningItem],
 ) -> list[ChatCompletionReasoningItem] | None:
@@ -346,14 +477,21 @@ def _fold_choices_into_one(
     tool_calls: Sequence[Mapping[str, object]],
     trailing_reasoning_content: str | None,
     trailing_reasoning_item: _BuiltReasoningItem | None,
+    tool_search_items: Sequence[Mapping[str, object]],
+    structured_output: bool = False,
 ) -> "Choices":
-    """Clients that only read choices[0] must still see every text and tool call."""
+    """Clients that only read choices[0] must still see every text and tool call.
+
+    Under a structured ``response_format`` each message item is a complete document, so
+    ``content`` is the first one, the same item the openai SDK returns as ``output_parsed``.
+    """
     from litellm.types.utils import Choices, Message
 
     messages: Final = tuple(choice.message for choice in choices)
-    texts: Final = tuple(str(message.content or "") for message in messages)
+    content_messages: Final = messages[:1] if structured_output else messages
+    texts: Final = tuple(str(message.content or "") for message in content_messages)
     offsets: Final = (0, *accumulate(len(text) for text in texts))
-    annotations: Final = tuple(_shifted_annotations(messages, offsets))
+    annotations: Final = tuple(_shifted_annotations(content_messages, offsets))
     reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content)
     reasoning_items: Final = (
         *chain.from_iterable(_message_reasoning_items(m) for m in messages),
@@ -367,12 +505,45 @@ def _fold_choices_into_one(
             tool_calls=list(tool_calls) or None,
             reasoning_content=" ".join(value for value in reasoning_contents if value) or None,
             reasoning_items=_as_chat_reasoning_items(reasoning_items),
+            provider_specific_fields=_tool_search_provider_fields(tool_search_items),
         ),
         finish_reason="tool_calls" if tool_calls else "stop",
         index=0,
     )
 
 
+def _responses_event_type(chunk: Mapping[str, object]) -> object:
+    event_type: Final = chunk.get("type")
+    return event_type.value if isinstance(event_type, ResponsesAPIStreamEvents) else event_type
+
+
+class _TypedPayload(BaseModel):
+    """The ``type`` discriminator of a response_format dict or a Responses output item."""
+
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    type: str | None = None
+
+
+_TYPED_PAYLOAD: Final = TypeAdapter(_TypedPayload)
+
+
+def _payload_type(payload: object) -> str | None:
+    try:
+        return _TYPED_PAYLOAD.validate_python(payload).type
+    except ValidationError:
+        return None
+
+
+def _is_message_item_done(chunk: Mapping[str, object]) -> bool:
+    if _responses_event_type(chunk) != ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
+        return False
+    return _payload_type(chunk.get("item")) == "message"
+
+
+def is_structured_output_request(optional_params: Mapping[str, object]) -> bool:
+    return _payload_type(optional_params.get("response_format")) in ("json_schema", "json_object")
+
+
 _ToolChoiceT = TypeVar("_ToolChoiceT")
 
 
@@ -419,7 +590,9 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
     provider_specific_fields: Mapping[str, object]
 
 
-def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
+def tool_call_dict_from_output_item(
+    item: Mapping[str, Any], index: int, *, include_namespace: bool = False
+) -> _ChatToolCallDict:
     """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
     completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
     string payload in ``input`` rather than ``arguments``; both map to
@@ -448,12 +621,24 @@ def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Cha
         provider_specific_fields = vars(raw_provider_fields)
     else:
         provider_specific_fields = None
-    if provider_specific_fields:
-        tool_call_dict["provider_specific_fields"] = provider_specific_fields
-        function_chunk["provider_specific_fields"] = provider_specific_fields
+    namespace: Final = item.get("namespace") if include_namespace else None
+    merged_provider_fields: Final = {
+        **(provider_specific_fields or {}),
+        **({"namespace": namespace} if isinstance(namespace, str) else {}),
+    }
+    if merged_provider_fields:
+        tool_call_dict["provider_specific_fields"] = merged_provider_fields
+        function_chunk["provider_specific_fields"] = merged_provider_fields
     return tool_call_dict
 
 
+def _tool_call_namespace(tool_call: object) -> str | None:
+    top_level: Final = _json_field(_json_field(tool_call, "provider_specific_fields"), "namespace")
+    function_fields: Final = _json_field(_json_field(tool_call, "function"), "provider_specific_fields")
+    nested: Final = _json_field(function_fields, "namespace")
+    return top_level if isinstance(top_level, str) else nested if isinstance(nested, str) else None
+
+
 def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFunctionParam | ToolChoiceCustomParam:
     if choice_type == "custom":
         return ToolChoiceCustomParam(type="custom", name=name)
@@ -479,6 +664,50 @@ def _reasoning_item_to_response_input(
     }
 
 
+def _requested_defer_loading(tool: Mapping[str, object], function_tool: Mapping[str, object]) -> object:
+    tool_level: Final = tool.get("defer_loading")
+    return tool_level if tool_level is not None else function_tool.get("defer_loading")
+
+
+def _defer_loading_field(
+    tool: Mapping[str, object],
+    function_tool: Mapping[str, object],
+    *,
+    tool_search: Literal["hosted", "client", "absent"],
+    drop_params: bool,
+) -> Mapping[str, bool]:
+    defer_loading: Final = _requested_defer_loading(tool, function_tool)
+    if defer_loading is None or (defer_loading is False and tool_search == "absent"):
+        return {}
+    if tool_search == "client":
+        verbose_logger.debug(
+            "Chat provider: dropping defer_loading=%r beside a client-executed tool_search", defer_loading
+        )
+        return {}
+    if isinstance(defer_loading, bool) and tool_search == "hosted":
+        return {"defer_loading": defer_loading}
+    if drop_params:
+        verbose_logger.debug("Chat provider: dropping defer_loading=%r under drop_params", defer_loading)
+        return {}
+    problem: Final = (
+        'defer_loading needs a {"type": "tool_search"} tool in the same request'
+        if isinstance(defer_loading, bool)
+        else f"defer_loading must be a boolean, got {defer_loading!r}"
+    )
+    raise litellm.UnsupportedParamsError(
+        message=f"{problem}. To drop it, set `litellm.drop_params=True` or pass `drop_params=True`."
+    )
+
+
+def _without_dropped_defer_loading(
+    tool: Mapping[str, object], *, tool_search: Literal["hosted", "client", "absent"], drop_params: bool
+) -> Mapping[str, object]:
+    if tool.get("type") != "function" or "defer_loading" not in tool:
+        return tool
+    kept: Final = _defer_loading_field(tool, {}, tool_search=tool_search, drop_params=drop_params)
+    return {**{key: value for key, value in tool.items() if key != "defer_loading"}, **kept}
+
+
 class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
     """
     Handler for transforming /chat/completions api requests to litellm.responses requests
@@ -559,9 +788,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         *,
         drop_params: bool = False,
         keep_prompt_cache_breakpoints: bool = False,
+        replay_tool_search: bool = False,
+        include_namespace: bool = False,
     ) -> tuple[list[object], str | None]:
         converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
-            messages, drop_params=drop_params
+            messages,
+            drop_params=drop_params,
+            replay_tool_search=replay_tool_search,
+            include_namespace=include_namespace,
         )
         return (
             converted_input_items
@@ -571,7 +805,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
 
     def _convert_chat_completion_messages_to_responses_input(
-        self, messages: list["AllMessageValues"], *, drop_params: bool = False
+        self,
+        messages: list["AllMessageValues"],
+        *,
+        drop_params: bool = False,
+        replay_tool_search: bool = False,
+        include_namespace: bool = False,
     ) -> tuple[list[object], str | None]:
         input_items: Final[list[object]] = []
         instructions: str | None = None
@@ -653,7 +892,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                         )
                     )
             elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
-                input_items.extend(_reasoning_input_items(msg))
+                input_items.extend(_assistant_replay_items(msg, replay_tool_search))
                 if content:
                     input_items.append(
                         {
@@ -676,6 +915,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                             input_tool_call["name"] = function["name"]
                         if "arguments" in function:
                             input_tool_call["arguments"] = function["arguments"]
+                        if include_namespace and (
+                            namespace := _tool_call_namespace(cast(object, tool_call))  # cast-ok: untyped
+                        ):
+                            input_tool_call["namespace"] = namespace
                         input_items.append(input_tool_call)
                     elif isinstance(custom, dict):
                         input_items.append(
@@ -690,7 +933,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                         raise ValueError(f"tool call not supported: {tool_call}")
             elif content is not None:
                 if role == "assistant":
-                    input_items.extend(_reasoning_input_items(msg))
+                    input_items.extend(_assistant_replay_items(msg, replay_tool_search))
                 input_items.append(
                     {
                         "type": "message",
@@ -701,7 +944,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     }
                 )
             elif role == "assistant":
-                input_items.extend(_reasoning_input_items(msg))
+                input_items.extend(_assistant_replay_items(msg, replay_tool_search))
 
         return input_items, instructions
 
@@ -709,6 +952,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         self,
         optional_params: dict,
         responses_api_request: "ResponsesAPIOptionalRequestParams",
+        drop_params: bool = False,
     ) -> None:
         """Map optional_params into responses_api_request (mutates in place)."""
         for key, value in optional_params.items():
@@ -718,7 +962,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                 responses_api_request["max_output_tokens"] = value
             elif key == "tools" and value is not None:
                 responses_api_request["tools"] = self._convert_tools_to_responses_format(
-                    cast(list[dict[str, object]], value)
+                    cast(list[dict[str, object]], value), drop_params=drop_params
                 )
             elif key == "response_format":
                 text_format = self._transform_response_format_to_text_format(value)
@@ -807,6 +1051,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         litellm_logging_obj: "LiteLLMLoggingObj",
         client: object | None = None,
     ) -> dict:
+        optional_params = self._extract_extra_body_params(optional_params)
         provider: Final = _served_provider(litellm_params)
         base_model: Final = litellm_params.get("base_model")
         supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model, provider) or (
@@ -818,6 +1063,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             messages,
             drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
             keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
+            replay_tool_search=_requests_tool_search(cast(object, optional_params)),  # cast-ok: upcast
+            include_namespace=True,
         )
         # OpenAI's Responses API rejects an empty input. For a system-only
         # request, carry the system message as a system-role input item instead
@@ -842,8 +1089,6 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
         instructions: Final = None if is_system_only_request else converted_instructions
 
-        optional_params = self._extract_extra_body_params(optional_params)
-
         # Build responses API request using the reverse transformation logic
         responses_api_request: Final = ResponsesAPIOptionalRequestParams()
 
@@ -851,7 +1096,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         if instructions:
             responses_api_request["instructions"] = instructions
 
-        self._map_optional_params_to_responses_api_request(optional_params, responses_api_request)
+        self._map_optional_params_to_responses_api_request(
+            optional_params,
+            responses_api_request,
+            drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
+        )
 
         stream: Final = optional_params.get("stream") or litellm_params.get("stream", False)
         verbose_logger.debug("Chat provider: Stream parameter: %s", stream)
@@ -897,6 +1146,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
     def _convert_response_output_to_choices(
         output_items: Sequence[object],
         handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None,
+        structured_output: bool = False,
     ) -> list["Choices"]:
         """
         Convert Responses API output items to chat completion choices.
@@ -904,6 +1154,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         Args:
             output_items: List of items from ResponsesAPIResponse.output
             handle_raw_dict_callback: Optional callback for handling raw dict items
+            structured_output: The chat request carried a json response_format, so only the
+                first message item becomes content
 
         Returns:
             List of Choices objects
@@ -1010,7 +1262,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     # Tool calls accumulate into the single trailing tool_calls choice
                     # like the typed branches above; a choice per call would hide every
                     # call after choices[0] from chat clients
-                    accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
+                    accumulated_tool_calls.append(
+                        tool_call_dict_from_output_item(raw_item, tool_call_index, include_namespace=True)
+                    )
                     tool_call_index += 1
                 elif handle_raw_dict_callback is not None:
                     choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@@ -1019,21 +1273,31 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             else:
                 pass  # don't fail request if item in list is not supported
 
+        tool_search_items: Final = _tool_search_items_from_output_items(output_items)
         if not choices and not accumulated_tool_calls:
-            return choices
+            if not tool_search_items:
+                return choices
+            return [
+                LiteLLMResponsesTransformationHandler._build_empty_choice(
+                    output_items, "stop", tool_search_items=tool_search_items
+                )
+            ]
         return [
             _fold_choices_into_one(
                 choices=choices,
                 tool_calls=accumulated_tool_calls,
                 trailing_reasoning_content=reasoning_content,
                 trailing_reasoning_item=pending_reasoning_item,
+                tool_search_items=tool_search_items,
+                structured_output=structured_output,
             )
         ]
 
     @staticmethod
-    def _build_empty_incomplete_choice(
+    def _build_empty_choice(
         output_items: Sequence[object],
-        finish_reason: Literal["length", "content_filter"],
+        finish_reason: Literal["stop", "length", "content_filter"],
+        tool_search_items: Sequence[Mapping[str, object]] = (),
     ) -> "Choices":
         from litellm.types.utils import Choices, Message
 
@@ -1045,9 +1309,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             if summary_block.get("text")
         )
         message: Final = Message(
-            content="",
+            content=None if tool_search_items else "",
             reasoning_content=reasoning_content if reasoning_content else None,
             reasoning_items=_as_chat_reasoning_items(reasoning_items),
+            provider_specific_fields=_tool_search_provider_fields(tool_search_items),
         )
         return Choices(message=message, finish_reason=finish_reason, index=0)
 
@@ -1157,6 +1422,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         choices: Final = self._convert_response_output_to_choices(
             output_items=output_items,
             handle_raw_dict_callback=self._handle_raw_dict_response_item,
+            structured_output=is_structured_output_request(
+                cast("Mapping[str, object]", optional_params)  # cast-ok: optional_params is an untyped dict
+            ),
         )
 
         response_is_incomplete: Final = raw_response.status == "incomplete" or (
@@ -1171,7 +1439,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                 raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None
             )
             if len(choices) == 0:
-                choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason))
+                choices.append(self._build_empty_choice(output_items, incomplete_finish_reason))
             else:
                 for choice in choices:
                     choice.finish_reason = incomplete_finish_reason
@@ -1219,8 +1487,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
         sync_stream: bool,
         json_mode: bool | None = False,
+        structured_output: bool = False,
     ) -> BaseModelResponseIterator:
-        return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
+        return OpenAiResponsesToChatCompletionStreamIterator(
+            streaming_response, sync_stream, json_mode, structured_output=structured_output
+        )
 
     def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
         if role in ("user", "system", "developer", "tool"):
@@ -1369,14 +1640,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             return result
 
     def _convert_tools_to_responses_format(
-        self, tools: list[dict[str, object]]
+        self, tools: list[dict[str, object]], drop_params: bool = False
     ) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]:
         """Convert chat completion tools to responses API tools format"""
+        tool_search: Final[Literal["hosted", "client", "absent"]] = (
+            "client"
+            if any(tool.get("type") == "tool_search" and tool.get("execution") == "client" for tool in tools)
+            else "hosted"
+            if any(tool.get("type") == "tool_search" or is_anthropic_tool_search_chat_tool(tool) for tool in tools)
+            else "absent"
+        )
         responses_tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = []
         for tool in tools:
             # convert function tool from chat completion to responses API format
-            if tool.get("type") == "function" and isinstance(tool.get("function"), dict):
-                function_tool = cast(ChatCompletionToolParamFunctionChunk, tool.get("function"))
+            if tool.get("type") == "function" and isinstance(function_payload := tool.get("function"), dict):
+                function_tool = cast(ChatCompletionToolParamFunctionChunk, function_payload)  # cast-ok: client tool
                 responses_tools.append(
                     FunctionToolParam(
                         name=function_tool["name"],
@@ -1384,6 +1662,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                         strict=function_tool.get("strict"),
                         type="function",
                         description=function_tool.get("description"),
+                        **_defer_loading_field(tool, function_tool, tool_search=tool_search, drop_params=drop_params),
                     )
                 )
             elif tool.get("type") == "custom" and isinstance(custom_payload := tool.get("custom"), dict):
@@ -1398,7 +1677,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     flat_custom["format"] = convert_custom_tool_format_to_responses_shape(custom_payload["format"])
                 responses_tools.append(flat_custom)
             else:
-                responses_tools.append(tool)
+                responses_tools.append(
+                    _without_dropped_defer_loading(tool, tool_search=tool_search, drop_params=drop_params)
+                )
 
         return cast(list["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools)
 
@@ -1587,11 +1868,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
         streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
         sync_stream: bool,
         json_mode: bool | None = False,
+        structured_output: bool = False,
     ):
         super().__init__(streaming_response, sync_stream, json_mode)
         self._chat_completion_id: str | None = None
         self._served_service_tier: str | None = None
         self._tool_call_index_map: dict[int, int] = {}  # mutable-ok: per-stream accumulator state
+        self._structured_output: Final = structured_output
+        self._first_message_item_done: bool = False
 
     def _handle_string_chunk(
         self, str_line: Union[str, "BaseModel"]
@@ -1692,7 +1976,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
             # New output item added
             output_item = parsed_chunk.get("item", {})
             if output_item.get("type") in ("function_call", "custom_tool_call"):
-                converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
+                converted: Final = tool_call_dict_from_output_item(
+                    output_item, parsed_chunk.get("output_index", 0), include_namespace=True
+                )
                 provider_specific_fields: Final = converted.get("provider_specific_fields")
 
                 function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@@ -1705,14 +1991,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index(
                     tool_call_index_map, parsed_chunk.get("output_index", 0)
                 )
-                tool_call_chunk: Final = ChatCompletionToolCallChunk(
+                tool_call_chunk: Final = _ChatToolCallDict(
                     id=converted["id"],
                     index=tool_call_index,
                     type="function",
                     function=function_chunk,
+                    **({"provider_specific_fields": provider_specific_fields} if provider_specific_fields else {}),
                 )
-                if provider_specific_fields:
-                    tool_call_chunk.provider_specific_fields = provider_specific_fields
 
                 return ModelResponseStream(
                     choices=[
@@ -1768,7 +2053,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                                 delta=Delta(
                                     tool_calls=(
                                         tool_call_dict_from_output_item(
-                                            output_item, parsed_chunk.get("output_index", 0)
+                                            output_item, parsed_chunk.get("output_index", 0), include_namespace=True
                                         ),
                                     )
                                 ),
@@ -1863,6 +2148,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                         delta=Delta(
                             content="",
                             reasoning_items=terminal_reasoning_items_typed,
+                            provider_specific_fields=_tool_search_provider_fields(
+                                _tool_search_items_from_output_items(cast(object, output_items))  # cast-ok: untyped
+                            ),
                         ),
                         finish_reason=finish_reason,
                     )
@@ -1903,13 +2191,31 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
         """
         verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk)
         self._remember_served_service_tier(chunk)
-        return self._with_served_service_tier(
-            self._with_stream_scoped_id(
-                OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
-                    chunk, tool_call_index_map=self._tool_call_index_map
-                )
+        translated: Final = (
+            self._empty_content_chunk()
+            if self._is_text_delta_after_structured_answer(chunk)
+            else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
+                chunk, tool_call_index_map=self._tool_call_index_map
             )
         )
+        if self._structured_output and _is_message_item_done(chunk):
+            self._first_message_item_done = True
+        return self._with_served_service_tier(self._with_stream_scoped_id(translated))
+
+    def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool:
+        """Under a structured response_format the first message item is the whole answer, so
+        text from any later message item is dropped, matching the non-stream fold."""
+        return (
+            self._structured_output
+            and self._first_message_item_done
+            and _responses_event_type(chunk) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA.value
+        )
+
+    @staticmethod
+    def _empty_content_chunk() -> "ModelResponseStream":
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+        return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason=None)])
 
     def _remember_served_service_tier(self, chunk: dict[str, object]) -> None:
         response_payload: Final = chunk.get("response")

litellm/cost_calculator.py

@@ -10,7 +10,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 
 from httpx import Response
-from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict, assert_never
 
 import litellm
@@ -20,6 +20,7 @@ from litellm.constants import (
     DEFAULT_MAX_LRU_CACHE_SIZE,
     DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND,
     OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES,
+    REALTIME_METERED_USAGE_KEY,
 )
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.cost_discount import resolve_cost_discount
@@ -117,6 +118,7 @@ from litellm.types.llms.openai import (
     ResponseAPIUsage,
     ResponsesAPIResponse,
 )
+from litellm.types.llms.xai import RealtimeMeteredUsage
 from litellm.types.rerank import RerankBilledUnits, RerankResponse
 from litellm.types.utils import (
     CallTypesLiteral,
@@ -554,12 +556,11 @@ def cost_per_token(
 
     assert custom_llm_provider is not None  # caller-supplied or get_llm_provider
 
-    model_without_prefix = model
     model_parts: Final = model.split("/", 1)
-    if len(model_parts) > 1:
-        model_without_prefix = model_parts[1]
-    else:
-        model_without_prefix = model
+    model_is_provider_relative: Final = (
+        model_with_provider == f"{custom_llm_provider}/{model}" and model_with_provider in model_cost_ref
+    )
+    model_without_prefix: Final = model if len(model_parts) == 1 or model_is_provider_relative else model_parts[1]
     """
     Code block that formats model to lookup in litellm.model_cost
     Option1. model = "bedrock/ap-northeast-1/anthropic.claude-instant-v1". This is the most accurate since it is region based. Should always be option 1
@@ -828,7 +829,14 @@ def _get_hidden_str_for_cost_calc(hidden_params: object, key: str) -> str | None
 
 
 _NON_TOKEN_RATE_FIELDS: Final = frozenset(
-    {"cost_per_second", "input_cost_per_second", "output_cost_per_second", "input_cost_per_query", "tiered_pricing"}
+    {
+        "cost_per_second",
+        "input_cost_per_second",
+        "output_cost_per_second",
+        "input_cost_per_query",
+        "input_cost_per_text_message",
+        "tiered_pricing",
+    }
 )
 
 
@@ -1394,7 +1402,7 @@ def completion_cost(
     total_time: float | None = 0.0,  # used for replicate, sagemaker
     call_type: CallTypesLiteral | None = None,
     ### REGION ###
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     region_name=None,  # used for bedrock pricing
     ### IMAGE GEN ###
     size: str | None = None,
@@ -1516,6 +1524,9 @@ def completion_cost(
 
         explicit_pricing: Final = custom_pricing is True or base_model is not None
         custom_priced_request_model: Final = model if custom_pricing is True else None
+        configured_cost_per_query: Final = (
+            _configured_search_cost_per_query(litellm_logging_obj) if custom_pricing else None
+        )
         base_model_provider: Final = (
             _base_model_pricing_provider(
                 base_model,
@@ -1537,6 +1548,7 @@ def completion_cost(
             region_name=region_name,
         )
 
+        requested_model: Final = model
         potential_model_names: Final = [
             selected_model,
             _get_response_model(completion_response),
@@ -1696,15 +1708,17 @@ def completion_cost(
                                 video_generation_cost,
                             )
 
-                            return (
-                                video_generation_cost(
-                                    model=model,
-                                    duration_seconds=duration_seconds,
-                                    custom_llm_provider=custom_llm_provider,
-                                    model_info=_video_model_info,
-                                    video_resolution=video_resolution,
-                                )
-                                * video_count
+                            return video_generation_cost(
+                                model=model,
+                                duration_seconds=duration_seconds,
+                                custom_llm_provider=custom_llm_provider,
+                                model_info=_video_model_info,
+                                video_resolution=video_resolution,
+                            ) * video_count + _video_input_image_cost(
+                                requested_model or model,
+                                custom_llm_provider,
+                                _video_model_info,
+                                getattr(completion_response, "usage", None),
                             )
                     # Fallback to default video cost calculation if no duration available
                     return default_video_cost_calculator(
@@ -1762,11 +1776,15 @@ def completion_cost(
                     (
                         prompt_cost,
                         completion_cost_result,
-                    ) = search_provider_cost_per_query(
-                        model=search_model,
-                        custom_llm_provider=custom_llm_provider,
-                        number_of_queries=number_of_queries,
-                        optional_params=optional_params,
+                    ) = (
+                        (number_of_queries * configured_cost_per_query, 0.0)
+                        if configured_cost_per_query is not None
+                        else search_provider_cost_per_query(
+                            model=search_model,
+                            custom_llm_provider=custom_llm_provider,
+                            number_of_queries=number_of_queries,
+                            optional_params=optional_params,
+                        )
                     )
 
                     # Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost)
@@ -2152,6 +2170,38 @@ def response_cost_calculator(
         raise e
 
 
+def _video_input_image_cost(
+    model: str,
+    custom_llm_provider: str | None,
+    model_info: ModelInfo | None,
+    usage: object,
+) -> float:
+    match usage:
+        case {"billable_input_image_count": int(billable_input_images)} if billable_input_images > 0:
+            deployment_rate: Final = model_info.get("input_cost_per_image") if model_info is not None else None
+            image_rate: Final = (
+                deployment_rate
+                if deployment_rate is not None
+                else litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get(
+                    "input_cost_per_image"
+                )
+            )
+            return billable_input_images * image_rate if isinstance(image_rate, (int, float)) else 0.0
+        case _:
+            return 0.0
+
+
+_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
+
+
+def _configured_search_cost_per_query(litellm_logging_obj: LitellmLoggingObject | None) -> float | None:
+    logged_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
+    if logged_params is None:
+        return None
+    cost_per_query: Final = _LOGGED_LITELLM_PARAMS.validate_python(logged_params).get("input_cost_per_query")
+    return float(cost_per_query) if isinstance(cost_per_query, (int, float)) else None
+
+
 def _deployment_model_info(
     litellm_logging_obj: LitellmLoggingObject | None,
     custom_pricing: bool | None,
@@ -3101,31 +3151,43 @@ def _live_session_wall_clock_seconds(litellm_logging_obj: LitellmLoggingObject |
     return 0.0
 
 
-def _live_input_cost_per_second(model_name: str, custom_llm_provider: str) -> float | None:
+def _live_model_rate(
+    model_name: str,
+    custom_llm_provider: str,
+    rate_field: Literal["input_cost_per_second", "input_cost_per_text_message"],
+) -> float | None:
     model_info: Final = _lookup_model_info_or_none(model_name, custom_llm_provider)
     if model_info is None:
         return None
-    rate: Final = model_info.get("input_cost_per_second")
+    rate: Final = model_info.get(rate_field)
     if isinstance(rate, (int, float)) and not isinstance(rate, bool) and math.isfinite(rate):
         return float(rate)
     return None
 
 
-def _first_live_input_cost_per_second(
+def _first_live_model_rate(
     potential_model_names: Sequence[str | None],
     custom_llm_provider: str,
+    rate_field: Literal["input_cost_per_second", "input_cost_per_text_message"],
 ) -> float:
     return next(
         (
             rate
             for model_name in potential_model_names
             if model_name is not None
-            and (rate := _live_input_cost_per_second(model_name, custom_llm_provider)) is not None
+            and (rate := _live_model_rate(model_name, custom_llm_provider, rate_field)) is not None
         ),
         0.0,
     )
 
 
+def _xai_metered_usage(litellm_logging_obj: LitellmLoggingObject | None) -> RealtimeMeteredUsage | None:
+    if litellm_logging_obj is None:
+        return None
+    usage: Final = litellm_logging_obj.model_call_details.get(REALTIME_METERED_USAGE_KEY)
+    return usage if isinstance(usage, RealtimeMeteredUsage) else None
+
+
 def _live_response_token_costs(
     response: _LiveResponsePayload,
     custom_llm_provider: str,
@@ -3285,9 +3347,28 @@ def handle_realtime_stream_cost_calculation(
     live_model_names: Final = (
         (custom_pricing_model, base_pricing_model, received_model, litellm_model_name) if live_event_seen else ()
     )
-    live_voice_cost: Final = voice_seconds * _first_live_input_cost_per_second(
+    live_voice_cost: Final = voice_seconds * _first_live_model_rate(
         live_model_names,
         custom_llm_provider,
+        "input_cost_per_second",
+    )
+    xai_metered: Final = (
+        _xai_metered_usage(litellm_logging_obj) if custom_llm_provider == LlmProviders.XAI.value else None
+    )
+    xai_audio_seconds: Final = (
+        xai_metered.audio_seconds_sent + xai_metered.audio_seconds_received
+        if xai_metered is not None and xai_metered.push_to_talk
+        else _live_session_wall_clock_seconds(litellm_logging_obj)
+        if custom_llm_provider == LlmProviders.XAI.value
+        else 0.0
+    )
+    audio_duration_cost: Final = xai_audio_seconds * _first_live_model_rate(
+        potential_model_names,
+        custom_llm_provider,
+        "input_cost_per_second",
+    )
+    text_message_cost: Final = (xai_metered.text_messages if xai_metered is not None else 0) * _first_live_model_rate(
+        potential_model_names, custom_llm_provider, "input_cost_per_text_message"
     )
     live_response_costs: Final = tuple(
         _live_response_token_costs(
@@ -3301,10 +3382,17 @@ def handle_realtime_stream_cost_calculation(
     live_response_output_cost: Final = sum(costs[1] for costs in live_response_costs)
     input_cost: Final = input_cost_per_token + live_response_input_cost
     output_cost: Final = output_cost_per_token + live_response_output_cost
-    total_cost: Final = input_cost + output_cost + transcription_cost + live_voice_cost
+    total_cost: Final = (
+        input_cost + output_cost + transcription_cost + live_voice_cost + audio_duration_cost + text_message_cost
+    )
     additional_costs: Final = tuple(
         (cost_name, cost)
-        for cost_name, cost in (("transcription_cost", transcription_cost), ("live_voice_cost", live_voice_cost))
+        for cost_name, cost in (
+            ("transcription_cost", transcription_cost),
+            ("live_voice_cost", live_voice_cost),
+            ("audio_duration_cost", audio_duration_cost),
+            ("text_message_cost", text_message_cost),
+        )
         if cost > 0
     )
 

litellm/decisions/call.py

@@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field, replace
 from functools import partial
 from types import MappingProxyType
-from typing import Final, TypeAlias
+from typing import Final, Literal, TypeAlias
 
 import httpx
 from pydantic import ConfigDict, TypeAdapter, ValidationError
@@ -10,6 +10,7 @@ from typing_extensions import assert_never
 
 import litellm
 from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
+from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
     async_convert_urls_to_base64,
@@ -31,10 +32,17 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir
 from litellm.types.decisions import (
     DecisionQuestion,
+    DecisionsAsset,
+    DecisionsAssetReference,
+    DecisionsAssets,
     DecisionsAudio,
     DecisionsBase64Image,
     DecisionsImage,
+    DecisionsIRAnswer,
+    DecisionsIRBoxQuestion,
     DecisionsIRMessages,
+    DecisionsIRQuestion,
+    DecisionsIRRefusal,
     DecisionsIRRequest,
     DecisionsIRResponse,
     DecisionsIRState,
@@ -48,12 +56,13 @@ from litellm.types.decisions import (
     OpenAIDecisionResponse,
     UnsupportedDecisionsRequest,
 )
-from litellm.types.utils import LlmProviders
+from litellm.types.utils import LlmProviders, ProviderSpecificHeader
 from litellm.utils import ProviderConfigManager
 
 SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]]
 OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]]
 DecisionsImages: TypeAlias = Sequence[str | DecisionsBase64Image | Mapping[str, object]]
+DecisionsAssetEntries: TypeAlias = Mapping[str, DecisionsAsset | Mapping[str, object]]
 DecisionsAudioClips: TypeAlias = Sequence[DecisionsAudio | Mapping[str, object]]
 DecisionsVideoClips: TypeAlias = Sequence[DecisionsVideo | Mapping[str, object]]
 DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequestBody
@@ -66,12 +75,22 @@ _SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(
 _IMAGES_ADAPTER: Final[TypeAdapter[tuple[DecisionsImage, ...] | None]] = TypeAdapter(
     tuple[DecisionsImage, ...] | None, config=ConfigDict(title="images")
 )
+_PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[
+    TypeAdapter[ProviderSpecificHeader | tuple[ProviderSpecificHeader, ...] | None]
+] = TypeAdapter(ProviderSpecificHeader | tuple[ProviderSpecificHeader, ...] | None)
 _HANDLER: Final = BaseLLMHTTPHandler()
 _BASE64_MARKER: Final = ";base64,"
 
 
 _IMAGE_ADAPTER: Final[TypeAdapter[DecisionsImage]] = TypeAdapter(DecisionsImage, config=ConfigDict(title="images"))
 _IMAGE_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
+_ASSETS_ADAPTER: Final[TypeAdapter[DecisionsAssets | None]] = TypeAdapter(
+    DecisionsAssets | None, config=ConfigDict(title="assets")
+)
+_ASSET_ENTRY_ADAPTER: Final[TypeAdapter[tuple[DecisionsAssetReference, DecisionsAsset]]] = TypeAdapter(
+    tuple[DecisionsAssetReference, DecisionsAsset]
+)
+_ASSET_ENTRIES_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
 
 
 def _valid_image(image: object) -> DecisionsImage | None:
@@ -96,6 +115,29 @@ def _no_image(_image: DecisionsImage) -> bool:
     return False
 
 
+def _valid_asset(entry: tuple[str, object]) -> tuple[str, DecisionsAsset] | None:
+    try:
+        return _ASSET_ENTRY_ADAPTER.validate_python(entry)
+    except ValidationError:
+        return None
+
+
+def _asset_entries(assets: object) -> Mapping[str, object]:
+    try:
+        return _ASSET_ENTRIES_ADAPTER.validate_python(assets)
+    except ValidationError:
+        return MappingProxyType({})
+
+
+def _request_assets(assets: object, drops_params: bool) -> Mapping[str, DecisionsAsset] | None:
+    if not drops_params:
+        return _ASSETS_ADAPTER.validate_python(assets)
+    valid: Final = MappingProxyType(
+        dict(entry for entry in map(_valid_asset, _asset_entries(assets).items()) if entry is not None)
+    )
+    return valid or None
+
+
 def _request_images(images: object, drops_params: bool) -> tuple[DecisionsImage, ...]:
     if not drops_params:
         return _IMAGES_ADAPTER.validate_python(images) or ()
@@ -107,6 +149,7 @@ class SystemOneRequestFields:
     state: DecisionsJSON | None
     questions: SystemOneQuestions | None
     images: object = None
+    assets: object = None
     audio: DecisionsAudioClips | None = None
     videos: DecisionsVideoClips | None = None
 
@@ -116,6 +159,7 @@ class SystemOneRequestFields:
                 "state": self.state,
                 "questions": self.questions,
                 "images": _request_images(self.images, drops_params) or None,
+                "assets": _request_assets(self.assets, drops_params),
                 "audio": self.audio,
                 "videos": self.videos,
             }
@@ -158,6 +202,7 @@ class DecisionsCall:
     headers: Mapping[str, str]
     timeout: float | httpx.Timeout | None
     litellm_params: Mapping[str, object] = field(repr=False)
+    caller_questions: tuple[DecisionsIRQuestion, ...]
 
 
 def supported_decisions_providers() -> tuple[str, ...]:
@@ -229,8 +274,9 @@ class _UnsupportedParams:
 
 
 @dataclass(frozen=True, slots=True)
-class _NoInputLeftAfterDrop:
+class _NothingLeftAfterDrop:
     names: tuple[str, ...]
+    missing: Literal["input to decide on", "questions to answer"]
 
 
 @dataclass(frozen=True, slots=True)
@@ -239,12 +285,22 @@ class _MalformedImages:
 
 
 def _unsupported_params(
-    ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig, model: str, custom_llm_provider: str
+    ir_request: DecisionsIRRequest,
+    provider_config: BaseDecisionsConfig,
+    model: str,
+    custom_llm_provider: str,
+    headers: Mapping[str, str],
 ) -> tuple[str, ...]:
     refuses_safety_identifier: Final = (
         ir_request.safety_identifier is not None and not provider_config.supports_safety_identifier
     )
-    refused_image_param: Final = None if provider_config.supports_images else image_param(ir_request.input)
+    refused_image_param: Final = (
+        None if provider_config.supports_images_in(ir_request.input, headers) else image_param(ir_request.input)
+    )
+    refuses_assets: Final = bool(ir_request.assets) and not provider_config.supports_assets(headers)
+    refuses_box_questions: Final = _has_box_question(ir_request.questions) and not (
+        provider_config.supports_box_questions(model, custom_llm_provider, headers)
+    )
     refuses_audio: Final = bool(ir_request.audio) and not provider_config.supports_audio_input(
         model, custom_llm_provider
     )
@@ -254,18 +310,42 @@ def _unsupported_params(
     return (
         *(("safety_identifier",) if refuses_safety_identifier else ()),
         *(() if refused_image_param is None else (refused_image_param,)),
+        *(("assets",) if refuses_assets else ()),
+        *(("box questions",) if refuses_box_questions else ()),
         *(("audio",) if refuses_audio else ()),
         *(("videos",) if refuses_videos else ()),
     )
 
 
+def _has_box_question(questions: Sequence[DecisionsIRQuestion]) -> bool:
+    return any(isinstance(question, DecisionsIRBoxQuestion) for question in questions)
+
+
+def _is_not_box_question(question: DecisionsIRQuestion) -> bool:
+    return not isinstance(question, DecisionsIRBoxQuestion)
+
+
 def _without_unsupported_params(
-    ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig
+    ir_request: DecisionsIRRequest,
+    provider_config: BaseDecisionsConfig,
+    model: str,
+    custom_llm_provider: str,
+    headers: Mapping[str, str],
 ) -> DecisionsIRRequest:
     return replace(
         ir_request,
-        input=ir_request.input if provider_config.supports_images else with_images_where(ir_request.input, _no_image),
+        input=(
+            ir_request.input
+            if provider_config.supports_images_in(ir_request.input, headers)
+            else with_images_where(ir_request.input, _no_image)
+        ),
         safety_identifier=ir_request.safety_identifier if provider_config.supports_safety_identifier else None,
+        assets=ir_request.assets if provider_config.supports_assets(headers) else MappingProxyType({}),
+        questions=(
+            ir_request.questions
+            if provider_config.supports_box_questions(model, custom_llm_provider, headers)
+            else tuple(filter(_is_not_box_question, ir_request.questions))
+        ),
     )
 
 
@@ -286,28 +366,39 @@ def _without_malformed_images(
 
 
 def _provider_ir_request(
-    request: DecisionsRequestFormat,
+    ir_request: DecisionsIRRequest,
     *,
-    safety_identifier: str | None,
     provider_config: BaseDecisionsConfig,
     model: str,
     custom_llm_provider: str,
+    headers: Mapping[str, str],
     kwargs: Mapping[str, object],
-) -> DecisionsIRRequest | _UnsupportedParams | _NoInputLeftAfterDrop | _MalformedImages:
-    ir_request: Final = _without_malformed_images(_ir_request(request, safety_identifier), _drops_params(kwargs))
-    if isinstance(ir_request, _MalformedImages):
-        return ir_request
-    unsupported: Final = _unsupported_params(ir_request, provider_config, model, custom_llm_provider)
+) -> DecisionsIRRequest | _UnsupportedParams | _NothingLeftAfterDrop:
+    unsupported: Final = _unsupported_params(ir_request, provider_config, model, custom_llm_provider, headers)
     if not unsupported:
         return ir_request
     if not _drops_params(kwargs):
         return _UnsupportedParams(names=unsupported)
-    supported: Final = _without_unsupported_params(ir_request, provider_config)
+    supported: Final = _without_unsupported_params(ir_request, provider_config, model, custom_llm_provider, headers)
     if _leaves_no_input(supported.input):
-        return _NoInputLeftAfterDrop(names=unsupported)
+        return _NothingLeftAfterDrop(names=unsupported, missing="input to decide on")
+    if not supported.questions:
+        return _NothingLeftAfterDrop(names=unsupported, missing="questions to answer")
     return supported
 
 
+def _outbound_headers(
+    extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object]
+) -> Mapping[str, str]:
+    forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
+        provider_specific_header=_PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(
+            kwargs.get("provider_specific_header")
+        ),
+        custom_llm_provider=custom_llm_provider,
+    )
+    return MappingProxyType({**(extra_headers or {}), **forwarded})
+
+
 def prepare_call(
     *,
     model: str,
@@ -338,6 +429,7 @@ def prepare_call(
     except ValueError as error:
         raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
     try:
+        headers: Final = _outbound_headers(extra_headers, provider, kwargs)
         request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs)
         request: Final = _fields_without_unsupported_media(
             fields, provider_config, canonical_model, provider, kwargs
@@ -349,14 +441,18 @@ def prepare_call(
             llm_provider=provider,
         ) from error
 
-    resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
+    try:
+        resolved_api_base, resolved_api_key = provider_config.resolve_connection(
+            dynamic_api_base or api_base, dynamic_api_key or api_key
+        )
+    except ValueError as error:
+        raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
     if resolved_api_base is None:
         raise litellm.BadRequestError(
             message=provider_config.missing_api_base_message(provider),
             model=model,
             llm_provider=provider,
         )
-    resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
     if resolved_api_key is None and provider_config.api_key_required:
         raise litellm.AuthenticationError(
             message=f"Missing API key for Decisions provider '{provider}'",
@@ -364,12 +460,24 @@ def prepare_call(
             llm_provider=provider,
         )
 
+    caller_ir_request: Final = _without_malformed_images(
+        _ir_request(request, request_safety_identifier), _drops_params(kwargs)
+    )
+    if isinstance(caller_ir_request, _MalformedImages):
+        raise litellm.BadRequestError(
+            message=(
+                f"Invalid Decisions request: {caller_ir_request.count} input_image part(s) have an image_url that is "
+                "neither a data:image/...;base64, URL nor an http(s) URL"
+            ),
+            model=model,
+            llm_provider=provider,
+        )
     ir_request: Final = _provider_ir_request(
-        request,
-        safety_identifier=request_safety_identifier,
+        caller_ir_request,
         provider_config=provider_config,
         model=canonical_model,
         custom_llm_provider=provider,
+        headers=headers,
         kwargs=kwargs,
     )
     if isinstance(ir_request, _UnsupportedParams):
@@ -382,20 +490,11 @@ def prepare_call(
             model=model,
             llm_provider=provider,
         )
-    if isinstance(ir_request, _MalformedImages):
-        raise litellm.BadRequestError(
-            message=(
-                f"Invalid Decisions request: {ir_request.count} input_image part(s) have an image_url that is "
-                "neither a data:image/...;base64, URL nor an http(s) URL"
-            ),
-            model=model,
-            llm_provider=provider,
-        )
-    if isinstance(ir_request, _NoInputLeftAfterDrop):
+    if isinstance(ir_request, _NothingLeftAfterDrop):
         raise litellm.BadRequestError(
             message=(
                 f"{provider} does not support parameters: {list(ir_request.names)}, for model={model}, "
-                "and dropping them leaves no input to decide on"
+                f"and dropping them leaves no {ir_request.missing}"
             ),
             model=model,
             llm_provider=provider,
@@ -421,9 +520,10 @@ def prepare_call(
         api_base=resolved_api_base,
         api_key=resolved_api_key,
         logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
-        headers=extra_headers or {},
+        headers=headers,
         timeout=timeout,
         litellm_params=kwargs,
+        caller_questions=caller_ir_request.questions,
     )
 
 
@@ -557,14 +657,24 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci
     return formatted
 
 
+def _caller_answers(call: DecisionsCall, answers: Sequence[DecisionsIRAnswer]) -> tuple[DecisionsIRAnswer, ...]:
+    sent_answers: Final = iter(answers)
+    return tuple(
+        next(sent_answers) if question in call.ir_request.questions else DecisionsIRRefusal()
+        for question in call.caller_questions
+    )
+
+
 def _formatted_response(
     response: DecisionsIRResponse, call: DecisionsCall
 ) -> DecisionsResponse | OpenAIDecisionResponse:
+    caller_request: Final = replace(call.ir_request, questions=call.caller_questions)
+    caller_response: Final = replace(response, answers=_caller_answers(call, response.answers))
     match call.request:
         case DecisionsRequestBody():
-            return ir_to_systemone_response(response, call.ir_request)
+            return ir_to_systemone_response(caller_response, caller_request)
         case OpenAIDecisionRequestBody():
-            return ir_to_openai_response(response, call.ir_request, call.requested_model)
+            return ir_to_openai_response(caller_response, caller_request, call.requested_model)
         case _:
             assert_never(call.request)
 

litellm/decisions/main.py

@@ -27,7 +27,7 @@ _SYSTEMONE_QUESTIONS_ADAPTER: Final[TypeAdapter[SystemOneQuestions | None]] = Ty
 _OPENAI_QUESTIONS_ADAPTER: Final[TypeAdapter[OpenAIDecisionQuestions | None]] = TypeAdapter(
     OpenAIDecisionQuestions | None
 )
-_SYSTEMONE_FIELDS: Final = frozenset({"state", "images"})
+_SYSTEMONE_FIELDS: Final = frozenset({"state", "images", "assets"})
 
 
 def _request_fields(
@@ -55,6 +55,7 @@ def _request_fields(
         state=_STATE_ADAPTER.validate_python(kwargs["state"]),
         questions=_SYSTEMONE_QUESTIONS_ADAPTER.validate_python(questions),
         images=kwargs.get("images"),
+        assets=kwargs.get("assets"),
         audio=audio,
         videos=videos,
     )

litellm/exceptions.py

@@ -10,7 +10,7 @@
 ## LiteLLM versions of the OpenAI Exception Types
 
 import enum
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
 import httpx
@@ -468,7 +468,7 @@ class RateLimitError(openai.RateLimitError):
         category: str | RateLimitErrorCategory = (RateLimitErrorCategory.VENDOR_RATE_LIMIT),
         rate_limit_type: str | RateLimitType | None = None,
         headers: dict[str, str] | None = None,
-        detail: Any = None,
+        detail: object = None,
         body: object | None = None,
     ):
         self.status_code = 429
@@ -1267,6 +1267,10 @@ class MidStreamFallbackError(ServiceUnavailableError):
         self.message = _saved_message
         self.args = (_saved_message,)
 
+        original_headers: Final = getattr(original_exception, "litellm_response_headers", None)
+        if original_headers is not None:
+            self.litellm_response_headers = original_headers
+
     def __str__(self):
         _message = self.message
         if self.num_retries:
@@ -1300,8 +1304,8 @@ class ModifyResponseException(Exception):
         model: str,
         request_data: dict[str, Any],
         guardrail_name: str | None = None,
-        detection_info: dict[str, Any] | None = None,
-        original_response: Any | None = None,
+        detection_info: Mapping[str, object] | None = None,
+        original_response: object = None,
     ):
         self.message = message
         self.model = model
@@ -1334,7 +1338,7 @@ class SensitiveDataRouteException(Exception):
         route_to_model: str,
         session_id: str,
         guardrail_name: str | None = None,
-        detection_info: dict[str, Any] | None = None,
+        detection_info: Mapping[str, object] | None = None,
         message: str | None = None,
         sticky_session_routing: bool = True,
     ):

litellm/files/main.py

@@ -11,7 +11,7 @@ import time
 import uuid as uuid_module
 from collections.abc import Coroutine
 from functools import partial
-from typing import Any, Final, Literal, cast
+from typing import Final, Literal, cast
 
 import httpx
 from openai import AsyncOpenAI, OpenAI
@@ -163,7 +163,7 @@ def create_file(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]:
+) -> OpenAIFileObject | Coroutine[object, None, OpenAIFileObject]:
     """
     Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API.
 
@@ -366,7 +366,15 @@ def file_retrieve(
 
         _is_async: Final = kwargs.pop("is_async", False) is True
 
-        if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
+        if custom_llm_provider == LlmProviders.XAI.value and is_xai_batch_results_id(file_id):
+            response = xai_batch_results_instance.batch_results_file(  # rebind-ok: shares the result slot with the provider branches below
+                _is_async=_is_async,
+                batch_id=file_id,
+                api_base=optional_params.api_base,
+                api_key=optional_params.api_key,
+                timeout=timeout,
+            )
+        elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
             openai_creds: Final = get_openai_credentials(
                 api_base=optional_params.api_base,
                 api_key=optional_params.api_key,
@@ -859,8 +867,8 @@ def file_content(
 ) -> (
     HttpxBinaryResponseContent
     | FileContentStreamingResult
-    | Coroutine[Any, Any, HttpxBinaryResponseContent]
-    | Coroutine[Any, Any, FileContentStreamingResult]
+    | Coroutine[object, None, HttpxBinaryResponseContent]
+    | Coroutine[object, None, FileContentStreamingResult]
 ):
     """
     Returns the contents of the specified file.

litellm/google_genai/adapters/handler.py

@@ -1,7 +1,10 @@
 from collections.abc import AsyncIterator, Coroutine, Mapping
 from typing import Final, cast
 
+from openai.types.chat import ChatCompletionStreamOptionsParam
+
 import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.types.google_genai.adapters import GenerateContentCompletionKwargs
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import ModelResponse
@@ -12,6 +15,23 @@ from .transformation import GoogleGenAIAdapter
 GOOGLE_GENAI_ADAPTER: Final = GoogleGenAIAdapter()
 
 
+def _mark_failure_logged_by_inner_completion(adapter_kwargs: Mapping[str, object]) -> None:
+    outer_logging_obj: Final = adapter_kwargs.get("litellm_logging_obj")
+    if isinstance(outer_logging_obj, LiteLLMLoggingObj):
+        outer_logging_obj.mark_logging_complete(event_type="sync_failure")
+        outer_logging_obj.mark_logging_complete(event_type="async_failure")
+
+
+def _adapter_stream_options(requested: object) -> ChatCompletionStreamOptionsParam | None:
+    match requested:
+        case {"include_usage": True}:
+            return {"include_usage": True}
+        case Mapping():
+            return None
+        case _:
+            return {"include_usage": True}
+
+
 class GenerateContentToCompletionHandler:
     """Handler for transforming generate_content calls to completion format when provider config is None"""
 
@@ -47,9 +67,14 @@ class GenerateContentToCompletionHandler:
                 completion_kwargs["extra_headers"] = extra_kwargs["extra_headers"]
             if "proxy_server_request" in extra_kwargs:
                 completion_kwargs["proxy_server_request"] = extra_kwargs["proxy_server_request"]
+            if "litellm_metadata" in extra_kwargs:
+                completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"]
 
         if stream:
             completion_kwargs["stream"] = stream
+            stream_options: Final = _adapter_stream_options((extra_kwargs or {}).get("stream_options"))
+            if stream_options is not None:
+                completion_kwargs["stream_options"] = stream_options
 
         return GenerateContentCompletionKwargs(**completion_kwargs)
 
@@ -75,33 +100,33 @@ class GenerateContentToCompletionHandler:
 
         try:
             completion_response: Final = await litellm.acompletion(**completion_kwargs)
+        except Exception:
+            _mark_failure_logged_by_inner_completion(kwargs)
+            raise
 
-            if stream:
-                # Check if completion_response is actually a stream or a ModelResponse
-                # This can happen in error cases or when stream is not properly supported
-                if not hasattr(completion_response, "__aiter__"):
-                    # If it's not a stream, treat it as a regular response
-                    generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
-                        cast(ModelResponse, completion_response)
-                    )
-                    return generate_content_response
-                else:
-                    # Transform streaming completion response to generate_content format
-                    transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
-                        completion_response
-                    )
-                    if transformed_stream is not None:
-                        return transformed_stream
-                    raise ValueError("Failed to transform streaming response")
-            else:
-                # Transform completion response back to generate_content format
+        if stream:
+            # Check if completion_response is actually a stream or a ModelResponse
+            # This can happen in error cases or when stream is not properly supported
+            if not hasattr(completion_response, "__aiter__"):
+                # If it's not a stream, treat it as a regular response
                 generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
                     cast(ModelResponse, completion_response)
                 )
                 return generate_content_response
-
-        except Exception as e:
-            raise ValueError(f"Error calling litellm.acompletion for generate_content: {e}")
+            else:
+                # Transform streaming completion response to generate_content format
+                transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
+                    completion_response
+                )
+                if transformed_stream is not None:
+                    return transformed_stream
+                raise ValueError("Failed to transform streaming response")
+        else:
+            # Transform completion response back to generate_content format
+            generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
+                cast(ModelResponse, completion_response)
+            )
+            return generate_content_response
 
     @staticmethod
     def generate_content_handler(
@@ -136,30 +161,30 @@ class GenerateContentToCompletionHandler:
 
         try:
             completion_response: Final = litellm.completion(**completion_kwargs)
+        except Exception:
+            _mark_failure_logged_by_inner_completion(kwargs)
+            raise
 
-            if stream:
-                # Check if completion_response is actually a stream or a ModelResponse
-                # This can happen in error cases or when stream is not properly supported
-                if not hasattr(completion_response, "__iter__"):
-                    # If it's not a stream, treat it as a regular response
-                    generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
-                        cast(ModelResponse, completion_response)
-                    )
-                    return generate_content_response
-                else:
-                    # Transform streaming completion response to generate_content format
-                    transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
-                        completion_response
-                    )
-                    if transformed_stream is not None:
-                        return transformed_stream
-                    raise ValueError("Failed to transform streaming response")
-            else:
-                # Transform completion response back to generate_content format
+        if stream:
+            # Check if completion_response is actually a stream or a ModelResponse
+            # This can happen in error cases or when stream is not properly supported
+            if not hasattr(completion_response, "__iter__"):
+                # If it's not a stream, treat it as a regular response
                 generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
                     cast(ModelResponse, completion_response)
                 )
                 return generate_content_response
-
-        except Exception as e:
-            raise ValueError(f"Error calling litellm.completion for generate_content: {e}")
+            else:
+                # Transform streaming completion response to generate_content format
+                transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
+                    completion_response
+                )
+                if transformed_stream is not None:
+                    return transformed_stream
+                raise ValueError("Failed to transform streaming response")
+        else:
+            # Transform completion response back to generate_content format
+            generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
+                cast(ModelResponse, completion_response)
+            )
+            return generate_content_response

litellm/google_genai/adapters/transformation.py

@@ -158,8 +158,17 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
         self.sent_first_chunk = False
         self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
         self._returned_response = False
+        self.finish_chunk_awaiting_usage: Mapping[str, object] | None = None
         super().__init__(completion_stream)
 
+    def hold_finish_chunk_for_usage(self, finish_chunk: Mapping[str, object]) -> None:
+        self.finish_chunk_awaiting_usage = finish_chunk
+
+    def take_finish_chunk_awaiting_usage(self) -> Mapping[str, object] | None:
+        finish_chunk: Final = self.finish_chunk_awaiting_usage
+        self.finish_chunk_awaiting_usage = None
+        return finish_chunk
+
     def __next__(self):
         try:
             if not hasattr(self.completion_stream, "__iter__"):
@@ -176,6 +185,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                 if transformed_chunk:
                     return transformed_chunk
 
+            finish_chunk: Final = self.take_finish_chunk_awaiting_usage()
+            if finish_chunk:
+                return finish_chunk
             raise StopIteration
         except StopIteration:
             raise
@@ -198,6 +210,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                 if transformed_chunk:
                     return transformed_chunk
 
+            finish_chunk: Final = self.take_finish_chunk_awaiting_usage()
+            if finish_chunk:
+                return finish_chunk
+
             # After the stream is exhausted, check for any remaining accumulated tool calls
             if self.accumulated_tool_calls:
                 try:
@@ -286,6 +302,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                 else:
                     yield str(chunk).encode()
 
+        finish_chunk: Final = self.take_finish_chunk_awaiting_usage()
+        if finish_chunk:
+            yield f"data: {json.dumps(finish_chunk)}\n\n".encode()
+
 
 class GoogleGenAIAdapter:
     """Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format"""
@@ -657,8 +677,7 @@ class GoogleGenAIAdapter:
         # Extract the main response content from streaming chunk
         choice: Final = response.choices[0] if response.choices else None
         if not choice:
-            # Return empty chunk if no choices
-            return None
+            return self._usage_only_streaming_chunk(response, wrapper)
 
         # Handle streaming choice
         if isinstance(choice, StreamingChoices):
@@ -675,7 +694,7 @@ class GoogleGenAIAdapter:
 
         # Only create response chunk if we have parts or it's the final chunk
         if not parts and not finish_reason:
-            return None
+            return self._usage_only_streaming_chunk(response, wrapper)
 
         # Create Google GenAI streaming format response
         streaming_chunk: Final[_JsonDict] = {
@@ -689,20 +708,34 @@ class GoogleGenAIAdapter:
             ]
         }
 
-        # Add usage metadata only in the final chunk (when finish_reason is present)
-        if finish_reason:
-            usage_metadata: Final = (
-                self._map_usage(getattr(response, "usage", None))
-                if hasattr(response, "usage") and getattr(response, "usage", None)
-                else {
-                    "promptTokenCount": 0,
-                    "candidatesTokenCount": 0,
-                    "totalTokenCount": 0,
-                }
-            )
-            streaming_chunk["usageMetadata"] = usage_metadata
+        if not finish_reason:
+            return streaming_chunk
 
-        return streaming_chunk
+        finish_usage: Final[object] = getattr(response, "usage", None)
+        if finish_usage:
+            return {**streaming_chunk, "usageMetadata": self._map_usage(finish_usage)}
+
+        wrapper.hold_finish_chunk_for_usage(
+            {
+                **streaming_chunk,
+                "usageMetadata": {"promptTokenCount": 0, "candidatesTokenCount": 0, "totalTokenCount": 0},
+            }
+        )
+        return None
+
+    def _usage_only_streaming_chunk(
+        self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper
+    ) -> Mapping[str, object] | None:
+        usage: Final[object] = getattr(response, "usage", None)
+        if not usage:
+            return None
+        finish_chunk: Final = wrapper.take_finish_chunk_awaiting_usage()
+        if finish_chunk:
+            return {**finish_chunk, "usageMetadata": self._map_usage(usage)}
+        return {
+            "candidates": [{"content": {"parts": [{"text": ""}], "role": "model"}, "index": 0}],
+            "usageMetadata": self._map_usage(usage),
+        }
 
     def _transform_openai_message_to_google_genai_parts(
         self,

litellm/google_genai/main.py

@@ -145,6 +145,7 @@ class GenerateContentHelper:
             ProviderConfigManager.get_provider_google_genai_generate_content_config(
                 model=model,
                 provider=litellm.LlmProviders(custom_llm_provider),
+                litellm_params=litellm_params.model_dump(include={"base_model", "api_base"}),
             )
         )
 

litellm/images/main.py

@@ -5,6 +5,8 @@ from collections.abc import Coroutine
 from functools import partial
 from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload
 
+from pydantic import TypeAdapter, ValidationError
+
 if TYPE_CHECKING:
     from litellm.images.utils import ImageEditRequestUtils
 
@@ -17,6 +19,9 @@ from litellm import client
 from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL
 from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
 from litellm.exceptions import LiteLLMUnknownProvider
+from litellm.litellm_core_utils.get_llm_provider_logic import (
+    invalid_provider_routing_hints_message,
+)
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.llm_request_utils import flatten_form_field_values
@@ -52,6 +57,7 @@ from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
     LITELLM_IMAGE_VARIATION_PROVIDERS,
     LlmProviders,
+    ProviderRoutingHints,
 )
 from litellm.utils import (
     ImageResponse,
@@ -60,6 +66,8 @@ from litellm.utils import (
     get_optional_params_image_gen,
 )
 
+_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
+
 # Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls
 _ImageEditRequestUtils_cache: Optional["ImageEditRequestUtils"] = None
 
@@ -92,6 +100,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
     model: Final = args[0] if len(args) > 0 else kwargs["model"]
     ### PASS ARGS TO Image Generation ###
     kwargs["aimg_generation"] = True
+    try:
+        provider_hints: Final = _PROVIDER_ROUTING_HINTS.validate_python(kwargs)
+    except ValidationError as error:
+        raise litellm.BadRequestError(
+            message=invalid_provider_routing_hints_message(error), model=model, llm_provider=""
+        ) from error
     custom_llm_provider = None
     try:
         # Use a partial function to pass your keyword arguments
@@ -101,7 +115,11 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
         ctx: Final = contextvars.copy_context()
         func_with_context: Final = partial(ctx.run, func)
 
-        _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None))
+        _, custom_llm_provider, _, _ = get_llm_provider(
+            model=model,
+            custom_llm_provider=provider_hints.get("custom_llm_provider"),
+            api_base=provider_hints.get("api_base"),
+        )
 
         # Await normally
         init_response: Final = await loop.run_in_executor(None, func_with_context)

litellm/integrations/anthropic_cache_control_hook.py

@@ -12,7 +12,7 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and
 import copy
 import os
 import re
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Callable, Container, Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
 
@@ -26,6 +26,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     cache_control_as_prompt_cache_breakpoint,
     with_prompt_cache_breakpoint,
 )
+from litellm.litellm_core_utils.prompt_templates.factory import tool_call_ids_answered_by_client
 from litellm.llms.anthropic.common_utils import (
     is_claude_code_one_shot_subagent_request,
     supports_anthropic_cache_control,
@@ -227,12 +228,12 @@ def _as_object_list(value: object | None) -> list[object] | None:
     return _validated_object_list(value)
 
 
-def _tool_call_carries_cache_breakpoint(tool_call: object, message: object) -> bool:
+def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answered_tool_ids: Container[str]) -> bool:
     if _attribute_or_key(tool_call, "cache_control") is None:
         return False
 
     return not tool_call_is_rebuilt_as_server_tool_use(
-        _attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields")
+        _attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields"), answered_tool_ids
     )
 
 
@@ -424,7 +425,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         system_blocks: Final = (
             sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
         )
-        return system_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages)
+        message_list: Final = tuple(messages)
+        answered_tool_ids: Final = tool_call_ids_answered_by_client(message_list)
+        return system_blocks + sum(
+            AnthropicCacheControlHook._count_cache_control_blocks(msg, answered_tool_ids) for msg in message_list
+        )
 
     @staticmethod
     def count_external_cache_breakpoints(
@@ -509,6 +514,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         so earlier points win when slots are scarce.
         """
         used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
+        answered_tool_ids: Final = tool_call_ids_answered_by_client(messages)
 
         limit_reached = False
         for point in points:
@@ -525,14 +531,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                     limit_reached = True
                     break
 
-                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
+                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], answered_tool_ids):
                     # Client already marked this message; don't overwrite it.
                     continue
 
                 messages[target_index] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
                     messages[target_index], control, openai_dialect
                 )
-                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
+                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], answered_tool_ids):
                     used_blocks += 1
 
             if limit_reached:
@@ -586,22 +592,26 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         return []
 
     @staticmethod
-    def _count_cache_control_blocks(message: object) -> int:
+    def _count_cache_control_blocks(message: object, answered_tool_ids: Container[str] = frozenset()) -> int:
         message_count: Final = 1 if _carries_cache_breakpoint(message) else 0
         content: Final = _as_object_list(_attribute_or_key(message, "content"))
         content_count: Final = sum(1 for block in content if _carries_cache_breakpoint(block)) if content else 0
         tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls"))
         tool_call_count: Final = (
-            sum(1 for tool_call in tool_calls if _tool_call_carries_cache_breakpoint(tool_call, message))
+            sum(
+                1
+                for tool_call in tool_calls
+                if _tool_call_carries_cache_breakpoint(tool_call, message, answered_tool_ids)
+            )
             if tool_calls
             else 0
         )
         return message_count + content_count + tool_call_count
 
     @staticmethod
-    def _message_has_cache_control(message: AllMessageValues) -> bool:
+    def _message_has_cache_control(message: AllMessageValues, answered_tool_ids: Container[str]) -> bool:
         """Return True if the message already carries any cache_control."""
-        return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0
+        return AnthropicCacheControlHook._count_cache_control_blocks(message, answered_tool_ids) > 0
 
     @staticmethod
     def _safe_insert_cache_control_in_message(
@@ -799,12 +809,16 @@ class AnthropicCacheControlHook(CustomPromptManagement):
 
         The hook sees neither the tools nor the request kwargs, so the target dialect
         and the client's breakpoint count outside the message list ride on the points.
-        Builds copies because config-owned point dicts are shared across requests.
+        A zero count replaces an earlier layer's stamp (the chat bridge runs behind the
+        Responses layer, whose outside marks sit in the bridge's own message list) and is
+        otherwise left off. Builds copies because config-owned point dicts are shared
+        across requests.
         """
         with_dialect: Final = AnthropicCacheControlHook._stamped_with_dialect(
             points, model, custom_llm_provider, api_base, prompt_cache_options
         )
-        if external_breakpoints == 0:
+        stamped_earlier: Final = any(EXTERNAL_BREAKPOINTS_STAMP in point for point in with_dialect)
+        if external_breakpoints == 0 and not stamped_earlier:
             return with_dialect
         return AnthropicCacheControlHook._stamped(with_dialect, EXTERNAL_BREAKPOINTS_STAMP, external_breakpoints)
 

litellm/integrations/focus/destinations/ternary_destination.py

@@ -9,12 +9,14 @@ from urllib.parse import quote, urlparse
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 from litellm._logging import verbose_logger
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPEventHooks
+from litellm.types.llms.custom_http import VerifyTypes
 
 from .base import FocusDestination, FocusTimeWindow
 
 if TYPE_CHECKING:
     import httpx
+    from aiohttp import ClientSession
 
 TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0
 
@@ -57,10 +59,17 @@ class _NoRedirectHTTPHandler(AsyncHTTPHandler):
 
     def create_client(
         self,
-        *args: object,
-        **kwargs: object,  # kwargs-ok: passthrough to AsyncHTTPHandler.create_client
+        timeout: float | httpx.Timeout | None,
+        event_hooks: HTTPEventHooks | None,
+        ssl_verify: VerifyTypes | None = None,
+        shared_session: ClientSession | None = None,
     ) -> httpx.AsyncClient:
-        client: Final = super().create_client(*args, **kwargs)  # pyright: ignore[reportArgumentType]  # passthrough
+        client: Final = super().create_client(
+            timeout=timeout,
+            event_hooks=event_hooks,
+            ssl_verify=ssl_verify,
+            shared_session=shared_session,
+        )
         client.follow_redirects = False
         return client
 

litellm/integrations/focus/focus_logger.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import os
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Final, cast
 
 import litellm
 from litellm._logging import verbose_logger
@@ -28,6 +28,8 @@ DEFAULT_DRY_RUN_LIMIT: Final = 500
 class FocusLogger(CustomLogger):
     """Coordinates Focus export jobs across transformer/serializer/destination layers."""
 
+    export_job_name: ClassVar[str] = FOCUS_USAGE_DATA_JOB_NAME
+
     def __init__(
         self,
         *,
@@ -120,14 +122,14 @@ class FocusLogger(CustomLogger):
                 pod_lock_manager = getattr(writer, "pod_lock_manager", None)
 
         if pod_lock_manager and pod_lock_manager.redis_cache:
-            acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=FOCUS_USAGE_DATA_JOB_NAME)
+            acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=self.export_job_name)
             if not acquired:
                 verbose_logger.debug("Focus export: unable to acquire pod lock")
                 return
             try:
                 await self._run_scheduled_export()
             finally:
-                await pod_lock_manager.release_lock(cronjob_id=FOCUS_USAGE_DATA_JOB_NAME)
+                await pod_lock_manager.release_lock(cronjob_id=self.export_job_name)
         else:
             await self._run_scheduled_export()
 

litellm/integrations/opentelemetry_utils/gen_ai_semconv.py

@@ -29,21 +29,25 @@ Events:
   included only when content capture is enabled).
 """
 
+from collections.abc import Callable
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Final, Protocol
 
 from litellm.integrations.otel.mappers.utils import drop_none
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
 if TYPE_CHECKING:
+    from opentelemetry._logs import SeverityNumber as _SeverityNumber
     from opentelemetry.trace import Span as _Span
 
     from litellm.integrations.opentelemetry import OpenTelemetryConfig
 
-    Span = _Span | Any
+    Span = _Span
 else:
-    Span = Any
+
+    class Span(Protocol):
+        def set_attribute(self, key: str, value: object) -> None: ...
 
 
 # OTEL_SEMCONV_STABILITY_OPT_IN is a comma-separated list of category-specific
@@ -128,7 +132,7 @@ class OTELGenAISemconvMixin:
 
         def _to_ns(self, dt: datetime) -> int: ...
 
-        def _otel_log_types(self) -> tuple[Any, Any]: ...
+        def _otel_log_types(self) -> tuple[Callable[..., object], type[_SeverityNumber]]: ...
 
     @property
     def _gen_ai_semconv_latest_experimental(self) -> bool:

litellm/integrations/otel/README.md

@@ -194,7 +194,10 @@ matching the FastAPI server span); only a genuine error sets `ERROR`.
 - **Server spans** (one per HTTP route) are created by the
   `opentelemetry-instrumentation-fastapi` package. It stamps `http.*` attributes
   and extracts inbound `traceparent` headers. This package does **not** create
-  or modify server spans — request routes never touch spans.
+  server spans. It renames passthrough server spans, sets `litellm.call_id` on
+  them, and returns each sampled server span to the client in a `traceresponse`
+  response header (`trace_response.py`). A relayed upstream `traceresponse` is
+  kept, and the server span links to it when that upstream restarted the trace.
 - **Gen-AI spans** (LLM calls, guardrails, internal service calls) are created
   by this package from LiteLLM's logging callbacks. Request-level spans parent to
   the server span via the captured anchor; DB/service spans parent to the active

litellm/integrations/otel/logger.py

@@ -292,7 +292,7 @@ class OpenTelemetryV2(CustomLogger):
     #  closed here. See ``log_pre_api_call``.
     # ====================================================================== #
 
-    def log_pre_api_call(self, model, messages, kwargs):
+    def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
         """Open the LLM-call span at the call boundary.
 
         Runs synchronously inside the request task, before the upstream call —
@@ -354,7 +354,13 @@ class OpenTelemetryV2(CustomLogger):
             _LLMCallSpan(span=span, start_time_ns=start_time_ns, provider=route.provider),
         )
 
-    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+    async def async_log_success_event(
+        self,
+        kwargs: Mapping[str, object],
+        response_obj: object,
+        start_time: datetime | float | None,
+        end_time: datetime | float | None,
+    ) -> None:
         if self._emit_mcp_tool_call(kwargs, start_time, end_time):
             return
         if self._emit_mcp_list_tools(kwargs, start_time, end_time):
@@ -396,7 +402,13 @@ class OpenTelemetryV2(CustomLogger):
         except Exception as exc:
             verbose_logger.debug("OpenTelemetryV2: metric recording failed: %s", exc)
 
-    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+    async def async_log_failure_event(
+        self,
+        kwargs: Mapping[str, object],
+        response_obj: object,
+        start_time: datetime | float | None,
+        end_time: datetime | float | None,
+    ) -> None:
         if self._emit_mcp_tool_call(kwargs, start_time, end_time):
             return
         if self._emit_mcp_list_tools(kwargs, start_time, end_time):

litellm/integrations/otel/mount.py

@@ -12,6 +12,8 @@ when the feature gate is off.
 """
 
 import os
+from collections.abc import Mapping
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Protocol
 
 from litellm._logging import verbose_logger
@@ -19,6 +21,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled
 
 if TYPE_CHECKING:
     from fastapi import FastAPI
+    from opentelemetry.trace import Span, TracerProvider
 
 # Routes excluded from server-span tracing by default: high-frequency pollers and
 # static UI/docs assets, none of which are LLM traffic. Entries are substring-matched
@@ -76,7 +79,7 @@ class _RenameableSpan(Protocol):
     def set_attribute(self, key: str, value: str) -> None: ...
 
 
-def _passthrough_span_name_hook(span: "_RenameableSpan | None", scope: dict) -> None:
+def _passthrough_span_name_hook(span: "_RenameableSpan | None", scope: Mapping[str, object]) -> None:
     """FastAPI ``server_request_hook``: give passthrough server spans a useful name.
 
     The instrumentation matches the route at span creation, so both the span name
@@ -89,8 +92,8 @@ def _passthrough_span_name_hook(span: "_RenameableSpan | None", scope: dict) ->
     try:
         if span is None or not span.is_recording():
             return
-        path: Final = scope.get("path") or ""
-        method: Final = scope.get("method") or ""
+        path: Final = str(scope.get("path") or "")
+        method: Final = str(scope.get("method") or "")
         first_segment: Final = path.lstrip("/").split("/", 1)[0]
         if first_segment in PASSTHROUGH_PREFIXES:
             span.update_name(f"{method} {path}".strip())
@@ -99,7 +102,7 @@ def _passthrough_span_name_hook(span: "_RenameableSpan | None", scope: dict) ->
         pass
 
 
-def instrument_fastapi_app(app: "FastAPI") -> None:
+def instrument_fastapi_app(app: "FastAPI", tracer_provider: "TracerProvider | None" = None) -> None:
     """Attach OTel server-span instrumentation to the proxy FastAPI app.
 
     Safe no-op when the V2 gate is off or ``opentelemetry-instrumentation-fastapi``
@@ -111,6 +114,14 @@ def instrument_fastapi_app(app: "FastAPI") -> None:
     ``ProxyTracerProvider``; the proxy publishes the real provider as the global
     after config load (see ``proxy_startup_event``), and the proxy delegates to it.
     That way server spans and gen-ai spans share one provider and the same trace.
+
+    Once instrumented, every sampled HTTP response this app sends carries a
+    ``traceresponse`` header naming its server span, unless the response already
+    relays an upstream's ``traceresponse`` or another response propagator is already
+    registered. Other OTel-instrumented apps in the same process get no header. A relayed
+    ``traceresponse`` from an upstream that restarted the trace becomes a link on the
+    server span. The server span also gets ``litellm.call_id`` from the response's
+    ``x-litellm-call-id`` header. ``tracer_provider`` defaults to the global provider.
     """
     try:
         if not is_otel_v2_enabled():
@@ -122,6 +133,20 @@ def instrument_fastapi_app(app: "FastAPI") -> None:
         # ``proxy_server``'s unconditional ``import`` of this module crash when the
         # package is absent, even with the gate off.
         from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+        from opentelemetry.instrumentation.propagators import (
+            get_global_response_propagator,
+            set_global_response_propagator,  # pyright: ignore[reportUnknownVariableType]  # upstream leaves its parameter unannotated
+        )
+
+        from litellm.integrations.otel.trace_response import (
+            UpstreamFirstTraceResponsePropagator,
+            claim_server_span,
+        )
+
+        def server_request_hook(span: "Span | None", scope: Mapping[str, object]) -> None:
+            if span is not None:
+                claim_server_span(span)
+            _passthrough_span_name_hook(span, scope)
 
         excluded_urls: Final = (
             os.environ.get("OTEL_PYTHON_FASTAPI_EXCLUDED_URLS")
@@ -130,12 +155,47 @@ def instrument_fastapi_app(app: "FastAPI") -> None:
         )
         FastAPIInstrumentor.instrument_app(
             app,
+            tracer_provider=tracer_provider,
             excluded_urls=excluded_urls,
-            server_request_hook=_passthrough_span_name_hook,
+            server_request_hook=server_request_hook,
             # Drop the ASGI "http receive"/"http send" lifecycle sub-spans: they
             # are low-value noise and (for passthrough) carry the catch-all route
             # template in their name, which can't be rewritten from a hook.
             exclude_spans=["receive", "send"],
         )
+        if get_global_response_propagator() is None:
+            set_global_response_propagator(UpstreamFirstTraceResponsePropagator())
     except Exception as e:
         verbose_logger.debug("Skipping OTel V2 FastAPI instrumentation: %s", e)
+
+
+def trace_response_expose_headers() -> tuple[str, ...]:
+    return ("traceresponse",) if is_otel_v2_enabled() else ()
+
+
+def trace_response_headers(span: "Span | None") -> Mapping[str, str]:
+    """Response headers naming ``span`` for a response built outside the FastAPI instrumentation.
+
+    Starlette runs the ``Exception`` handler outside every middleware, so its 500 never reaches the response
+    propagator and the handler has to add these itself.
+    """
+    if span is None or not is_otel_v2_enabled():
+        return MappingProxyType({})
+    try:
+        from opentelemetry import trace
+        from opentelemetry.instrumentation.asgi import (  # pyright: ignore[reportMissingTypeStubs]  # the ASGI instrumentation ships no py.typed marker
+            asgi_getter,
+            asgi_setter,
+        )
+        from opentelemetry.instrumentation.propagators import get_global_response_propagator
+
+        from litellm.integrations.otel.trace_response import claim_server_span
+    except ImportError:
+        return MappingProxyType({})
+    claim_server_span(span)
+    propagator: Final = get_global_response_propagator()  # pyright: ignore[reportUnknownVariableType]  # upstream leaves inject unannotated
+    if propagator is None:
+        return MappingProxyType({})
+    start: Final[dict[str, object]] = {"type": "http.response.start"}  # mutable-ok: the propagator writes into it
+    propagator.inject(start, context=trace.set_span_in_context(span), setter=asgi_setter)  # pyright: ignore[reportUnknownMemberType]  # upstream leaves inject unannotated
+    return MappingProxyType({key: ", ".join(asgi_getter.get(start, key) or ()) for key in asgi_getter.keys(start)})  # pyright: ignore[reportUnknownMemberType]  # the ASGI instrumentation ships no py.typed marker

litellm/integrations/otel/trace_response.py

@@ -0,0 +1,74 @@
+from contextvars import ContextVar
+from typing import Final
+
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.instrumentation.asgi import (  # pyright: ignore[reportMissingTypeStubs]  # the ASGI instrumentation ships no py.typed marker
+    asgi_getter,
+)
+from opentelemetry.instrumentation.propagators import ResponsePropagator
+from opentelemetry.propagators.textmap import CarrierT, Setter, default_setter
+from opentelemetry.trace import Span, format_span_id, format_trace_id
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from litellm.integrations.otel.model.semconv import LiteLLM
+
+_LITELLM_SERVER_SPAN: Final[ContextVar[Span | None]] = ContextVar(
+    "litellm_otel_traceresponse_server_span", default=None
+)
+
+
+def claim_server_span(span: Span) -> None:
+    _LITELLM_SERVER_SPAN.set(span)
+
+
+def _response_header(message: dict[str, object], name: str) -> str | None:  # mutable-ok: asgi_getter takes a dict
+    values: Final = asgi_getter.get(message, name)  # pyright: ignore[reportUnknownMemberType]  # the ASGI instrumentation ships no py.typed marker
+    return values[0] if values else None
+
+
+def _http_response_start(message: object) -> dict[str, object] | None:  # mutable-ok: the propagator writes into it
+    if isinstance(message, dict) and message.get("type") == "http.response.start":  # pyright: ignore[reportUnknownMemberType]  # an ASGI message is an untyped dict
+        return message  # pyright: ignore[reportUnknownVariableType]  # an ASGI message is an untyped dict
+    return None
+
+
+def _link_a_restarted_upstream_trace(span: Span, upstream_traceresponse: str) -> None:
+    upstream: Final = trace.get_current_span(
+        TraceContextTextMapPropagator().extract({"traceparent": upstream_traceresponse})
+    ).get_span_context()
+    if upstream.trace_id != span.get_span_context().trace_id:
+        span.add_link(upstream)
+
+
+class UpstreamFirstTraceResponsePropagator(ResponsePropagator):
+    def inject(
+        self,
+        carrier: CarrierT,
+        context: Context | None = None,
+        setter: Setter[CarrierT] = default_setter,
+    ) -> None:
+        start: Final = _http_response_start(carrier)
+        if start is None:
+            return
+        span: Final = trace.get_current_span(context)
+        if span is not _LITELLM_SERVER_SPAN.get():
+            return
+        call_id: Final = _response_header(start, "x-litellm-call-id")
+        if call_id:
+            span.set_attribute(LiteLLM.CALL_ID, call_id)
+        upstream_traceresponse: Final = _response_header(start, "traceresponse")
+        if upstream_traceresponse is not None:
+            _link_a_restarted_upstream_trace(span, upstream_traceresponse)
+            return
+        own: Final = span.get_span_context()
+        if not own.trace_flags.sampled:
+            return
+        sent: Final = start.get("headers")
+        if isinstance(sent, tuple):
+            start["headers"] = [*sent]
+        setter.set(
+            carrier,
+            "traceresponse",
+            f"00-{format_trace_id(own.trace_id)}-{format_span_id(own.span_id)}-{own.trace_flags:02x}",
+        )

litellm/integrations/prometheus_services.py

@@ -3,7 +3,9 @@
 #    On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers)
 
 
-from typing import Final
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING, Final
 
 import litellm
 from litellm._logging import print_verbose, verbose_logger
@@ -17,6 +19,42 @@ from litellm.types.services import (
 
 FAILED_REQUESTS_LABELS: Final = ["error_class", "function_name"]
 
+if TYPE_CHECKING:
+    from prometheus_client import Counter as PrometheusCounter
+    from prometheus_client import Gauge
+
+
+@dataclass(frozen=True, slots=True)
+class _RedisBreakerCollectors:
+    state: "Gauge"
+    transitions: "PrometheusCounter"
+    failures: "PrometheusCounter"
+
+
+@lru_cache(maxsize=1)
+def _redis_breaker_collectors() -> _RedisBreakerCollectors:
+    from prometheus_client import Counter as PrometheusCounter
+    from prometheus_client import Gauge
+
+    return _RedisBreakerCollectors(
+        state=Gauge(
+            "litellm_redis_circuit_breaker_state",
+            "Number of Redis circuit breakers currently in each state",
+            labelnames=("state",),
+            multiprocess_mode="livesum",
+        ),
+        transitions=PrometheusCounter(
+            "litellm_redis_circuit_breaker_transitions",
+            "Redis circuit breaker state transitions",
+            labelnames=("state",),
+        ),
+        failures=PrometheusCounter(
+            "litellm_redis_circuit_breaker_failures",
+            "Redis health failures counted by the circuit breaker",
+            labelnames=("failure_class",),
+        ),
+    )
+
 
 class PrometheusServicesLogger:
     # Class variables or attributes
@@ -41,6 +79,11 @@ class PrometheusServicesLogger:
             self.Counter = Counter
             self.Gauge = Gauge
             self.REGISTRY = REGISTRY
+            breaker_collectors: Final = _redis_breaker_collectors()
+            self.redis_breaker_state = breaker_collectors.state
+            self.redis_breaker_transitions = breaker_collectors.transitions
+            self.redis_breaker_failures = breaker_collectors.failures
+            self._refresh_redis_breaker_gauge()
 
             verbose_logger.debug("in init prometheus services metrics")
 
@@ -98,6 +141,20 @@ class PrometheusServicesLogger:
             return DEFAULT_METRICS
         return metrics
 
+    def _refresh_redis_breaker_gauge(self) -> None:
+        from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts
+
+        counts: Final = redis_breaker_state_counts()
+        for state in (RedisCircuitBreaker.CLOSED, RedisCircuitBreaker.OPEN, RedisCircuitBreaker.HALF_OPEN):
+            self.redis_breaker_state.labels(state).set(counts.get(state, 0))
+
+    def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
+        self._refresh_redis_breaker_gauge()
+        if transitions_to is not None:
+            self.redis_breaker_transitions.labels(transitions_to).inc()
+        if failure_class is not None:
+            self.redis_breaker_failures.labels(failure_class).inc()
+
     def is_metric_registered(self, metric_name) -> bool:
         # Use _names_to_collectors (O(1)) instead of REGISTRY.collect() (O(n)) to avoid
         # perf regression when a new Router is created per request (e.g. router_settings in DB).

litellm/integrations/s3_v2.py

@@ -10,13 +10,14 @@ import asyncio
 import contextvars
 import logging
 import os
+import random
 import re
 import time
 from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import partial
-from typing import TYPE_CHECKING, Final, Literal, cast
+from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast
 from urllib.parse import quote
 from uuid import uuid4
 
@@ -30,7 +31,9 @@ from litellm.constants import (
     DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY,
     DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
     S3_PARTITION_GRANULARITY_ENV_VAR,
+    S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS,
 )
+from litellm.exceptions import Timeout as LiteLLMTimeout
 from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
 from litellm.integrations.s3 import (
     get_s3_object_download_filename,
@@ -79,8 +82,13 @@ _TERMINAL_ERROR_CODES: Final = frozenset(
     }
 )
 _BODY_CODED_STATUSES: Final = frozenset({400, 403})
-_RETRYABLE_STATUSES: Final = frozenset({403, 500, 503})
+_RETRYABLE_STATUSES: Final = frozenset({403, 429, 500, 503})
 _S3_ERROR_CODE: Final = re.compile(r"<Code>([^<]+)</Code>")
+_TRANSIENT_ERRORS: Final = (httpx.TransportError, LiteLLMTimeout)
+_UPLOAD_ATTEMPTS: Final = 3
+_UPLOAD_ATTEMPT_TIMEOUT: Final = httpx.Timeout(S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS)
+
+PutOutcome: TypeAlias = httpx.Response | Exception
 
 
 @dataclass(frozen=True, slots=True)
@@ -105,6 +113,35 @@ def _s3_key_parent(s3_object_key: str) -> str:
     return s3_object_key.rsplit("/", 1)[0] if "/" in s3_object_key else ""
 
 
+def _backoff_seconds(attempt: int) -> float:
+    return (1 << attempt) * (0.5 + random.random())
+
+
+def _log_retry(level: int, attempt: int, reason: str, wait: float, s3_object_key: str) -> None:
+    verbose_logger.log(
+        level,
+        "S3 upload attempt %s/%s failed with %s, retrying in %.2fs key=%s",
+        attempt + 1,
+        _UPLOAD_ATTEMPTS,
+        reason,
+        wait,
+        s3_object_key,
+    )
+
+
+def _response_or_raise(outcome: PutOutcome) -> httpx.Response:
+    if isinstance(outcome, Exception):
+        raise outcome
+    outcome.raise_for_status()
+    return outcome
+
+
+def _throttled(outcome: PutOutcome) -> bool:
+    if isinstance(outcome, Exception):
+        return True
+    return outcome.status_code in (429, 503) or _s3_error_code(outcome) == "SlowDown"
+
+
 class S3BatchUploadError(Exception):
     def __init__(self, failed: int, total: int) -> None:
         self.failed = failed
@@ -566,7 +603,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
 
             url: Final = self._build_object_url(batch_logging_element.s3_object_key)
 
-            async def signed_put(prepared: _PreparedPut) -> httpx.Response:
+            async def signed_put(prepared: _PreparedPut) -> PutOutcome:
                 credentials: Final = await asyncified_get_credentials(
                     aws_access_key_id=self.s3_aws_access_key_id,
                     aws_secret_access_key=self.s3_aws_secret_access_key,
@@ -582,34 +619,16 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
                     self._sign_put, credentials, url, prepared.json_string, prepared.headers
                 )
                 try:
-                    return await self.async_httpx_client.put(url, data=prepared.json_string, headers=signed_headers)
+                    return await self.async_httpx_client.put(
+                        url, data=prepared.json_string, headers=signed_headers, timeout=_UPLOAD_ATTEMPT_TIMEOUT
+                    )
                 except httpx.HTTPStatusError as error:
                     return error.response
+                except _TRANSIENT_ERRORS as error:
+                    return error
 
-            max_retries: Final = 3
             prepared: Final = self._prepare_put(batch_logging_element)
-            for attempt in range(max_retries):
-                response = await self._recorded_put(partial(signed_put, prepared))
-                if (
-                    response.status_code in _RETRYABLE_STATUSES
-                    and not (self.s3_drop_on_terminal_error and _is_terminal(response))
-                    and attempt < max_retries - 1
-                ):
-                    wait_time = 1 << attempt  # 1s, 2s
-                    verbose_logger.log(
-                        logging.DEBUG if _in_flush.get() else logging.WARNING,
-                        "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
-                        response.status_code,
-                        wait_time,
-                        attempt + 1,
-                        max_retries,
-                        batch_logging_element.s3_object_key,
-                    )
-                    self._flush_retries += 1
-                    await asyncio.sleep(wait_time)
-                    continue
-                response.raise_for_status()
-                break
+            await self._put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
         except Exception as e:
             verbose_logger.exception("Error uploading to s3: %s", e)
             self.handle_callback_failure(callback_name="S3Logger")
@@ -678,9 +697,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
         )
         dropped: Final = len(failed) - len(requeued)
         if dropped:
+            terminal: Final = sum(1 for _, outcome in failed if outcome == "dropped")
             verbose_logger.warning(
-                "s3 logging: %s uploads dropped (terminal or retrying longer than s3_max_retry_age_seconds=%s)",
+                "s3 logging: %s uploads dropped this flush: %s rejected as terminal by S3, "
+                "%s retrying longer than s3_max_retry_age_seconds=%s",
                 dropped,
+                terminal,
+                dropped - terminal,
                 self.s3_max_retry_age_seconds,
             )
         if not requeued:
@@ -728,22 +751,46 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
             return "dropped"
         return "retry"
 
-    async def _recorded_put(self, signed_put: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
+    async def _recorded_put(self, signed_put: Callable[[], Awaitable[PutOutcome]]) -> PutOutcome:
+        outcome: Final = await signed_put()
         limiter: Final = self._upload_limiter
-        adaptive: Final = limiter if isinstance(limiter, AdaptiveConcurrencyLimiter) else None
-        try:
-            response: Final = await signed_put()
-        except Exception:
-            if adaptive is not None:
-                adaptive.record(PutSample(throttled=True))
-            raise
-        if adaptive is not None:
-            adaptive.record(
-                PutSample(
-                    throttled=response.status_code in (429, 503) or _s3_error_code(response) == "SlowDown",
-                )
-            )
-        return response
+        if isinstance(limiter, AdaptiveConcurrencyLimiter):
+            limiter.record(PutSample(throttled=_throttled(outcome)))
+        return outcome
+
+    def _retry_reason(self, outcome: PutOutcome) -> str | None:
+        if isinstance(outcome, Exception):
+            return type(outcome).__name__
+        if outcome.status_code not in _RETRYABLE_STATUSES:
+            return None
+        if self.s3_drop_on_terminal_error and _is_terminal(outcome):
+            return None
+        return str(outcome.status_code)
+
+    async def _put_with_retries(
+        self, signed_put: Callable[[], Awaitable[PutOutcome]], s3_object_key: str
+    ) -> httpx.Response:
+        for attempt in range(_UPLOAD_ATTEMPTS - 1):
+            outcome = await self._recorded_put(signed_put)
+            reason = self._retry_reason(outcome)
+            if reason is None:
+                return _response_or_raise(outcome)
+            wait = _backoff_seconds(attempt)
+            _log_retry(logging.DEBUG if _in_flush.get() else logging.WARNING, attempt, reason, wait, s3_object_key)
+            self._flush_retries += 1
+            await asyncio.sleep(wait)
+        return _response_or_raise(await self._recorded_put(signed_put))
+
+    def _sync_put_with_retries(self, signed_put: Callable[[], PutOutcome], s3_object_key: str) -> httpx.Response:
+        for attempt in range(_UPLOAD_ATTEMPTS - 1):
+            outcome = signed_put()
+            reason = self._retry_reason(outcome)
+            if reason is None:
+                return _response_or_raise(outcome)
+            wait = _backoff_seconds(attempt)
+            _log_retry(logging.WARNING, attempt, reason, wait, s3_object_key)
+            time.sleep(wait)
+        return _response_or_raise(signed_put())
 
     def _batch_file_elements(self, batch: tuple[s3BatchLoggingElement, ...]) -> tuple[s3BatchLoggingElement, ...]:
         now: Final = datetime.now(timezone.utc)
@@ -879,7 +926,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
                 params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None)
             )
 
-            def signed_put(prepared_put: _PreparedPut) -> httpx.Response:
+            def signed_put(prepared_put: _PreparedPut) -> PutOutcome:
                 credentials: Final = self.get_credentials(
                     aws_access_key_id=self.s3_aws_access_key_id,
                     aws_secret_access_key=self.s3_aws_secret_access_key,
@@ -887,29 +934,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
                     aws_region_name=self.s3_region_name,
                 )
                 signed_headers: Final = self._sign_put(credentials, url, prepared_put.json_string, prepared_put.headers)
-                return httpx_client.put(url, data=prepared_put.json_string, headers=signed_headers)
-
-            max_retries: Final = 3
-            for attempt in range(max_retries):
-                response = signed_put(prepared)
-                if (
-                    response.status_code in _RETRYABLE_STATUSES
-                    and not (self.s3_drop_on_terminal_error and _is_terminal(response))
-                    and attempt < max_retries - 1
-                ):
-                    wait_time = 1 << attempt  # 1s, 2s
-                    verbose_logger.warning(
-                        "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
-                        response.status_code,
-                        wait_time,
-                        attempt + 1,
-                        max_retries,
-                        batch_logging_element.s3_object_key,
+                try:
+                    return httpx_client.put(
+                        url, data=prepared_put.json_string, headers=signed_headers, timeout=_UPLOAD_ATTEMPT_TIMEOUT
                     )
-                    time.sleep(wait_time)
-                    continue
-                response.raise_for_status()
-                break
+                except _TRANSIENT_ERRORS as error:
+                    return error
+
+            self._sync_put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
         except Exception as e:
             verbose_logger.exception("Error uploading to s3: %s", e)
             self.handle_callback_failure(callback_name="S3Logger")

litellm/integrations/ternary/export_engine.py

@@ -222,16 +222,24 @@ class TernaryExportEngine(FocusExportEngine):
         if data.is_empty():
             verbose_logger.debug("Ternary export: no usage data for window %s", window)
             return
-        failed: Final[list[str]] = []  # mutable-ok: days whose upload failed, reported once after all days run
-        for day, day_data in _partition_by_day(data):
-            try:
-                await self._deliver_day(day_data, _day_window(day, window.frequency))
-            except Exception as e:  # noqa: BLE001  # one day's failure must not block the rest; logged and re-raised
-                verbose_logger.warning("Ternary export: upload for %s failed: %s", day, e)
-                failed.append(day)
+        failed: Final = tuple(
+            [
+                day
+                for day, day_data in _partition_by_day(data)
+                if not await TernaryExportEngine._deliver_day_or_warn(self, day, day_data, window.frequency)
+            ]
+        )
         if failed:
             raise RuntimeError(f"Ternary export: {len(failed)} day(s) failed to upload: {', '.join(failed)}")
 
+    async def _deliver_day_or_warn(self, day: str, day_data: _Frame, frequency: str) -> bool:
+        try:
+            await self._deliver_day(day_data, _day_window(day, frequency))
+        except Exception as e:  # noqa: BLE001  # one day's failure must not block the rest; logged and re-raised
+            verbose_logger.warning("Ternary export: upload for %s failed: %s", day, e)
+            return False
+        return True
+
     async def _deliver_day(self, data: _Frame, window: FocusTimeWindow) -> None:
         payload: Final = _day_payload(data, transformer=self._transformer, serializer=self._serializer)
         if not payload:

litellm/integrations/ternary/ternary_logger.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import os
 from collections.abc import Awaitable, Callable
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Final, Protocol
+from typing import TYPE_CHECKING, ClassVar, Final, Protocol
 
 import litellm
 from litellm._logging import verbose_logger
@@ -51,6 +51,8 @@ class TernaryLogger(FocusLogger):
     TERNARY_EXPORT_FREQUENCY ("daily" default / "interval"), TERNARY_EXPORT_INTERVAL_SECONDS.
     """
 
+    export_job_name: ClassVar[str] = TERNARY_USAGE_DATA_JOB_NAME
+
     def __init__(
         self,
         *,
@@ -80,13 +82,15 @@ class TernaryLogger(FocusLogger):
                 "TERNARY_EXPORT_INTERVAL_SECONDS must be a positive integer when TERNARY_EXPORT_FREQUENCY is 'interval'"
             )
 
-        destination_config: Final[dict[str, str]] = {}  # mutable-ok: built from the config values present below
-        if resolved_api_key:
-            destination_config["api_key"] = resolved_api_key
-        if resolved_connection_id:
-            destination_config["connection_id"] = resolved_connection_id
-        if resolved_base_url:
-            destination_config["base_url"] = resolved_base_url
+        destination_config: Final = {
+            key: value
+            for key, value in (
+                ("api_key", resolved_api_key),
+                ("connection_id", resolved_connection_id),
+                ("base_url", resolved_base_url),
+            )
+            if value
+        }
 
         super().__init__(
             provider="ternary",
@@ -123,25 +127,6 @@ class TernaryLogger(FocusLogger):
             )
         return self._engine
 
-    async def initialize_focus_export_job(self) -> None:
-        """Run the scheduled export under a Ternary-specific pod lock (distinct from FocusLogger's)."""
-        from litellm.proxy.proxy_server import proxy_logging_obj
-
-        writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) if proxy_logging_obj else None
-        pod_lock_manager: Final = getattr(writer, "pod_lock_manager", None) if writer is not None else None
-
-        if pod_lock_manager and pod_lock_manager.redis_cache:
-            acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME)
-            if not acquired:
-                verbose_logger.debug("Ternary export: unable to acquire pod lock")
-                return
-            try:
-                await self._run_scheduled_export()
-            finally:
-                await pod_lock_manager.release_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME)
-        else:
-            await self._run_scheduled_export()
-
     @staticmethod
     async def init_ternary_background_job(
         scheduler: _JobScheduler,

litellm/litellm_core_utils/env_utils.py

@@ -4,23 +4,65 @@ Utility helpers for reading and parsing environment variables.
 
 import logging
 import os
-from typing import Final
+from typing import Final, TypeVar
+
+_NumberT: Final = TypeVar("_NumberT", int, float)
+
+# Marked before logging: the LiteLLM loggers' own filters read env knobs through these helpers, so
+# warning first would re-enter the helper from inside the log call and recurse until RecursionError.
+_reported_unparsable_values: Final[set[tuple[str, str]]] = set()  # mutable-ok: process-wide warn-once registry
+
+
+def _warn_unparsable_once(env_var: str, raw: str, number_type: type[int] | type[float]) -> None:
+    key: Final = (env_var, raw)
+    if key in _reported_unparsable_values:
+        return
+    _reported_unparsable_values.add(key)
+    logging.getLogger("LiteLLM").warning(
+        "%s=%r is not a valid %s. Treating it as unset.", env_var, raw, number_type.__name__
+    )
+
+
+def _get_env_number_or_none(env_var: str, number_type: type[_NumberT]) -> _NumberT | None:
+    raw: Final = os.getenv(env_var)
+    if raw is None or not raw.strip():
+        return None
+    try:
+        return number_type(raw.strip())
+    except ValueError:
+        _warn_unparsable_once(env_var, raw, number_type)
+        return None
+
+
+def get_env_int_or_none(env_var: str) -> int | None:
+    """Parse an environment variable as an integer, returning None when it is unset, blank or unparsable.
+
+    Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
+    from "left at the default", for example when an override should take precedence over a
+    value resolved from somewhere else.
+    """
+    return _get_env_number_or_none(env_var, int)
+
+
+def get_env_float_or_none(env_var: str) -> float | None:
+    """Float counterpart of `get_env_int_or_none`."""
+    return _get_env_number_or_none(env_var, float)
 
 
 def get_env_int(env_var: str, default: int) -> int:
-    """Parse an environment variable as an integer, falling back to default on invalid values.
+    """Parse an environment variable as an integer, falling back to default when it is unset, blank or unparsable.
 
-    Handles empty strings, whitespace, and non-numeric values gracefully
-    so that misconfiguration doesn't crash the process at import time.
+    An unparsable value is warned about once, so misconfiguration neither crashes the process at
+    import time nor goes unnoticed.
     """
-    raw = os.getenv(env_var)
-    if raw is None:
-        return default
-    raw = raw.strip()
-    try:
-        return int(raw)
-    except (ValueError, TypeError):
-        return default
+    value: Final = get_env_int_or_none(env_var)
+    return default if value is None else value
+
+
+def get_env_float(env_var: str, default: float) -> float:
+    """Float counterpart of `get_env_int`."""
+    value: Final = get_env_float_or_none(env_var)
+    return default if value is None else value
 
 
 def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) -> int:
@@ -41,19 +83,3 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int)
         default,
     )
     return default
-
-
-def get_env_int_or_none(env_var: str) -> int | None:
-    """Parse an environment variable as an integer, returning None when it is unset or unusable.
-
-    Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
-    from "left at the default", for example when an override should take precedence over a
-    value resolved from somewhere else.
-    """
-    raw: Final = os.getenv(env_var)
-    if raw is None:
-        return None
-    try:
-        return int(raw.strip())
-    except (ValueError, TypeError):
-        return None

litellm/litellm_core_utils/exception_mapping_utils.py

@@ -113,6 +113,9 @@ class ExceptionCheckers:
         if "maximum input length is" in _error_str_lowercase and "tokens" in _error_str_lowercase:
             return True
 
+        if "input too long:" in _error_str_lowercase and "input tokens, limit is" in _error_str_lowercase:
+            return True
+
         return False
 
     @staticmethod

litellm/litellm_core_utils/get_llm_provider_logic.py

@@ -1,6 +1,8 @@
 from typing import Final, cast
 from urllib.parse import urlparse
 
+from pydantic import ValidationError
+
 import litellm
 from litellm.constants import (
     NADIR_DEFAULT_API_BASE,
@@ -16,6 +18,12 @@ from litellm.secret_managers.main import get_secret, get_secret_str
 from ..types.router import GenericLiteLLMParams, LiteLLM_Params
 
 
+def invalid_provider_routing_hints_message(error: ValidationError) -> str:
+    return "; ".join(
+        f"Invalid {'.'.join(str(part) for part in detail['loc'])}: {detail['msg']}" for detail in error.errors()
+    )
+
+
 def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
     """
     Match a registered openai-compatible endpoint against a caller-supplied

litellm/litellm_core_utils/get_model_cost_map.py

@@ -21,6 +21,7 @@ from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass, replace
 from datetime import datetime, timezone
 from importlib.resources import files
+from itertools import chain
 from pathlib import Path
 from types import MappingProxyType
 from typing import Final, Protocol
@@ -93,16 +94,23 @@ class GetModelCostMap:
         return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map
 
     _loaded_catalog: Mapping[str, Mapping[str, object]] = MappingProxyType({})
+    _loaded_catalog_fields: frozenset[str] = frozenset()
 
     @classmethod
     def loaded_model_cost_map(cls) -> Mapping[str, Mapping[str, object]]:
         """The catalog as last loaded (bundled or remote), untouched by ``register_model`` or router registrations."""
         return cls._loaded_catalog
 
+    @classmethod
+    def loaded_model_cost_map_fields(cls) -> frozenset[str]:
+        """Every field name that at least one entry of the loaded catalog carries."""
+        return cls._loaded_catalog_fields
+
     @classmethod
     def _snapshot_loaded_catalog(cls, model_cost: Mapping[str, object]) -> None:
         raw: Final = _CATALOG_ADAPTER.validate_python(model_cost)
         cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()})
+        cls._loaded_catalog_fields = frozenset(chain.from_iterable(entry.keys() for entry in raw.values()))
 
     @classmethod
     def get_backup_model_count(cls) -> int:

litellm/litellm_core_utils/get_provider_specific_headers.py

@@ -1,4 +1,4 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from typing import Final
 
 from litellm.types.utils import ProviderSpecificHeader
@@ -9,7 +9,7 @@ class ProviderSpecificHeaderUtils:
     def get_provider_specific_headers(
         provider_specific_header: ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None,
         custom_llm_provider: str | None,
-    ) -> dict:
+    ) -> Mapping[str, str]:
         """
         Get the provider specific headers for the given custom llm provider.
 

litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py

@@ -6,7 +6,7 @@ import traceback
 from collections.abc import Mapping, Sequence
 from typing import Final, Literal, cast
 
-from pydantic import ConfigDict, TypeAdapter, with_config
+from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -18,6 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
 )
 from litellm.types.llms.databricks import DatabricksTool
 from litellm.types.llms.openai import (
+    ChatCompletionReasoningItem,
     ChatCompletionThinkingBlock,
     ImageURLListItem,
     OpenAIModerationResponse,
@@ -156,6 +157,21 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
         del choice.enhancements
 
 
+def _first_replay_delta(
+    message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None
+) -> Delta:
+    return Delta(
+        content=message.get("content"),
+        role=message["role"],
+        function_call=message.get("function_call"),
+        tool_calls=message.get("tool_calls"),
+        reasoning_items=None if reasoning_items is None else list(reasoning_items),
+        provider_specific_fields=_cached_provider_specific_fields(
+            _CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
+        ),
+    )
+
+
 def _invalid_choices_message(response_object: Mapping[str, object]) -> str:
     raw_keys: Final = list(response_object.keys())
     if "choices" not in response_object:
@@ -217,12 +233,7 @@ async def convert_to_streaming_response_async(
                     t["index"] = index
                 pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t))
             choice["message"]["tool_calls"] = pydantic_tool_calls
-        delta = Delta(
-            content=choice["message"].get("content", None),
-            role=choice["message"]["role"],
-            function_call=choice["message"].get("function_call", None),
-            tool_calls=choice["message"].get("tool_calls", None),
-        )
+        delta = _first_replay_delta(choice["message"], choice["message"].get("reasoning_items"))
         finish_reason = choice.get("finish_reason", None)
 
         if finish_reason is None:
@@ -243,6 +254,9 @@ async def convert_to_streaming_response_async(
                 completion_tokens=response_object["usage"].get("completion_tokens", 0),
                 prompt_tokens=response_object["usage"].get("prompt_tokens", 0),
                 total_tokens=response_object["usage"].get("total_tokens", 0),
+                completion_tokens_details=_cached_completion_tokens_details(
+                    _CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
+                ),
             ),
         )
 
@@ -334,6 +348,9 @@ def convert_to_streaming_response(
         model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0)
         model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0)
         model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0)
+        model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details(
+            _CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
+        )
 
     if "id" in response_object:
         model_response_object.id = response_object["id"]
@@ -390,6 +407,35 @@ class _ParallelToolUseArgs(TypedDict):
 _PARALLEL_TOOL_USE_ARGS: Final = TypeAdapter(_ParallelToolUseArgs)
 
 
+class _CachedMessageProviderFields(TypedDict, total=False):
+    provider_specific_fields: ReadOnly[object]
+
+
+class _CachedUsageDetails(TypedDict, total=False):
+    completion_tokens_details: ReadOnly[object]
+
+
+_CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields)
+_CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails)
+_FIELD_MAP: Final = TypeAdapter(dict[str, object])
+_COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper)
+
+
+def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None:
+    fields: Final = message.get("provider_specific_fields")
+    return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else None
+
+
+def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionTokensDetailsWrapper | None:
+    details: Final = usage.get("completion_tokens_details")
+    if not isinstance(details, dict):
+        return None
+    try:
+        return _COMPLETION_TOKENS_DETAILS.validate_python(details)
+    except ValidationError:
+        return None
+
+
 def handle_invalid_parallel_tool_calls(
     tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall],
 ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None:
@@ -716,6 +762,7 @@ def convert_to_model_response_object(
                         provider_specific_fields=provider_specific_fields,
                         reasoning_content=reasoning_content,
                         thinking_blocks=thinking_blocks,
+                        reasoning_items=choice["message"].get("reasoning_items", None),
                         annotations=choice["message"].get("annotations", None),
                         images=_normalize_images_for_message(choice["message"].get("images", None)),
                     )

litellm/litellm_core_utils/llm_response_utils/response_metadata.py

@@ -115,8 +115,19 @@ class ResponseMetadata:
         """Check if response type supports timing metrics"""
         return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse))
 
-    def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None:
-        """Set hidden parameters on the response"""
+    def set_hidden_params(
+        self,
+        logging_obj: LiteLLMLoggingObject,
+        model: str | None,
+        kwargs: dict,
+        cache_hit: bool | None = None,
+    ) -> None:
+        """Set hidden parameters on the response.
+
+        ``cache_hit`` is passed for a response served from the cache: ``logging_obj`` only
+        learns of the hit later, in the success handler, so its cost calculator would
+        otherwise price the cached response as a fresh call.
+        """
 
         ## ADD OTHER HIDDEN PARAMS
         model_info: Final = kwargs.get("model_info", {}) or {}
@@ -126,7 +137,10 @@ class ResponseMetadata:
             "api_base": get_api_base(model=model or "", optional_params=kwargs),
             "model_id": model_id,
             "response_cost": logging_obj.response_cost_calculator(
-                result=self.result, litellm_model_name=model, router_model_id=model_id
+                result=self.result,
+                litellm_model_name=model,
+                router_model_id=model_id,
+                cache_hit=cache_hit,
             ),
             "additional_headers": process_response_headers(
                 self._get_additional_headers_from_hidden_params() or {},
@@ -228,6 +242,7 @@ def update_response_metadata(
     start_time: datetime.datetime,
     end_time: datetime.datetime,
     include_overhead: bool = True,
+    cache_hit: bool | None = None,
 ) -> None:
     """
     Updates response metadata including hidden params and timing metrics
@@ -252,5 +267,5 @@ def update_response_metadata(
 
     metadata: Final = ResponseMetadata(result)
     metadata.set_timing_metrics(start_time, end_time, logging_obj, include_overhead)
-    metadata.set_hidden_params(logging_obj, model, kwargs)
+    metadata.set_hidden_params(logging_obj, model, kwargs, cache_hit=cache_hit)
     metadata.apply()

litellm/litellm_core_utils/logging_utils.py

@@ -5,7 +5,7 @@ import re
 import time
 from collections.abc import Iterator, Mapping, Sequence
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast
 
 from litellm._logging import format_base64_size, verbose_logger
 from litellm.constants import (
@@ -28,7 +28,7 @@ if TYPE_CHECKING:
     )
 
     LiteLLMModelResponse = _ModelResponse
-    Span = _Span | Any
+    Span: TypeAlias = _Span
 else:
     LiteLLMModelResponse = Any
     LiteLLMLoggingObject = Any

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -2,6 +2,8 @@
 Common utility functions used for translating messages across providers
 """
 
+import base64
+import binascii
 import io
 import json
 import mimetypes
@@ -1805,18 +1807,62 @@ def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT:
     )
 
 
-LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"})
+LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS: Final = frozenset(
+    {"thinking_blocks", "reasoning_items", "provider_specific_fields"}
+)
+LITELLM_INTERNAL_MESSAGE_FIELDS: Final = LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS | {"reasoning_content"}
+
+
+def _tool_call_without_provider_specific_fields(tool_call: object) -> object:
+    if not isinstance(tool_call, dict):
+        return tool_call
+    fields: Final = cast(dict[str, object], tool_call)  # cast-ok: untyped JSON tool call, keys forwarded verbatim
+    function_fields: Final = (
+        cast("dict[str, object]", fields.get("function"))  # cast-ok: untyped JSON function chunk
+        if isinstance(fields.get("function"), dict)
+        else None
+    )
+    if "provider_specific_fields" not in fields and (
+        function_fields is None or "provider_specific_fields" not in function_fields
+    ):
+        return cast("dict[str, object]", tool_call)  # cast-ok: untyped JSON tool call, no strip needed
+    return {
+        key: (
+            {fkey: fvalue for fkey, fvalue in function_fields.items() if fkey != "provider_specific_fields"}
+            if key == "function" and function_fields is not None
+            else value
+        )
+        for key, value in fields.items()
+        if key != "provider_specific_fields"
+    }
+
+
+def _message_without_fields(message: AllMessageValues, dropped: frozenset[str]) -> AllMessageValues:
+    kept: Final = {key: value for key, value in message.items() if key not in dropped}
+    tool_calls: Final = message.get("tool_calls")
+    if not isinstance(tool_calls, list):
+        return cast(AllMessageValues, kept)  # cast-ok: same TypedDict minus dropped keys
+    return cast(  # cast-ok: same TypedDict minus dropped keys, tool calls minus their provider_specific_fields
+        AllMessageValues,
+        {
+            **kept,
+            "tool_calls": [
+                _tool_call_without_provider_specific_fields(tool_call)
+                for tool_call in cast(list[object], tool_calls)  # cast-ok: untyped JSON tool calls
+            ],
+        },
+    )
+
+
+def strip_litellm_response_only_fields(message: AllMessageValues) -> AllMessageValues:
+    """Drop the fields litellm attaches to a returned message (e.g. Anthropic thinking blocks or Responses
+    API items) that no OpenAI-compatible endpoint accepts back, so a reply can be sent to a strict one."""
+    return _message_without_fields(message, LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS)
 
 
 def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessageValues:
-    """Drop the fields litellm attaches to assistant messages (e.g. when translating Anthropic thinking
-    blocks) that OpenAI-compatible endpoints with strict schemas reject as extra inputs."""
-    if LITELLM_INTERNAL_MESSAGE_FIELDS.isdisjoint(message):
-        return message
-    return cast(  # cast-ok: same TypedDict minus internal keys
-        AllMessageValues,
-        {key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS},
-    )
+    """``strip_litellm_response_only_fields`` plus ``reasoning_content``, for endpoints that reject it too."""
+    return _message_without_fields(message, LITELLM_INTERNAL_MESSAGE_FIELDS)
 
 
 def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any:
@@ -2132,6 +2178,32 @@ def is_unsignable_thinking_block(block: object) -> bool:
     return not (isinstance(thinking_text, str) and len(thinking_text.strip()) > 0)
 
 
+def native_redacted_thinking_data(data: str) -> str:
+    """Claude's own `data` string for a redacted_thinking block.
+
+    Bedrock Converse carries it base64-wrapped as `redactedContent`, and LiteLLM passed that
+    wrapper through as `data` until it started decoding Converse replies, so histories saved
+    before then still hold it. Claude's own value never decodes twice, since its first byte
+    decodes to a protobuf tag outside the base64 alphabet.
+    """
+    return _converse_wrapped_payload(data) or data
+
+
+def converse_redacted_content(data: str) -> str:
+    """The `redactedContent` Bedrock Converse expects for a redacted_thinking block's `data`."""
+    return base64.b64encode(native_redacted_thinking_data(data).encode()).decode()
+
+
+def _converse_wrapped_payload(data: str) -> str | None:
+    try:
+        payload: Final = base64.b64decode(data, validate=True).decode("ascii")
+        _ = base64.b64decode(payload, validate=True)
+    except (binascii.Error, ValueError):
+        return None
+    else:
+        return payload
+
+
 def strip_encrypted_reasoning_from_messages(
     messages: object,
     *,

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -6,13 +6,13 @@ import json
 import mimetypes
 import re
 import xml.etree.ElementTree as ET
-from collections.abc import Container, Iterator, Mapping, Sequence
+from collections.abc import Container, Iterable, Iterator, Mapping, Sequence
 from enum import Enum
 from types import MappingProxyType
 from typing import Any, Final, TypeAlias, TypedDict, cast, overload
 
 from jinja2.sandbox import ImmutableSandboxedEnvironment
-from pydantic import BaseModel
+from pydantic import BaseModel, JsonValue
 
 import litellm
 import litellm.types
@@ -20,7 +20,13 @@ import litellm.types.llms
 from litellm import verbose_logger
 from litellm._uuid import uuid
 from litellm.constants import REDACTED_BY_LITELLM
-from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import anthropic_system_messages
+from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
+    anthropic_system_messages,
+    find_anthropic_server_tool_result,
+    is_anthropic_tool_search_use,
+    message_field,
+    parts_of,
+)
 from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
 from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
 from litellm.types.files import get_file_extension_from_mime_type
@@ -50,10 +56,13 @@ from litellm.types.llms.vertex_ai import PartType as VertexPartType
 from litellm.types.utils import GenericImageParsingChunk
 
 from .common_utils import (
+    converse_redacted_content,
     convert_content_list_to_str,
     infer_content_type_from_url_and_content,
+    is_encrypted_reasoning_block,
     is_non_content_values_set,
     is_unsignable_thinking_block,
+    native_redacted_thinking_data,
     parse_tool_call_arguments,
     salvage_concatenated_tool_arguments,
 )
@@ -1431,6 +1440,45 @@ def convert_to_gemini_tool_call_invoke(
         raise Exception(f"Unable to convert openai tool calls={message} to gemini tool calls. Received error={e}")
 
 
+def _json_children(node: JsonValue) -> Iterable[JsonValue]:
+    if isinstance(node, dict):
+        return node.values()
+    if isinstance(node, list):
+        return (child for child in node if not isinstance(child, list))
+    return ()
+
+
+def _next_json_level(level: tuple[JsonValue, ...], _: None) -> tuple[JsonValue, ...]:
+    return tuple(itertools.chain.from_iterable(_json_children(node) for node in level))
+
+
+def _is_part_reference(node: JsonValue) -> bool:
+    if not isinstance(node, dict):
+        return False
+    ref: Final[JsonValue] = node.get("$ref")
+    return isinstance(ref, str) and ref != ""
+
+
+_REF_CHAR_ESCAPE: Final = re.compile(r"\\u00(?:24|72|65|66)")
+
+
+def _may_spell_ref_key(raw_json: str) -> bool:
+    """JSON text can only spell the key $ref literally or by \\u-escaping one of its four characters"""
+    return '"$ref"' in raw_json or _REF_CHAR_ESCAPE.search(raw_json) is not None
+
+
+def _contains_part_reference(payload: JsonValue) -> bool:
+    """Gemini reads a {"$ref": "<non-empty string>"} object in a function_response as a pointer to one of its parts,
+    unless it sits under a list nested directly in a list
+
+    Google does not document this. It was observed on the Gemini API and Vertex AI on 2026-10-15, see issue #38223
+    """
+    levels: Final = itertools.takewhile(
+        bool, itertools.accumulate(itertools.repeat(None), _next_json_level, initial=(payload,))
+    )
+    return any(_is_part_reference(node) for node in itertools.chain.from_iterable(levels))
+
+
 def convert_to_gemini_tool_call_result(
     message: ChatCompletionToolMessage | ChatCompletionFunctionMessage,
     last_message_with_tool_calls: dict | None,
@@ -1575,8 +1623,8 @@ def convert_to_gemini_tool_call_result(
     try:
         if content_str.strip().startswith("{") or content_str.strip().startswith("["):
             # Try to parse as JSON (for Computer Use structured responses)
-            parsed: Final = json.loads(content_str)
-            if isinstance(parsed, dict):
+            parsed: Final[JsonValue] = json.loads(content_str)
+            if isinstance(parsed, dict) and not (_may_spell_ref_key(content_str) and _contains_part_reference(parsed)):
                 response_data = parsed  # Use the parsed JSON directly
             else:
                 response_data = {"content": content_str}
@@ -1815,23 +1863,60 @@ def convert_function_to_anthropic_tool_invoke(
         raise e
 
 
-ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX: Final = "srvtoolu_"
-
-
-def find_anthropic_server_tool_result(
-    tool_id: str,
-    web_search_results: Sequence[object] | None,
-    tool_results: Sequence[object] | None,
-) -> dict[str, object] | None:
-    if not tool_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX):
-        return None
-    candidates: Final = (*(web_search_results or ()), *(tool_results or ()))
-    return next(
-        (result for result in candidates if isinstance(result, dict) and result.get("tool_use_id") == tool_id),
-        None,
+def _tool_call_ids_in_tool_messages(messages: Sequence[object]) -> frozenset[str]:
+    return frozenset(
+        tool_call_id for message in messages if isinstance(tool_call_id := message_field(message, "tool_call_id"), str)
     )
 
 
+def _messages_until_the_next_assistant(messages: Sequence[object], index: int) -> tuple[object, ...]:
+    return tuple(
+        itertools.takewhile(
+            lambda message: message_field(message, "role") != "assistant",
+            (messages[following_index] for following_index in range(index + 1, len(messages))),
+        )
+    )
+
+
+def _tool_call_ids_with_a_kept_server_result(message: object) -> frozenset[str]:
+    from litellm.llms.anthropic.common_utils import tool_call_is_rebuilt_as_server_tool_use
+
+    provider_specific_fields: Final = message_field(message, "provider_specific_fields")
+    return frozenset(
+        tool_call_id
+        for tool_call in parts_of(message_field(message, "tool_calls"))
+        if isinstance(tool_call_id := message_field(tool_call, "id"), str)
+        and tool_call_is_rebuilt_as_server_tool_use(tool_call_id, provider_specific_fields)
+    )
+
+
+def _server_tool_answers_replaced_after(messages: Sequence[object], index: int) -> frozenset[str]:
+    kept_server_result_ids: Final = _tool_call_ids_with_a_kept_server_result(messages[index])
+    if not kept_server_result_ids:
+        return frozenset()
+    following: Final = _messages_until_the_next_assistant(messages, index)
+    a_user_turn_remains: Final = any(
+        message_field(message, "role") in ("user", "tool", "function")
+        and message_field(message, "tool_call_id") not in kept_server_result_ids
+        for message in following
+    )
+    return _tool_call_ids_in_tool_messages(following) & kept_server_result_ids if a_user_turn_remains else frozenset()
+
+
+def server_tool_answers_replaced_by_their_result(messages: Sequence[object]) -> frozenset[str]:
+    return frozenset(
+        itertools.chain.from_iterable(
+            _server_tool_answers_replaced_after(messages, index)
+            for index, message in enumerate(messages)
+            if message_field(message, "role") == "assistant"
+        )
+    )
+
+
+def tool_call_ids_answered_by_client(messages: Sequence[object]) -> frozenset[str]:
+    return _tool_call_ids_in_tool_messages(messages) - server_tool_answers_replaced_by_their_result(messages)
+
+
 _AnthropicToolInvokeItem: TypeAlias = AnthropicMessagesToolUseParam | dict[str, Any]
 
 
@@ -1839,6 +1924,7 @@ def convert_to_anthropic_tool_invoke(
     tool_calls: list[ChatCompletionAssistantToolCall],
     web_search_results: Sequence[object] | None = None,
     tool_results: Sequence[object] | None = None,
+    answered_tool_ids: Container[str] = frozenset(),
 ) -> list[_AnthropicToolInvokeItem]:
     """
     OpenAI tool invokes:
@@ -1899,7 +1985,11 @@ def convert_to_anthropic_tool_invoke(
             context="Anthropic tool invoke",
         )
 
-        server_tool_result = find_anthropic_server_tool_result(tool_id, web_search_results, tool_results)
+        server_tool_result = (
+            None
+            if tool_id in answered_tool_ids
+            else find_anthropic_server_tool_result(tool_id, web_search_results, tool_results)
+        )
         if server_tool_result is not None:
             anthropic_tool_invoke.append(
                 {
@@ -1910,6 +2000,8 @@ def convert_to_anthropic_tool_invoke(
                 }
             )
             anthropic_tool_invoke.append(server_tool_result)
+        elif is_anthropic_tool_search_use(tool_id, tool_name) and tool_id not in answered_tool_ids:
+            continue
         else:
             sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id)
             _anthropic_tool_use_param = AnthropicMessagesToolUseParam(
@@ -2423,10 +2515,36 @@ def sanitize_messages_for_tool_calling(
     return sanitized_messages
 
 
-def _drop_unsignable_thinking_blocks(
+def _replayable_thinking_blocks(
     thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
 ) -> list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock]:
-    return [block for block in thinking_blocks if not is_unsignable_thinking_block(block)]
+    return [
+        _without_stream_index(_with_native_redacted_data(block))
+        for block in thinking_blocks
+        if not is_unsignable_thinking_block(block)
+    ]
+
+
+def _without_stream_index(
+    block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
+) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock:
+    if "index" not in block:
+        return block
+    return cast(  # cast-ok: dropping one key keeps the block's TypedDict shape
+        ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
+        {key: value for key, value in block.items() if key != "index"},
+    )
+
+
+def _with_native_redacted_data(
+    block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
+) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock:
+    match block:
+        case {"type": "redacted_thinking", "data": str() as data}:
+            native: Final[ChatCompletionRedactedThinkingBlock] = {**block, "data": native_redacted_thinking_data(data)}
+            return native
+        case _:
+            return block
 
 
 _AnthropicMessageList: TypeAlias = list[AllAnthropicPassThroughMessageValues]
@@ -2462,6 +2580,8 @@ def anthropic_messages_pt(
     # (Cases A/B/D in `sanitize_messages_for_tool_calling`) remains gated on
     # `modify_params` because it actually mutates conversation structure.
     messages = [_sanitize_empty_text_content(m) for m in messages]
+    replaced_server_tool_answers: Final = server_tool_answers_replaced_by_their_result(messages)
+    client_answered_tool_ids: Final = _tool_call_ids_in_tool_messages(messages) - replaced_server_tool_answers
 
     # add role=tool support to allow function call result/error submission
     user_message_types: Final = {"user", "tool", "function"}
@@ -2583,7 +2703,9 @@ def anthropic_messages_pt(
 
                     user_content.append(_anthropic_content_text_element)
 
-            elif user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function":
+            elif (
+                user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function"
+            ) and message_field(user_message_types_block, "tool_call_id") not in replaced_server_tool_answers:
                 # OpenAI's tool message content will always be a string
                 user_content.append(
                     convert_to_anthropic_tool_result(user_message_types_block, force_base64=force_base64)
@@ -2612,7 +2734,7 @@ def anthropic_messages_pt(
 
             _raw_thinking_blocks = assistant_content_block.get("thinking_blocks", None)
             thinking_blocks = (
-                _drop_unsignable_thinking_blocks(_raw_thinking_blocks) if _raw_thinking_blocks is not None else None
+                _replayable_thinking_blocks(_raw_thinking_blocks) if _raw_thinking_blocks is not None else None
             )
 
             # Check if tool_calls contain server tool calls (web search, etc.)
@@ -2651,6 +2773,7 @@ def anthropic_messages_pt(
                     assistant_tool_calls,
                     web_search_results=_web_search_results_tc,
                     tool_results=_tool_results_tc,
+                    answered_tool_ids=client_answered_tool_ids,
                 )
 
                 # Group tool invoke results into (server_tool_use, result) pairs
@@ -2834,6 +2957,7 @@ def anthropic_messages_pt(
                     assistant_tool_calls,
                     web_search_results=_web_search_results,
                     tool_results=_tool_results,
+                    answered_tool_ids=client_answered_tool_ids,
                 )
 
                 # Prevent "tool_use ids must be unique" errors by filtering duplicates
@@ -3714,6 +3838,7 @@ class BedrockImageProcessor:
 def _convert_to_bedrock_tool_call_invoke(
     tool_calls: list,
     model: str | None = None,
+    deployment_model: str | None = None,
 ) -> list[BedrockContentBlock]:
     """
     OpenAI tool invokes:
@@ -3795,6 +3920,7 @@ def _convert_to_bedrock_tool_call_invoke(
                                     {"cache_control": tool["cache_control"]},
                                     block_type="content_block",
                                     model=model,
+                                    deployment_model=deployment_model,
                                 )
                                 if _cache_point_block is not None:
                                     _parts_list.append(_cache_point_block)
@@ -3814,6 +3940,7 @@ def _convert_to_bedrock_tool_call_invoke(
                         {"cache_control": tool["cache_control"]},
                         block_type="content_block",
                         model=model,
+                        deployment_model=deployment_model,
                     )
                     if cache_point_block is not None:
                         _parts_list.append(cache_point_block)
@@ -4454,9 +4581,14 @@ class BedrockConverseMessagesProcessor:
         llm_provider: str,
         user_continue_message: ChatCompletionUserMessage | None = None,
         assistant_continue_message: str | ChatCompletionAssistantMessage | None = None,
+        deployment_model: str | None = None,
+        base_model: str | None = None,
     ) -> list[BedrockMessageBlock]:
+        from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking
+
         contents: list[BedrockMessageBlock] = []
         msg_i = 0
+        to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model)
 
         messages = BedrockConverseMessagesProcessor._initial_message_setup(
             messages, model, llm_provider, user_continue_message
@@ -4512,6 +4644,7 @@ class BedrockConverseMessagesProcessor:
                                 message_block=element,
                                 block_type="content_block",
                                 model=model,
+                                deployment_model=deployment_model,
                             )
                             if _cache_point_block is not None:
                                 _parts.append(_cache_point_block)
@@ -4519,7 +4652,7 @@ class BedrockConverseMessagesProcessor:
                 elif message_content and isinstance(message_content, str):
                     _part = BedrockContentBlock(text=message_content)
                     _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block(
-                        message_block, block_type="content_block", model=model
+                        message_block, block_type="content_block", model=model, deployment_model=deployment_model
                     )
                     user_content.append(_part)
                     if _cache_point_block is not None:
@@ -4569,6 +4702,7 @@ class BedrockConverseMessagesProcessor:
                         {"cache_control": tool_msg_cache_control},
                         block_type="content_block",
                         model=model,
+                        deployment_model=deployment_model,
                     )
                     if cache_point_block is not None:
                         tool_content.append(cache_point_block)
@@ -4601,15 +4735,12 @@ class BedrockConverseMessagesProcessor:
                     assistant_continue_message=assistant_continue_message,
                 )
                 _assistant_content = assistant_message_block.get("content", None)
-                thinking_blocks = cast(
-                    list[ChatCompletionThinkingBlock] | None,
-                    assistant_message_block.get("thinking_blocks"),
-                )
+                thinking_blocks = assistant_message_block.get("thinking_blocks")
 
                 if thinking_blocks is not None:
                     converted_thinking_blocks = (
                         BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
-                            thinking_blocks
+                            thinking_blocks, to_claude=to_claude
                         )
                     )
                     assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
@@ -4621,9 +4752,9 @@ class BedrockConverseMessagesProcessor:
                     assistants_parts: list[BedrockContentBlock] = []
                     for element in _assistant_content:
                         if isinstance(element, dict):
-                            if element["type"] == "thinking":
+                            if element["type"] == "thinking" or element["type"] == "redacted_thinking":
                                 thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
-                                    thinking_blocks=[element]
+                                    thinking_blocks=[element], to_claude=to_claude
                                 )
                                 assistants_parts = (
                                     BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
@@ -4648,6 +4779,7 @@ class BedrockConverseMessagesProcessor:
                             message_block=cast(OpenAIMessageContentListBlock, element),
                             block_type="content_block",
                             model=model,
+                            deployment_model=deployment_model,
                         )
                         if _cache_point_block is not None:
                             assistants_parts.append(_cache_point_block)
@@ -4659,14 +4791,21 @@ class BedrockConverseMessagesProcessor:
                     # If content is empty/whitespace, skip it (don't add a placeholder)
                     # Add cache point block for assistant string content
                     _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block(
-                        assistant_message_block, block_type="content_block", model=model
+                        assistant_message_block,
+                        block_type="content_block",
+                        model=model,
+                        deployment_model=deployment_model,
                     )
                     if _cache_point_block is not None:
                         assistant_content.append(_cache_point_block)
 
                 _tool_calls = assistant_message_block.get("tool_calls", [])
                 if _tool_calls:
-                    assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls, model=model))
+                    assistant_content.extend(
+                        _convert_to_bedrock_tool_call_invoke(
+                            _tool_calls, model=model, deployment_model=deployment_model
+                        )
+                    )
 
                 msg_i += 1
 
@@ -4687,23 +4826,49 @@ class BedrockConverseMessagesProcessor:
 
     @staticmethod
     def translate_thinking_blocks_to_reasoning_content_blocks(
-        thinking_blocks: list[ChatCompletionThinkingBlock],
+        thinking_blocks: Sequence[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
+        *,
+        to_claude: bool = False,
     ) -> list[BedrockContentBlock]:
-        reasoning_content_blocks: Final[list[BedrockContentBlock]] = []
-        for thinking_block in thinking_blocks:
-            reasoning_text = thinking_block.get("thinking")
-            reasoning_signature = thinking_block.get("signature")
-            text_block = BedrockConverseReasoningTextBlock(
-                text=reasoning_text or "",
+        return [
+            BedrockContentBlock(
+                reasoningContent=BedrockConverseMessagesProcessor._to_reasoning_content_block(
+                    thinking_block, to_claude=to_claude
+                )
             )
-            if reasoning_signature is not None:
-                text_block["signature"] = reasoning_signature
-            reasoning_content_block = BedrockConverseReasoningContentBlock(
-                reasoningText=text_block,
-            )
-            bedrock_content_block = BedrockContentBlock(reasoningContent=reasoning_content_block)
-            reasoning_content_blocks.append(bedrock_content_block)
-        return reasoning_content_blocks
+            for thinking_block in thinking_blocks
+            if not (to_claude and is_encrypted_reasoning_block(thinking_block))
+        ]
+
+    @staticmethod
+    def _to_reasoning_content_block(
+        thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
+        *,
+        to_claude: bool,
+    ) -> BedrockConverseReasoningContentBlock:
+        match thinking_block:
+            case {"type": "redacted_thinking", "data": str() as data} if to_claude and data.strip():
+                return BedrockConverseReasoningContentBlock(redactedContent=converse_redacted_content(data))
+            case {"signature": signature} if signature is not None:
+                return BedrockConverseReasoningContentBlock(
+                    reasoningText=BedrockConverseReasoningTextBlock(
+                        text=BedrockConverseMessagesProcessor._thinking_text(thinking_block), signature=signature
+                    )
+                )
+            case _:
+                return BedrockConverseReasoningContentBlock(
+                    reasoningText=BedrockConverseReasoningTextBlock(
+                        text=BedrockConverseMessagesProcessor._thinking_text(thinking_block)
+                    )
+                )
+
+    @staticmethod
+    def _thinking_text(thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock) -> str:
+        match thinking_block:
+            case {"thinking": str() as text}:
+                return text
+            case _:
+                return ""
 
     @staticmethod
     def process_file_message(message: ChatCompletionFileObject) -> BedrockContentBlock:
@@ -4821,6 +4986,8 @@ def _bedrock_converse_messages_pt(
     llm_provider: str,
     user_continue_message: ChatCompletionUserMessage | None = None,
     assistant_continue_message: str | ChatCompletionAssistantMessage | None = None,
+    deployment_model: str | None = None,
+    base_model: str | None = None,
 ) -> list[BedrockMessageBlock]:
     """
     Converts given messages from OpenAI format to Bedrock format
@@ -4829,9 +4996,11 @@ def _bedrock_converse_messages_pt(
     - Please ensure that function response turn comes immediately after a function call turn
     - Conversation blocks and tool result blocks cannot be provided in the same turn. Issue: https://github.com/BerriAI/litellm/issues/6053
     """
+    from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking
 
     contents: list[BedrockMessageBlock] = []
     msg_i = 0
+    to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model)
 
     messages = BedrockConverseMessagesProcessor._initial_message_setup(
         messages, model, llm_provider, user_continue_message
@@ -4886,6 +5055,7 @@ def _bedrock_converse_messages_pt(
                             message_block=element,
                             block_type="content_block",
                             model=model,
+                            deployment_model=deployment_model,
                         )
                         if _cache_point_block is not None:
                             _parts.append(_cache_point_block)
@@ -4893,7 +5063,7 @@ def _bedrock_converse_messages_pt(
             elif message_content and isinstance(message_content, str):
                 _part = BedrockContentBlock(text=message_content)
                 _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block(
-                    message_block, block_type="content_block", model=model
+                    message_block, block_type="content_block", model=model, deployment_model=deployment_model
                 )
                 user_content.append(_part)
                 if _cache_point_block is not None:
@@ -4945,6 +5115,7 @@ def _bedrock_converse_messages_pt(
                     {"cache_control": tool_msg_cache_control},
                     block_type="content_block",
                     model=model,
+                    deployment_model=deployment_model,
                 )
                 if cache_point_block is not None:
                     tool_content.append(cache_point_block)
@@ -4977,15 +5148,12 @@ def _bedrock_converse_messages_pt(
                 assistant_continue_message=assistant_continue_message,
             )
             _assistant_content = assistant_message_block.get("content", None)
-            thinking_blocks = cast(
-                list[ChatCompletionThinkingBlock] | None,
-                assistant_message_block.get("thinking_blocks"),
-            )
+            thinking_blocks = assistant_message_block.get("thinking_blocks")
 
             if thinking_blocks is not None:
                 converted_thinking_blocks = (
                     BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
-                        thinking_blocks
+                        thinking_blocks, to_claude=to_claude
                     )
                 )
                 assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
@@ -4997,10 +5165,10 @@ def _bedrock_converse_messages_pt(
                 assistants_parts: list[BedrockContentBlock] = []
                 for element in _assistant_content:
                     if isinstance(element, dict):
-                        if element["type"] == "thinking":
+                        if element["type"] == "thinking" or element["type"] == "redacted_thinking":
                             thinking_block = (
                                 BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
-                                    thinking_blocks=[element]
+                                    thinking_blocks=[element], to_claude=to_claude
                                 )
                             )
                             assistants_parts = (
@@ -5027,6 +5195,7 @@ def _bedrock_converse_messages_pt(
                             message_block=cast(OpenAIMessageContentListBlock, element),
                             block_type="content_block",
                             model=model,
+                            deployment_model=deployment_model,
                         )
                         if _cache_point_block is not None:
                             assistants_parts.append(_cache_point_block)
@@ -5037,13 +5206,15 @@ def _bedrock_converse_messages_pt(
                     assistant_content.append(BedrockContentBlock(text=_assistant_content))
                 # Add cache point block for assistant string content
                 _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block(
-                    assistant_message_block, block_type="content_block", model=model
+                    assistant_message_block, block_type="content_block", model=model, deployment_model=deployment_model
                 )
                 if _cache_point_block is not None:
                     assistant_content.append(_cache_point_block)
             _tool_calls = assistant_message_block.get("tool_calls", [])
             if _tool_calls:
-                assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls, model=model))
+                assistant_content.extend(
+                    _convert_to_bedrock_tool_call_invoke(_tool_calls, model=model, deployment_model=deployment_model)
+                )
 
             msg_i += 1
 
@@ -5090,20 +5261,22 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str:
     return valid_string
 
 
-def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockToolBlock | None:
+def add_cache_point_tool_block(
+    tool: dict, model: str | None = None, deployment_model: str | None = None
+) -> BedrockToolBlock | None:
     from litellm.llms.bedrock.common_utils import (
+        bedrock_cache_point_ttl_supported,
         bedrock_model_accepts_cache_points,
-        supports_bedrock_prompt_cache_ttl,
     )
 
     cache_control: Final = tool.get("cache_control", None)
-    if cache_control is not None and bedrock_model_accepts_cache_points(model):
+    if cache_control is not None and bedrock_model_accepts_cache_points(model, deployment_model):
         cache_point: Final = cache_control.get("type", "ephemeral")
         if cache_point == "ephemeral":
             cache_point_block: Final[CachePointBlock] = {"type": "default"}
             if isinstance(cache_control, dict) and "ttl" in cache_control:
                 ttl: Final = cache_control["ttl"]
-                if ttl in ["5m", "1h"] and model is not None and supports_bedrock_prompt_cache_ttl(model):
+                if ttl in ["5m", "1h"] and bedrock_cache_point_ttl_supported(model, deployment_model):
                     cache_point_block["ttl"] = ttl
             return {"cachePoint": cache_point_block}
     return None
@@ -5133,7 +5306,9 @@ def _is_bedrock_tool_block(tool: dict) -> bool:
     return isinstance(tool, dict) and ("systemTool" in tool or "toolSpec" in tool or "cachePoint" in tool)
 
 
-def _bedrock_tools_pt(tools: list, model: str | None = None) -> list[BedrockToolBlock]:
+def _bedrock_tools_pt(
+    tools: list, model: str | None = None, deployment_model: str | None = None
+) -> list[BedrockToolBlock]:
     """
     OpenAI tools looks like:
     tools = [
@@ -5253,7 +5428,7 @@ def _bedrock_tools_pt(tools: list, model: str | None = None) -> list[BedrockTool
         tool_block_list.append(tool_block)
 
         ## ADD CACHE POINT TOOL BLOCK ##
-        cache_point_tool_block = add_cache_point_tool_block(tool, model=model)
+        cache_point_tool_block = add_cache_point_tool_block(tool, model=model, deployment_model=deployment_model)
         if cache_point_tool_block is not None:
             tool_block_list.append(cache_point_tool_block)
 

litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py

@@ -29,11 +29,17 @@ prompt cache readable across turns. Messages are handled in OpenAI format; the
 Anthropic wire shape is built later by ``anthropic_messages_pt``.
 """
 
-from collections.abc import Iterator, Mapping, Sequence
-from itertools import chain, groupby
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from itertools import chain, dropwhile, groupby, takewhile
+from types import MappingProxyType
 from typing import Final, Literal, TypeAlias, TypeVar
 
-from litellm.types.llms.anthropic import AnthropicMessagesSystemMessageParam, AnthropicSystemMessageContent
+from litellm.types.llms.anthropic import (
+    ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
+    ANTHROPIC_TOOL_SEARCH_TOOL_NAMES,
+    AnthropicMessagesSystemMessageParam,
+    AnthropicSystemMessageContent,
+)
 from litellm.types.llms.openai import (
     AllMessageValues,
     ChatCompletionCachedContent,
@@ -50,6 +56,7 @@ CONVERTED_SYSTEM_NOTE: Final = (
 
 _USER_TYPE_ROLES: Final = frozenset({"user", "tool", "function"})
 _TOOL_ROLES: Final = frozenset({"tool", "function"})
+_NO_IDS: Final[frozenset[str]] = frozenset()
 _RENDERED_PART_TYPES: Final = frozenset({"text", "image_url", "document", "file"})
 _RENDERED_ASSISTANT_PART_TYPES: Final = frozenset({"text", "server_tool_use"})
 _THINKING_BLOCK_TYPES: Final = frozenset({"thinking", "redacted_thinking"})
@@ -283,13 +290,112 @@ def _separate_thinking_blocks_render(message: object, parts: Sequence[object]) -
     return any(_thinking_block_renders(block) for block in parts_of(message_field(message, "thinking_blocks")))
 
 
-def _assistant_renders(message: object) -> bool:
+def is_anthropic_tool_search_use(tool_id: str, tool_name: str) -> bool:
+    return tool_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX) and tool_name in ANTHROPIC_TOOL_SEARCH_TOOL_NAMES
+
+
+def find_anthropic_server_tool_result(
+    tool_id: str,
+    web_search_results: Sequence[object] | None,
+    tool_results: Sequence[object] | None,
+) -> dict[str, object] | None:
+    if not tool_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX):
+        return None
+    candidates: Final = (*(web_search_results or ()), *(tool_results or ()))
+    return next(
+        (result for result in candidates if isinstance(result, dict) and result.get("tool_use_id") == tool_id),
+        None,
+    )
+
+
+def _tool_result_ids(answers: Iterable[object]) -> Iterator[str]:
+    for message in answers:
+        if isinstance(tool_call_id := message_field(message, "tool_call_id"), str) and tool_call_id:
+            yield tool_call_id
+
+
+def answered_tool_call_ids(messages: Sequence[object], assistant_index: int) -> frozenset[str]:
+    from_assistant: Final = (messages[index] for index in range(assistant_index, len(messages)))
+    following: Final = dropwhile(lambda message: message_field(message, "role") == "assistant", from_assistant)
+    answers: Final = takewhile(lambda message: message_field(message, "role") in _TOOL_ROLES, following)
+    return frozenset(_tool_result_ids(answers))
+
+
+def _role_group(message: object) -> str:
+    role: Final = message_field(message, "role")
+    if role == "assistant":
+        return "assistant"
+    return "tool" if role in _TOOL_ROLES else "other"
+
+
+def _answered_by_assistant_index(
+    groups: Sequence[tuple[str, tuple[tuple[int, object], ...]]], group_index: int
+) -> Iterator[tuple[int, frozenset[str]]]:
+    following: Final = groups[group_index + 1] if group_index + 1 < len(groups) else ("other", ())
+    answered: Final = (
+        frozenset(_tool_result_ids(message for _, message in following[1])) if following[0] == "tool" else _NO_IDS
+    )
+    return ((index, answered) for index, _ in groups[group_index][1])
+
+
+def _answered_ids_by_index(messages: Sequence[object]) -> Mapping[int, frozenset[str]]:
+    """``answered_tool_call_ids`` for every assistant message, computed in one pass over the messages."""
+    groups: Final = tuple(
+        (role_group, tuple(group))
+        for role_group, group in groupby(enumerate(messages), key=lambda pair: _role_group(pair[1]))
+    )
+    return MappingProxyType(
+        dict(
+            chain.from_iterable(
+                _answered_by_assistant_index(groups, group_index)
+                for group_index, (role_group, _) in enumerate(groups)
+                if role_group == "assistant"
+            )
+        )
+    )
+
+
+def _unresolved_tool_search_id(tool_call: object, message: object) -> str | None:
+    tool_id: Final = message_field(tool_call, "id")
+    tool_name: Final = message_field(message_field(tool_call, "function"), "name")
+    if not isinstance(tool_id, str) or not isinstance(tool_name, str):
+        return None
+    if not is_anthropic_tool_search_use(tool_id, tool_name):
+        return None
+    provider_fields: Final = message_field(message, "provider_specific_fields")
+    server_result: Final = find_anthropic_server_tool_result(
+        tool_id,
+        parts_of(message_field(provider_fields, "web_search_results")),
+        parts_of(message_field(provider_fields, "tool_results")),
+    )
+    return tool_id if server_result is None else None
+
+
+def _unresolved_tool_search_ids(message: object) -> tuple[str, ...]:
+    return tuple(
+        tool_id
+        for tool_call in parts_of(message_field(message, "tool_calls"))
+        if (tool_id := _unresolved_tool_search_id(tool_call, message)) is not None
+    )
+
+
+def _tool_calls_render(message: object, answered: frozenset[str]) -> bool:
+    tool_calls: Final = parts_of(message_field(message, "tool_calls"))
+    search_ids: Final = _unresolved_tool_search_ids(message)
+    if len(search_ids) < len(tool_calls):
+        return True
+    return bool(search_ids) and not answered.isdisjoint(search_ids)
+
+
+def _assistant_renders(message: object, answered: frozenset[str]) -> bool:
     """Whether ``anthropic_messages_pt`` puts a block on the wire for this assistant message.
 
     String content (the converter pads an empty one with a placeholder), a text part,
     a signed thinking part, a server tool part, tool calls, a function call, a kept
-    thinking block and compaction blocks each render. An assistant message with none
-    of them, such as ``content: None`` or an empty list, vanishes from the wire.
+    thinking block and compaction blocks each render. A tool search use with neither a
+    server result nor a caller's answer is dropped, so it renders nothing. An assistant
+    message with none of them, such as ``content: None`` or an empty list, vanishes
+    from the wire.
     """
     content: Final = message_field(message, "content")
     if isinstance(content, str):
@@ -298,13 +404,13 @@ def _assistant_renders(message: object) -> bool:
     return (
         any(_assistant_part_renders(part) for part in parts)
         or _separate_thinking_blocks_render(message, parts)
-        or bool(message_field(message, "tool_calls"))
+        or _tool_calls_render(message, answered)
         or bool(message_field(message, "function_call"))
         or bool(message_field(message_field(message, "provider_specific_fields"), "compaction_blocks"))
     )
 
 
-def _renders(message: object) -> bool:
+def _renders(message: object, answered: frozenset[str]) -> bool:
     """Whether ``anthropic_messages_pt`` puts a block on the wire for this message.
 
     A tool message always becomes a ``tool_result`` and a user message with string
@@ -316,26 +422,31 @@ def _renders(message: object) -> bool:
     if role in _TOOL_ROLES:
         return True
     if role == "assistant":
-        return _assistant_renders(message)
+        return _assistant_renders(message, answered)
     content: Final = message_field(message, "content")
     return isinstance(content, str) or any(
         message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content)
     )
 
 
+def _rendered_indices(messages: Sequence[object]) -> frozenset[int]:
+    answered: Final = _answered_ids_by_index(messages)
+    return frozenset(index for index, message in enumerate(messages) if _renders(message, answered.get(index, _NO_IDS)))
+
+
 def _rendered_block(
     message_index: int,
-    messages: Sequence[AllMessageValues],
+    rendered: frozenset[int],
     blocks: Sequence[tuple[bool, tuple[int, ...]]],
 ) -> int | None:
     block_index: Final = _block_containing(message_index, blocks)
     _, indices = blocks[block_index]
-    return block_index if any(_renders(messages[index]) for index in indices) else None
+    return block_index if not rendered.isdisjoint(indices) else None
 
 
 def _system_may_follow(
     block_index: int,
-    messages: Sequence[AllMessageValues],
+    rendered: frozenset[int],
     blocks: Sequence[tuple[bool, tuple[int, ...]]],
 ) -> bool:
     """Whether a system message behind this block precedes an assistant turn or ends the array on the wire.
@@ -344,11 +455,7 @@ def _system_may_follow(
     first later block that puts anything on the wire is an assistant block.
     """
     return next(
-        (
-            not is_user
-            for is_user, indices in blocks[block_index + 1 :]
-            if any(_renders(messages[index]) for index in indices)
-        ),
+        (not is_user for is_user, indices in blocks[block_index + 1 :] if not rendered.isdisjoint(indices)),
         True,
     )
 
@@ -356,6 +463,7 @@ def _system_may_follow(
 def _anchor_block(
     run: Sequence[int],
     messages: Sequence[AllMessageValues],
+    rendered: frozenset[int],
     blocks: Sequence[tuple[bool, tuple[int, ...]]],
 ) -> int | None:
     """The user-type block a system run must follow, or ``None`` when it converts in place.
@@ -373,8 +481,8 @@ def _anchor_block(
     neighbour: Final = previous if _is_user_type(messages[previous]) else run[-1] + 1
     if neighbour >= len(messages) or not _is_user_type(messages[neighbour]):
         return None
-    block_index: Final = _rendered_block(neighbour, messages, blocks)
-    if block_index is None or not _system_may_follow(block_index, messages, blocks):
+    block_index: Final = _rendered_block(neighbour, rendered, blocks)
+    if block_index is None or not _system_may_follow(block_index, rendered, blocks):
         return None
     return block_index
 
@@ -389,7 +497,8 @@ def _placed_for_flagged_model(messages: Sequence[AllMessageValues]) -> tuple[All
     move it. Runs that share a user turn merge into one system message.
     """
     blocks: Final = _user_type_blocks(messages)
-    anchors: Final = tuple((run, _anchor_block(run, messages, blocks)) for run in _system_runs(messages))
+    rendered: Final = _rendered_indices(messages)
+    anchors: Final = tuple((run, _anchor_block(run, messages, rendered, blocks)) for run in _system_runs(messages))
 
     def messages_of(run: tuple[int, ...]) -> tuple[AllMessageValues, ...]:
         return tuple(messages[index] for index in run)

litellm/litellm_core_utils/realtime_streaming.py

@@ -10,7 +10,11 @@ from typing_extensions import ReadOnly
 
 import litellm
 from litellm._logging import redact_internal_details_from_client_message, verbose_logger
-from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY
+from litellm.constants import (
+    REALTIME_METERED_USAGE_KEY,
+    REALTIME_SESSION_FAILURE_LOGGED_KEY,
+    REALTIME_SESSION_SUCCESS_LOGGED_KEY,
+)
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
 from litellm.types.llms.openai import (
@@ -114,6 +118,12 @@ class RealtimeEventNormalizer(Protocol):
     def patch_outgoing_session(self, session: dict) -> dict: ...
 
 
+class RealtimeUsageMeter(Protocol):
+    def observe_client_event(self, event: Mapping[str, object]) -> None: ...
+    def observe_backend_event(self, event: Mapping[str, object]) -> None: ...
+    def snapshot(self) -> object: ...
+
+
 DefaultLoggedRealTimeEventTypes: Final = [
     "session.created",
     "response.create",
@@ -136,6 +146,7 @@ class RealTimeStreaming:
         backend_uses_beta_protocol: bool | None = None,
         force_transcription_model: str | None = None,
         event_normalizer: RealtimeEventNormalizer | None = None,
+        usage_meter: RealtimeUsageMeter | None = None,
         logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER,
     ):
         self.websocket: _ClientWebSocket = websocket
@@ -198,6 +209,7 @@ class RealTimeStreaming:
         self._is_transcription_session: bool = force_transcription_model is not None
         # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer).
         self._event_normalizer = event_normalizer
+        self._usage_meter = usage_meter
 
     @property
     def _session_created_sent_to_client(self) -> bool:
@@ -416,6 +428,8 @@ class RealTimeStreaming:
     async def log_messages(self):
         """Log messages in list"""
         if self.logging_obj:
+            if self._usage_meter is not None:
+                self.logging_obj.model_call_details[REALTIME_METERED_USAGE_KEY] = self._usage_meter.snapshot()
             if self.input_messages:
                 self.logging_obj.model_call_details["messages"] = self.input_messages
             if self.session_tools or self.tool_calls:
@@ -481,6 +495,10 @@ class RealTimeStreaming:
                     sent = True
             return sent
         await self.backend_ws.send(message)
+        if self._usage_meter is not None:
+            sent_event: Final = self._parse_backend_event(message)
+            if sent_event is not None:
+                self._usage_meter.observe_client_event(sent_event)
         return True
 
     def _enforce_transcription_session_model(self, message: str) -> str:
@@ -1113,6 +1131,9 @@ class RealTimeStreaming:
                 if self._should_drop_event_from_client(event):
                     continue
 
+                if self._usage_meter is not None:
+                    self._usage_meter.observe_backend_event(event)
+
                 if await self._handle_raw_backend_message(event, raw_response):
                     continue