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

# Conflicts:
#	litellm/caching/redis_cache.py
#	litellm/constants.py
#	litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
#	tests/unit/proxy/management_endpoints/test_organization_endpoints.py
#	tests/unit/proxy/proxy_server/test_spend_counters.py
#	tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py

.circleci/config.yml

@@ -1293,6 +1293,19 @@ jobs:
                   uv venv --python 3.12 .venv-mcp-peer
                   uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
                   echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
+      - when:
+          condition:
+            equal: [sdk, << parameters.suite >>]
+          steps:
+            - run:
+                name: Install redis-server 7.2.16
+                command: |
+                  mkdir -p /tmp/redis-bin
+                  cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
+                  docker cp "${cid}:/usr/local/bin/redis-server" /tmp/redis-bin/redis-server
+                  docker rm "$cid"
+                  /tmp/redis-bin/redis-server --version | grep -q 'v=7.2.16'
+                  echo 'export PATH="/tmp/redis-bin:$PATH"' >> "$BASH_ENV"
       - start_postgres:
           image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
           server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"

.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-linting.yml

@@ -21,7 +21,7 @@ jobs:
       pull-requests: read
       actions: read
     runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 20
     services:
       redis:
         image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
@@ -600,6 +600,9 @@ jobs:
       - name: check_unsafe_enterprise_import
         run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
 
+      - name: check_telemetry_imports
+        run: uv run --no-sync python ./tests/code_coverage_tests/check_telemetry_imports.py
+
       - name: ban_copy_deepcopy_kwargs
         run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
 
@@ -615,6 +618,9 @@ jobs:
       - name: check_migrations_no_data_rewrites
         run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
 
+      - name: check_migrations_lock_timeout
+        run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_lock_timeout.py
+
       - name: check_no_publicly_known_master_key
         run: uv run --no-sync python ./tests/code_coverage_tests/check_no_publicly_known_master_key.py
 

.github/workflows/test-unit.yml

@@ -233,6 +233,7 @@ jobs:
               tests/unit/llms/base_llm
               tests/unit/llms/baseten
               tests/unit/llms/black_forest_labs
+              tests/unit/llms/brave
               tests/unit/llms/bytez
               tests/unit/llms/cerebras
               tests/unit/llms/chat
@@ -268,6 +269,7 @@ jobs:
               tests/unit/llms/gdc
               tests/unit/llms/gemini
               tests/unit/llms/gigachat
+              tests/unit/llms/google_pse
               tests/unit/llms/github_copilot
               tests/unit/llms/gradient_ai
               tests/unit/llms/groq
@@ -300,6 +302,7 @@ jobs:
               tests/unit/llms/nebius
               tests/unit/llms/neosantara
               tests/unit/llms/nimble
+              tests/unit/llms/nlp_cloud
               tests/unit/llms/novita
               tests/unit/llms/nscale
               tests/unit/llms/nvidia_nim
@@ -310,6 +313,7 @@ jobs:
               tests/unit/llms/oobabooga
               tests/unit/llms/openai_like
               tests/unit/llms/opencode
+              tests/unit/llms/opensandbox
               tests/unit/llms/openrouter
               tests/unit/llms/ovhcloud
               tests/unit/llms/parallel_ai
@@ -342,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
@@ -412,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"
@@ -446,13 +453,14 @@ jobs:
               tests/unit/repositories
               tests/unit/secret_managers
               tests/unit/skills/test_skills_main.py
+              tests/unit/telemetry tests/unit/proxy/telemetry
               tests/unit/tracing
             python-version: "3.12"
-            workers: 2
+            workers: 4
             reruns: 0
             timeout-minutes: 20
             job-timeout-minutes: 60
-            dist: loadscope
+            dist: load
 
           - shard: proxy-auth
             test-path: |-
@@ -557,7 +565,6 @@ jobs:
 
           - shard: mcp-elicitation
             test-path: |-
-              --override-ini=pythonpath=tests
               tests/unit/experimental_mcp_client/test_mcp_client.py
               tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
               tests/unit/proxy/_experimental/mcp_server/test_interactions.py
@@ -565,9 +572,8 @@ jobs:
               tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
               tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
               tests/unit/proxy/_experimental/mcp_server/test_operations.py
-              tests/integration/mcp/test_interactions.py
             python-version: "3.12"
-            workers: 2
+            workers: 4
             reruns: 0
             timeout-minutes: 20
             job-timeout-minutes: 60
@@ -1445,7 +1451,6 @@ jobs:
             tests/unit/caching/test_redis_connection_pool.py \
             tests/unit/caching/test_redis_cluster_cache.py \
             tests/unit/caching/test_evicted_client_closer.py \
-            tests/integration/sdk/test_redis_cluster_iam_auth.py \
             --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
             --cov=./litellm --cov-report=xml:coverage.xml
 

AGENTS.md

@@ -23,7 +23,7 @@ When adding new features, add meaningful tests. Don't add tests that don't check
 
 Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
 
-Never test structure of code only function of it
+Never test structure or implementation details of code, only function of it. A test should catch a change that breaks behavior by accident. If the only way it can fail is someone deliberately undoing the change, don't write it
 
 A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent". Assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
 

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:

backend/routes/allowlist.py

@@ -32,6 +32,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
     "/token",
     "/onboarding/",
     "/audit",
+    "/telemetry/",
     "/oauth/",
     "/invitation/",
     "/jwt/",

ci_cd/generate_model_prices_schema.py

@@ -189,6 +189,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = {
         **NONNEG_INTEGER,
         "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
     },
+    "token_count_framing": {
+        **NONNEG_INTEGER,
+        "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
+    },
     "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
     "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
 }
@@ -283,6 +287,11 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
             "enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
         },
         "comment": STRING,
+        "token_count_fallback_model": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Model that counts tokens for this one when the provider cannot count it.",
+        },
         "audio_transcription_config": STRING,
         "vertex_ai_audio_api": {
             "type": "string",

ci_cd/run_migration.py

@@ -322,7 +322,12 @@ def create_migration(
 
                     # Write the SQL to migration.sql
                     migration_file = migration_dir / "migration.sql"
-                    migration_file.write_text(diff_sql)
+                    migration_sql: Final = (
+                        diff_sql
+                        if "CONCURRENTLY" in diff_sql.upper()
+                        else "SET LOCAL lock_timeout = '10s';\n\n" + diff_sql
+                    )
+                    migration_file.write_text(migration_sql)
 
                     print(f"Created migration in {migration_dir}")
                     return True

cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json

@@ -6563,6 +6563,87 @@
       ],
       "title": "litellm_project_model_rate_limit_used_metric",
       "type": "timeseries"
+    },
+    {
+      "collapsed": false,
+      "gridPos": {
+        "h": 1,
+        "w": 24,
+        "x": 0,
+        "y": 447
+      },
+      "id": 115,
+      "panels": [],
+      "title": "Agent tracing export (needs LITELLM_LENS_URL)",
+      "type": "row"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "${DS_PROMETHEUS}"
+      },
+      "description": "Request records the Lens exporter delivered to the Lens service against the records it dropped before the service accepted them",
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "drawStyle": "line",
+            "fillOpacity": 10,
+            "lineWidth": 1,
+            "showPoints": "never",
+            "spanNulls": false
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 24,
+        "x": 0,
+        "y": 448
+      },
+      "id": 116,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "multi",
+          "sort": "desc"
+        }
+      },
+      "targets": [
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "editorMode": "code",
+          "expr": "sum(rate(litellm_lens_export_rows_written_total[$__rate_interval]))",
+          "legendFormat": "written",
+          "range": true,
+          "refId": "A"
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "editorMode": "code",
+          "expr": "sum(rate(litellm_lens_export_rows_dropped_total[$__rate_interval]))",
+          "legendFormat": "dropped",
+          "range": true,
+          "refId": "B"
+        }
+      ],
+      "title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate",
+      "type": "timeseries"
     }
   ],
   "preload": false,

cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md

@@ -1,10 +1,10 @@
 # LiteLLM All Prometheus Metrics dashboard
 
-Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
+Every `litellm_*` metric family the proxy can expose on `/metrics` (143 families across 98 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes), project model rate limits, and the agent tracing export to Lens. Panel titles are the metric names so you can grep the JSON for the metric you care about
 
 Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
 
-The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
+The first eleven rows and the project model rate limits row need only `callbacks: ["prometheus"]`. The other rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, the agent tracing export row needs `LITELLM_LENS_URL`, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
 
 ## Pre-requisites
 

docker/docker-compose.quickstart.yml

@@ -11,7 +11,7 @@
 # beyond local evaluation, pin the image to a specific release tag.
 services:
   litellm:
-    image: docker.litellm.ai/berriai/litellm:main-stable
+    image: docker.litellm.ai/berriai/litellm:latest
     ports:
       # LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart
       # script sets it to "127.0.0.1:" so new installs listen on this machine only.

enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py

@@ -1,14 +1,16 @@
-from typing import Any, Optional
+from collections.abc import Awaitable, Callable
 
 from fastapi import Request
 
 from litellm._logging import verbose_proxy_logger
-from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+from litellm.proxy._types import ProxyException
 
 
 async def enterprise_custom_auth(
-    request: Request, api_key: str, user_custom_auth: Optional[Any]
-) -> Optional[UserAPIKeyAuth]:
+    request: Request,
+    api_key: str,
+    user_custom_auth: Callable[[Request, str], Awaitable[object]] | None,
+) -> object:
     from litellm_enterprise.proxy.proxy_server import custom_auth_settings
 
     if user_custom_auth is None:

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -2,27 +2,34 @@
 Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
 """
 
-from collections.abc import Sequence
+import asyncio
+from collections.abc import Mapping, Sequence
 from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta, timezone
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
 
+from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
+    ManagedObjectCreatorAttribution,
+    ManagedObjectRow,
+)
+
+import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.constants import (
-    CLI_SESSION_KEY_PREFIX,
+    BATCH_LINE_ITEMS_MAX_PART_BYTES,
     MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
     MAX_OBJECTS_PER_POLL_CYCLE,
 )
 from litellm.repositories.table_repositories import ManagedObjectRepository
 from litellm.repositories.team_repository import TeamRepository
-from litellm.repositories.user_repository import UserRepository
 from litellm.repositories.verification_token_repository import VerificationTokenRepository
 
 if TYPE_CHECKING:
     from prisma import models as prisma_models
 
+    from litellm.batches.batch_line_items import BatchFileProvider
     from litellm.integrations.prometheus import PrometheusLogger
     from litellm.proxy.utils import PrismaClient, ProxyLogging
     from litellm.repositories.prisma_protocols import TableActions
@@ -47,29 +54,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
 )
 
 
-class _ManagedObjectRow(Protocol):
-    @property
-    def id(self) -> str: ...
-
-    @property
-    def unified_object_id(self) -> str: ...
-
-    @property
-    def created_by(self) -> str | None: ...
-
-    @property
-    def file_object(self) -> object: ...
-
-    @property
-    def org_id(self) -> str | None: ...
-
-    @property
-    def api_key(self) -> str | None: ...
-
-    @property
-    def team_id(self) -> str | None: ...
-
-
 class _ReadableFileContent(Protocol):
     async def read(self) -> bytes: ...
 
@@ -87,14 +71,10 @@ async def _file_content_bytes(file_content: object) -> bytes:
     return cast(bytes, file_content)
 
 
-def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
+def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
     return ManagedObjectRepository(prisma_client).table
 
 
-def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]":
-    return UserRepository(prisma_client).table
-
-
 def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
     return VerificationTokenRepository(prisma_client).table
 
@@ -103,6 +83,11 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
     return TeamRepository(prisma_client).table
 
 
+def _row_metadata(row: object) -> dict[str, object]:
+    metadata: Final = getattr(row, "metadata", None)
+    return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {}  # cast-ok: prisma Json object column
+
+
 class CheckBatchCost:
     def __init__(
         self,
@@ -118,10 +103,15 @@ class CheckBatchCost:
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
         self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
+        self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
         # Cached after the first poll cycle. Once we know the column is absent we skip
         # the guaranteed-failing primary query on every subsequent cycle.
         self._has_batch_processed_column: bool = True
         self.batch_processed_support_confirmed: bool = False
+        self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
+        self._line_item_store_turn: Final = asyncio.Lock()
+        self._poll_idle: Final = asyncio.Event()
+        self._poll_idle.set()
 
     @staticmethod
     def _is_missing_batch_processed_column_error(err: Exception) -> bool:
@@ -149,133 +139,6 @@ class CheckBatchCost:
             return
         self.batch_processed_support_confirmed = True
 
-    async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
-        """
-        Look up user email and key alias by user_id for enriching the S3 callback metadata.
-        Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
-        Returns an empty dict when user_id is None: batches created by a team or service
-        account key carry no user id, and find_unique(where={"user_id": None}) raises.
-        """
-        if not user_id:
-            return {}
-        try:
-            user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique(
-                where={"user_id": user_id}
-            )
-            if user_row is None:
-                return {}
-            return {
-                "user_api_key_user_email": getattr(user_row, "user_email", None),
-                "user_api_key_alias": getattr(user_row, "user_alias", None),
-            }
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
-            return {}
-
-    async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None:
-        """Resolve the creating virtual key's alias from its hashed token."""
-        if not api_key:
-            return None
-        if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}":
-            return api_key
-        try:
-            key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
-                self.prisma_client
-            ).find_unique(where={"token": api_key})
-            return getattr(key_row, "key_alias", None) if key_row is not None else None
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
-            return None
-
-    async def _get_team_alias(self, team_id: str | None) -> str | None:
-        """Resolve a team's alias from its id."""
-        if not team_id:
-            return None
-        try:
-            team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
-                where={"team_id": team_id}
-            )
-            return getattr(team_row, "team_alias", None) if team_row is not None else None
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
-            return None
-
-    async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
-        org_id: Final = getattr(job, "org_id", None)
-        if org_id:
-            return org_id
-        api_key: Final = getattr(job, "api_key", None)
-        team_id: Final = getattr(job, "team_id", None)
-        if api_key:
-            try:
-                key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
-                    self.prisma_client
-                ).find_unique(where={"token": api_key})
-                key_org_id: Final = (
-                    cast(str | None, getattr(key_row, "organization_id", None))
-                    if key_row is not None
-                    else None
-                )
-                if key_org_id:
-                    return key_org_id
-            except Exception as e:
-                verbose_proxy_logger.error(
-                    f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
-                    f"still trying the team's: {e}"
-                )
-        if not team_id:
-            return None
-        try:
-            team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
-                where={"team_id": team_id}
-            )
-            return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
-            return None
-
-    async def _build_creator_attribution_metadata(
-        self, job: "_ManagedObjectRow", batch_id: str
-    ) -> dict[str, object]:
-        """
-        Rebuild the spend-tracking metadata for the key, team, and tags that created the
-        batch so the batch-cost spend log is attributed the same way a non-batch request
-        is. Rows created before api_key and request_tags were persisted carry only
-        created_by and team_id, and fall back to those. A named creating key owns
-        user_api_key_alias; when it has no alias, or the key has since been rotated or
-        deleted, the field keeps the creating user's alias that _get_user_info filled in,
-        because a resolvable name is more useful on the spend row than a null.
-
-        user_api_key_org_id must be resolved here too: the spend update writer reads it
-        off this metadata to increment organization spend, so leaving it out silently
-        drops batch cost from org accounting for keys and teams that belong to one.
-        """
-        api_key = getattr(job, "api_key", None)
-        team_id = getattr(job, "team_id", None)
-        request_tags = getattr(job, "request_tags", None)
-
-        metadata: dict[str, object] = {
-            "user_api_key_user_id": job.created_by,
-            "user_api_key": api_key,
-            "user_api_key_hash": api_key,
-            "user_api_key_team_id": team_id,
-            **(await self._get_user_info(batch_id, job.created_by)),
-        }
-
-        key_alias = await self._get_key_alias(batch_id, api_key, job.created_by)
-        if key_alias is not None:
-            metadata["user_api_key_alias"] = key_alias
-        team_alias = await self._get_team_alias(team_id)
-        if team_alias is not None:
-            metadata["user_api_key_team_alias"] = team_alias
-        org_id: Final = await self._get_org_id(job, batch_id)
-        if org_id is not None:
-            metadata["user_api_key_org_id"] = org_id
-        if isinstance(request_tags, list) and request_tags:
-            metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
-
-        return metadata
-
     async def _cleanup_stale_managed_objects(self) -> None:
         """
         Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@@ -317,7 +180,7 @@ class CheckBatchCost:
                 f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
             )
 
-    async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
+    async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
         """Query batch jobs without the batch_processed filter (for older schemas)."""
         return await _managed_object_table(self.prisma_client).find_many(
             where={
@@ -337,7 +200,7 @@ class CheckBatchCost:
             order={"created_at": "asc"},
         )
 
-    async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
+    async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
         """
         Take a row that can never be costed out of the poll page. Leaving it selectable
         would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
@@ -364,7 +227,7 @@ class CheckBatchCost:
             "so it will no longer be polled"
         )
 
-    async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
+    async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
         """
         Atomically flip batch_processed from false to true, returning whether this pod won
         the row. Every pod and uvicorn worker schedules its own poller against the shared
@@ -390,7 +253,7 @@ class CheckBatchCost:
             return False
         return claimed > 0
 
-    async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
+    async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
         """Give a claimed row back once billing it failed, so a later poll cycle retries it.
 
         Safe to match on batch_processed=True: while this poller is active the retrieve
@@ -411,7 +274,7 @@ class CheckBatchCost:
             )
 
     @staticmethod
-    def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool:
+    def _has_unified_id_without_model(job: "ManagedObjectRow") -> bool:
         """A unified id that decodes but carries no model_id can never be routed."""
         from litellm.proxy.openai_files_endpoints.common_utils import (
             convert_b64_uid_to_unified_uid,
@@ -460,10 +323,10 @@ class CheckBatchCost:
         return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
 
     async def _finalize_unbilled_terminal_job(
-        self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
-    ) -> None:
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch"
+    ) -> bool:
         """Persist a terminal batch that has nothing billable, converting any raw
-        provider file ids to managed ids, and take it out of the poll page."""
+        provider file ids to managed ids, and take it out of the poll page. Returns whether it was persisted."""
         try:
             from litellm.proxy.openai_files_endpoints.common_utils import (
                 _is_base64_encoded_unified_file_id,
@@ -491,10 +354,12 @@ class CheckBatchCost:
             verbose_proxy_logger.info(
                 f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
             )
+            return True
         except Exception as db_err:
             verbose_proxy_logger.error(
                 f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
             )
+            return False
 
     @staticmethod
     def _record_error(
@@ -505,7 +370,7 @@ class CheckBatchCost:
 
     def _resolve_job_routing(
         self,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         prom_logger: Optional["PrometheusLogger"],
     ) -> Optional[Tuple[str, str]]:
         """
@@ -582,7 +447,7 @@ class CheckBatchCost:
 
     def _resolve_unmanaged_provider_routing(
         self,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         prom_logger: Optional["PrometheusLogger"],
         llm_provider: str,
         bare_model_name: str,
@@ -678,7 +543,7 @@ class CheckBatchCost:
     @classmethod
     def _get_managed_file_model_name(
         cls,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         deployment_info: "Deployment",
     ) -> Optional[str]:
         """
@@ -698,7 +563,7 @@ class CheckBatchCost:
         )
 
     @staticmethod
-    def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
+    def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
         import json
 
         from litellm.types.utils import LiteLLMBatch
@@ -719,7 +584,7 @@ class CheckBatchCost:
 
     async def _track_completed_batch_cost(
         self,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         response: "LiteLLMBatch",
         model_id: str,
         batch_id: str,
@@ -735,9 +600,9 @@ class CheckBatchCost:
 
         """
         from litellm.batches.batch_utils import (
+            calculate_batch_cost_and_usage,
             count_error_file_failed_requests,
             get_file_content_as_dictionary,
-            calculate_batch_cost_and_usage,
         )
         from litellm.files.main import afile_content
         from litellm.files.types import FileContentCallOptions
@@ -911,7 +776,7 @@ class CheckBatchCost:
                 },
                 **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
                 "metadata": {
-                    **(await self._build_creator_attribution_metadata(job, batch_id)),
+                    **(await self._creator_attribution.batch_metadata(job, batch_id)),
                     # spend logs read the deployment identity off these metadata keys, so
                     # without them the batch cost row carries no model_id or model_group
                     "model_info": {"id": model_id},
@@ -956,7 +821,191 @@ class CheckBatchCost:
 
         return model_name, str(llm_provider) if llm_provider else None
 
-    async def check_batch_cost(self):
+    async def _retire_unbilled_terminal_job(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str
+    ) -> None:
+        """Persist a terminal batch that has nothing billable, then store its lines, reading its input file by the id
+        the provider reported rather than the managed id that was persisted."""
+        provider_input_file_id: Final = response.input_file_id
+        if await self._finalize_unbilled_terminal_job(job, response):
+            response.input_file_id = provider_input_file_id
+            self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=False)
+
+    def _store_batch_line_items_in_background(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
+    ) -> None:
+        """Store a retired batch's lines off the poll loop, one batch at a time and only between polls, so neither
+        pricing the lines nor a slow or hung bucket holds up billing or retiring the other batches. While a poll
+        cycle's worth of batches is still waiting, the lines of the next one are dropped with an error."""
+        if not litellm.batch_line_item_callbacks:
+            return
+        if len(self._line_item_stores) >= MAX_OBJECTS_PER_POLL_CYCLE:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: {len(self._line_item_stores)} batches are still waiting to store their line items, "
+                f"so the line items of batch {batch_id} (job {job.id}) were not stored"
+            )
+            return
+        store: Final = asyncio.create_task(
+            self._store_batch_line_items_in_turn(job, response, model_id, batch_id, billed=billed)
+        )
+        self._line_item_stores.add(store)
+        store.add_done_callback(self._line_item_stores.discard)
+
+    async def _store_batch_line_items_in_turn(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
+    ) -> None:
+        async with self._line_item_store_turn:
+            await self._poll_idle.wait()
+            await self._store_batch_line_items(job, response, model_id, batch_id, billed=billed)
+
+    async def _store_batch_line_items(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
+    ) -> None:
+        """Send every line of a finished batch to batch_line_item_callbacks, one part at a time. A billed batch's
+        output file is read again here, so nothing waiting for its turn holds a batch's files. Failures, including an
+        output file that can no longer be read, are logged and never raised, so the batch stays billed and retired
+        whether or not its lines were stored."""
+        if not litellm.batch_line_item_callbacks:
+            return
+        from litellm.batches.batch_line_items import (
+            BatchFileProvider,
+            BatchLineItemSource,
+            batch_line_item_parts,
+            dispatch_batch_line_items,
+            line_item_metadata,
+        )
+        from litellm.batches.batch_utils import fetch_batch_managed_file_content
+        from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
+        from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
+
+        try:
+            deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
+            if deployment_info is None:
+                verbose_proxy_logger.error(
+                    f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
+                    f"{batch_id} (job {job.id}) were not stored"
+                )
+                return
+            litellm_model_name: Final = deployment_info.litellm_params.model
+            model_name, llm_provider, _, _ = get_llm_provider(
+                model=litellm_model_name,
+                custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
+            )
+            cost_provider: Final = cast(BatchFileProvider, llm_provider)
+            credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
+            file_access: Final = {
+                **credentials,
+                "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
+            }
+            output_content: Final = (
+                await fetch_batch_managed_file_content(
+                    response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access
+                )
+                if billed and response.output_file_id
+                else b""
+            )
+            error_content: Final = await self._batch_file_content(response.error_file_id, cost_provider, file_access)
+            if not output_content and not error_content:
+                return
+            redact_bodies: Final = await self._creator_turned_off_message_logging(job)
+            source: Final = BatchLineItemSource(
+                batch_id=job.unified_object_id,
+                finished_at=response.completed_at
+                or response.failed_at
+                or response.expired_at
+                or response.cancelled_at
+                or response.created_at,
+                model=model_name,
+                custom_llm_provider=llm_provider,
+                cost_provider=cost_provider,
+                model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
+                metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
+            )
+            input_content: Final = (
+                b""
+                if redact_bodies
+                else await self._batch_file_content(response.input_file_id, cost_provider, file_access)
+            )
+            parts: Final = batch_line_item_parts(
+                source, input_content, (output_content, error_content), BATCH_LINE_ITEMS_MAX_PART_BYTES
+            )
+            handed = 0
+            # pricing a part's lines is CPU-bound, so each part is built off the event loop
+            while (payload := await asyncio.to_thread(next, parts, None)) is not None:
+                await dispatch_batch_line_items(payload, redact_bodies=redact_bodies)
+                handed += len(payload["lines"])
+            if handed:
+                verbose_proxy_logger.info(
+                    f"CheckBatchCost: handed the {handed} line items of batch {batch_id} to batch_line_item_callbacks"
+                )
+        except Exception as error:  # noqa: BLE001  # storing the lines is best effort, the batch is already billed
+            verbose_proxy_logger.exception(
+                f"CheckBatchCost: the line items of batch {batch_id} (job {job.id}) were not stored: {error}"
+            )
+
+    @staticmethod
+    async def _batch_file_content(
+        file_id: str | None, custom_llm_provider: "BatchFileProvider", file_access: Mapping[str, object]
+    ) -> bytes:
+        """A batch file's raw JSONL, or nothing when the batch has no such file or the provider will not serve it."""
+        from litellm.batches.batch_utils import fetch_batch_managed_file_content
+
+        if not file_id:
+            return b""
+        try:
+            return await fetch_batch_managed_file_content(
+                file_id, custom_llm_provider=custom_llm_provider, litellm_params=file_access
+            )
+        except Exception as error:  # noqa: BLE001  # a missing input or error file leaves those lines out, not the rest
+            verbose_proxy_logger.warning(f"CheckBatchCost: could not read batch file {file_id}: {error}")
+            return b""
+
+    async def _creator_turned_off_message_logging(self, job: "ManagedObjectRow") -> bool:
+        """Whether the creating key's or team's logging settings turn off message logging, falling back to the
+        proxy-wide setting. A batch whose saved credential is not a virtual key's hash, such as one created under JWT
+        or OAuth2 auth or with the master key, has only its team's and the proxy's settings. A key or team that can no
+        longer be read counts as turned off."""
+        from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
+        from litellm.proxy._types import UserAPIKeyAuth
+        from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata
+        from litellm.proxy.proxy_server import proxy_config
+        from litellm.secret_managers.main import str_to_bool
+
+        api_key: Final = getattr(job, "api_key", None)
+        team_id: Final = getattr(job, "team_id", None)
+        has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
+        try:
+            key_row: Final = (
+                await _token_table(self.prisma_client).find_unique(where={"token": api_key}) if has_key_row else None
+            )
+            team_row: Final = (
+                await _team_table(self.prisma_client).find_unique(where={"team_id": team_id}) if team_id else None
+            )
+        except Exception as error:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: could not read the logging settings of batch job {job.id}'s creator, "
+                f"so its line items are stored without request and response: {error}"
+            )
+            return True
+        if (has_key_row and key_row is None) or (team_id and team_row is None):
+            return True
+        settings: Final = get_dynamic_logging_metadata(
+            UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
+            proxy_config,
+        )
+        callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
+        creator_setting: Final = str_to_bool(callback_vars.get("turn_off_message_logging"))
+        return creator_setting if creator_setting is not None else litellm.turn_off_message_logging is True
+
+    async def check_batch_cost(self) -> None:
+        """Run one poll cycle. Batch line items wait until it has billed and retired every batch it found."""
+        self._poll_idle.clear()
+        try:
+            await self._poll_batches()
+        finally:
+            self._poll_idle.set()
+
+    async def _poll_batches(self):
         """
         Check if the batch JOB has been tracked.
         - get all status="validating" and file_purpose="batch" jobs
@@ -1070,7 +1119,7 @@ class CheckBatchCost:
                             f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
                             f"does not exist at the provider; retiring job {job.id} unbilled"
                         )
-                        await self._finalize_unbilled_terminal_job(job, response)
+                        await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
                         continue
                     verbose_proxy_logger.error(
                         f"CheckBatchCost: failed to track cost for batch {batch_id} "
@@ -1100,6 +1149,10 @@ class CheckBatchCost:
                     verbose_proxy_logger.error(
                         f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
                     )
+                    if not self._has_batch_processed_column:
+                        # nothing claimed the row, so it stays in the poll page and its lines wait for that retry
+                        continue
+                self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
 
             elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
                 from litellm.proxy.openai_files_endpoints.common_utils import (
@@ -1112,7 +1165,7 @@ class CheckBatchCost:
                         f"has not appeared yet; leaving job {job.id} for the next poll cycle"
                     )
                     continue
-                await self._finalize_unbilled_terminal_job(job, response)
+                await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
 
         # Record polling run metrics (always, even if nothing was processed)
         if prom_logger:

enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py

@@ -1,12 +1,36 @@
 """
 Polls LiteLLM_ManagedObjectTable to check if the response is complete.
-Cost tracking is handled by the get-responses call, which prices normally only because the
-poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
-same route are non-inference and free.
+A background response is billed exactly once: whichever client read or poll first sees it
+finished moves its row out of a pending status, then logs the body that read already fetched as
+one spend event stamped with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN and the creating key's
+attribution, priced against the deployment that served it. The poll's status checks are stamped
+with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
+managed files hook, so both stay free. The spend event carries the id the client holds, the row's
+unified_object_id, so its spend log is found under that id with the _background_cost suffix.
+A read that finds no row, because the create could not store one, stores it under the reading key
+and bills it the same way, so a response the create left untracked is still billed once. A streamed
+background create is free on its own; its stream stores the row under the creating key and bills
+the finished response through it before the client sees response.completed.
+A read the caller retargets with client-side credentials, and a delete of a still-pending response,
+settle the bill from a fetch through the configured deployment instead: the client controls what a
+retargeted read returns, and a deleted response leaves the poll nothing to fetch.
+A row stored before this billing model keeps the encrypted create id as its model_object_id. The
+first lookup by that id rekeys it by the provider id, and the poll retires such a row without
+billing it once a read has already tracked the response under the provider id.
 """
 
+import uuid
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Dict, Final, Protocol, cast
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, cast
+
+from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
+    ManagedObjectCreatorAttribution,
+    ManagedObjectRow,
+)
+from pydantic import BaseModel, ConfigDict, TypeAdapter
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -14,36 +38,124 @@ from litellm.constants import (
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
     MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
     MAX_OBJECTS_PER_POLL_CYCLE,
+    PENDING_RESPONSE_STATUSES,
     STALE_OBJECT_CLEANUP_BATCH_SIZE,
+    TERMINAL_RESPONSE_STATUSES,
 )
+from litellm.proxy.hooks.responses_id_security import decrypted_response_id
 from litellm.repositories.table_repositories import ManagedObjectRepository
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import ResponsesAPIResponse
-from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
+from litellm.types.router import Deployment
+from litellm.types.utils import (
+    BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
+    BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
+)
 
 if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
     from litellm.proxy.utils import PrismaClient, ProxyLogging
     from litellm.repositories.prisma_protocols import TableActions
     from litellm.router import Router
 
-TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
+_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
+_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
 
 
-class _ManagedObjectRow(Protocol):
-    @property
-    def id(self) -> str: ...
-
-    @property
-    def unified_object_id(self) -> str: ...
-
-    @property
-    def created_by(self) -> str | None: ...
-
-    @property
-    def file_object(self) -> object: ...
+@dataclass(frozen=True, slots=True)
+class _PollTarget:
+    response_id: str
+    model_metadata: Mapping[str, object]
+    deployment: Deployment | None
 
 
-def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
+@dataclass(frozen=True, slots=True)
+class BilledBackgroundResponse:
+    response: ResponsesAPIResponse
+    model: str
+    custom_llm_provider: str | None
+    litellm_params: Mapping[str, object]
+    metadata: Mapping[str, object]
+
+
+@dataclass(frozen=True, slots=True)
+class BilledCharge:
+    cost: float | None
+    logging_obj: "LiteLLMLogging"
+
+
+BilledResponseLogger = Callable[[BilledBackgroundResponse], Awaitable[BilledCharge]]
+UntrackedResponseTracker = Callable[[ResponsesAPIResponse], Awaitable[bool]]
+
+
+def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
+    unpriced: Final = response.model_copy(deep=True)
+    unpriced.hidden_params = {key: value for key, value in response.hidden_params.items() if key != "response_cost"}
+    return unpriced
+
+
+async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+    logging_obj: Final = LiteLLMLogging(
+        model=billed.model,
+        messages=[{"role": "user", "content": "<retrieve_response>"}],
+        stream=False,
+        call_type=_RETRIEVE_RESPONSE_CALL_TYPE,
+        start_time=datetime.now(),
+        litellm_call_id=str(uuid.uuid4()),
+        function_id=str(uuid.uuid4()),
+    )
+    logging_obj.update_environment_variables(
+        litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
+        optional_params={},
+        custom_llm_provider=billed.custom_llm_provider,
+    )
+    unpriced: Final = _without_read_price(billed.response)
+    cost: Final = logging_obj.response_cost_calculator(result=unpriced)
+    end_time: Final = datetime.now()
+    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
+        async_coroutine=logging_obj.async_success_handler(
+            result=unpriced, start_time=logging_obj.start_time, end_time=end_time
+        )
+    )
+    logging_obj.handle_sync_success_callbacks_for_async_calls(
+        result=unpriced, start_time=logging_obj.start_time, end_time=end_time
+    )
+    return BilledCharge(cost=cost, logging_obj=logging_obj)
+
+
+class _DeploymentBaseModel(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    base_model: str | None = None
+
+
+def _deployment_litellm_params(deployment: Deployment) -> Mapping[str, object]:
+    from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
+
+    api_base: Final = deployment.litellm_params.api_base
+    base_model: Final = _DeploymentBaseModel.model_validate(deployment.litellm_params, from_attributes=True).base_model
+    return MappingProxyType(
+        {
+            **({"api_base": mask_api_base_credentials(api_base)} if api_base else {}),
+            **({"base_model": base_model} if base_model else {}),
+        }
+    )
+
+
+def _deployment_metadata(deployment: Deployment) -> Mapping[str, object]:
+    return MappingProxyType(
+        {
+            "model_info": deployment.model_info.model_dump(exclude_none=True),
+            "model_group": deployment.model_name,
+            "deployment": deployment.litellm_params.model,
+        }
+    )
+
+
+def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
     return ManagedObjectRepository(prisma_client).table
 
 
@@ -51,12 +163,17 @@ def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -
     return getattr(error, "status_code", None) == 404 and provider_response_id in str(error)
 
 
+def _keyed_by_issued_id(job: ManagedObjectRow) -> bool:
+    return decrypted_response_id(job.model_object_id) != job.model_object_id
+
+
 class CheckResponsesCost:
     def __init__(
         self,
         proxy_logging_obj: "ProxyLogging",
         prisma_client: "PrismaClient",
         llm_router: "Router",
+        log_billed_response: BilledResponseLogger = log_billed_background_response,
     ):
         from litellm.proxy.utils import PrismaClient, ProxyLogging
         from litellm.router import Router
@@ -64,15 +181,17 @@ class CheckResponsesCost:
         self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
+        self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
+        self._log_billed_response: Final = log_billed_response
 
-    def _resolve_deployment(self, response_id: str) -> bool:
+    def _resolve_deployment(self, response_id: str) -> Deployment | None:
         model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
-        return model_id is not None and self.llm_router.get_deployment(model_id=model_id) is not None
+        return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
 
     async def _get_response(
         self,
         response_id: str,
-        litellm_metadata: Dict[str, str],
+        litellm_metadata: dict[str, object],
         via_router: bool,
     ) -> ResponsesAPIResponse:
         """Fetch the upstream response, using deployment credentials when available.
@@ -91,9 +210,17 @@ class CheckResponsesCost:
         )
         return cast(ResponsesAPIResponse, router_response)
 
-    async def _expire_stale_rows(
-        self, cutoff: datetime, batch_size: int
-    ) -> int:
+    async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
+        return await self._get_response(
+            response_id=target.response_id,
+            litellm_metadata={
+                **target.model_metadata,
+                INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
+            },
+            via_router=target.deployment is not None,
+        )
+
+    async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int:
         """Execute the bounded UPDATE that marks stale rows as 'stale_expired'.
 
         Isolated so it can be swapped / mocked in tests without touching the
@@ -140,14 +267,154 @@ class CheckResponsesCost:
                 f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
             )
 
+    def _target(self, response_id: str, model_name: str | None) -> _PollTarget:
+        return _PollTarget(
+            response_id=response_id,
+            model_metadata=MappingProxyType({"model": model_name, "model_group": model_name} if model_name else {}),
+            deployment=self._resolve_deployment(response_id),
+        )
+
+    def _poll_target(self, job: ManagedObjectRow) -> _PollTarget:
+        stored_model: Final = _STORED_RESPONSE_ADAPTER.validate_python(job.file_object).get("model")
+        return self._target(
+            decrypted_response_id(job.unified_object_id), stored_model if isinstance(stored_model, str) else None
+        )
+
+    async def _claim_for_billing(self, job: ManagedObjectRow) -> bool:
+        claimed: Final = await _managed_object_table(self.prisma_client).update_many(
+            # bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
+            where={"id": job.id, "status": {"in": list(PENDING_RESPONSE_STATUSES)}},
+            data={"status": "completed"},
+        )
+        return claimed > 0
+
+    async def _release_billing_claim(self, job: ManagedObjectRow) -> None:
+        await _managed_object_table(self.prisma_client).update_many(
+            where={"id": job.id, "status": "completed"},
+            data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
+        )
+
+    async def _billed_response(
+        self,
+        job: ManagedObjectRow,
+        target: _PollTarget,
+        response: ResponsesAPIResponse,
+    ) -> BilledBackgroundResponse:
+        deployment: Final = target.deployment
+        return BilledBackgroundResponse(
+            response=response.model_copy(update={"id": job.unified_object_id}),
+            model=deployment.litellm_params.model if deployment is not None else (response.model or ""),
+            custom_llm_provider=ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
+                "custom_llm_provider"
+            ),
+            litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
+            metadata=MappingProxyType(
+                {
+                    **(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
+                    INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
+                    **target.model_metadata,
+                    **(_deployment_metadata(deployment) if deployment is not None else {}),
+                }
+            ),
+        )
+
+    async def _bill_once(
+        self,
+        job: ManagedObjectRow,
+        target: _PollTarget,
+        response: ResponsesAPIResponse,
+    ) -> BilledCharge | None:
+        if not await self._claim_for_billing(job):
+            verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
+            return None
+        try:
+            charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
+        except BaseException:
+            await self._release_billing_claim(job)
+            raise
+        verbose_proxy_logger.info(f"Response {job.unified_object_id} reached a terminal status and was billed")
+        return charge
+
+    async def _job_for_provider_id(self, provider_response_id: str) -> ManagedObjectRow | None:
+        return await _managed_object_table(self.prisma_client).find_first(
+            where={"model_object_id": provider_response_id, "file_purpose": "response"}
+        )
+
+    async def _rekey_by_provider_id(self, job: ManagedObjectRow) -> None:
+        await _managed_object_table(self.prisma_client).update_many(
+            where={"id": job.id}, data={"model_object_id": decrypted_response_id(job.model_object_id)}
+        )
+
+    async def _tracked_job(self, response_id: str) -> ManagedObjectRow | None:
+        if tracked := await self._job_for_provider_id(decrypted_response_id(response_id)):
+            return tracked
+        stored_by_issued_id: Final = await _managed_object_table(self.prisma_client).find_first(
+            where={"unified_object_id": response_id, "file_purpose": "response"}
+        )
+        if stored_by_issued_id is not None and _keyed_by_issued_id(stored_by_issued_id):
+            await self._rekey_by_provider_id(stored_by_issued_id)
+        return stored_by_issued_id
+
+    async def _superseded_by_a_provider_id_row(self, job: ManagedObjectRow) -> bool:
+        if not _keyed_by_issued_id(job):
+            return False
+        if await self._job_for_provider_id(decrypted_response_id(job.model_object_id)) is None:
+            await self._rekey_by_provider_id(job)
+            return False
+        await self._claim_for_billing(job)
+        return True
+
+    async def _track_from_deployment(
+        self, response_id: str, track_untracked: UntrackedResponseTracker | None
+    ) -> BilledCharge | None:
+        target: Final = self._target(decrypted_response_id(response_id), None)
+        if track_untracked is None or target.deployment is None:
+            return None
+        return await self.bill_if_finished(response_id, await self._fetch_unbilled(target), track_untracked)
+
+    async def _start_tracking(
+        self, response_id: str, response: ResponsesAPIResponse, track_untracked: UntrackedResponseTracker | None
+    ) -> ManagedObjectRow | None:
+        if track_untracked is None or not await track_untracked(response):
+            return None
+        return await self._tracked_job(response_id)
+
+    async def bill_if_finished(
+        self,
+        response_id: str,
+        response: ResponsesAPIResponse,
+        track_untracked: UntrackedResponseTracker | None = None,
+    ) -> BilledCharge | None:
+        if response.status not in TERMINAL_RESPONSE_STATUSES:
+            return None
+        job: Final = await self._tracked_job(response_id) or await self._start_tracking(
+            response_id, response, track_untracked
+        )
+        if job is None or job.status not in PENDING_RESPONSE_STATUSES:
+            return None
+        return await self._bill_once(job, self._poll_target(job), response)
+
+    async def bill_from_deployment(
+        self, response_id: str, track_untracked: UntrackedResponseTracker | None = None
+    ) -> BilledCharge | None:
+        job: Final = await self._tracked_job(response_id)
+        if job is None:
+            return await self._track_from_deployment(response_id, track_untracked)
+        if job.status not in PENDING_RESPONSE_STATUSES:
+            return None
+        target: Final = self._poll_target(job)
+        response: Final = await self._fetch_unbilled(target)
+        if response.status not in TERMINAL_RESPONSE_STATUSES:
+            return None
+        return await self._bill_once(job, target, response)
+
     async def check_responses_cost(self):
         """
         Check if background responses are complete and track their cost.
         - Get all status="queued" or "in_progress" and file_purpose="response" jobs
-        - Query the provider to check if response is complete
-        - Cost is tracked by the get-responses call, billed because the poll is stamped
-          with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
-        - Mark responses in a terminal state as complete in the database
+        - Query the provider, unbilled, to check if the response is complete
+        - Claim each terminal job by moving its row to completed, then log the body that check fetched
+          as its one billed event, so two pollers racing over the same job bill it once
         - Mark responses the provider no longer has (404 through a resolved
           deployment) as stale_expired
         """
@@ -160,99 +427,53 @@ class CheckResponsesCost:
 
         jobs = await _managed_object_table(self.prisma_client).find_many(
             where={
-                "status": {"in": ["queued", "in_progress"]},
+                # bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
+                "status": {"in": list(PENDING_RESPONSE_STATUSES)},
                 "file_purpose": "response",
             },
             take=MAX_OBJECTS_PER_POLL_CYCLE,
             order={"created_at": "asc"},
         )
-        
+
         verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
-        completed_jobs: Final[list[_ManagedObjectRow]] = []
-        expired_jobs: Final[list[_ManagedObjectRow]] = []
+        expired_jobs: Final[list[ManagedObjectRow]] = []
 
         for job in jobs:
             unified_object_id = job.unified_object_id
 
             try:
-                from litellm.proxy.hooks.responses_id_security import (
-                    ResponsesIDSecurity,
-                )
-
-                # Get the stored response object to extract model information
-                stored_response = job.file_object
-                model_name = stored_response.get("model", None)
-
-                # Decrypt the response ID
-                responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
-
-                # Prepare metadata with model information for cost tracking
-                litellm_metadata = {
-                    "user_api_key_user_id": job.created_by or "default-user-id",
-                    INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
-                }
-
-                # Add model information if available
-                if model_name:
-                    litellm_metadata["model"] = model_name
-                    litellm_metadata["model_group"] = model_name  # Use same value for model_group
-                via_router = self._resolve_deployment(responses_id_security)
+                if await self._superseded_by_a_provider_id_row(job):
+                    verbose_proxy_logger.info(f"Response {unified_object_id} is billed under its provider id row")
+                    continue
+                target = self._poll_target(job)
             except Exception as e:
-                verbose_proxy_logger.warning(
-                    f"Skipping job {unified_object_id} due to error: {e}"
-                )
+                verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
                 continue
 
-            provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(
-                responses_id_security
-            ).get("response_id", responses_id_security)
+            provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
+                "response_id", target.response_id
+            )
             try:
-                response = await self._get_response(
-                    response_id=responses_id_security,
-                    litellm_metadata=litellm_metadata,
-                    via_router=via_router,
-                )
+                response = await self._fetch_unbilled(target)
 
-                verbose_proxy_logger.debug(
-                    f"Response {unified_object_id} status: {response.status}, model: {model_name}"
-                )
+                verbose_proxy_logger.debug(f"Response {unified_object_id} status: {response.status}")
+                if response.status in TERMINAL_RESPONSE_STATUSES:
+                    await self._bill_once(job, target, response)
 
             except Exception as e:
-                if via_router and _is_response_gone_at_provider(e, provider_response_id):
+                if target.deployment is not None and _is_response_gone_at_provider(e, provider_response_id):
                     verbose_proxy_logger.info(
                         f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}"
                     )
                     expired_jobs.append(job)
                     continue
-                verbose_proxy_logger.warning(
-                    f"Skipping job {unified_object_id} due to error: {e}"
-                )
+                verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
                 continue
 
-            if response.status in TERMINAL_RESPONSE_STATUSES:
-                verbose_proxy_logger.info(
-                    f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
-                )
-                completed_jobs.append(job)
-
-        # Mark completed jobs in the database
-        if len(completed_jobs) > 0:
-            await _managed_object_table(self.prisma_client).update_many(
-                # bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
-                where={"id": {"in": [job.id for job in completed_jobs]}},
-                data={"status": "completed"},
-            )
-            verbose_proxy_logger.info(
-                f"Marked {len(completed_jobs)} response jobs as completed"
-            )
-
         if len(expired_jobs) > 0:
             await _managed_object_table(self.prisma_client).update_many(
                 # bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
                 where={"id": {"in": [job.id for job in expired_jobs]}},
                 data={"status": "stale_expired"},
             )
-            verbose_proxy_logger.info(
-                f"Marked {len(expired_jobs)} response jobs as stale_expired"
-            )
-
+            verbose_proxy_logger.info(f"Marked {len(expired_jobs)} response jobs as stale_expired")

enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py

@@ -0,0 +1,218 @@
+"""Creator attribution for the spend row a managed-object cost poll writes.
+
+A batch or a background response is billed by a poll job that runs long after the creating
+request returned, so the managed object row's creator columns are the only record of which
+key, team, organization, user, end user, and tags carry the spend.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final, Protocol, cast
+
+from pydantic import TypeAdapter
+
+from litellm._logging import verbose_proxy_logger
+from litellm.constants import CLI_SESSION_KEY_PREFIX
+from litellm.repositories.team_repository import TeamRepository
+from litellm.repositories.user_repository import UserRepository
+from litellm.repositories.verification_token_repository import VerificationTokenRepository
+
+if TYPE_CHECKING:
+    from prisma import models as prisma_models
+
+    from litellm.proxy.utils import PrismaClient
+    from litellm.repositories.prisma_protocols import TableActions
+
+
+class ManagedObjectRow(Protocol):
+    @property
+    def id(self) -> str: ...
+
+    @property
+    def unified_object_id(self) -> str: ...
+
+    @property
+    def model_object_id(self) -> str: ...
+
+    @property
+    def created_by(self) -> str | None: ...
+
+    @property
+    def status(self) -> str | None: ...
+
+    @property
+    def file_object(self) -> object: ...
+
+    @property
+    def org_id(self) -> str | None: ...
+
+    @property
+    def api_key(self) -> str | None: ...
+
+    @property
+    def team_id(self) -> str | None: ...
+
+    @property
+    def request_tags(self) -> object: ...
+
+    @property
+    def end_user_id(self) -> str | None: ...
+
+
+_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
+_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
+
+
+def _string_tags(request_tags: object) -> list[str] | None:
+    if not isinstance(request_tags, list):
+        return None
+    return [tag for tag in _TAG_LIST_ADAPTER.validate_python(request_tags) if isinstance(tag, str)] or None
+
+
+def _creator_user_id(created_by: str | None) -> str | None:
+    if created_by is None or created_by.startswith(_KEY_OWNER_SCOPE_PREFIX):
+        return None
+    return created_by
+
+
+def _user_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_UserTable]:
+    return UserRepository(prisma_client).table
+
+
+def _token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_VerificationToken]:
+    return VerificationTokenRepository(prisma_client).table
+
+
+def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_TeamTable]:
+    return TeamRepository(prisma_client).table
+
+
+@dataclass(frozen=True, slots=True)
+class _CreatorRows:
+    key: prisma_models.LiteLLM_VerificationToken | None
+    team: prisma_models.LiteLLM_TeamTable | None
+    user: prisma_models.LiteLLM_UserTable | None
+
+
+class ManagedObjectCreatorAttribution:
+    """Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
+    job's spend row is attributed like the creating request: key hash, team, organization,
+    user, aliases, and tags.
+
+    Rows created before api_key and request_tags were persisted carry only created_by and
+    team_id, and fall back to those. A named creating key owns user_api_key_alias; when it
+    has no alias, or the key has since been rotated or deleted, the field keeps the creating
+    user's alias, because a resolvable name is more useful on the spend row than a null.
+    user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
+    to increment organization spend, so leaving it out drops the cost from org accounting.
+    A background response also carries its end user and the key, team, and user budgets, the
+    way the read that billed it before did, and never names a key-scoped creator as its user.
+    """
+
+    def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
+        self._prisma_client: Final = prisma_client
+        self._poller_name: Final = poller_name
+
+    async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
+        rows: Final = await self._rows(job, object_id)
+        return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
+
+    async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
+        rows: Final = await self._rows(job, object_id)
+        end_user_id: Final = getattr(job, "end_user_id", None)
+        return {
+            "user_api_key_user_id": _creator_user_id(job.created_by),
+            **_shared_metadata(job, rows),
+            **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
+            **_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
+            **_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
+            **_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
+        }
+
+    async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
+        return _CreatorRows(
+            key=await self._key_row(object_id, getattr(job, "api_key", None), job.created_by),
+            team=await self._team_row(object_id, getattr(job, "team_id", None)),
+            user=await self._user_row(object_id, _creator_user_id(job.created_by)),
+        )
+
+    async def _user_row(self, object_id: str, user_id: str | None) -> prisma_models.LiteLLM_UserTable | None:
+        """Objects created by a team or service account key carry no user id, and
+        find_unique(where={"user_id": None}) raises, so those resolve no row."""
+        if not user_id:
+            return None
+        try:
+            return await _user_table(self._prisma_client).find_unique(where={"user_id": user_id})
+        except Exception as e:
+            verbose_proxy_logger.error(f"{self._poller_name}: could not look up user {user_id} for {object_id}: {e}")
+            return None
+
+    async def _key_row(
+        self, object_id: str, api_key: str | None, created_by: str | None
+    ) -> prisma_models.LiteLLM_VerificationToken | None:
+        if not api_key or _is_cli_session_key(api_key, created_by):
+            return None
+        try:
+            return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
+        except Exception as e:
+            verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
+            return None
+
+    async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
+        if not team_id:
+            return None
+        try:
+            return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
+        except Exception as e:
+            verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
+            return None
+
+
+def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
+    api_key: Final = getattr(job, "api_key", None)
+    team_id: Final = getattr(job, "team_id", None)
+    key_alias: Final = _key_alias(api_key, job.created_by, rows.key)
+    team_alias: Final = getattr(rows.team, "team_alias", None)
+    org_id: Final = getattr(job, "org_id", None) or _organization_id(rows.key) or _organization_id(rows.team)
+    tags: Final = _string_tags(getattr(job, "request_tags", None))
+    return {
+        "user_api_key": api_key,
+        "user_api_key_hash": api_key,
+        "user_api_key_team_id": team_id,
+        **(
+            {
+                "user_api_key_user_email": getattr(rows.user, "user_email", None),
+                "user_api_key_alias": getattr(rows.user, "user_alias", None),
+            }
+            if rows.user is not None
+            else {}
+        ),
+        **({"user_api_key_alias": key_alias} if key_alias is not None else {}),
+        **({"user_api_key_team_alias": team_alias} if team_alias is not None else {}),
+        **({"user_api_key_org_id": org_id} if org_id is not None else {}),
+        **({"tags": tags} if tags else {}),
+    }
+
+
+def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
+    return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
+
+
+def _key_alias(
+    api_key: str | None, created_by: str | None, key_row: prisma_models.LiteLLM_VerificationToken | None
+) -> str | None:
+    if api_key and _is_cli_session_key(api_key, created_by):
+        return api_key
+    return getattr(key_row, "key_alias", None) if key_row is not None else None
+
+
+def _organization_id(row: object) -> str | None:
+    return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
+
+
+def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
+    if row is None:
+        return {}
+    return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -24,7 +24,9 @@ from uuid import NAMESPACE_URL, uuid5
 
 import httpx
 from fastapi import HTTPException
-from pydantic import ValidationError
+from openai import APIConnectionError, AsyncOpenAI
+from openai.types.file_deleted import FileDeleted
+from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, Unpack
 
 import litellm
@@ -35,17 +37,17 @@ from litellm.caching.caching import DualCache
 from litellm.constants import (
     BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
     BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
+    MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
     MAX_FILE_LIST_LIMIT,
+    PENDING_RESPONSE_STATUSES,
 )
 from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
+from litellm.litellm_core_utils.internal_call_metadata import requests_background
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     extract_file_metadata,
 )
-from openai import APIConnectionError, AsyncOpenAI
-from openai.types.file_deleted import FileDeleted
-
 from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
 from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
 from litellm.llms.base_llm.managed_resources.isolation import (
@@ -54,6 +56,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,
@@ -87,6 +90,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
     request_tags_from_metadata,
 )
 from litellm.repositories.managed_file_repository import ManagedFileRepository
+from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccessIssue]
     AllMessageValues,
     AsyncCursorPage,
@@ -99,9 +103,11 @@ from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccess
     ResponsesAPIResponse,
 )
 from litellm.types.utils import (
+    ADDRESSED_RESPONSE_ID_FIELD,
     CallTypesLiteral,
     LiteLLMBatch,
     LiteLLMFineTuningJob,
+    LlmProviders,
     LLMResponseTypes,
     SpecialEnums,
 )
@@ -194,6 +200,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 +284,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
@@ -301,6 +350,8 @@ class _ManagedObjectTableActions(Protocol):
 
     async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
 
+    async def create(self, data: Mapping[str, object]) -> "PrismaManagedObjectRow": ...
+
 
 class _ManagedResourceDatabase(Protocol):
     @property
@@ -347,6 +398,39 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
     return prisma_client.db.litellm_managedobjecttable
 
 
+_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
+
+
+def _storable_text(text: str) -> str:
+    return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
+
+
+def _storable_json_value(value: JsonValue) -> JsonValue:
+    if isinstance(value, str):
+        return _storable_text(value)
+    if isinstance(value, list):
+        return [_storable_json_value(item) for item in value]
+    if isinstance(value, dict):
+        return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
+    return value
+
+
+def _postgres_storable_json(file_object: BaseModel) -> str:
+    dumped: Final = file_object.model_dump_json()
+    if "\\u0000" not in dumped:
+        return dumped
+    return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
+
+
+def _without_model_output(
+    file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
+) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
+    return (
+        file_object.model_copy(update={"output": []}) if isinstance(file_object, ResponsesAPIResponse) else file_object
+    )
+
+
 def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
     hidden_params: Final = cast(  # cast-ok: _hidden_params is an untyped attribute the upload path sets
         "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@@ -360,6 +444,40 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
     )
 
 
+_MANAGED_RESPONSE_READ_CALL_TYPES: Final = frozenset({"aget_responses", "acancel_responses"})
+_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object])
+
+
+def _addresses_a_proxy_issued_response(data: Mapping[str, object]) -> bool:
+    from litellm.proxy.hooks.responses_id_security import decrypted_response_id
+
+    addressed_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) or data.get("response_id")
+    return (
+        isinstance(addressed_id, str)
+        and ResponsesAPIRequestUtils.get_model_id_from_response_id(decrypted_response_id(addressed_id)) is not None
+    )
+
+
+def _is_streamed_background_create(request: Mapping[str, object]) -> bool:
+    return (
+        request.get("stream") is True
+        and isinstance(request.get("litellm_metadata"), dict)
+        and requests_background(request)
+    )
+
+
+def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, object]:
+    metadata_field: Final = "litellm_metadata" if isinstance(request.get("litellm_metadata"), dict) else "metadata"
+    existing_metadata: Final = request.get(metadata_field)
+    return {
+        **request,
+        metadata_field: {
+            **(_REQUEST_FIELDS.validate_python(existing_metadata) if isinstance(existing_metadata, dict) else {}),
+            MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True,
+        },
+    }
+
+
 _MANAGED_FILES_TARGET: Final = "managed_files"
 _PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
 
@@ -484,11 +602,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
     ) -> None:
         """Persist a managed object row, caching it and upserting it in the DB.
 
-        persist_attribution is set only by the batch create, which is the one caller
-        that can speak for the creator; it gates the api_key and request_tags columns
-        that CheckBatchCost bills against, so a later poll or retrieve of the same
-        batch cannot record itself as the paying key. Like created_by and team_id,
-        both are written only in the upsert create branch, never on update.
+        persist_attribution is set only by the batch and background response creates,
+        the callers that can speak for the creator; it gates the api_key, request_tags, and,
+        for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
+        retrieve of the same object cannot record itself as the paying key. Like created_by and team_id,
+        they are written only in the upsert create branch, never on update.
+
+        A stored response keeps no model output: billing reads only its status and model from
+        the row, so the row never holds generated content.
 
         create_if_missing is cleared by callers that observe a batch they did not
         create, such as a poll. They still refresh status and file_object, but a
@@ -501,11 +622,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         It is written only in the upsert create branch.
         """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
+        stored_object: Final = _without_model_output(file_object)
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
             model_object_id=model_object_id,
             file_purpose=file_purpose,
-            file_object=file_object,
+            file_object=stored_object,
         )
         await self.internal_usage_cache.async_set_cache(
             key=unified_object_id,
@@ -516,32 +638,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         from prisma import Json
 
         api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
+        end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None
         attribution_columns = (
             {
                 **({"api_key": api_key} if api_key is not None else {}),
-                **({"request_tags": Json(list(request_tags))} if request_tags else {}),
+                **({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
+                **({"end_user_id": end_user_id} if end_user_id else {}),
             }
             if persist_attribution
             else {}
         )
         # FIX: Update status and file_object on every operation to keep state in sync
         update_columns: Final = {
-            "file_object": file_object.model_dump_json(),
+            "file_object": _postgres_storable_json(stored_object),
             "status": file_object.status,
             "updated_by": user_api_key_dict.user_id,
         }
+        managed_object_table: Final = _managed_object_table(self.prisma_client)
         if not create_if_missing:
-            await _managed_object_table(self.prisma_client).update_many(
+            await managed_object_table.update_many(
                 where={"unified_object_id": unified_object_id},
                 data=update_columns,
             )
             return
-        await _managed_object_table(self.prisma_client).upsert(
+        await managed_object_table.upsert(
             where={"unified_object_id": unified_object_id},
             data={
                 "create": {
                     "unified_object_id": unified_object_id,
-                    "file_object": file_object.model_dump_json(),
+                    "file_object": _postgres_storable_json(stored_object),
                     "model_object_id": model_object_id,
                     "file_purpose": file_purpose,
                     "created_by": resolve_resource_owner_id(user_api_key_dict),
@@ -556,6 +681,42 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             },
         )
 
+    async def track_untracked_background_response(
+        self,
+        response_id: str,
+        response: ResponsesAPIResponse,
+        user_api_key_dict: UserAPIKeyAuth,
+        request_tags: Sequence[str] | None = None,
+    ) -> bool:
+        if not _addresses_a_proxy_issued_response({"response_id": response_id}):
+            return False
+        from prisma import Json
+        from prisma.errors import UniqueViolationError
+
+        from litellm.proxy.hooks.responses_id_security import decrypted_response_id
+
+        api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
+        try:
+            await _managed_object_table(self.prisma_client).create(
+                data={
+                    "unified_object_id": response_id,
+                    "model_object_id": decrypted_response_id(response_id),
+                    "file_purpose": "response",
+                    "file_object": _postgres_storable_json(_without_model_output(response)),
+                    "status": PENDING_RESPONSE_STATUSES[0],
+                    "created_by": resolve_resource_owner_id(user_api_key_dict),
+                    "team_id": user_api_key_dict.team_id,
+                    "org_id": await self._resolve_creator_org_id(user_api_key_dict),
+                    "updated_by": user_api_key_dict.user_id,
+                    **({"api_key": api_key} if api_key is not None else {}),
+                    **({"end_user_id": user_api_key_dict.end_user_id} if user_api_key_dict.end_user_id else {}),
+                    **({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
+                }
+            )
+        except UniqueViolationError:
+            return False
+        return True
+
     @with_service_target(_MANAGED_FILES_TARGET)
     async def get_unified_file_id(
         self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
@@ -845,6 +1006,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 +1097,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 +1148,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,
@@ -1258,6 +1491,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools)
                 if unified_vs_ids:
                     await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict)
+            create_request: Final = _REQUEST_FIELDS.validate_python(data)
+            if _is_streamed_background_create(create_request):
+                return _with_managed_background_marker(create_request)
         elif call_type == CallTypes.afile_content.value:
             retrieve_file_id = cast(Optional[str], data.get("file_id"))
             potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
@@ -1315,6 +1551,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
             elif retrieve_object_id and accessor_key == "batch_id":
                 await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
+        elif call_type in _MANAGED_RESPONSE_READ_CALL_TYPES:
+            request: Final = _REQUEST_FIELDS.validate_python(data)
+            if _addresses_a_proxy_issued_response(request):
+                return _with_managed_background_marker(request)
         elif call_type == CallTypes.acreate_fine_tuning_job.value:
             fine_tuning_file_ids: Final = tuple(
                 file_id
@@ -1943,8 +2183,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 "

enterprise/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-enterprise"
-version = "0.1.76"
+version = "0.1.77"
 description = "Package for LiteLLM Enterprise features"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.1.76"
+version = "0.1.77"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-enterprise==",

gateway/routes/allowlist.py

@@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
     # Realtime / streaming
     "/v1/realtime",
     "/realtime",
+    "/v1/live/sessions",
+    "/live/sessions",
     # Health & ops
     "/health",
     "/metrics",

helm/litellm-helm/README.md

@@ -126,6 +126,11 @@ Set `billingMetrics.caSecretName` only when the collector is a private or test o
 | `db.secret.name`          | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials.                                                                                                                                                                                               | `postgres`                                                                                 |
 | `db.secret.usernameKey`   | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance.                                                                                                                                    | `username`                                                                                 |
 | `db.secret.passwordKey`   | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user.                                                                                                                                                   | `password`                                                                                 |
+| `db.username`             | If `db.useExisting` is `true`, the database user as a plain value. When set, `db.secret.usernameKey` is not read.                                                                                                                                                                        | `""`                                                                                       |
+| `db.port`                 | If `db.useExisting` is `true` with token auth, the Postgres port used to build the connection URL. Password auth takes the port from `db.url`.                                                                                                                                           | `5432`                                                                                     |
+| `db.schema`               | With token auth, the Postgres schema holding the LiteLLM tables, rendered as `DATABASE_SCHEMA` when the tables are not in `public`. Ignored with password auth, where the schema stays in `db.url` as `?schema=`.                                                                        | `""`                                                                                       |
+| `db.useIAMAuth`           | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with an AWS RDS IAM token instead of a password. See [Token auth](#token-auth-to-an-existing-postgres-rds-iam-or-microsoft-entra-id).                                                                  | `false`                                                                                    |
+| `db.useAzureEntraAuth`    | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with a Microsoft Entra ID token instead of a password. Mutually exclusive with `db.useIAMAuth`.                                                                                                        | `false`                                                                                    |
 | `db.useStackgresOperator` | Not yet implemented.                                                                                                                                                                                                                                                                      | `false`                                                                                    |
 | `db.deployStandalone`     | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups.                                                                                                  | `true`                                                                                     |
 | `postgresql.*`            | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml)                                                           |
@@ -156,6 +161,22 @@ data:
 type: Opaque
 ```
 
+#### Token auth to an existing Postgres (RDS IAM or Microsoft Entra ID)
+
+When the database accepts short-lived tokens instead of a password (AWS RDS or Aurora with IAM database authentication, or Azure Database for PostgreSQL with Microsoft Entra ID), there is no password to put in a Secret. Turn on `db.useIAMAuth` or `db.useAzureEntraAuth` and the chart stops referencing `db.secret.passwordKey` and `db.url`. It sets `IAM_TOKEN_DB_AUTH` (or `AZURE_POSTGRESQL_AUTH`) and emits `DATABASE_HOST`, `DATABASE_PORT`, `DATABASE_USERNAME` and `DATABASE_NAME` for the proxy, the migrations job and the collector sidecar. Each of them mints its own token at startup and builds `DATABASE_URL` from those parts, and the in-container pgbouncer (`db.connectionPool`) refreshes the token before it expires.
+
+```yaml
+db:
+  useExisting: true
+  deployStandalone: false
+  useIAMAuth: true
+  endpoint: litellm.cluster-abc123.us-west-1.rds.amazonaws.com
+  database: litellm
+  username: litellm_iam   # the role granted rds_iam; no db.secret needed
+```
+
+The pod needs cloud credentials for the token: on EKS use IRSA or Pod Identity on the chart's `serviceAccount` (and `migrationJob.serviceAccountName` when the Helm hook runs before the ServiceAccount exists), on AKS use workload identity. Set `AWS_RDS_REGION` under `envVars` when the RDS hostname is not a standard regional endpoint (the token must be signed for the region the database lives in). TLS options that would otherwise go in `db.url` are set as `DATABASE_SSLMODE` and `DATABASE_SSLROOTCERT` under `envVars`. Leaving `db.username` empty reads the user from `db.secret.usernameKey` as before.
+
 #### Examples for `environmentSecrets` and `environemntConfigMaps`
 
 ```yaml

helm/litellm-helm/templates/_helpers.tpl

@@ -162,6 +162,77 @@ taken before the change, which by that point no longer exists.
 {{- end -}}
 {{- end -}}
 
+{{/*
+Whether the existing Postgres is reached with a short-lived token (AWS RDS IAM
+or Microsoft Entra ID) instead of a password.
+*/}}
+{{- define "litellm.db.tokenAuth" -}}
+{{- if or .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}true{{- end -}}
+{{- end -}}
+
+{{- define "litellm.validateDbTokenAuth" -}}
+{{- if and .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}
+{{- fail "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" -}}
+{{- end -}}
+{{- if and (include "litellm.db.tokenAuth" .) (or (not .Values.db.useExisting) .Values.db.deployStandalone) -}}
+{{- fail "db.useIAMAuth and db.useAzureEntraAuth need db.useExisting: true and db.deployStandalone: false (token auth targets an existing RDS or Azure Database for PostgreSQL; the bundled postgresql subchart only supports password auth)" -}}
+{{- end -}}
+{{- end -}}
+
+{{/*
+Database env for an existing Postgres, shared by the proxy container, the
+collector sidecar and the migrations Job. Under token auth every entrypoint
+mints its own token and assembles DATABASE_URL itself, so db.url is left out.
+*/}}
+{{- define "litellm.existingDbEnv" -}}
+{{- $tokenAuth := include "litellm.db.tokenAuth" . -}}
+- name: DATABASE_USERNAME
+  {{- if .Values.db.username }}
+  value: {{ .Values.db.username | quote }}
+  {{- else }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Values.db.secret.name }}
+      key: {{ .Values.db.secret.usernameKey }}
+  {{- end }}
+{{- if .Values.db.useIAMAuth }}
+- name: IAM_TOKEN_DB_AUTH
+  value: "true"
+{{- else if .Values.db.useAzureEntraAuth }}
+- name: AZURE_POSTGRESQL_AUTH
+  value: "true"
+{{- else }}
+- name: DATABASE_PASSWORD
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Values.db.secret.name }}
+      key: {{ .Values.db.secret.passwordKey }}
+{{- end }}
+- name: DATABASE_HOST
+  {{- if .Values.db.secret.endpointKey }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Values.db.secret.name }}
+      key: {{ .Values.db.secret.endpointKey }}
+  {{- else }}
+  value: {{ .Values.db.endpoint }}
+  {{- end }}
+{{- if $tokenAuth }}
+- name: DATABASE_PORT
+  value: {{ .Values.db.port | quote }}
+{{- end }}
+- name: DATABASE_NAME
+  value: {{ .Values.db.database }}
+{{- if and $tokenAuth .Values.db.schema }}
+- name: DATABASE_SCHEMA
+  value: {{ .Values.db.schema | quote }}
+{{- end }}
+{{- if not $tokenAuth }}
+- name: DATABASE_URL
+  value: {{ .Values.db.url | quote }}
+{{- end }}
+{{- end -}}
+
 {{/*
 Environment shared by the proxy container and the opt-in collector sidecar:
 database, pgbouncer, master key, redis, user envVars. Both containers must see
@@ -169,6 +240,7 @@ the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
 and the same spend transaction buffer.
 */}}
 {{- define "litellm.proxyEnv" -}}
+{{- include "litellm.validateDbTokenAuth" . -}}
 - name: HOST
   value: "{{ .Values.listen | default "0.0.0.0" }}"
 - name: PORT
@@ -189,29 +261,7 @@ and the same spend transaction buffer.
 - name: DATABASE_NAME
   value: litellm
 {{- else if .Values.db.useExisting }}
-- name: DATABASE_USERNAME
-  valueFrom:
-    secretKeyRef:
-      name: {{ .Values.db.secret.name }}
-      key: {{ .Values.db.secret.usernameKey }}
-- name: DATABASE_PASSWORD
-  valueFrom:
-    secretKeyRef:
-      name: {{ .Values.db.secret.name }}
-      key: {{ .Values.db.secret.passwordKey }}
-- name: DATABASE_HOST
-  {{- if .Values.db.secret.endpointKey }}
-  valueFrom:
-    secretKeyRef:
-      name: {{ .Values.db.secret.name }}
-      key: {{ .Values.db.secret.endpointKey }}
-  {{- else }}
-  value: {{ .Values.db.endpoint }}
-  {{- end }}
-- name: DATABASE_NAME
-  value: {{ .Values.db.database }}
-- name: DATABASE_URL
-  value: {{ .Values.db.url | quote }}
+{{ include "litellm.existingDbEnv" . }}
 {{- end }}
 {{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
 - name: DATABASE_READER_HOST

helm/litellm-helm/templates/migrations-job.yaml

@@ -68,30 +68,9 @@ spec:
           {{- end }}
           workingDir: "/app"
           env:
+            {{- include "litellm.validateDbTokenAuth" . -}}
             {{- if .Values.db.useExisting }}
-            - name: DATABASE_USERNAME
-              valueFrom:
-                secretKeyRef:
-                  name: {{ .Values.db.secret.name }}
-                  key: {{ .Values.db.secret.usernameKey }}
-            - name: DATABASE_PASSWORD
-              valueFrom:
-                secretKeyRef:
-                  name: {{ .Values.db.secret.name }}
-                  key: {{ .Values.db.secret.passwordKey }}
-            - name: DATABASE_HOST
-              {{- if .Values.db.secret.endpointKey }}
-              valueFrom:
-                secretKeyRef:
-                  name: {{ .Values.db.secret.name }}
-                  key: {{ .Values.db.secret.endpointKey }}
-              {{- else }}
-              value: {{ .Values.db.endpoint }}
-              {{- end }}
-            - name: DATABASE_NAME
-              value: {{ .Values.db.database }}
-            - name: DATABASE_URL
-              value: {{ .Values.db.url | quote }}
+            {{- include "litellm.existingDbEnv" . | nindent 12 }}
             {{- else if .Values.db.deployStandalone }}
             - name: DATABASE_URL
               value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }}

helm/litellm-helm/tests/database_auth_tests.yaml

@@ -0,0 +1,288 @@
+suite: test existing database token auth
+templates:
+  - deployment.yaml
+  - migrations-job.yaml
+  - configmap-litellm.yaml
+tests:
+  - it: should keep password auth from the secret when no token auth is set
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            valueFrom:
+              secretKeyRef:
+                name: postgres
+                key: username
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+            valueFrom:
+              secretKeyRef:
+                name: postgres
+                key: password
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+          any: true
+
+  - it: should render RDS IAM auth without a password secret on the proxy and the collector sidecar
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.username: litellm_iam
+      db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
+      db.port: 5433
+      collector.enabled: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            value: litellm_iam
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST
+            value: writer.cluster-abc.us-west-1.rds.amazonaws.com
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+            value: "5433"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_NAME
+            value: litellm
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+          any: true
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+
+  - it: should give the migrations job the same token auth env as the proxy
+    template: migrations-job.yaml
+    set:
+      migrationJob.enabled: true
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.username: litellm_iam
+      db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            value: litellm_iam
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST
+            value: writer.cluster-abc.us-west-1.rds.amazonaws.com
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+            value: "5432"
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+          any: true
+
+  - it: should carry db.schema as DATABASE_SCHEMA under token auth and omit it when empty
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.username: litellm_iam
+      db.schema: tenant1
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+            value: tenant1
+
+  - it: should leave the schema to db.url under password auth even when db.schema is set
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.schema: tenant1
+      db.url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+          any: true
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
+
+  - it: should not emit DATABASE_SCHEMA when db.schema is empty
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+          any: true
+
+  - it: should read the user from the secret when db.username is empty under token auth
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.secret.name: rds-identity
+      db.secret.usernameKey: user
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            valueFrom:
+              secretKeyRef:
+                name: rds-identity
+                key: user
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+
+  - it: should render Entra auth with AZURE_POSTGRESQL_AUTH and no password
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useAzureEntraAuth: true
+      db.username: litellm@example.onmicrosoft.com
+      db.endpoint: litellm.postgres.database.azure.com
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            value: litellm@example.onmicrosoft.com
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+          any: true
+
+  - it: should fail when both token sources are enabled
+    template: deployment.yaml
+    set:
+      migrationJob.enabled: false
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.useAzureEntraAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
+
+  - it: should fail when token auth is set while the bundled postgres is still deployed
+    template: deployment.yaml
+    set:
+      migrationJob.enabled: false
+      db.useExisting: true
+      db.useIAMAuth: true
+    asserts:
+      - failedTemplate:
+          errorPattern: "need db.useExisting: true and db.deployStandalone: false"
+
+  - it: should fail when token auth is set without useExisting
+    template: migrations-job.yaml
+    set:
+      migrationJob.enabled: true
+      db.deployStandalone: false
+      db.useAzureEntraAuth: true
+    asserts:
+      - failedTemplate:
+          errorPattern: "need db.useExisting: true and db.deployStandalone: false"

helm/litellm-helm/values.yaml

@@ -376,6 +376,30 @@ db:
   endpoint: localhost
   database: litellm
   url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
+  # Authenticate with a short-lived token instead of a password. useIAMAuth is
+  # AWS RDS / Aurora IAM database authentication (the pod needs AWS credentials,
+  # e.g. IRSA or EKS Pod Identity, and the database user needs GRANT rds_iam);
+  # useAzureEntraAuth is Azure Database for PostgreSQL with a Microsoft Entra ID
+  # identity (workload identity). Mutually exclusive, and both need
+  # useExisting: true with deployStandalone: false. The chart then sets
+  # IAM_TOKEN_DB_AUTH or AZURE_POSTGRESQL_AUTH, emits DATABASE_HOST, DATABASE_PORT,
+  # DATABASE_USERNAME and DATABASE_NAME and skips DATABASE_PASSWORD and db.url:
+  # the proxy, the migrations job and the collector sidecar each mint their own
+  # token and build DATABASE_URL from those parts. Set DATABASE_SSLMODE /
+  # DATABASE_SSLROOTCERT under envVars for TLS options you would otherwise put
+  # in db.url.
+  useIAMAuth: false
+  useAzureEntraAuth: false
+  # Postgres port, used to build the token-auth URL (password auth takes the
+  # port from db.url)
+  port: 5432
+  # Postgres schema for the LiteLLM tables, rendered as DATABASE_SCHEMA under
+  # token auth only (password auth carries it as ?schema= in db.url)
+  schema: ""
+  # Database user as a plain value. Empty means read it from
+  # db.secret.usernameKey. With token auth the user name is the only credential
+  # the chart needs, so setting it here removes the need for db.secret entirely.
+  username: ""
   secret:
     name: postgres
     usernameKey: username
@@ -440,10 +464,10 @@ db:
 # Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
 # the pod NotReady and blocks only until in-flight requests actually finish
 # (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
-# waiting the worst-case duration. The drain runs once (the preStop hook and
-# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
-# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
-# SIGKILL.
+# waiting the worst-case duration. After SIGTERM the server waits up to
+# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels
+# them, so set terminationGracePeriodSeconds a few seconds above twice
+# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL.
 #
 # /health/drain is off by default; enable it with
 # general_settings.enable_drain_endpoint: true. The kubelet calls preStop

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,32 @@ 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 }}
+{{- $sslCertReqs := $root.Values.redis.sslCertReqs }}
+{{- if and $sslCertReqs (not (has $sslCertReqs (list "none" "optional" "required"))) }}
+{{- fail (printf "redis.sslCertReqs must be one of none, optional, required (got %q)" $sslCertReqs) }}
+{{- end }}
+{{- if and $urlKey (or $root.Values.redis.ssl $sslCertReqs $root.Values.redis.sslCaCerts) }}
+{{- fail "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead" }}
+{{- end }}
+{{- if or $root.Values.redis.host $hostKey }}
+{{- if $root.Values.redis.ssl }}
+- name: REDIS_SSL
+  value: "true"
+{{- end }}
+{{- if $sslCertReqs }}
+- name: REDIS_SSL_CERT_REQS
+  value: {{ $sslCertReqs | quote }}
+{{- end }}
+{{- if $root.Values.redis.sslCaCerts }}
+- name: REDIS_SSL_CA_CERTS
+  value: {{ $root.Values.redis.sslCaCerts | quote }}
 {{- 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

@@ -3,6 +3,7 @@ templates:
   - gateway/deployment.yaml
   - gateway/configmap.yaml
   - backend/deployment.yaml
+  - migrations-job.yaml
 values:
   - ./values/required.yaml
 tests:
@@ -21,6 +22,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 +90,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 +150,339 @@ 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"
+
+  - it: redis omits TLS env vars by default
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+          any: true
+
+  - it: redis ssl false omits REDIS_SSL
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: false
+      redis.sslCertReqs: required
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+          any: true
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "required"
+
+  - it: gateway renders redis TLS env vars
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: none
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "none"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+            value: "/etc/redis-tls/ca.crt"
+
+  - it: backend renders redis TLS env vars
+    template: backend/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: required
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "required"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+            value: "/etc/redis-tls/ca.crt"
+
+  - it: migration job renders redis TLS env vars
+    template: migrations-job.yaml
+    set:
+      db.deployStandalone: false
+      db.useExisting: true
+      db.endpoint: postgres.example.com
+      db.secret.name: pg-secret
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: optional
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "optional"
+
+  - it: redis TLS values reject a secret-backed URL
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.urlKey: connection-url
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead"
+
+  - it: hostKey mode renders redis TLS env vars
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.hostKey: hostname
+      redis.ssl: true
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+            value: "/etc/redis-tls/ca.crt"
+
+  - it: redis TLS env vars are omitted when redis is not configured
+    template: gateway/deployment.yaml
+    set:
+      redis.ssl: true
+      redis.sslCertReqs: none
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+          any: true
+
+  - it: redis sslCertReqs rejects unknown values
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: strict
+    asserts:
+      - failedTemplate:
+          errorMessage: 'redis.sslCertReqs must be one of none, optional, required (got "strict")'

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`
@@ -267,13 +275,26 @@ database:
 # self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from
 # `host` / `port` as the single seed; the cluster client discovers the
 # remaining nodes from CLUSTER SLOTS at startup.
+#
+# For a TLS Redis endpoint set `ssl: true`, which the chart emits as REDIS_SSL.
+# `sslCertReqs` (none / optional / required) becomes REDIS_SSL_CERT_REQS and
+# `sslCaCerts` (path to a CA bundle mounted in the pod, e.g. via
+# `extraVolumeMounts`) becomes REDIS_SSL_CA_CERTS. All three are omitted when
+# unset. They apply to `host` / `hostKey` mode only. With `passwordSecret.urlKey`
+# the proxy reads TLS settings from the URL itself, so use `rediss://` and the
+# `ssl_cert_reqs` / `ssl_ca_certs` query parameters there instead.
 redis:
   cluster: false
   host: ""
   port: 6379
+  ssl: false
+  sslCertReqs: ""
+  sslCaCerts: ""
   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 +514,7 @@ gateway:
 backend:
   enabled: true
   logLevel: INFO
+  gatewayUrl: ""
   extraEnv: []
   envConfigMaps: []
   envSecrets: []
@@ -575,9 +597,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/_redaction.py

@@ -0,0 +1,50 @@
+import functools
+import os
+import re
+from collections.abc import Callable
+from typing import Final
+from urllib.parse import unquote, urlsplit
+
+_REDACTED: Final = "REDACTED"
+_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
+
+
+@functools.cache
+def _secret_shape_redactor() -> Callable[[str], str]:
+    try:
+        from litellm._logging import redact_secrets
+    except ImportError:
+        return lambda text: text
+    return redact_secrets
+
+
+def _url_passwords(url: str) -> frozenset[str]:
+    try:
+        parts: Final = urlsplit(url)
+    except ValueError:
+        return frozenset()
+    query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
+    raw_query_passwords: Final = tuple(
+        value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
+    )
+    raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
+    return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
+
+
+def _configured_database_passwords() -> frozenset[str]:
+    database_url: Final = os.getenv("DATABASE_URL")
+    direct_url: Final = os.getenv("DIRECT_URL")
+    database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
+    direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
+    return database_passwords | direct_passwords
+
+
+def _redact_credentials(text: str) -> str:
+    """Mask configured database passwords before passing the text to LiteLLM redaction."""
+    passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
+    alternation: Final = "|".join(re.escape(password) for password in passwords)
+    password_pattern: Final = (
+        re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
+    )
+    result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
+    return _secret_shape_redactor()(result)

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -8,14 +8,34 @@ from typing import TYPE_CHECKING, Final
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras._redaction import _redact_credentials
 from litellm_proxy_extras.prisma_toolchain import (
     MIGRATION_LOCK_TIMEOUT_ENV_VAR,
-    migration_ddl_lock_timeout,
     migration_lock_timeout,
 )
 
 MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
-_LOCK_TIMEOUT_PINNED_RE: Final = re.compile(r"(?:-c\s*|--)lock_timeout=")
+PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
+_SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
+
+
+@dataclass(frozen=True, slots=True)
+class AdvisoryLock:
+    key: int
+    name: str
+    holder: str
+
+
+COORDINATOR_LOCK: Final = AdvisoryLock(
+    MIGRATION_LOCK_KEY,
+    "v2 migration coordinator lock",
+    "another v2 migration resolver",
+)
+PRISMA_LOCK: Final = AdvisoryLock(
+    PRISMA_ADVISORY_LOCK_KEY,
+    "Prisma migration lock",
+    "another Prisma migration",
+)
 
 if TYPE_CHECKING:
     import psycopg
@@ -30,29 +50,12 @@ def _migration_url(database_url: str, direct_url: "str | None") -> str:
     return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
 
 
-def _with_ddl_lock_timeout(url: str) -> str:
-    split: Final = urlsplit(url)
-    pairs: Final = tuple(parse_qsl(split.query))
-    if any(key == "pgbouncer" and value == "true" for key, value in pairs):
-        return url
-    existing: Final = next((value for key, value in pairs if key == "options"), "")
-    if _LOCK_TIMEOUT_PINNED_RE.search(existing):
-        return url
-    timeout_ms: Final = max(1, int(migration_ddl_lock_timeout() * 1000))
-    options: Final = f"{existing} -c lock_timeout={timeout_ms}".strip()
-    rewritten: Final = tuple((key, value) for key, value in pairs if key != "options") + (("options", options),)
-    return urlunsplit(split._replace(query=urlencode(rewritten)))
-
-
 def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
     database_url: Final = environment.get("DATABASE_URL")
     if not database_url:
         return environment
     migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
-    return {
-        **environment,
-        "DATABASE_URL": _with_ddl_lock_timeout(migration_url),
-    }
+    return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
 
 
 @dataclass(frozen=True, slots=True)
@@ -68,6 +71,28 @@ def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = M
     return row is not None and row.acquired
 
 
+def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting: bool) -> str:
+    import psycopg
+
+    prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
+    raw: Final = str(exc)
+    unquoted: Final = re.sub(r"\"[^\"]*\"|'[^']*'", "", raw)
+    detail: Final = _redact_credentials(raw).strip()
+    if (
+        not connecting
+        or not isinstance(exc, psycopg.OperationalError)
+        or any(marker in unquoted for marker in _SERVER_ANSWERED_MARKERS)
+    ):
+        return prefix + detail
+    parsed_url: Final = urlsplit(database_url)
+    try:
+        port: Final = parsed_url.port
+    except ValueError:
+        return prefix + f"P1001: Can't reach database server. {detail}"
+    host: Final = parsed_url.hostname or "localhost"
+    return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
+
+
 @dataclass(frozen=True, slots=True)
 class MigrationCoordinator:
     connection: "psycopg.Connection[tuple[object, ...]]"
@@ -75,10 +100,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(
@@ -89,32 +119,49 @@ class MigrationCoordinator:
 
 
 @contextmanager
-def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
+def migration_lock(
+    database_url: str,
+    lock: AdvisoryLock = COORDINATOR_LOCK,
+) -> Generator[MigrationCoordinator, None, None]:
     import psycopg
 
     wait_seconds: Final = migration_lock_timeout()
     deadline: Final = time.monotonic() + wait_seconds
     try:
-        with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
+        connection: Final = psycopg.connect(database_url, connect_timeout=10, autocommit=True)
+    except psycopg.Error as exc:
+        raise RuntimeError(_coordination_failure(database_url, exc, connecting=True)) from None
+    try:
+        with connection:
             coordinator: Final = MigrationCoordinator(connection)
-            logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
+            logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds)
             while time.monotonic() < deadline:
                 with connection.transaction():
-                    if _try_lock(connection):
-                        logger.info("Acquired the v2 migration coordinator lock")
+                    if _try_lock(connection, lock.key):
+                        logger.info("Acquired the %s", lock.name)
 
                         yield coordinator
                         coordinator.check_connection()
                         return
                 time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
     except psycopg.Error as exc:
-        raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
+        raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
     raise RuntimeError(
-        f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
+        f"Timed out waiting for {lock.holder} after {wait_seconds}s. "
         f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
     )
 
 
+@contextmanager
+def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
+    """Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled.
+
+    This keeps deploy independent of the coordinator key held by request-log index builds."""
+    with migration_lock(database_url, PRISMA_LOCK) as coordinator:
+        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/20261009000100_telemetry_reports/migration.sql

@@ -0,0 +1,14 @@
+SET LOCAL lock_timeout = '10s';
+
+CREATE TABLE IF NOT EXISTS "LiteLLM_TelemetryReport" (
+    "id" TEXT NOT NULL,
+    "window_start" DOUBLE PRECISION NOT NULL,
+    "window_end" DOUBLE PRECISION NOT NULL,
+    "report" JSONB NOT NULL,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    CONSTRAINT "LiteLLM_TelemetryReport_pkey" PRIMARY KEY ("id")
+);
+
+CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_window_end_idx" ON "LiteLLM_TelemetryReport"("window_end");
+
+CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_created_at_idx" ON "LiteLLM_TelemetryReport"("created_at");

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/migrations/20261015000000_managed_object_end_user_id/migration.sql

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

litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py

@@ -60,13 +60,11 @@ PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
 PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
 PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
 MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
-MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_DDL_LOCK_TIMEOUT"
 NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
 
 DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
 DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
 DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
-DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT = 10.0
 
 BOOTSTRAP_ARG = "--version"
 PRISMA_CONSOLE_SCRIPT = "prisma"
@@ -113,11 +111,6 @@ def migration_lock_timeout() -> float:
     return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
 
 
-def migration_ddl_lock_timeout() -> float:
-    """Seconds one migration statement may wait on a table lock before Postgres cancels it."""
-    return _timeout_from_env(MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR, DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT)
-
-
 def prisma_bootstrap_timeout() -> float:
     """Seconds the one-time Node toolchain install may run for."""
     return _timeout_from_env(

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
   api_key String?
   request_tags Json? @default("[]")
+  end_user_id String?
   updated_at DateTime @updatedAt
   updated_by String?
 
@@ -1588,7 +1589,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([])
@@ -2045,3 +2048,14 @@ model LiteLLM_LensTraceSignal {
 
   @@id([trace_id, trace_ref])
 }
+
+model LiteLLM_TelemetryReport {
+  id           String   @id
+  window_start Float
+  window_end   Float
+  report       Json
+  created_at   DateTime @default(now())
+
+  @@index([window_end])
+  @@index([created_at])
+}

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -1,4 +1,3 @@
-import functools
 import glob
 import os
 import random
@@ -9,16 +8,16 @@ 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
-from urllib.parse import unquote, urlsplit
 
 from litellm_proxy_extras import prisma_toolchain
 from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras._redaction import _redact_credentials
 from litellm_proxy_extras.migration_lock import held_migration_lock
 from litellm_proxy_extras.prisma_toolchain import (
-    MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR,
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
     PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
     ensure_prisma_toolchain,
@@ -212,51 +211,6 @@ def _max_migration_timestamp(names) -> int:
     return max(_migration_timestamp(n) for n in names)
 
 
-_REDACTED: Final = "REDACTED"
-_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
-
-
-@functools.cache
-def _secret_shape_redactor() -> Callable[[str], str]:
-    try:
-        from litellm._logging import redact_secrets
-    except ImportError:
-        return lambda text: text
-    return redact_secrets
-
-
-def _url_passwords(url: str) -> frozenset[str]:
-    try:
-        parts: Final = urlsplit(url)
-    except ValueError:
-        return frozenset()
-    query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
-    raw_query_passwords: Final = tuple(
-        value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
-    )
-    raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
-    return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
-
-
-def _configured_database_passwords() -> frozenset[str]:
-    database_url: Final = os.getenv("DATABASE_URL")
-    direct_url: Final = os.getenv("DIRECT_URL")
-    database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
-    direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
-    return database_passwords | direct_passwords
-
-
-def _redact_credentials(text: str) -> str:
-    """Mask configured database passwords before passing the text to LiteLLM redaction."""
-    passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
-    alternation: Final = "|".join(re.escape(password) for password in passwords)
-    password_pattern: Final = (
-        re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
-    )
-    result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
-    return _secret_shape_redactor()(result)
-
-
 def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
     if isinstance(command, tuple):
         return tuple(_redact_credentials(str(argument)) for argument in command)
@@ -789,16 +743,23 @@ class ProxyExtrasDBManager:
         (connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
         translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
         client certificate where Prisma reads it as the CA."""
-        from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
+        from urllib.parse import parse_qsl, quote, unquote, urlencode, urlparse, urlunparse
 
         parsed: Final = urlparse(url)
+        netloc: Final = parsed.netloc
+        userinfo, separator, hostport = netloc.rpartition("@")
+        user, password_separator, password = userinfo.partition(":")
+        canonical_userinfo: Final = quote(unquote(user), safe="") + (
+            password_separator + quote(unquote(password), safe="") if password_separator else ""
+        )
+        canonical_netloc: Final = f"{canonical_userinfo}{separator}{hostport}" if separator else netloc
         if not parsed.query:
-            return url
+            return url if canonical_netloc == netloc else urlunparse(parsed._replace(netloc=canonical_netloc))
         pairs: Final = tuple(parse_qsl(parsed.query))
         kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
         sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
         libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
-        return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
+        return urlunparse(parsed._replace(netloc=canonical_netloc, query=urlencode(libpq_pairs, quote_via=quote)))
 
     @staticmethod
     def _libpq_tls_params(
@@ -1022,7 +983,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,
@@ -1056,7 +1021,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
 
@@ -1065,6 +1032,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).
@@ -1127,11 +1095,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
 
@@ -1156,6 +1125,7 @@ class ProxyExtrasDBManager:
                         e,
                         budget,
                         time.monotonic() - attempt_started,
+                        deploy_lock=deploy_lock,
                     )
 
                 if next_budget.attempts_left < budget.attempts_left:
@@ -1165,10 +1135,9 @@ class ProxyExtrasDBManager:
             raise RuntimeError(
                 f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
                 "attempts that made no progress (timeouts, lock timeouts or deadlock retries). "
-                "Check database connectivity, load, long-running transactions on the migrated tables, "
-                "and _prisma_migrations ledger state. Raise "
-                f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out, or "
-                f"{MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if they timed out waiting for a table lock."
+                "Check database connectivity, load and _prisma_migrations ledger state. "
+                "If a migration hit SET LOCAL lock_timeout, check for long-running transactions on the migrated tables. "
+                f"Raise {PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
             )
         finally:
             os.chdir(original_dir)
@@ -1189,15 +1158,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
 
@@ -1206,6 +1179,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.
 
@@ -1227,14 +1201,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(
@@ -1278,17 +1252,17 @@ 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:
                 ProxyExtrasDBManager._log_migration_lock_holders(migration_name)
                 logger.warning(
-                    "Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying. "
-                    f"Raise {MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if the database needs longer.",
+                    "Migration %s timed out waiting for a table lock after hitting SET LOCAL lock_timeout, rolling "
+                    "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-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-proxy-extras"
-version = "0.4.108"
+version = "0.4.109"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 readme = "README.md"
 requires-python = ">=3.9"
@@ -30,7 +30,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.4.108"
+version = "0.4.109"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-proxy-extras==",

litellm-rust/AGENTS.md

@@ -15,6 +15,7 @@ Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `st
 - `crates/<crate>/tests/` is only for tests that drive several modules together end to end, such as a full request through a route. A test of one type or function never goes there
 - A helper shared by several inline test modules goes in a `#[cfg(test)] mod test_support;` at the crate root, and only once a second module needs it
 - A test for another crate's item belongs in that crate, not in a downstream one
+- Tests that call a real provider go in `crates/<crate>/tests/live/<provider>.rs` behind a `tests/live/main.rs`. Each reads the provider's key from the environment through an `Option` fixture and returns early when it is unset, never `#[ignore]`, so `cargo test` stays offline until the key is exported
 - Never set `autotests = false` or hand-list `[[test]]` targets. Every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests/<name>/mod.rs` or `tests/<subject>/support.rs` so it is not picked up as a test crate of its own
 
 ## Test fixtures and cases

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",
 ]
 
@@ -4085,6 +4089,7 @@ dependencies = [
  "litellm-cache",
  "litellm-cache-memory",
  "litellm-cache-response",
+ "litellm-core-utils",
  "litellm-host",
  "litellm-host-native",
  "litellm-http",

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/framer/src/sse.rs

@@ -1,6 +1,6 @@
 use std::str;
 
-use bytes::{Buf, BufMut, BytesMut};
+use bytes::{Buf, BufMut, Bytes, BytesMut};
 use tokio_util::codec::{Decoder, Encoder};
 
 use crate::SseError;
@@ -29,38 +29,107 @@ impl Decoder for SseCodec {
             return Ok(None);
         }
         while let Some(end) = block_end(src) {
-            let block = src.split_to(end);
-            let pending = lines(&block)
-                .map(|(line, _)| line)
-                .take_while(|line| !line.is_empty())
-                .try_fold(Pending::default(), Pending::apply)?;
-            if let Some(event) = pending.dispatch() {
+            if let Some(event) = event_in(&src.split_to(end))? {
                 return Ok(Some(event));
             }
         }
         Ok(None)
     }
 
-    fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
-        Ok(None)
+    /// An event the stream ended inside of is dropped, as the SSE spec requires.
+    fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
+        self.decode(pending)
     }
 }
 
+pub fn decode_all(wire: &[u8]) -> Result<Vec<SseEvent>, SseError> {
+    let mut codec = SseCodec::default();
+    let mut buffer = BytesMut::from(wire);
+    let mut events = Vec::new();
+    while let Some(event) = codec.decode_eof(&mut buffer)? {
+        events.push(event);
+    }
+    Ok(events)
+}
+
 impl SseCodec {
     fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
-        if self.past_bom {
-            return true;
-        }
-        if src.starts_with(BOM) {
-            src.advance(BOM.len());
-        } else if BOM.starts_with(src) {
-            return false;
-        }
-        self.past_bom = true;
-        true
+        skip_bom(&mut self.past_bom, src)
     }
 }
 
+fn event_in(block: &[u8]) -> Result<Option<SseEvent>, SseError> {
+    Ok(lines(block)
+        .map(|(line, _)| line)
+        .take_while(|line| !line.is_empty())
+        .try_fold(Pending::default(), Pending::apply)?
+        .dispatch())
+}
+
+/// One event together with the exact wire bytes that carried it, for relays that forward a
+/// provider's bytes untouched but still need to see each event go by.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct SseBlock {
+    pub bytes: Bytes,
+    pub event: SseEvent,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct SseBlockCodec {
+    past_bom: bool,
+}
+
+impl Decoder for SseBlockCodec {
+    type Item = SseBlock;
+    type Error = SseError;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
+        self.decode_block(src, false)
+    }
+
+    fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
+        self.decode_block(pending, true)
+    }
+}
+
+impl SseBlockCodec {
+    fn decode_block(
+        &mut self,
+        src: &mut BytesMut,
+        eof: bool,
+    ) -> Result<Option<SseBlock>, SseError> {
+        if !skip_bom(&mut self.past_bom, src) {
+            return Ok(None);
+        }
+        while let Some(end) = block_end(src) {
+            if !eof && end == src.len() && src.ends_with(b"\r") {
+                return Ok(None);
+            }
+            let block = src.split_to(end);
+            if let Some(event) = event_in(&block)? {
+                return Ok(Some(SseBlock {
+                    bytes: block.freeze(),
+                    event,
+                }));
+            }
+        }
+        Ok(None)
+    }
+}
+
+fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool {
+    if *past_bom {
+        return true;
+    }
+    if src.starts_with(BOM) {
+        src.advance(BOM.len());
+    } else if BOM.starts_with(src) {
+        return false;
+    }
+    *past_bom = true;
+    true
+}
+
 fn block_end(bytes: &[u8]) -> Option<usize> {
     lines(bytes)
         .find(|(line, _)| line.is_empty())
@@ -168,3 +237,179 @@ impl Encoder<SseEvent> for SseCodec {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use futures_util::{StreamExt, TryStreamExt, stream};
+    use rstest::rstest;
+
+    use super::*;
+    use crate::frames;
+
+    fn pieces(wire: &[u8], size: usize) -> Vec<Bytes> {
+        wire.chunks(size).map(Bytes::copy_from_slice).collect()
+    }
+
+    async fn blocks(wire: &[u8], size: usize) -> Result<Vec<SseBlock>, SseError> {
+        frames(
+            stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)),
+            SseBlockCodec::default(),
+        )
+        .try_collect()
+        .await
+    }
+
+    #[rstest]
+    #[case::lf(
+        b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n",
+        b"event: delta\ndata: a\ndata: b\n\n"
+    )]
+    #[case::crlf(
+        b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n",
+        b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n"
+    )]
+    #[case::cr(
+        b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r",
+        b"event: delta\rdata: a\rdata: b\r\r"
+    )]
+    #[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")]
+    #[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")]
+    #[case::extensions(
+        b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n",
+        b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n"
+    )]
+    #[tokio::test]
+    async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch(
+        #[case] wire: &[u8],
+        #[case] expected: &[u8],
+        #[values(1, 3, 4096)] size: usize,
+    ) {
+        let blocks = blocks(wire, size).await.unwrap();
+        let expected_events = decode_all(expected).unwrap();
+        assert_eq!(expected_events.len(), 1);
+        assert_eq!(
+            blocks,
+            vec![SseBlock {
+                bytes: Bytes::copy_from_slice(expected),
+                event: expected_events[0].clone(),
+            }]
+        );
+        assert_eq!(expected_events, decode_all(wire).unwrap());
+    }
+
+    #[rstest]
+    #[case::empty(b"", vec![])]
+    #[case::unterminated(b"data: partial\n", vec![])]
+    #[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])]
+    #[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])]
+    #[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])]
+    #[tokio::test]
+    async fn blocks_dispatch_complete_events_and_drop_a_cut_tail(
+        #[case] wire: &[u8],
+        #[case] expected: Vec<&[u8]>,
+    ) {
+        let mut blocks = Box::pin(frames(
+            stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)),
+            SseBlockCodec::default(),
+        ));
+        for expected in expected {
+            assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected);
+        }
+        assert!(blocks.next().await.is_none());
+    }
+
+    #[rstest]
+    fn bulk_decode_preserves_every_field() {
+        let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap();
+        assert_eq!(
+            events,
+            vec![SseEvent {
+                event: Some("delta".into()),
+                data: "a\nb".into(),
+                id: Some("1".into()),
+                retry: Some(20),
+            }]
+        );
+    }
+
+    #[rstest]
+    #[case::only_a_tail(b"data: partial\n", vec![])]
+    #[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])]
+    #[case::no_terminator_at_all(b"data: partial", vec![])]
+    fn bulk_decode_drops_the_event_a_stream_ends_inside(
+        #[case] wire: &[u8],
+        #[case] expected: Vec<&str>,
+    ) {
+        let data: Vec<String> = decode_all(wire)
+            .unwrap()
+            .into_iter()
+            .map(|e| e.data)
+            .collect();
+        assert_eq!(data, expected);
+    }
+
+    #[rstest]
+    #[case::trailing_blank_lines(b"data: x\n\n\n\n")]
+    #[case::trailing_comment(b"data: x\n\n: bye\n\n")]
+    #[case::lone_cr_at_eof(b"data: x\r\r")]
+    fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) {
+        assert_eq!(decode_all(wire).unwrap().len(), 1);
+    }
+
+    #[rstest]
+    fn bulk_decode_rejects_invalid_utf8() {
+        assert!(matches!(
+            decode_all(b"data: \xff\n\n"),
+            Err(SseError::InvalidUtf8(_))
+        ));
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn block_framing_preserves_transport_errors() {
+        let input = stream::iter([
+            Ok(Bytes::from_static(b"data: first\n\n")),
+            Err(std::io::Error::new(
+                std::io::ErrorKind::ConnectionReset,
+                "reset",
+            )),
+        ]);
+        let mut blocks = Box::pin(frames(input, SseBlockCodec::default()));
+        assert_eq!(
+            blocks.next().await.unwrap().unwrap().bytes,
+            Bytes::from_static(b"data: first\n\n")
+        );
+        let Some(Err(SseError::Body(error))) = blocks.next().await else {
+            panic!("expected body error");
+        };
+        let cause = error
+            .get_ref()
+            .unwrap()
+            .downcast_ref::<std::io::Error>()
+            .unwrap();
+        assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
+        assert!(blocks.next().await.is_none());
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn invalid_utf8_stops_block_framing_after_prior_events() {
+        let mut blocks = Box::pin(frames(
+            stream::iter(
+                pieces(b"data: first\n\ndata: \xff\n\n", 1)
+                    .into_iter()
+                    .map(Ok::<_, std::io::Error>),
+            ),
+            SseBlockCodec::default(),
+        ));
+        assert_eq!(
+            blocks.next().await.unwrap().unwrap().bytes,
+            Bytes::from_static(b"data: first\n\n")
+        );
+        assert!(matches!(
+            blocks.next().await,
+            Some(Err(SseError::InvalidUtf8(_)))
+        ));
+        assert!(blocks.next().await.is_none());
+    }
+}

litellm-rust/crates/framer/tests/sse.rs

@@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected:
 }
 
 #[rstest]
-#[case::unterminated_single(b"data: partial\n", vec![])]
-#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
 #[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
-#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
+#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])]
+#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])]
 #[tokio::test]
-async fn eof_dispatches_only_terminated_events(
+async fn eof_between_events_ends_the_stream_cleanly(
     #[case] wire: &[u8],
     #[case] expected: Vec<SseEvent>,
 ) {
@@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events(
     );
 }
 
+#[rstest]
+#[case::unterminated_single(b"data: partial\n", vec![])]
+#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
+#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
+#[case::partial_bom(b"\xEF\xBB", vec![])]
+#[tokio::test]
+async fn eof_inside_an_event_drops_it_after_the_terminated_events(
+    #[case] wire: &[u8],
+    #[case] expected: Vec<SseEvent>,
+) {
+    assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
+}
+
 #[rstest]
 #[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
 #[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]

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/inference-messages/src/handler.rs

@@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError;
 use litellm_llms::base_llm::{
     auth::{Authenticated, resolve_auth},
     messages::{
-        streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
+        streaming::{self, ByteStream, StreamDecoder},
         transformation::BaseMessagesConfig,
     },
 };
@@ -245,7 +245,7 @@ fn decoded_chunks(
         let Some(event) = events.try_next().await? else {
             return Ok(None);
         };
-        let chunk = encode_anthropic_sse(&event)?;
+        let chunk = streaming::encode(&event)?;
         log_chunk(provider, "client_response", &chunk);
         Ok(Some((chunk, events)))
     })
@@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
 
 #[cfg(test)]
 mod tests {
-    use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream;
     use rstest::rstest;
     use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
 
@@ -300,7 +299,7 @@ mod tests {
             .await
             .unwrap();
         let MessagesCallResponse::Stream { mut chunks, .. } =
-            streaming_response(response, Some(anthropic_sse_event_stream), "test")
+            streaming_response(response, Some(streaming::event_stream), "test")
         else {
             panic!("a streaming response returns chunks");
         };

litellm-rust/crates/inference-messages/src/prepare.rs

@@ -124,7 +124,7 @@ fn prepare_provider_request(
         url,
         body: transformed,
         environment,
-        timeout,
+        timeout: timeout.or_else(|| shaping.settings.request_timeout()),
         api_key: api_key.map(SecretValue::new),
     })
 }
@@ -253,6 +253,41 @@ mod tests {
         );
     }
 
+    #[rstest]
+    #[case::nothing_configured(None, None, None)]
+    #[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
+    #[case::call_timeout_wins(
+        Some(Duration::from_secs(30)),
+        Some(1800.0),
+        Some(Duration::from_secs(30))
+    )]
+    #[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
+    #[case::unusable_setting_ignored(None, Some(-1.0), None)]
+    #[case::zero_setting_ignored(None, Some(0.0), None)]
+    fn timeout_falls_back_to_the_configured_request_timeout(
+        shaping: MessagesShaping,
+        #[case] call_timeout: Option<Duration>,
+        #[case] request_timeout: Option<f64>,
+        #[case] expected: Option<Duration>,
+    ) {
+        let prepared = prepare(MessagesCall {
+            body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
+            api_key: Some("sk-test".into()),
+            api_base: Some("https://anthropic.test".into()),
+            custom_llm_provider: Some("anthropic".into()),
+            litellm_params: Default::default(),
+            extra_headers: None,
+            provider_specific_header: None,
+            timeout: call_timeout,
+            shaping: MessagesShaping {
+                settings: MessagesSettings { request_timeout, ..shaping.settings },
+                ..shaping
+            },
+        })
+        .unwrap();
+        assert_eq!(prepared.timeout, expected);
+    }
+
     fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
         prepare(MessagesCall {
             body: body(fields),

litellm-rust/crates/inference-messages/src/types.rs

@@ -60,6 +60,16 @@ pub struct MessagesSettings {
     pub reasoning_auto_summary: bool,
     #[serde(default)]
     pub additional_drop_params: Vec<String>,
+    #[serde(default)]
+    pub request_timeout: Option<f64>,
+}
+
+impl MessagesSettings {
+    pub fn request_timeout(&self) -> Option<Duration> {
+        self.request_timeout
+            .and_then(|secs| Duration::try_from_secs_f64(secs).ok())
+            .filter(|timeout| !timeout.is_zero())
+    }
 }
 
 #[cfg(test)]
@@ -120,13 +130,15 @@ mod tests {
             },
             "drop_params": true,
             "reasoning_auto_summary": true,
-            "additional_drop_params": ["metadata.user_id", "thinking"]
+            "additional_drop_params": ["metadata.user_id", "thinking"],
+            "request_timeout": 1800.0
         }),
         MessagesShaping {
             settings: MessagesSettings {
                 drop_params: true,
                 reasoning_auto_summary: true,
                 additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
+                request_timeout: Some(1800.0),
             },
             capabilities: MessagesModelCapabilities {
                 supports_reasoning: true,

litellm-rust/crates/inference-responses/Cargo.toml

@@ -10,6 +10,7 @@ bytes.workspace = true
 futures-util.workspace = true
 litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
 litellm-cache-response.workspace = true
+litellm-core-utils.workspace = true
 litellm-host.workspace = true
 litellm-http.workspace = true
 litellm-inference.workspace = true

litellm-rust/crates/inference-responses/src/prepare.rs

@@ -1,7 +1,11 @@
-use litellm_host::interceptors::RequestContext;
+use litellm_core_utils::get_llm_provider_logic::{
+    CustomLlmProvider, LlmProviders, get_custom_llm_provider,
+};
+use litellm_host::interceptors::{ProviderIdentity, RequestContext};
 use litellm_llms::{
     base_llm::responses::transformation::BaseResponsesApiConfig,
     openai::responses::transformation::OpenAiResponsesApiConfig,
+    openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
 };
 use litellm_secrets::source::SecretSource;
 use serde_json::Value;
@@ -16,10 +20,10 @@ pub(super) async fn prepare(
     call: ResponsesCall,
     secrets: &dyn SecretSource,
 ) -> Result<ProviderResponsesRequest, Error> {
-    let identity = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
+    let ResponsesProvider { identity, config } =
+        resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
     let provider = identity.provider.as_str();
     let model = identity.model.as_str();
-    let config: &'static dyn BaseResponsesApiConfig = &OpenAiResponsesApiConfig;
     let snapshot = secrets
         .resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
         .await?;
@@ -52,20 +56,107 @@ pub(super) async fn prepare(
     })
 }
 
+pub(super) struct ResponsesProvider {
+    pub identity: ProviderIdentity,
+    pub config: &'static dyn BaseResponsesApiConfig,
+}
+
+/// Python's `get_llm_provider`: a bare model is OpenAI's, and with an explicit OpenRouter
+/// provider a model without a vendor segment keeps the `openrouter/` prefix as its wire id.
 pub(super) fn resolve_provider(
     model: &str,
     custom_llm_provider: Option<&str>,
-) -> Result<litellm_host::interceptors::ProviderIdentity, Error> {
-    let provider = custom_llm_provider.unwrap_or("openai");
-    if provider != "openai" {
-        return Err(Error::Unsupported("native HTTP responses provider"));
-    }
-    let resolved = model.strip_prefix("openai/").unwrap_or(model);
-    if resolved.is_empty() || resolved.contains('/') {
+) -> Result<ResponsesProvider, Error> {
+    let explicit = custom_llm_provider.filter(|provider| !provider.is_empty());
+    let (provider, resolved) = match get_custom_llm_provider(model, explicit) {
+        Some(CustomLlmProvider {
+            model,
+            custom_llm_provider,
+        }) => (custom_llm_provider, model),
+        None => ("openai", model),
+    };
+    if resolved.is_empty() {
         return Err(Error::InvalidProvider(model.into()));
     }
-    Ok(litellm_host::interceptors::ProviderIdentity {
-        model: resolved.into(),
-        provider: provider.into(),
+    let (config, wire_model): (&'static dyn BaseResponsesApiConfig, String) =
+        match provider.parse::<LlmProviders>() {
+            Ok(LlmProviders::Openai) if resolved.contains('/') => {
+                return Err(Error::InvalidProvider(model.into()));
+            }
+            Ok(LlmProviders::Openai) => (&OpenAiResponsesApiConfig, resolved.into()),
+            Ok(LlmProviders::Openrouter) if explicit.is_some() && !resolved.contains('/') => (
+                &OPENROUTER_RESPONSES_CONFIG,
+                format!("openrouter/{resolved}"),
+            ),
+            Ok(LlmProviders::Openrouter) => (&OPENROUTER_RESPONSES_CONFIG, resolved.into()),
+            _ => return Err(Error::Unsupported("native HTTP responses provider")),
+        };
+    Ok(ResponsesProvider {
+        identity: ProviderIdentity {
+            model: wire_model,
+            provider: provider.into(),
+        },
+        config,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    /// Pinned against Python's `get_llm_provider` on 2026-10-11.
+    #[rstest]
+    #[case::bare("gpt-5", None, "openai", "gpt-5")]
+    #[case::openai_prefix("openai/gpt-5", None, "openai", "gpt-5")]
+    #[case::openai_explicit("gpt-5", Some("openai"), "openai", "gpt-5")]
+    #[case::openrouter_prefix(
+        "openrouter/anthropic/claude",
+        None,
+        "openrouter",
+        "anthropic/claude"
+    )]
+    #[case::openrouter_explicit(
+        "openrouter/anthropic/claude",
+        Some("openrouter"),
+        "openrouter",
+        "anthropic/claude"
+    )]
+    #[case::openrouter_explicit_unprefixed(
+        "anthropic/claude",
+        Some("openrouter"),
+        "openrouter",
+        "anthropic/claude"
+    )]
+    #[case::openrouter_native_id(
+        "openrouter/auto",
+        Some("openrouter"),
+        "openrouter",
+        "openrouter/auto"
+    )]
+    #[case::openrouter_native_id_implicit("openrouter/auto", None, "openrouter", "auto")]
+    fn providers_and_wire_models_match_python(
+        #[case] model: &str,
+        #[case] custom_llm_provider: Option<&str>,
+        #[case] provider: &str,
+        #[case] wire_model: &str,
+    ) {
+        let resolved = resolve_provider(model, custom_llm_provider).unwrap();
+        assert_eq!(resolved.identity.provider, provider);
+        assert_eq!(resolved.identity.model, wire_model);
+    }
+
+    #[rstest]
+    #[case::other_prefix("other/model", None)]
+    #[case::other_explicit("model", Some("other"))]
+    #[case::openai_with_vendor_segment("openai/other/model", None)]
+    #[case::empty_model("openai/", None)]
+    #[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
+    fn everything_else_is_rejected(#[case] model: &str, #[case] custom_llm_provider: Option<&str>) {
+        assert!(matches!(
+            resolve_provider(model, custom_llm_provider),
+            Err(Error::InvalidProvider(_) | Error::Unsupported(_))
+        ));
+    }
+}

litellm-rust/crates/inference-responses/src/websocket.rs

@@ -2,7 +2,6 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use futures_util::{SinkExt, StreamExt};
 use litellm_http::websocket::{UpstreamWebSocket, connect_upstream};
-use litellm_llms_types::formats::responses::streaming_websocket::ResponsesWsEventType;
 use tokio::sync::Mutex;
 use tokio_tungstenite::tungstenite::{
     Message,
@@ -12,17 +11,6 @@ use tokio_tungstenite::tungstenite::{
 
 use super::Error;
 
-pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool {
-    matches!(
-        event_type,
-        ResponsesWsEventType::ResponseCreated
-            | ResponsesWsEventType::ResponseCompleted
-            | ResponsesWsEventType::ResponseFailed
-            | ResponsesWsEventType::ResponseIncomplete
-            | ResponsesWsEventType::Error
-    )
-}
-
 #[derive(Clone)]
 pub struct ResponsesWebSocketConnection {
     socket: Arc<Mutex<Option<UpstreamWebSocket>>>,

litellm-rust/crates/inference-responses/tests/live/main.rs

@@ -0,0 +1,7 @@
+//! Real provider calls, one module per provider. Each test skips when its provider's key is not
+//! in the environment, so `cargo test` stays offline by default.
+
+#[path = "../support/mod.rs"]
+mod support;
+
+mod openrouter;

litellm-rust/crates/inference-responses/tests/live/openrouter.rs

@@ -0,0 +1,98 @@
+use std::{sync::Arc, time::Duration};
+
+use futures_util::TryStreamExt;
+use litellm_inference_responses::types::{ResponsesCall, ResponsesOutput};
+use litellm_inference_testing::RecordingSecrets;
+use litellm_llms::base_llm::responses::streaming::decode;
+use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
+use rstest::{fixture, rstest};
+use serde_json::json;
+
+use crate::support::responses_route;
+
+const DEFAULT_MODEL: &str = "openrouter/openai/gpt-5.5";
+
+/// `OPENROUTER_API_KEY` as a secret source; tests skip when it is unset.
+#[fixture]
+fn secrets() -> Option<Arc<RecordingSecrets>> {
+    let key = std::env::var("OPENROUTER_API_KEY").ok()?;
+    Some(Arc::new(RecordingSecrets::new([(
+        "OPENROUTER_API_KEY",
+        key.as_str(),
+    )])))
+}
+
+fn call(stream: bool) -> ResponsesCall {
+    ResponsesCall {
+        model: std::env::var("OPENROUTER_LIVE_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()),
+        input: json!("Reply with the single word: pong"),
+        optional_params: json!({"max_output_tokens": 32, "stream": stream})
+            .as_object()
+            .unwrap()
+            .clone(),
+        api_key: None,
+        api_base: None,
+        custom_llm_provider: None,
+        extra_headers: None,
+        timeout: Some(Duration::from_secs(60)),
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_complete_response_carries_output_and_usage(secrets: Option<Arc<RecordingSecrets>>) {
+    let Some(secrets) = secrets else { return };
+    let output = responses_route(secrets)
+        .execute(call(false), &(), None)
+        .await
+        .unwrap();
+    let ResponsesOutput::Complete(response) = output else {
+        panic!("expected a complete response");
+    };
+    let body = serde_json::to_value(response).unwrap();
+    let text = body["output"]
+        .as_array()
+        .into_iter()
+        .flatten()
+        .filter_map(|item| item["content"].as_array())
+        .flatten()
+        .filter_map(|part| part["text"].as_str())
+        .collect::<String>();
+    assert!(text.to_ascii_lowercase().contains("pong"), "{body}");
+    assert!(
+        body["usage"]["total_tokens"].as_u64().unwrap() > 0,
+        "{body}"
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_stream_ends_with_response_completed(secrets: Option<Arc<RecordingSecrets>>) {
+    let Some(secrets) = secrets else { return };
+    let output = responses_route(secrets)
+        .execute(call(true), &(), None)
+        .await
+        .unwrap();
+    let ResponsesOutput::Stream { chunks, .. } = output else {
+        panic!("expected SSE");
+    };
+    let wire = chunks.try_collect::<Vec<_>>().await.unwrap().concat();
+    let events = decode(&wire).unwrap();
+    assert_eq!(
+        events.first().map(|event| &event.event_type),
+        Some(&ResponsesEventType::ResponseCreated)
+    );
+    assert!(
+        events
+            .iter()
+            .any(|event| event.event_type == ResponsesEventType::OutputTextDelta)
+    );
+    let last = events.last().unwrap();
+    assert_eq!(last.event_type, ResponsesEventType::ResponseCompleted);
+    assert!(
+        last.data["response"]["usage"]["total_tokens"]
+            .as_u64()
+            .unwrap()
+            > 0
+    );
+}

litellm-rust/crates/inference-responses/tests/openrouter.rs

@@ -0,0 +1,150 @@
+use std::sync::Arc;
+
+use futures_util::TryStreamExt;
+use litellm_inference_responses::{
+    Error,
+    types::{ResponsesCall, ResponsesOutput},
+};
+use litellm_inference_testing::{RecordingSecrets, no_secrets};
+use rstest::{fixture, rstest};
+use serde_json::json;
+use wiremock::ResponseTemplate;
+
+mod support;
+use support::*;
+
+#[fixture]
+fn call() -> ResponsesCall {
+    ResponsesCall {
+        model: "openrouter/anthropic/test-model".into(),
+        input: json!("hello"),
+        optional_params: Default::default(),
+        api_key: None,
+        api_base: None,
+        custom_llm_provider: None,
+        extra_headers: None,
+        timeout: None,
+    }
+}
+
+fn response_body() -> serde_json::Value {
+    json!({"id": "resp_1", "model": "anthropic/test-model", "output": [{"type": "message", "content": []}], "usage": {"total_tokens": 7, "cost": 0.00017}})
+}
+
+#[rstest]
+#[case::explicit_key(Some("explicit"), &[], "Bearer explicit")]
+#[case::managed_key(None, &[("OPENROUTER_API_KEY", "managed")], "Bearer managed")]
+#[case::legacy_name(None, &[("OR_API_KEY", "legacy")], "Bearer legacy")]
+#[tokio::test]
+async fn the_payload_goes_to_the_openrouter_responses_endpoint_with_a_bearer(
+    call: ResponsesCall,
+    #[case] api_key: Option<&str>,
+    #[case] secrets: &[(&str, &str)],
+    #[case] expected: &str,
+) {
+    let upstream = upstream([json_response(response_body())]).await;
+    let output = responses_route(Arc::new(RecordingSecrets::new(secrets.iter().copied())))
+        .execute(
+            ResponsesCall {
+                api_key: api_key.map(str::to_owned),
+                api_base: Some(format!("{}/api/v1", upstream.uri())),
+                optional_params: json!({"max_output_tokens": 16})
+                    .as_object()
+                    .unwrap()
+                    .clone(),
+                ..call
+            },
+            &(),
+            None,
+        )
+        .await
+        .unwrap();
+    let ResponsesOutput::Complete(response) = output else {
+        panic!("expected a complete response");
+    };
+    assert_eq!(serde_json::to_value(response).unwrap(), response_body());
+    let request = only_request(&upstream).await;
+    assert_eq!(request.url.path(), "/api/v1/responses");
+    assert_eq!(request.header("authorization"), Some(expected));
+    assert_eq!(
+        request.json(),
+        json!({"max_output_tokens": 16, "model": "anthropic/test-model", "input": "hello"})
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn the_secret_manager_supplies_the_base_and_key(call: ResponsesCall) {
+    let upstream = upstream([json_response(response_body())]).await;
+    let base = format!("{}/api/v1/", upstream.uri());
+    let secrets = Arc::new(RecordingSecrets::new([
+        ("OPENROUTER_API_KEY", "managed"),
+        ("OPENROUTER_API_BASE", base.as_str()),
+    ]));
+    responses_route(secrets.clone())
+        .execute(call, &(), None)
+        .await
+        .unwrap();
+    assert_eq!(
+        secrets.requested(),
+        ["OPENROUTER_API_KEY", "OR_API_KEY", "OPENROUTER_API_BASE"]
+    );
+    let request = only_request(&upstream).await;
+    assert_eq!(request.url.path(), "/api/v1/responses");
+    assert_eq!(request.header("authorization"), Some("Bearer managed"));
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_missing_key_fails_before_transport(call: ResponsesCall) {
+    let upstream = upstream([json_response(response_body())]).await;
+    let result = responses_route(no_secrets())
+        .execute(
+            ResponsesCall {
+                api_base: Some(upstream.uri()),
+                ..call
+            },
+            &(),
+            None,
+        )
+        .await;
+    assert!(matches!(
+        result,
+        Err(Error::Auth(litellm_auth::Error::MissingApiKey {
+            provider: "OpenRouter",
+            ..
+        }))
+    ));
+    assert!(received(&upstream).await.is_empty());
+}
+
+#[rstest]
+#[tokio::test]
+async fn streaming_relays_every_openrouter_event_untouched(call: ResponsesCall) {
+    let sse = "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"usage\":{\"cost\":0.0002}}}\n\n";
+    let upstream =
+        upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await;
+    let output = responses_route(no_secrets())
+        .execute(
+            ResponsesCall {
+                api_key: Some("key".into()),
+                api_base: Some(upstream.uri()),
+                optional_params: json!({"stream": true}).as_object().unwrap().clone(),
+                ..call
+            },
+            &(),
+            None,
+        )
+        .await
+        .unwrap();
+    let ResponsesOutput::Stream { chunks, .. } = output else {
+        panic!("expected SSE");
+    };
+    assert_eq!(
+        chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
+        sse.as_bytes()
+    );
+    let request = only_request(&upstream).await;
+    assert_eq!(request.json()["stream"], true);
+    assert_eq!(request.json()["model"], "anthropic/test-model");
+}

litellm-rust/crates/llms-types/src/formats/chat_completions/streaming.rs

@@ -33,16 +33,82 @@ pub struct ChatCompletionStreamingChoice {
     pub logprobs: Option<Value>,
 }
 
+/// Compatible hosts omit `id`, `created` and `object` on chunks, and Python's
+/// `ModelResponseStream` fills them in, so none of the envelope fields is required here.
 #[macro_rules_attribute::apply(crate::wire_type)]
+#[derive(Default)]
 pub struct ChatCompletionChunk {
-    pub id: String,
-    pub created: u64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub created: Option<u64>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub model: Option<String>,
-    pub object: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub object: Option<String>,
+    #[serde(default)]
     pub choices: Vec<ChatCompletionStreamingChoice>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub usage: Option<ChatCompletionsUsage>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub provider_specific_fields: Option<Map<String, Value>>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use serde_json::{Value, json};
+
+    use super::*;
+
+    #[rstest]
+    #[case::openai(json!({
+        "id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m",
+        "choices": [{"index": 0, "delta": {"role": "assistant", "content": "hi", "future": 1}, "finish_reason": null, "logprobs": null}],
+        "usage": null
+    }))]
+    #[case::usage_only(json!({
+        "id": "chatcmpl-1", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [],
+        "usage": {
+            "prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3,
+            "prompt_tokens_details": {"cached_tokens": 0, "cache_creation_tokens": 0, "text_tokens": 1}
+        }
+    }))]
+    #[case::bare_host(json!({"choices": [{"index": 0, "delta": {"content": "hi"}}]}))]
+    #[case::host_extensions(json!({"id": "x", "choices": [], "citations": ["a"], "future": {"k": null}}))]
+    fn chunks_round_trip_without_losing_fields(#[case] wire: Value) {
+        let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
+        assert_eq!(serde_json::to_value(chunk).unwrap(), without_nulls(wire));
+    }
+
+    /// Typed optional fields serialize by omission; extension fields keep their nulls.
+    fn without_nulls(value: Value) -> Value {
+        let Value::Object(fields) = value else {
+            return value;
+        };
+        Value::Object(
+            fields
+                .into_iter()
+                .filter(|(_, value)| !value.is_null())
+                .map(|(key, value)| match (key.as_str(), value) {
+                    ("choices", Value::Array(choices)) => (
+                        key,
+                        Value::Array(choices.into_iter().map(without_nulls).collect()),
+                    ),
+                    (_, value) => (key, value),
+                })
+                .collect(),
+        )
+    }
+
+    #[rstest]
+    #[case::choices_not_a_list(json!({"choices": {}}))]
+    #[case::choice_without_delta(json!({"choices": [{"index": 0}]}))]
+    #[case::created_not_a_number(json!({"created": "now", "choices": []}))]
+    #[case::non_object(json!([]))]
+    fn malformed_chunks_are_rejected(#[case] wire: Value) {
+        assert!(serde_json::from_value::<ChatCompletionChunk>(wire).is_err());
+    }
 }

litellm-rust/crates/llms-types/src/formats/messages/streaming.rs

@@ -109,7 +109,9 @@ pub struct MessagesStreamError {
 }
 
 #[macro_rules_attribute::apply(crate::wire_type)]
+#[derive(strum::AsRefStr)]
 #[serde(tag = "type", rename_all = "snake_case")]
+#[strum(serialize_all = "snake_case")]
 pub enum MessagesStreamEvent {
     MessageStart {
         message: MessagesStreamMessage,
@@ -175,6 +177,23 @@ mod tests {
         assert_eq!(serde_json::to_value(event).unwrap(), wire);
     }
 
+    #[rstest]
+    #[case::start(json!({"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "x", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {}}}))]
+    #[case::block_start(json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}))]
+    #[case::block_delta(json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "a"}}))]
+    #[case::block_stop(json!({"type": "content_block_stop", "index": 0}))]
+    #[case::delta(json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}))]
+    #[case::stop(json!({"type": "message_stop"}))]
+    #[case::ping(json!({"type": "ping"}))]
+    #[case::error(json!({"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}))]
+    fn the_event_name_is_the_serialized_type_tag(#[case] wire: Value) {
+        let event: MessagesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
+        assert_eq!(
+            Some(event.as_ref()),
+            wire.get("type").and_then(Value::as_str)
+        );
+    }
+
     #[rstest]
     #[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))]
     #[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))]

litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md

@@ -1,7 +1,7 @@
 # rules
 
 - Responses originated at OpenAI and is now offered by other hosts and gateways. Types here describe the format, not OpenAI's deployment of it
-- Served over HTTP and over the WebSocket mode. `ResponsesWsEvent` is a wire event and lives here
+- Served over HTTP and over the WebSocket mode. `ResponsesStreamEvent` and `ResponsesEventType` in `streaming.rs` are the wire event for both
 - `ResponsesWsTransformResult` wraps a provider transformation's output, so it lives in `llms/src/base_llm/responses/`
 - Hosts that accept only part of the format, or add output item types, are handled in `llms/src/<provider>/responses/`. Keep unknown output items passing through instead of rejecting them
 - Responses-to-Chat emulation and model-specific parameter rewriting are transformations, not format rules

litellm-rust/crates/llms-types/src/formats/responses/mod.rs

@@ -1,5 +1,6 @@
 mod output;
 mod response;
+pub mod streaming;
 pub mod streaming_websocket;
 
 pub use output::*;

litellm-rust/crates/llms-types/src/formats/responses/streaming.rs

@@ -0,0 +1,305 @@
+use serde_json::{Map, Value};
+
+/// Every event name on https://developers.openai.com/api/reference/resources/responses/streaming-events
+/// as of 2026-10-11, plus `response.create` from the WebSocket mode. Names OpenAI adds later ride in
+/// `Other` untouched.
+#[derive(
+    Clone,
+    Debug,
+    PartialEq,
+    Eq,
+    strum::AsRefStr,
+    strum::EnumIter,
+    strum::EnumString,
+    strum::Display,
+    serde_with::DeserializeFromStr,
+    serde_with::SerializeDisplay,
+)]
+#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
+#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))]
+pub enum ResponsesEventType {
+    #[strum(serialize = "response.create")]
+    ResponseCreate,
+    #[strum(serialize = "response.created")]
+    ResponseCreated,
+    #[strum(serialize = "response.queued")]
+    ResponseQueued,
+    #[strum(serialize = "response.in_progress")]
+    ResponseInProgress,
+    #[strum(serialize = "response.completed")]
+    ResponseCompleted,
+    #[strum(serialize = "response.failed")]
+    ResponseFailed,
+    #[strum(serialize = "response.incomplete")]
+    ResponseIncomplete,
+    #[strum(serialize = "error")]
+    Error,
+    #[strum(serialize = "response.audio.delta")]
+    AudioDelta,
+    #[strum(serialize = "response.audio.done")]
+    AudioDone,
+    #[strum(serialize = "response.audio.transcript.delta")]
+    AudioTranscriptDelta,
+    #[strum(serialize = "response.audio.transcript.done")]
+    AudioTranscriptDone,
+    #[strum(serialize = "response.code_interpreter_call.completed")]
+    CodeInterpreterCallCompleted,
+    #[strum(serialize = "response.code_interpreter_call.in_progress")]
+    CodeInterpreterCallInProgress,
+    #[strum(serialize = "response.code_interpreter_call.interpreting")]
+    CodeInterpreterCallInterpreting,
+    #[strum(serialize = "response.code_interpreter_call_code.delta")]
+    CodeInterpreterCallCodeDelta,
+    #[strum(serialize = "response.code_interpreter_call_code.done")]
+    CodeInterpreterCallCodeDone,
+    #[strum(serialize = "response.compaction.compacting")]
+    CompactionCompacting,
+    #[strum(serialize = "response.content_part.added")]
+    ContentPartAdded,
+    #[strum(serialize = "response.content_part.done")]
+    ContentPartDone,
+    #[strum(serialize = "response.custom_tool_call_input.delta")]
+    CustomToolCallInputDelta,
+    #[strum(serialize = "response.custom_tool_call_input.done")]
+    CustomToolCallInputDone,
+    #[strum(serialize = "response.file_search_call.completed")]
+    FileSearchCallCompleted,
+    #[strum(serialize = "response.file_search_call.in_progress")]
+    FileSearchCallInProgress,
+    #[strum(serialize = "response.file_search_call.searching")]
+    FileSearchCallSearching,
+    #[strum(serialize = "response.function_call_arguments.delta")]
+    FunctionCallArgumentsDelta,
+    #[strum(serialize = "response.function_call_arguments.done")]
+    FunctionCallArgumentsDone,
+    #[strum(serialize = "response.image_generation_call.completed")]
+    ImageGenerationCallCompleted,
+    #[strum(serialize = "response.image_generation_call.generating")]
+    ImageGenerationCallGenerating,
+    #[strum(serialize = "response.image_generation_call.in_progress")]
+    ImageGenerationCallInProgress,
+    #[strum(serialize = "response.image_generation_call.partial_image")]
+    ImageGenerationCallPartialImage,
+    #[strum(serialize = "response.mcp_call.completed")]
+    McpCallCompleted,
+    #[strum(serialize = "response.mcp_call.failed")]
+    McpCallFailed,
+    #[strum(serialize = "response.mcp_call.in_progress")]
+    McpCallInProgress,
+    #[strum(serialize = "response.mcp_call_arguments.delta")]
+    McpCallArgumentsDelta,
+    #[strum(serialize = "response.mcp_call_arguments.done")]
+    McpCallArgumentsDone,
+    #[strum(serialize = "response.mcp_list_tools.completed")]
+    McpListToolsCompleted,
+    #[strum(serialize = "response.mcp_list_tools.failed")]
+    McpListToolsFailed,
+    #[strum(serialize = "response.mcp_list_tools.in_progress")]
+    McpListToolsInProgress,
+    #[strum(serialize = "response.output_item.added")]
+    OutputItemAdded,
+    #[strum(serialize = "response.output_item.done")]
+    OutputItemDone,
+    #[strum(serialize = "response.output_text.annotation.added")]
+    OutputTextAnnotationAdded,
+    #[strum(serialize = "response.output_text.delta")]
+    OutputTextDelta,
+    #[strum(serialize = "response.output_text.done")]
+    OutputTextDone,
+    #[strum(serialize = "response.reasoning_summary_part.added")]
+    ReasoningSummaryPartAdded,
+    #[strum(serialize = "response.reasoning_summary_part.done")]
+    ReasoningSummaryPartDone,
+    #[strum(serialize = "response.reasoning_summary_text.delta")]
+    ReasoningSummaryTextDelta,
+    #[strum(serialize = "response.reasoning_summary_text.done")]
+    ReasoningSummaryTextDone,
+    #[strum(serialize = "response.reasoning_text.delta")]
+    ReasoningTextDelta,
+    #[strum(serialize = "response.reasoning_text.done")]
+    ReasoningTextDone,
+    #[strum(serialize = "response.refusal.delta")]
+    RefusalDelta,
+    #[strum(serialize = "response.refusal.done")]
+    RefusalDone,
+    #[strum(serialize = "response.shell_call_command.added")]
+    ShellCallCommandAdded,
+    #[strum(serialize = "response.shell_call_command.delta")]
+    ShellCallCommandDelta,
+    #[strum(serialize = "response.shell_call_command.done")]
+    ShellCallCommandDone,
+    #[strum(serialize = "response.shell_call_output_content.delta")]
+    ShellCallOutputContentDelta,
+    #[strum(serialize = "response.shell_call_output_content.done")]
+    ShellCallOutputContentDone,
+    #[strum(serialize = "response.web_search_call.completed")]
+    WebSearchCallCompleted,
+    #[strum(serialize = "response.web_search_call.in_progress")]
+    WebSearchCallInProgress,
+    #[strum(serialize = "response.web_search_call.searching")]
+    WebSearchCallSearching,
+    #[strum(default, transparent)]
+    Other(String),
+}
+
+impl ResponsesEventType {
+    /// The last event of a response. `response.completed`, `response.incomplete` and
+    /// `response.failed` each close the response and carry its final body; `error` aborts it.
+    pub fn is_terminal(&self) -> bool {
+        matches!(
+            self,
+            Self::ResponseCompleted | Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
+        )
+    }
+
+    /// A terminal event after which the response is not a success.
+    pub fn is_failure(&self) -> bool {
+        matches!(
+            self,
+            Self::ResponseFailed | Self::ResponseIncomplete | Self::Error
+        )
+    }
+}
+
+#[macro_rules_attribute::apply(crate::wire_type)]
+pub struct ResponsesStreamEvent {
+    #[serde(rename = "type")]
+    pub event_type: ResponsesEventType,
+    #[serde(flatten)]
+    pub data: Map<String, Value>,
+}
+
+impl ResponsesStreamEvent {
+    pub fn model(&self) -> Option<&str> {
+        let model = self.data.get("model").and_then(Value::as_str);
+        if model.is_some() {
+            return model;
+        }
+        self.data
+            .get("response")
+            .and_then(Value::as_object)
+            .and_then(|response| response.get("model"))
+            .and_then(Value::as_str)
+    }
+
+    pub fn is_response_create(&self) -> bool {
+        self.event_type == ResponsesEventType::ResponseCreate
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use serde_json::{Value, json};
+    use strum::IntoEnumIterator;
+
+    use super::*;
+
+    #[rstest]
+    fn every_named_event_type_round_trips() {
+        let named: Vec<_> = ResponsesEventType::iter()
+            .filter(|event_type| !matches!(event_type, ResponsesEventType::Other(_)))
+            .collect();
+        assert!(named.len() > 50);
+        for event_type in named {
+            let wire = event_type.to_string();
+            assert_eq!(wire, event_type.as_ref());
+            assert_eq!(serde_json::to_value(&event_type).unwrap(), json!(wire));
+            assert_eq!(wire.parse::<ResponsesEventType>().unwrap(), event_type);
+            assert_eq!(
+                serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
+                event_type
+            );
+        }
+    }
+
+    #[rstest]
+    #[case::unknown("future.event", ResponsesEventType::Other("future.event".into()))]
+    #[case::empty("", ResponsesEventType::Other(String::new()))]
+    #[case::case_sensitive(
+        "Response.Completed",
+        ResponsesEventType::Other("Response.Completed".into())
+    )]
+    #[case::escaped("future\"\\\n", ResponsesEventType::Other("future\"\\\n".into()))]
+    fn unknown_event_types_round_trip_through_other(
+        #[case] wire: &str,
+        #[case] expected: ResponsesEventType,
+    ) {
+        assert_eq!(
+            serde_json::from_value::<ResponsesEventType>(json!(wire)).unwrap(),
+            expected
+        );
+        assert_eq!(serde_json::to_value(&expected).unwrap(), json!(wire));
+        assert_eq!(expected.as_ref(), wire);
+        assert_eq!(expected.to_string(), wire);
+    }
+
+    /// Python's `ResponsesAPIStreamingIterator` stores the response on `response.completed`,
+    /// `response.incomplete` and `response.failed`, and raises on `response.failed` and `error`.
+    #[rstest]
+    #[case::completed(ResponsesEventType::ResponseCompleted, true, false)]
+    #[case::incomplete(ResponsesEventType::ResponseIncomplete, true, true)]
+    #[case::failed(ResponsesEventType::ResponseFailed, true, true)]
+    #[case::error(ResponsesEventType::Error, true, true)]
+    #[case::create(ResponsesEventType::ResponseCreate, false, false)]
+    #[case::created(ResponsesEventType::ResponseCreated, false, false)]
+    #[case::in_progress(ResponsesEventType::ResponseInProgress, false, false)]
+    #[case::delta(ResponsesEventType::OutputTextDelta, false, false)]
+    #[case::item_done(ResponsesEventType::OutputItemDone, false, false)]
+    #[case::unknown(ResponsesEventType::Other("response.future".into()), false, false)]
+    #[case::tool_failed(ResponsesEventType::McpCallFailed, false, false)]
+    fn terminal_and_failure_events_match_the_python_iterator(
+        #[case] event_type: ResponsesEventType,
+        #[case] terminal: bool,
+        #[case] failure: bool,
+    ) {
+        assert_eq!(event_type.is_terminal(), terminal);
+        assert_eq!(event_type.is_failure(), failure);
+    }
+
+    #[rstest]
+    #[case::number("17")]
+    #[case::boolean("true")]
+    #[case::null("null")]
+    #[case::array("[]")]
+    #[case::object("{}")]
+    fn event_types_reject_non_strings(#[case] wire: &str) {
+        assert!(serde_json::from_str::<ResponsesEventType>(wire).is_err());
+    }
+
+    #[rstest]
+    #[case::text(json!({"type":"response.output_text.delta","delta":"hello","sequence_number":3,"future":null}))]
+    #[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1","output_index":0}))]
+    #[case::unknown(json!({"type":"future.event","payload":{"items":[1,null]},"future":true}))]
+    #[case::terminal(json!({"type":"response.completed","response":{"id":"r","output":[],"future":null}}))]
+    fn events_preserve_every_payload_field(#[case] wire: Value) {
+        let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
+        assert_eq!(serde_json::to_value(event).unwrap(), wire);
+    }
+
+    #[rstest]
+    #[case::missing_type(json!({"delta":"hello"}))]
+    #[case::numeric_type(json!({"type":7}))]
+    #[case::null_type(json!({"type":null}))]
+    #[case::non_object(json!([]))]
+    fn malformed_events_are_rejected(#[case] wire: Value) {
+        assert!(serde_json::from_value::<ResponsesStreamEvent>(wire).is_err());
+    }
+
+    #[rstest]
+    #[case::flat(json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
+    #[case::nested(json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
+    #[case::none(json!({"type":"response.created","response":{}}), None)]
+    fn model_reads_flat_and_nested_shapes(#[case] wire: Value, #[case] expected: Option<&str>) {
+        let event: ResponsesStreamEvent = serde_json::from_value(wire).unwrap();
+        assert_eq!(event.model(), expected);
+    }
+
+    #[cfg(feature = "schema")]
+    #[rstest]
+    fn event_type_schema_accepts_unknown_strings() {
+        let schema = schemars::schema_for!(ResponsesEventType).to_value();
+        assert_eq!(schema.get("type"), Some(&json!("string")));
+        assert!(schema.get("enum").is_none());
+    }
+}

litellm-rust/crates/llms-types/src/formats/responses/streaming_websocket.rs

@@ -2,68 +2,6 @@ use super::ResponsesOutputItem;
 use crate::recognized::Recognized;
 use serde_json::{Map, Value};
 
-#[derive(
-    Clone,
-    Debug,
-    PartialEq,
-    Eq,
-    strum::AsRefStr,
-    strum::EnumString,
-    strum::Display,
-    serde_with::DeserializeFromStr,
-    serde_with::SerializeDisplay,
-)]
-#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
-#[cfg_attr(feature = "schema", schemars(with = "String"))]
-pub enum ResponsesWsEventType {
-    #[strum(serialize = "response.create")]
-    ResponseCreate,
-    #[strum(serialize = "response.created")]
-    ResponseCreated,
-    #[strum(serialize = "response.completed")]
-    ResponseCompleted,
-    #[strum(serialize = "response.failed")]
-    ResponseFailed,
-    #[strum(serialize = "response.incomplete")]
-    ResponseIncomplete,
-    #[strum(serialize = "error")]
-    Error,
-    #[strum(default, transparent)]
-    Other(String),
-}
-
-impl ResponsesWsEventType {
-    pub fn as_str(&self) -> &str {
-        self.as_ref()
-    }
-}
-
-#[macro_rules_attribute::apply(crate::wire_type)]
-pub struct ResponsesWsEvent {
-    #[serde(rename = "type")]
-    pub event_type: ResponsesWsEventType,
-    #[serde(flatten)]
-    pub data: Map<String, Value>,
-}
-
-impl ResponsesWsEvent {
-    pub fn model(&self) -> Option<&str> {
-        let model = self.data.get("model").and_then(Value::as_str);
-        if model.is_some() {
-            return model;
-        }
-        self.data
-            .get("response")
-            .and_then(Value::as_object)
-            .and_then(|response| response.get("model"))
-            .and_then(Value::as_str)
-    }
-
-    pub fn is_response_create(&self) -> bool {
-        self.event_type == ResponsesWsEventType::ResponseCreate
-    }
-}
-
 #[macro_rules_attribute::apply(crate::wire_type)]
 #[derive(Eq)]
 pub struct ResponsesErrorFrame {
@@ -112,8 +50,7 @@ mod tests {
     use crate::formats::responses::ResponsesOutputItem;
 
     use crate::{
-        formats::responses::streaming_websocket::{ResponsesEventResponse, ResponsesWsEventType},
-        recognized::Recognized,
+        formats::responses::streaming_websocket::ResponsesEventResponse, recognized::Recognized,
         test_support::*,
     };
 
@@ -134,66 +71,6 @@ mod tests {
         );
     }
 
-    #[rstest]
-    #[case::flat(serde_json::json!({"type":"response.create","model":"model-a"}), Some("model-a"))]
-    #[case::nested(serde_json::json!({"type":"response.create","response":{"model":"model-b"}}), Some("model-b"))]
-    fn model_reads_flat_and_nested_create_shapes(
-        #[case] payload: serde_json::Value,
-        #[case] expected: Option<&str>,
-    ) {
-        let event: ResponsesWsEvent = serde_json::from_value(payload).expect("valid event");
-        assert_eq!(event.model(), expected);
-    }
-
-    #[rstest]
-    #[case::create("response.create", ResponsesWsEventType::ResponseCreate)]
-    #[case::created("response.created", ResponsesWsEventType::ResponseCreated)]
-    #[case::completed("response.completed", ResponsesWsEventType::ResponseCompleted)]
-    #[case::failed("response.failed", ResponsesWsEventType::ResponseFailed)]
-    #[case::incomplete("response.incomplete", ResponsesWsEventType::ResponseIncomplete)]
-    #[case::error("error", ResponsesWsEventType::Error)]
-    #[case::unknown(
-        "response.output_text.delta",
-        ResponsesWsEventType::Other("response.output_text.delta".to_string())
-    )]
-    #[case::empty("", ResponsesWsEventType::Other(String::new()))]
-    #[case::case_sensitive(
-        "Response.Completed",
-        ResponsesWsEventType::Other("Response.Completed".into())
-    )]
-    #[case::escaped("future\"\\\n", ResponsesWsEventType::Other("future\"\\\n".into()))]
-    fn websocket_event_type_round_trips(
-        #[case] wire: &str,
-        #[case] expected: ResponsesWsEventType,
-    ) {
-        let serialized = serde_json::to_string(&expected).unwrap();
-        assert_eq!(serialized, serde_json::to_string(wire).unwrap());
-        assert_eq!(
-            serde_json::from_str::<ResponsesWsEventType>(&serialized).unwrap(),
-            expected
-        );
-    }
-
-    #[rstest]
-    #[case::number("17")]
-    #[case::boolean("true")]
-    #[case::null("null")]
-    #[case::array("[]")]
-    #[case::object("{}")]
-    fn websocket_event_type_rejects_non_strings(#[case] wire: &str) {
-        assert!(serde_json::from_str::<ResponsesWsEventType>(wire).is_err());
-    }
-
-    #[cfg(feature = "schema")]
-    #[rstest]
-    fn websocket_event_type_schema_is_open_string() {
-        let schema = schemars::schema_for!(ResponsesWsEventType);
-        assert_eq!(
-            schema.to_value().get("type"),
-            Some(&serde_json::json!("string"))
-        );
-    }
-
     #[rstest]
     fn nested_event_response_exposes_typed_output_and_preserves_extensions() {
         let wire = json!({

litellm-rust/crates/llms/AGENTS.md

@@ -32,6 +32,13 @@
   - Provider-wide docs (auth, errors, regions): `src/<provider>/`
   - One host's endpoint docs: `src/<provider>/<format>/`
 
+## Streaming
+
+- Each format's `src/base_llm/<format>/streaming.rs` exposes the same three names: `event_stream` (bytes to typed events), `decode` (a whole stream to typed events) and `encode` (one typed event to the SSE frame LiteLLM's proxy sends clients). Callers use the module path, never a vendor prefix
+- SSE framing, `[DONE]` and the generic helpers live in `src/base_llm/sse.rs`. Messages frames are `event:` plus `data:`; Chat Completions and Responses frames are bare `data:`, as Python's proxy writes them
+- Decoding is strict: a malformed event, or `[DONE]` on a Messages stream, is `InvalidResponse` and ends the stream. Python never fails a stream this way; the divergence is deliberate and `src/base_llm/sse.rs` lists what Python does per format. A stream cut inside an event drops that event, as the SSE spec requires
+- What ends or fails a stream belongs on the format's event type (`is_terminal`, `is_failure`). The cache's success check in `inference/src/caching.rs` still matches type strings; move it onto the typed events rather than adding new string matches
+
 ## Python pairs
 
 - Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure
@@ -51,6 +58,7 @@
 
 - `AnthropicConfig`, `AmazonConverseConfig` and the non-OCR base traits are partial. Preserve their acceptance gates, passthrough and host fallback contracts
 - `OpenAiResponsesApiConfig` implements WebSocket transformations and direct HTTP Responses, without Python's model-specific param rewriting or Responses-to-Chat emulation
+- OpenAI-compatible Responses hosts (`openrouter`) reuse the OpenAI request and response transforms as free functions and own only their credentials and base URL
 - Chat `supported_openai_param_mappings` lists OpenAI/provider name pairs, unlike Python's name-only list. Audio `map_transcription_params` is a Rust filtering helper
 - Auth resolution is split between configs and route preparation in `inference-<format>`
 

litellm-rust/crates/llms/src/anthropic/chat/transformation.rs

@@ -30,7 +30,7 @@ use crate::{
                 Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
             },
         },
-        messages::streaming::anthropic_sse_event_stream,
+        messages::streaming::event_stream as messages_event_stream,
     },
 };
 
@@ -196,7 +196,7 @@ impl BaseConfig for AnthropicConfig {
 
     fn model_response_iterator(&self, shape: StreamShape) -> Option<ChatStream> {
         Some(ChatStream::new(
-            anthropic_sse_event_stream,
+            messages_event_stream,
             ModelResponseIterator::new(shape),
         ))
     }

litellm-rust/crates/llms/src/base_llm/chat/streaming.rs

@@ -1,15 +1,32 @@
 use std::collections::HashMap;
 
+use bytes::Bytes;
 use futures_util::{StreamExt, stream::BoxStream};
 use litellm_llms_types::formats::chat_completions::ChatCompletionChunk;
 
 use crate::{
     Error,
-    base_llm::base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
+    base_llm::{
+        base_model_iterator::{ByteStream, StreamError, StreamTransformer, transform_stream},
+        sse::{self, is_done},
+    },
 };
 
 pub type ChatChunkStream = BoxStream<'static, Result<ChatCompletionChunk, Error>>;
 
+/// OpenAI-compatible SSE: bare `data:` frames closed by `data: [DONE]`.
+pub fn event_stream(bytes: ByteStream) -> ChatChunkStream {
+    sse::event_stream(bytes, is_done)
+}
+
+pub fn decode(wire: &[u8]) -> Result<Vec<ChatCompletionChunk>, Error> {
+    sse::decode(wire, is_done)
+}
+
+pub fn encode(chunk: &ChatCompletionChunk) -> Result<Bytes, Error> {
+    sse::encode_data(chunk)
+}
+
 /// What Python's `map_openai_params` decides about the stream and `completion`
 /// hands to `ModelResponseIterator`: it is settled while the request is built,
 /// never re-derived from the body.
@@ -52,3 +69,109 @@ impl ChatStream {
         (self.run)(bytes)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use futures_util::{StreamExt, TryStreamExt, stream};
+    use rstest::rstest;
+    use serde_json::{Value, json};
+
+    use super::*;
+
+    fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
+        let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
+        stream::iter(pieces.into_iter().map(Ok)).boxed()
+    }
+
+    #[rstest]
+    #[case::role(json!({"id":"c","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}))]
+    #[case::text(json!({"id":"c","choices":[{"index":0,"delta":{"content":"hi\nthere"}}]}))]
+    #[case::tool(json!({"id":"c","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":"{\"x\":"}}]}}]}))]
+    #[case::finish(json!({"id":"c","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}))]
+    #[case::usage(json!({"id":"c","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":0,"cache_creation_tokens":0,"text_tokens":1}}}))]
+    #[case::host_extension(json!({"id":"c","choices":[],"citations":[{"url":"u"}]}))]
+    fn encoded_chunks_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
+        let chunk: ChatCompletionChunk = serde_json::from_value(wire.clone()).unwrap();
+        let encoded = encode(&chunk).unwrap();
+        assert!(encoded.starts_with(b"data: "));
+        assert!(encoded.ends_with(b"\n\n"));
+        assert_eq!(
+            serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
+            wire
+        );
+        assert_eq!(decode(&encoded).unwrap(), vec![chunk]);
+    }
+
+    #[rstest]
+    #[case::lf("\n")]
+    #[case::crlf("\r\n")]
+    #[case::cr("\r")]
+    #[tokio::test]
+    async fn streams_split_anywhere_skip_comments_and_done(
+        #[case] newline: &str,
+        #[values(1, 5, 4096)] size: usize,
+    ) {
+        let wire = format!(
+            ": keepalive{newline}{newline}data: {{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"a\"}}}}]}}{newline}{newline}data:{{\"id\":\"c\",\"choices\":[{{\"index\":0,\"delta\":{{}},\"finish_reason\":\"stop\"}}]}}{newline}{newline}data: [DONE]{newline}{newline}"
+        );
+        let streamed = event_stream(in_pieces(wire.as_bytes(), size))
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
+        assert_eq!(streamed.len(), 2);
+        assert_eq!(streamed[0].choices[0].delta.content.as_deref(), Some("a"));
+        assert_eq!(
+            streamed[1].choices[0].finish_reason.as_deref(),
+            Some("stop")
+        );
+    }
+
+    #[rstest]
+    #[case::spaced(b"data: [DONE]\n\n")]
+    #[case::unspaced(b"data:[DONE]\n\n")]
+    #[case::trailing_space(b"data: [DONE] \n\n")]
+    fn every_done_spelling_closes_the_stream_without_a_chunk(#[case] wire: &[u8]) {
+        assert_eq!(decode(wire).unwrap(), vec![]);
+    }
+
+    #[rstest]
+    #[case::invalid_json(b"data: {\n\n")]
+    #[case::choices_not_a_list(b"data: {\"choices\":{}}\n\n")]
+    #[case::invalid_utf8(b"data: \xff\n\n")]
+    #[tokio::test]
+    async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
+        assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
+        let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
+        assert!(matches!(
+            streamed.last(),
+            Some(Err(Error::InvalidResponse(_))) | None
+        ));
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn a_body_error_ends_the_stream_after_delivered_chunks() {
+        let mut chunks = event_stream(
+            stream::iter([
+                Ok(Bytes::from_static(
+                    b"data: {\"id\":\"c\",\"choices\":[]}\n\n",
+                )),
+                Err(std::io::Error::new(
+                    std::io::ErrorKind::ConnectionReset,
+                    "reset",
+                )),
+            ])
+            .boxed(),
+        );
+        assert_eq!(
+            chunks.next().await.unwrap().unwrap().id.as_deref(),
+            Some("c")
+        );
+        assert!(matches!(
+            chunks.next().await,
+            Some(Err(Error::InvalidResponse(_)))
+        ));
+        assert!(chunks.next().await.is_none());
+    }
+}

litellm-rust/crates/llms/src/base_llm/messages/streaming.rs

@@ -1,38 +1,32 @@
 use bytes::Bytes;
-use futures_util::{StreamExt, stream::BoxStream};
-use litellm_framer::{frames, sse::SseCodec};
 use litellm_llms_types::formats::messages::streaming::MessagesStreamEvent;
 
-use crate::Error;
 pub use crate::base_llm::base_model_iterator::ByteStream;
+use crate::{Error, base_llm::sse};
 
-pub type EventStream = BoxStream<'static, Result<MessagesStreamEvent, Error>>;
+pub type EventStream = sse::EventStream<MessagesStreamEvent>;
 pub type StreamDecoder = fn(ByteStream) -> EventStream;
 
-pub fn anthropic_sse_event_stream(bytes: ByteStream) -> EventStream {
-    Box::pin(frames(bytes, SseCodec::default()).map(|event| {
-        let event = event.map_err(|error| {
-            Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error))
-        })?;
-        serde_json::from_str(&event.data).map_err(|error| {
-            Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
-        })
-    }))
+fn no_control_frames(_: &litellm_framer::sse::SseEvent) -> bool {
+    false
 }
 
-pub fn encode_anthropic_sse(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
-    let data = serde_json::to_value(event).map_err(|error| {
-        Error::InvalidResponse(crate::ErrorDetail::invalid("Anthropic stream event", error))
-    })?;
-    let name = data
-        .get("type")
-        .and_then(serde_json::Value::as_str)
-        .ok_or_else(|| {
-            Error::InvalidResponse(
-                "Anthropic stream event is invalid: stream event has no type".into(),
-            )
-        })?;
-    Ok(Bytes::from(format!("event: {name}\ndata: {data}\n\n")))
+pub fn event_stream(bytes: ByteStream) -> EventStream {
+    sse::event_stream(bytes, no_control_frames)
+}
+
+pub fn decode(wire: &[u8]) -> Result<Vec<MessagesStreamEvent>, Error> {
+    sse::decode(wire, no_control_frames)
+}
+
+/// Anthropic names every frame after the event's `type`, and LiteLLM's proxy does the same
+/// when it re-emits translated events.
+pub fn encode(event: &MessagesStreamEvent) -> Result<Bytes, Error> {
+    let data = serde_json::to_string(event).map_err(sse::invalid)?;
+    Ok(Bytes::from(format!(
+        "event: {}\ndata: {data}\n\n",
+        event.as_ref()
+    )))
 }
 
 #[cfg(test)]
@@ -41,6 +35,7 @@ mod tests {
     use litellm_llms_types::formats::messages::streaming::{
         MessagesContentBlockDelta, MessagesStreamUsage,
     };
+    use rstest::rstest;
     use serde_json::json;
 
     use super::*;
@@ -48,51 +43,11 @@ mod tests {
     const TEXT_DELTA: &str =
         r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
 
-    fn in_pieces(wire: &[u8]) -> ByteStream {
-        let pieces: Vec<Bytes> = wire.chunks(3).map(Bytes::copy_from_slice).collect();
+    fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
+        let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
         stream::iter(pieces.into_iter().map(Ok)).boxed()
     }
 
-    #[tokio::test]
-    async fn sse_frames_split_anywhere_decode_into_typed_events() {
-        let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
-        let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
-            .try_collect::<Vec<_>>()
-            .await
-            .unwrap();
-
-        assert_eq!(
-            events,
-            vec![MessagesStreamEvent::ContentBlockDelta {
-                index: 0,
-                delta: MessagesContentBlockDelta::TextDelta {
-                    text: "hello".into(),
-                },
-            }]
-        );
-    }
-
-    #[tokio::test]
-    async fn decodes_citations_delta_events() {
-        let wire = concat!(
-            "event: content_block_delta\n",
-            r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
-            "\n\n",
-        );
-        let events = anthropic_sse_event_stream(in_pieces(wire.as_bytes()))
-            .try_collect::<Vec<_>>()
-            .await
-            .unwrap();
-
-        assert!(matches!(
-            events.as_slice(),
-            [MessagesStreamEvent::ContentBlockDelta {
-                delta: MessagesContentBlockDelta::Citations { .. },
-                ..
-            }]
-        ));
-    }
-
     fn events() -> Vec<MessagesStreamEvent> {
         vec![
             MessagesStreamEvent::Ping,
@@ -110,27 +65,80 @@ mod tests {
         ]
     }
 
-    #[tokio::test]
-    async fn encoded_events_decode_back_to_themselves() {
-        let wire = events()
-            .iter()
-            .map(encode_anthropic_sse)
-            .collect::<Result<Vec<_>, _>>()
-            .unwrap();
+    fn wire(newline: &str) -> String {
+        format!(
+            "event: content_block_delta{newline}data: {TEXT_DELTA}{newline}{newline}: keepalive{newline}{newline}event: message_stop{newline}data: {{\"type\":\"message_stop\"}}{newline}{newline}"
+        )
+    }
 
-        let decoded = anthropic_sse_event_stream(stream::iter(wire.into_iter().map(Ok)).boxed())
+    #[rstest]
+    #[case::lf("\n")]
+    #[case::crlf("\r\n")]
+    #[case::cr("\r")]
+    #[tokio::test]
+    async fn frames_split_anywhere_decode_into_typed_events(
+        #[case] newline: &str,
+        #[values(1, 3, 4096)] size: usize,
+    ) {
+        let wire = wire(newline);
+        let expected = vec![
+            MessagesStreamEvent::ContentBlockDelta {
+                index: 0,
+                delta: MessagesContentBlockDelta::TextDelta {
+                    text: "hello".into(),
+                },
+            },
+            MessagesStreamEvent::MessageStop { usage: None },
+        ];
+        let streamed = event_stream(in_pieces(wire.as_bytes(), size))
             .try_collect::<Vec<_>>()
             .await
             .unwrap();
+        assert_eq!(streamed, expected);
+        assert_eq!(decode(wire.as_bytes()).unwrap(), expected);
+    }
 
-        assert_eq!(decoded, events());
+    #[tokio::test]
+    async fn decodes_citations_delta_events() {
+        let wire = concat!(
+            "event: content_block_delta\n",
+            r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#,
+            "\n\n",
+        );
+        let events = event_stream(in_pieces(wire.as_bytes(), 3))
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        assert!(matches!(
+            events.as_slice(),
+            [MessagesStreamEvent::ContentBlockDelta {
+                delta: MessagesContentBlockDelta::Citations { .. },
+                ..
+            }]
+        ));
+    }
+
+    #[tokio::test]
+    async fn encoded_events_decode_back_to_themselves() {
+        let wire: Vec<u8> = events()
+            .iter()
+            .map(encode)
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap()
+            .into_iter()
+            .flatten()
+            .collect();
+        let streamed = event_stream(in_pieces(&wire, 3))
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        assert_eq!(streamed, events());
+        assert_eq!(decode(&wire).unwrap(), events());
     }
 
     #[test]
     fn an_event_is_named_by_its_type() {
-        let encoded =
-            encode_anthropic_sse(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
-
+        let encoded = encode(&MessagesStreamEvent::MessageStop { usage: None }).unwrap();
         assert_eq!(
             encoded,
             Bytes::from(format!(
@@ -139,4 +147,68 @@ mod tests {
             ))
         );
     }
+
+    #[rstest]
+    #[case::json(b"data: {\n\n")]
+    #[case::unknown_event(b"data: {\"type\":\"future_event\"}\n\n")]
+    #[case::malformed_tool(b"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":3}}\n\n")]
+    #[case::utf8(b"data: \xff\n\n")]
+    #[case::done_is_not_an_anthropic_frame(b"data: [DONE]\n\n")]
+    #[tokio::test]
+    async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
+        assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
+        let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
+        assert!(matches!(
+            streamed.last(),
+            Some(Err(Error::InvalidResponse(_))) | None
+        ));
+    }
+
+    #[rstest]
+    fn decoding_keeps_tool_deltas_and_block_extensions() {
+        let wire = concat!(
+            "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"lookup\",\"input\":{},\"future\":null}}\n\n",
+            "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"x\\\":\"}}\n\n",
+        );
+        let decoded = decode(wire.as_bytes()).unwrap();
+        let [
+            MessagesStreamEvent::ContentBlockStart { content_block, .. },
+            MessagesStreamEvent::ContentBlockDelta {
+                delta: MessagesContentBlockDelta::InputJsonDelta { partial_json },
+                ..
+            },
+        ] = decoded.as_slice()
+        else {
+            panic!("expected tool events");
+        };
+        assert_eq!(
+            content_block.extra.get("future"),
+            Some(&serde_json::Value::Null)
+        );
+        assert_eq!(partial_json, "{\"x\":");
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn a_body_error_ends_the_stream_after_delivered_events() {
+        let mut events = event_stream(
+            stream::iter([
+                Ok(Bytes::from_static(b"data: {\"type\":\"ping\"}\n\n")),
+                Err(std::io::Error::new(
+                    std::io::ErrorKind::ConnectionReset,
+                    "reset",
+                )),
+            ])
+            .boxed(),
+        );
+        assert_eq!(
+            events.next().await.unwrap().unwrap(),
+            MessagesStreamEvent::Ping
+        );
+        assert!(matches!(
+            events.next().await,
+            Some(Err(Error::InvalidResponse(_)))
+        ));
+        assert!(events.next().await.is_none());
+    }
 }

litellm-rust/crates/llms/src/base_llm/mod.rs

@@ -5,3 +5,4 @@ pub mod chat;
 pub mod messages;
 pub mod ocr;
 pub mod responses;
+pub mod sse;

litellm-rust/crates/llms/src/base_llm/responses/mod.rs

@@ -1 +1,2 @@
+pub mod streaming;
 pub mod transformation;

litellm-rust/crates/llms/src/base_llm/responses/streaming.rs

@@ -0,0 +1,135 @@
+use bytes::Bytes;
+use litellm_llms_types::formats::responses::streaming::ResponsesStreamEvent;
+
+pub use crate::base_llm::base_model_iterator::ByteStream;
+use crate::{
+    Error,
+    base_llm::sse::{self, is_done},
+};
+
+pub type EventStream = sse::EventStream<ResponsesStreamEvent>;
+
+pub fn event_stream(bytes: ByteStream) -> EventStream {
+    sse::event_stream(bytes, is_done)
+}
+
+pub fn decode(wire: &[u8]) -> Result<Vec<ResponsesStreamEvent>, Error> {
+    sse::decode(wire, is_done)
+}
+
+/// A bare `data:` frame. Python's proxy serializes every Responses event this way and never
+/// writes an `event:` line, so clients see one shape whichever runtime served them.
+pub fn encode(event: &ResponsesStreamEvent) -> Result<Bytes, Error> {
+    sse::encode_data(event)
+}
+
+#[cfg(test)]
+mod tests {
+    use futures_util::{StreamExt, TryStreamExt, stream};
+    use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
+    use rstest::rstest;
+    use serde_json::{Value, json};
+
+    use super::*;
+
+    fn in_pieces(wire: &[u8], size: usize) -> ByteStream {
+        let pieces: Vec<Bytes> = wire.chunks(size).map(Bytes::copy_from_slice).collect();
+        stream::iter(pieces.into_iter().map(Ok)).boxed()
+    }
+
+    #[rstest]
+    #[case::created(json!({"type":"response.created","response":{"id":"r","output":[]}}))]
+    #[case::text(json!({"type":"response.output_text.delta","delta":"hello\nworld","sequence_number":3}))]
+    #[case::tool(json!({"type":"response.function_call_arguments.delta","delta":"{\"x\":","item_id":"fc_1"}))]
+    #[case::unknown(json!({"type":"future.event","payload":[1,null],"future":{"nested":true}}))]
+    #[case::completed(json!({"type":"response.completed","response":{"id":"r","usage":{"input_tokens":3}}}))]
+    #[case::failed(json!({"type":"response.failed","response":{"error":{"code":"failure"}}}))]
+    #[case::incomplete(json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"limit"}}}))]
+    #[case::error(json!({"type":"error","error":{"message":"boom"}}))]
+    fn encoded_events_are_bare_data_frames_that_decode_back(#[case] wire: Value) {
+        let event: ResponsesStreamEvent = serde_json::from_value(wire.clone()).unwrap();
+        let encoded = encode(&event).unwrap();
+        assert!(encoded.starts_with(b"data: "));
+        assert!(encoded.ends_with(b"\n\n"));
+        assert_eq!(
+            serde_json::from_slice::<Value>(&encoded[b"data: ".len()..encoded.len() - 2]).unwrap(),
+            wire
+        );
+        assert_eq!(decode(&encoded).unwrap(), vec![event]);
+    }
+
+    #[rstest]
+    #[case::lf("\n")]
+    #[case::crlf("\r\n")]
+    #[case::cr("\r")]
+    #[tokio::test]
+    async fn streams_split_anywhere_skip_comments_and_done(
+        #[case] newline: &str,
+        #[values(1, 5, 4096)] size: usize,
+    ) {
+        let wire = format!(
+            "\u{feff}event: response.created{newline}data: {{\"type\":\"response.created\"}}{newline}{newline}: keepalive{newline}{newline}event: future.event{newline}data: {{\"type\":\"future.event\",\"payload\":null}}{newline}{newline}data: [DONE]{newline}{newline}"
+        );
+        let streamed = event_stream(in_pieces(wire.as_bytes(), size))
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        assert_eq!(streamed, decode(wire.as_bytes()).unwrap());
+        assert_eq!(streamed.len(), 2);
+        assert_eq!(streamed[0].event_type, ResponsesEventType::ResponseCreated);
+        assert_eq!(
+            streamed[1].event_type,
+            ResponsesEventType::Other("future.event".into())
+        );
+        assert_eq!(streamed[1].data.get("payload"), Some(&Value::Null));
+    }
+
+    #[rstest]
+    #[case::spaced(b"data: [DONE]\n\n")]
+    #[case::unspaced(b"data:[DONE]\n\n")]
+    #[case::trailing_space(b"data: [DONE] \n\n")]
+    fn every_done_spelling_closes_the_stream_without_an_event(#[case] wire: &[u8]) {
+        assert_eq!(decode(wire).unwrap(), vec![]);
+    }
+
+    #[rstest]
+    #[case::invalid_json(b"data: {\n\n")]
+    #[case::missing_type(b"data: {}\n\n")]
+    #[case::numeric_type(b"data: {\"type\":3}\n\n")]
+    #[case::invalid_utf8(b"data: \xff\n\n")]
+    #[tokio::test]
+    async fn malformed_streams_are_invalid_responses(#[case] wire: &[u8]) {
+        assert!(matches!(decode(wire), Err(Error::InvalidResponse(_))));
+        let streamed = event_stream(in_pieces(wire, 1)).collect::<Vec<_>>().await;
+        assert!(matches!(
+            streamed.last(),
+            Some(Err(Error::InvalidResponse(_))) | None
+        ));
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn a_body_error_ends_the_stream_after_delivered_events() {
+        let mut events = event_stream(
+            stream::iter([
+                Ok(Bytes::from_static(
+                    b"data: {\"type\":\"response.created\"}\n\n",
+                )),
+                Err(std::io::Error::new(
+                    std::io::ErrorKind::ConnectionReset,
+                    "reset",
+                )),
+            ])
+            .boxed(),
+        );
+        assert_eq!(
+            events.next().await.unwrap().unwrap().event_type,
+            ResponsesEventType::ResponseCreated
+        );
+        assert!(matches!(
+            events.next().await,
+            Some(Err(Error::InvalidResponse(_)))
+        ));
+        assert!(events.next().await.is_none());
+    }
+}

litellm-rust/crates/llms/src/base_llm/responses/transformation.rs

@@ -1,5 +1,5 @@
 use litellm_llms_types::formats::responses::{
-    ResponsesApiResponse, streaming_websocket::ResponsesWsEvent,
+    ResponsesApiResponse, streaming::ResponsesStreamEvent,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
@@ -38,11 +38,11 @@ pub trait BaseResponsesApiConfig: Sync {
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct ResponsesWsTransformResult {
-    pub events: Vec<ResponsesWsEvent>,
+    pub events: Vec<ResponsesStreamEvent>,
 }
 
 impl ResponsesWsTransformResult {
-    pub fn passthrough(event: ResponsesWsEvent) -> Self {
+    pub fn passthrough(event: ResponsesStreamEvent) -> Self {
         Self {
             events: vec![event],
         }
@@ -58,18 +58,18 @@ pub trait ResponsesWebSocketProviderConfig: Sync {
 
     fn transform_ws_request(
         &self,
-        event: &ResponsesWsEvent,
+        event: &ResponsesStreamEvent,
         model: &str,
     ) -> Result<ResponsesWsTransformResult, Error>;
 
     fn transform_ws_response(
         &self,
-        event: &ResponsesWsEvent,
+        event: &ResponsesStreamEvent,
         model: &str,
     ) -> Result<ResponsesWsTransformResult, Error>;
 }
 
-pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent {
+pub fn enforce_model(event: &ResponsesStreamEvent, model: &str) -> ResponsesStreamEvent {
     if !event.is_response_create() {
         return event.clone();
     }
@@ -103,7 +103,7 @@ pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent
 mod tests {
     use super::*;
 
-    fn event(value: serde_json::Value) -> ResponsesWsEvent {
+    fn event(value: serde_json::Value) -> ResponsesStreamEvent {
         serde_json::from_value(value).expect("valid event")
     }
 

litellm-rust/crates/llms/src/base_llm/sse.rs

@@ -0,0 +1,102 @@
+//! SSE decoding and encoding shared by every format's `streaming` module.
+//!
+//! DIVERGENCE: decoding is strict. An event whose `data` is not a valid event of the format is
+//! `InvalidResponse` and ends the stream, and so is `[DONE]` on a Messages stream. Python
+//! never fails a stream on a malformed event:
+//! - Chat Completions (`BaseModelResponseIterator`) and Responses
+//!   (`BaseResponsesAPIStreamingIterator`) skip the event and keep going
+//! - Messages (`BaseAnthropicMessagesStreamingIterator`) forwards the provider's bytes to the
+//!   client without parsing them
+//! - Anthropic to Chat (`anthropic/chat/handler.py` `ModelResponseIterator`) buffers the invalid
+//!   lines on the guess that one payload was split across them
+//!
+//! Skipping loses content or a tool-call fragment with nothing to show for it, so Rust reports it.
+
+use bytes::Bytes;
+use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream};
+use litellm_framer::{
+    frames,
+    sse::{SseCodec, SseEvent, decode_all},
+};
+use serde::{Serialize, de::DeserializeOwned};
+
+use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream};
+
+pub type EventStream<E> = BoxStream<'static, Result<E, Error>>;
+
+/// Which frames are control frames rather than events of the format.
+pub type IsControl = fn(&SseEvent) -> bool;
+
+/// OpenAI's end-of-stream marker on Chat Completions and Responses streams. It closes the
+/// stream and is not an event; Anthropic Messages streams never carry it.
+pub const DONE_FRAME: &[u8] = b"data: [DONE]\n\n";
+
+pub fn is_done(event: &SseEvent) -> bool {
+    event.data.trim() == "[DONE]"
+}
+
+pub fn event_stream<E: DeserializeOwned + Send + 'static>(
+    bytes: ByteStream,
+    is_control: IsControl,
+) -> EventStream<E> {
+    frames(bytes, SseCodec::default())
+        .map_err(framing)
+        .try_filter(move |event| ready(!is_control(event)))
+        .and_then(|event| ready(parse(&event)))
+        .boxed()
+}
+
+/// Decodes a whole stream at once. An event the stream ends inside is dropped.
+pub fn decode<E: DeserializeOwned>(wire: &[u8], is_control: IsControl) -> Result<Vec<E>, Error> {
+    decode_all(wire)
+        .map_err(framing)?
+        .iter()
+        .filter(|event| !is_control(event))
+        .map(parse)
+        .collect()
+}
+
+/// A bare `data:` frame, which is what LiteLLM's proxy sends to clients on OpenAI formats.
+pub fn encode_data<E: Serialize>(event: &E) -> Result<Bytes, Error> {
+    let data = serde_json::to_string(event).map_err(invalid)?;
+    Ok(Bytes::from(format!("data: {data}\n\n")))
+}
+
+fn parse<E: DeserializeOwned>(event: &SseEvent) -> Result<E, Error> {
+    serde_json::from_str(&event.data).map_err(invalid)
+}
+
+pub fn invalid(error: serde_json::Error) -> Error {
+    Error::InvalidResponse(ErrorDetail::invalid("stream event", error))
+}
+
+fn framing(error: litellm_framer::SseError) -> Error {
+    Error::InvalidResponse(ErrorDetail::failed("stream framing", error))
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    #[rstest]
+    #[case::exact("[DONE]", true)]
+    #[case::padded(" [DONE] ", true)]
+    #[case::json("{\"type\":\"x\"}", false)]
+    #[case::mentions("[DONE] was sent", false)]
+    #[case::empty("", false)]
+    fn done_is_an_exact_marker(#[case] data: &str, #[case] done: bool) {
+        let event = SseEvent {
+            data: data.into(),
+            ..SseEvent::default()
+        };
+        assert_eq!(is_done(&event), done);
+    }
+
+    #[rstest]
+    fn the_done_frame_we_send_decodes_as_a_control_frame() {
+        let events: Vec<serde_json::Value> = decode(DONE_FRAME, is_done).unwrap();
+        assert!(events.is_empty());
+    }
+}