Merge branch 'main' into fix/mantle-geo-inference-profile-cost-lookup

.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/scripts/assert_ci_coverage.py

@@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)")
 DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
 COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
 GLOB_CHARS = frozenset("*?")
+DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"})
 
 # Trees whose jobs are sharded with no catch-all bucket, so every child that holds
 # tests has to be named by some shard or it runs nowhere. A child listed here is
@@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
     return frozenset(ast.literal_eval(values))
 
 
+def _owns(workflow: str, relative_path: str) -> bool:
+    owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow)
+    return owned_root is not None and relative_path.startswith(owned_root)
+
+
+def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]:
+    return frozenset[str]().union(
+        *(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path))
+    )
+
+
 def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
     runner: Final = repo_root / "tests/integration/run.py"
     if not runner.exists():
@@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         )
         != 1
     )
-    gha_tokens: Final = _invoked_test_tokens(
-        scalar
-        for path in (repo_root / ".github/workflows").glob("*.y*ml")
-        for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
+    workflow_tokens: Final = MappingProxyType(
+        {
+            path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name))
+            for path in (repo_root / ".github/workflows").glob("*.y*ml")
+        }
+    )
+    gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values())
+    circle_owned: Final = (
+        frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files
     )
     findings: Final = (
         tuple(
-            Finding(path, "integration contract is also selected by GitHub Actions")
-            for path in paths
+            Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions")
+            for path in sorted(circle_owned)
             if any(_token_covers(token, path) for token in gha_tokens)
         )
         + tuple(
@@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
             for path in sorted(github_files)
             if not (repo_root / path).is_file()
         )
+        + tuple(
+            Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow")
+            for path in sorted(github_files)
+            if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens))
+        )
     )
     browser_commands: Final = tuple(
         scalar.value
@@ -688,13 +710,9 @@ def main() -> int:
     scalars = _all_scalars()
 
     integration_paths, ownership_findings = _integration_ownership()
-    test_findings = (
-        _uncovered_tests(
-            allowlist,
-            _invoked_selections(scalars)
-            + (Selection(included=integration_paths, ignored=frozenset()),),
-        )
-        + ownership_findings
+    test_findings = _uncovered_tests(
+        allowlist,
+        _invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),),
     )
     dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
     stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
@@ -711,13 +729,25 @@ def main() -> int:
             test_findings,
             "Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
         )
+    if ownership_findings:
+        browser_remedy = (
+            " CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may "
+            "select them and the default Playwright config must exclude them."
+            if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings)
+            else ""
+        )
+        _report(
+            "integration contracts with the wrong CI owner",
+            ownership_findings,
+            "Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy,
+        )
     if dockerfile_findings:
         _report(
             "Dockerfiles that no CI job builds",
             dockerfile_findings,
             "Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
         )
-    if stale_findings or test_findings or dockerfile_findings:
+    if stale_findings or test_findings or ownership_findings or dockerfile_findings:
         return 1
 
     _write(

.github/workflows/auto_update_price_and_context_window.yml

@@ -24,12 +24,9 @@ jobs:
       - name: Update JSON Data
         run: |
           uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
-      - name: Regenerate JSON Schema
-        run: |
-          uv run --frozen python ci_cd/generate_model_prices_schema.py
       - name: Create Pull Request
         run: |
-          git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
+          git add model_prices_and_context_window.json
           git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')" 
           gh pr create --title "Update model_prices_and_context_window.json file" \
             --body "Automated update for model_prices_and_context_window.json" \

.github/workflows/sync-together-ai-models.yml

@@ -41,10 +41,6 @@ jobs:
           uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
         env:
           TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
-      - name: Regenerate the JSON schema
-        if: steps.existing.outputs.open_pr == ''
-        run: |
-          uv run --frozen python ci_cd/generate_model_prices_schema.py
       - name: Create a pull request when the registry changed
         if: steps.existing.outputs.open_pr == ''
         run: |
@@ -57,8 +53,7 @@ jobs:
           git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
           git checkout -b "$branch"
           git add model_prices_and_context_window.json \
-            litellm/model_prices_and_context_window_backup.json \
-            model_prices_and_context_window.schema.json
+            litellm/model_prices_and_context_window_backup.json
           git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
           gh auth setup-git
           git push origin "$branch"

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

@@ -16,6 +16,10 @@ on:
       - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
+      - "litellm/chat_completions/dispatch.py"
+      - "litellm/embeddings/dispatch.py"
+      - "litellm/messages/dispatch.py"
+      - "litellm/responses/dispatch.py"
       - "litellm/llms/base_llm/ocr/**"
       - "litellm/llms/custom_httpx/llm_http_handler.py"
       - "Makefile"
@@ -46,6 +50,10 @@ on:
       - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
+      - "litellm/chat_completions/dispatch.py"
+      - "litellm/embeddings/dispatch.py"
+      - "litellm/messages/dispatch.py"
+      - "litellm/responses/dispatch.py"
       - "litellm/llms/base_llm/ocr/**"
       - "litellm/llms/custom_httpx/llm_http_handler.py"
       - "Makefile"

.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
@@ -698,7 +696,6 @@ jobs:
               tests/unit/responses/litellm_completion_transformation
               tests/unit/responses/test_additional_tools.py
               tests/unit/responses/test_custom_tool_call.py
-              tests/unit/responses/test_dispatch.py
               tests/unit/responses/test_metadata_codex_callback.py
               tests/unit/responses/test_no_duplicate_spend_logs.py
               tests/unit/responses/test_null_test_fix.py
@@ -1453,7 +1450,6 @@ jobs:
             tests/unit/caching/test_redis_connection_pool.py \
             tests/unit/caching/test_redis_cluster_cache.py \
             tests/unit/caching/test_evicted_client_closer.py \
-            tests/integration/sdk/test_redis_cluster_iam_auth.py \
             --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
             --cov=./litellm --cov-report=xml:coverage.xml
 

AGENTS.md

@@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
 
 `tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
 
-End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
+End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it
 
 When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
 
@@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
 
 Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
 
-If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
+Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_
 
 Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
 

README.md

@@ -384,6 +384,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
 | [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ |  | ✅ |  |  |  |  |  |
 | [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
+| [Nace (`nace`)](https://console.nace.ai/docs/drex) |  |  |  |  |  |  |  |  |  |  |
 | [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ |  |  |  |  |  |  |
 | [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |

backend/Dockerfile

@@ -4,6 +4,33 @@ ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2
 
 FROM $UV_IMAGE AS uvbin
 
+# ---------- Rust extension ----------
+# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
+FROM $LITELLM_BUILD_IMAGE AS rust-native
+
+WORKDIR /app
+USER root
+
+COPY --from=uvbin /uv /uvx /usr/local/bin/
+
+RUN for i in 1 2 3; do \
+      apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
+      [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
+      sleep 5; \
+    done
+
+ENV UV_PYTHON_DOWNLOADS=0
+
+COPY pyproject.toml README.md rust-toolchain.toml ./
+COPY .cargo ./.cargo
+COPY litellm-rust ./litellm-rust
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    mkdir -p litellm/rust_bridge /wheels && \
+    touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
+    uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
+        python -c "import maturin; maturin.build_wheel('/wheels')"
+
 # ---------- Builder ----------
 FROM $LITELLM_BUILD_IMAGE AS builder
 
@@ -54,7 +81,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 COPY . .
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --frozen --no-default-groups --no-editable \
+    uv sync --frozen --no-install-project --no-default-groups --no-editable \
         --extra proxy \
         --extra proxy-runtime \
         --group admin-mcp \
@@ -63,6 +90,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra saml \
         --python python3.13
 
+RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
+    uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
+
 RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
 
 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

ci_cd/cost_map_guard.py

@@ -1,11 +1,11 @@
 """Guard the cost map on pull requests.
 
-Every pull request whose diff against its merge base touches one of the three cost map files gets the file
-checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the
-map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies
-and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot
-(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three
-files and may only add or update models.
+Every pull request whose diff against its merge base touches one of the two cost map files gets the file
+checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
+skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
+branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
+get the file checks and additionally may only touch those two files or the legacy schema
+path and may only add or update models.
 """
 
 from __future__ import annotations
@@ -19,13 +19,12 @@ from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final
 
-from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
-
 COST_MAP_PATH: Final = "model_prices_and_context_window.json"
 BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
-SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
-GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
+LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
+GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
 BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
+SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
 EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
 
 CostMap = dict[str, object]
@@ -35,7 +34,6 @@ CostMap = dict[str, object]
 class Snapshot:
     cost_map: str
     backup: str
-    schema: str
 
 
 def _parse_object(text: str, path: str) -> CostMap | str:
@@ -46,13 +44,6 @@ def _parse_object(text: str, path: str) -> CostMap | str:
     return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
 
 
-def _rendered_schema(cost_map: CostMap) -> str:
-    try:
-        return render(build_schema(cost_map))
-    except SystemExit as error:
-        return str(error)
-
-
 def _evaluation_output_price_failure(
     key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
 ) -> str | None:
@@ -78,29 +69,13 @@ def _file_failures(
     *,
     evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
 ) -> tuple[str, ...]:
-    schema_text: Final = _rendered_schema(head_map)
-    if not schema_text.startswith("{"):
-        return (schema_text,)
     backup_failure: Final = (
         ()
         if head.backup == head.cost_map
         else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
     )
-    schema_failure: Final = (
-        ()
-        if head.schema == schema_text
-        else (
-            f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
-            "run `python ci_cd/generate_model_prices_schema.py` and commit the result",
-        )
-    )
     return (
         *backup_failure,
-        *schema_failure,
-        *(
-            f"{COST_MAP_PATH} does not validate against its schema: {error}"
-            for error in validation_errors(head_map, json.loads(schema_text))[:20]
-        ),
         *(
             failure
             for key, entry in head_map.items()
@@ -135,7 +110,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
         *(
             f"bot PRs may only change the cost map files, not {path}"
             for path in changed_files
-            if path not in GUARDED_PATHS
+            if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
         ),
         *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
         *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),

ci_cd/generate_model_prices_schema.py

@@ -1,409 +0,0 @@
-from __future__ import annotations
-
-import json
-import sys
-from pathlib import Path
-from typing import Final, Optional
-
-import jsonschema
-
-REPO_ROOT = Path(__file__).parent.parent
-PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
-SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
-
-SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
-
-JsonSchema = dict
-
-NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
-NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
-BOOLEAN: JsonSchema = {"type": "boolean"}
-STRING: JsonSchema = {"type": "string"}
-TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
-WEEKDAY_PATTERN: Final = (
-    r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
-)
-
-EXTRA_BOOLEAN_KEYS = frozenset(
-    {
-        "gemini_native_audio",
-        "gemini_audio_only_live",
-        "uses_embed_content",
-        "use_openai_responses_path",
-        "bedrock_converse_supports_strict_tools",
-        "thinking_always_on",
-    }
-)
-
-HOURS_UTC: Final[JsonSchema] = {
-    "description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
-    "oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
-}
-
-OFF_PEAK_WINDOW: Final[JsonSchema] = {
-    "type": "object",
-    "properties": {
-        "hours_utc": HOURS_UTC,
-        "weekdays": {
-            "type": "array",
-            "description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
-            "items": {
-                "oneOf": [
-                    {"type": "integer", "minimum": 1, "maximum": 7},
-                    {"type": "string", "pattern": WEEKDAY_PATTERN},
-                ]
-            },
-            "minItems": 1,
-        },
-    },
-    "required": ["hours_utc"],
-    "additionalProperties": False,
-}
-
-OBJECT_KEYS: dict[str, JsonSchema] = {
-    "off_peak_pricing": {
-        "type": "object",
-        "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
-        "properties": {
-            "hours_utc": HOURS_UTC,
-            "windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
-            "weekday_timezone": {
-                "type": "string",
-                "description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
-            },
-            "input_cost_per_token": NONNEG_NUMBER,
-            "output_cost_per_token": NONNEG_NUMBER,
-            "output_cost_per_reasoning_token": NONNEG_NUMBER,
-            "cache_read_input_token_cost": NONNEG_NUMBER,
-            "cache_creation_input_token_cost": NONNEG_NUMBER,
-        },
-        "anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
-        "additionalProperties": False,
-    },
-    "search_context_cost_per_query": {
-        "type": "object",
-        "description": "USD cost per web search query, keyed by search context size.",
-        "properties": {
-            "search_context_size_low": NONNEG_NUMBER,
-            "search_context_size_medium": NONNEG_NUMBER,
-            "search_context_size_high": NONNEG_NUMBER,
-        },
-        "additionalProperties": False,
-    },
-    "guardrail_cost_per_unit": {
-        "type": "object",
-        "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
-        "additionalProperties": NONNEG_NUMBER,
-    },
-    "metadata": {
-        "type": "object",
-        "description": "Free-form notes about the entry (e.g. pricing derivation).",
-    },
-    "provider_specific_entry": {
-        "type": "object",
-        "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
-    },
-}
-
-ARRAY_KEYS: dict[str, JsonSchema] = {
-    "supported_audio_formats": {
-        "type": "array",
-        "description": "Audio container formats the model can return.",
-        "items": {"type": "string", "enum": ["mp3", "wav"]},
-    },
-    "supported_endpoints": {
-        "type": "array",
-        "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
-        "items": STRING,
-    },
-    "supported_modalities": {
-        "type": "array",
-        "description": "Input modalities the model accepts.",
-        "items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
-    },
-    "supported_output_modalities": {
-        "type": "array",
-        "description": "Output modalities the model can produce.",
-        "items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
-    },
-    "reasoning_effort_levels": {
-        "type": "array",
-        "description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
-        "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
-    },
-    "supported_regions": {
-        "type": "array",
-        "description": "Cloud regions the model is available in ('global' or region ids).",
-        "items": STRING,
-    },
-    "tiered_pricing": {
-        "type": "array",
-        "description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
-        "items": {
-            "type": "object",
-            "properties": {
-                "range": {
-                    "type": "array",
-                    "description": "[min, max] prompt-token span this tier applies to.",
-                    "items": NONNEG_NUMBER,
-                    "minItems": 2,
-                    "maxItems": 2,
-                },
-                "max_results_range": {
-                    "type": "array",
-                    "description": "[min, max] result-count span this tier applies to (search models).",
-                    "items": NONNEG_NUMBER,
-                    "minItems": 2,
-                    "maxItems": 2,
-                },
-                "input_cost_per_token": NONNEG_NUMBER,
-                "output_cost_per_token": NONNEG_NUMBER,
-                "output_cost_per_reasoning_token": NONNEG_NUMBER,
-                "cache_read_input_token_cost": NONNEG_NUMBER,
-                "cache_creation_input_token_cost": NONNEG_NUMBER,
-                "input_cost_per_query": NONNEG_NUMBER,
-            },
-            "additionalProperties": False,
-        },
-    },
-}
-
-INTEGER_KEYS: dict[str, JsonSchema] = {
-    "max_tokens": {
-        **NONNEG_INTEGER,
-        "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
-    },
-    "max_input_tokens": {
-        **NONNEG_INTEGER,
-        "description": "Maximum prompt/context tokens the model accepts.",
-    },
-    "max_output_tokens": {
-        **NONNEG_INTEGER,
-        "description": "Maximum tokens the model can generate in one response.",
-    },
-    "output_vector_size": {
-        **NONNEG_INTEGER,
-        "description": "Embedding dimension for embedding models.",
-    },
-    "prompt_cache_min_tokens": {
-        **NONNEG_INTEGER,
-        "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
-    },
-    "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
-    "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
-}
-
-NUMBER_KEYS: dict[str, JsonSchema] = {
-    "regional_processing_uplift_multiplier_eu": {
-        "type": "number",
-        "minimum": 1,
-        "description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
-    },
-    "regional_processing_uplift_multiplier_us": {
-        "type": "number",
-        "minimum": 1,
-        "description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
-    },
-    "regional_endpoint_uplift_multiplier": {
-        "type": "number",
-        "minimum": 1,
-        "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
-    },
-}
-
-COST_DESCRIPTIONS: dict[str, str] = {
-    "input_cost_per_token": "USD per prompt token.",
-    "output_cost_per_token": "USD per generated token.",
-    "output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
-    "google_maps_grounding_cost_per_query": (
-        "USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
-    ),
-    "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
-    "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
-    "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
-    "output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
-}
-
-
-def cost_description(key: str) -> Optional[str]:
-    if key in COST_DESCRIPTIONS:
-        return COST_DESCRIPTIONS[key]
-    if key.endswith("_flex"):
-        return "Flex service-tier rate for the same-named base field."
-    if key.endswith("_priority"):
-        return "Priority service-tier rate for the same-named base field."
-    if "_above_" in key:
-        return "Rate applied once the prompt exceeds the token threshold in the field name."
-    return None
-
-
-def cost_schema(key: str) -> JsonSchema:
-    description = cost_description(key)
-    return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
-
-
-def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
-    return {
-        "litellm_provider": {
-            "type": "string",
-            "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
-        },
-        "mode": {
-            "type": "string",
-            "description": "Primary API surface / task type of the model.",
-            "enum": list(modes),
-        },
-        "source": {
-            "type": "string",
-            "description": "URL of the provider pricing/model page this entry was taken from.",
-        },
-        "deprecation_date": {
-            "type": "string",
-            "description": "Date the provider deprecates the model, YYYY-MM-DD.",
-            "format": "date",
-            "pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
-        },
-        "web_search_billing_unit": {
-            "type": "string",
-            "description": "Whether web search is billed per query or per prompt.",
-            "enum": ["per_query", "per_prompt"],
-        },
-        "bedrock_output_config_effort_ceiling": {
-            "type": "string",
-            "description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
-            "enum": ["low", "medium", "high", "max", "xhigh"],
-        },
-        "default_reasoning_effort": {
-            "type": "string",
-            "description": (
-                "Reasoning effort the provider applies when the request omits reasoning_effort. "
-                "Gates whether a non-default temperature or the top_p/logprobs sampling params are "
-                "accepted, which hold only when the effort resolves to 'none'."
-            ),
-            "enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
-        },
-        "comment": STRING,
-        "audio_transcription_config": STRING,
-        "vertex_ai_audio_api": {
-            "type": "string",
-            "enum": ["lyria_predict", "lyria_interactions"],
-        },
-    }
-
-
-def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
-    curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
-    if key in curated:
-        return curated[key]
-    if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
-        return BOOLEAN
-    if "cost" in key:
-        return cost_schema(key)
-    return None
-
-
-def build_schema(prices: dict) -> JsonSchema:
-    entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
-    all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
-    modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
-    unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
-    if unclassified:
-        raise SystemExit(
-            f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
-            f"Add them to the key tables in {Path(__file__).name} and rerun it."
-        )
-    entry_properties = {key: classify(key, modes) for key in all_keys}
-    return {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "title": "LiteLLM model_prices_and_context_window.json",
-        "description": (
-            "Schema for LiteLLM's model price and context window registry "
-            "(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
-            "Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
-            "optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
-            "All costs are USD per unit. New optional fields are added regularly, so consumers should "
-            "ignore unknown fields rather than reject them."
-        ),
-        "type": "object",
-        "properties": {
-            "sample_spec": {
-                "type": "object",
-                "description": (
-                    "Documentation placeholder illustrating the entry shape; not a real model and not "
-                    "schema-conformant (several values are prose)."
-                ),
-            },
-            "fallback_generalizations": {
-                "type": "object",
-                "description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
-                "properties": {
-                    "rules": {
-                        "type": "array",
-                        "items": {
-                            "type": "object",
-                            "properties": {
-                                "name": STRING,
-                                "pattern": STRING,
-                                "description": STRING,
-                            },
-                            "required": ["name", "pattern"],
-                            "additionalProperties": True,
-                        },
-                    }
-                },
-                "additionalProperties": False,
-            },
-        },
-        "additionalProperties": {"$ref": "#/$defs/modelEntry"},
-        "$defs": {
-            "modelEntry": {
-                "type": "object",
-                "description": (
-                    "Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
-                    "are optional; boolean capability flags are simply omitted when unknown or false."
-                ),
-                "required": ["litellm_provider"],
-                "properties": entry_properties,
-                "additionalProperties": True,
-            }
-        },
-    }
-
-
-def render(schema: JsonSchema) -> str:
-    return json.dumps(schema, indent=2) + "\n"
-
-
-def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
-    validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
-    return tuple(
-        f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
-        for error in validator.iter_errors(prices)
-    )
-
-
-def main() -> int:
-    check = "--check" in sys.argv[1:]
-    prices = json.loads(PRICES_PATH.read_text())
-    rendered = render(build_schema(prices))
-    errors = validation_errors(prices, json.loads(rendered))
-    if errors:
-        print(f"{PRICES_PATH.name} does not validate against the generated schema:")
-        print("\n".join(errors[:20]))
-        return 1
-    if not check:
-        SCHEMA_PATH.write_text(rendered)
-        print(f"wrote {SCHEMA_PATH}")
-        return 0
-    if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
-        print(
-            f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
-            f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
-        )
-        return 1
-    print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

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
 

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

gateway/Dockerfile

@@ -20,6 +20,33 @@ RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${
     make -j"$(nproc)" pgbouncer && \
     install -m 0755 pgbouncer /usr/local/bin/pgbouncer
 
+# ---------- Rust extension ----------
+# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
+FROM $LITELLM_BUILD_IMAGE AS rust-native
+
+WORKDIR /app
+USER root
+
+COPY --from=uvbin /uv /uvx /usr/local/bin/
+
+RUN for i in 1 2 3; do \
+      apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
+      [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
+      sleep 5; \
+    done
+
+ENV UV_PYTHON_DOWNLOADS=0
+
+COPY pyproject.toml README.md rust-toolchain.toml ./
+COPY .cargo ./.cargo
+COPY litellm-rust ./litellm-rust
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    mkdir -p litellm/rust_bridge /wheels && \
+    touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
+    uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
+        python -c "import maturin; maturin.build_wheel('/wheels')"
+
 # ---------- Builder ----------
 FROM $LITELLM_BUILD_IMAGE AS builder
 
@@ -66,11 +93,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra bedrock-realtime \
         --python python3.13
 
-# Stage 2 — copy source and install the project + workspace members.
+# Stage 2 — copy source and install the workspace members.
 COPY . .
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --frozen --no-default-groups --no-editable \
+    uv sync --frozen --no-install-project --no-default-groups --no-editable \
         --extra proxy \
         --extra proxy-runtime \
         --group admin-mcp \
@@ -79,6 +106,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra bedrock-realtime \
         --python python3.13
 
+RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
+    uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
+
 # PYTHONPATH=/app makes the source tree shadow the installed package, so the
 # compiled Rust extension must live next to the source or it is never imported.
 RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/

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/migrations/20261015120000_add_environment_variables_revision_trigger/migration.sql

@@ -0,0 +1,23 @@
+SET LOCAL lock_timeout = '10s';
+
+CREATE SEQUENCE IF NOT EXISTS "litellm_environment_variables_revision_seq";
+GRANT USAGE ON SEQUENCE "litellm_environment_variables_revision_seq" TO PUBLIC;
+
+CREATE OR REPLACE FUNCTION litellm_bump_environment_variables_revision() RETURNS TRIGGER AS $$
+DECLARE
+  next_revision BIGINT := nextval(format('%I.litellm_environment_variables_revision_seq', TG_TABLE_SCHEMA)::regclass);
+BEGIN
+  IF TG_OP = 'UPDATE' THEN
+    NEW."reload_revision" := GREATEST(next_revision, OLD."reload_revision" + 1);
+  ELSE
+    NEW."reload_revision" := next_revision;
+  END IF;
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+DROP TRIGGER IF EXISTS "litellm_environment_variables_revision" ON "LiteLLM_Config";
+CREATE TRIGGER "litellm_environment_variables_revision"
+BEFORE INSERT OR UPDATE ON "LiteLLM_Config"
+FOR EACH ROW WHEN (NEW."param_name" = 'environment_variables')
+EXECUTE PROCEDURE litellm_bump_environment_variables_revision();

litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql

@@ -0,0 +1,2 @@
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at"
+ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at");

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?
 
@@ -1980,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
   outcome             String?
 
   @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
+  @@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
 }
 
 model LiteLLM_Lens {

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

@@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [
     "x-openrouter-experimental-metadata",
 ];
 const ANTHROPIC_API_HEADER_PROVIDERS: &str =
-    "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
+    "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek";
 
 pub async fn create(
     State(gateway): State<Arc<Gateway>>,

litellm-rust/crates/gateway-inference/tests/messages.rs

@@ -16,7 +16,7 @@ use wiremock::{
 #[rstest]
 #[case::anthropic("anthropic/test-model", "/v1/messages", true)]
 #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)]
-#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)]
+#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)]
 #[tokio::test]
 async fn messages_reaches_the_provider_and_preserves_json_or_sse(
     #[case] model: &str,

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-messages/tests/messages/request.rs

@@ -343,6 +343,77 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
     assert_eq!(sent, expected);
 }
 
+#[rstest]
+#[case::anthropic("anthropic", &[AnthropicBeta::DangerousToolUse20260903])]
+#[case::deepseek("deepseek", &[])]
+#[case::minimax("minimax", &[])]
+#[tokio::test]
+async fn safeguards_add_the_dangerous_tool_use_beta_only_where_the_beta_map_has_a_column(
+    call: MessagesCall,
+    #[case] provider: &str,
+    #[case] derived: &[AnthropicBeta],
+) {
+    let upstream = upstream([message_response()]).await;
+    let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]);
+
+    run_message(with_fields(
+        MessagesCall {
+            custom_llm_provider: Some(provider.into()),
+            api_key: Some("sk".into()),
+            api_base: Some(upstream.uri()),
+            extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
+            ..call
+        },
+        json!({"safeguards": safeguards}),
+    ))
+    .await;
+
+    let request = only_request(&upstream).await;
+    let expected: BetaSet = derived
+        .iter()
+        .cloned()
+        .chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
+        .collect();
+    assert_eq!(sent_betas(&request), expected);
+    assert_eq!(
+        request.body_json::<Value>().unwrap()["safeguards"],
+        safeguards
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn null_safeguards_are_dropped_instead_of_sent_without_their_beta(
+    call: MessagesCall,
+    #[values("anthropic", "azure_ai", "bedrock_mantle", "deepseek", "minimax")] provider: &str,
+) {
+    let upstream = upstream([message_response()]).await;
+
+    run_message(with_fields(
+        MessagesCall {
+            custom_llm_provider: Some(provider.into()),
+            api_key: Some("sk".into()),
+            api_base: Some(upstream.uri()),
+            extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
+            ..call
+        },
+        json!({"safeguards": null}),
+    ))
+    .await;
+
+    let request = only_request(&upstream).await;
+    assert_eq!(
+        sent_betas(&request),
+        BetaSet::from_iter([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
+    );
+    assert!(
+        !request
+            .body_json::<Map<String, Value>>()
+            .unwrap()
+            .contains_key("safeguards")
+    );
+}
+
 #[rstest]
 #[tokio::test]
 async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) {

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/anthropic/messages/transformation.rs

@@ -109,7 +109,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
     }
 
     fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
-        update_headers_with_anthropic_beta(headers, request)
+        merge_beta_headers(
+            headers,
+            provider_feature_betas(request, BetaProvider::Anthropic),
+        )
     }
 }
 
@@ -136,12 +139,20 @@ pub(crate) fn transform_messages_request(
         messages: strip_encrypted_reasoning_blocks(messages),
         params: MessagesOptionalParams {
             context_management,
+            extra: without_null_safeguards(request.params.extra),
             ..request.params
         },
         ..request
     })
 }
 
+fn without_null_safeguards(extra: Map<String, Value>) -> Map<String, Value> {
+    extra
+        .into_iter()
+        .filter(|(key, value)| key != "safeguards" || !value.is_null())
+        .collect()
+}
+
 pub(crate) fn update_headers_with_anthropic_beta(
     headers: Headers,
     request: &MessagesRequest,
@@ -149,12 +160,34 @@ pub(crate) fn update_headers_with_anthropic_beta(
     merge_beta_headers(headers, feature_betas(request))
 }
 
-/// The betas a request needs on a host other than the first-party API: the Anthropic set
-/// filtered and renamed through the host's `BetaProvider` policy.
+pub(crate) fn update_headers_with_safeguards_beta(
+    headers: Headers,
+    request: &MessagesRequest,
+    provider: BetaProvider,
+) -> Headers {
+    merge_beta_headers(
+        headers,
+        feature_betas(request).union(on_provider(mapped_host_betas(&request.params), provider)),
+    )
+}
+
+/// The betas a request needs on a host with a column in the beta map, filtered and renamed
+/// through that column.
 pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet {
-    feature_betas(request)
-        .iter()
-        .filter_map(|beta| beta.on(provider))
+    on_provider(
+        feature_betas(request).union(mapped_host_betas(&request.params)),
+        provider,
+    )
+}
+
+fn on_provider(betas: BetaSet, provider: BetaProvider) -> BetaSet {
+    betas.iter().filter_map(|beta| beta.on(provider)).collect()
+}
+
+fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet {
+    carries_safeguards(params)
+        .then_some(AnthropicBeta::DangerousToolUse20260903)
+        .into_iter()
         .collect()
 }
 
@@ -206,6 +239,13 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool {
             .is_some_and(|config| config.format.is_some())
 }
 
+fn carries_safeguards(params: &MessagesOptionalParams) -> bool {
+    params
+        .extra
+        .get("safeguards")
+        .is_some_and(|value| !value.is_null())
+}
+
 fn messages_carry_output_config(messages: &[Message]) -> bool {
     messages
         .iter()
@@ -1025,10 +1065,44 @@ mod tests {
         json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}),
         &["per-turn-control-2026-07-01"]
     )]
+    #[case::safeguards(
+        json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}),
+        &[]
+    )]
     fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) {
         assert_eq!(feature_betas(&request(fields)), betas(expected));
     }
 
+    #[rstest]
+    #[case::anthropic(BetaProvider::Anthropic, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])]
+    #[case::vertex_ai(BetaProvider::VertexAi, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])]
+    #[case::empty_list(BetaProvider::Anthropic, json!([]), &["dangerous-tool-use-2026-09-03"])]
+    #[case::not_a_list(BetaProvider::Anthropic, json!(1), &["dangerous-tool-use-2026-09-03"])]
+    #[case::null(BetaProvider::Anthropic, Value::Null, &[])]
+    #[case::bedrock_converse_has_a_null_cell(BetaProvider::BedrockConverse, json!([{"type": "dangerous_tool_use"}]), &[])]
+    fn safeguards_add_the_dangerous_tool_use_beta_where_the_beta_map_has_a_cell(
+        #[case] provider: BetaProvider,
+        #[case] safeguards: Value,
+        #[case] expected: &[&str],
+    ) {
+        assert_eq!(
+            provider_feature_betas(&request(json!({"safeguards": safeguards})), provider),
+            betas(expected)
+        );
+    }
+
+    #[rstest]
+    fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() {
+        let client = headers(&[("anthropic-beta", "claude-code-20250219")]);
+        assert_eq!(
+            update_headers_with_anthropic_beta(
+                client.clone(),
+                &request(json!({"safeguards": [{"type": "dangerous_tool_use"}]}))
+            ),
+            client
+        );
+    }
+
     #[rstest]
     #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])]
     #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])]
@@ -1067,6 +1141,11 @@ mod tests {
         json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
         &[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")],
     )]
+    #[case::safeguards_add_the_dangerous_tool_use_beta_to_the_client_betas(
+        &[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")],
+        json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
+        &[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14")],
+    )]
     fn request_headers_merge_the_feature_betas(
         #[case] input: &[(&str, &str)],
         #[case] fields: Value,

litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs

@@ -2,7 +2,7 @@ use litellm_auth::SecretValue;
 use litellm_http::request::{has_bearer_auth, has_header};
 use litellm_llms_types::{
     formats::messages::MessagesRequest,
-    providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
+    providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
 };
 use litellm_router_types::LitellmParams;
 
@@ -12,7 +12,7 @@ use crate::{
         common_utils::API_KEY_PLACEMENT,
         messages::{
             handler::shape_anthropic_messages_request,
-            transformation::{transform_messages_request, update_headers_with_anthropic_beta},
+            transformation::{transform_messages_request, update_headers_with_safeguards_beta},
         },
     },
     azure_ai::common_utils::{
@@ -104,7 +104,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
     }
 
     fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
-        update_headers_with_anthropic_beta(headers, request)
+        update_headers_with_safeguards_beta(headers, request, BetaProvider::AzureAi)
     }
 }
 
@@ -544,6 +544,12 @@ mod tests {
         &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")]
     )]
     #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])]
+    #[case::safeguards_need_the_dangerous_tool_use_beta(
+        json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
+        &[("anthropic-beta", "claude-code-20250219")],
+        &[("x-api-key", "k"), ("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")]
+    )]
+    #[case::null_safeguards_need_no_beta(json!({"safeguards": null}), &[], &[("x-api-key", "k")])]
     fn request_headers_carry_the_anthropic_feature_betas(
         #[case] fields: serde_json::Value,
         #[case] forwarded: &[(&str, &str)],

litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs

@@ -6,7 +6,7 @@ use litellm_core_utils::settings::resolve_non_empty;
 use litellm_http::request::{header_value, with_header, without_headers};
 use litellm_llms_types::{
     formats::messages::MessagesRequest,
-    providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
+    providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
 };
 use litellm_router_types::LitellmParams;
 use serde_json::Value;
@@ -15,7 +15,7 @@ use crate::{
     Error, ErrorDetail,
     anthropic::messages::{
         handler::shape_anthropic_messages_request,
-        transformation::{transform_messages_request, update_headers_with_anthropic_beta},
+        transformation::{transform_messages_request, update_headers_with_safeguards_beta},
     },
     base_llm::{
         auth::{AuthScheme, Headers, ValidatedEnvironment},
@@ -174,7 +174,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
             Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()),
             None => headers,
         };
-        update_headers_with_anthropic_beta(headers, request)
+        update_headers_with_safeguards_beta(headers, request, BetaProvider::Anthropic)
     }
 
     fn wire_body(&self, body: Value) -> Value {
@@ -356,6 +356,21 @@ mod tests {
         );
     }
 
+    #[rstest]
+    #[case::safeguards(json!([{"type": "dangerous_tool_use"}]), Some("dangerous-tool-use-2026-09-03"))]
+    #[case::null_safeguards(Value::Null, None)]
+    fn safeguards_need_the_dangerous_tool_use_beta(
+        #[case] safeguards: Value,
+        #[case] expected: Option<&str>,
+    ) {
+        let request = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": safeguards})).unwrap();
+        let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &request);
+        assert_eq!(
+            header_value(&headers, "anthropic-beta").map(|value| value.to_string()),
+            expected.map(str::to_string)
+        );
+    }
+
     #[rstest]
     #[case::default(json!({}), false)]
     #[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)]

litellm-rust/crates/llms/src/deepseek/messages/transformation.rs

@@ -451,4 +451,27 @@ mod tests {
             json!({"type": "enabled", "budget_tokens": 1024})
         );
     }
+
+    #[rstest]
+    #[case::without_client_betas(&[])]
+    #[case::with_client_betas(&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")])]
+    fn safeguards_and_client_betas_reach_deepseek_as_sent(#[case] client: &[(&str, &str)]) {
+        let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]);
+        let body = json!({
+            "model": "deepseek-v4-pro",
+            "max_tokens": 100,
+            "messages": [{"role": "user", "content": "hi"}],
+            "safeguards": safeguards
+        });
+        let headers: Headers = client
+            .iter()
+            .map(|(name, value)| (name.to_string(), value.to_string()))
+            .collect();
+        assert_eq!(
+            DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG
+                .request_headers(headers.clone(), &request_from(body.clone())),
+            headers
+        );
+        assert_eq!(transformed(body)["safeguards"], safeguards);
+    }
 }

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/llms/src/vertex_ai/messages/transformation.rs

@@ -213,23 +213,14 @@ fn has_web_search_tool(tools: Option<&[Recognized<MessagesTool>]>) -> bool {
 }
 
 /// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature
-/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled
-/// out.
+/// betas under Vertex's policy, plus the web search beta Vertex needs spelled out.
 fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet {
     let params = &request.params;
     provider_feature_betas(request, BetaProvider::VertexAi).union(
-        [
-            has_web_search_tool(params.tools.as_deref())
-                .then_some(AnthropicBeta::WebSearch20250305),
-            params
-                .extra
-                .get("safeguards")
-                .is_some_and(|value| !value.is_null())
-                .then_some(AnthropicBeta::DangerousToolUse20260903),
-        ]
-        .into_iter()
-        .flatten()
-        .collect(),
+        has_web_search_tool(params.tools.as_deref())
+            .then_some(AnthropicBeta::WebSearch20250305)
+            .into_iter()
+            .collect(),
     )
 }
 

litellm-rust/crates/model-catalog/AGENTS.md

@@ -2,5 +2,4 @@
 
 For `model_prices_and_context_window.json` validation, we should eventually:
 
-- Remove any schema file like `model_prices_and_context_window.schema.json`
 - Stop skipping this crate's tests

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

@@ -1,34 +1,13 @@
-### Hide pydantic namespace conflict warnings globally ###
 from __future__ import annotations
 
-import warnings
+from litellm._process_defaults import apply_process_defaults as _apply_process_defaults
+
+_apply_process_defaults()
 
-warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*")
-# Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances
-# This warning can accumulate during streaming and cause memory leaks
-warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*")
-# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it
-# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked
-warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*")
 ### INIT VARIABLES #########################
 import threading
 import os
 
-# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available
-import dotenv as _dotenv
-
-
-def _dev_env_hot_reload_enabled() -> bool:
-    """The proxy exports this flag when started with ``--reload``. A reloaded
-    worker is a fresh process that inherits the reloader's environment, so an
-    edited ``.env`` value stays masked by the stale inherited one unless we
-    let the file win; overriding makes the edit take effect on reload."""
-    return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
-
-
-if os.getenv("LITELLM_MODE", "DEV") == "DEV":
-    _dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
-
 from collections.abc import Mapping, Sequence
 from typing import (
     Any,
@@ -118,6 +97,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",
@@ -206,7 +186,7 @@ _async_failure_callback: List[Union[str, Callable[..., object], "CustomLogger"]]
     []
 )  # internal variable - async custom callbacks are routed here.
 pre_call_rules: List[Callable] = []
-post_call_rules: List[Callable] = []
+post_call_rules: List[Callable[..., object]] = []
 turn_off_message_logging: Optional[bool] = False
 standard_logging_payload_excluded_fields: Optional[List[str]] = (
     None  # Fields to exclude from StandardLoggingPayload before callbacks receive it
@@ -392,6 +372,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
@@ -1175,7 +1156,9 @@ def _build_models_by_provider() -> dict:
         "edenai": edenai_models,
         "datarobot": datarobot_models,
         "vertex_ai": vertex_chat_models
+        | vertex_code_chat_models
         | vertex_text_models
+        | vertex_code_text_models
         | vertex_anthropic_models
         | vertex_vision_models
         | vertex_language_models
@@ -1183,7 +1166,7 @@ def _build_models_by_provider() -> dict:
         | vertex_minimax_models
         | vertex_moonshot_models
         | vertex_zai_models,
-        "ai21": ai21_models,
+        "ai21": ai21_models | ai21_chat_models,
         "bedrock": bedrock_models | bedrock_converse_models,
         "petals": petals_models,
         "ollama": ollama_models,
@@ -1414,6 +1397,7 @@ from .exceptions import (
     InvalidRequestError,
     BadRequestError,
     ImageFetchError,
+    MalformedRequestFieldError as MalformedRequestFieldError,
     VectorStoreSearchError,
     NotFoundError,
     PermissionDeniedError,
@@ -1689,6 +1673,9 @@ if TYPE_CHECKING:
     from .llms.perplexity.decisions.transformation import (
         PerplexityDecisionsConfig as PerplexityDecisionsConfig,
     )
+    from .llms.nace.decisions.transformation import (
+        NaceDecisionsConfig as NaceDecisionsConfig,
+    )
     from .llms.typesafe.decisions.transformation import (
         TypeSafeDecisionsConfig as TypeSafeDecisionsConfig,
     )

litellm/_internal_context.py

@@ -37,6 +37,16 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families"
 _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
 
 
+@contextmanager
+def internal_sub_call() -> Generator[None]:
+    """A nested call the parent call logs and bills, so the parent must fold this call's cost into its own."""
+    token: Final = is_internal_call.set(True)
+    try:
+        yield
+    finally:
+        is_internal_call.reset(token)
+
+
 @contextmanager
 def post_response_phase() -> Generator[None]:
     """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns."""

litellm/_lazy_imports_registry.py

@@ -162,6 +162,7 @@ LLM_CONFIG_NAMES: Final = (
     "HostedVLLMRerankConfig",
     "PerplexityDecisionsConfig",
     "TypeSafeDecisionsConfig",
+    "NaceDecisionsConfig",
     "OpenRouterDecisionsConfig",
     "FireworksAIDecisionsConfig",
     "CloudflareDecisionsConfig",
@@ -729,6 +730,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
         "HostedVLLMRerankConfig",
     ),
     "PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"),
+    "NaceDecisionsConfig": (".llms.nace.decisions.transformation", "NaceDecisionsConfig"),
     "TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"),
     "OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
     "FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"),

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/_process_defaults.py

@@ -0,0 +1,24 @@
+import os
+import warnings
+
+import dotenv
+
+
+def _dev_env_hot_reload_enabled() -> bool:
+    """The proxy exports this flag when started with ``--reload``. A reloaded
+    worker is a fresh process that inherits the reloader's environment, so an
+    edited ``.env`` value stays masked by the stale inherited one unless we
+    let the file win; overriding makes the edit take effect on reload."""
+    return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
+
+
+def apply_process_defaults() -> None:
+    warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*")
+    # Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances
+    # This warning can accumulate during streaming and cause memory leaks
+    warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*")
+    # ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it
+    # cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked
+    warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*")
+    if os.getenv("LITELLM_MODE", "DEV") == "DEV":
+        dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())

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_beta_headers_manager.py

@@ -22,17 +22,28 @@ Environment Variables:
 - LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional)
 """
 
+import asyncio
 import json
 import os
+import threading
+from collections.abc import Mapping
+from concurrent.futures import ThreadPoolExecutor
+from functools import cache
 from importlib.resources import files
-from typing import Final
+from typing import Final, TypeAlias
 
 import httpx
+from pydantic import TypeAdapter
 
 from litellm._logging import verbose_logger
 
 # Cache for the loaded configuration
 _BETA_HEADERS_CONFIG: dict | None = None
+_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock()
+_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config")
+
+BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]]
+_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig)
 
 
 class GetAnthropicBetaHeadersConfig:
@@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict:
     if _BETA_HEADERS_CONFIG is not None:
         return _BETA_HEADERS_CONFIG
 
-    # Get the URL from environment or use default
-    from litellm import anthropic_beta_headers_url
+    with _BETA_HEADERS_CONFIG_LOCK:
+        if _BETA_HEADERS_CONFIG is not None:
+            return _BETA_HEADERS_CONFIG
 
-    _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
-    verbose_logger.debug("Loaded and cached beta headers config")
+        # Get the URL from environment or use default
+        from litellm import anthropic_beta_headers_url
 
-    return _BETA_HEADERS_CONFIG
+        _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
+        verbose_logger.debug("Loaded and cached beta headers config")
+
+        return _BETA_HEADERS_CONFIG
+
+
+async def aload_beta_headers_config() -> None:
+    if _BETA_HEADERS_CONFIG is None:
+        await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config)
+
+
+@cache
+def _bundled_beta_headers_config() -> BetaHeadersConfig:
+    return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config())
 
 
 def reload_beta_headers_config() -> dict:
@@ -288,6 +313,14 @@ def is_beta_header_supported(
     return beta_header in provider_mapping and provider_mapping[beta_header] is not None
 
 
+def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool:
+    config: Final = _bundled_beta_headers_config()
+    aliases: Final = config.get("provider_aliases", {})
+    canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider
+    provider_mapping: Final = config.get(canonical_provider, {})
+    return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None
+
+
 def get_provider_beta_header(
     anthropic_beta_header: str,
     provider: str,

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/assistants/main.py

@@ -129,6 +129,7 @@ def get_assistants(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_assistants=aget_assistants,
         )
@@ -312,6 +313,7 @@ def create_assistants(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             create_assistant_data=create_assistant_data,
             client=client,
             async_create_assistants=async_create_assistants,
@@ -461,6 +463,7 @@ def delete_assistant(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             assistant_id=assistant_id,
             client=client,
             async_delete_assistants=async_delete_assistants,
@@ -637,6 +640,7 @@ def create_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             acreate_thread=acreate_thread,
         )
@@ -783,6 +787,7 @@ def get_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_thread=aget_thread,
         )
@@ -960,6 +965,7 @@ def add_message(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             a_add_message=a_add_message,
         )
@@ -1111,6 +1117,7 @@ def get_messages(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_messages=aget_messages,
         )
@@ -1308,6 +1315,7 @@ def run_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             arun_thread=arun_thread,
             event_handler=event_handler,

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
@@ -17,7 +19,7 @@ from litellm.llms.vertex_ai.batches.transformation import (
 )
 from litellm.types.llms.openai import Batch
 from litellm.types.utils import ModelInfo, Usage
-from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS
+from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS
 from litellm.utils import token_counter
 
 
@@ -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",
@@ -553,6 +569,7 @@ def extract_file_access_credentials(
             # A federated deployment holds no api_key, so without these the fetch that reads a
             # finished batch's output has nothing to authenticate with and its cost is never billed.
             *sorted(ANTHROPIC_WIF_KWARGS_KEYS),
+            *sorted(OPENAI_WIF_KWARGS_KEYS),
         )
         for key in credential_keys:
             if key in litellm_params:
@@ -604,14 +621,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,
@@ -309,6 +309,7 @@ def create_batch(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -471,6 +472,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
             organization=organization,
             timeout=timeout,
             max_retries=optional_params.max_retries,
+            litellm_params=litellm_params,
         )
     elif custom_llm_provider == "azure":
         api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -568,7 +570,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.
 
@@ -819,6 +821,7 @@ def list_batches(
                 organization=organization,
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -941,7 +944,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.
 
@@ -1016,6 +1019,7 @@ def cancel_batch(
                 organization=organization,
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")

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/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/caching_handler.py

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
 from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
 
 import litellm
-from litellm._internal_context import post_response_phase
+from litellm._internal_context import is_internal_call, post_response_phase
 from litellm._logging import print_verbose, verbose_logger
 from litellm.caching import InMemoryCache
 from litellm.caching.caching import S3Cache, response_cache_phase
@@ -511,8 +511,10 @@ class LLMCachingHandler:
                         custom_llm_provider=custom_llm_provider,
                     )
 
-                    if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
-                        logging_obj.handle_sync_success_callbacks_for_async_calls(
+                    if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks(
+                        cached_result=cached_result
+                    ):
+                        logging_obj.handle_sync_cache_hit_success_callbacks(
                             result=cached_result,
                             start_time=start_time,
                             end_time=end_time,
@@ -820,7 +822,12 @@ class LLMCachingHandler:
             start_time (datetime): The start time of the operation.
             end_time (datetime): The end time of the operation.
             cache_hit (bool): Whether it was a cache hit.
+
+        An internal sub-call's hit fires nothing, since the parent call that folds it in logs it.
         """
+        if is_internal_call.get():
+            return
+
         from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 
         GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(

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/completion_extras/litellm_responses_transformation/transformation.py

@@ -197,7 +197,7 @@ def _drops_audio_input(model: str, litellm_params: Mapping[str, object]) -> bool
     return not _supports_audio_input(model, litellm_params)
 
 
-def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
+def responses_provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
     return MappingProxyType(
         {
             key: value
@@ -1455,7 +1455,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
 
         model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
-        for key, value in _provider_metadata(raw_response.model_extra).items():
+        for key, value in responses_provider_metadata(raw_response.model_extra).items():
             setattr(model_response, key, value)
 
         # Preserve hidden params from the ResponsesAPIResponse, especially the headers
@@ -2139,7 +2139,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 from litellm.responses.utils import ResponseAPILoggingUtils
 
                 usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_data.get("usage"))
-            provider_metadata: Final = _provider_metadata(response_data)
+            provider_metadata: Final = responses_provider_metadata(response_data)
             served_service_tier: Final = response_data.get("service_tier")
             return ModelResponseStream(
                 choices=[

litellm/constants.py

@@ -1,4 +1,5 @@
 import os
+import re
 import sys
 from enum import Enum
 from types import MappingProxyType
@@ -17,7 +18,10 @@ 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"})
+CLAUDE_CODE_PICKER_PATTERN: Final = re.compile(r"claude|anthropic", re.IGNORECASE)
+GATEWAY_CLIENT_HEADER: Final = "x-gateway-client"
+CLAUDE_CODE_CLIENT: Final = "claude-code"
 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"
@@ -57,6 +61,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
         "model_group_retry_policy",
         "model_group_alias",
         "enable_weighted_failover",
+        "enable_mid_stream_fallback_continuation",
         "enable_tag_filtering",
         "tag_routing_prefix",
         "optional_pre_call_checks",
@@ -106,6 +111,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
@@ -145,6 +151,7 @@ HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
 # Streams exceeding this duration are terminated with a Timeout error.
 # None (default) = no limit.  Set env var to a number of seconds to enable globally.
 LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS")
+MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0
 
 # Maximum number of base64 characters to keep in logging payloads.
 # Data URIs exceeding this are replaced with a size placeholder.
@@ -379,6 +386,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
@@ -387,7 +395,7 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
 NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1"
 DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
 
-BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
+BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message"
 BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
 BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
 BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime"
@@ -845,6 +853,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 +975,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",
@@ -1609,6 +1619,7 @@ CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = get_env_int("CLOUDZERO_EXPORT_INTERVA
 MCP_TOOL_NAME_PREFIX: Final = "mcp_tool"
 MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = get_env_int("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)
 PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096
+INTERNAL_CALLBACK_NAMES: Final = frozenset({"cache", "vector_store_pre_call_hook"})
 
 # Headers to control callbacks
 X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
@@ -1622,8 +1633,10 @@ 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"
+CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id"
 LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
 LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
     "Truncation is a DB storage safeguard. "
@@ -1700,6 +1713,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"
@@ -1850,6 +1864,7 @@ SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_L
 PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25))
 PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0))
 PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0))
+TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS: Final = 30
 DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = get_env_int("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)  # 1 minute
 PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
 RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, get_env_int("RESET_BUDGET_JOB_BATCH_SIZE", 500))
@@ -1861,7 +1876,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))
@@ -1883,6 +1904,7 @@ _background_interaction_cost_polling_env: Final = os.getenv(
     "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true"
 ).lower()
 BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
+BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0
 PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
 PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10)  # in seconds, increased from 10
 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)

litellm/containers/README.md

@@ -120,6 +120,7 @@ class AzureContainerConfig(BaseContainerConfig):
         self,
         headers: dict,
         api_key: Optional[str] = None,
+        litellm_params: Optional[GenericLiteLLMParams] = None,
     ) -> dict:
         """Azure uses api-key header instead of Bearer token."""
         import litellm

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

@@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
 
 from httpx import Response
 from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
@@ -86,6 +86,9 @@ from litellm.llms.openai.cost_calculation import (
 from litellm.llms.openai.cost_calculation import (
     video_output_cost_per_second,
 )
+from litellm.llms.openai.transcriptions.gpt_transformation import (
+    OpenAIGPTAudioTranscriptionConfig,
+)
 from litellm.llms.perplexity.cost_calculator import (
     cost_per_token as perplexity_cost_per_token,
 )
@@ -147,6 +150,7 @@ from litellm.utils import (
     TranscriptionResponse,
     cached_get_model_info_helper,
     token_counter,
+    unmapped_model_generalization,
 )
 
 if TYPE_CHECKING:
@@ -325,6 +329,15 @@ def _transcription_usage_has_token_details(
     return (prompt_tokens_val > 0) or (completion_tokens_val > 0)
 
 
+def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool:
+    if model is None or custom_llm_provider != LlmProviders.OPENAI.value:
+        return False
+    transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config(
+        model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI
+    )
+    return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig)
+
+
 OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"]
 
 
@@ -863,6 +876,16 @@ def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | Non
     return own_model_info["litellm_provider"]
 
 
+def _pricing_provider_for_cost_candidate(
+    candidate_provider: str | None,
+    request_provider: str | None,
+    current_provider: object,
+) -> str | None:
+    if candidate_provider != request_provider:
+        return candidate_provider
+    return current_provider if isinstance(current_provider, str) else None
+
+
 def select_model_name_for_cost_calc(
     model: str | None,
     completion_response: object | None,
@@ -1426,32 +1449,15 @@ def completion_cost(
     vertex_location: str | None = None,  # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global")
 ) -> float:
     """
-    Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm.
+    Return the USD cost of a LiteLLM call.
 
-    Parameters:
-        completion_response (litellm.ModelResponses): [Required] The response received from a LiteLLM completion request.
+    Pass `completion_response` to price a response LiteLLM returned. Without one, pass `model` plus
+    `messages` / `prompt` and `completion`, and token counts are computed from the text. Prices come from
+    `litellm.model_cost` unless `custom_cost_per_token` or `custom_cost_per_second` is given. `total_time`
+    is in milliseconds and only matters for per-second pricing such as Replicate and SageMaker.
 
-        [OPTIONAL PARAMS]
-        model (str): Optional. The name of the language model used in the completion calls
-        prompt (str): Optional. The input prompt passed to the llm
-        completion (str): Optional. The output completion text from the llm
-        total_time (float, int): Optional. (Only used for Replicate LLMs) The total time used for the request in seconds
-        custom_cost_per_token: Optional[CostPerToken]: the cost per input + output token for the llm api call.
-        custom_cost_per_second: Optional[float]: the cost per second for the llm api call.
-
-    Returns:
-        float: The cost in USD dollars for the completion based on the provided parameters.
-
-    Exceptions:
-        Raises exception if model not in the litellm model cost map. Register model, via custom pricing or PR - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json
-
-
-    Note:
-        - If completion_response is provided, the function extracts token information and the model name from it.
-        - If completion_response is not provided, the function calculates token counts based on the model and input text.
-        - The cost is calculated based on the model, prompt tokens, and completion tokens.
-        - For certain models containing "togethercomputer" in the name, prices are based on the model size.
-        - For un-mapped Replicate models, the cost is calculated based on the total time used for the request.
+    Raises if no candidate model name resolves to pricing. Register the model via `litellm.register_model`
+    or add it to model_prices_and_context_window.json.
     """
     try:
         call_type = infer_call_type(call_type, completion_response) or "completion"
@@ -1549,14 +1555,20 @@ def completion_cost(
         )
 
         requested_model: Final = model
-        potential_model_names: Final = [
-            selected_model,
-            _get_response_model(completion_response),
-        ]
-        if model is not None:
-            potential_model_names.append(model)
+        request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None
+        raw_candidates: Final = (
+            (
+                selected_model,
+                base_model_provider if base_model_provider is not None else request_pricing_provider,
+            ),
+            (_get_response_model(completion_response), request_pricing_provider),
+            *(((model, request_pricing_provider),) if model is not None else ()),
+        )
+        potential_model_names: Final = (
+            raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates)
+        )
 
-        for idx, model in enumerate(potential_model_names):
+        for idx, (model, pricing_provider) in enumerate(potential_model_names):
             try:
                 if verbose_logger.isEnabledFor(logging.DEBUG):
                     verbose_logger.debug("selected model name for cost calculation: %s", model)
@@ -1739,6 +1751,12 @@ def completion_cost(
                     audio_transcription_file_duration = _hidden.get(
                         "audio_transcription_duration",
                         getattr(completion_response, "duration", 0.0),
+                    ) or (
+                        completion_response.usage.seconds
+                        if isinstance(completion_response, TranscriptionResponse)
+                        and isinstance(completion_response.usage, TranscriptionUsageDurationObject)
+                        and _bills_provider_transcription_seconds(requested_model, custom_llm_provider)
+                        else 0.0
                     )
                 elif call_type in _RERANK_CALL_TYPES:
                     if completion_response is not None and isinstance(completion_response, RerankResponse):
@@ -1895,13 +1913,6 @@ def completion_cost(
                 if litellm_logging_obj is not None:
                     request_model_for_cost = litellm_logging_obj.model
 
-                pricing_provider: str | None = (
-                    base_model_provider
-                    if base_model_provider is not None and idx == 0
-                    else custom_llm_provider
-                    if isinstance(custom_llm_provider, str)
-                    else None
-                )
                 (
                     prompt_tokens_cost_usd_dollar,
                     completion_tokens_cost_usd_dollar,
@@ -1909,7 +1920,11 @@ def completion_cost(
                     model=model,
                     prompt_tokens=prompt_tokens or 0,
                     completion_tokens=completion_tokens or 0,
-                    custom_llm_provider=pricing_provider,
+                    custom_llm_provider=_pricing_provider_for_cost_candidate(
+                        candidate_provider=pricing_provider,
+                        request_provider=request_pricing_provider,
+                        current_provider=custom_llm_provider,
+                    ),
                     response_time_ms=response_time_ms,
                     region_name=None if explicit_pricing else region_name,
                     custom_cost_per_second=custom_cost_per_second,
@@ -2008,7 +2023,11 @@ def completion_cost(
                     if cost_per_token_usage_object is not None and model:
                         _token_type_breakdown = get_token_type_cost_breakdown(
                             model=model,
-                            custom_llm_provider=pricing_provider,
+                            custom_llm_provider=_pricing_provider_for_cost_candidate(
+                                candidate_provider=pricing_provider,
+                                request_provider=request_pricing_provider,
+                                current_provider=custom_llm_provider,
+                            ),
                             usage=cost_per_token_usage_object,
                             service_tier=service_tier,
                             data_residency=data_residency,
@@ -2048,7 +2067,7 @@ def completion_cost(
                 )
                 if idx == len(potential_model_names) - 1:
                     raise e
-        raise Exception(f"Unable to calculat cost for received potential model names - {potential_model_names}")
+        raise Exception(f"Unable to calculate cost for received potential model names - {potential_model_names}")
     except Exception as e:
         raise e
 
@@ -2063,13 +2082,20 @@ def get_response_cost_from_hidden_params(
 
     additional_headers: Final = _hidden_params_dict.get("additional_headers", {})
     if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers:
-        response_cost: Final = additional_headers["llm_provider-x-litellm-response-cost"]
-        if response_cost is None:
-            return None
-        return float(additional_headers["llm_provider-x-litellm-response-cost"])
+        return _billable_header_cost(additional_headers["llm_provider-x-litellm-response-cost"])
     return None
 
 
+def _billable_header_cost(value: object) -> float | None:
+    if not isinstance(value, (str, int, float)):
+        return None
+    try:
+        cost: Final = float(value)
+    except ValueError:
+        return None
+    return cost if math.isfinite(cost) and cost >= 0 else None
+
+
 def response_cost_calculator(
     response_object: ModelResponse
     | EmbeddingResponse
@@ -2263,6 +2289,25 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn
         return None
 
 
+_CostCandidate: TypeAlias = tuple[str | None, str | None]
+
+
+def _is_unpriced_generalized_candidate(candidate: _CostCandidate) -> bool:
+    model, pricing_provider = candidate
+    if model is None:
+        return False
+    rule_info: Final = unmapped_model_generalization(model, pricing_provider)
+    return rule_info is not None and not _cost_map_entry_prices_anything(rule_info)
+
+
+def _priced_candidates_before_unpriced_rules(candidates: tuple[_CostCandidate, ...]) -> tuple[_CostCandidate, ...]:
+    unpriced: Final = tuple(_is_unpriced_generalized_candidate(candidate) for candidate in candidates)
+    return (
+        *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if not is_unpriced),
+        *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if is_unpriced),
+    )
+
+
 def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None:
     raw_entry: Final = litellm.model_cost.get(key)
     return raw_entry if isinstance(raw_entry, Mapping) else None
@@ -3008,6 +3053,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
 
@@ -3059,6 +3105,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
@@ -57,7 +57,7 @@ from litellm.types.decisions import (
     UnsupportedDecisionsRequest,
 )
 from litellm.types.utils import LlmProviders, ProviderSpecificHeader
-from litellm.utils import ProviderConfigManager
+from litellm.utils import ProviderConfigManager, get_model_cost_key
 
 SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]]
 OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]]
@@ -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:
@@ -442,7 +458,7 @@ def prepare_call(
         ) from error
 
     try:
-        resolved_api_base, resolved_api_key = provider_config.resolve_connection(
+        resolved_api_base, static_api_key = provider_config.resolve_connection(
             dynamic_api_base or api_base, dynamic_api_key or api_key
         )
     except ValueError as error:
@@ -453,6 +469,7 @@ def prepare_call(
             model=model,
             llm_provider=provider,
         )
+    resolved_api_key: Final = provider_config.resolve_credential(static_api_key, resolved_api_base, kwargs)
     if resolved_api_key is None and provider_config.api_key_required:
         raise litellm.AuthenticationError(
             message=f"Missing API key for Decisions provider '{provider}'",
@@ -650,7 +667,10 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci
         {
             "model": f"{call.custom_llm_provider}/{call.model}",
             "custom_llm_provider": call.custom_llm_provider,
-            "provider_response_model": f"{call.custom_llm_provider}/{call.model}",
+            "provider_response_model": (
+                get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
+                or f"{call.custom_llm_provider}/{call.model}"
+            ),
             **({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}),
         }
     )

litellm/decisions/main.py

@@ -1,5 +1,7 @@
+import asyncio
 import warnings
 from collections.abc import Mapping
+from functools import partial
 from typing import Final
 
 import httpx
@@ -104,19 +106,22 @@ async def adecisions(
     videos: DecisionsVideoClips | None = None,
     **kwargs: object,
 ) -> OpenAIDecisionResponse | DecisionsResponse:
-    call: Final = _prepare_decisions_call(
-        model=model,
-        decision_input=input,
-        questions=questions,
-        safety_identifier=safety_identifier,
-        audio=audio,
-        videos=videos,
-        api_key=api_key,
-        api_base=api_base,
-        timeout=timeout,
-        custom_llm_provider=custom_llm_provider,
-        extra_headers=extra_headers,
-        kwargs=kwargs,
+    call: Final = await asyncio.to_thread(
+        partial(
+            _prepare_decisions_call,
+            model=model,
+            decision_input=input,
+            questions=questions,
+            safety_identifier=safety_identifier,
+            audio=audio,
+            videos=videos,
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            custom_llm_provider=custom_llm_provider,
+            extra_headers=extra_headers,
+            kwargs=kwargs,
+        )
     )
     return await asend(call)
 

litellm/exceptions.py

@@ -313,6 +313,10 @@ class VectorStoreSearchError(BadRequestError):
         )
 
 
+class MalformedRequestFieldError(BadRequestError):
+    pass
+
+
 class UnprocessableEntityError(openai.UnprocessableEntityError):
     def __init__(
         self,
@@ -1217,6 +1221,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
         num_retries: int | None = None,
         generated_content: str = "",
         is_pre_first_chunk: bool = False,
+        emitted_disqualifying_content: bool = False,
     ):
         original_status: Final = getattr(original_exception, "status_code", None)
         self.status_code = int(original_status) if original_status is not None else 503
@@ -1229,6 +1234,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
         self.num_retries = num_retries
         self.generated_content = generated_content
         self.is_pre_first_chunk = is_pre_first_chunk
+        self.emitted_disqualifying_content = emitted_disqualifying_content
 
         # Create a response if one wasn't provided
         if response is None:

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:
@@ -477,6 +498,7 @@ class MCPClient:
         elicitation_callback: "ElicitationCallback | None" = None,
         logging_callback: Callable | None = None,
         protocol_version: MCPUpstreamProtocol = "auto",
+        upstream_transport: httpx2.AsyncBaseTransport | None = None,
     ):
         self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol](
             MCPUpstreamProtocol
@@ -503,6 +525,7 @@ class MCPClient:
         self._sampling_callback: Callable | None = sampling_callback
         self._elicitation_callback: ElicitationCallback | None = elicitation_callback
         self._logging_callback: Callable | None = logging_callback
+        self._upstream_transport: httpx2.AsyncBaseTransport | None = upstream_transport
         self._session_callbacks_override: _SessionCallbackMap | None = None
         # handle the basic auth value if provided
         if auth_value:
@@ -533,6 +556,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
 
@@ -578,7 +612,11 @@ class MCPClient:
             return stdio_client(server_params), None
         if self.transport_type == MCPTransport.sse:
             headers = self._get_auth_headers()
-            httpx_client_factory = self._create_httpx_client_factory()
+            httpx_client_factory = (
+                self._create_httpx_client_factory()
+                if self._upstream_transport is None
+                else self._create_httpx_client_factory(transport=self._upstream_transport)
+            )
             return (
                 sse_client(
                     url=self.server_url,
@@ -590,7 +628,11 @@ class MCPClient:
             )
         # HTTP transport (default)
         headers = self._get_auth_headers()
-        httpx_client_factory = self._create_httpx_client_factory()
+        httpx_client_factory = (
+            self._create_httpx_client_factory()
+            if self._upstream_transport is None
+            else self._create_httpx_client_factory(transport=self._upstream_transport)
+        )
         verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
         http_client = httpx_client_factory(
             headers=headers,
@@ -696,6 +738,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 +759,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 +770,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 +828,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 +839,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 +858,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 +1131,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 +1163,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 +1621,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 +1763,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

@@ -259,6 +259,7 @@ def create_file(
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
                 create_file_data=_create_file_request,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -388,6 +389,7 @@ def file_retrieve(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -569,6 +571,7 @@ def file_delete(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -775,6 +778,7 @@ def file_list(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -972,6 +976,7 @@ def file_content(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -1093,6 +1098,7 @@ def file_content_streaming(
             organization=openai_creds.organization,
             chunk_size=chunk_size,
             client=client if isinstance(client, (OpenAI, AsyncOpenAI)) else None,
+            litellm_params=optional_params.model_dump(exclude_none=True),
         )
     elif custom_llm_provider == LlmProviders.VERTEX_AI.value:
         if not _is_async:

litellm/fine_tuning/main.py

@@ -239,6 +239,7 @@ def create_fine_tuning_job(
                 client=kwargs.get(
                     "client", None
                 ),  # note, when we add this to `GenericLiteLLMParams` it impacts a lot of other tests + linting
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -436,6 +437,7 @@ def cancel_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -591,6 +593,7 @@ def list_fine_tuning_jobs(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -730,6 +733,7 @@ def retrieve_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":