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

# Conflicts:
#	litellm/integrations/langfuse/langfuse_otel.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"

.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

.github/workflows/test-unit.yml

@@ -456,11 +456,11 @@ jobs:
               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: |-
@@ -565,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
@@ -573,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
@@ -1453,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
 

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:

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",

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 BaseModel, 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 (
@@ -88,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,
@@ -100,6 +103,7 @@ from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccess
     ResponsesAPIResponse,
 )
 from litellm.types.utils import (
+    ADDRESSED_RESPONSE_ID_FIELD,
     CallTypesLiteral,
     LiteLLMBatch,
     LiteLLMFineTuningJob,
@@ -346,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
@@ -392,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 {}
@@ -405,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)
 
@@ -529,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
@@ -546,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,
@@ -561,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),
@@ -601,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
@@ -1375,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
@@ -1432,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

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

helm/litellm/templates/_helpers.tpl

@@ -475,6 +475,27 @@ harmless no-op for the Job and authoritative for the app pods.
   {{- 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 }}
 {{ toYaml . }}
 {{- end }}

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:
@@ -320,3 +321,168 @@ tests:
     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

@@ -275,10 +275,21 @@ 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

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

@@ -1,4 +1,5 @@
 import random
+import re
 import time
 from collections.abc import Generator, Mapping
 from contextlib import contextmanager
@@ -7,6 +8,7 @@ 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_lock_timeout,
@@ -14,6 +16,26 @@ from litellm_proxy_extras.prisma_toolchain import (
 
 MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
 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
@@ -49,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,39 +119,45 @@ 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 the coordinator lock and Prisma's own lock key as transaction-level locks on one
-    pinned connection while a Prisma command runs with its session advisory lock disabled, so a
-    transaction pooler can never strand that lock on an idle backend after the command exits."""
-    with migration_lock(database_url) as coordinator:
-        coordinator.acquire_prisma_lock()
+    """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
 

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/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?
 

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -1,4 +1,3 @@
-import functools
 import glob
 import os
 import random
@@ -13,10 +12,10 @@ 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 (
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@@ -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(

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

@@ -4089,6 +4089,7 @@ dependencies = [
  "litellm-cache",
  "litellm-cache-memory",
  "litellm-cache-response",
+ "litellm-core-utils",
  "litellm-host",
  "litellm-host-native",
  "litellm-http",

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/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/AGENTS.md

@@ -58,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/openai/responses/transformation.rs

@@ -165,34 +165,46 @@ impl BaseResponsesApiConfig for OpenAiResponsesApiConfig {
         input: Value,
         params: Map<String, Value>,
     ) -> Result<Value, Error> {
-        if !input.is_string() && !input.is_array() {
-            return Err(Error::InvalidRequest(
-                "responses input must be a string or an array".into(),
-            ));
-        }
-        if params
-            .get("stream")
-            .is_some_and(|stream| !stream.is_boolean())
-        {
-            return Err(Error::InvalidRequest("stream must be a boolean".into()));
-        }
-        Ok(Value::Object(
-            params
-                .into_iter()
-                .chain([
-                    ("model".into(), Value::String(model.into())),
-                    ("input".into(), input),
-                ])
-                .collect(),
-        ))
+        responses_request(model, input, params)
     }
 
     fn transform_response_api_response(&self, body: Value) -> Result<ResponsesApiResponse, Error> {
-        serde_json::from_value(body)
-            .map_err(|error| Error::InvalidResponse(error.to_string().into()))
+        responses_response(body)
     }
 }
 
+/// The OpenAI Responses request body, shared by every OpenAI-compatible host.
+pub fn responses_request(
+    model: &str,
+    input: Value,
+    params: Map<String, Value>,
+) -> Result<Value, Error> {
+    if !input.is_string() && !input.is_array() {
+        return Err(Error::InvalidRequest(
+            "responses input must be a string or an array".into(),
+        ));
+    }
+    if params
+        .get("stream")
+        .is_some_and(|stream| !stream.is_boolean())
+    {
+        return Err(Error::InvalidRequest("stream must be a boolean".into()));
+    }
+    Ok(Value::Object(
+        params
+            .into_iter()
+            .chain([
+                ("model".into(), Value::String(model.into())),
+                ("input".into(), input),
+            ])
+            .collect(),
+    ))
+}
+
+pub fn responses_response(body: Value) -> Result<ResponsesApiResponse, Error> {
+    serde_json::from_value(body).map_err(|error| Error::InvalidResponse(error.to_string().into()))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -1 +1,2 @@
 pub mod messages;
+pub mod responses;

litellm-rust/crates/llms/src/openrouter/responses/AGENTS.md

@@ -0,0 +1,10 @@
+# rules
+
+- Send the OpenAI Responses payload to `/api/v1/responses` and relay the OpenAI SSE events
+- Authenticate with a bearer token from `OPENROUTER_API_KEY`, falling back to `OR_API_KEY`
+- Only `OPENROUTER_API_BASE` overrides the default base, never the OpenAI base variables
+- No native WebSocket mode
+
+# references
+
+- https://openrouter.ai/docs/api_reference/responses/overview

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

@@ -0,0 +1 @@
+pub mod transformation;

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

@@ -0,0 +1,174 @@
+use litellm_auth::{CredentialPlacement, SecretValue};
+use litellm_core_utils::settings::resolve_non_empty;
+use litellm_llms_types::formats::responses::ResponsesApiResponse;
+use serde_json::{Map, Value};
+
+use crate::{
+    Error,
+    base_llm::{
+        auth::{AuthScheme, ValidatedEnvironment},
+        responses::transformation::BaseResponsesApiConfig,
+    },
+    openai::responses::transformation::{responses_request, responses_response},
+};
+
+const API_KEY_ENVS: [&str; 2] = ["OPENROUTER_API_KEY", "OR_API_KEY"];
+const API_BASE_ENV: &str = "OPENROUTER_API_BASE";
+const DEFAULT_BASE: &str = "https://openrouter.ai/api/v1";
+const RESPONSES_PATH: &str = "/responses";
+
+pub struct OpenRouterResponsesApiConfig;
+pub const OPENROUTER_RESPONSES_CONFIG: OpenRouterResponsesApiConfig = OpenRouterResponsesApiConfig;
+
+impl BaseResponsesApiConfig for OpenRouterResponsesApiConfig {
+    fn secret_names(
+        &self,
+        api_key: Option<&str>,
+        api_base: Option<&str>,
+    ) -> &'static [&'static str] {
+        match (
+            api_key.is_some_and(|key| !key.trim().is_empty()),
+            api_base.is_some_and(|base| !base.trim().is_empty()),
+        ) {
+            (true, true) => &[],
+            (true, false) => &[API_BASE_ENV],
+            (false, true) => &API_KEY_ENVS,
+            (false, false) => &["OPENROUTER_API_KEY", "OR_API_KEY", API_BASE_ENV],
+        }
+    }
+
+    fn validate_environment(
+        &self,
+        headers: Vec<(String, String)>,
+        api_key: Option<&str>,
+        lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<ValidatedEnvironment, Error> {
+        let key = resolve_non_empty(api_key, lookup, &API_KEY_ENVS).ok_or(
+            litellm_auth::Error::MissingApiKey {
+                provider: "OpenRouter",
+                environment_variable: API_KEY_ENVS[0],
+            },
+        )?;
+        Ok(ValidatedEnvironment {
+            headers: litellm_http::request::with_default_headers(
+                headers,
+                &[("content-type", "application/json")],
+            ),
+            auth: AuthScheme::Credential {
+                placement: CredentialPlacement::Bearer,
+                secret: SecretValue::new(key),
+            },
+        })
+    }
+
+    fn get_complete_url(
+        &self,
+        api_base: Option<&str>,
+        lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> String {
+        let base = resolve_non_empty(api_base, lookup, &[API_BASE_ENV])
+            .unwrap_or_else(|| DEFAULT_BASE.into());
+        format!("{}{RESPONSES_PATH}", base.trim_end_matches('/'))
+    }
+
+    fn transform_responses_api_request(
+        &self,
+        model: &str,
+        input: Value,
+        params: Map<String, Value>,
+    ) -> Result<Value, Error> {
+        responses_request(model, input, params)
+    }
+
+    fn transform_response_api_response(&self, body: Value) -> Result<ResponsesApiResponse, Error> {
+        responses_response(body)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    fn env(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option<String> {
+        move |name| {
+            values
+                .iter()
+                .find(|(key, _)| *key == name)
+                .map(|(_, value)| (*value).to_owned())
+        }
+    }
+
+    #[rstest]
+    #[case::default(None, &[], "https://openrouter.ai/api/v1/responses")]
+    #[case::explicit(Some("https://proxy.test/api/v1"), &[], "https://proxy.test/api/v1/responses")]
+    #[case::trailing_slash(Some("https://proxy.test/api/v1/"), &[], "https://proxy.test/api/v1/responses")]
+    #[case::blank_falls_to_secret(Some(" "), &[("OPENROUTER_API_BASE", "https://secret.test/v1")], "https://secret.test/v1/responses")]
+    #[case::openai_base_is_ignored(None, &[("OPENAI_BASE_URL", "https://openai.test/v1")], "https://openrouter.ai/api/v1/responses")]
+    fn url_is_the_base_plus_responses(
+        #[case] api_base: Option<&str>,
+        #[case] secrets: &'static [(&'static str, &'static str)],
+        #[case] expected: &str,
+    ) {
+        assert_eq!(
+            OPENROUTER_RESPONSES_CONFIG.get_complete_url(api_base, &env(secrets)),
+            expected
+        );
+    }
+
+    #[rstest]
+    #[case::explicit(Some("explicit"), &[("OPENROUTER_API_KEY", "managed")], "explicit")]
+    #[case::managed(None, &[("OPENROUTER_API_KEY", "managed"), ("OR_API_KEY", "legacy")], "managed")]
+    #[case::legacy_name(None, &[("OR_API_KEY", "legacy")], "legacy")]
+    #[case::blank_explicit(Some(" "), &[("OR_API_KEY", "legacy")], "legacy")]
+    fn bearer_comes_from_the_call_then_either_secret_name(
+        #[case] api_key: Option<&str>,
+        #[case] secrets: &'static [(&'static str, &'static str)],
+        #[case] expected: &str,
+    ) {
+        let environment = OPENROUTER_RESPONSES_CONFIG
+            .validate_environment(vec![], api_key, &env(secrets))
+            .unwrap();
+        let AuthScheme::Credential { placement, secret } = environment.auth else {
+            panic!("expected a credential");
+        };
+        assert_eq!(placement, CredentialPlacement::Bearer);
+        assert_eq!(secret.expose(), expected);
+        assert_eq!(
+            environment.headers,
+            vec![("content-type".to_owned(), "application/json".to_owned())]
+        );
+    }
+
+    #[rstest]
+    #[case::nothing(&[])]
+    #[case::openai_key_only(&[("OPENAI_API_KEY", "openai")])]
+    fn a_missing_key_names_the_primary_variable(
+        #[case] secrets: &'static [(&'static str, &'static str)],
+    ) {
+        assert!(matches!(
+            OPENROUTER_RESPONSES_CONFIG.validate_environment(vec![], None, &env(secrets)),
+            Err(Error::Auth(litellm_auth::Error::MissingApiKey {
+                provider: "OpenRouter",
+                environment_variable: "OPENROUTER_API_KEY",
+            }))
+        ));
+    }
+
+    #[rstest]
+    #[case::both_given(Some("key"), Some("base"), &[])]
+    #[case::key_given(Some("key"), None, &["OPENROUTER_API_BASE"])]
+    #[case::base_given(None, Some("base"), &["OPENROUTER_API_KEY", "OR_API_KEY"])]
+    #[case::neither(None, None, &["OPENROUTER_API_KEY", "OR_API_KEY", "OPENROUTER_API_BASE"])]
+    fn only_missing_settings_are_looked_up(
+        #[case] api_key: Option<&str>,
+        #[case] api_base: Option<&str>,
+        #[case] expected: &[&str],
+    ) {
+        assert_eq!(
+            OPENROUTER_RESPONSES_CONFIG.secret_names(api_key, api_base),
+            expected
+        );
+    }
+}

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

@@ -23,7 +23,7 @@ fn run_chat_completions(
 ) -> PyResult<Py<PyAny>> {
     let host = InferenceHost::new(
         call.view()?,
-        "litellm.rust_bridge.chat_completions.route_host",
+        "litellm.rust_bridge.host.routes.chat_completions",
     );
     run_inference::<ChatCompletionsRoute, _>(
         py,

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

@@ -21,10 +21,7 @@ use serde_json::{Map, Value};
 use super::{NativeCall, RequestView};
 use crate::{
     errors::{RustUpstreamError, route_error_to_pyerr},
-    marshal::{
-        RouteOptions, optional_timeout, project_optional_fields, public_response,
-        python_timeout_seconds,
-    },
+    marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
 };
 
 pub(super) struct InferenceHost {
@@ -115,10 +112,7 @@ impl InferenceHost {
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> PyResult<Map<String, Value>> {
-        let names: Vec<String> = py.import(self.module)?.getattr("PARAMETERS")?.extract()?;
-        project_optional_fields(names.iter().map(String::as_str), |name| {
-            self.argument(py, arguments, name)
-        })
+        self.view.body_params(py, arguments, self.module)
     }
 
     pub fn response(&self, py: Python<'_>, response: &impl Serialize) -> PyResult<Py<PyAny>> {

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

@@ -26,34 +26,9 @@ use crate::{
     routes::RequestView,
 };
 
-const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host";
+const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages";
 const REQUEST_ERROR_MARKER: &str = "messages_request_error";
 
-const BODY_FIELDS: [&str; 22] = [
-    "max_tokens",
-    "metadata",
-    "stop_sequences",
-    "stream",
-    "system",
-    "temperature",
-    "thinking",
-    "tool_choice",
-    "tools",
-    "top_k",
-    "inference_geo",
-    "top_p",
-    "mcp_servers",
-    "context_management",
-    "compaction",
-    "container",
-    "output_format",
-    "speed",
-    "output_config",
-    "cache_control",
-    "reasoning_effort",
-    "safeguards",
-];
-
 fn merge_headers(
     forwarded: Option<Map<String, Value>>,
     extra_headers: Option<Map<String, Value>>,
@@ -143,7 +118,7 @@ impl MessagesPythonHost {
         let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?;
         let messages =
             argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?;
-        let fields = project_optional_fields(BODY_FIELDS, argument)?;
+        let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?;
         let body = [
             ("model".to_string(), Value::String(model.clone())),
             ("messages".to_string(), from_py(&messages)?),

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

@@ -13,6 +13,7 @@ use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCal
 use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
 use litellm_host_python::{
     CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
+    from_py,
 };
 use litellm_host_python::{Error as HostError, present};
 use pyo3::{
@@ -20,6 +21,7 @@ use pyo3::{
     prelude::*,
     types::{PyDict, PyMapping, PyTuple},
 };
+use serde_json::{Map, Value};
 
 /// The public call as Python bound it: `base` holds the positional arguments by name plus
 /// the signature defaults, `kwargs` the caller's keyword dict.
@@ -92,6 +94,22 @@ impl RequestView {
         present(arguments, self.request(py)?, name)
     }
 
+    /// The body params `module.body_params` keeps from the keywords the call would send, so
+    /// each route's Python host decides which params reach the provider.
+    pub(crate) fn body_params(
+        &self,
+        py: Python<'_>,
+        arguments: &Bound<'_, PyDict>,
+        module: &str,
+    ) -> PyResult<Map<String, Value>> {
+        let request = effective_py_args(self.request(py)?, arguments)?;
+        from_py(
+            &py.import(module)?
+                .getattr("body_params")?
+                .call1((request,))?,
+        )
+    }
+
     pub(crate) fn close(&mut self) {
         self.0 = None;
     }

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

@@ -77,7 +77,7 @@ impl OcrPythonHost {
             _ => "",
         };
         let mapped = py
-            .import("litellm.rust_bridge.ocr.route_host")
+            .import("litellm.rust_bridge.host.routes.ocr")
             .and_then(|module| module.getattr("map_failure"))
             .and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider)))
             .and_then(|mapped| mapped.extract::<Py<PyBaseException>>().map_err(PyErr::from));
@@ -115,7 +115,7 @@ impl PythonBinding for OcrPythonHost {
         py: Python<'_>,
         response: LiteLLMOcrResponse,
     ) -> PyResult<Py<PyAny>> {
-        public_response(py, "litellm.rust_bridge.ocr.route_host", &response)
+        public_response(py, "litellm.rust_bridge.host.routes.ocr", &response)
     }
 
     fn encode_stream_head(

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

@@ -21,7 +21,7 @@ use super::{
 };
 use crate::errors::RustBridgeDeclined;
 
-const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host";
+const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
 
 fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
     let resolved = call.resolved()?;

litellm/__init__.py

@@ -118,6 +118,7 @@ success_callback: List[CALLBACK_TYPES] = []
 failure_callback: List[CALLBACK_TYPES] = []
 service_callback: List[CALLBACK_TYPES] = []
 audit_log_callbacks: List[CALLBACK_TYPES] = []
+batch_line_item_callbacks: Sequence[CALLBACK_TYPES] = ()
 # logging_callback_manager is lazy-loaded via __getattr__
 _custom_logger_compatible_callbacks_literal = Literal[
     "lago",
@@ -392,6 +393,7 @@ key_alias_pattern: str | None = None
 enable_gemini_default_thinking_level_low: bool = (
     False  # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping
 )
+gemini_include_thoughts_default: bool = True
 ####################
 logging: bool = True
 enable_loadbalancing_on_batch_endpoints: Optional[bool] = None

litellm/_logging.py

@@ -6,12 +6,14 @@ import logging
 import os
 import re
 import sys
-from collections.abc import Iterator
+from collections.abc import Iterator, Mapping
 from datetime import datetime
 from logging import Formatter
 from typing import Final, TextIO
 from urllib.parse import unquote
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm.constants import (
     LITELLM_TRUNCATED_PAYLOAD_FIELD,
@@ -84,6 +86,7 @@ _redact_string = redact_string
 
 _REDACTED_RECORD_ATTR: Final = "litellm_redacted"
 _REDACTED_STAMP: Final = object()
+_DECODED_LOG_PAYLOAD: Final = TypeAdapter(object)
 _UNREDACTED_SCALAR_TYPES: Final = (bool, int, float, type(None))
 
 
@@ -91,6 +94,10 @@ def _is_redacted(record: logging.LogRecord) -> bool:
     return getattr(record, _REDACTED_RECORD_ATTR, None) is _REDACTED_STAMP
 
 
+def _record_attributes(record: logging.LogRecord) -> Mapping[str, object]:
+    return record.__dict__
+
+
 def _scrubbing_changed_nothing(scrubbed: object, original: object) -> bool:
     try:
         return bool(scrubbed == original)
@@ -438,7 +445,7 @@ def _process_record(record: logging.LogRecord, *, base64_limit: int, text_limit:
                     key, substituted_color if key == "color_message" and substituted_color is not None else value
                 ),
             )
-            for key, value in record.__dict__.items()
+            for key, value in _record_attributes(record).items()
             if key not in _STANDARD_RECORD_ATTRS
             and key != _REDACTED_RECORD_ATTR
             and not isinstance(value, _UNREDACTED_SCALAR_TYPES)
@@ -669,7 +676,7 @@ def _try_parse_json_message(message: str) -> dict[str, object] | None:
     msg_stripped: Final = message.strip()
     if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")):
         return None
-    parsed: Final[object] = safe_json_loads(message, default=None)
+    parsed: Final = _DECODED_LOG_PAYLOAD.validate_python(safe_json_loads(message, default=None))
     if parsed is None or not isinstance(parsed, dict):
         return None
     return parsed
@@ -699,7 +706,7 @@ def _try_parse_embedded_python_dict(message: str) -> dict[str, object] | None:
                 if depth == 0:
                     substr = message[start : j + 1]
                     try:
-                        result: object = ast.literal_eval(substr)
+                        result: object = _DECODED_LOG_PAYLOAD.validate_python(ast.literal_eval(substr))
                         if isinstance(result, dict) and len(result) > 0:
                             return result
                     except (ValueError, SyntaxError, TypeError):
@@ -757,7 +764,7 @@ class JsonFormatter(Formatter):
                     json_record[key] = value
 
         # Include extra attributes passed via logger.debug("msg", extra={...})
-        for key, value in record.__dict__.items():
+        for key, value in _record_attributes(record).items():
             if key not in _NON_EXTRA_RECORD_ATTRS and key not in json_record:
                 json_record[key] = value
 

litellm/_redis.py

@@ -19,6 +19,7 @@ from urllib.parse import urlsplit, urlunsplit
 
 import redis
 import redis.asyncio as async_redis
+from pydantic import TypeAdapter
 from redis.credentials import CredentialProvider
 
 from litellm import get_secret, get_secret_str
@@ -40,6 +41,7 @@ from ._logging import verbose_logger
 _generate_gcp_iam_access_token = generate_gcp_iam_access_token
 
 AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
+_UNWRAPPED_INIT: Final = TypeAdapter(object)
 
 _AWS_IAM_KWARG_NAMES: Final = (
     "aws_iam_auth",
@@ -64,7 +66,7 @@ def _unwrapped_init_args(cls: type) -> frozenset[str]:
     ``@deprecated_args`` too, which the same class of bug as ``_init_arg_names``
     would otherwise silently empty this allowlist through (see its docstring).
     """
-    spec: Final = inspect.getfullargspec(inspect.unwrap(cls.__init__))
+    spec: Final = inspect.getfullargspec(_UNWRAPPED_INIT.validate_python(inspect.unwrap(cls.__init__)))
     return frozenset(spec.args + spec.kwonlyargs)
 
 
@@ -112,7 +114,7 @@ def _init_arg_names(cls: type) -> frozenset[str]:
         name
         for klass in inspect.getmro(cls)
         if klass is not object
-        for spec in (inspect.getfullargspec(inspect.unwrap(klass.__init__)),)
+        for spec in (inspect.getfullargspec(_UNWRAPPED_INIT.validate_python(inspect.unwrap(klass.__init__))),)
         for name in spec.args + spec.kwonlyargs
     )
 
@@ -935,6 +937,10 @@ def get_redis_connection_pool(
     return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
 
 
+def _mask_for_display(masker: SensitiveDataMasker, redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
+    return masker.mask_dict(redis_kwargs)
+
+
 def _redis_kwargs_for_logging(redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
     return {
         key: "<credential provider>"
@@ -966,7 +972,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
         masker = SensitiveDataMasker()
 
         # Mask sensitive data in redis_kwargs
-        masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
+        masked_redis_kwargs = _mask_for_display(masker, redis_kwargs_for_logging)
 
         # Create main panel title
         title: Final = Text("Redis Configuration", style="bold blue")
@@ -1029,7 +1035,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
     except ImportError:
         # Fallback to simple logging if rich is not available
         masker = SensitiveDataMasker()
-        masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
+        masked_redis_kwargs = _mask_for_display(masker, redis_kwargs_for_logging)
         verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
     except Exception as e:
         verbose_logger.error("Error pretty printing Redis configuration: %s", e)

litellm/a2a_protocol/litellm_completion_bridge/handler.py

@@ -11,6 +11,7 @@ A2A Streaming Events (in order):
 """
 
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
+from types import MappingProxyType
 from typing import Any, Final
 
 import litellm
@@ -21,6 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
     A2AStreamingContext,
 )
 from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
+from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS
 from litellm.interactions.agents.utils import merge_agent_headers
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.utils import ModelResponse
@@ -44,6 +46,32 @@ _AGENT_ONLY_PARAMS: Final = frozenset(
 )
 
 
+def bridge_model_name(litellm_params: Mapping[str, object]) -> str:
+    """The litellm model string a bridge agent completes with: ``<custom_llm_provider>/<model>``.
+
+    A provider bridged on chat names a real upstream endpoint with ``model``, so a missing one stays empty for that
+    provider to reject or ignore instead of becoming a made-up ``agent`` endpoint.
+    """
+    custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
+    default_model: Final = "" if custom_llm_provider in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS else "agent"
+    configured_model: Final = litellm_params.get("model")
+    model: Final = configured_model if isinstance(configured_model, str) else default_model
+    if not custom_llm_provider or model.startswith(f"{custom_llm_provider}/"):
+        return model
+    return f"{custom_llm_provider}/{model}"
+
+
+def agent_completion_kwargs(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
+    """The agent's litellm_params that are valid litellm.acompletion kwargs (api_key, api_base, ...)."""
+    return MappingProxyType(
+        {
+            key: value
+            for key, value in litellm_params.items()
+            if key not in ("model", "custom_llm_provider") and key not in _AGENT_ONLY_PARAMS
+        }
+    )
+
+
 class A2ACompletionBridgeHandler:
     """
     Static methods for handling A2A requests via LiteLLM completion.
@@ -64,16 +92,7 @@ class A2ACompletionBridgeHandler:
         # Transform A2A message to OpenAI format
         openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
 
-        # Get completion params
-        custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
-        model: Final[str] = litellm_params.get("model", "agent")
-
-        # Build full model string if provider specified
-        # Skip prepending if model already starts with the provider prefix
-        if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
-            full_model = f"{custom_llm_provider}/{model}"
-        else:
-            full_model = model
+        full_model: Final = bridge_model_name(litellm_params)
 
         if stream:
             verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
@@ -87,13 +106,7 @@ class A2ACompletionBridgeHandler:
             "api_base": api_base,
             "stream": stream,
         }
-        # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
-        litellm_params_to_add: Final = {
-            k: v
-            for k, v in litellm_params.items()
-            if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
-        }
-        completion_params.update(litellm_params_to_add)
+        completion_params.update(agent_completion_kwargs(litellm_params))
         # Apply forward metadata AFTER the litellm_params merge so the helper
         # sees any agent-owner-configured ``extra_body.metadata`` and can keep
         # those keys authoritative over the client-supplied A2A metadata.

litellm/a2a_protocol/main.py

@@ -14,7 +14,9 @@ import datetime
 import uuid
 from collections.abc import AsyncIterator, Coroutine, Mapping
 from types import MappingProxyType, ModuleType
-from typing import TYPE_CHECKING, Any, Final, Optional, cast
+from typing import TYPE_CHECKING, Final, Optional, cast
+
+from typing_extensions import TypedDict, Unpack
 
 import litellm
 from litellm._logging import verbose_logger, verbose_proxy_logger
@@ -89,6 +91,10 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
 A2ACardResolver: Final = LiteLLMA2ACardResolver
 
 
+class _SendMessageOptions(TypedDict, total=False, extra_items=object):
+    pass
+
+
 def _set_usage_on_logging_obj(
     kwargs: Mapping[str, object],
     prompt_tokens: int,
@@ -563,7 +569,7 @@ async def asend_message(
 def send_message(
     a2a_client: "A2AClientType",
     request: "SendMessageRequest",
-    **kwargs: Any,
+    **kwargs: Unpack[_SendMessageOptions],
 ) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]:
     """
     Sync: Send a message to an A2A agent.

litellm/a2a_protocol/streaming_iterator.py

@@ -7,16 +7,21 @@ from collections.abc import AsyncIterator
 from datetime import datetime
 from typing import TYPE_CHECKING, Final
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
 from litellm.a2a_protocol.utils import A2ARequestUtils
 from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
 if TYPE_CHECKING:
     from a2a.compat.v0_3.types import SendStreamingMessageRequest, SendStreamingMessageResponse
 
+_OPAQUE_PAYLOAD: Final = TypeAdapter(object)
+
 
 class A2AStreamingIterator:
     """
@@ -160,8 +165,11 @@ class A2AStreamingIterator:
         # Add final chunk result if available
         if self.final_chunk:
             try:
-                chunk_dict: Final = self.final_chunk.model_dump(mode="json", exclude_none=True)
-                result["result"] = chunk_dict.get("result", {})
+                chunk_dict: Final = as_str_mapping(
+                    _OPAQUE_PAYLOAD.validate_python(self.final_chunk.model_dump(mode="json", exclude_none=True))
+                )
+                if chunk_dict is not None:
+                    result["result"] = chunk_dict.get("result", {})
             except Exception:
                 pass
 

litellm/anthropic_interface/messages/__init__.py

@@ -11,7 +11,6 @@ This is an __init__.py file to allow the following interface
 """
 
 from collections.abc import AsyncIterator, Coroutine, Iterator
-from typing import Any
 
 from litellm.messages import (
     anthropic_messages as _async_anthropic_messages,
@@ -102,8 +101,8 @@ def create(
 ) -> (
     AnthropicMessagesResponse
     | Iterator[bytes]
-    | AsyncIterator[Any]
-    | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]]
+    | AsyncIterator[object]
+    | Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[object] | Iterator[bytes]]
 ):
     """
     Async wrapper for Anthropic's messages API

litellm/batches/batch_line_items.py

@@ -0,0 +1,250 @@
+"""Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks."""
+
+import asyncio
+import hashlib
+import re
+from collections import Counter
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import dataclass
+from itertools import accumulate, chain, count
+from typing import Final, Literal, TypeAlias
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.batches.batch_utils import batch_cost_and_usage, iter_batch_input_lines, parse_batch_output_line
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class, is_valid_sha256_hash
+from litellm.llms.vertex_ai.files.transformation import vertex_batch_row_custom_id
+from litellm.types.utils import (
+    ModelInfo,
+    StandardBatchLineItem,
+    StandardBatchLineItemMetadata,
+    StandardBatchLineItemsPayload,
+    StandardBatchLineItemUsage,
+)
+
+BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"]
+
+_VERTEX_AI_UNKNOWN_CUSTOM_ID: Final = "unknown"  # what a Vertex AI batch row that carries no custom_id is read as
+_LINE: Final = re.compile(rb"[^\n]+")
+
+
+@dataclass(frozen=True, slots=True)
+class BatchLineItemSource:
+    batch_id: str
+    finished_at: int
+    model: str
+    custom_llm_provider: str | None
+    cost_provider: BatchFileProvider
+    model_info: ModelInfo | None
+    metadata: StandardBatchLineItemMetadata
+
+
+def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata:
+    return StandardBatchLineItemMetadata(
+        user_api_key_hash=_creator_key_hash(attribution),
+        user_api_key_alias=_text(attribution, "user_api_key_alias"),
+        user_api_key_team_id=_text(attribution, "user_api_key_team_id"),
+        user_api_key_team_alias=_text(attribution, "user_api_key_team_alias"),
+        user_api_key_user_id=_text(attribution, "user_api_key_user_id"),
+        user_api_key_org_id=_text(attribution, "user_api_key_org_id"),
+    )
+
+
+def _text(attribution: Mapping[str, object], name: str) -> str | None:
+    value: Final = attribution.get(name)
+    return value if isinstance(value, str) else None
+
+
+def _creator_key_hash(attribution: Mapping[str, object]) -> str | None:
+    """The creating credential as a SHA-256 hash. A virtual key is saved hashed already, while a credential saved as
+    is, such as an OAuth2 access token, is hashed here so the object never holds it."""
+    credential: Final = _text(attribution, "user_api_key_hash")
+    if credential is None or is_valid_sha256_hash(credential):
+        return credential
+    return hashlib.sha256(credential.encode()).hexdigest()
+
+
+def _line_key(entry: Mapping[str, object]) -> str | None:
+    """A line's custom_id. Bedrock calls it recordId, and native Vertex rows carry it in the row key or request labels."""
+    key: Final = entry.get("custom_id", entry.get("recordId"))
+    return key if isinstance(key, str) else vertex_batch_row_custom_id(entry)
+
+
+def batch_line_item_parts(
+    source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes], max_part_bytes: int
+) -> Iterator[StandardBatchLineItemsPayload]:
+    """Pair every output and error line with the input line that shares its custom_id, priced one line at a time, in
+    parts that each close once their raw request and result lines reach max_part_bytes. Input lines stay unparsed until
+    their result line is reached, so only one part's records are held at a time."""
+    sized: Final = _sized_line_items(source, input_content, result_contents)
+    for part in count():
+        lines = tuple(_until_full(sized, max_part_bytes))
+        if not lines:
+            return
+        yield StandardBatchLineItemsPayload(
+            batch_id=source.batch_id, finished_at=source.finished_at, part=part, lines=lines
+        )
+
+
+def _sized_line_items(
+    source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes]
+) -> Iterator[tuple[StandardBatchLineItem, int]]:
+    """Each output and error line's record, with the raw bytes of that line and of the input line it pairs with."""
+    request_spans: Final = _request_spans_by_custom_id(input_content)
+    for raw_result in chain.from_iterable(map(iter_batch_input_lines, result_contents)):
+        entry = parse_batch_output_line(raw_result)
+        if entry is None:
+            continue
+        custom_id = _line_key(entry)
+        span = _request_span_for(source, request_spans, custom_id)
+        request = None if span is None else parse_batch_output_line(input_content[span[0] : span[1]])
+        yield (
+            _line_item(source, custom_id, request, entry),
+            len(raw_result) + (0 if span is None else span[1] - span[0]),
+        )
+
+
+def _until_full(
+    sized: Iterator[tuple[StandardBatchLineItem, int]], max_part_bytes: int
+) -> Iterator[StandardBatchLineItem]:
+    """Take records until their raw bytes reach max_part_bytes, leaving the rest in sized for the next part."""
+    for line, part_bytes in accumulate(sized, _with_bytes_so_far):
+        yield line
+        if part_bytes >= max_part_bytes:
+            return
+
+
+def _with_bytes_so_far(
+    so_far: tuple[StandardBatchLineItem, int], line: tuple[StandardBatchLineItem, int]
+) -> tuple[StandardBatchLineItem, int]:
+    return line[0], so_far[1] + line[1]
+
+
+def _request_spans_by_custom_id(input_content: bytes) -> Mapping[str | None, tuple[int, int]]:
+    """Where each input line starts and ends, by its custom_id. A custom_id that several input lines share, as when a
+    Vertex AI embeddings entry is split into one row per input, is paired with none of them."""
+    keyed: Final = tuple(_keyed_spans(input_content))
+    counts: Final = Counter(key for key, _ in keyed)
+    return {key: span for key, span in keyed if counts[key] == 1}
+
+
+def _keyed_spans(input_content: bytes) -> Iterator[tuple[str | None, tuple[int, int]]]:
+    for line in _LINE.finditer(input_content):
+        entry = parse_batch_output_line(line[0]) if line[0].strip() else None
+        if entry is not None:
+            yield _line_key(entry), line.span()
+
+
+def _line_item(
+    source: BatchLineItemSource, custom_id: str | None, request: Mapping[str, object] | None, entry: dict[str, object]
+) -> StandardBatchLineItem:
+    priced: Final = batch_cost_and_usage(
+        (entry,), source.cost_provider, model_name=source.model, model_info=source.model_info
+    )
+    succeeded: Final = priced.successful_requests > 0
+    return StandardBatchLineItem(
+        batch_id=source.batch_id,
+        custom_id=custom_id,
+        status="success" if succeeded else "failure",
+        request=request,
+        response=entry,
+        usage=StandardBatchLineItemUsage(
+            prompt_tokens=priced.usage.prompt_tokens,
+            completion_tokens=priced.usage.completion_tokens,
+            total_tokens=priced.usage.total_tokens,
+        )
+        if succeeded
+        else None,
+        response_cost=priced.cost,
+        model=source.model,
+        custom_llm_provider=source.custom_llm_provider,
+        metadata=source.metadata,
+    )
+
+
+def _request_span_for(
+    source: BatchLineItemSource, request_spans: Mapping[str | None, tuple[int, int]], custom_id: str | None
+) -> tuple[int, int] | None:
+    """Where the input line that shares a result line's custom_id is. A Vertex AI line read as "unknown" pairs with
+    none."""
+    if custom_id is None or (source.cost_provider == "vertex_ai" and custom_id == _VERTEX_AI_UNKNOWN_CUSTOM_ID):
+        return None
+    return request_spans.get(custom_id)
+
+
+def without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload:
+    return StandardBatchLineItemsPayload(
+        batch_id=payload["batch_id"],
+        finished_at=payload["finished_at"],
+        part=payload["part"],
+        lines=tuple(
+            StandardBatchLineItem(
+                batch_id=line["batch_id"],
+                custom_id=line["custom_id"],
+                status=line["status"],
+                request=None,
+                response=None,
+                usage=line["usage"],
+                response_cost=line["response_cost"],
+                model=line["model"],
+                custom_llm_provider=line["custom_llm_provider"],
+                metadata=line["metadata"],
+            )
+            for line in payload["lines"]
+        ),
+    )
+
+
+async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None:
+    """Hand one part's lines to every batch_line_item_callbacks logger. A failing logger is logged and never stops the
+    rest."""
+    await asyncio.gather(
+        *(_deliver(callback, payload, redact_bodies) for callback in litellm.batch_line_item_callbacks)
+    )
+
+
+async def _deliver(
+    callback: litellm.CALLBACK_TYPES, payload: StandardBatchLineItemsPayload, redact_bodies: bool
+) -> None:
+    try:
+        logger: Final = _resolve_logger(callback)
+        if logger is None:
+            verbose_logger.error(
+                "batch_line_item_callbacks: %r cannot store batch line items, part %d of batch %s was not stored",
+                callback,
+                payload["part"],
+                payload["batch_id"],
+            )
+            return
+        await logger.async_log_batch_line_items(
+            without_request_and_response(payload)
+            if redact_bodies or _logger_turned_off_message_logging(logger)
+            else payload
+        )
+    except Exception as error:  # noqa: BLE001  # one failing sink must not stop the other sinks or the cost poller
+        verbose_logger.exception(
+            "batch_line_item_callbacks: %r failed to store the %d lines of part %d of batch %s: %s",
+            callback,
+            len(payload["lines"]),
+            payload["part"],
+            payload["batch_id"],
+            error,
+        )
+
+
+def _resolve_logger(callback: litellm.CALLBACK_TYPES) -> CustomLogger | None:
+    match callback:
+        case CustomLogger():
+            return callback
+        case "s3_v2" | "gcs_bucket":
+            return init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
+        case _:
+            return None
+
+
+def _logger_turned_off_message_logging(logger: CustomLogger) -> bool:
+    return (
+        getattr(logger, "turn_off_message_logging", False) is True
+        or getattr(logger, "message_logging", True) is not True
+    )

litellm/batches/batch_utils.py

@@ -1,10 +1,12 @@
 import json
-from collections.abc import Iterable, Iterator, Mapping
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from dataclasses import dataclass
 from dataclasses import replace as dataclasses_replace
 from enum import Enum
 from typing import Any, Final, Literal, cast
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
@@ -35,6 +37,7 @@ class BatchCostUsageResult:
 
 
 _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
+_DECODED_BATCH_LINE: Final = TypeAdapter(object)
 
 
 def _uses_native_vertex_output(
@@ -70,7 +73,17 @@ def batch_cost_is_final(batch: Batch) -> bool:
 
 
 async def calculate_batch_cost_and_usage(
-    file_content_dictionary: list[dict[str, object]],
+    file_content_dictionary: Sequence[dict[str, object]],
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
+    model_name: str | None = None,
+    model_info: ModelInfo | None = None,
+) -> BatchCostUsageResult:
+    """Calculate the cost and usage of a batch, see batch_cost_and_usage."""
+    return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info)
+
+
+def batch_cost_and_usage(
+    file_content_dictionary: Sequence[dict[str, object]],  # mutable-ok: the rows calculate_batch_cost_and_usage takes
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
     model_name: str | None = None,
     model_info: ModelInfo | None = None,
@@ -436,7 +449,7 @@ def _provider_output_file_id(output_file_id: str) -> str:
     return extracted
 
 
-async def _fetch_batch_managed_file_content(
+async def fetch_batch_managed_file_content(
     file_id: str,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
     litellm_params: Mapping[str, object] | None = None,
@@ -467,6 +480,9 @@ async def _fetch_batch_managed_file_content(
     return _file_content.content
 
 
+_fetch_batch_managed_file_content: Final = fetch_batch_managed_file_content
+
+
 async def _fetch_batch_output_file_content(
     batch: Batch,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
@@ -604,14 +620,14 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict[str, object
     whole batch's cost accounting.
     """
     for line in iter_batch_input_lines(file_content):
-        entry = _parse_batch_output_line(line)
+        entry = parse_batch_output_line(line)
         if entry is not None:
             yield entry
 
 
-def _parse_batch_output_line(line: bytes) -> dict[str, object] | None:
+def parse_batch_output_line(line: bytes) -> dict[str, object] | None:
     try:
-        parsed: Final[object] = json.loads(line)
+        parsed: Final = _DECODED_BATCH_LINE.validate_python(json.loads(line))
     except ValueError as e:
         verbose_logger.warning("skipping malformed batch output line: %s", str(e))
         return None

litellm/batches/main.py

@@ -127,7 +127,7 @@ async def acreate_batch(
     metadata: dict[str, str] | None = None,
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
-    output_expires_after: dict[str, Any] | None = None,
+    output_expires_after: FileExpiresAfter | None = None,
     **kwargs,
 ) -> LiteLLMBatch:
     """
@@ -191,9 +191,9 @@ def create_batch(
     metadata: dict[str, str] | None = None,
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
-    output_expires_after: dict[str, Any] | None = None,
+    output_expires_after: FileExpiresAfter | None = None,
     **kwargs,
-) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
+) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
     """
     Creates and executes a batch from an uploaded file of request
 
@@ -246,7 +246,7 @@ def create_batch(
             extra_body=extra_body,
         )
         if output_expires_after is not None:
-            _create_batch_request["output_expires_after"] = cast(FileExpiresAfter, output_expires_after)
+            _create_batch_request["output_expires_after"] = output_expires_after
         if model is not None:
             provider_config = ProviderConfigManager.get_provider_batches_config(
                 model=model,
@@ -568,7 +568,7 @@ def retrieve_batch(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
+) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
     """
     Retrieves a batch.
 
@@ -941,7 +941,7 @@ def cancel_batch(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
+) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
     """
     Cancels a batch.
 

litellm/budget_manager.py

@@ -13,6 +13,8 @@ import threading
 import time
 from typing import Final, Literal
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm.constants import (
     DAYS_IN_A_MONTH,
@@ -22,6 +24,8 @@ from litellm.constants import (
 )
 from litellm.utils import ModelResponse
 
+_BUDGET_SERVER_REPLY: Final = TypeAdapter(object)
+
 
 class BudgetManager:
     def __init__(
@@ -211,6 +215,5 @@ class BudgetManager:
         elif self.client_type == "hosted":
             url: Final = self.api_base + "/set_budget"
             data: Final = {"project_name": self.project_name, "user_dict": self.user_dict}
-            response = litellm.module_level_client.post(url, headers=self.headers, json=data)
-            response = response.json()
-            return response
+            response: Final = litellm.module_level_client.post(url, headers=self.headers, json=data)
+            return _BUDGET_SERVER_REPLY.validate_python(response.json())

litellm/caching/azure_blob_cache.py

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

litellm/caching/caching.py

@@ -18,14 +18,14 @@ from contextlib import contextmanager
 from enum import Enum
 from typing import Any, Final, Literal, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 import litellm
 from litellm._internal_context import current_service_target, service_target
 from litellm._logging import verbose_logger
 from litellm.constants import CACHED_STREAMING_CHUNK_DELAY
 from litellm.integrations.otel.runtime import phase_span
-from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
+from litellm.litellm_core_utils.model_param_helper import ModelParamHelper, peek_reasoning_summary_aliases
 from litellm.types.caching import *
 from litellm.types.utils import EmbeddingResponse, is_litellm_owned_kwarg
 
@@ -42,6 +42,8 @@ from .redis_semantic_cache import RedisSemanticCache
 from .s3_cache import S3Cache
 
 RESPONSE_CACHE_TARGET: Final = "llm_response"
+_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
+_DECODED_LITERAL: Final = TypeAdapter(object)
 
 
 @contextmanager
@@ -374,7 +376,7 @@ class Cache:
         )
         return "".join(f"{field}: {value}" for field, value in scope_values if value is not None)
 
-    def get_cache_key(self, **kwargs) -> str:
+    def get_cache_key(self, **kwargs: object) -> str:
         """
         Get the cache key for the given arguments.
 
@@ -398,21 +400,26 @@ class Cache:
         for param in kwargs:
             if param in scope_excluded_params:
                 continue
-            if param in combined_kwargs:
+            if param in combined_kwargs or param in _ALWAYS_KEYED_PROVIDER_PARAMS:
                 param_value: str | None = self._get_param_value(param, kwargs)
+                if param in _ALWAYS_KEYED_PROVIDER_PARAMS and not param_value:
+                    continue
                 if param_value is not None:
                     cache_key += f"{param}: {param_value}"
             elif not is_litellm_owned_kwarg(param):
                 if litellm.enable_caching_on_provider_specific_optional_params is True:  # feature flagged for now
                     if kwargs[param] is None:
                         continue  # ignore None params
-                    param_value = kwargs[param]
-                    cache_key += f"{param}: {param_value}"
+                    cache_key += f"{param}: {kwargs[param]}"
 
         if is_semantic_cache:
             cache_key += self._get_semantic_cache_tenant_scope(kwargs)
 
-        hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
+        reasoning_summary: Final = peek_reasoning_summary_aliases(kwargs)
+        reasoning_summary_material: Final = (
+            "" if reasoning_summary is None else f"reasoning_summary: {reasoning_summary}"
+        )
+        hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material)
         hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
         verbose_logger.debug(
             "\nCreated cache key: %s (source material length: %d)",
@@ -599,9 +606,11 @@ class Cache:
                 if isinstance(cached_response, dict):
                     pass
                 else:
-                    cached_response = json.loads(cached_response)  # Convert string to dictionary
+                    cached_response = _DECODED_LITERAL.validate_python(
+                        json.loads(cached_response)
+                    )  # Convert string to dictionary
             except Exception:
-                cached_response = ast.literal_eval(cached_response)
+                cached_response = _DECODED_LITERAL.validate_python(ast.literal_eval(cached_response))
             return cached_response
         return cached_result
 

litellm/caching/gcs_cache.py

@@ -7,6 +7,8 @@ import json
 from typing import Final
 from urllib.parse import quote
 
+from pydantic import TypeAdapter
+
 from litellm._logging import print_verbose, verbose_logger
 from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
 from litellm.llms.custom_httpx.http_handler import (
@@ -17,6 +19,8 @@ from litellm.llms.custom_httpx.http_handler import (
 
 from .base_cache import BaseCache
 
+_DECODED_JSON: Final = TypeAdapter(object)
+
 
 class GCSCache(BaseCache):
     def __init__(
@@ -70,7 +74,7 @@ class GCSCache(BaseCache):
             url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
             response: Final = self.sync_client.get(url=url, headers=headers)
             if response.status_code == 200:
-                cached_response: Final = json.loads(response.text)
+                cached_response: Final = _DECODED_JSON.validate_python(json.loads(response.text))
                 verbose_logger.debug(
                     "Got GCS Cache: key: %s, cached_response %s. Type Response %s",
                     key,
@@ -90,7 +94,7 @@ class GCSCache(BaseCache):
             url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
             response: Final = await self.async_client.get(url=url, headers=headers)
             if response.status_code == 200:
-                return json.loads(response.text)
+                return _DECODED_JSON.validate_python(json.loads(response.text))
             return None
         except Exception as e:
             verbose_logger.error("GCS Caching: async_get_cache() - Got exception from GCS: %s", e)

litellm/caching/qdrant_semantic_cache.py

@@ -15,6 +15,8 @@ import os
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Protocol, cast
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import print_verbose
 from litellm.constants import (
@@ -38,6 +40,7 @@ from ._embedding_router import (
 from .base_cache import BaseCache
 
 _WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"})
+_EMBEDDING_VECTOR: Final = TypeAdapter(object)
 
 if TYPE_CHECKING:
     from litellm.router import Router
@@ -295,7 +298,7 @@ class QdrantSemanticCache(BaseCache):
         )
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         value = str(value)
         assert isinstance(value, str)
@@ -334,7 +337,7 @@ class QdrantSemanticCache(BaseCache):
         )
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         data: Final = {
             "vector": embedding,
@@ -404,7 +407,7 @@ class QdrantSemanticCache(BaseCache):
         embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         value = str(value)
         assert isinstance(value, str)
@@ -440,7 +443,7 @@ class QdrantSemanticCache(BaseCache):
         embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         data: Final = {
             "vector": embedding,

litellm/caching/redis_cluster_cache.py

@@ -9,7 +9,7 @@ import asyncio
 import contextvars
 import os
 import weakref
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar
 
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 
     pipeline = Pipeline
     async_redis_client = Redis
-    Span = _Span | Any
+    Span = _Span
 else:
     pipeline = Any
     async_redis_client = Any
@@ -126,13 +126,13 @@ class RedisClusterCache(RedisCache):
             return call()
         return self._connection_gate.run(call, self._get_async_client_cache_key())
 
-    def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
+    def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
         """
         Overrides `_run_redis_mget_operation` in redis_cache.py
         """
         return self.redis_client.mget_nonatomic(keys=keys)
 
-    async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
+    async def _async_run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
         """
         Overrides `_async_run_redis_mget_operation` in redis_cache.py
         """

litellm/caching/redis_semantic_cache.py

@@ -117,7 +117,8 @@ class RedisSemanticCache(BaseCache):
                 password = password or os.environ["REDIS_PASSWORD"]
             except KeyError as e:
                 # Raise a more informative exception if any of the required keys are missing
-                missing_var: Final[object] = e.args[0]
+                key_error_args: Final[tuple[object, ...]] = e.args
+                missing_var: Final = key_error_args[0]
                 raise ValueError(
                     f"Missing required Redis configuration: {missing_var}. Provide {missing_var} or redis_url."
                 ) from e
@@ -310,9 +311,9 @@ class RedisSemanticCache(BaseCache):
                         return
             return
 
-        content = getattr(value, "content", None)
-        if content is not None:
-            cls._collect_responses_input_text(content, prompt_parts)
+        attribute_content: Final[object] = getattr(value, "content", None)
+        if attribute_content is not None:
+            cls._collect_responses_input_text(attribute_content, prompt_parts)
             return
 
         for text_key in ("text", "output", "input_text", "output_text"):

litellm/caching/valkey_semantic_cache.py

@@ -21,6 +21,7 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any, Final
 
+from pydantic import TypeAdapter
 from redis import Redis
 from redis.asyncio import Redis as AsyncRedis
 from redis.commands.search.field import TagField, VectorField
@@ -34,6 +35,8 @@ from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
 from ._embedding_router import resolve_embedding_timeout
 from .redis_semantic_cache import RedisSemanticCache
 
+_OPAQUE_SEARCH_FIELD: Final = TypeAdapter(object)
+
 
 @dataclass(frozen=True, slots=True)
 class _ValkeyCacheHit:
@@ -226,7 +229,7 @@ class ValkeySemanticCache(RedisSemanticCache):
         if not docs:
             return None
         doc: Final = docs[0]
-        response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME)
+        response_field: Final = _OPAQUE_SEARCH_FIELD.validate_python(getattr(doc, cls.RESPONSE_FIELD_NAME))
         distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME)
         return _ValkeyCacheHit(
             response=str(response_field),

litellm/constants.py

@@ -17,7 +17,7 @@ OAUTH_TOKEN_EXCHANGE_CACHE_SAFETY_MARGIN_SECONDS: Final = 60
 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_PROFILE: Final = "jwt_bearer_obo"
 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_SCOPE: Final = "https://graph.microsoft.com/.default"
 # Replies depend on the caller's delegated identity, which response-cache keys do not include.
-RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot"})
+RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot", "databricks_agent"})
 MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
 GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
 GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
@@ -106,6 +106,7 @@ DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY: Final = get_env_int("DEFAULT_S3_MAX_ADAPTIV
 S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: Final = get_env_int_in_range("S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", 30, 1, 3600)
 # https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
 MAX_S3_OBJECT_KEY_BYTES: Final = 1024
+MAX_S3_OBJECT_KEY_FILENAME_BYTES: Final = 255
 S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
 S3_PREFIX_DIGEST_CHARS: Final = 16
 # s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
@@ -379,6 +380,7 @@ REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = get_env_int_or_none("REALTIME
 REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float(
     "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0
 )
+BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0)
 
 # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
 WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
@@ -845,6 +847,8 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
     }
 )
 
+A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"})
+
 LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [
     "openai",
     "azure",
@@ -965,7 +969,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = {
     "context_management": None,
 }
 
-openai_compatible_endpoints: Final[list] = [
+openai_compatible_endpoints: Final[list[str]] = [
     "api.perplexity.ai",
     "api.endpoints.anyscale.com/v1",
     "api.deepinfra.com/v1/openai",
@@ -1131,6 +1135,7 @@ _openai_like_providers: Final[list[str]] = [
     "databricks",
     "lemonade",
     "watsonx",
+    "inferenceapis",
 ]  # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
 # well supported replicate llms
 replicate_models: Final[set] = set(
@@ -1621,6 +1626,7 @@ CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
 ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
 ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
 INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
+MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
 SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
 SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
 LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
@@ -1699,6 +1705,7 @@ AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:tran
 AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com"
 AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com"
 AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key"
+AWS_ERROR_TYPE_HEADER: Final = "x-amzn-errortype"
 AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio"
 AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription"
 AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription"
@@ -1860,7 +1867,13 @@ RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
 RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
 PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600)
 MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50))
+# raw request and result bytes per batch line-item object, so a large batch is built and stored one part at a time
+BATCH_LINE_ITEMS_MAX_PART_BYTES: Final = get_env_int_in_range(
+    "BATCH_LINE_ITEMS_MAX_PART_BYTES", 16 * 1024 * 1024, 1, 1024 * 1024 * 1024
+)
 MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
+PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress")
+TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"})
 BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0
 BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60
 STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))

litellm/containers/endpoint_factory.py

@@ -36,6 +36,7 @@ RESPONSE_TYPES: Final[dict[str, type]] = {
     "DeleteContainerFileResponse": DeleteContainerFileResponse,
 }
 _ENDPOINTS_DOCUMENT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_ENDPOINT_RESULT: Final = TypeAdapter(object)
 
 
 def _load_endpoints_config() -> dict:
@@ -173,7 +174,7 @@ def create_async_endpoint_function(
             init_response: Final = await loop.run_in_executor(None, func_with_context)
 
             if asyncio.iscoroutine(init_response):
-                response = await init_response
+                response = _ENDPOINT_RESULT.validate_python(await init_response)
             else:
                 response = init_response
 

litellm/containers/utils.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import Any, Final, TypeVar
 
 from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
@@ -99,7 +100,7 @@ class ContainerRequestUtils:
         response_obj: T,
         custom_llm_provider: str | None,
         litellm_metadata: dict[str, Any] | None = None,
-        extra_body: dict[str, Any] | None = None,
+        extra_body: Mapping[str, object] | None = None,
     ) -> T:
         """
         Encode container_id in response object with provider/model metadata for routing.

litellm/cost_calculator.py

@@ -556,12 +556,11 @@ def cost_per_token(
 
     assert custom_llm_provider is not None  # caller-supplied or get_llm_provider
 
-    model_without_prefix = model
     model_parts: Final = model.split("/", 1)
-    if len(model_parts) > 1:
-        model_without_prefix = model_parts[1]
-    else:
-        model_without_prefix = model
+    model_is_provider_relative: Final = (
+        model_with_provider == f"{custom_llm_provider}/{model}" and model_with_provider in model_cost_ref
+    )
+    model_without_prefix: Final = model if len(model_parts) == 1 or model_is_provider_relative else model_parts[1]
     """
     Code block that formats model to lookup in litellm.model_cost
     Option1. model = "bedrock/ap-northeast-1/anthropic.claude-instant-v1". This is the most accurate since it is region based. Should always be option 1
@@ -3009,6 +3008,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re
 
 
 class _ResponsesWsEventResponse(LiteLLMBaseModel):
+    id: str | None = None
     usage: Mapping[str, object] | None = None
     service_tier: str | None = None
 
@@ -3060,6 +3060,23 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
         )
         return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects))
 
+    @staticmethod
+    def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None:
+        response_ids: Final = tuple(
+            response.id for _, response in _billable_responses_ws_events(results) if response.id
+        )
+        return response_ids[-1] if response_ids else None
+
+    @staticmethod
+    def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject:
+        tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results)
+        return LiteLLMRealtimeStreamLoggingObject(
+            id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
+            usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
+            results=cast(OpenAIRealtimeStreamList, list(results)),  # cast-ok: raw WS frames, unvalidated field
+            service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None,
+        )
+
 
 _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
 

litellm/decisions/call.py

@@ -5,7 +5,7 @@ from types import MappingProxyType
 from typing import Final, Literal, TypeAlias
 
 import httpx
-from pydantic import ConfigDict, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
 import litellm
@@ -290,6 +290,7 @@ def _unsupported_params(
     model: str,
     custom_llm_provider: str,
     headers: Mapping[str, str],
+    supports_box_questions: bool,
 ) -> tuple[str, ...]:
     refuses_safety_identifier: Final = (
         ir_request.safety_identifier is not None and not provider_config.supports_safety_identifier
@@ -298,9 +299,7 @@ def _unsupported_params(
         None if provider_config.supports_images_in(ir_request.input, headers) else image_param(ir_request.input)
     )
     refuses_assets: Final = bool(ir_request.assets) and not provider_config.supports_assets(headers)
-    refuses_box_questions: Final = _has_box_question(ir_request.questions) and not (
-        provider_config.supports_box_questions(model, custom_llm_provider, headers)
-    )
+    refuses_box_questions: Final = _has_box_question(ir_request.questions) and not supports_box_questions
     refuses_audio: Final = bool(ir_request.audio) and not provider_config.supports_audio_input(
         model, custom_llm_provider
     )
@@ -328,9 +327,8 @@ def _is_not_box_question(question: DecisionsIRQuestion) -> bool:
 def _without_unsupported_params(
     ir_request: DecisionsIRRequest,
     provider_config: BaseDecisionsConfig,
-    model: str,
-    custom_llm_provider: str,
     headers: Mapping[str, str],
+    supports_box_questions: bool,
 ) -> DecisionsIRRequest:
     return replace(
         ir_request,
@@ -343,7 +341,7 @@ def _without_unsupported_params(
         assets=ir_request.assets if provider_config.supports_assets(headers) else MappingProxyType({}),
         questions=(
             ir_request.questions
-            if provider_config.supports_box_questions(model, custom_llm_provider, headers)
+            if supports_box_questions
             else tuple(filter(_is_not_box_question, ir_request.questions))
         ),
     )
@@ -365,6 +363,19 @@ def _without_malformed_images(
     return well_formed
 
 
+class _DeploymentBoxQuestions(BaseModel):
+    model_config = ConfigDict(frozen=True, from_attributes=True)
+
+    supports_box_questions: StrictBool | None = None
+
+
+def _deployment_supports_box_questions(kwargs: Mapping[str, object]) -> bool | None:
+    try:
+        return _DeploymentBoxQuestions.model_validate(kwargs.get("model_info") or {}).supports_box_questions
+    except ValidationError:
+        return None
+
+
 def _provider_ir_request(
     ir_request: DecisionsIRRequest,
     *,
@@ -374,12 +385,17 @@ def _provider_ir_request(
     headers: Mapping[str, str],
     kwargs: Mapping[str, object],
 ) -> DecisionsIRRequest | _UnsupportedParams | _NothingLeftAfterDrop:
-    unsupported: Final = _unsupported_params(ir_request, provider_config, model, custom_llm_provider, headers)
+    supports_box_questions: Final = provider_config.supports_box_questions(
+        model, custom_llm_provider, headers, _deployment_supports_box_questions(kwargs)
+    )
+    unsupported: Final = _unsupported_params(
+        ir_request, provider_config, model, custom_llm_provider, headers, supports_box_questions
+    )
     if not unsupported:
         return ir_request
     if not _drops_params(kwargs):
         return _UnsupportedParams(names=unsupported)
-    supported: Final = _without_unsupported_params(ir_request, provider_config, model, custom_llm_provider, headers)
+    supported: Final = _without_unsupported_params(ir_request, provider_config, headers, supports_box_questions)
     if _leaves_no_input(supported.input):
         return _NothingLeftAfterDrop(names=unsupported, missing="input to decide on")
     if not supported.questions:
@@ -387,20 +403,13 @@ def _provider_ir_request(
     return supported
 
 
-def _provider_specific_header(
-    value: object,
-) -> ProviderSpecificHeader | tuple[ProviderSpecificHeader, ...] | None:
-    try:
-        return _PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(value)
-    except ValidationError:
-        return None
-
-
 def _outbound_headers(
     extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object]
 ) -> Mapping[str, str]:
     forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
-        provider_specific_header=_provider_specific_header(kwargs.get("provider_specific_header")),
+        provider_specific_header=_PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(
+            kwargs.get("provider_specific_header")
+        ),
         custom_llm_provider=custom_llm_provider,
     )
     return MappingProxyType({**(extra_headers or {}), **forwarded})

litellm/exceptions.py

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

litellm/experimental_mcp_client/client.py

@@ -5,6 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
 import asyncio
 import base64
 import hashlib
+import inspect
 import json
 import os
 import time
@@ -37,6 +38,9 @@ _TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
 _SessionCallbackName: TypeAlias = Literal["sampling_callback", "elicitation_callback", "logging_callback"]
 
 
+SESSION_CLEANUP_TIMEOUT_MESSAGE: Final = "MCP session cleanup timed out"
+
+
 class _SessionCallbackMap(TypedDict):
     sampling_callback: ReadOnly[SamplingFnT | None]
     elicitation_callback: ReadOnly[ElicitationFnT | None]
@@ -185,6 +189,23 @@ def _strip_header_whitespace(headers: dict[str, str]) -> dict[str, str]:
     }
 
 
+async def _receive_session_message(
+    message: ServerNotification | Exception,
+    *,
+    stream_error: asyncio.Future[Exception],
+    on_notification: Callable[[ServerNotification], Awaitable[None]] | None,
+) -> None:
+    if not isinstance(message, Exception):
+        if on_notification is not None:
+            await on_notification(message)
+        return
+    if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)):
+        return
+    if not stream_error.done():
+        stream_error.set_result(message)
+    raise RuntimeError("MCP response stream failed")
+
+
 def _first_non_cancelled_cause(exc: BaseException, cleanup_errors: tuple[Exception, ...] = ()) -> BaseException | None:
     queue: Final[list[BaseException]] = [exc]
     while queue:
@@ -533,6 +554,17 @@ class MCPClient:
     async def discovery_auth_fingerprint(self) -> str:
         return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity())
 
+    async def subscription_auth_fingerprint(self) -> str:
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import (
+            ClientCredentialsBearerAuth,
+        )
+
+        request: Final = await self.prepare_request_auth()
+        auth: Final = self._resolved_auth
+        if isinstance(auth, ClientCredentialsBearerAuth) and auth.subscription_identity is not None:
+            request.headers[auth.header_name] = "managed-grant:" + auth.subscription_identity
+        return self._hash_discovery_auth(request, self._aws_identity())
+
     def _aws_identity(self) -> tuple[str, str, str | None, str | None] | None:
         return self._aws_auth.identity if isinstance(self._aws_auth, MCPSigV4Auth) else None
 
@@ -696,6 +728,7 @@ class MCPClient:
         *,
         on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
         on_cleanup: Callable[[], None] | None = None,
+        on_notification: Callable[[ServerNotification], Awaitable[None]] | None = None,
     ) -> TSessionResult:
         """
         Execute an operation within a transport and session context.
@@ -716,16 +749,6 @@ class MCPClient:
                     stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
                     _notify_stream_error(on_stream_error, stream_error)
 
-                    async def receive_message(
-                        message: ServerNotification | Exception,
-                    ) -> None:
-                        if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)):
-                            return
-                        if not stream_error.done():
-                            stream_error.set_result(message)
-                        # The SDK closes pending requests when its message handler raises.
-                        raise RuntimeError("MCP response stream failed")
-
                     session_kwargs: Final = (
                         self._session_callbacks_override
                         if self._session_callbacks_override is not None
@@ -737,7 +760,9 @@ class MCPClient:
                         read_stream,
                         write_stream,
                         read_timeout_seconds=self.timeout,
-                        message_handler=receive_message,
+                        message_handler=partial(
+                            _receive_session_message, stream_error=stream_error, on_notification=on_notification
+                        ),
                         sampling_callback=session_kwargs["sampling_callback"],
                         elicitation_callback=session_kwargs["elicitation_callback"],
                         logging_callback=session_kwargs["logging_callback"],
@@ -793,7 +818,7 @@ class MCPClient:
             raise (
                 in_flight_error
                 if in_flight_error is not None
-                else asyncio.CancelledError("MCP session cleanup timed out")
+                else asyncio.CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE)
             )
         return result
 
@@ -804,6 +829,7 @@ class MCPClient:
         quiet_on_error: bool = False,
         on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
         on_cleanup: Callable[[], None] | None = None,
+        on_notification: Callable[[ServerNotification], Awaitable[None]] | None = None,
     ) -> TSessionResult:
         """Open a session, run the provided coroutine, and clean up.
 
@@ -822,8 +848,9 @@ class MCPClient:
                     http_client=http_client,
                     on_stream_error=on_stream_error,
                     on_cleanup=on_cleanup,
+                    on_notification=on_notification,
                 )
-                if on_stream_error is not None or on_cleanup is not None
+                if on_stream_error is not None or on_cleanup is not None or on_notification is not None
                 else self._execute_session_operation(transport_ctx, operation, http_client=http_client)
             )
             result: Final = await execution
@@ -1094,7 +1121,7 @@ class MCPClient:
     async def call_tool(
         self,
         call_tool_request_params: MCPCallToolRequestParams,
-        host_progress_callback: Callable | None = None,
+        host_progress_callback: Callable[..., Awaitable[None]] | None = None,
         raise_on_error: bool = False,
         allow_input_required: bool = False,
         persistent_session: "PersistentMCPSession | None" = None,
@@ -1126,7 +1153,12 @@ class MCPClient:
             # Forward to Host if callback provided
             if host_progress_callback:
                 try:
-                    await host_progress_callback(progress, total)
+                    try:
+                        inspect.signature(host_progress_callback).bind(progress, total, message)
+                    except (TypeError, ValueError):
+                        await host_progress_callback(progress, total)
+                    else:
+                        await host_progress_callback(progress, total, message)
                 except Exception as e:
                     verbose_logger.warning("Failed to forward to Host: %s", e)
 
@@ -1579,6 +1611,20 @@ def _deliver_persistent_outcome(
     )
 
 
+async def _await_persistent_operation(operation: Awaitable[object], result: asyncio.Future[object]) -> object:
+    running: Final = asyncio.ensure_future(operation)
+
+    def cancel_operation(_: asyncio.Future[object]) -> None:
+        running.cancel()
+
+    result.add_done_callback(cancel_operation)
+    try:
+        (outcome,) = await asyncio.gather(running, return_exceptions=True)
+        return outcome
+    finally:
+        result.remove_done_callback(cancel_operation)
+
+
 class UpstreamSessionClosedError(RuntimeError):
     def __init__(self) -> None:
         super().__init__("upstream MCP session closed")
@@ -1707,12 +1753,14 @@ class PersistentMCPSession:
                     if callbacks is not None:
                         self._current_callbacks = callbacks
                     async with self._admission() if self._admission is not None else nullcontext():
-                        if not future.done():
-                            (outcome,) = await asyncio.gather(operation(session), return_exceptions=True)
-                            if _deliver_persistent_outcome(future, outcome, self._stream_error):
-                                self._active = None
-                                self._ended = True
-                                return
+                        if not future.done() and _deliver_persistent_outcome(
+                            future,
+                            await _await_persistent_operation(operation(session), future),
+                            self._stream_error,
+                        ):
+                            self._active = None
+                            self._ended = True
+                            return
                     self._active = None
                 if self._retiring and self._queue.empty():
                     return

litellm/files/main.py

@@ -11,7 +11,7 @@ import time
 import uuid as uuid_module
 from collections.abc import Coroutine
 from functools import partial
-from typing import Any, Final, Literal, cast
+from typing import Final, Literal, cast
 
 import httpx
 from openai import AsyncOpenAI, OpenAI
@@ -163,7 +163,7 @@ def create_file(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]:
+) -> OpenAIFileObject | Coroutine[object, None, OpenAIFileObject]:
     """
     Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API.
 
@@ -867,8 +867,8 @@ def file_content(
 ) -> (
     HttpxBinaryResponseContent
     | FileContentStreamingResult
-    | Coroutine[Any, Any, HttpxBinaryResponseContent]
-    | Coroutine[Any, Any, FileContentStreamingResult]
+    | Coroutine[object, None, HttpxBinaryResponseContent]
+    | Coroutine[object, None, FileContentStreamingResult]
 ):
     """
     Returns the contents of the specified file.

litellm/google_genai/adapters/handler.py

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

litellm/google_genai/adapters/transformation.py

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

litellm/google_genai/main.py

@@ -1,11 +1,10 @@
 import asyncio
 import contextvars
-from collections.abc import Iterator
 from functools import partial
 from typing import TYPE_CHECKING, Any, ClassVar, Final
 
 import httpx
-from pydantic import ConfigDict
+from pydantic import ConfigDict, TypeAdapter
 
 import litellm
 from litellm.constants import request_timeout
@@ -39,6 +38,7 @@ else:
 ####### ENVIRONMENT VARIABLES ###################
 # Initialize any necessary instances or variables here
 base_llm_http_handler = BaseLLMHTTPHandler()
+_GENERATE_CONTENT_RESULT: Final = TypeAdapter(object)
 #################################################
 
 
@@ -145,6 +145,7 @@ class GenerateContentHelper:
             ProviderConfigManager.get_provider_google_genai_generate_content_config(
                 model=model,
                 provider=litellm.LlmProviders(custom_llm_provider),
+                litellm_params=litellm_params.model_dump(include={"base_model", "api_base"}),
             )
         )
 
@@ -247,7 +248,7 @@ async def agenerate_content(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Any:
+) -> object:
     """
     Async: Generate content using Google GenAI
     """
@@ -285,7 +286,7 @@ async def agenerate_content(
         init_response: Final = await loop.run_in_executor(None, func_with_context)
 
         if asyncio.iscoroutine(init_response):
-            response = await init_response
+            response = _GENERATE_CONTENT_RESULT.validate_python(await init_response)
         else:
             response = init_response
 
@@ -315,7 +316,7 @@ def generate_content(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Any:
+) -> object:
     """
     Generate content using Google GenAI
     """
@@ -361,22 +362,24 @@ def generate_content(
             )
 
         # Call the standard handler
-        response: Final = base_llm_http_handler.generate_content_handler(
-            model=setup_result.model,
-            contents=contents,
-            tools=tools,
-            generate_content_provider_config=setup_result.generate_content_provider_config,
-            generate_content_config_dict=setup_result.generate_content_config_dict,
-            custom_llm_provider=setup_result.custom_llm_provider,
-            litellm_params=setup_result.litellm_params,
-            logging_obj=setup_result.litellm_logging_obj,
-            extra_headers=extra_headers,
-            extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
-            timeout=timeout or request_timeout,
-            _is_async=_is_async,
-            client=kwargs.get("client"),
-            litellm_metadata=kwargs.get("litellm_metadata", {}),
-            system_instruction=system_instruction,
+        response: Final = _GENERATE_CONTENT_RESULT.validate_python(
+            base_llm_http_handler.generate_content_handler(
+                model=setup_result.model,
+                contents=contents,
+                tools=tools,
+                generate_content_provider_config=setup_result.generate_content_provider_config,
+                generate_content_config_dict=setup_result.generate_content_config_dict,
+                custom_llm_provider=setup_result.custom_llm_provider,
+                litellm_params=setup_result.litellm_params,
+                logging_obj=setup_result.litellm_logging_obj,
+                extra_headers=extra_headers,
+                extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
+                timeout=timeout or request_timeout,
+                _is_async=_is_async,
+                client=kwargs.get("client"),
+                litellm_metadata=kwargs.get("litellm_metadata", {}),
+                system_instruction=system_instruction,
+            )
         )
 
         return response
@@ -405,7 +408,7 @@ async def agenerate_content_stream(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Any:
+) -> object:
     """
     Async: Generate content using Google GenAI with streaming response
     """
@@ -456,23 +459,25 @@ async def agenerate_content_stream(
 
         # Call the handler with async enabled and streaming
         # Return the coroutine directly for the router to handle
-        return await base_llm_http_handler.generate_content_handler(
-            model=setup_result.model,
-            contents=contents,
-            generate_content_provider_config=setup_result.generate_content_provider_config,
-            generate_content_config_dict=setup_result.generate_content_config_dict,
-            tools=tools,
-            custom_llm_provider=setup_result.custom_llm_provider,
-            litellm_params=setup_result.litellm_params,
-            logging_obj=setup_result.litellm_logging_obj,
-            extra_headers=extra_headers,
-            extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
-            timeout=timeout or request_timeout,
-            _is_async=True,
-            client=kwargs.get("client"),
-            stream=True,
-            litellm_metadata=kwargs.get("litellm_metadata", {}),
-            system_instruction=system_instruction,
+        return _GENERATE_CONTENT_RESULT.validate_python(
+            await base_llm_http_handler.generate_content_handler(
+                model=setup_result.model,
+                contents=contents,
+                generate_content_provider_config=setup_result.generate_content_provider_config,
+                generate_content_config_dict=setup_result.generate_content_config_dict,
+                tools=tools,
+                custom_llm_provider=setup_result.custom_llm_provider,
+                litellm_params=setup_result.litellm_params,
+                logging_obj=setup_result.litellm_logging_obj,
+                extra_headers=extra_headers,
+                extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
+                timeout=timeout or request_timeout,
+                _is_async=True,
+                client=kwargs.get("client"),
+                stream=True,
+                litellm_metadata=kwargs.get("litellm_metadata", {}),
+                system_instruction=system_instruction,
+            )
         )
 
     except Exception as e:
@@ -500,7 +505,7 @@ def generate_content_stream(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Iterator[Any]:
+) -> object:
     """
     Generate content using Google GenAI with streaming response
     """
@@ -545,23 +550,25 @@ def generate_content_stream(
             )
 
         # Call the handler with streaming enabled (sync version)
-        return base_llm_http_handler.generate_content_handler(
-            model=setup_result.model,
-            contents=contents,
-            generate_content_provider_config=setup_result.generate_content_provider_config,
-            generate_content_config_dict=setup_result.generate_content_config_dict,
-            tools=tools,
-            custom_llm_provider=setup_result.custom_llm_provider,
-            litellm_params=setup_result.litellm_params,
-            logging_obj=setup_result.litellm_logging_obj,
-            extra_headers=extra_headers,
-            extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
-            timeout=timeout or request_timeout,
-            _is_async=_is_async,
-            client=kwargs.get("client"),
-            stream=True,
-            litellm_metadata=kwargs.get("litellm_metadata", {}),
-            system_instruction=system_instruction,
+        return _GENERATE_CONTENT_RESULT.validate_python(
+            base_llm_http_handler.generate_content_handler(
+                model=setup_result.model,
+                contents=contents,
+                generate_content_provider_config=setup_result.generate_content_provider_config,
+                generate_content_config_dict=setup_result.generate_content_config_dict,
+                tools=tools,
+                custom_llm_provider=setup_result.custom_llm_provider,
+                litellm_params=setup_result.litellm_params,
+                logging_obj=setup_result.litellm_logging_obj,
+                extra_headers=extra_headers,
+                extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body),
+                timeout=timeout or request_timeout,
+                _is_async=_is_async,
+                client=kwargs.get("client"),
+                stream=True,
+                litellm_metadata=kwargs.get("litellm_metadata", {}),
+                system_instruction=system_instruction,
+            )
         )
 
     except Exception as e:

litellm/google_genai/streaming_iterator.py

@@ -67,7 +67,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
         request_body: dict,
         model: str,
         custom_llm_provider: str,
-        hidden_params: dict[str, Any] | None = None,
+        hidden_params: dict[str, object] | None = None,
     ):
         self.litellm_logging_obj = litellm_logging_obj
         self.request_body = request_body
@@ -118,7 +118,7 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
         litellm_metadata: dict,
         custom_llm_provider: str,
         request_body: dict | None = None,
-        hidden_params: dict[str, Any] | None = None,
+        hidden_params: dict[str, object] | None = None,
     ):
         super().__init__(
             litellm_logging_obj=logging_obj,
@@ -169,7 +169,7 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
         litellm_metadata: dict,
         custom_llm_provider: str,
         request_body: dict | None = None,
-        hidden_params: dict[str, Any] | None = None,
+        hidden_params: dict[str, object] | None = None,
     ):
         super().__init__(
             litellm_logging_obj=logging_obj,

litellm/harness/handlers/deepagents_handler.py

@@ -13,8 +13,8 @@ import asyncio
 import importlib
 from collections.abc import AsyncIterator, Callable, Mapping
 from dataclasses import dataclass
-from types import MappingProxyType, ModuleType
-from typing import TYPE_CHECKING, Any
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Protocol
 
 from litellm.harness.context import SessionContext
 from litellm.harness.errors import HarnessError, HarnessInstallFailed
@@ -51,10 +51,17 @@ if TYPE_CHECKING:
     from langgraph.types import Command
 
     from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
+    from litellm.llms.deepagents.harness.sandbox_backend import SandboxBackend, ToolFilterMiddleware, UsageCallback
 
 _MODEL_NODE = "model"
 
 
+class _SandboxBackendModule(Protocol):
+    SandboxBackend: type[SandboxBackend]
+    UsageCallback: type[UsageCallback]
+    ToolFilterMiddleware: type[ToolFilterMiddleware]
+
+
 @dataclass(frozen=True)
 class DeepAgentsDeps:
     """The optional-dependency entrypoints this handler uses."""
@@ -65,7 +72,7 @@ class DeepAgentsDeps:
     command_cls: Callable[..., Command]
     subagent_defaults: Mapping[str, object]
     convert_to_openai_messages: Any
-    backend: ModuleType
+    backend: _SandboxBackendModule
 
 
 def load_deps() -> DeepAgentsDeps:
@@ -77,7 +84,7 @@ def load_deps() -> DeepAgentsDeps:
         memory = importlib.import_module("langgraph.checkpoint.memory")
         lg_types = importlib.import_module("langgraph.types")
         messages = importlib.import_module("langchain_core.messages")
-        backend = importlib.import_module("litellm.llms.deepagents.harness.sandbox_backend")
+        from litellm.llms.deepagents.harness import sandbox_backend as backend
     except ImportError as e:
         raise HarnessInstallFailed(f"{INSTALL_HINT} ({e})") from e
     return DeepAgentsDeps(

litellm/integrations/anthropic_cache_control_hook.py

@@ -10,6 +10,7 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and
 """
 
 import copy
+import itertools
 import os
 import re
 from collections.abc import Callable, Container, Iterable, Mapping, Sequence
@@ -26,7 +27,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     cache_control_as_prompt_cache_breakpoint,
     with_prompt_cache_breakpoint,
 )
-from litellm.litellm_core_utils.prompt_templates.factory import tool_call_ids_answered_by_client
+from litellm.litellm_core_utils.prompt_templates.factory import (
+    tool_call_ids_answered_by_client,
+    tool_call_ids_answered_in,
+)
 from litellm.llms.anthropic.common_utils import (
     is_claude_code_one_shot_subagent_request,
     supports_anthropic_cache_control,
@@ -39,6 +43,7 @@ from litellm.types.integrations.anthropic_cache_control_hook import (
     CacheControlMessageInjectionPoint,
 )
 from litellm.types.llms.anthropic import (
+    ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
     ANTHROPIC_TOOL_SEARCH_TOOL_TYPES,
     AllAnthropicToolsValues,
     AnthropicSystemMessageContent,
@@ -75,6 +80,7 @@ OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
 _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
 _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
 _CONTENT_PART_ADAPTER: Final = TypeAdapter(dict[str, object])
+_ATTRIBUTE_VALUE: Final = TypeAdapter(object)
 
 AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
 
@@ -216,7 +222,7 @@ def _carries_cache_breakpoint(block: object) -> bool:
 
 def _attribute_or_key(value: object, key: str) -> object | None:
     if hasattr(value, key):
-        return getattr(value, key)
+        return _ATTRIBUTE_VALUE.validate_python(getattr(value, key))
     if isinstance(value, Mapping):
         return value.get(key)
     return None
@@ -237,6 +243,77 @@ def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answ
     )
 
 
+def _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
+    messages: Sequence[object], model: str | None = None, custom_llm_provider: str | None = None
+) -> frozenset[str]:
+    if not _marks_a_server_tool_call(messages):
+        return frozenset()
+    if model is not None and not _builds_anthropic_messages(model, custom_llm_provider):
+        return tool_call_ids_answered_in(messages)
+    return tool_call_ids_answered_by_client(messages, _keeps_mid_conversation_system(model, custom_llm_provider))
+
+
+def _marks_a_server_tool_call(messages: Sequence[object]) -> bool:
+    tool_calls: Final = itertools.chain.from_iterable(
+        _as_object_list(_attribute_or_key(message, "tool_calls")) or () for message in messages
+    )
+    return any(_is_marked_server_tool_call(tool_call) for tool_call in tool_calls)
+
+
+def _builds_anthropic_messages(model: str, custom_llm_provider: str | None) -> bool:
+    from litellm.exceptions import BadRequestError
+    from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
+    from litellm.llms.anthropic.chat.transformation import AnthropicConfig
+    from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
+    from litellm.types.utils import LlmProviders
+    from litellm.utils import ProviderConfigManager
+
+    if declared_authenticating_provider(model, custom_llm_provider) is not None:
+        return False
+    try:
+        served_model, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
+        chat_config: Final = ProviderConfigManager.get_provider_chat_config(
+            model=served_model, provider=LlmProviders(provider)
+        )
+    except (BadRequestError, ValueError):
+        return True
+    return isinstance(chat_config, AnthropicConfig) and not isinstance(chat_config, OpenAIGPTConfig)
+
+
+def _custom_llm_provider(non_default_params: Mapping[str, object], point: Mapping[str, object]) -> str | None:
+    provider: Final = non_default_params.get("custom_llm_provider") or point.get(CUSTOM_LLM_PROVIDER_STAMP)
+    return provider if isinstance(provider, str) else None
+
+
+def _keeps_mid_conversation_system(model: str | None, custom_llm_provider: str | None) -> bool:
+    if model is None:
+        return True
+    from litellm.utils import supports_mid_conversation_system
+
+    if supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider):
+        return True
+    return _model_info_flags_mid_conversation_system(model, custom_llm_provider)
+
+
+def _model_info_flags_mid_conversation_system(model: str, custom_llm_provider: str | None) -> bool:
+    from litellm.utils import get_model_info
+
+    try:
+        model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
+    except Exception:
+        return False
+    return model_info.get("supports_mid_conversation_system") is True
+
+
+def _is_marked_server_tool_call(tool_call: object) -> bool:
+    tool_call_id: Final = _attribute_or_key(tool_call, "id")
+    return (
+        isinstance(tool_call_id, str)
+        and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
+        and _attribute_or_key(tool_call, "cache_control") is not None
+    )
+
+
 def _tool_carries_cache_breakpoint(tool: object) -> bool:
     return _carries_cache_breakpoint(tool) or (
         isinstance(tool, dict) and _carries_cache_breakpoint(tool.get("function"))
@@ -262,6 +339,7 @@ def _may_carry_prompt_cache_breakpoint(message: Mapping[str, object]) -> bool:
 CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_points"
 
 EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints"
+CUSTOM_LLM_PROVIDER_STAMP: Final = "_litellm_custom_llm_provider"
 
 
 class AnthropicCacheControlHook(CustomPromptManagement):
@@ -335,13 +413,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         )
         remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message")
 
+        custom_llm_provider: Final = _custom_llm_provider(non_default_params, injection_points[0])
         stamped_dialect: Final = injection_points[0].get("_litellm_openai_dialect")
         openai_dialect: Final = (
             stamped_dialect
             if isinstance(stamped_dialect, bool)
             else AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(
                 model,
-                non_default_params.get("custom_llm_provider"),
+                custom_llm_provider,
                 non_default_params.get("api_base") or non_default_params.get("base_url"),
                 non_default_params.get("prompt_cache_options"),
             )
@@ -364,16 +443,25 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages(
             remaining_points, external_breakpoints, openai_dialect
         )
-        breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
+        answered_tool_ids: Final = _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
+            processed_messages, model, custom_llm_provider
+        )
+        breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(
+            processed_messages, answered_tool_ids=answered_tool_ids
+        )
         processed_messages = self._apply_message_injections(
             points=applied_message_points,
             messages=processed_messages,
             max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
             openai_dialect=openai_dialect,
+            answered_tool_ids=answered_tool_ids,
         )
         if (
             openai_dialect
-            and AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) > breakpoints_before
+            and AnthropicCacheControlHook.count_request_cache_breakpoints(
+                processed_messages, answered_tool_ids=answered_tool_ids
+            )
+            > breakpoints_before
             and non_default_params.get("prompt_cache_options") is None
         ):
             non_default_params["prompt_cache_options"] = PromptCacheOptions(mode="implicit")
@@ -386,7 +474,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         carried_points: Final[Sequence[CacheControlInjectionPoint]] = (
             *AnthropicCacheControlHook._points_with_a_slot_left(
                 remaining_points,
-                AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) + external_breakpoints,
+                AnthropicCacheControlHook.count_request_cache_breakpoints(
+                    processed_messages, answered_tool_ids=answered_tool_ids
+                )
+                + external_breakpoints,
                 openai_dialect,
             ),
             *carried_message_points,
@@ -421,14 +512,20 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         return _served_provider_for_model(model)
 
     @staticmethod
-    def count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
+    def count_request_cache_breakpoints(
+        messages: Iterable[object], system: object = None, answered_tool_ids: frozenset[str] | None = None
+    ) -> int:
         system_blocks: Final = (
             sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
         )
         message_list: Final = tuple(messages)
-        answered_tool_ids: Final = tool_call_ids_answered_by_client(message_list)
+        client_answers: Final = (
+            _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(message_list)
+            if answered_tool_ids is None
+            else answered_tool_ids
+        )
         return system_blocks + sum(
-            AnthropicCacheControlHook._count_cache_control_blocks(msg, answered_tool_ids) for msg in message_list
+            AnthropicCacheControlHook._count_cache_control_blocks(msg, client_answers) for msg in message_list
         )
 
     @staticmethod
@@ -503,6 +600,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         messages: list[AllMessageValues],
         max_blocks: int,
         openai_dialect: bool = False,
+        answered_tool_ids: frozenset[str] | None = None,
     ) -> list[AllMessageValues]:
         """Apply message-level cache control injection points in order.
 
@@ -513,8 +611,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         ``max_blocks`` is reached. Injection points are honored in config order,
         so earlier points win when slots are scarce.
         """
-        used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
-        answered_tool_ids: Final = tool_call_ids_answered_by_client(messages)
+        client_answers: Final = (
+            _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(messages)
+            if answered_tool_ids is None
+            else answered_tool_ids
+        )
+        used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(
+            messages, answered_tool_ids=client_answers
+        )
 
         limit_reached = False
         for point in points:
@@ -531,14 +635,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                     limit_reached = True
                     break
 
-                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], answered_tool_ids):
+                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], client_answers):
                     # Client already marked this message; don't overwrite it.
                     continue
 
                 messages[target_index] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
                     messages[target_index], control, openai_dialect
                 )
-                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], answered_tool_ids):
+                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], client_answers):
                     used_blocks += 1
 
             if limit_reached:
@@ -983,7 +1087,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
 
     @staticmethod
     def maybe_seed_default_injection_points(
-        non_default_params: dict[str, Any],
+        non_default_params: dict[str, object],
         messages: list[AllMessageValues],
         model: str,
         custom_llm_provider: str | None,
@@ -1007,7 +1111,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
             tools_keeping_marks: Final = tuple(
                 tool for tool in tools or () if not _chat_transform_drops_tool_cache_control(tool)
             )
-            non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_for_prompt_hook(
+            stamped: Final = AnthropicCacheControlHook._stamped_for_prompt_hook(
                 configured,
                 AnthropicCacheControlHook.count_external_cache_breakpoints(
                     tools_keeping_marks, non_default_params.get("cache_control"), non_default_params
@@ -1017,6 +1121,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 api_base,
                 non_default_params.get("prompt_cache_options"),
             )
+            non_default_params["cache_control_injection_points"] = (
+                AnthropicCacheControlHook._stamped(stamped, CUSTOM_LLM_PROVIDER_STAMP, custom_llm_provider)
+                if custom_llm_provider is not None and _marks_a_server_tool_call(messages)
+                else stamped
+            )
             return
         points: Final = AnthropicCacheControlHook.get_default_injection_points(
             messages=messages,
@@ -1109,7 +1218,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
     def maybe_inject_cache_control(
         messages: list[dict],
         system: str | list | None,
-        kwargs: dict[str, Any],
+        kwargs: dict[str, object],
         model: str | None = None,
         custom_llm_provider: str | None = None,
         tools: list[dict] | None = None,

litellm/integrations/arize/arize_phoenix.py

@@ -1,6 +1,7 @@
 import os
 import threading
 from collections import OrderedDict
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final
 
 from litellm._logging import verbose_logger
@@ -22,7 +23,7 @@ if TYPE_CHECKING:
 
     Protocol = _Protocol
     OpenTelemetryConfig = _OpenTelemetryConfig
-    Span = _Span | Any
+    Span = _Span
     OpenTelemetry = _OpenTelemetry
     LITELLM_TRACER_NAME: str
 else:
@@ -100,7 +101,7 @@ class ArizePhoenixLogger(OpenTelemetry):
         if getattr(self, "_use_injected_tracer_provider", False):
             return
 
-        shared_processor: Final = getattr(self, "_shared_span_processor", None)
+        shared_processor: Final[SpanProcessor | None] = getattr(self, "_shared_span_processor", None)
         if shared_processor is not None:
             try:
                 shared_processor.force_flush()
@@ -111,7 +112,8 @@ class ArizePhoenixLogger(OpenTelemetry):
                 )
 
         with getattr(self, "_project_providers_lock", threading.Lock()):
-            providers: Final = list(getattr(self, "_project_providers", {}).values())
+            project_providers: Final[Mapping[str, TracerProvider]] = getattr(self, "_project_providers", {})
+            providers: Final = list(project_providers.values())
 
         for provider in providers:
             try:
@@ -193,7 +195,7 @@ class ArizePhoenixLogger(OpenTelemetry):
         primary ``otel`` callback which handles proxy-level parent spans.
         """
 
-    def set_attributes(self, span: Span, kwargs, response_obj: Any | None):
+    def set_attributes(self, span: Span, kwargs, response_obj: object):
         ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj)
 
     @staticmethod

litellm/integrations/arize/arize_phoenix_client.py

@@ -68,7 +68,7 @@ class ArizePhoenixClient:
 
         try:
             # Use the underlying httpx client directly to avoid query param extraction
-            response = self.http_handler.get(url, headers=self.headers)
+            response: Final = self.http_handler.get(url, headers=self.headers)
             response.raise_for_status()
 
             data: Final = response.json()
@@ -76,15 +76,16 @@ class ArizePhoenixClient:
 
         except Exception as e:
             # Check if it's an HTTP error
-            response = getattr(e, "response", None)
-            if response is not None and hasattr(response, "status_code"):
-                if response.status_code == 404:
+            error_response: Final[object] = getattr(e, "response", None)
+            if error_response is not None and hasattr(error_response, "status_code"):
+                status_code: Final[object] = getattr(error_response, "status_code", None)
+                if status_code == 404:
                     return None
-                elif response.status_code == 403:
+                elif status_code == 403:
                     raise Exception(
                         f"Access denied to prompt version '{prompt_version_id}'. Check your Arize Phoenix permissions."
                     )
-                elif response.status_code == 401:
+                elif status_code == 401:
                     raise Exception("Authentication failed. Check your Arize Phoenix API key and permissions.")
                 else:
                     raise Exception(f"Failed to fetch prompt version '{prompt_version_id}': {e}")

litellm/integrations/batch_line_item_objects.py

@@ -0,0 +1,29 @@
+"""How the line items of one finished batch are laid out as a single object in a storage bucket."""
+
+import json
+from datetime import datetime, timezone
+from typing import Final
+
+from litellm.types.utils import StandardBatchLineItemsPayload
+
+
+def batch_line_items_ndjson(payload: StandardBatchLineItemsPayload) -> str:
+    return "".join(f"{_ndjson_line(line)}\n" for line in payload["lines"])
+
+
+def _ndjson_line(line: object) -> str:
+    """One line as JSON, with non-ASCII text kept readable. A line holding a lone surrogate, which a provider's
+    JSON escapes can carry, has no UTF-8 form, so that line alone is written with every non-ASCII character escaped."""
+    text: Final = json.dumps(line, default=str, ensure_ascii=False)
+    try:
+        text.encode("utf-8")
+    except UnicodeEncodeError:
+        return json.dumps(line, default=str)
+    return text
+
+
+def batch_line_items_object_name(payload: StandardBatchLineItemsPayload) -> str:
+    """batch_line_items/<finished day, UTC>/<batch id>/part-<part>.jsonl, the same name on every retry so a rewrite
+    replaces it."""
+    day: Final = datetime.fromtimestamp(payload["finished_at"], tz=timezone.utc).strftime("%Y-%m-%d")
+    return f"batch_line_items/{day}/{payload['batch_id'].replace('/', '_')}/part-{payload['part']:05d}.jsonl"

litellm/integrations/bitbucket/bitbucket_client.py

@@ -47,6 +47,10 @@ def _sanitize_file_path(file_path: str) -> str:
     return "/".join(urllib.parse.quote(part, safe="") for part in parts)
 
 
+def _content_type(headers: Mapping[str, str]) -> str:
+    return headers.get("content-type", "")
+
+
 class BitBucketClient:
     """
     Client for interacting with BitBucket API to fetch .prompt files.
@@ -119,7 +123,7 @@ class BitBucketClient:
             response.raise_for_status()
 
             # BitBucket returns file content as base64 encoded
-            if response.headers.get("content-type", "").startswith("text/"):
+            if _content_type(response.headers).startswith("text/"):
                 return response.text
             else:
                 # For binary files or when content-type is not text, try to decode as base64

litellm/integrations/clickhouse/clickhouse_spend_logger.py

@@ -136,6 +136,7 @@ def spend_log_row_from_payload(
     trace_id, span_id = _find_traceparent(metadata, kwargs)
     request_id = str(payload.get("id") or "")
     redact = litellm.turn_off_message_logging is True
+    track_end_user: Final = not litellm.disable_end_user_cost_tracking
     completion_start_ms = _to_ms(payload.get("completionStartTime"))
     response_cost: Final = payload.get("response_cost")
     unknown_success_cost: Final[bool] = payload.get("status") == "success" and kwargs.get("response_cost") is None
@@ -171,7 +172,7 @@ def spend_log_row_from_payload(
         team_alias=metadata.get("user_api_key_team_alias") or metadata.get("team_alias") or "",
         organization_id=metadata.get("user_api_key_org_id") or "",
         user=metadata.get("user_api_key_user_id") or "",
-        end_user=payload.get("end_user") or metadata.get("user_api_key_end_user_id") or "",
+        end_user=(payload.get("end_user") or metadata.get("user_api_key_end_user_id") or "") if track_end_user else "",
         model=payload.get("model") or "",
         model_group=payload.get("model_group") or "",
         model_id=payload.get("model_id") or "",
@@ -194,7 +195,14 @@ def spend_log_row_from_payload(
         span_id=span_id,
         request_tags=_request_tags(payload.get("request_tags")),
         metadata=_json_mapping(
-            MappingProxyType({**request_metadata, **metadata, "litellm_lens_internal": is_lens_analysis()})
+            MappingProxyType(
+                {
+                    **request_metadata,
+                    **metadata,
+                    "litellm_lens_internal": is_lens_analysis(),
+                    **({} if track_end_user else {"user_api_key_end_user_id": None}),
+                }
+            )
         ),
         messages="" if redact else _json(payload.get("messages")),
         response="" if redact else _json(payload.get("response")),

litellm/integrations/custom_guardrail.py

@@ -145,6 +145,10 @@ def _unified_hook_fields(guardrail: "CustomGuardrail", request_data: Mapping[str
     }
 
 
+def _instance_attributes(value: object) -> Mapping[str, object]:
+    return vars(value)
+
+
 def _strict_guardrail_modes_enabled() -> bool:
     """Whether guardrail-mode validation raises (default) or logs a warning.
 
@@ -396,7 +400,7 @@ class CustomGuardrail(CustomLogger):
 
     def inject_advisory_message(
         self,
-        data: dict[str, Any],  # mutable-ok: caller's dict is mutated in place, matching mark_pre_call_hook_ran
+        data: dict[str, object],  # mutable-ok: caller's dict is mutated in place, matching mark_pre_call_hook_ran
         message: str,
     ) -> bool:
         """
@@ -473,8 +477,8 @@ class CustomGuardrail(CustomLogger):
     def raise_sensitive_data_route_exception(
         self,
         route_to_model: str,
-        request_data: dict[str, Any],
-        detection_info: dict[str, Any] | None = None,
+        request_data: dict[str, object],
+        detection_info: dict[str, object] | None = None,
     ) -> None:
         """
         Raise an exception to reroute the request to a different model.
@@ -620,8 +624,8 @@ class CustomGuardrail(CustomLogger):
 
     def handle_sensitive_data_detection(
         self,
-        request_data: dict[str, Any],
-        detection_info: dict[str, Any] | None = None,
+        request_data: dict[str, object],
+        detection_info: dict[str, object] | None = None,
     ) -> None:
         """
         Handle sensitive data detection based on guardrail configuration.
@@ -691,9 +695,7 @@ class CustomGuardrail(CustomLogger):
         supported_event_hooks: list[GuardrailEventHooks],
     ) -> None:
         allowed_hooks: Final = frozenset(supported_event_hooks) | (
-            frozenset((GuardrailEventHooks.logging_only,))
-            if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
-            else frozenset()
+            frozenset((GuardrailEventHooks.logging_only,)) if self.accepts_logging_only_mode() else frozenset()
         )
 
         def _validate_event_hook_list_is_in_supported_event_hooks(
@@ -865,8 +867,13 @@ class CustomGuardrail(CustomLogger):
                     return True
         return False
 
-    def uses_apply_guardrail_interface(self) -> bool:
-        return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
+    @classmethod
+    def uses_apply_guardrail_interface(cls) -> bool:
+        return cls.apply_guardrail is not CustomGuardrail.apply_guardrail
+
+    @classmethod
+    def accepts_logging_only_mode(cls) -> bool:
+        return cls.uses_apply_guardrail_interface() and not cls.use_native_lifecycle_hooks
 
     @classmethod
     def supports_logging_only_scope(cls) -> bool:
@@ -888,7 +895,9 @@ class CustomGuardrail(CustomLogger):
             ) from e
         return unified_guardrail
 
-    async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
+    async def async_pre_call_deployment_hook(
+        self, kwargs: dict[str, object], call_type: CallTypes | None
+    ) -> dict | None:
         # should run guardrail
         litellm_guardrails: Final = kwargs.get("guardrails")
         if litellm_guardrails is None or not isinstance(litellm_guardrails, list):
@@ -1023,7 +1032,7 @@ class CustomGuardrail(CustomLogger):
             )
             return kwargs, result
         litellm_params: Final = kwargs.get("litellm_params") or {}
-        scratch_metadata: Final = {
+        scratch_metadata: Final[dict[str, object]] = {  # mutable-ok: apply_guardrail records its verdict here
             key: value
             for key, value in (litellm_params.get("metadata") or {}).items()
             if key != "standard_logging_guardrail_information"
@@ -1077,11 +1086,11 @@ class CustomGuardrail(CustomLogger):
 
     async def _scan_logged_call(
         self,
-        kwargs: dict,  # mutable-ok: CustomLogger.async_logging_hook contract
+        kwargs: dict[str, object],  # mutable-ok: CustomLogger.async_logging_hook contract
         response: object | None,
         translation: "BaseTranslation",
         output_translation: "BaseTranslation",
-        scratch_metadata: dict,  # mutable-ok: apply_guardrail records its verdict into request metadata
+        scratch_metadata: dict[str, object],  # mutable-ok: apply_guardrail records its verdict into request metadata
     ) -> None:
         scratch_fields: Final = self._copy_scratch_request_fields(kwargs)
         scratch_input, scratch_tools = scratch_fields or (None, None)
@@ -1579,7 +1588,7 @@ class CustomGuardrail(CustomLogger):
         """
         Update the guardrails litellm params in memory
         """
-        for key, value in vars(litellm_params).items():
+        for key, value in _instance_attributes(litellm_params).items():
             setattr(self, key, value)
 
     def get_guardrails_messages_for_call_type(

litellm/integrations/custom_logger.py

@@ -22,6 +22,7 @@ from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
     StandardAuditLogPayload,
+    StandardBatchLineItemsPayload,
     StandardCallbackDynamicParams,
     StandardLoggingPayload,
 )
@@ -195,6 +196,9 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
     async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"):
         """Called when an audit log is created. Override in subclasses to handle."""
 
+    async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
+        """Called once per part of a finished batch's lines, when this logger is listed in batch_line_item_callbacks."""
+
     #### PROMPT MANAGEMENT HOOKS ####
 
     async def async_get_chat_completion_prompt(
@@ -252,7 +256,7 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
         self,
         model: str,
         request_kwargs: dict,
-        messages: list[dict[str, Any]] | None = None,
+        messages: list[dict[str, object]] | None = None,
         input: str | list | None = None,
         specific_deployment: bool | None = False,
     ) -> PreRoutingHookResponse | None:

litellm/integrations/datadog/datadog_cost_management.py

@@ -3,7 +3,7 @@ import os
 import time
 from collections.abc import Mapping
 from datetime import datetime
-from typing import Any, Final, cast
+from typing import Final
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
@@ -45,6 +45,10 @@ _RESERVED_TAG_KEYS: Final[frozenset] = frozenset(
 )
 
 
+def _as_mapping(metadata: Mapping[str, object]) -> Mapping[str, object]:
+    return metadata
+
+
 class DatadogCostManagementLogger(CustomBatchLogger):
     def __init__(self, cost_tag_keys: list[str] | None = None, **kwargs):
         self.cost_tag_keys: list[str] = list(cost_tag_keys) if cost_tag_keys else []
@@ -180,9 +184,7 @@ class DatadogCostManagementLogger(CustomBatchLogger):
         self._add_tag(tags, "model", log.get("model"))
         self._add_tag(tags, "model_id", log.get("model_id"))
 
-        # cast because StandardLoggingMetadata is a TypedDict; we iterate it
-        # as a generic mapping below.
-        metadata: Final[Mapping[str, object]] = cast(dict[str, Any], log.get("metadata") or {})
+        metadata: Final = _as_mapping(log.get("metadata") or {})
 
         # Backwards-compat: team/user/model_group preserved regardless of allowlist.
         if metadata.get("user_api_key_alias"):

litellm/integrations/datadog/datadog_llm_obs.py

@@ -62,6 +62,7 @@ _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
 _EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
 _MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
 _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_DECODED_RESPONSE: Final = TypeAdapter(object)
 _SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset(
     {"agent", "assistant", "developer", "function", "model", "system", "tool", "user"}
 )
@@ -815,11 +816,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
             try:
                 import ast
 
-                response_obj = ast.literal_eval(response_obj)
+                response_obj = _DECODED_RESPONSE.validate_python(ast.literal_eval(response_obj))
             except (ValueError, SyntaxError):
                 try:
                     # fallback to json parsing
-                    response_obj = json.loads(str(response_obj))
+                    response_obj = _DECODED_RESPONSE.validate_python(json.loads(str(response_obj)))
                 except json.JSONDecodeError:
                     return ()
 

litellm/integrations/gcs_bucket/gcs_bucket.py

@@ -10,13 +10,14 @@ from litellm._logging import verbose_logger
 from litellm._uuid import uuid
 from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
 from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils
+from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name
 from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
 from litellm.litellm_core_utils.cloud_storage_security import (
     sanitize_cloud_object_component,
 )
 from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
 from litellm.types.integrations.gcs_bucket import *
-from litellm.types.utils import StandardLoggingPayload
+from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload
 
 if TYPE_CHECKING:
     from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
@@ -95,6 +96,20 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
         except Exception as e:
             verbose_logger.exception("GCS Bucket logging error: %s", e)
 
+    async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None:
+        """Write one part of a batch's lines as an NDJSON object under batch_line_items/ in the default bucket."""
+        gcs_logging_config: Final = await self.get_gcs_logging_config()
+        headers: Final = await self.construct_request_headers(
+            vertex_instance=gcs_logging_config["vertex_instance"],
+            service_account_json=gcs_logging_config["path_service_account"],
+        )
+        await self._log_json_data_on_gcs(
+            headers=headers,
+            bucket_name=gcs_logging_config["bucket_name"],
+            object_name=batch_line_items_object_name(payload),
+            logging_payload=await asyncio.to_thread(batch_line_items_ndjson, payload),
+        )
+
     def _drain_queue_batch(self) -> list[GCSLogQueueItem]:
         """
         Drain items from the queue (non-blocking), respecting batch_size limit.

litellm/integrations/gcs_bucket/gcs_bucket_base.py

@@ -2,6 +2,8 @@ import json
 import os
 from typing import TYPE_CHECKING, Any, Final
 
+from pydantic import TypeAdapter
+
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.integrations.gcs_bucket.gcs_bucket_mock_client import (
@@ -25,6 +27,7 @@ if TYPE_CHECKING:
 else:
     VertexBase = Any
 IAM_AUTH_KEY: Final = "IAM_AUTH"
+_UPLOAD_RESPONSE: Final = TypeAdapter(object)
 
 
 class GCSBucketBase(CustomBatchLogger):
@@ -297,7 +300,7 @@ class GCSBucketBase(CustomBatchLogger):
         bucket_name: str,
         object_name: str,
         logging_payload: StandardLoggingPayload | str,
-    ):
+    ) -> object:
         """
         Helper function to make POST request to GCS Bucket in the specified bucket.
         """
@@ -325,4 +328,4 @@ class GCSBucketBase(CustomBatchLogger):
         verbose_logger.debug("GCS Bucket status code %s", response.status_code)
         verbose_logger.debug("GCS Bucket response.text %s", response.text)
 
-        return response.json()
+        return _UPLOAD_RESPONSE.validate_python(response.json())

litellm/integrations/generic_api/generic_api_callback.py

@@ -252,7 +252,7 @@ class GenericAPILogger(CustomBatchLogger):
         if self.retry_delay <= 0:
             return
 
-        delay: Final = self.retry_delay * (2**attempt)
+        delay: Final = self.retry_delay * (2.0**attempt)
         await asyncio.sleep(delay)
 
     async def _post_with_retries(self, data: str) -> httpx.Response:

litellm/integrations/gitlab/gitlab_prompt_manager.py

@@ -219,7 +219,7 @@ class GitLabTemplateManager:
                 elif value.replace(".", "").isdigit():
                     try:
                         result[key] = float(value)
-                    except Exception:
+                    except ValueError:
                         result[key] = value
                 else:
                     result[key] = value.strip("\"'")
@@ -297,7 +297,7 @@ class GitLabPromptManager(CustomPromptManagement):
 
     def __init__(
         self,
-        gitlab_config: dict[str, Any],
+        gitlab_config: dict[str, object],
         prompt_id: str | None = None,
         ref: str | None = None,  # tag/branch/SHA override
         gitlab_client: GitLabClient | None = None,
@@ -662,7 +662,7 @@ class GitLabPromptCache:
 
     def __init__(
         self,
-        gitlab_config: dict[str, Any],
+        gitlab_config: dict[str, object],
         *,
         ref: str | None = None,
         gitlab_client: GitLabClient | None = None,

litellm/integrations/lago.py

@@ -147,9 +147,10 @@ class LagoLogger(CustomLogger):
 
             response.raise_for_status()
         except Exception as e:
-            error_response: Final = getattr(e, "response", None)
+            error_response: Final[object] = getattr(e, "response", None)
             if error_response is not None and hasattr(error_response, "text"):
-                verbose_logger.debug("\nError Message: %s", error_response.text)
+                error_text: Final[object] = getattr(error_response, "text", None)
+                verbose_logger.debug("\nError Message: %s", error_text)
             raise e
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):

litellm/integrations/langfuse/langfuse_otel.py

@@ -25,7 +25,7 @@ from litellm.types.utils import StandardCallbackDynamicParams
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
 
-    Span = _Span | Any
+    Span = _Span
 else:
     Span = Any
 

litellm/integrations/langsmith.py

@@ -10,7 +10,7 @@ from datetime import datetime, timezone
 from typing import Any, Final
 
 import httpx
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 import litellm
 from litellm._logging import verbose_logger
@@ -28,6 +28,8 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.types.integrations.langsmith import *
 from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload
 
+_RUN_RESPONSE: Final = TypeAdapter(object)
+
 
 def is_serializable(value):
     non_serializable_types: Final = (
@@ -524,7 +526,7 @@ class LangsmithLogger(CustomBatchLogger):
             # If we can't get an event loop, create a new one
             asyncio.run(self.async_send_batch())
 
-    def get_run_by_id(self, run_id):
+    def get_run_by_id(self, run_id) -> object:
         langsmith_api_key: Final = self.default_credentials["LANGSMITH_API_KEY"]
         langsmith_api_base: Final = self.default_credentials["LANGSMITH_BASE_URL"]
         langsmith_tenant_id: Final = self.default_credentials.get("LANGSMITH_TENANT_ID")
@@ -538,7 +540,7 @@ class LangsmithLogger(CustomBatchLogger):
             headers=headers,
         )
 
-        return response.json()
+        return _RUN_RESPONSE.validate_python(response.json())
 
     def make_dot_order(self, run_id: str):
         st: Final = datetime.now(timezone.utc)

litellm/integrations/literal_ai.py

@@ -2,6 +2,7 @@
 # This file contains the LiteralAILogger class which is used to log steps to the LiteralAI observability platform.
 import asyncio
 import os
+from collections.abc import Mapping
 from typing import Final
 
 import httpx
@@ -181,7 +182,8 @@ class LiteralAILogger(CustomBatchLogger):
 
         if messages and isinstance(messages, list) and isinstance(messages[0], dict):
             for message in messages:
-                if literal_prompt := getattr(message, "__literal_prompt__", None):
+                literal_prompt: Mapping[str, object] | None = getattr(message, "__literal_prompt__", None)
+                if literal_prompt:
                     prompt_id = literal_prompt.get("prompt_id")
                     variables = literal_prompt.get("variables")
                     message["uuid"] = literal_prompt.get("uuid")

litellm/integrations/lunary.py

@@ -3,10 +3,15 @@
 import importlib
 import traceback
 from datetime import datetime, timezone
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import packaging
 
+if TYPE_CHECKING:
+    from openai.types.chat.chat_completion_message_custom_tool_call import Custom
+
+    from litellm.types.utils import ChatCompletionCustomToolCallPayload, ChatCompletionDeltaCustomToolCallPayload
+
 
 # convert to {completion: xx, tokens: xx}
 def parse_usage(usage):
@@ -21,7 +26,9 @@ def parse_tool_calls(tool_calls):
         return None
 
     def clean_tool_call(tool_call):
-        custom: Final = getattr(tool_call, "custom", None)
+        custom: Final[
+            ChatCompletionCustomToolCallPayload | ChatCompletionDeltaCustomToolCallPayload | Custom | None
+        ] = getattr(tool_call, "custom", None)
         if custom is not None:
             name, arguments = custom.name, custom.input
         else:

litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py

@@ -24,6 +24,8 @@ from collections.abc import Callable
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Final
 
+from typing_extensions import TypedDict, Unpack
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import MAVVRIK_FOCUS_EXPORT_JOB_NAME
@@ -82,10 +84,16 @@ def _is_empty_metrics_marker(marker: object | None) -> bool:
     return False
 
 
+class _CustomLoggerOptions(TypedDict, total=False, extra_items=object):
+    pass
+
+
 class MavvrikFocusLogger(FocusLogger):
     """FOCUS-based export logger that routes to the Mavvrik destination."""
 
-    def __init__(self, *, clock: Callable[[], datetime] = get_utc_datetime, **kwargs: Any) -> None:
+    def __init__(
+        self, *, clock: Callable[[], datetime] = get_utc_datetime, **kwargs: Unpack[_CustomLoggerOptions]
+    ) -> None:
         self._clock: Final = clock
         frequency: Final = os.getenv("MAVVRIK_FOCUS_FREQUENCY", "daily").lower()
         if frequency != "daily":

litellm/integrations/opentelemetry_utils/gen_ai_semconv.py

@@ -29,21 +29,26 @@ Events:
   included only when content capture is enabled).
 """
 
+from collections.abc import Callable
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Final, Protocol
 
 from litellm.integrations.otel.mappers.utils import drop_none
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
 if TYPE_CHECKING:
+    from opentelemetry._logs import SeverityNumber as _SeverityNumber
     from opentelemetry.trace import Span as _Span
+    from opentelemetry.trace import SpanContext
 
     from litellm.integrations.opentelemetry import OpenTelemetryConfig
 
-    Span = _Span | Any
+    Span = _Span
 else:
-    Span = Any
+
+    class Span(Protocol):
+        def set_attribute(self, key: str, value: object) -> None: ...
 
 
 # OTEL_SEMCONV_STABILITY_OPT_IN is a comma-separated list of category-specific
@@ -128,7 +133,7 @@ class OTELGenAISemconvMixin:
 
         def _to_ns(self, dt: datetime) -> int: ...
 
-        def _otel_log_types(self) -> tuple[Any, Any]: ...
+        def _otel_log_types(self) -> tuple[Callable[..., object], type[_SeverityNumber]]: ...
 
     @property
     def _gen_ai_semconv_latest_experimental(self) -> bool:
@@ -227,7 +232,7 @@ class OTELGenAISemconvMixin:
         response_obj: dict,
         provider: str | None,
         otel_logger,
-        parent_ctx,
+        parent_ctx: "SpanContext",
     ) -> None:
         """Emit the consolidated ``gen_ai.client.inference.operation.details``
         log event, correlated to the request span via ``parent_ctx``.

litellm/integrations/otel/logger.py

@@ -292,7 +292,7 @@ class OpenTelemetryV2(CustomLogger):
     #  closed here. See ``log_pre_api_call``.
     # ====================================================================== #
 
-    def log_pre_api_call(self, model, messages, kwargs):
+    def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
         """Open the LLM-call span at the call boundary.
 
         Runs synchronously inside the request task, before the upstream call —
@@ -354,7 +354,13 @@ class OpenTelemetryV2(CustomLogger):
             _LLMCallSpan(span=span, start_time_ns=start_time_ns, provider=route.provider),
         )
 
-    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+    async def async_log_success_event(
+        self,
+        kwargs: Mapping[str, object],
+        response_obj: object,
+        start_time: datetime | float | None,
+        end_time: datetime | float | None,
+    ) -> None:
         if self._emit_mcp_tool_call(kwargs, start_time, end_time):
             return
         if self._emit_mcp_list_tools(kwargs, start_time, end_time):
@@ -396,7 +402,13 @@ class OpenTelemetryV2(CustomLogger):
         except Exception as exc:
             verbose_logger.debug("OpenTelemetryV2: metric recording failed: %s", exc)
 
-    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+    async def async_log_failure_event(
+        self,
+        kwargs: Mapping[str, object],
+        response_obj: object,
+        start_time: datetime | float | None,
+        end_time: datetime | float | None,
+    ) -> None:
         if self._emit_mcp_tool_call(kwargs, start_time, end_time):
             return
         if self._emit_mcp_list_tools(kwargs, start_time, end_time):
@@ -964,12 +976,16 @@ def _excluded_db_systems(logger: "OpenTelemetryV2") -> frozenset[str]:
     callback folds into the preset, whose config is env-only.
     """
     otel_settings: Final = (litellm.callback_settings or {}).get("otel")
-    configured: Final = otel_settings.get("excluded_services") if isinstance(otel_settings, dict) else None
+    configured: Final = _excluded_services_setting(otel_settings) if isinstance(otel_settings, dict) else None
     if configured is None:
         return logger.config.excluded_services
     return excluded_db_systems_from(configured)
 
 
+def _excluded_services_setting(otel_settings: Mapping[str, object]) -> object:
+    return otel_settings.get("excluded_services")
+
+
 def _v2_configs(in_memory_loggers: Sequence[object], logger: "OpenTelemetryV2") -> tuple[OpenTelemetryV2Config, ...]:
     """Every v2 logger's config, the published logger's first.