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

# Conflicts:
#	tests/test_litellm/proxy/auth/test_auth_checks.py

.circleci/config.yml

@@ -257,7 +257,7 @@ commands:
       - install_rust
       - restore_cache:
           keys:
-            - v1-uv-cache-{{ checksum "uv.lock" }}
+            - v3-integration-uv-cache-{{ checksum "uv.lock" }}
       - run:
           name: Install Dependencies
           command: |
@@ -266,7 +266,7 @@ commands:
       - save_cache:
           paths:
             - ~/.cache/uv
-          key: v1-uv-cache-{{ checksum "uv.lock" }}
+          key: v3-integration-uv-cache-{{ checksum "uv.lock" }}
 
 jobs:
   # Add Windows testing job
@@ -2955,6 +2955,32 @@ jobs:
     working_directory: ~/project
     steps:
       - setup_litellm_test_deps
+      - when:
+          condition:
+            equal: [browser, << parameters.suite >>]
+          steps:
+            - install_node
+            - restore_cache:
+                keys:
+                  - integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
+            - run:
+                name: Install locked browser dependencies
+                command: |
+                  cd ui/litellm-dashboard
+                  npm ci
+                  cd ../../tests/e2e/ui
+                  npm ci
+                  sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l \
+                    timeout --signal=TERM --kill-after=20s 6m node node_modules/@playwright/test/cli.js install-deps chromium
+                  timeout --signal=TERM --kill-after=20s 3m node node_modules/@playwright/test/cli.js install chromium
+            - save_cache:
+                key: integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
+                paths:
+                  - ~/.npm
+                  - ~/.cache/ms-playwright
+            - run:
+                name: Build the candidate dashboard
+                command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
       - start_postgres:
           image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
       - start_redis
@@ -2983,7 +3009,7 @@ workflows:
           name: integration-<< matrix.suite >>
           matrix:
             parameters:
-              suite: [management, accounting, providers]
+              suite: [management, accounting, database, providers, extensions, browser]
           filters:
             branches:
               only:

.circleci/scripts/classify_changes.sh

@@ -1,12 +1,14 @@
 #!/usr/bin/env bash
 set -uo pipefail
 
-category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
+category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
 
 has_client=false
 has_backend=false
 has_ci=false
 has_provider_harness=false
+has_cost_map=false
+outside_cost_map_set=false
 while IFS= read -r file || [ -n "$file" ]; do
   [ -n "$file" ] || continue
   case "$file" in
@@ -20,9 +22,18 @@ while IFS= read -r file || [ -n "$file" ]; do
     .github/* | .circleci/*) has_ci=true; has_backend=true ;;
     *) has_backend=true ;;
   esac
+  case "$file" in
+    model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
+      has_cost_map=true ;;
+    tests/test_litellm/* | tests/proxy_unit_tests/*) : ;;
+    *) outside_cost_map_set=true ;;
+  esac
 done
 
 case "$category" in
+  cost-map-only)
+    { [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
+    ;;
   provider-harness)
     [ "$has_provider_harness" = true ] && echo run || echo skip
     ;;

.circleci/scripts/run_integration.sh

@@ -1,6 +1,11 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
+if [ "${GITHUB_ACTIONS:-}" = true ]; then
+  echo "Integration contracts are owned by CircleCI" >&2
+  exit 1
+fi
+
 suite="${1:?integration suite required}"
 results="test-results/integration-${suite}"
 mkdir -p "$results"
@@ -65,7 +70,13 @@ export INTEGRATION_PROXY_URL=http://127.0.0.1:4000
 export INTEGRATION_PEER_URL=""
 export INTEGRATION_UPSTREAM_URL=http://127.0.0.1:8190
 export INTEGRATION_MASTER_KEY="$LITELLM_MASTER_KEY"
-export INTEGRATION_SEED="$((16#$(git rev-parse --short=8 HEAD)))"
+export LITELLM_UI_PATH="$PWD/litellm/proxy/_experimental/out"
+if [ "$suite" = browser ]; then
+  export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
+  test -f "$LITELLM_UI_PATH/index.html"
+fi
+export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
+export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
 
 uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
 
@@ -102,7 +113,7 @@ start_proxy() {
   local log_name="$2"
   setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
     DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
-    LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" \
+    LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
     LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
     AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
     .venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
@@ -131,6 +142,19 @@ if [ "$suite" = providers ]; then
     --junitxml="$results/replay-controls.xml"
 fi
 
+if [ "$suite" = browser ]; then
+  export E2E_UI_BASE_URL="$INTEGRATION_PROXY_URL" E2E_UI_ARTIFACT_DIR="$PWD/$results"
+  export INTEGRATION_PYTHON="$PWD/.venv/bin/python"
+  timeout --signal=TERM --kill-after=20s 3m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
+    INTEGRATION_RUN_ID="$integration_identity" DATABASE_URL="$DATABASE_URL" \
+    INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" INTEGRATION_PYTHON="$INTEGRATION_PYTHON" \
+    E2E_UI_BASE_URL="$E2E_UI_BASE_URL" E2E_UI_ARTIFACT_DIR="$E2E_UI_ARTIFACT_DIR" \
+    LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" CI=true \
+    node tests/e2e/ui/node_modules/@playwright/test/cli.js test --config tests/e2e/ui/integration.config.ts
+  .venv/bin/python .circleci/scripts/verify_integration_browser.py "$results/browser-results.json"
+  exit 0
+fi
+
 timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
   INTEGRATION_RUN_ID="$integration_identity" \
   DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
@@ -138,5 +162,6 @@ timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTH
   INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
   INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
   INTEGRATION_SEED="$INTEGRATION_SEED" \
+  INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
   LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
   .venv/bin/python tests/integration/run.py "$suite" --results "$results"

.circleci/scripts/verify_integration_browser.py

@@ -0,0 +1,60 @@
+import json
+import sys
+from pathlib import Path
+from typing import Final
+
+from pydantic import TypeAdapter
+from typing_extensions import NotRequired, ReadOnly, TypedDict
+
+
+class BrowserAttempt(TypedDict):
+    status: ReadOnly[str]
+    retry: ReadOnly[int]
+
+
+class BrowserTest(TypedDict):
+    results: ReadOnly[list[BrowserAttempt]]
+
+
+class BrowserSpec(TypedDict):
+    file: ReadOnly[str]
+    title: ReadOnly[str]
+    tests: ReadOnly[list[BrowserTest]]
+
+
+class BrowserSuite(TypedDict):
+    specs: NotRequired[ReadOnly[list[BrowserSpec]]]
+    suites: NotRequired[ReadOnly[list["BrowserSuite"]]]
+
+
+def main() -> None:
+    result: Final = json.loads(Path(sys.argv[1]).read_text())
+    assert not result.get("errors"), result.get("errors")
+    expected: Final = json.loads(
+        (Path(__file__).resolve().parents[2] / "tests/integration/contracts.json").read_text()
+    )["browser"]
+    assert expected and result["stats"]["expected"] == len(expected)
+    assert all(result["stats"][name] == 0 for name in ("unexpected", "flaky", "skipped"))
+
+    def cases(suite: BrowserSuite) -> tuple[BrowserSpec, ...]:
+        return tuple(suite.get("specs", ())) + tuple(spec for child in suite.get("suites", ()) for spec in cases(child))
+
+    suites: Final = TypeAdapter(list[BrowserSuite]).validate_python(result["suites"], strict=True)
+    specs: Final = tuple(spec for suite in suites for spec in cases(suite))
+    repository: Final = Path(__file__).resolve().parents[2]
+    report_root: Final = Path(result["config"]["rootDir"])
+    assert report_root.is_absolute(), "Playwright rootDir must be explicit"
+    observed: Final = tuple(
+        str((report_root / spec["file"]).resolve().relative_to(repository)) + "::" + spec["title"] for spec in specs
+    )
+    assert sorted(observed) == sorted(expected)
+    for spec in specs:
+        tests: Final = spec["tests"]
+        assert len(tests) == 1 and len(tests[0]["results"]) == 1
+        assert tests[0]["results"][0]["status"] == "passed" and tests[0]["results"][0]["retry"] == 0
+
+    sys.stdout.write("One canonical browser contract passed once without skips or retries\n")
+
+
+if __name__ == "__main__":
+    main()

.github/scripts/assert_ci_coverage.py

@@ -505,6 +505,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         return frozenset(), ()
     entries: Final = json.loads(manifest.read_text())
     paths: Final = frozenset(node.split("::", 1)[0] for node in entries["tests"])
+    browser_paths: Final = frozenset(node.split("::", 1)[0] for node in entries.get("browser", {}))
     circle_path: Final = repo_root / ".circleci/config.yml"
     circle: Final = yaml.safe_load(circle_path.read_text()) if circle_path.exists() else {}
     steps: Final = circle.get("jobs", {}).get("integration_contracts", {}).get("steps", ())
@@ -523,7 +524,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         .get("suite", (job["integration_contracts"].get("suite"),))
         if isinstance(suite, str)
     )
-    required: Final = frozenset(
+    required: Final = (frozenset({"browser"}) if browser_paths else frozenset()) | frozenset(
         group
         for group, folders in entries["groups"].items()
         if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
@@ -551,6 +552,40 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         for path in paths
         if not (repo_root / path).is_file()
     )
+    browser_commands: Final = tuple(
+        scalar.value
+        for path in (repo_root / ".github/workflows").glob("*.y*ml")
+        for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
+        if scalar.key in {"run", "command"}
+    )
+    browser_findings: Final = tuple(
+        Finding(path, "browser integration contract is explicitly selected by GitHub Actions")
+        for path in browser_paths
+        if any(
+            path in command
+            or pathlib.Path(path).name in command
+            or "integrationCritical" in command
+            or "integration.config.ts" in command
+            or ("run_integration.sh" in command and "browser" in command)
+            for command in browser_commands
+        )
+    ) + tuple(
+        Finding(path, "canonical browser integration file is missing")
+        for path in browser_paths
+        if not (repo_root / path).is_file()
+    )
+    default_browser: Final = repo_root / "tests/e2e/ui/playwright.config.ts"
+    exclusion_findings: Final = (
+        (
+            Finding(
+                str(default_browser.relative_to(repo_root)),
+                "default Playwright selection must exclude integrationCritical",
+            ),
+        )
+        if browser_paths
+        and (not default_browser.exists() or "**/integrationCritical/**" not in default_browser.read_text())
+        else ()
+    )
     group_findings: Final = tuple(
         Finding(group, "canonical integration group is not scheduled by CircleCI")
         for group in sorted(required - scheduled)
@@ -559,7 +594,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         return frozenset(), findings + (
             Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
         )
-    return paths, findings + group_findings
+    return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings
 
 
 def main() -> int:

.github/scripts/auto_merge_price_sync.py

@@ -0,0 +1,465 @@
+"""Auto-merge the provider-info-sync bot's cost-map pull requests.
+
+Evaluates every gate (author allowlist, cost-map-only diff, required and
+non-required checks, Greptile confidence, Bugbot review, human reviews) and
+merges with a merge commit when all of them hold. Every hold reason is
+logged; the process exits 0 on hold and 1 only on API or programming errors.
+``DRY_RUN=1`` prints the verdict without calling the merge endpoint.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import subprocess
+import sys
+import time
+import urllib.error
+import urllib.request
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Final
+
+REPO_ROOT: Final = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classify_changes.sh")
+API_ROOT: Final = "https://api.github.com"
+CHANGED_FILE_CEILING: Final = 3000
+OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"})
+GREPTILE_LOGIN: Final = "greptile-apps[bot]"
+BUGBOT_LOGIN: Final = "cursor[bot]"
+GREPTILE_SCORE_RE: Final = re.compile(r"Confidence Score:\s*(\d)/5")
+BUGBOT_REVIEW_MARKER: Final = "<!-- BUGBOT_REVIEW -->"
+BUGBOT_STALE_MARKER: Final = "<!-- BUGBOT_REVIEW_STALE -->"
+BUGBOT_CLEAN: Final = "found no new issues"
+
+
+@dataclass(frozen=True, slots=True)
+class PullRequest:
+    number: int
+    title: str
+    author_login: str
+    state: str
+    draft: bool
+    mergeable: bool | None
+    mergeable_state: str
+    head_sha: str
+
+
+@dataclass(frozen=True, slots=True)
+class CheckRun:
+    name: str
+    status: str
+    conclusion: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class CommitStatus:
+    context: str
+    state: str
+
+
+@dataclass(frozen=True, slots=True)
+class IssueComment:
+    author_login: str
+    body: str
+    updated_at: datetime
+
+
+@dataclass(frozen=True, slots=True)
+class Review:
+    author_login: str
+    state: str
+    body: str
+    commit_id: str
+    submitted_at: datetime
+
+
+@dataclass(frozen=True, slots=True)
+class Verdict:
+    merge: bool
+    reasons: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class EvaluationInputs:
+    pr: PullRequest
+    changed_files: tuple[str, ...]
+    required_contexts: frozenset[str]
+    check_runs: tuple[CheckRun, ...]
+    statuses: tuple[CommitStatus, ...]
+    comments: tuple[IssueComment, ...]
+    reviews: tuple[Review, ...]
+    head_commit_date: datetime
+    self_check_name: str
+    author_allowlist: frozenset[str]
+
+
+def _is_bot_login(login: str) -> bool:
+    return login.lower().endswith("[bot]")
+
+
+def _classify(changed_files: Sequence[str]) -> str:
+    result: Final = subprocess.run(
+        ["bash", CLASSIFY_SCRIPT, "cost-map-only"],
+        input="\n".join(changed_files),
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if result.returncode != 0:
+        return "error"
+    return result.stdout.strip()
+
+
+def evaluate(
+    inputs: EvaluationInputs,
+    *,
+    classify: Callable[[Sequence[str]], str] = _classify,
+) -> Verdict:
+    pr: Final = inputs.pr
+    reasons: list[str] = []
+
+    if pr.author_login.lower() not in {login.lower() for login in inputs.author_allowlist}:
+        reasons.append(f"author {pr.author_login!r} not in allowlist")
+    if pr.state != "open":
+        reasons.append("pr not open")
+    if pr.draft:
+        reasons.append("pr is a draft")
+    if pr.mergeable is None:
+        reasons.append("mergeability unknown")
+    elif not pr.mergeable:
+        reasons.append("pr not mergeable")
+    if pr.mergeable_state == "dirty":
+        reasons.append("pr has merge conflicts")
+
+    if len(inputs.changed_files) > CHANGED_FILE_CEILING:
+        reasons.append(f"changed file count {len(inputs.changed_files)} over {CHANGED_FILE_CEILING} ceiling")
+    else:
+        decision: Final = classify(inputs.changed_files)
+        if decision != "run":
+            reasons.append("changed files outside the cost-map-only set")
+
+    green_runs: Final = frozenset(run.name for run in inputs.check_runs if run.conclusion in OK_CHECK_CONCLUSIONS)
+    green_statuses: Final = frozenset(status.context for status in inputs.statuses if status.state == "success")
+    for context in sorted(inputs.required_contexts):
+        if context not in green_runs and context not in green_statuses:
+            reasons.append(f"required check {context!r} not green")
+    for run in inputs.check_runs:
+        if run.name == inputs.self_check_name:
+            continue
+        if run.status != "completed" or run.conclusion not in OK_CHECK_CONCLUSIONS:
+            reasons.append(f"check run {run.name!r} is {run.status}/{run.conclusion}")
+    for status in inputs.statuses:
+        if status.state != "success":
+            reasons.append(f"commit status {status.context!r} is {status.state}")
+
+    greptile: Final = tuple(
+        comment
+        for comment in inputs.comments
+        if comment.author_login == GREPTILE_LOGIN and GREPTILE_SCORE_RE.search(comment.body)
+    )
+    if not greptile:
+        reasons.append("greptile score not available")
+    else:
+        latest: Final = max(greptile, key=lambda comment: comment.updated_at)
+        match: Final = GREPTILE_SCORE_RE.search(latest.body)
+        score: Final = int(match.group(1)) if match else 0
+        if latest.updated_at < inputs.head_commit_date:
+            reasons.append("greptile score older than head commit")
+        elif score != 5:
+            reasons.append(f"greptile score {score}/5 below 5")
+
+    bugbot: Final = tuple(
+        review
+        for review in inputs.reviews
+        if review.author_login == BUGBOT_LOGIN
+        and BUGBOT_REVIEW_MARKER in review.body
+        and BUGBOT_STALE_MARKER not in review.body
+        and review.commit_id == pr.head_sha
+    )
+    if not bugbot:
+        reasons.append("bugbot review not available")
+    else:
+        latest_review: Final = max(bugbot, key=lambda review: review.submitted_at)
+        if BUGBOT_CLEAN not in latest_review.body:
+            reasons.append("bugbot reported issues")
+
+    latest_state_by_reviewer: Final[dict[str, str]] = {}
+    for review in sorted(inputs.reviews, key=lambda review: review.submitted_at):
+        if _is_bot_login(review.author_login):
+            continue
+        latest_state_by_reviewer[review.author_login] = review.state
+    for reviewer, state in latest_state_by_reviewer.items():
+        if state == "CHANGES_REQUESTED":
+            reasons.append(f"changes requested by {reviewer}")
+
+    return Verdict(merge=not reasons, reasons=tuple(reasons))
+
+
+def _request(token: str, method: str, path: str, body: Mapping[str, object] | None = None) -> object:
+    url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
+    data: Final = None if body is None else json.dumps(body).encode("utf-8")
+    request: Final = urllib.request.Request(
+        url,
+        data=data,
+        method=method,
+        headers={
+            "Accept": "application/vnd.github+json",
+            "Authorization": f"Bearer {token}",
+            "X-GitHub-Api-Version": "2022-11-28",
+        },
+    )
+    with urllib.request.urlopen(request) as response:
+        return json.loads(response.read().decode("utf-8"))
+
+
+def _request_allow_fail(
+    token: str, method: str, path: str, body: Mapping[str, object] | None = None
+) -> tuple[int, object | None]:
+    url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
+    data: Final = None if body is None else json.dumps(body).encode("utf-8")
+    request: Final = urllib.request.Request(
+        url,
+        data=data,
+        method=method,
+        headers={
+            "Accept": "application/vnd.github+json",
+            "Authorization": f"Bearer {token}",
+            "X-GitHub-Api-Version": "2022-11-28",
+        },
+    )
+    try:
+        with urllib.request.urlopen(request) as response:
+            return response.status, json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        return exc.code, None
+
+
+def _items(payload: object, key: str | None = None) -> tuple[object, ...]:
+    source: Final = payload.get(key) if key and isinstance(payload, Mapping) else payload
+    if not isinstance(source, list):
+        return ()
+    return tuple(source)
+
+
+def _paginate(token: str, path: str, key: str | None = None) -> list[object]:
+    separator: Final = "&" if "?" in path else "?"
+    results: list[object] = []
+    for page in range(1, 10_000):
+        batch: Final = _items(_request(token, "GET", f"{path}{separator}per_page=100&page={page}"), key)
+        results.extend(batch)
+        if len(batch) < 100:
+            return results
+    return results
+
+
+def _text(value: object) -> str:
+    return value if isinstance(value, str) else ""
+
+
+def _int(value: object) -> int:
+    return value if isinstance(value, int) else 0
+
+
+def _bool(value: object) -> bool:
+    return value is True
+
+
+def _nested(value: object, *keys: str) -> object:
+    current: object = value
+    for key in keys:
+        if not isinstance(current, Mapping):
+            return None
+        current = current.get(key)
+    return current
+
+
+def _parse_time(value: object) -> datetime:
+    text: Final = _text(value)
+    if not text:
+        return datetime.min.replace(tzinfo=timezone.utc)
+    return datetime.fromisoformat(text.replace("Z", "+00:00"))
+
+
+def _load_pr(token: str, repo: str, number: int) -> PullRequest:
+    data: Final = _request(token, "GET", f"/repos/{repo}/pulls/{number}")
+    if not isinstance(data, Mapping):
+        raise RuntimeError(f"unexpected pull payload for #{number}")
+    return PullRequest(
+        number=number,
+        title=_text(data.get("title")),
+        author_login=_text(_nested(data, "user", "login")),
+        state=_text(data.get("state")),
+        draft=_bool(data.get("draft")),
+        mergeable=data.get("mergeable") if isinstance(data.get("mergeable"), bool) else None,
+        mergeable_state=_text(data.get("mergeable_state")),
+        head_sha=_text(_nested(data, "head", "sha")),
+    )
+
+
+def _list_candidate_prs(token: str, repo: str, base: str, allowlist: frozenset[str]) -> list[int]:
+    candidates: Final = _paginate(token, f"/repos/{repo}/pulls?state=open&base={base}")
+    return [
+        _int(item.get("number"))
+        for item in candidates
+        if isinstance(item, Mapping) and _text(_nested(item, "user", "login")).lower() in allowlist
+    ]
+
+
+def _changed_files(token: str, repo: str, number: int) -> tuple[str, ...]:
+    files: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/files")
+    return tuple(_text(item.get("filename")) for item in files if isinstance(item, Mapping))
+
+
+def _required_contexts(token: str, repo: str, base: str) -> frozenset[str]:
+    payload: Final = _request(token, "GET", f"/repos/{repo}/rules/branches/{base}")
+    contexts: set[str] = set()
+    for rule in _items(payload):
+        if not isinstance(rule, Mapping) or rule.get("type") != "required_status_checks":
+            continue
+        checks: Final = _nested(rule, "parameters", "required_status_checks")
+        for check in _items(checks):
+            if isinstance(check, Mapping):
+                context: Final = _text(check.get("context"))
+                if context:
+                    contexts.add(context)
+    return frozenset(contexts)
+
+
+def _check_runs(token: str, repo: str, sha: str) -> tuple[CheckRun, ...]:
+    runs: Final = _paginate(token, f"/repos/{repo}/commits/{sha}/check-runs", key="check_runs")
+    return tuple(
+        CheckRun(
+            name=_text(item.get("name")),
+            status=_text(item.get("status")),
+            conclusion=item.get("conclusion") if isinstance(item.get("conclusion"), str) else None,
+        )
+        for item in runs
+        if isinstance(item, Mapping)
+    )
+
+
+def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]:
+    payload: Final = _request(token, "GET", f"/repos/{repo}/commits/{sha}/status")
+    return tuple(
+        CommitStatus(context=_text(item.get("context")), state=_text(item.get("state")))
+        for item in _items(payload, "statuses")
+        if isinstance(item, Mapping)
+    )
+
+
+def _comments(token: str, repo: str, number: int) -> tuple[IssueComment, ...]:
+    comments: Final = _paginate(token, f"/repos/{repo}/issues/{number}/comments")
+    return tuple(
+        IssueComment(
+            author_login=_text(_nested(item, "user", "login")),
+            body=_text(item.get("body")),
+            updated_at=_parse_time(item.get("updated_at")),
+        )
+        for item in comments
+        if isinstance(item, Mapping)
+    )
+
+
+def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
+    reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews")
+    return tuple(
+        Review(
+            author_login=_text(_nested(item, "user", "login")),
+            state=_text(item.get("state")),
+            body=_text(item.get("body")),
+            commit_id=_text(item.get("commit_id")),
+            submitted_at=_parse_time(item.get("submitted_at")),
+        )
+        for item in reviews
+        if isinstance(item, Mapping)
+    )
+
+
+def _head_commit_date(token: str, repo: str, number: int) -> datetime:
+    commits: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/commits")
+    if not commits:
+        return datetime.min.replace(tzinfo=timezone.utc)
+    last: Final = commits[-1]
+    if not isinstance(last, Mapping):
+        return datetime.min.replace(tzinfo=timezone.utc)
+    return _parse_time(_nested(last, "commit", "committer", "date"))
+
+
+def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest:
+    if pr.mergeable is not None:
+        return pr
+    time.sleep(5)
+    return _load_pr(token, repo, pr.number)
+
+
+def _gather_inputs(
+    token: str,
+    repo: str,
+    number: int,
+    base: str,
+    self_check_name: str,
+    allowlist: frozenset[str],
+) -> EvaluationInputs:
+    pr: Final = _mergeable_or_refetch(token, repo, _load_pr(token, repo, number))
+    return EvaluationInputs(
+        pr=pr,
+        changed_files=_changed_files(token, repo, number),
+        required_contexts=_required_contexts(token, repo, base),
+        check_runs=_check_runs(token, repo, pr.head_sha),
+        statuses=_statuses(token, repo, pr.head_sha),
+        comments=_comments(token, repo, number),
+        reviews=_reviews(token, repo, number),
+        head_commit_date=_head_commit_date(token, repo, number),
+        self_check_name=self_check_name,
+        author_allowlist=allowlist,
+    )
+
+
+def merge_request_body(pr: PullRequest) -> dict[str, str]:
+    return {"merge_method": "merge", "commit_title": f"{pr.title} (#{pr.number})", "sha": pr.head_sha}
+
+
+def _merge(token: str, repo: str, pr: PullRequest) -> None:
+    status, _ = _request_allow_fail(token, "PUT", f"/repos/{repo}/pulls/{pr.number}/merge", merge_request_body(pr))
+    if status in (200, 405, 409):
+        print(f"auto-merge-price-sync: PR #{pr.number} merge call returned {status}")
+        return
+    raise RuntimeError(f"merge call for PR #{pr.number} returned {status}")
+
+
+def main() -> int:
+    token: Final = os.environ.get("GH_TOKEN", "")
+    repo: Final = os.environ.get("REPO", "")
+    base: Final = os.environ.get("BASE_BRANCH", "main")
+    dry_run: Final = os.environ.get("DRY_RUN", "") != ""
+    self_check_name: Final = os.environ.get("SELF_CHECK_NAME", "auto-merge-price-sync")
+    allowlist: Final = frozenset(login.lower() for login in os.environ.get("PR_AUTHOR_ALLOWLIST", "").split() if login)
+    if not token:
+        print("auto-merge-price-sync: app credentials not configured")
+        return 0
+    if not repo:
+        print("auto-merge-price-sync: REPO not set", file=sys.stderr)
+        return 1
+
+    pr_number_env: Final = os.environ.get("PR_NUMBER", "")
+    candidates: Final = [int(pr_number_env)] if pr_number_env else _list_candidate_prs(token, repo, base, allowlist)
+    for number in candidates:
+        inputs: Final = _gather_inputs(token, repo, number, base, self_check_name, allowlist)
+        verdict: Final = evaluate(inputs)
+        for reason in verdict.reasons:
+            print(f"auto-merge-price-sync: PR #{number} hold: {reason}")
+        if not verdict.merge:
+            continue
+        print(f"auto-merge-price-sync: PR #{number} all gates green")
+        if dry_run:
+            print(f"auto-merge-price-sync: DRY_RUN merge suppressed for PR #{number}")
+            continue
+        _merge(token, repo, inputs.pr)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/auto-merge-price-sync.yml

@@ -0,0 +1,61 @@
+name: auto-merge-price-sync
+
+on:
+  issue_comment:
+    types: [created, edited]
+  check_suite:
+    types: [completed]
+  status: {}
+  schedule:
+    - cron: "*/30 * * * *"
+  workflow_dispatch:
+    inputs:
+      pr-number:
+        description: "Evaluate only this PR number (empty = scan all open sync-bot PRs)"
+        required: false
+        default: ""
+
+permissions:
+  contents: read
+  pull-requests: read
+  checks: read
+  statuses: read
+
+concurrency:
+  group: auto-merge-price-sync
+  cancel-in-progress: false
+
+jobs:
+  auto-merge-price-sync:
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    env:
+      PROVIDER_INFO_SYNC_APP_ID: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
+      PROVIDER_INFO_SYNC_APP_PRIVATE_KEY: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - name: Mint app token
+        id: app-token
+        if: ${{ env.PROVIDER_INFO_SYNC_APP_ID != '' && env.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY != '' }}
+        uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
+        with:
+          app-id: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
+          private-key: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
+
+      - name: Auto-merge eligible sync PRs
+        env:
+          GH_TOKEN: ${{ steps.app-token.outputs.token }}
+          REPO: ${{ github.repository }}
+          PR_NUMBER: ${{ (github.event.issue.pull_request && github.event.issue.number) || github.event.inputs.pr-number || '' }}
+          BASE_BRANCH: main
+          PR_AUTHOR_ALLOWLIST: "berriai-litellm-provider-info-sync[bot]"
+          SELF_CHECK_NAME: auto-merge-price-sync
+        run: python3 .github/scripts/auto_merge_price_sync.py

.github/workflows/test-rust.yml

@@ -70,7 +70,7 @@ env:
 jobs:
   rust-lint:
     runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
     defaults:
       run:
         working-directory: litellm-rust
@@ -81,24 +81,48 @@ jobs:
 
       - run: rustup toolchain install --no-self-update
 
-      - run: cargo fmt --check
+      - run: cargo fmt --all --check
 
-      - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
         with:
-          path: |
-            ~/.cargo/registry
-            ~/.cargo/git
-            litellm-rust/target
-          key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-cargo-${{ github.job }}-
+          workspaces: litellm-rust
+          cache-on-failure: true
 
       - run: cargo clippy --workspace --all-targets --locked -- -D warnings
 
   rust-test:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 20
+    defaults:
+      run:
+        working-directory: litellm-rust
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
 
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - run: rustup toolchain install --no-self-update
+
+      - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8
+        with:
+          tool: cargo-nextest@0.9.143
+
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
+        with:
+          workspaces: litellm-rust
+          cache-on-failure: true
+
+      - run: cargo nextest run --workspace --locked
+
+      - run: cargo test --workspace --doc --locked
+
+  rust-wheel:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
         with:
@@ -114,18 +138,10 @@ jobs:
 
       - run: rustup toolchain install --no-self-update
 
-      - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
         with:
-          path: |
-            ~/.cargo/registry
-            ~/.cargo/git
-            litellm-rust/target
-          key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-cargo-${{ github.job }}-
-
-      - run: cargo test --workspace --locked
-        working-directory: litellm-rust
+          workspaces: litellm-rust
+          cache-on-failure: true
 
       - run: uv build --wheel --out-dir dist
 

ARCHITECTURE.md

@@ -45,7 +45,7 @@ sequenceDiagram
     ProxyServer->>Auth: user_api_key_auth()
     Auth->>Redis: Check API key cache
     Redis-->>Auth: Key info + spend limits
-    ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
+    ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
     Hooks->>Redis: Check/increment rate limit counters
     ProxyServer->>Router: route_request()
     Router->>Main: litellm.acompletion()
@@ -145,7 +145,6 @@ graph TD
 
 | Hook | File | Purpose |
 |------|------|---------|
-| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
 | `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
 | `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
 | `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
     key_alias   String?
     soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
     spend      Float    @default(0.0)
+    total_spend Float   @default(0.0)
     expires    DateTime?
     models     String[]
     aliases    Json  @default("{}")
@@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
     key_alias               String?
     soft_budget_cooldown    Boolean  @default(false)
     spend                   Float    @default(0.0)
+    total_spend             Float    @default(0.0)
     expires                 DateTime?
     models                  String[]
     aliases                 Json     @default("{}")

litellm-rust/Cargo.lock

@@ -70,6 +70,12 @@ dependencies = [
  "rustversion",
 ]
 
+[[package]]
+name = "arcstr"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
+
 [[package]]
 name = "async-compression"
 version = "0.4.46"
@@ -262,6 +268,17 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "aws-smithy-eventstream"
+version = "0.61.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
+dependencies = [
+ "aws-smithy-types",
+ "bytes",
+ "crc32fast",
+]
+
 [[package]]
 name = "aws-smithy-http"
 version = "0.64.0"
@@ -1837,6 +1854,12 @@ version = "0.2.186"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
 
+[[package]]
+name = "linux-raw-sys"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
+
 [[package]]
 name = "litellm-auth"
 version = "0.1.0"
@@ -1915,6 +1938,17 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "litellm-cache-redis"
+version = "0.1.0"
+dependencies = [
+ "litellm-cache",
+ "redis",
+ "redis-test",
+ "serde_json",
+ "tokio",
+]
+
 [[package]]
 name = "litellm-core"
 version = "0.1.0"
@@ -1947,10 +1981,25 @@ dependencies = [
  "veil",
 ]
 
+[[package]]
+name = "litellm-framing"
+version = "0.1.0"
+dependencies = [
+ "aws-smithy-eventstream",
+ "aws-smithy-types",
+ "bytes",
+ "futures-util",
+ "rstest",
+ "sse-stream",
+ "thiserror 2.0.19",
+ "tokio",
+]
+
 [[package]]
 name = "litellm-python-bridge"
 version = "0.1.0"
 dependencies = [
+ "bytes",
  "criterion",
  "futures-util",
  "litellm-auth",
@@ -2139,6 +2188,16 @@ dependencies = [
  "minimal-lexical",
 ]
 
+[[package]]
+name = "num-bigint"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
+dependencies = [
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-conv"
 version = "0.2.2"
@@ -2655,6 +2714,36 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "redis"
+version = "1.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
+dependencies = [
+ "arcstr",
+ "combine",
+ "itoa",
+ "num-bigint",
+ "percent-encoding",
+ "ryu",
+ "sha1_smol",
+ "socket2 0.6.5",
+ "url",
+ "xxhash-rust",
+]
+
+[[package]]
+name = "redis-test"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
+dependencies = [
+ "rand 0.9.5",
+ "redis",
+ "socket2 0.6.5",
+ "tempfile",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.5.18"
@@ -2845,6 +2934,19 @@ dependencies = [
  "semver",
 ]
 
+[[package]]
+name = "rustix"
+version = "1.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
+dependencies = [
+ "bitflags",
+ "errno",
+ "libc",
+ "linux-raw-sys",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "rustls"
 version = "0.21.12"
@@ -3095,6 +3197,12 @@ dependencies = [
  "digest 0.10.7",
 ]
 
+[[package]]
+name = "sha1_smol"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
+
 [[package]]
 name = "sha2"
 version = "0.10.9"
@@ -3199,6 +3307,19 @@ dependencies = [
  "unicode-segmentation",
 ]
 
+[[package]]
+name = "sse-stream"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
+dependencies = [
+ "bytes",
+ "futures-util",
+ "http-body 1.1.0",
+ "http-body-util",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "stable_deref_trait"
 version = "1.2.1"
@@ -3298,6 +3419,19 @@ version = "0.13.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
 
+[[package]]
+name = "tempfile"
+version = "3.27.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
+dependencies = [
+ "fastrand",
+ "getrandom 0.4.3",
+ "once_cell",
+ "rustix",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.69"
@@ -4181,6 +4315,12 @@ version = "0.13.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
 
+[[package]]
+name = "xxhash-rust"
+version = "0.8.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
+
 [[package]]
 name = "yoke"
 version = "0.8.3"

litellm-rust/crates/cache-redis/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "litellm-cache-redis"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-cache.workspace = true
+redis = "1.7.0"
+serde_json.workspace = true
+tokio.workspace = true
+
+[dev-dependencies]
+redis-test = "1.0.4"

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

@@ -0,0 +1,315 @@
+use std::sync::{Arc, Mutex, MutexGuard};
+use std::time::Duration;
+
+use litellm_cache::{
+    BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
+    Error,
+};
+use redis::Commands;
+
+const DEFAULT_TTL: Duration = Duration::from_secs(600);
+const KEY_PREFIX: &str = "litellm-cache:";
+
+pub struct RedisCache<C = redis::Connection> {
+    connection: Arc<Mutex<C>>,
+    default_ttl: Duration,
+}
+
+impl RedisCache<redis::Connection> {
+    pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
+        let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
+        let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
+        Ok(Self::with_connection(connection, default_ttl))
+    }
+}
+
+impl<C> RedisCache<C>
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
+        Self {
+            connection: Arc::new(Mutex::new(connection)),
+            default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
+        }
+    }
+
+    fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
+        self.connection.lock().map_err(|_| Error::Unavailable)
+    }
+
+    fn namespaced_key(key: &str) -> String {
+        format!("{KEY_PREFIX}{key}")
+    }
+
+    fn namespaced_pattern() -> &'static str {
+        const PATTERN: &str = "litellm-cache:*";
+        PATTERN
+    }
+
+    fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
+        serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
+    }
+
+    fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
+        serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
+    }
+
+    fn ttl_seconds(ttl: Duration) -> u64 {
+        ttl.as_secs()
+            .saturating_add(u64::from(ttl.subsec_nanos() > 0))
+            .max(1)
+    }
+
+    fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
+    where
+        T: Send + 'static,
+        F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
+    {
+        Box::pin(async move {
+            tokio::task::spawn_blocking(move || {
+                let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
+                operation(&mut connection)
+            })
+            .await
+            .map_err(|_| Error::Unavailable)?
+        })
+    }
+}
+
+impl<C> BaseCache for RedisCache<C>
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type Value = CacheEntry;
+
+    fn default_ttl(&self) -> Duration {
+        self.default_ttl
+    }
+
+    fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
+        let payload = Self::encode(&value)?;
+        let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
+        self.connection()?
+            .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
+            .map_err(|_| Error::Unavailable)
+    }
+
+    fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
+        self.connection()?
+            .get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
+            .map_err(|_| Error::Unavailable)?
+            .map(Self::decode)
+            .transpose()
+    }
+
+    fn delete_cache(&self, key: &str) -> Result<(), Error> {
+        self.connection()?
+            .del::<_, ()>(Self::namespaced_key(key))
+            .map_err(|_| Error::Unavailable)
+    }
+
+    fn flush_cache(&self) -> Result<(), Error> {
+        let mut connection = self.connection()?;
+        let keys = connection
+            .scan_match(Self::namespaced_pattern())
+            .map_err(|_| Error::Unavailable)?
+            .collect::<redis::RedisResult<Vec<String>>>()
+            .map_err(|_| Error::Unavailable)?;
+        if keys.is_empty() {
+            return Ok(());
+        }
+        connection
+            .del::<_, usize>(keys)
+            .map(|_| ())
+            .map_err(|_| Error::Unavailable)
+    }
+
+    fn async_set_cache<'a>(
+        &'a self,
+        key: &'a str,
+        value: Self::Value,
+        kwargs: CacheKwargs,
+    ) -> CacheFuture<'a, ()> {
+        let payload = Self::encode(&value);
+        let key = Self::namespaced_key(key);
+        let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
+        Self::run_blocking(Arc::clone(&self.connection), move |connection| {
+            connection
+                .set_ex::<_, _, ()>(key, payload?, ttl)
+                .map_err(|_| Error::Unavailable)
+        })
+    }
+
+    fn async_get_cache<'a>(
+        &'a self,
+        key: &'a str,
+        _: &'a CacheKwargs,
+    ) -> CacheFuture<'a, Option<Self::Value>> {
+        let key = Self::namespaced_key(key);
+        Box::pin(async move {
+            Self::run_blocking(Arc::clone(&self.connection), move |connection| {
+                connection
+                    .get::<_, Option<Vec<u8>>>(key)
+                    .map_err(|_| Error::Unavailable)
+            })
+            .await?
+            .map(Self::decode)
+            .transpose()
+        })
+    }
+
+    fn async_set_cache_pipeline<'a>(
+        &'a self,
+        cache_list: Vec<(String, Self::Value)>,
+        kwargs: CacheKwargs,
+    ) -> CacheFuture<'a, ()> {
+        let entries = cache_list
+            .into_iter()
+            .map(|(key, value)| {
+                Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
+            })
+            .collect::<Result<Vec<_>, _>>();
+        let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
+        Self::run_blocking(Arc::clone(&self.connection), move |connection| {
+            for (key, payload) in entries? {
+                connection
+                    .set_ex::<_, _, ()>(key, payload, ttl)
+                    .map_err(|_| Error::Unavailable)?;
+            }
+            Ok(())
+        })
+    }
+
+    fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
+        let key = Self::namespaced_key(key);
+        Self::run_blocking(Arc::clone(&self.connection), move |connection| {
+            connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
+        })
+    }
+
+    fn disconnect(&self) -> CacheFuture<'_, ()> {
+        Box::pin(async { Ok(()) })
+    }
+
+    fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
+        Box::pin(async move {
+            Self::run_blocking(Arc::clone(&self.connection), |connection| {
+                redis::cmd("PING")
+                    .query::<String>(connection)
+                    .map_err(|_| Error::Unavailable)
+            })
+            .await?;
+            Ok(CacheConnectionResult {
+                status: CacheConnectionStatus::Success,
+                message: "Redis cache connection test successful".into(),
+                error: None,
+            })
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::RedisCache;
+    use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
+    use redis_test::{MockCmd, MockRedisConnection};
+    use serde_json::json;
+    use std::time::Duration;
+
+    fn entry() -> CacheEntry {
+        CacheEntry {
+            timestamp: 123.0,
+            response: json!({"choices": [{"text": "cached"}]}),
+        }
+    }
+
+    #[test]
+    fn cache_entries_round_trip_through_json() {
+        let entry = entry();
+        let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
+        assert_eq!(
+            RedisCache::<redis::Connection>::decode(encoded).unwrap(),
+            entry
+        );
+    }
+
+    #[test]
+    fn invalid_json_is_rejected() {
+        assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
+    }
+
+    #[test]
+    fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
+        assert_eq!(
+            RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
+            1
+        );
+        assert_eq!(
+            RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
+            2
+        );
+        assert_eq!(
+            RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
+            15
+        );
+    }
+
+    #[test]
+    fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
+        let value = entry();
+        let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
+        let connection = MockRedisConnection::new([
+            MockCmd::new(
+                redis::cmd("SETEX")
+                    .arg("litellm-cache:key")
+                    .arg(600)
+                    .arg(payload.clone()),
+                Ok("OK"),
+            ),
+            MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
+            MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
+        ])
+        .assert_all_commands_consumed();
+        let cache = RedisCache::with_connection(connection, None);
+
+        cache
+            .set_cache("key", value.clone(), CacheKwargs::default())
+            .unwrap();
+        assert_eq!(
+            cache.get_cache("key", &CacheKwargs::default()).unwrap(),
+            Some(value)
+        );
+        cache.delete_cache("key").unwrap();
+    }
+
+    #[test]
+    fn flush_scans_and_deletes_only_cache_keys() {
+        let connection = MockRedisConnection::new([
+            MockCmd::new(
+                redis::cmd("SCAN")
+                    .cursor_arg(0)
+                    .arg("MATCH")
+                    .arg("litellm-cache:*"),
+                Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
+            ),
+            MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
+        ])
+        .assert_all_commands_consumed();
+        let cache = RedisCache::with_connection(connection, None);
+
+        cache.flush_cache().unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_connection_runs_ping_off_executor() {
+        let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
+            .assert_all_commands_consumed();
+        let cache = RedisCache::with_connection(connection, None);
+
+        assert_eq!(
+            cache.test_connection().await.unwrap().status,
+            litellm_cache::CacheConnectionStatus::Success
+        );
+    }
+}

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

@@ -0,0 +1,3 @@
+mod cache;
+
+pub use cache::RedisCache;

litellm-rust/crates/cache-redis/tests/cache.rs

@@ -0,0 +1,6 @@
+use litellm_cache_redis::RedisCache;
+
+#[test]
+fn constructor_rejects_invalid_urls() {
+    assert!(RedisCache::new("not a redis url", None).is_err());
+}

litellm-rust/crates/core/src/ocr/client.rs

@@ -54,9 +54,16 @@ impl OcrClient {
             match call.resume(result.take()).await? {
                 OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
                     result = Some(OcrHostResult::Request(Ok((
-                        Box::new(request.take().ok_or_else(|| {
-                            Error::InvalidRequest("OCR request was already projected".into())
-                        })?),
+                        Box::new(
+                            request
+                                .take()
+                                .ok_or_else(|| {
+                                    Error::InvalidRequest(
+                                        "OCR request was already projected".into(),
+                                    )
+                                })?
+                                .into(),
+                        ),
                         false,
                     ))))
                 }

litellm-rust/crates/core/src/ocr/document.rs

@@ -1,3 +1,6 @@
+use std::io::Read;
+use std::path::Path;
+
 use base64::{Engine, engine::general_purpose::STANDARD};
 use data_url::mime::Mime;
 use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
@@ -5,12 +8,52 @@ use reqwest::Url;
 use serde_json::Map;
 
 use super::error::{OcrError, OcrRequestError, OcrResponseError};
-use super::types::{OcrConnection, OcrDocument};
+use super::types::{OcrConnection, OcrDocument, OcrDocumentInput};
 use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
 use crate::media::Error as MediaError;
 use crate::media::{DownloadPolicy, MediaFetcher};
 use crate::transport::Error as TransportError;
 
+pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, super::Error> {
+    match input {
+        OcrDocumentInput::Document(document) => Ok(document),
+        OcrDocumentInput::Path { path, mime_type } => {
+            read_path_document(&path, mime_type.as_deref())
+        }
+        OcrDocumentInput::Bytes {
+            bytes,
+            file_name,
+            mime_type,
+        } => Ok(encode_file_document(
+            &bytes,
+            file_name.as_deref(),
+            mime_type.as_deref(),
+        )?),
+        OcrDocumentInput::HostReader { .. } => Err(super::Error::InvalidRequest(
+            "OCR file reader was not read by the host".into(),
+        )),
+    }
+}
+
+pub fn read_path_document(
+    path: &Path,
+    mime_type: Option<&str>,
+) -> Result<OcrDocument, super::Error> {
+    let mut bytes = Vec::new();
+    std::fs::File::open(path)
+        .and_then(|file| {
+            file.take(OCR_INLINE_MAX_BYTES as u64 + 1)
+                .read_to_end(&mut bytes)
+        })
+        .map_err(|source| super::Error::FileRead {
+            path: path.to_owned(),
+            kind: source.kind(),
+            message: source.to_string(),
+        })?;
+    let name = path.file_name().map(|name| name.to_string_lossy());
+    Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?)
+}
+
 pub fn encode_file_document(
     bytes: &[u8],
     file_name: Option<&str>,
@@ -75,18 +118,6 @@ pub fn mime_type_for_name(name: &str) -> &'static str {
     }
 }
 
-pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str {
-    match content_type
-        .and_then(|value| value.split(';').next())
-        .map(str::trim)
-    {
-        Some(value) if !value.is_empty() && value != "application/octet-stream" => value,
-        _ => file_name
-            .map(mime_type_for_name)
-            .unwrap_or("application/octet-stream"),
-    }
-}
-
 pub(crate) struct InlineDocument<'a>(DataUrl<'a>);
 
 impl<'a> InlineDocument<'a> {
@@ -230,24 +261,65 @@ mod tests {
     }
 
     #[test]
-    fn upload_mime_mapping_matches_python() {
+    fn path_documents_are_read_and_named_by_core() {
+        let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
+        std::fs::create_dir_all(&dir).unwrap();
+        let path = dir.join("scan.png");
+        std::fs::write(&path, b"abc").unwrap();
         assert_eq!(
-            upload_mime_type(Some("report.pdf"), Some("application/octet-stream")),
-            "application/pdf"
-        );
-        assert_eq!(upload_mime_type(Some("image.png"), None), "image/png");
-        assert_eq!(upload_mime_type(None, None), "application/octet-stream");
-        assert_eq!(
-            upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")),
-            "application/pdf"
+            prepare_document(OcrDocumentInput::Path {
+                path: path.clone(),
+                mime_type: None,
+            })
+            .unwrap(),
+            OcrDocument::ImageUrl {
+                image_url: "data:image/png;base64,YWJj".into(),
+                extra_fields: Map::new(),
+            }
         );
         assert_eq!(
-            upload_mime_type(
-                Some("img.png"),
-                Some("image/png; charset=utf-8; boundary=something")
-            ),
-            "image/png"
+            prepare_document(OcrDocumentInput::Path {
+                path: path.clone(),
+                mime_type: Some("application/pdf".into()),
+            })
+            .unwrap(),
+            document("data:application/pdf;base64,YWJj")
         );
+        std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap();
+        assert_eq!(
+            prepare_document(OcrDocumentInput::Path {
+                path: path.clone(),
+                mime_type: None,
+            }),
+            Err(OcrRequestError::InlineDocumentTooLarge.into())
+        );
+        std::fs::remove_dir_all(&dir).unwrap();
+
+        let missing = dir.join("missing.pdf");
+        let Err(super::super::Error::FileRead { path, kind, .. }) =
+            prepare_document(OcrDocumentInput::Path {
+                path: missing.clone(),
+                mime_type: None,
+            })
+        else {
+            panic!("missing paths must surface a file read error");
+        };
+        assert_eq!(path, missing);
+        assert_eq!(kind, std::io::ErrorKind::NotFound);
+    }
+
+    #[test]
+    fn byte_documents_are_encoded_and_host_readers_must_be_read_first() {
+        assert_eq!(
+            prepare_document(OcrDocumentInput::Bytes {
+                bytes: b"abc".as_slice().into(),
+                file_name: Some("scan.pdf".into()),
+                mime_type: None,
+            })
+            .unwrap(),
+            document("data:application/pdf;base64,YWJj")
+        );
+        assert!(prepare_document(OcrDocumentInput::HostReader { mime_type: None }).is_err());
     }
 
     #[test]

litellm-rust/crates/core/src/ocr/error.rs

@@ -50,6 +50,12 @@ pub enum Error {
     Connect(String),
     #[error("routing error: {0}")]
     Routing(String),
+    #[error("Failed to read OCR file {}: {message}", path.display())]
+    FileRead {
+        path: std::path::PathBuf,
+        kind: std::io::ErrorKind,
+        message: String,
+    },
     /// The request is outside the surface this route covers in Rust. Hosts that
     /// keep a reference implementation treat this as "fall back", not "fail".
     #[error("unsupported by the rust path: {0}")]

litellm-rust/crates/core/src/ocr/lifecycle.rs

@@ -9,6 +9,7 @@ use super::hooks::{
     OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest,
     OcrPreCallRequest,
 };
+use super::types::{OcrDocumentInput, OcrFileContent};
 use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
 use crate::call_lifecycle::host::{
     HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase,
@@ -52,6 +53,7 @@ impl OcrAdmission {
 #[derive(Clone, Debug)]
 pub enum OcrHostOperation {
     ProjectRequest,
+    ReadDocument,
     Lifecycle(HostPhase),
     ConstructResponse(Arc<LiteLLMOcrResponse>),
     MapFailure(Error),
@@ -83,7 +85,8 @@ impl OcrHostOperation {
 }
 
 pub enum OcrHostResult {
-    Request(Result<(Box<LiteLLMOcrRequest>, bool), Error>),
+    Request(Result<(Box<LiteLLMOcrRequest<OcrDocumentInput>>, bool), Error>),
+    Document(Result<OcrFileContent, Error>),
     Lifecycle(Result<(), HostFailure<Error>>),
     AzureAdToken(Result<ResolvedCredential, AuthError>),
     PreCall(Result<OcrPreCallRequest, Error>),
@@ -313,7 +316,7 @@ struct PendingOperation {
 
 struct OcrExecution {
     client: Option<OcrClient>,
-    request: Option<LiteLLMOcrRequest>,
+    request: Option<LiteLLMOcrRequest<OcrDocumentInput>>,
     operations_tx: mpsc::UnboundedSender<PendingOperation>,
     operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
     pending_result: Option<oneshot::Sender<OcrHostResult>>,
@@ -397,12 +400,14 @@ impl OcrExecution {
                 },
             )));
         }
-        request.hooks = Arc::new(ProtocolHooks {
+        let hooks = Arc::new(ProtocolHooks {
             operations: self.operations_tx.clone(),
             intercepts_requests,
             terminal: self.terminal.clone(),
         });
+        request.hooks = hooks.clone();
         self.execution = Some(tokio::spawn(async move {
+            let request = prepare_request_document(request, &hooks).await?;
             perform_ocr_request(&client, request).await
         }));
     }
@@ -423,6 +428,39 @@ impl OcrExecution {
     }
 }
 
+async fn prepare_request_document(
+    request: LiteLLMOcrRequest<OcrDocumentInput>,
+    hooks: &ProtocolHooks,
+) -> Result<LiteLLMOcrRequest, Error> {
+    let request = match &request.document {
+        OcrDocumentInput::HostReader { mime_type } => {
+            let mime_type = mime_type.clone();
+            let content = match hooks.invoke(OcrHostOperation::ReadDocument).await? {
+                OcrHostResult::Document(result) => result?,
+                _ => {
+                    return Err(Error::InvalidRequest(
+                        "invalid OCR document read host result".into(),
+                    ));
+                }
+            };
+            request.with_document(OcrDocumentInput::Bytes {
+                bytes: content.bytes,
+                file_name: content.file_name,
+                mime_type,
+            })
+        }
+        _ => request,
+    };
+    if let OcrDocumentInput::Document(_) = &request.document {
+        return request.map_document(super::document::prepare_document);
+    }
+    tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document))
+        .await
+        .map_err(|error| {
+            Error::InvalidRequest(format!("OCR document preparation task failed: {error}"))
+        })?
+}
+
 impl Drop for OcrExecution {
     fn drop(&mut self) {
         if let Some(execution) = &self.execution {
@@ -567,6 +605,9 @@ impl OcrHost for NoopOcrHost {
                 OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
                     Error::InvalidRequest("OCR host has no request projection".into()),
                 )),
+                OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
+                    Error::InvalidRequest("OCR host has no document reader".into()),
+                )),
                 OcrHostOperation::Lifecycle(_)
                 | OcrHostOperation::ConstructResponse(_)
                 | OcrHostOperation::MapFailure(_)
@@ -602,6 +643,9 @@ impl OcrHost for OcrHookHost {
                 OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err(
                     Error::InvalidRequest("OCR hook host has no request projection".into()),
                 )),
+                OcrHostOperation::ReadDocument => OcrHostResult::Document(Err(
+                    Error::InvalidRequest("OCR hook host has no document reader".into()),
+                )),
                 OcrHostOperation::Success {
                     context,
                     response,

litellm-rust/crates/core/src/ocr/mod.rs

@@ -13,12 +13,15 @@ pub mod types;
 pub mod wire;
 
 pub use client::{OcrClient, ocr};
-pub use document::{encode_file_document, mime_type_for_name, upload_mime_type};
+pub use document::{encode_file_document, mime_type_for_name, read_path_document};
 pub use lifecycle::{
     NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline,
     OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult,
 };
-pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
+pub use types::{
+    LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrDocumentInput,
+    OcrFileContent,
+};
 
 #[cfg(test)]
 #[path = "../../tests/azure_ai_ocr.rs"]

litellm-rust/crates/core/src/ocr/types.rs

@@ -1,7 +1,10 @@
 use std::collections::BTreeMap;
+use std::convert::Infallible;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
 
+use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 
@@ -50,6 +53,35 @@ impl OcrDocument {
     }
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub enum OcrDocumentInput {
+    Document(OcrDocument),
+    Path {
+        path: PathBuf,
+        mime_type: Option<String>,
+    },
+    Bytes {
+        bytes: Bytes,
+        file_name: Option<String>,
+        mime_type: Option<String>,
+    },
+    HostReader {
+        mime_type: Option<String>,
+    },
+}
+
+impl From<OcrDocument> for OcrDocumentInput {
+    fn from(document: OcrDocument) -> Self {
+        Self::Document(document)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct OcrFileContent {
+    pub bytes: Bytes,
+    pub file_name: Option<String>,
+}
+
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "lowercase")]
 pub enum OcrResponseFormat {
@@ -89,9 +121,9 @@ impl Default for OcrConnection {
     }
 }
 
-pub struct LiteLLMOcrRequest {
+pub struct LiteLLMOcrRequest<D = OcrDocument> {
     pub model: String,
-    pub document: OcrDocument,
+    pub document: D,
     pub connection: OcrConnection,
     pub hooks: Arc<dyn OcrHooks>,
     pub litellm_call_id: Option<String>,
@@ -101,10 +133,10 @@ pub struct LiteLLMOcrRequest {
     pub(crate) adapter: OcrAdapterKind,
 }
 
-impl LiteLLMOcrRequest {
+impl<D> LiteLLMOcrRequest<D> {
     pub fn new(
         model: String,
-        document: OcrDocument,
+        document: D,
         custom_llm_provider: Option<&str>,
         optional_params: Map<String, Value>,
     ) -> Result<Self, Error> {
@@ -151,6 +183,36 @@ impl LiteLLMOcrRequest {
             ..self
         }
     }
+
+    pub fn map_document<T, E>(
+        self,
+        map: impl FnOnce(D) -> Result<T, E>,
+    ) -> Result<LiteLLMOcrRequest<T>, E> {
+        Ok(LiteLLMOcrRequest {
+            model: self.model,
+            document: map(self.document)?,
+            connection: self.connection,
+            hooks: self.hooks,
+            litellm_call_id: self.litellm_call_id,
+            optional_params: self.optional_params,
+            input_sources: self.input_sources,
+            azure_ad_token_provider: self.azure_ad_token_provider,
+            adapter: self.adapter,
+        })
+    }
+
+    pub fn with_document<T>(self, document: T) -> LiteLLMOcrRequest<T> {
+        let Ok(request) = self.map_document(|_| Ok::<T, Infallible>(document));
+        request
+    }
+}
+
+impl From<LiteLLMOcrRequest> for LiteLLMOcrRequest<OcrDocumentInput> {
+    fn from(request: LiteLLMOcrRequest) -> Self {
+        let Ok(request) = request
+            .map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document)));
+        request
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

litellm-rust/crates/core/src/ocr/wire.rs

@@ -68,9 +68,9 @@ pub struct DecodedOcrResponse<T> {
 
 #[derive(Deserialize)]
 #[serde(deny_unknown_fields)]
-pub struct OcrWireRequest {
+pub struct OcrWireRequest<D = Value> {
     pub model: String,
-    pub document: Value,
+    pub document: D,
     pub api_key: Option<String>,
     pub api_base: Option<String>,
     pub custom_llm_provider: Option<String>,
@@ -141,10 +141,34 @@ pub fn consumed_optional_params(
 }
 
 pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
+    let OcrWireRequest {
+        model,
+        document,
+        api_key,
+        api_base,
+        custom_llm_provider,
+        extra_headers,
+        optional_params,
+        input_sources,
+        timeout_seconds,
+    } = wire;
+    decode_request_input(OcrWireRequest {
+        model,
+        document: decode_document(document)?,
+        api_key,
+        api_base,
+        custom_llm_provider,
+        extra_headers,
+        optional_params,
+        input_sources,
+        timeout_seconds,
+    })
+}
+
+pub fn decode_request_input<D>(wire: OcrWireRequest<D>) -> Result<LiteLLMOcrRequest<D>, Error> {
     let api_key_source = source_for(&wire.input_sources, "api_key");
     let api_base_source = source_for(&wire.input_sources, "api_base");
     let extra_headers_source = source_for(&wire.input_sources, "extra_headers");
-    let document = decode_document(wire.document)?;
     let headers = wire
         .extra_headers
         .unwrap_or_default()
@@ -183,7 +207,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
         .unwrap_or(defaults.max_response_bytes);
     let request = LiteLLMOcrRequest::new(
         wire.model,
-        document,
+        wire.document,
         wire.custom_llm_provider.as_deref(),
         wire.optional_params
             .into_iter()
@@ -209,14 +233,14 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
     })
 }
 
-fn decode_document(value: Value) -> Result<OcrDocument, OcrRequestError> {
+pub fn decode_document(value: Value) -> Result<OcrDocument, Error> {
     let kind = value.get("type").and_then(Value::as_str);
     let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none()
         || matches!(kind, Some("image_url")) && value.get("image_url").is_none();
     if missing_url {
-        return Err(OcrRequestError::MissingDocumentUrl);
+        return Err(OcrRequestError::MissingDocumentUrl.into());
     }
-    decode_request_value(value, "document")
+    Ok(decode_request_value(value, "document")?)
 }
 
 fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
@@ -334,10 +358,7 @@ mod tests {
             serde_json::json!({"type": "document_url"}),
             serde_json::json!({"type": "image_url"}),
         ] {
-            assert_eq!(
-                decode_document(document),
-                Err(OcrRequestError::MissingDocumentUrl)
-            );
+            assert_eq!(decode_document(document), Err(Error::MissingDocumentUrl));
         }
     }
 }

litellm-rust/crates/core/tests/ocr.rs

@@ -348,13 +348,14 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() {
                     }
                     OcrHostOperation::ProjectRequest => {
                         result = Some(OcrHostResult::Request(Ok((
-                            Box::new(request.take().unwrap()),
+                            Box::new(request.take().unwrap().into()),
                             false,
                         ))))
                     }
                     OcrHostOperation::AcquireAzureAdToken => {
                         panic!("test request has no token provider")
                     }
+                    OcrHostOperation::ReadDocument => panic!("test request has no file reader"),
                     OcrHostOperation::PreCall(request) => {
                         phases.push("pre");
                         result = Some(OcrHostResult::PreCall(if failure_phase == "pre" {
@@ -405,7 +406,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure()
         match call.resume(result.take()).await {
             Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
                 result = Some(OcrHostResult::Request(Ok((
-                    Box::new(request.take().unwrap()),
+                    Box::new(request.take().unwrap().into()),
                     false,
                 ))));
             }
@@ -467,9 +468,10 @@ async fn direct_native_host_drives_the_same_state_machine() {
                     _ => panic!("unexpected OCR operation"),
                 });
                 result = Some(match operation {
-                    OcrHostOperation::ProjectRequest => {
-                        OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
-                    }
+                    OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok((
+                        Box::new(request.take().unwrap().into()),
+                        false,
+                    ))),
                     operation => host.invoke(operation).await,
                 });
             }
@@ -501,6 +503,137 @@ async fn direct_native_host_drives_the_same_state_machine() {
     ));
 }
 
+async fn drive_native_file_call(
+    request: super::LiteLLMOcrRequest<super::OcrDocumentInput>,
+    content: Result<super::OcrFileContent, crate::ocr::Error>,
+) -> (Result<super::LiteLLMOcrResponse, crate::ocr::Error>, usize) {
+    let NativeOutcome::Completed(mut call) =
+        OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all())
+    else {
+        panic!("supported call declined")
+    };
+    let mut request = Some(request);
+    let mut content = Some(content);
+    let mut result = None;
+    let mut reads = 0;
+    let outcome = loop {
+        match call.resume(result.take()).await {
+            Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => {
+                result = Some(OcrHostResult::Request(Ok((
+                    Box::new(request.take().unwrap()),
+                    false,
+                ))));
+            }
+            Ok(OcrCallStep::Host(OcrHostOperation::ReadDocument)) => {
+                reads += 1;
+                result = Some(OcrHostResult::Document(content.take().unwrap()));
+            }
+            Ok(OcrCallStep::Host(operation)) => result = Some(NoopOcrHost.invoke(operation).await),
+            Ok(OcrCallStep::Complete(response)) => break Ok(response),
+            Err(error) => break Err(error),
+        }
+    };
+    (outcome, reads)
+}
+
+#[tokio::test]
+async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "pages":[{"index":0,"markdown":"file"}]
+    }))])
+    .await;
+    let request = wire_request("mistral/model", &base, json!({})).with_document(
+        super::OcrDocumentInput::HostReader {
+            mime_type: Some("application/pdf".into()),
+        },
+    );
+    let (response, reads) = drive_native_file_call(
+        request,
+        Ok(super::OcrFileContent {
+            bytes: b"abc".as_slice().into(),
+            file_name: Some("scan.png".into()),
+        }),
+    )
+    .await;
+    server.await.unwrap();
+    assert_eq!(response.unwrap().pages[0]["markdown"], "file");
+    assert_eq!(reads, 1);
+    assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj"));
+}
+
+#[tokio::test]
+async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() {
+    let (base, seen, _server) = mock_server(vec![]).await;
+    let request = wire_request("mistral/model", &base, json!({}));
+    let failure = crate::ocr::Error::InvalidRequest("reader exploded".into());
+    let (response, reads) = drive_native_file_call(
+        request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
+        Err(failure.clone()),
+    )
+    .await;
+    assert_eq!(response.unwrap_err(), failure);
+    assert_eq!(reads, 1);
+
+    let request = wire_request("mistral/model", &base, json!({}));
+    let (response, _) = drive_native_file_call(
+        request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
+        Ok(super::OcrFileContent {
+            bytes: Default::default(),
+            file_name: None,
+        }),
+    )
+    .await;
+    assert!(matches!(
+        response.unwrap_err(),
+        crate::ocr::Error::InvalidRequest(_)
+    ));
+    assert!(seen.lock().unwrap().is_empty());
+}
+
+#[tokio::test]
+async fn path_documents_are_read_by_core_without_a_host_operation() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "pages":[{"index":0,"markdown":"path"}]
+    }))])
+    .await;
+    let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::<u64>()));
+    std::fs::create_dir_all(&dir).unwrap();
+    let path = dir.join("scan.png");
+    std::fs::write(&path, b"abc").unwrap();
+    let request = wire_request("mistral/model", &base, json!({})).with_document(
+        super::OcrDocumentInput::Path {
+            path: path.clone(),
+            mime_type: None,
+        },
+    );
+    let (response, reads) = drive_native_file_call(
+        request,
+        Err(crate::ocr::Error::InvalidRequest("unused".into())),
+    )
+    .await;
+    server.await.unwrap();
+    std::fs::remove_dir_all(&dir).unwrap();
+    assert_eq!(response.unwrap().pages[0]["markdown"], "path");
+    assert_eq!(reads, 0);
+    assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj"));
+
+    let (base, seen, _server) = mock_server(vec![]).await;
+    let request = wire_request("mistral/model", &base, json!({}));
+    let (response, _) = drive_native_file_call(
+        request.with_document(super::OcrDocumentInput::Path {
+            path: path.clone(),
+            mime_type: None,
+        }),
+        Err(crate::ocr::Error::InvalidRequest("unused".into())),
+    )
+    .await;
+    assert!(matches!(
+        response.unwrap_err(),
+        crate::ocr::Error::FileRead { path: failed, kind: std::io::ErrorKind::NotFound, .. } if failed == path
+    ));
+    assert!(seen.lock().unwrap().is_empty());
+}
+
 #[tokio::test]
 async fn public_finalization_failure_never_dispatches_success_or_replays_provider() {
     use crate::call_lifecycle::host::{HostFailure, HostPhase};
@@ -543,9 +676,10 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide
                     | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => {
                         panic!("finalization failure used provider/success dispatch")
                     }
-                    OcrHostOperation::ProjectRequest => {
-                        OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false)))
-                    }
+                    OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok((
+                        Box::new(request.take().unwrap().into()),
+                        false,
+                    ))),
                     operation => host.invoke(operation).await,
                 });
             }
@@ -582,7 +716,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption
             OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break,
             OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
                 result = Some(OcrHostResult::Request(Ok((
-                    Box::new(request.take().unwrap()),
+                    Box::new(request.take().unwrap().into()),
                     false,
                 ))))
             }
@@ -799,7 +933,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_
                     _ = entered.notified() => break,
                     step = call.resume(result.take()) => {
                         result = Some(match step.unwrap() {
-                            OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))),
+                            OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap().into()), false))),
                             OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await,
                             OcrCallStep::Complete(_) => panic!("pending provider completed"),
                         });

litellm-rust/crates/framer/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "litellm-framing"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[features]
+default = ["aws", "sse"]
+aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
+sse = ["dep:sse-stream"]
+
+[dependencies]
+aws-smithy-eventstream = { version = "=0.61.1", optional = true }
+aws-smithy-types = { version = "1.6.1", optional = true }
+bytes = "1"
+futures-util.workspace = true
+sse-stream = { version = "=0.2.6", optional = true }
+thiserror.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true
+tokio.workspace = true

litellm-rust/crates/framer/src/aws_event_stream.rs

@@ -0,0 +1,66 @@
+use bytes::{Buf, Bytes, BytesMut};
+use futures_util::{Stream, StreamExt};
+
+use aws_smithy_eventstream::frame::read_message_from;
+use aws_smithy_types::event_stream::Header;
+
+use crate::{Error, Framer};
+
+const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct AwsEventStreamFrame {
+    pub headers: Vec<Header>,
+    pub payload: Bytes,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct AwsEventStreamFramer;
+
+impl Framer for AwsEventStreamFramer {
+    type Frame = AwsEventStreamFrame;
+
+    fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
+    where
+        S: Stream<Item = Result<B, E>> + Send,
+        B: Buf + Send,
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        futures_util::stream::try_unfold(
+            (Box::pin(input), BytesMut::new()),
+            |(mut input, mut buffer)| async move {
+                loop {
+                    if buffer.len() >= 4 {
+                        let length = (&buffer[..4]).get_u32() as usize;
+                        if !(16..=MAX_FRAME_BYTES).contains(&length) {
+                            return Err(Error::InvalidLength(length));
+                        }
+                        if buffer.len() >= length {
+                            let raw = buffer.split_to(length).freeze();
+                            let message = read_message_from(raw)?;
+                            let frame = AwsEventStreamFrame {
+                                headers: message.headers().to_vec(),
+                                payload: message.payload().clone(),
+                            };
+                            return Ok(Some((frame, (input, buffer))));
+                        }
+                    }
+                    match input.next().await {
+                        Some(Ok(mut chunk)) => {
+                            while chunk.has_remaining() {
+                                let bytes = chunk.chunk();
+                                buffer.extend_from_slice(bytes);
+                                let length = bytes.len();
+                                chunk.advance(length);
+                            }
+                        }
+                        Some(Err(error)) => return Err(Error::Body(Box::new(error))),
+                        None if buffer.is_empty() => return Ok(None),
+                        None => return Err(Error::Truncated),
+                    }
+                }
+            },
+        )
+        .fuse()
+    }
+}

litellm-rust/crates/framer/src/error.rs

@@ -0,0 +1,17 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[cfg(feature = "sse")]
+    #[error("SSE framing failed: {0}")]
+    Sse(#[from] sse_stream::Error),
+    #[cfg(feature = "aws")]
+    #[error("AWS EventStream framing failed: {0}")]
+    Aws(#[from] aws_smithy_eventstream::error::Error),
+    #[error("body stream failed: {0}")]
+    Body(#[source] Box<dyn std::error::Error + Send + Sync>),
+    #[cfg(feature = "aws")]
+    #[error("invalid AWS EventStream frame length: {0}")]
+    InvalidLength(usize),
+    #[cfg(feature = "aws")]
+    #[error("truncated AWS EventStream frame")]
+    Truncated,
+}

litellm-rust/crates/framer/src/framer.rs

@@ -0,0 +1,13 @@
+use futures_util::Stream;
+
+use crate::Error;
+
+pub trait Framer: Send {
+    type Frame: Send;
+
+    fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
+    where
+        S: Stream<Item = Result<B, E>> + Send,
+        B: bytes::Buf + Send,
+        E: std::error::Error + Send + Sync + 'static;
+}

litellm-rust/crates/framer/src/lib.rs

@@ -0,0 +1,10 @@
+mod error;
+mod framer;
+
+pub use error::*;
+pub use framer::*;
+
+#[cfg(feature = "aws")]
+pub mod aws_event_stream;
+#[cfg(feature = "sse")]
+pub mod sse;

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

@@ -0,0 +1,43 @@
+use futures_util::{Stream, StreamExt};
+
+use crate::{Error, Framer};
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct SseFrame {
+    pub event: Option<String>,
+    pub data: Option<String>,
+    pub id: Option<String>,
+    pub retry: Option<u64>,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct SseFramer;
+
+impl Framer for SseFramer {
+    type Frame = SseFrame;
+
+    fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<SseFrame, Error>> + Send
+    where
+        S: Stream<Item = Result<B, E>> + Send,
+        B: bytes::Buf + Send,
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input));
+        futures_util::stream::try_unfold(frames, |mut frames| async move {
+            let Some(frame) = frames.next().await else {
+                return Ok(None);
+            };
+            let frame = frame?;
+            Ok(Some((
+                SseFrame {
+                    event: frame.event,
+                    data: frame.data,
+                    id: frame.id,
+                    retry: frame.retry,
+                },
+                frames,
+            )))
+        })
+        .fuse()
+    }
+}

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

@@ -0,0 +1,92 @@
+#![cfg(feature = "aws")]
+
+mod support;
+
+use std::io;
+
+use futures_util::TryStreamExt;
+use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
+use litellm_framing::{Error, Framer};
+use rstest::{fixture, rstest};
+
+use support::encode;
+
+async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result<Vec<AwsEventStreamFrame>, Error> {
+    AwsEventStreamFramer
+        .frame(futures_util::stream::iter(
+            bytes.chunks(chunk_size).map(Ok::<_, io::Error>),
+        ))
+        .try_collect()
+        .await
+}
+
+#[fixture]
+fn two_frames() -> Vec<u8> {
+    [encode(b"\xff\x00"), encode(b"second")].concat()
+}
+
+#[fixture]
+fn payload_frame() -> Vec<u8> {
+    encode(b"payload")
+}
+
+#[rstest]
+#[case(1)]
+#[case(3)]
+#[case(12)]
+#[case(usize::MAX)]
+#[tokio::test]
+async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads(
+    two_frames: Vec<u8>,
+    #[case] chunk_size: usize,
+) {
+    let chunk_size = chunk_size.min(two_frames.len());
+    let frames = collect_aws(&two_frames, chunk_size).await.unwrap();
+    assert_eq!(frames.len(), 2);
+    assert_eq!(frames[0].payload, &b"\xff\x00"[..]);
+    assert_eq!(frames[1].payload, "second");
+    assert_eq!(
+        frames[0].headers[0].value().as_string().unwrap().as_str(),
+        "payload"
+    );
+    assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7));
+}
+
+#[rstest]
+#[case(8)]
+#[case(0)]
+#[tokio::test]
+async fn rejects_corrupt_crcs(payload_frame: Vec<u8>, #[case] index: usize) {
+    let corrupt_index = if index == 0 {
+        payload_frame.len() - 1
+    } else {
+        index
+    };
+    let mut corrupt = payload_frame;
+    corrupt[corrupt_index] ^= 1;
+    assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_))));
+}
+
+#[rstest]
+#[case(0_u32)]
+#[case(15)]
+#[case(u32::MAX)]
+#[tokio::test]
+async fn rejects_invalid_lengths(#[case] length: u32) {
+    assert!(matches!(
+        collect_aws(&length.to_be_bytes(), 1).await,
+        Err(Error::InvalidLength(_))
+    ));
+}
+
+#[rstest]
+#[case(1)]
+#[case(3)]
+#[case(5)]
+#[tokio::test]
+async fn rejects_truncation(payload_frame: Vec<u8>, #[case] end: usize) {
+    assert!(matches!(
+        collect_aws(&payload_frame[..end], 1).await,
+        Err(Error::Truncated)
+    ));
+}

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

@@ -0,0 +1,29 @@
+#![cfg(all(feature = "aws", feature = "sse"))]
+
+mod support;
+
+use std::io;
+
+use futures_util::TryStreamExt;
+use litellm_framing::Framer;
+use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
+use litellm_framing::sse::SseFramer;
+
+use support::encode;
+
+#[tokio::test]
+async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() {
+    let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat();
+    let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter(
+        bytes.chunks(3).map(Ok::<_, io::Error>),
+    ));
+    let frames = SseFramer
+        .frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload))
+        .try_collect::<Vec<_>>()
+        .await
+        .unwrap();
+    assert_eq!(frames.len(), 1);
+    assert_eq!(frames[0].event.as_deref(), Some("delta"));
+    assert_eq!(frames[0].data.as_deref(), Some("hello"));
+    assert_eq!(frames[0].id.as_deref(), Some("7"));
+}

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

@@ -0,0 +1,67 @@
+#![cfg(feature = "sse")]
+
+use std::io;
+
+use futures_util::{StreamExt, TryStreamExt};
+use litellm_framing::sse::{SseFrame, SseFramer};
+use litellm_framing::{Error, Framer};
+use rstest::rstest;
+
+async fn collect_sse(chunks: &[&[u8]]) -> Result<Vec<SseFrame>, Error> {
+    SseFramer
+        .frame(futures_util::stream::iter(
+            chunks.iter().copied().map(Ok::<_, io::Error>),
+        ))
+        .try_collect()
+        .await
+}
+
+#[rstest]
+#[case(
+    &[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]],
+    vec![
+        SseFrame {
+            event: Some("delta".into()),
+            data: Some("€\nnext".into()),
+            id: Some("7".into()),
+            retry: Some(10),
+        },
+        SseFrame {
+            event: None,
+            data: Some("[DONE]".into()),
+            id: None,
+            retry: None,
+        },
+    ]
+)]
+#[tokio::test]
+async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel(
+    #[case] chunks: &[&[u8]],
+    #[case] expected: Vec<SseFrame>,
+) {
+    assert_eq!(collect_sse(chunks).await.unwrap(), expected);
+}
+
+#[tokio::test]
+async fn eof_does_not_dispatch_an_unterminated_frame() {
+    assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty());
+}
+
+#[rstest]
+#[case(io::ErrorKind::ConnectionReset)]
+#[case(io::ErrorKind::UnexpectedEof)]
+#[tokio::test]
+async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) {
+    let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([
+        Err(io::Error::new(kind, "reset")),
+        Ok(&b"data: later\n\n"[..]),
+    ])));
+    let error = frames.next().await.unwrap().unwrap_err();
+    assert!(matches!(
+        error,
+        Error::Sse(sse_stream::Error::Body(ref cause))
+            if cause.downcast_ref::<io::Error>().unwrap().kind() == kind
+    ));
+    assert!(frames.next().await.is_none());
+    assert!(frames.next().await.is_none());
+}

litellm-rust/crates/framer/tests/support/mod.rs

@@ -0,0 +1,15 @@
+use aws_smithy_eventstream::frame::write_message_to;
+use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
+use bytes::Bytes;
+
+pub fn encode(payload: &'static [u8]) -> Vec<u8> {
+    let message = Message::new(Bytes::from_static(payload))
+        .add_header(Header::new(
+            ":event-type",
+            HeaderValue::String("payload".into()),
+        ))
+        .add_header(Header::new("sequence", HeaderValue::Int32(7)));
+    let mut bytes = Vec::new();
+    write_message_to(&message, &mut bytes).unwrap();
+    bytes
+}

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

@@ -16,6 +16,7 @@ extension-module = ["pyo3/extension-module"]
 panic-test = []
 
 [dependencies]
+bytes.workspace = true
 futures-util.workspace = true
 litellm-core.workspace = true
 litellm-auth.workspace = true

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

@@ -599,6 +599,34 @@ mod tests {
 
     static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
 
+    fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
+        py.run(
+            pyo3::ffi::c_str!(
+                r#"
+import sys
+import types
+
+sys.modules.setdefault('litellm', types.ModuleType('litellm'))
+sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge'))
+"#
+            ),
+            None,
+            None,
+        )
+        .unwrap();
+        let source = std::ffi::CString::new(include_str!(
+            "../../../../../litellm/rust_bridge/lifecycle.py"
+        ))
+        .unwrap();
+        PyModule::from_code(
+            py,
+            &source,
+            pyo3::ffi::c_str!("lifecycle.py"),
+            pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"),
+        )
+        .unwrap()
+    }
+
     fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> {
         py.import("litellm.litellm_core_utils.logging_worker")?
             .setattr("GLOBAL_LOGGING_WORKER", worker)
@@ -773,17 +801,7 @@ mod tests {
             .unwrap_or_else(|error| error.into_inner());
         Python::initialize();
         Python::attach(|py| {
-            let source = std::ffi::CString::new(include_str!(
-                "../../../../../litellm/rust_bridge/lifecycle.py"
-            ))
-            .unwrap();
-            PyModule::from_code(
-                py,
-                &source,
-                pyo3::ffi::c_str!("lifecycle.py"),
-                pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"),
-            )
-            .unwrap();
+            install_lifecycle_module(py);
             let route = SyntheticRoute(
                 PythonCallState::new(
                     py,
@@ -819,17 +837,7 @@ mod tests {
         Python::initialize();
         Python::attach(|py| {
             py.import("asyncio").unwrap();
-            let source = std::ffi::CString::new(include_str!(
-                "../../../../../litellm/rust_bridge/lifecycle.py"
-            ))
-            .unwrap();
-            let module = PyModule::from_code(
-                py,
-                &source,
-                pyo3::ffi::c_str!("lifecycle.py"),
-                pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"),
-            )
-            .unwrap();
+            let module = install_lifecycle_module(py);
             let locals = PyDict::new(py);
             locals
                 .set_item("drive", module.getattr("drive").unwrap())

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

@@ -190,6 +190,7 @@ mod tests {
 
     #[test]
     fn required_shapes_preserve_nested_values_and_existing_errors() {
+        Python::initialize();
         let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]);
         assert_eq!(
             Value::Array(required_array("messages", nested.clone()).unwrap()),

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

@@ -1,97 +1,56 @@
-use std::io::Read;
 use std::path::PathBuf;
 
-use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError};
+use bytes::Bytes;
+use pyo3::exceptions::{PyTypeError, PyValueError};
+use pyo3::gc::{PyTraverseError, PyVisit};
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedBytes;
-#[cfg(test)]
-use pyo3::types::PyDict;
 use pyo3::types::{PyBytes, PyString};
 
-use litellm_core::constants::OCR_INLINE_MAX_BYTES;
-use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type};
-use litellm_python_interop::to_py_preserving_errors;
+use litellm_core::ocr::{OcrDocumentInput, OcrFileContent};
 
-enum FileBytes {
-    Python(PyBackedBytes),
-    Native(Vec<u8>),
+#[derive(Debug)]
+pub(super) struct PythonFileReader {
+    reader: Py<PyAny>,
+    name: Option<String>,
 }
 
-impl AsRef<[u8]> for FileBytes {
-    fn as_ref(&self) -> &[u8] {
-        match self {
-            Self::Python(bytes) => bytes,
-            Self::Native(bytes) => bytes,
-        }
+impl PythonFileReader {
+    pub(super) fn read(&self, py: Python<'_>) -> PyResult<OcrFileContent> {
+        let value = self.reader.bind(py).call0()?;
+        let bytes = if value.is_instance_of::<PyString>() {
+            Bytes::from(value.extract::<String>()?)
+        } else if value.is_instance_of::<PyBytes>() {
+            extract_bytes(&value)?
+        } else {
+            return Err(PyTypeError::new_err(format!(
+                "OCR file read must return bytes or str, got {}",
+                value.get_type(),
+            )));
+        };
+        Ok(OcrFileContent {
+            bytes,
+            file_name: self.name.clone(),
+        })
+    }
+
+    pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        visit.call(&self.reader)
     }
 }
 
-fn read_file_input(
-    py: Python<'_>,
-    file: &Bound<'_, PyAny>,
-) -> PyResult<(FileBytes, Option<String>)> {
-    if file.is_instance_of::<PyString>() {
-        return Err(PyValueError::new_err(
-            "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
-        ));
+fn extract_bytes(value: &Bound<'_, PyAny>) -> PyResult<Bytes> {
+    if value.is_exact_instance_of::<PyBytes>() {
+        return Ok(Bytes::from_owner(value.extract::<PyBackedBytes>()?));
     }
-    if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
-        let path: PathBuf = file.extract()?;
-        let name = path
-            .file_name()
-            .map(|value| value.to_string_lossy().into_owned());
-        let bytes = py
-            .detach(|| {
-                let mut bytes = Vec::new();
-                std::fs::File::open(&path)?
-                    .take(OCR_INLINE_MAX_BYTES as u64 + 1)
-                    .read_to_end(&mut bytes)?;
-                Ok::<_, std::io::Error>(bytes)
-            })
-            .map_err(|error| {
-                if error.kind() == std::io::ErrorKind::NotFound {
-                    PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
-                } else {
-                    error.into()
-                }
-            })?;
-        return Ok((FileBytes::Native(bytes), name));
-    }
-    if file.is_instance_of::<PyBytes>() {
-        return Ok((FileBytes::Python(file.extract()?), None));
-    }
-    let reader = file
-        .getattr_opt("read")?
-        .filter(|value| value.is_callable());
-    let Some(reader) = reader else {
-        return Err(PyValueError::new_err(format!(
-            "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.",
-            file.get_type(),
-        )));
-    };
-    let name = file
-        .getattr_opt("name")?
-        .filter(|value| !value.is_none())
-        .map(|value| value.extract::<String>())
-        .transpose()?;
-    let value = reader.call0()?;
-    let bytes = if value.is_instance_of::<PyString>() {
-        FileBytes::Native(value.extract::<String>()?.into_bytes())
-    } else if value.is_instance_of::<PyBytes>() {
-        FileBytes::Python(value.extract()?)
-    } else {
-        return Err(PyTypeError::new_err(format!(
-            "OCR file read must return bytes or str, got {}",
-            value.get_type(),
-        )));
-    };
-    Ok((bytes, name))
+    Ok(Bytes::copy_from_slice(
+        value.extract::<PyBackedBytes>()?.as_ref(),
+    ))
 }
 
 pub(super) struct FileDocumentInput {
-    bytes: FileBytes,
-    name: Option<String>,
-    mime_type: Option<String>,
+    pub input: OcrDocumentInput,
+    pub reader: Option<PythonFileReader>,
 }
 
 impl FromPyObject<'_, '_> for FileDocumentInput {
@@ -104,79 +63,79 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
             Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => None,
             Err(error) => return Err(error),
         };
+        let missing = || {
+            PyValueError::new_err(
+                "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
+            )
+        };
         let file = document.get_item("file").map_err(|error| {
             if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) {
-                PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes")
+                missing()
             } else {
                 error
             }
         })?;
         if file.is_none() {
+            return Err(missing());
+        }
+        if file.is_instance_of::<PyString>() {
             return Err(PyValueError::new_err(
-                "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes",
+                "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
             ));
         }
-        let (bytes, name) = read_file_input(py, &file)?;
+        if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
+            return Ok(Self {
+                input: OcrDocumentInput::Path {
+                    path: file.extract::<PathBuf>()?,
+                    mime_type,
+                },
+                reader: None,
+            });
+        }
+        if file.is_instance_of::<PyBytes>() {
+            return Ok(Self {
+                input: OcrDocumentInput::Bytes {
+                    bytes: extract_bytes(&file)?,
+                    file_name: None,
+                    mime_type,
+                },
+                reader: None,
+            });
+        }
+        let reader = file
+            .getattr_opt("read")?
+            .filter(|value| value.is_callable());
+        let Some(reader) = reader else {
+            return Err(PyValueError::new_err(format!(
+                "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.",
+                file.get_type(),
+            )));
+        };
+        let name = file
+            .getattr_opt("name")?
+            .filter(|value| !value.is_none())
+            .map(|value| value.extract::<String>())
+            .transpose()?;
         Ok(Self {
-            bytes,
-            name,
-            mime_type,
+            input: OcrDocumentInput::HostReader { mime_type },
+            reader: Some(PythonFileReader {
+                reader: reader.unbind(),
+                name,
+            }),
         })
     }
 }
 
-pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult<OcrDocument> {
-    py.detach(|| {
-        encode_file_document(
-            document.bytes.as_ref(),
-            document.name.as_deref(),
-            document.mime_type.as_deref(),
-        )
-    })
-    .map_err(|error| PyValueError::new_err(error.to_string()))
-}
-
-#[pyfunction]
-fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
-    to_py_preserving_errors(py, &file_document(py, document.extract()?)?)
-}
-
-#[pyfunction]
-fn _ocr_mime_type(file_name: &str) -> String {
-    mime_type_for_name(file_name).into()
-}
-
-#[pyfunction]
-#[pyo3(signature = (file_content, file_name=None, content_type=None))]
-fn _ocr_upload_document(
-    py: Python<'_>,
-    file_content: &Bound<'_, PyBytes>,
-    file_name: Option<&str>,
-    content_type: Option<&str>,
-) -> PyResult<Py<PyAny>> {
-    let bytes: PyBackedBytes = file_content.extract()?;
-    let document = py
-        .detach(|| {
-            encode_file_document(
-                &bytes,
-                None,
-                Some(upload_mime_type(file_name, content_type)),
-            )
-        })
-        .map_err(|error| PyValueError::new_err(error.to_string()))?;
-    to_py_preserving_errors(py, &document)
-}
-
-pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?;
-    module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?;
-    module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?;
-    module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
+    use pyo3::types::PyDict;
+
+    fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
+        let locals = PyDict::new(py);
+        py.run(source, Some(&locals), Some(&locals)).unwrap();
+        locals
+    }
 
     #[test]
     fn extraction_validates_required_file_and_optional_mime_type() {
@@ -196,69 +155,148 @@ mod tests {
                 let error = document.extract::<FileDocumentInput>().err().unwrap();
                 assert!(error.is_instance_of::<PyTypeError>(py));
             }
-            let document = py.eval(c"{'file': b'abc'}", None, None).unwrap();
+            let error = py
+                .eval(c"{'file': 'scan.pdf'}", None, None)
+                .unwrap()
+                .extract::<FileDocumentInput>()
+                .err()
+                .unwrap();
+            assert!(error.is_instance_of::<PyValueError>(py));
+            assert!(error.to_string().contains("bare str"));
+            let document = py
+                .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None)
+                .unwrap();
             let input: FileDocumentInput = document.extract().unwrap();
-            assert_eq!(input.bytes.as_ref(), b"abc");
-            assert_eq!(input.name, None);
-            assert_eq!(input.mime_type, None);
+            assert!(input.reader.is_none());
+            assert_eq!(
+                input.input,
+                OcrDocumentInput::Bytes {
+                    bytes: b"abc".as_slice().into(),
+                    file_name: None,
+                    mime_type: Some("image/png".into()),
+                }
+            );
         });
     }
 
     #[test]
-    fn extraction_validates_mime_type_before_consuming_file() {
+    fn paths_and_readers_are_projected_without_io() {
         Python::initialize();
         Python::attach(|py| {
-            let locals = PyDict::new(py);
-            py.run(
-                c"class Reader:
+            let locals = eval(
+                py,
+                c"from pathlib import Path
+class Reader:
+    name = 'scan.png'
     def __init__(self):
         self.reads = 0
     def read(self):
         self.reads += 1
         return b'abc'
 reader = Reader()
-document = {'file': reader, 'mime_type': 7}",
-                Some(&locals),
-                Some(&locals),
-            )
-            .unwrap();
+document = {'file': reader, 'mime_type': 7}
+reader_document = {'file': reader}
+path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_type': 'image/png'}",
+            );
             let document = locals.get_item("document").unwrap().unwrap();
             let error = document.extract::<FileDocumentInput>().err().unwrap();
             assert!(error.is_instance_of::<PyTypeError>(py));
-            let reads: usize = locals
-                .get_item("reader")
-                .unwrap()
-                .unwrap()
-                .getattr("reads")
-                .unwrap()
-                .extract()
-                .unwrap();
-            assert_eq!(reads, 0);
+
+            let document = locals.get_item("reader_document").unwrap().unwrap();
+            let input: FileDocumentInput = document.extract().unwrap();
+            assert_eq!(
+                input.input,
+                OcrDocumentInput::HostReader { mime_type: None }
+            );
+            let reads = || {
+                locals
+                    .get_item("reader")
+                    .unwrap()
+                    .unwrap()
+                    .getattr("reads")
+                    .unwrap()
+                    .extract::<usize>()
+                    .unwrap()
+            };
+            assert_eq!(reads(), 0);
+            let content = input.reader.unwrap().read(py).unwrap();
+            assert_eq!(reads(), 1);
+            assert_eq!(
+                content,
+                OcrFileContent {
+                    bytes: b"abc".as_slice().into(),
+                    file_name: Some("scan.png".into()),
+                }
+            );
+
+            let document = locals.get_item("path_document").unwrap().unwrap();
+            let input: FileDocumentInput = document.extract().unwrap();
+            assert!(input.reader.is_none());
+            assert_eq!(
+                input.input,
+                OcrDocumentInput::Path {
+                    path: PathBuf::from("/nonexistent/ocr-projection-test.pdf"),
+                    mime_type: Some("image/png".into()),
+                }
+            );
         });
     }
 
     #[test]
-    fn extraction_preserves_reader_key_error_identity() {
+    fn reader_results_are_normalized_and_exceptions_keep_their_identity() {
         Python::initialize();
         Python::attach(|py| {
-            let locals = PyDict::new(py);
-            py.run(
+            let locals = eval(
+                py,
                 c"failure = KeyError('reader failed')
-class Reader:
+class Raising:
     def read(self):
         raise failure
-document = {'file': Reader()}",
-                Some(&locals),
-                Some(&locals),
-            )
-            .unwrap();
-            let document = locals.get_item("document").unwrap().unwrap();
-            let error = document.extract::<FileDocumentInput>().err().unwrap();
+class Text:
+    def read(self):
+        return 'héllo'
+class Wrong:
+    def read(self):
+        return 7
+raising = {'file': Raising()}
+text = {'file': Text()}
+wrong = {'file': Wrong()}",
+            );
+            let reader = |name: &str| {
+                locals
+                    .get_item(name)
+                    .unwrap()
+                    .unwrap()
+                    .extract::<FileDocumentInput>()
+                    .unwrap()
+                    .reader
+                    .unwrap()
+            };
+            let error = reader("raising").read(py).unwrap_err();
             assert!(
                 error
                     .value(py)
                     .is(locals.get_item("failure").unwrap().unwrap())
             );
+            assert_eq!(
+                reader("text").read(py).unwrap().bytes.as_ref(),
+                "héllo".as_bytes()
+            );
+            let error = reader("wrong").read(py).unwrap_err();
+            assert!(error.is_instance_of::<PyTypeError>(py));
+            assert!(error.to_string().contains("bytes or str"));
         });
     }
+
+    #[test]
+    fn exact_python_bytes_transfer_without_copying_and_outlive_the_input() {
+        Python::initialize();
+        let (bytes, pointer) = Python::attach(|py| {
+            let value = PyBytes::new(py, b"document bytes");
+            let pointer = value.as_bytes().as_ptr() as usize;
+            (extract_bytes(value.as_any()).unwrap(), pointer)
+        });
+        assert_eq!(bytes.as_ptr() as usize, pointer);
+        assert_eq!(bytes.as_ref(), b"document bytes");
+    }
 }

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

@@ -1,4 +1,5 @@
 use litellm_core::ocr::Error;
+use pyo3::exceptions::{PyFileNotFoundError, PyOSError};
 use pyo3::prelude::*;
 
 use crate::errors::{RustUpstreamError, core_error_to_pyerr};
@@ -7,6 +8,12 @@ pub(super) fn to_pyerr(error: Error) -> PyErr {
     let status = error.http_status_code();
     let mapped = match error {
         Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
+        Error::FileRead {
+            path,
+            kind: std::io::ErrorKind::NotFound,
+            ..
+        } => PyFileNotFoundError::new_err(format!("File not found: {}", path.display())),
+        Error::FileRead { message, .. } => PyOSError::new_err(message),
         other => core_error_to_pyerr(other.into()),
     };
     attach_status(mapped, status)

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

@@ -66,13 +66,27 @@ impl PythonOcrHost {
                 retained_fields.set_item(name, value)?;
             }
         }
-        retained_fields.set_item("document", &self.projected()?.fields.document)?;
         let projected = self.projected_mut()?;
+        let document = match &projected.fields.document {
+            Some(document) => document.clone_ref(py),
+            None => to_py(py, &request.document)?,
+        };
+        retained_fields.set_item("document", &document)?;
+        projected.fields.document = Some(document);
         projected.retained_fields = Some(retained_fields.unbind());
         projected.pre_call = Some((&request).into());
         Ok(request)
     }
 
+    fn read_document(&self, py: Python<'_>) -> PyResult<litellm_core::ocr::OcrFileContent> {
+        self.projected()?
+            .fields
+            .reader
+            .as_ref()
+            .ok_or_else(missing_state)?
+            .read(py)
+    }
+
     fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
         let provider = self
             .projected()?
@@ -193,7 +207,7 @@ impl PythonRoute for PythonOcrHost {
                 let OcrHostData::Unprojected { request } = &self.data else {
                     return Err(missing_state());
                 };
-                let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?;
+                let projected = project_request(request.bind(py), self.state.kwargs.bind(py))?;
                 let has_token_provider = projected.fields.azure_ad_token_provider.is_some();
                 let request = projected.request;
                 self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost {
@@ -205,6 +219,7 @@ impl PythonRoute for PythonOcrHost {
                 }));
                 OcrHostResult::Request(Ok((Box::new(request), has_token_provider)))
             }
+            OcrHostOperation::ReadDocument => OcrHostResult::Document(Ok(self.read_document(py)?)),
             OcrHostOperation::AcquireAzureAdToken => {
                 OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?))
             }
@@ -258,6 +273,9 @@ impl PythonRoute for PythonOcrHost {
             OcrHostData::Projected(projected) => {
                 visit.call(&projected.fields.boundary_request)?;
                 visit.call(&projected.fields.document)?;
+                if let Some(reader) = &projected.fields.reader {
+                    reader.traverse(visit)?;
+                }
                 visit.call(&projected.fields.api_key)?;
                 if let Some(provider) = &projected.fields.azure_ad_token_provider {
                     provider.traverse(visit)?;

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

@@ -9,6 +9,5 @@ use pyo3::prelude::*;
 
 pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
     value::register(module)?;
-    document::register(module)?;
     lifecycle::register(module)
 }

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

@@ -1,14 +1,15 @@
 use std::sync::Arc;
 
-use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request};
-use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall};
-use litellm_python_interop::{
-    from_py_preserving_errors as from_py, to_py_preserving_errors as to_py,
+use litellm_core::ocr::wire::{
+    OcrWireRequest, consumed_optional_params, decode_document, decode_request_input,
 };
+use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall, OcrDocumentInput};
+use litellm_python_interop::from_py_preserving_errors as from_py;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use serde_json::{Map, Value};
 
+use super::document::{FileDocumentInput, PythonFileReader};
 use super::errors::to_pyerr as ocr_error_to_pyerr;
 use super::lifecycle::BridgeOcrHooks;
 use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider};
@@ -17,7 +18,8 @@ use crate::marshal::{project_optional_fields, python_timeout_seconds, request_in
 
 pub(super) struct ProjectedOcrFields {
     pub boundary_request: Py<PyAny>,
-    pub document: Py<PyAny>,
+    pub document: Option<Py<PyAny>>,
+    pub reader: Option<PythonFileReader>,
     pub api_key: Py<PyAny>,
     pub azure_ad_token_provider: Option<PythonTokenProvider>,
     pub provider: &'static str,
@@ -25,7 +27,7 @@ pub(super) struct ProjectedOcrFields {
 }
 
 pub(super) struct ProjectedOcrCall {
-    pub request: LiteLLMOcrRequest,
+    pub request: LiteLLMOcrRequest<OcrDocumentInput>,
     pub fields: ProjectedOcrFields,
 }
 
@@ -80,12 +82,12 @@ impl<'py> OcrArguments<'_, 'py> {
 }
 
 enum ProjectedDocument {
-    File { wire: Value, retained: Py<PyAny> },
+    File(FileDocumentInput),
     Other { wire: Value, retained: Py<PyAny> },
 }
 
 impl ProjectedDocument {
-    fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult<Self> {
+    fn project(document: &Bound<'_, PyAny>) -> PyResult<Self> {
         let kind: String = document.get_item("type")?.extract()?;
         if kind != "file" {
             return Ok(Self::Other {
@@ -93,25 +95,28 @@ impl ProjectedDocument {
                 retained: document.clone().unbind(),
             });
         }
-        let input = document.extract()?;
-        let encoded = super::document::file_document(py, input)?;
-        let wire = serde_json::to_value(encoded)
-            .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?;
-        Ok(Self::File {
-            retained: to_py(py, &wire)?,
-            wire,
-        })
+        Ok(Self::File(document.extract()?))
     }
 
-    fn into_parts(self) -> (Value, Py<PyAny>) {
+    fn into_parts(
+        self,
+    ) -> PyResult<(
+        OcrDocumentInput,
+        Option<Py<PyAny>>,
+        Option<PythonFileReader>,
+    )> {
         match self {
-            Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained),
+            Self::File(FileDocumentInput { input, reader }) => Ok((input, None, reader)),
+            Self::Other { wire, retained } => Ok((
+                decode_document(wire).map_err(ocr_error_to_pyerr)?.into(),
+                Some(retained),
+                None,
+            )),
         }
     }
 }
 
 pub(super) fn project_request(
-    py: Python<'_>,
     request: &Bound<'_, PyAny>,
     kwargs: &Bound<'_, PyDict>,
 ) -> PyResult<ProjectedOcrCall> {
@@ -119,8 +124,7 @@ pub(super) fn project_request(
     let arguments = OcrArguments { request, kwargs };
     let model = arguments.model()?;
     let custom_llm_provider = arguments.custom_llm_provider()?;
-    let (wire_document, retained_document) =
-        ProjectedDocument::project(py, &arguments.document()?)?.into_parts();
+    let document = ProjectedDocument::project(&arguments.document()?)?;
     let api_key = arguments.api_key()?;
     let specs = consumed_optional_params(&model, custom_llm_provider.as_deref())
         .map_err(ocr_error_to_pyerr)?;
@@ -136,9 +140,10 @@ pub(super) fn project_request(
     let azure_ad_token_provider = kwargs
         .get_item("azure_ad_token_provider")?
         .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER));
+    let (document, retained_document, reader) = document.into_parts()?;
     let wire = OcrWireRequest {
         model,
-        document: wire_document,
+        document,
         api_key: api_key.extract()?,
         api_base: arguments.api_base()?,
         custom_llm_provider,
@@ -147,13 +152,14 @@ pub(super) fn project_request(
         input_sources,
         timeout_seconds: arguments.timeout_seconds()?,
     };
-    let request = decode_request(wire).map_err(ocr_error_to_pyerr)?;
+    let request = decode_request_input(wire).map_err(ocr_error_to_pyerr)?;
     let provider = request.provider_name();
     Ok(ProjectedOcrCall {
         request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None),
         fields: ProjectedOcrFields {
             boundary_request,
             document: retained_document,
+            reader,
             api_key: api_key.unbind(),
             azure_ad_token_provider,
             provider,
@@ -197,10 +203,21 @@ mod tests {
     }
 
     fn project_document(
-        py: Python<'_>,
         document: &Bound<'_, PyAny>,
-    ) -> PyResult<(Value, Py<PyAny>)> {
-        ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts)
+    ) -> PyResult<(
+        OcrDocumentInput,
+        Option<Py<PyAny>>,
+        Option<PythonFileReader>,
+    )> {
+        ProjectedDocument::project(document)?.into_parts()
+    }
+
+    fn url_document(url: &str) -> OcrDocumentInput {
+        litellm_core::ocr::OcrDocument::DocumentUrl {
+            document_url: url.into(),
+            extra_fields: Map::new(),
+        }
+        .into()
     }
 
     fn stub_timeout_conversion(py: Python<'_>) {
@@ -374,7 +391,7 @@ kwargs = {}
     }
 
     #[test]
-    fn document_reader_mutations_are_visible_to_later_field_reads() {
+    fn document_readers_are_not_consumed_during_projection() {
         Python::initialize();
         Python::attach(|py| {
             stub_timeout_conversion(py);
@@ -406,7 +423,12 @@ kwargs = {}
                 .unwrap();
             let arguments = arguments(&request, &kwargs);
             let document = arguments.document().unwrap();
-            project_document(py, &document).unwrap();
+            let (input, retained, reader) = project_document(&document).unwrap();
+            assert_eq!(input, OcrDocumentInput::HostReader { mime_type: None });
+            assert!(retained.is_none());
+            assert_eq!(arguments.api_base().unwrap().as_deref(), Some("original"));
+            assert_eq!(arguments.timeout_seconds().unwrap(), Some(1.0));
+            reader.unwrap().read(py).unwrap();
             assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated"));
             assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0));
         });
@@ -444,7 +466,7 @@ kwargs = {'api_key': key}
     }
 
     #[test]
-    fn file_documents_are_encoded_and_other_documents_keep_the_python_object() {
+    fn file_documents_become_typed_inputs_and_other_documents_keep_the_python_object() {
         Python::initialize();
         Python::attach(|py| {
             let file = py
@@ -454,13 +476,17 @@ kwargs = {'api_key': key}
                     None,
                 )
                 .unwrap();
+            let (input, retained, reader) = project_document(&file).unwrap();
             assert_eq!(
-                project_document(py, &file).unwrap().0,
-                serde_json::json!({
-                    "type": "document_url",
-                    "document_url": "data:application/pdf;base64,JVBERi0xLjQ=",
-                })
+                input,
+                OcrDocumentInput::Bytes {
+                    bytes: b"%PDF-1.4".as_slice().into(),
+                    file_name: None,
+                    mime_type: Some("application/pdf".into()),
+                }
             );
+            assert!(retained.is_none());
+            assert!(reader.is_none());
 
             let original = py
                 .eval(
@@ -469,44 +495,21 @@ kwargs = {'api_key': key}
                     None,
                 )
                 .unwrap();
-            let (wire, retained) = project_document(py, &original).unwrap();
-            assert_eq!(
-                wire,
-                serde_json::json!({
-                    "type": "document_url",
-                    "document_url": "https://example.com/a.pdf",
-                })
-            );
-            assert!(retained.bind(py).is(&original));
+            let (input, retained, _) = project_document(&original).unwrap();
+            assert_eq!(input, url_document("https://example.com/a.pdf"));
+            assert!(retained.unwrap().bind(py).is(&original));
         });
     }
 
     #[test]
-    fn unknown_document_types_reach_existing_downstream_validation() {
+    fn unknown_document_types_reach_existing_core_validation() {
         Python::initialize();
         Python::attach(|py| {
             let document = py
                 .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None)
                 .unwrap();
-            let wire_document = project_document(py, &document).unwrap().0;
-            assert_eq!(
-                wire_document,
-                serde_json::json!({"type": "mystery", "mystery": "x"})
-            );
-            let error = match decode_request(OcrWireRequest {
-                model: "mistral/mistral-ocr-latest".into(),
-                document: wire_document,
-                api_key: None,
-                api_base: None,
-                custom_llm_provider: None,
-                extra_headers: None,
-                optional_params: Map::new(),
-                input_sources: Default::default(),
-                timeout_seconds: None,
-            }) {
-                Ok(_) => panic!("unknown discriminators belong to core validation"),
-                Err(error) => error,
-            };
+            let error = project_document(&document).unwrap_err();
+            assert!(error.is_instance_of::<PyValueError>(py));
             assert!(error.to_string().contains("document"));
         });
     }
@@ -517,14 +520,14 @@ kwargs = {'api_key': key}
         Python::attach(|py| {
             let missing = py.eval(c"{}", None, None).unwrap();
             assert!(
-                project_document(py, &missing)
+                project_document(&missing)
                     .unwrap_err()
                     .is_instance_of::<PyKeyError>(py)
             );
 
             let non_string = py.eval(c"{'type': 1}", None, None).unwrap();
             assert!(
-                project_document(py, &non_string)
+                project_document(&non_string)
                     .unwrap_err()
                     .is_instance_of::<PyTypeError>(py)
             );
@@ -540,7 +543,7 @@ document = Document()
 ",
             );
             let error =
-                project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err();
+                project_document(&locals.get_item("document").unwrap().unwrap()).unwrap_err();
             assert!(
                 error
                     .value(py)
@@ -569,9 +572,9 @@ document = Document()
 ",
             );
             let document = locals.get_item("document").unwrap().unwrap();
-            let (wire, retained) = project_document(py, &document).unwrap();
-            assert_eq!(wire["type"], "document_url");
-            assert!(!retained.bind(py).is(&document));
+            let (input, retained, _) = project_document(&document).unwrap();
+            assert!(matches!(input, OcrDocumentInput::Bytes { .. }));
+            assert!(retained.is_none());
             let reads: Vec<String> = document.getattr("reads").unwrap().extract().unwrap();
             assert_eq!(reads, ["type", "mime_type", "file"]);
         });

litellm-rust/crates/token-counter/src/tiktoken.rs

@@ -195,14 +195,21 @@ mod tests {
     }
 
     #[test]
-    fn long_repeated_runs_stay_cheap() {
+    fn long_repeated_runs_cost_close_to_linear() {
         let ranks = ranks();
         let mut scratch = MergeScratch::default();
-        let piece = vec![b' '; 1 << 20];
-        let started = std::time::Instant::now();
-        let count = ranks.count_piece(&piece, &mut scratch);
-        assert!(count > 0);
-        assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed());
+        let mut time = |len: usize| {
+            let piece = vec![b' '; len];
+            let started = std::time::Instant::now();
+            assert!(ranks.count_piece(&piece, &mut scratch) > 0);
+            started.elapsed()
+        };
+        let small = (0..3).map(|_| time(1 << 14)).min().unwrap();
+        let large = time(1 << 18);
+        assert!(
+            large < small * 64,
+            "{small:?} for 2^14 bytes, {large:?} for 2^18"
+        );
     }
 
     #[test]

litellm/__init__.py

@@ -244,6 +244,7 @@ telemetry = True
 max_tokens: int = DEFAULT_MAX_TOKENS  # OpenAI Defaults
 drop_params = drop_params_env_flag(os.environ, verbose_logger)
 modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
+bedrock_neutralize_orphaned_tool_blocks: bool = True
 use_chat_completions_url_for_anthropic_messages: bool = bool(
     os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
 )  # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
@@ -1818,6 +1819,9 @@ if TYPE_CHECKING:
     from .llms.azure.responses.o_series_transformation import (
         AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
     )
+    from .llms.azure_ai.responses.transformation import (
+        AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig,
+    )
     from .llms.xai.responses.transformation import (
         XAIResponsesAPIConfig as XAIResponsesAPIConfig,
     )

litellm/_lazy_imports_registry.py

@@ -234,6 +234,7 @@ LLM_CONFIG_NAMES: Final = (
     "OpenAIResponsesAPIConfig",
     "AzureOpenAIResponsesAPIConfig",
     "AzureOpenAIOSeriesResponsesAPIConfig",
+    "AzureAIResponsesAPIConfig",
     "XAIResponsesAPIConfig",
     "LiteLLMProxyResponsesAPIConfig",
     "HostedVLLMResponsesAPIConfig",
@@ -946,6 +947,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
         ".llms.azure.responses.o_series_transformation",
         "AzureOpenAIOSeriesResponsesAPIConfig",
     ),
+    "AzureAIResponsesAPIConfig": (
+        ".llms.azure_ai.responses.transformation",
+        "AzureAIResponsesAPIConfig",
+    ),
     "XAIResponsesAPIConfig": (
         ".llms.xai.responses.transformation",
         "XAIResponsesAPIConfig",

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -502,7 +502,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             elif key == "response_format":
                 text_format = self._transform_response_format_to_text_format(value)
                 if text_format:
-                    responses_api_request["text"] = text_format
+                    responses_api_request["text"] = self._merge_text(responses_api_request, text_format)
+            elif key == "verbosity":
+                responses_api_request["text"] = self._merge_text(
+                    responses_api_request,
+                    MappingProxyType({"verbosity": value}),  # pyright: ignore[reportUnknownArgumentType]  # untyped value
+                )
             elif key == "tool_choice":
                 responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value)
             elif key == "stream_options":
@@ -518,6 +523,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             elif key == "web_search_options":
                 self._add_web_search_tool(responses_api_request, value)
 
+    @staticmethod
+    def _merge_text(
+        responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object]
+    ) -> "ResponseText":
+        existing: Final = cast(  # cast-ok: text field is a ResponseText | dict[str, Any] | None union
+            "dict[str, object]",
+            dict(responses_api_request).get("text") or {},  # mutable-ok: one-shot merge seed
+        )
+        return cast(  # cast-ok: merged mapping is a valid ResponseText shape
+            "ResponseText",
+            {**existing, **update},  # mutable-ok: one-shot merged payload
+        )
+
     def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
         """Build sanitized litellm_params with merged metadata."""
         responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())

litellm/constants.py

@@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
         "router_general_settings",
         "ignore_invalid_deployments",
         "fallback_access_check",
+        "fallback_budget_check",
         "auto_router_capability_limit",
     }
 )
@@ -53,6 +54,7 @@ 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
 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
+S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
 MAX_FILE_LIST_LIMIT: Final = 10000
 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
 DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
@@ -1499,6 +1501,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
 CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
 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"
 SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
 SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"

litellm/experimental_mcp_client/client.py

@@ -346,13 +346,17 @@ class MCPClient:
             self.update_auth_value(auth_value)
 
     async def discovery_auth_fingerprint(self) -> str:
+        return self._hash_discovery_auth(await self.prepare_request_auth())
+
+    async def prepare_request_auth(self) -> httpx.Request:
+        """Preview the authenticated request without sending it, closing the auth flow afterwards."""
         request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
         if self._resolved_auth is None:
-            return self._hash_discovery_auth(request)
+            return request
         flow: Final = self._resolved_auth.async_auth_flow(request)
         try:
             authenticated: Final = await flow.__anext__()
-            return self._hash_discovery_auth(authenticated)
+            return authenticated
         finally:
             await flow.aclose()
 

litellm/integrations/callback_configs.json

@@ -446,6 +446,12 @@
         "ui_name": "S3 Path Prefix",
         "description": "Path prefix within the bucket for organizing logs",
         "required": false
+      },
+      "s3_log_prompts_only": {
+        "type": "boolean",
+        "ui_name": "Log Prompts Only",
+        "description": "Log request messages to S3 but drop the model response from each logged object",
+        "required": false
       }
     },
     "description": "S3 Bucket (AWS) Logging Integration"

litellm/integrations/custom_logger.py

@@ -118,6 +118,7 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
 
         alias_map: Final = {
             "langfuse_otel": "langfuse",
+            "s3_v2": "s3",
         }
         lookup_name: Final = alias_map.get(normalized_name, normalized_name)
 

litellm/integrations/opentelemetry.py

@@ -5,6 +5,7 @@ from collections.abc import Callable, Iterable, Mapping
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from datetime import datetime
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
 
 import litellm
@@ -20,7 +21,10 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
     OTELSemconvCategory,
     parse_semconv_opt_in,
 )
+from litellm.integrations.otel.mappers.utils import drop_none
+from litellm.integrations.otel.model.baggage import promoted_metadata
 from litellm.integrations.otel.model.db_endpoint import db_span_attributes
+from litellm.integrations.otel.model.metadata import flatten_metadata
 from litellm.integrations.otel.model.semconv import Metric
 from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@@ -205,6 +209,20 @@ def _resolve_metric_attribute_filter(
     )
 
 
+def _provider_label(custom_llm_provider: object) -> str | None:
+    """The provider label for one call's metrics and events, or None when the
+    call carries no provider.
+
+    Every attribute set drops None before export, so the label is simply absent
+    in that case: the OTLP encoder rejects a None attribute value outright, and a
+    placeholder would mint a permanent metric series that no operator can act
+    on. Mirrors the v2 integration's ``_provider_attributes``.
+    """
+    if not isinstance(custom_llm_provider, str) or not custom_llm_provider:
+        return None
+    return custom_llm_provider
+
+
 def _normalize_team_metadata_keys(value: str | Iterable[object] | None) -> list[str]:
     """Coerce a team-metadata allowlist from a list or comma-separated string.
 
@@ -288,6 +306,7 @@ class OpenTelemetryConfig:
     # under ``litellm.team.metadata``. Empty by default so none of a team's
     # metadata leaves the process until explicitly allowlisted.
     baggage_team_metadata_keys: list[str] = field(default_factory=list)
+    baggage_metadata_keys: list[str] = field(default_factory=list)
     # Prometheus-style include/exclude control over which attributes are stamped
     # on emitted metrics, to cap metric cardinality.
     attributes: OTELMetricAttributeFilter | None = None
@@ -314,6 +333,9 @@ class OpenTelemetryConfig:
         self.baggage_team_metadata_keys = _normalize_team_metadata_keys(
             self.baggage_team_metadata_keys
         ) or _normalize_team_metadata_keys(os.getenv("LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS"))
+        self.baggage_metadata_keys = _normalize_team_metadata_keys(
+            self.baggage_metadata_keys
+        ) or _normalize_team_metadata_keys(os.getenv("LITELLM_OTEL_BAGGAGE_METADATA_KEYS"))
 
     @classmethod
     def from_env(cls):
@@ -366,11 +388,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         **kwargs,
     ):
         team_metadata_keys_override: Final = kwargs.pop("baggage_team_metadata_keys", None)
+        metadata_keys_override: Final = kwargs.pop("baggage_metadata_keys", None)
         metric_attributes_override: Final = kwargs.pop("attributes", None)
         if config is None:
             config = OpenTelemetryConfig.from_env()
         if team_metadata_keys_override is not None:
             config.baggage_team_metadata_keys = _normalize_team_metadata_keys(team_metadata_keys_override)
+        if metadata_keys_override is not None:
+            config.baggage_metadata_keys = _normalize_team_metadata_keys(metadata_keys_override)
         if metric_attributes_override is not None:
             config.attributes = _build_metric_attribute_filter(metric_attributes_override)
 
@@ -1542,6 +1567,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         if team_metadata:
             self.safe_set_attribute(span=span, key=TEAM_METADATA_ATTRIBUTE, value=team_metadata)
 
+        if self.config.baggage_metadata_keys:
+            flat_metadata: Final = MappingProxyType(dict(flatten_metadata(metadata)))
+            for key, value in promoted_metadata(flat_metadata, tuple(self.config.baggage_metadata_keys)).items():
+                self.safe_set_attribute(span=span, key=key, value=value)
+
         model_group: Final = standard_logging_payload.get("model_group")
         if model_group:
             self.safe_set_attribute(span=span, key=MODEL_GROUP_ATTRIBUTE, value=model_group)
@@ -1601,19 +1631,22 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         ) = _resolve_metric_attribute_filter(attributes)
         self._metric_attr_filter_resolved = True
 
-    def _filter_metric_attributes(self, attrs: dict[str, str]) -> dict[str, str]:
+    def _filter_metric_attributes(self, attrs: Mapping[str, str | None]) -> dict[str, str]:
         if not self._metric_attr_filter_resolved:
             self._ensure_metric_attribute_filter()
+        return {k: v for k, v in attrs.items() if v is not None and self._metric_attribute_allowed(k)}
+
+    def _metric_attribute_allowed(self, key: str) -> bool:
         if self._metric_attr_include is not None:
-            return {k: v for k, v in attrs.items() if k in self._metric_attr_include}
+            return key in self._metric_attr_include
         if self._metric_attr_exclude is not None:
-            return {k: v for k, v in attrs.items() if k not in self._metric_attr_exclude}
-        return attrs
+            return key not in self._metric_attr_exclude
+        return True
 
     def _record_metrics(self, kwargs, response_obj, start_time, end_time):
         duration_s: Final = (end_time - start_time).total_seconds()
         params: Final = kwargs.get("litellm_params") or {}
-        provider: Final = params.get("custom_llm_provider", "Unknown")
+        provider: Final = _provider_label(params.get("custom_llm_provider"))
 
         common_attrs = {
             "gen_ai.operation.name": (
@@ -1857,7 +1890,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         otel_logger: Final = self._logger_provider.get_logger(LITELLM_LOGGER_NAME)
 
         parent_ctx: Final = span.get_span_context()
-        provider: Final = (kwargs.get("litellm_params") or {}).get("custom_llm_provider", "Unknown")
+        provider: Final = _provider_label((kwargs.get("litellm_params") or {}).get("custom_llm_provider"))
 
         if self._gen_ai_semconv_latest_experimental:
             self._emit_inference_details_event(
@@ -1894,7 +1927,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
                 severity_number=SeverityNumber.INFO,
                 severity_text="INFO",
                 body=body,
-                attributes=attrs,
+                attributes=drop_none(attrs),
             )
             otel_logger.emit(log_record)
 
@@ -1926,7 +1959,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
                 severity_number=SeverityNumber.INFO,
                 severity_text="INFO",
                 body=body,
-                attributes=attrs,
+                attributes=drop_none(attrs),
             )
             otel_logger.emit(log_record)
 
@@ -2932,16 +2965,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         )
 
         propagator: Final = TraceContextTextMapPropagator()
-        carrier: Final = {"traceparent": _traceparent}
+        carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None}
         _parent_context: Final = propagator.extract(carrier=carrier)
 
         return _parent_context
 
     def _get_span_context(self, kwargs, default_span: Span | None = None):
         from opentelemetry import context, trace
-        from opentelemetry.trace.propagation.tracecontext import (
-            TraceContextTextMapPropagator,
-        )
 
         litellm_params: Final = kwargs.get("litellm_params", {}) or {}
         proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {}
@@ -2965,11 +2995,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
         # Priority 2: HTTP traceparent header
         if traceparent is not None:
             verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation")
-            carrier: Final = {"traceparent": traceparent}
-            return (
-                TraceContextTextMapPropagator().extract(carrier=carrier),
-                None,
-            )
+            return self.get_traceparent_from_header(headers=headers), None
 
         # Priority 3: Active span from global context (auto-detection)
         try:

litellm/integrations/opentelemetry_utils/gen_ai_semconv.py

@@ -33,6 +33,7 @@ from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Final
 
+from litellm.integrations.otel.mappers.utils import drop_none
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
 if TYPE_CHECKING:
@@ -195,13 +196,16 @@ class OTELGenAISemconvMixin:
             if value:
                 self.safe_set_attribute(span=span, key=semconv_key, value=value)
 
-    def _build_inference_details_attrs(self, kwargs: dict, response_obj: dict, provider: str) -> dict[str, str]:
+    def _build_inference_details_attrs(
+        self, kwargs: dict, response_obj: dict, provider: str | None
+    ) -> dict[str, str | None]:
         """Build the attribute payload for the inference-details event.
 
-        Always includes provider/operation; input/output messages are added
+        Always includes operation and provider (None when the call carries none,
+        dropped before the event is emitted); input/output messages are added
         only when content capture is enabled and non-empty. Mixin-internal.
         """
-        attrs: Final[dict[str, str]] = {
+        attrs: Final[dict[str, str | None]] = {
             "event_name": _INFERENCE_DETAILS_EVENT_NAME,
             "gen_ai.provider.name": provider,
             "gen_ai.operation.name": self._gen_ai_operation_name(kwargs),
@@ -221,7 +225,7 @@ class OTELGenAISemconvMixin:
         self,
         kwargs: dict,
         response_obj: dict,
-        provider: str,
+        provider: str | None,
         otel_logger,
         parent_ctx,
     ) -> None:
@@ -239,6 +243,6 @@ class OTELGenAISemconvMixin:
             severity_number=SeverityNumber.INFO,
             severity_text="INFO",
             body=None,
-            attributes=self._build_inference_details_attrs(kwargs, response_obj, provider),
+            attributes=drop_none(self._build_inference_details_attrs(kwargs, response_obj, provider)),
         )
         otel_logger.emit(log_record)

litellm/integrations/otel/logger.py

@@ -33,6 +33,7 @@ from litellm.integrations.otel.model.metadata import (
     LLMCallEvent,
     RequestIdentity,
     auth_metadata,
+    metadata_from_request_data,
     model_from_request_data,
 )
 from litellm.integrations.otel.model.payloads import (
@@ -679,7 +680,12 @@ class OpenTelemetryV2(CustomLogger):
     #  / errors are the FastAPI instrumentor's job, so we don't touch it here.
     # ====================================================================== #
 
-    def seed_request_identity(self, user_api_key_dict: object, model: str | None = None) -> None:
+    def seed_request_identity(
+        self,
+        user_api_key_dict: object,
+        model: str | None = None,
+        request_metadata: Mapping[str, object] | None = None,
+    ) -> None:
         """Attach request-identity Baggage to the current context + server span.
 
         Seeding identity into Baggage makes **every** span emitted afterwards for
@@ -691,7 +697,7 @@ class OpenTelemetryV2(CustomLogger):
         isn't determined yet, which is correct.
         """
         try:
-            identity: Final = RequestIdentity.from_user_api_key_auth(user_api_key_dict)
+            identity: Final = RequestIdentity.from_user_api_key_auth(user_api_key_dict, request_metadata)
             bag: Final = promoted_baggage(
                 identity,
                 model,
@@ -743,6 +749,7 @@ class OpenTelemetryV2(CustomLogger):
         self.seed_request_identity(
             user_api_key_dict,
             model=model_from_request_data(data),
+            request_metadata=metadata_from_request_data(data),
         )
         return data
 

litellm/integrations/otel/model/baggage.py

@@ -15,9 +15,10 @@ never promoted whole.
 
 import json
 from collections.abc import Callable, Mapping
+from types import MappingProxyType
 from typing import Final
 
-from litellm.integrations.otel.model.metadata import RequestIdentity
+from litellm.integrations.otel.model.metadata import REQUESTER_METADATA_PATH, RequestIdentity
 from litellm.integrations.otel.model.semconv import GenAI, LiteLLM
 
 # Attribute key -> value extractor over (identity, request_model,
@@ -79,17 +80,23 @@ def promoted_baggage(
     ``team_metadata_keys`` selects sub-keys of the team's metadata to promote
     under ``litellm.team.metadata``. Empty values are dropped.
     """
-    out: Final[dict[str, str]] = {}
-    for key, extract in _PROMOTABLE.items():
-        if key in promoted_keys:
-            value = extract(identity, request_model, team_metadata_keys)
-            if value:
-                out[key] = value
-    for meta_key in metadata_keys:
-        value = identity.metadata.get(meta_key)
-        if value:
-            out[f"{LiteLLM.METADATA_PREFIX}{meta_key}"] = value
-    return out
+    identity_values: Final = {
+        key: value
+        for key, extract in _PROMOTABLE.items()
+        if key in promoted_keys and (value := extract(identity, request_model, team_metadata_keys))
+    }
+    return {**identity_values, **promoted_metadata(identity.metadata, metadata_keys)}
+
+
+def promoted_metadata(metadata: Mapping[str, str], metadata_keys: tuple[str, ...]) -> Mapping[str, str]:
+    """Allowlisted entries of a flattened metadata mapping under ``litellm.metadata.*``."""
+    return MappingProxyType(
+        {
+            f"{LiteLLM.METADATA_PREFIX}{meta_key.removeprefix(REQUESTER_METADATA_PATH)}": value
+            for meta_key in metadata_keys
+            if (value := metadata.get(meta_key))
+        }
+    )
 
 
 def _filtered_team_metadata_json(

litellm/integrations/otel/model/config.py

@@ -210,7 +210,10 @@ class OpenTelemetryV2Config(BaseSettings):
         validation_alias=AliasChoices("baggage_metadata_keys", "LITELLM_OTEL_BAGGAGE_METADATA_KEYS"),
         description=(
             "Metadata sub-keys promoted under the ``litellm.metadata.*`` "
-            "namespace. Configure via the ``LITELLM_OTEL_BAGGAGE_METADATA_KEYS`` "
+            "namespace. A dotted path such as ``requester_metadata.trace_id`` "
+            "reads the caller's nested ``metadata.trace_id`` and is promoted as "
+            "``litellm.metadata.trace_id``; other dotted keys keep their full path. "
+            "Configure via the ``LITELLM_OTEL_BAGGAGE_METADATA_KEYS`` "
             "env var (comma-separated) or "
             "``callback_settings.otel.baggage_metadata_keys`` in config.yaml."
         ),

litellm/integrations/otel/model/metadata.py

@@ -49,6 +49,8 @@ if TYPE_CHECKING:
     from litellm.types.utils import StandardLoggingPayload
 
 LANGFUSE_TRACE_NAME_HEADER: Final = "langfuse_trace_name"
+REQUESTER_METADATA_KEY: Final = "requester_metadata"
+REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}."
 
 
 @dataclass(frozen=True)
@@ -78,7 +80,7 @@ class RequestIdentity:
         model, not just the user-facing one.
         """
         raw_meta: Final = cast(Mapping[str, object], payload.get("metadata") or {})
-        metadata = {key: str(value) for key, value in raw_meta.items() if isinstance(value, (str, bool, int, float))}
+        metadata: Final = MappingProxyType(dict(flatten_metadata(raw_meta)))
         return cls(
             call_id=as_str(payload.get("litellm_call_id")) or as_str(payload.get("id")),
             # StandardLoggingMetadata's canonical key is ``user_api_key_team_id``;
@@ -95,7 +97,9 @@ class RequestIdentity:
         )
 
     @classmethod
-    def from_user_api_key_auth(cls, auth: object) -> RequestIdentity:
+    def from_user_api_key_auth(
+        cls, auth: object, request_metadata: Mapping[str, object] | None = None
+    ) -> RequestIdentity:
         """Identity from a ``UserAPIKeyAuth`` (duck-typed to keep this module
         free of a proxy import).
 
@@ -103,11 +107,13 @@ class RequestIdentity:
         guardrail, or service span is created — so the whole request's spans
         inherit identity, not just the LLM-call span. Metadata sub-keys use the
         ``user_api_key_*`` names that ``baggage.DEFAULT_BAGGAGE_METADATA_KEYS``
-        promotes.
+        promotes; ``request_metadata`` (the caller's ``requester_metadata``
+        snapshot) is flattened to dotted keys so ``requester_metadata.<key>``
+        resolves too.
         """
         get: Final = lambda name: getattr(auth, name, None)  # noqa: E731
-        metadata: Final = {
-            meta_key: str(value)
+        auth_meta: Final = tuple(
+            (meta_key, str(value))
             for meta_key, attr in (
                 ("user_api_key_user_id", "user_id"),
                 ("user_api_key_org_id", "org_id"),
@@ -115,7 +121,9 @@ class RequestIdentity:
                 ("user_api_key_end_user_id", "end_user_id"),
             )
             if (value := get(attr))
-        }
+        )
+        request_meta: Final = flatten_metadata(request_metadata) if request_metadata is not None else ()
+        metadata: Final = MappingProxyType(dict((*request_meta, *auth_meta)))
         return cls(
             team_id=as_str(get("team_id")),
             team_alias=as_str(get("team_alias")),
@@ -351,6 +359,35 @@ def model_from_request_data(data: object) -> str | None:
     return None
 
 
+def metadata_from_request_data(data: object) -> Mapping[str, object] | None:
+    """The caller's ``requester_metadata`` snapshot from a pre-call ``data`` dict, keyed under its wrapper.
+
+    The proxy stores it under ``metadata`` or ``litellm_metadata`` depending on the route;
+    the proxy-owned siblings (``user_api_key_*``, ``requester_ip_address``) are not read.
+    """
+    top: Final = _as_str_mapping(data)
+    if top is None:
+        return None
+    snapshots: Final = tuple(
+        snapshot
+        for name in ("metadata", "litellm_metadata")
+        if (nested := _as_str_mapping(top.get(name))) is not None
+        and (snapshot := _as_str_mapping(nested.get(REQUESTER_METADATA_KEY))) is not None
+    )
+    return MappingProxyType({REQUESTER_METADATA_KEY: snapshots[0]}) if snapshots else None
+
+
+def flatten_metadata(raw: Mapping[str, object]) -> Iterator[tuple[str, str]]:
+    """Scalar leaves of a nested metadata mapping, keyed by their dotted path."""
+    stack: Final = list(tuple(raw.items())[::-1])  # mutable-ok: iterative worklist keeps the walk off the call stack
+    while stack:
+        key, value = stack.pop()
+        if (nested := _as_str_mapping(value)) is not None:
+            stack.extend(tuple((f"{key}.{sub_key}", sub_value) for sub_key, sub_value in nested.items())[::-1])
+        elif isinstance(value, (str, bool, int, float)):
+            yield key, str(value)
+
+
 def resolve_provider_model(payload: StandardLoggingPayload) -> str | None:
     """The model litellm dispatched to the provider, from the payload.
 

litellm/integrations/otel/plumbing/context.py

@@ -26,6 +26,7 @@ if TYPE_CHECKING:
     from litellm.integrations.otel.model.destination import OtelDestination
 
 _PROPAGATOR: Final = TraceContextTextMapPropagator()
+_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate"))
 
 # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
 # proxy first resolves it, so request-level spans (the LLM call, guardrails) can
@@ -310,6 +311,37 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
     return _PROPAGATOR.extract(carrier)
 
 
+def _outgoing_trace_context(parent_span: object) -> Context | None:
+    if isinstance(parent_span, Span) and is_recordable_span(parent_span):
+        return context_from_span(parent_span)
+
+    root: Final = request_root_span()
+    if root is not None:
+        return context_from_span(root)
+
+    current: Final = get_current()
+    if is_recordable_span(get_current_span(current)):
+        return current
+    return None
+
+
+def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
+    """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
+
+    Parent preference: ``parent_span`` (the request span auth stashed on the key), then
+    the anchored request root span, then the ambient active span. Only trace context is
+    injected, never Baggage. Unchanged when no valid span exists anywhere.
+    """
+    context: Final = _outgoing_trace_context(parent_span)
+    if context is None:
+        return dict(headers)  # mutable-ok: OpenTelemetry propagator requires a mutable carrier
+    carrier: Final = {  # mutable-ok: OpenTelemetry propagator requires a mutable carrier
+        key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
+    }
+    _PROPAGATOR.inject(carrier, context=context)
+    return carrier
+
+
 # The OTLP destinations this request's key or team pointed its traces at, resolved
 # once during auth. A ``ContextVar`` for the same reason the root span above is one:
 # it rides the request task's context into the ``asyncio.create_task`` children that

litellm/integrations/prometheus.py

@@ -13,6 +13,7 @@ from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast
 
 from pydantic import BaseModel
+from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm._logging import print_verbose, verbose_logger
@@ -36,6 +37,7 @@ from litellm.litellm_core_utils.core_helpers import (
 from litellm.litellm_core_utils.service_tier_utils import (
     get_service_tier_from_standard_logging_payload,
 )
+from litellm.models.end_user import LiteLLM_EndUserTable
 from litellm.proxy._types import (
     LiteLLM_DeletedVerificationToken,
     LiteLLM_TeamTable,
@@ -43,7 +45,9 @@ from litellm.proxy._types import (
     UserAPIKeyAuth,
 )
 from litellm.repositories.base_repository import BaseRepository
+from litellm.repositories.budget_repository import BudgetRepository
 from litellm.repositories.organization_repository import OrganizationRepository
+from litellm.repositories.table_repositories import EndUserRepository
 from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.user_repository import UserRepository
 from litellm.types.guardrails import GuardrailEventHooks
@@ -66,13 +70,26 @@ from litellm.types.utils import (
 
 if TYPE_CHECKING:
     from apscheduler.schedulers.asyncio import AsyncIOScheduler
+    from prisma.types import (
+        LiteLLM_BudgetTableWhereUniqueInput,
+        LiteLLM_EndUserTableInclude,
+        LiteLLM_EndUserTableOrderByInput,
+    )
     from prometheus_client import Gauge
     from prometheus_client.metrics import MetricWrapperBase
 
+    from litellm.proxy.utils import PrismaClient
     from litellm.router import Router
 else:
     AsyncIOScheduler = Any
 
+_IsNotNull = TypedDict("_IsNotNull", {"not": ReadOnly[None]})
+
+
+class _BudgetedCustomerFilter(TypedDict):
+    budget_id: ReadOnly[_IsNotNull]
+
+
 _BudgetRowT: Final = TypeVar("_BudgetRowT")
 _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
 
@@ -116,8 +133,8 @@ def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrisma
     )
 
 
-class _OrgBudgetRow(Protocol):
-    """The budget columns joined onto an organization row."""
+class _JoinedBudgetRow(Protocol):
+    """The budget columns joined onto an organization or customer row."""
 
     @property
     def max_budget(self) -> float | None: ...
@@ -126,6 +143,23 @@ class _OrgBudgetRow(Protocol):
     def budget_reset_at(self) -> datetime | None: ...
 
 
+class _CustomerBudgetRow(Protocol):
+    """The columns of a customer (end user) row that budget gauges read."""
+
+    @property
+    def user_id(self) -> str: ...
+
+    @property
+    def spend(self) -> float: ...
+
+    @property
+    def litellm_budget_table(self) -> _JoinedBudgetRow | None: ...
+
+
+def _customer_budget_metrics_enabled() -> bool:
+    return litellm.enable_end_user_cost_tracking_prometheus_only is True and not litellm.disable_end_user_cost_tracking
+
+
 class _ExcludedLabelMetric:
     """Proxies a prometheus metric whose declared ``labelnames`` had globally
     excluded labels removed, dropping those labels from every ``labels(...)``
@@ -471,6 +505,24 @@ class PrometheusLogger(CustomLogger):
                 labelnames=self.get_labels_for_metric("litellm_user_budget_remaining_hours_metric"),
             )
 
+            self.litellm_remaining_customer_budget_metric = self._gauge_factory(
+                "litellm_remaining_customer_budget_metric",
+                "Remaining budget for customer (end user)",
+                labelnames=self.get_labels_for_metric("litellm_remaining_customer_budget_metric"),
+            )
+
+            self.litellm_customer_max_budget_metric = self._gauge_factory(
+                "litellm_customer_max_budget_metric",
+                "Maximum budget set for customer (end user)",
+                labelnames=self.get_labels_for_metric("litellm_customer_max_budget_metric"),
+            )
+
+            self.litellm_customer_budget_remaining_hours_metric = self._gauge_factory(
+                "litellm_customer_budget_remaining_hours_metric",
+                "Remaining hours for customer (end user) budget to be reset",
+                labelnames=self.get_labels_for_metric("litellm_customer_budget_remaining_hours_metric"),
+            )
+
             ########################################
             # LiteLLM Virtual API KEY metrics
             ########################################
@@ -1334,7 +1386,7 @@ class PrometheusLogger(CustomLogger):
         self,
         metric: Any,
         metric_name: DEFINED_PROMETHEUS_METRICS,
-        labels: dict[str, str | None],
+        labels: Mapping[str, str | None],
     ) -> None:
         """
         Cap the cardinality of metrics that include the ``end_user`` label.
@@ -1501,6 +1553,7 @@ class PrometheusLogger(CustomLogger):
             response_cost=response_cost,
             user_id=user_id,
             user_api_key_org_id=user_api_key_org_id,
+            end_user_id=end_user_id,
         )
 
         # set proxy virtual key rpm/tpm metrics
@@ -1930,12 +1983,14 @@ class PrometheusLogger(CustomLogger):
         response_cost: float,
         user_id: str | None = None,
         user_api_key_org_id: str | None = None,
+        end_user_id: str | None = None,
     ):
         if (
             isinstance(self.litellm_remaining_team_budget_metric, NoOpMetric)
             and isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric)
             and isinstance(self.litellm_remaining_user_budget_metric, NoOpMetric)
             and isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric)
+            and self._customer_budget_gauges_are_noop()
         ):
             return
 
@@ -1990,6 +2045,10 @@ class PrometheusLogger(CustomLogger):
                 carried=OrgBudgetSnapshot.from_metadata(_metadata),
                 org_alias=_org_alias if isinstance(_org_alias, str) else None,
             ),
+            self._set_customer_budget_metrics_after_api_request(
+                end_user_id=end_user_id,
+                response_cost=response_cost,
+            ),
             return_exceptions=True,
         )
         try:
@@ -2006,7 +2065,7 @@ class PrometheusLogger(CustomLogger):
             if isinstance(r, Exception):
                 verbose_logger.debug(
                     "[Non-Blocking] Prometheus: Budget metric lookup %s failed: %s",
-                    ["key", "team", "user", "org"][i],
+                    ("key", "team", "user", "org", "customer")[i],
                     r,
                 )
 
@@ -3574,9 +3633,9 @@ class PrometheusLogger(CustomLogger):
 
     async def _initialize_budget_metrics(
         self,
-        data_fetch_function: Callable[..., Awaitable[tuple[list[_BudgetRowT], int | None]]],
-        set_metrics_function: Callable[[list[_BudgetRowT]], Awaitable[None]],
-        data_type: Literal["teams", "keys", "users", "orgs"],
+        data_fetch_function: Callable[..., Awaitable[tuple[Sequence[_BudgetRowT], int | None]]],
+        set_metrics_function: Callable[[Sequence[_BudgetRowT]], Awaitable[None]],
+        data_type: Literal["teams", "keys", "users", "orgs", "customers"],
     ):
         """
         Generic method to initialize budget metrics for teams or API keys.
@@ -3735,6 +3794,49 @@ class PrometheusLogger(CustomLogger):
             data_type="orgs",
         )
 
+    async def _initialize_customer_budget_metrics(self):
+        from litellm.proxy.proxy_server import prisma_client
+
+        if prisma_client is None:
+            verbose_logger.debug("Prometheus: skipping customer metrics initialization, DB not initialized")
+            return
+
+        if self._customer_budget_gauges_are_noop():
+            return
+
+        if not _customer_budget_metrics_enabled():
+            verbose_logger.debug("Prometheus: skipping customer metrics initialization, end_user tracking disabled")
+            return
+
+        default_budget: Final = await self._get_default_customer_budget(prisma_client)
+        customers_table: Final = EndUserRepository(prisma_client).table
+        with_persisted_budget: Final[_BudgetedCustomerFilter] = {"budget_id": {"not": None}}
+        budgeted_customers: Final = None if default_budget is not None else with_persisted_budget
+        by_user_id: Final[LiteLLM_EndUserTableOrderByInput] = {"user_id": "asc"}
+        with_budget: Final[LiteLLM_EndUserTableInclude] = {"litellm_budget_table": True}
+
+        async def fetch_customers(page_size: int, page: int) -> tuple[Sequence[_CustomerBudgetRow], int | None]:
+            skip: Final = (page - 1) * page_size
+            customers: Final = await customers_table.find_many(
+                skip=skip,
+                take=page_size,
+                where=budgeted_customers,
+                order=by_user_id,
+                include=with_budget,
+            )
+            total_count: Final = await customers_table.count(where=budgeted_customers) if page == 1 else None
+            return customers, total_count
+
+        async def set_customer_metrics(customers: Sequence[_CustomerBudgetRow]) -> None:
+            for customer in customers:
+                self._set_customer_budget_metrics_from_row(customer, default_budget=default_budget)
+
+        await self._initialize_budget_metrics(
+            data_fetch_function=fetch_customers,
+            set_metrics_function=set_customer_metrics,
+            data_type="customers",
+        )
+
     async def initialize_remaining_budget_metrics(self):
         """
         Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies.
@@ -3765,11 +3867,12 @@ class PrometheusLogger(CustomLogger):
         """
         Helper to initialize remaining budget metrics for all teams, API keys, and users.
         """
-        verbose_logger.debug("Emitting key, team, user, org budget metrics....")
+        verbose_logger.debug("Emitting key, team, user, org, customer budget metrics....")
         await self._initialize_team_budget_metrics()
         await self._initialize_api_key_budget_metrics()
         await self._initialize_user_budget_metrics()
         await self._initialize_org_budget_metrics()
+        await self._initialize_customer_budget_metrics()
         await self._initialize_user_and_team_count_metrics()
 
     async def _initialize_user_and_team_count_metrics(self):
@@ -3805,27 +3908,27 @@ class PrometheusLogger(CustomLogger):
             verbose_logger.exception("Error initializing user/team count metrics: %s", e)
 
     async def _set_key_list_budget_metrics(
-        self, keys: list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]
+        self, keys: Sequence[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]
     ) -> None:
         """Helper function to set budget metrics for a list of keys"""
         for key in keys:
             if isinstance(key, UserAPIKeyAuth):
                 self._set_key_budget_metrics(key)
 
-    async def _set_team_list_budget_metrics(self, teams: list[LiteLLM_TeamTable]):
+    async def _set_team_list_budget_metrics(self, teams: Sequence[LiteLLM_TeamTable]):
         """Helper function to set budget metrics for a list of teams"""
         for team in teams:
             self._set_team_budget_metrics(team)
 
-    async def _set_user_list_budget_metrics(self, users: list[LiteLLM_UserTable]):
+    async def _set_user_list_budget_metrics(self, users: Sequence[LiteLLM_UserTable]):
         """Helper function to set budget metrics for a list of users"""
         for user in users:
             self._set_user_budget_metrics(user)
 
-    async def _set_org_list_budget_metrics(self, orgs: list):
+    async def _set_org_list_budget_metrics(self, orgs: Sequence):
         """Helper function to set budget metrics for a list of orgs"""
         for org in orgs:
-            budget_table: _OrgBudgetRow | None = getattr(org, "litellm_budget_table", None)
+            budget_table: _JoinedBudgetRow | None = getattr(org, "litellm_budget_table", None)
             self._set_org_budget_metrics(
                 org_id=org.organization_id or "",
                 org_alias=org.organization_alias or "",
@@ -3834,6 +3937,19 @@ class PrometheusLogger(CustomLogger):
                 budget_reset_at=(getattr(budget_table, "budget_reset_at", None) if budget_table else None),
             )
 
+    def _set_customer_budget_metrics_from_row(
+        self, customer: _CustomerBudgetRow, default_budget: _JoinedBudgetRow | None
+    ):
+        budget_table: Final = (
+            customer.litellm_budget_table if customer.litellm_budget_table is not None else default_budget
+        )
+        self._set_customer_budget_metrics(
+            end_user_id=customer.user_id,
+            spend=customer.spend,
+            max_budget=budget_table.max_budget if budget_table is not None else None,
+            budget_reset_at=budget_table.budget_reset_at if budget_table is not None else None,
+        )
+
     async def _set_team_budget_metrics_after_api_request(
         self,
         user_api_team: str | None,
@@ -4083,6 +4199,98 @@ class PrometheusLogger(CustomLogger):
                 self._get_remaining_hours_for_budget_reset(budget_reset_at=budget_reset_at)
             )
 
+    async def _set_customer_budget_metrics_after_api_request(
+        self,
+        end_user_id: str | None,
+        response_cost: float,
+    ):
+        if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled():
+            return
+
+        if not end_user_id:
+            return
+
+        from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key
+        from litellm.proxy.proxy_server import user_api_key_cache
+
+        try:
+            cached_customer: Final = await user_api_key_cache.async_get_cache(
+                key=end_user_cache_key(end_user_id),
+                model_type=LiteLLM_EndUserTable,
+            )
+        except Exception as e:
+            verbose_logger.debug("[Non-Blocking] Prometheus: Error getting customer info: %s", e)
+            return
+
+        if cached_customer is None:
+            return
+
+        budget_table: Final = cached_customer.litellm_budget_table
+        self._set_customer_budget_metrics(
+            end_user_id=end_user_id,
+            spend=cached_customer.spend + response_cost,
+            max_budget=budget_table.max_budget if budget_table is not None else None,
+            budget_reset_at=None,
+        )
+
+    async def _get_default_customer_budget(self, prisma_client: PrismaClient) -> _JoinedBudgetRow | None:
+        default_budget_id: Final = litellm.max_end_user_budget_id
+        if default_budget_id is None:
+            return None
+        default_budget_key: Final[LiteLLM_BudgetTableWhereUniqueInput] = {"budget_id": default_budget_id}
+        try:
+            return await BudgetRepository(prisma_client).table.find_unique(where=default_budget_key)
+        except Exception as e:
+            verbose_logger.debug("[Non-Blocking] Prometheus: Error getting default customer budget: %s", e)
+            return None
+
+    def _customer_budget_gauges_are_noop(self) -> bool:
+        return (
+            isinstance(self.litellm_remaining_customer_budget_metric, NoOpMetric)
+            and isinstance(self.litellm_customer_max_budget_metric, NoOpMetric)
+            and isinstance(self.litellm_customer_budget_remaining_hours_metric, NoOpMetric)
+        )
+
+    def _set_customer_budget_metrics(
+        self,
+        end_user_id: str,
+        spend: float,
+        max_budget: float | None,
+        budget_reset_at: datetime | None,
+    ):
+        _labels: Final[dict[str, str | None]] = prometheus_label_factory(
+            supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_customer_budget_metric"),
+            enum_values=UserAPIKeyLabelValues(end_user=end_user_id),
+        )
+        if _labels.get(UserAPIKeyLabelNames.END_USER.value) is None:
+            return
+
+        self.litellm_remaining_customer_budget_metric.labels(**_labels).set(
+            self._safe_get_remaining_budget(
+                max_budget=max_budget,
+                spend=spend,
+            )
+        )
+        self._track_end_user_metric_series(
+            self.litellm_remaining_customer_budget_metric, "litellm_remaining_customer_budget_metric", _labels
+        )
+
+        if max_budget is not None:
+            self.litellm_customer_max_budget_metric.labels(**_labels).set(max_budget)
+            self._track_end_user_metric_series(
+                self.litellm_customer_max_budget_metric, "litellm_customer_max_budget_metric", _labels
+            )
+
+        if budget_reset_at is not None:
+            self.litellm_customer_budget_remaining_hours_metric.labels(**_labels).set(
+                self._get_remaining_hours_for_budget_reset(budget_reset_at=budget_reset_at)
+            )
+            self._track_end_user_metric_series(
+                self.litellm_customer_budget_remaining_hours_metric,
+                "litellm_customer_budget_remaining_hours_metric",
+                _labels,
+            )
+
     def _set_key_budget_metrics(self, user_api_key_dict: UserAPIKeyAuth):
         """
         Set virtual key budget metrics

litellm/integrations/s3.py

@@ -2,19 +2,42 @@
 #    On success + failure, log events to Supabase
 
 import hashlib
+import os
+from collections.abc import Mapping
 from datetime import datetime
 from typing import Final, cast
 
+from pydantic import TypeAdapter, ValidationError
+
 import litellm
 from litellm._logging import print_verbose, verbose_logger
 from litellm.constants import (
     MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES,
     MAX_S3_OBJECT_KEY_BYTES,
     S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
+    S3_LOG_PROMPTS_ONLY_ENV_VAR,
     S3_PREFIX_DIGEST_CHARS,
 )
 from litellm.types.utils import StandardLoggingPayload
 
+_S3_LOG_PROMPTS_ONLY: Final = TypeAdapter(bool)
+
+
+def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | None = None) -> bool:
+    env: Final = os.environ if environ is None else environ
+    raw: Final = env.get(S3_LOG_PROMPTS_ONLY_ENV_VAR) if configured is None else configured
+    if raw is None or raw == "":
+        return False
+    try:
+        return _S3_LOG_PROMPTS_ONLY.validate_python(raw.strip() if isinstance(raw, str) else raw)
+    except ValidationError:
+        verbose_logger.warning("s3 logging: s3_log_prompts_only=%r is not a boolean, logging prompts only", raw)
+        return True
+
+
+def prompts_only_payload(payload: StandardLoggingPayload) -> StandardLoggingPayload:
+    return {**payload, "response": None}
+
 
 class S3Logger:
     # Class variables or attributes
@@ -33,6 +56,7 @@ class S3Logger:
         s3_config=None,
         s3_server_side_encryption: str | None = None,
         s3_sse_kms_key_id: str | None = None,
+        s3_log_prompts_only: bool | None = None,
         **kwargs,
     ):
         import boto3
@@ -41,29 +65,30 @@ class S3Logger:
             verbose_logger.debug("in init s3 logger - s3_callback_params %s", litellm.s3_callback_params)
 
             s3_use_team_prefix = False
+            params: Final = {
+                key: litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value
+                for key, value in (litellm.s3_callback_params or {}).items()
+            }
 
             if litellm.s3_callback_params is not None:
-                # read in .env variables - example os.environ/AWS_BUCKET_NAME
-                for key, value in litellm.s3_callback_params.items():
-                    if isinstance(value, str) and value.startswith("os.environ/"):
-                        litellm.s3_callback_params[key] = litellm.get_secret(value)
-                # now set s3 params from litellm.s3_logger_params
-                s3_bucket_name = litellm.s3_callback_params.get("s3_bucket_name")
-                s3_region_name = litellm.s3_callback_params.get("s3_region_name")
-                s3_api_version = litellm.s3_callback_params.get("s3_api_version")
-                s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl", True)
-                s3_verify = litellm.s3_callback_params.get("s3_verify")
-                s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url")
-                s3_aws_access_key_id = litellm.s3_callback_params.get("s3_aws_access_key_id")
-                s3_aws_secret_access_key = litellm.s3_callback_params.get("s3_aws_secret_access_key")
-                s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token")
-                s3_config = litellm.s3_callback_params.get("s3_config")
-                s3_path = litellm.s3_callback_params.get("s3_path")
-                s3_server_side_encryption = litellm.s3_callback_params.get("s3_server_side_encryption")
-                s3_sse_kms_key_id = litellm.s3_callback_params.get("s3_sse_kms_key_id")
-                # done reading litellm.s3_callback_params
-                s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False))
+                s3_bucket_name = params.get("s3_bucket_name")
+                s3_region_name = params.get("s3_region_name")
+                s3_api_version = params.get("s3_api_version")
+                s3_use_ssl = params.get("s3_use_ssl", True)
+                s3_verify = params.get("s3_verify")
+                s3_endpoint_url = params.get("s3_endpoint_url")
+                s3_aws_access_key_id = params.get("s3_aws_access_key_id")
+                s3_aws_secret_access_key = params.get("s3_aws_secret_access_key")
+                s3_aws_session_token = params.get("s3_aws_session_token")
+                s3_config = params.get("s3_config")
+                s3_path = params.get("s3_path")
+                s3_server_side_encryption = params.get("s3_server_side_encryption")
+                s3_sse_kms_key_id = params.get("s3_sse_kms_key_id")
+                s3_use_team_prefix = bool(params.get("s3_use_team_prefix", False))
             self.s3_use_team_prefix = s3_use_team_prefix
+            self.s3_log_prompts_only: object = (
+                params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only
+            )
             self.bucket_name = s3_bucket_name
             self.s3_path = s3_path
             self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
@@ -144,7 +169,9 @@ class S3Logger:
 
             from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
-            payload_str: Final = safe_dumps(payload)
+            payload_str: Final = safe_dumps(
+                prompts_only_payload(payload) if resolve_s3_log_prompts_only(self.s3_log_prompts_only) else payload
+            )
 
             print_verbose(f"\ns3 Logger - Logging payload = {payload_str}")
 

litellm/integrations/s3_v2.py

@@ -21,6 +21,8 @@ from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_S
 from litellm.integrations.s3 import (
     get_s3_object_download_filename,
     get_s3_object_key,
+    prompts_only_payload,
+    resolve_s3_log_prompts_only,
     resolve_sse_params,
 )
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
@@ -68,6 +70,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
         s3_use_virtual_hosted_style: bool = False,
         s3_server_side_encryption: str | None = None,
         s3_sse_kms_key_id: str | None = None,
+        s3_log_prompts_only: bool | None = None,
         s3_callback_params_override: dict | None = None,
         **kwargs,
     ):
@@ -108,6 +111,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
                 s3_use_virtual_hosted_style=s3_use_virtual_hosted_style,
                 s3_server_side_encryption=s3_server_side_encryption,
                 s3_sse_kms_key_id=s3_sse_kms_key_id,
+                s3_log_prompts_only=s3_log_prompts_only,
             )
             verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url)
 
@@ -163,6 +167,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
         s3_use_virtual_hosted_style: bool = False,
         s3_server_side_encryption: str | None = None,
         s3_sse_kms_key_id: str | None = None,
+        s3_log_prompts_only: bool | None = None,
         params_source: dict | None = None,
     ):
         """
@@ -212,6 +217,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
             bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style
         )
 
+        self.s3_log_prompts_only: object = (
+            params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only
+        )
+
         self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params(
             params.get("s3_server_side_encryption") or s3_server_side_encryption,
             params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id,
@@ -489,8 +498,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
 
         s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"])
 
+        payload: Final = (
+            prompts_only_payload(standard_logging_payload)
+            if resolve_s3_log_prompts_only(self.s3_log_prompts_only)
+            else standard_logging_payload
+        )
         return s3BatchLoggingElement(
-            payload=dict(standard_logging_payload),
+            payload=dict(payload),
             s3_object_key=s3_object_key,
             s3_object_download_filename=s3_object_download_filename,
         )

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -148,6 +148,7 @@ class _ToolCallChunk(TypedDict):
 class _UsageBearingChunk(TypedDict, total=False):
     usage: Usage | None
     _hidden_params: Mapping[str, str]
+    choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
 
 
 class _UsageSummary(TypedDict):
@@ -921,21 +922,22 @@ class ChunkProcessor:
 
         prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
 
-        completion_tokens = self._reset_anthropic_cursor_completion_tokens(
+        recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
             chunks=chunks,
             completion_tokens=completion_tokens,
             completion_usage_updates=completion_usage_updates,
         )
+        cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
 
         return UsagePerChunk(
             prompt_tokens=prompt_tokens,
-            completion_tokens=completion_tokens,
+            completion_tokens=recovered_completion_tokens,
             cache_creation_input_tokens=cache_creation_input_tokens,
             cache_read_input_tokens=cache_read_input_tokens,
             server_tool_use=server_tool_use,
             web_search_requests=web_search_requests,
             google_maps_grounding_requests=google_maps_grounding_requests,
-            completion_tokens_details=completion_tokens_details,
+            completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
             prompt_tokens_details=prompt_tokens_details,
             cost=cost,
             inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
@@ -960,6 +962,30 @@ class ChunkProcessor:
         ]
         return values[-1] if values else None
 
+    @staticmethod
+    def _finish_reason_of_choice(choice: object) -> str | None:
+        match choice:
+            case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
+                return reason
+            case {"finish_reason": str() as reason}:
+                return reason
+            case _:
+                return None
+
+    @staticmethod
+    def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]:
+        if isinstance(chunk, dict):
+            return chunk.get("choices", ())
+        return getattr(chunk, "choices", ())
+
+    @staticmethod
+    def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
+        return any(
+            ChunkProcessor._finish_reason_of_choice(choice) is not None
+            for chunk in chunks
+            for choice in ChunkProcessor._chunk_choices(chunk)
+        )
+
     @staticmethod
     def _reset_anthropic_cursor_completion_tokens(
         chunks: Sequence["_UsageBearingChunk | ModelResponse"],
@@ -970,18 +996,18 @@ class ChunkProcessor:
 
         See the ``completion_usage_updates`` comment in
         ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
-        cursor when either it is > 1 (definitely not a placeholder) or we saw
-        >= 2 completion-bearing usage events (positive evidence ``message_delta``
-        arrived). Otherwise — the only completion update we ever saw was the
-        Anthropic ``message_start`` cursor (=1) — reset to 0 so
-        ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
-        from the actually-received completion text instead of trusting the
-        placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
-        heuristic (which encodes Anthropic's specific message_start SSE shape)
-        does not silently affect other providers that may legitimately report
-        ``completion_tokens=1`` from a single usage event.
+        cursor when we saw >= 2 completion-bearing usage events or any chunk
+        carried a ``finish_reason`` (positive evidence ``message_delta``
+        arrived). Otherwise the only completion update we ever saw was the
+        Anthropic ``message_start`` cursor, a small placeholder whose magnitude
+        varies per request (1 and 8 both observed live), so reset to 0 and let
+        ``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
+        the actually-received text and reasoning instead. Gated on
+        ``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
+        Anthropic's specific message_start SSE shape) does not silently affect
+        other providers that legitimately report usage from a single event.
         """
-        saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
+        saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
         if saw_non_cursor_completion:
             return completion_tokens
 
@@ -995,7 +1021,7 @@ class ChunkProcessor:
             if isinstance(hp, dict):
                 custom_llm_provider = hp.get("custom_llm_provider")
 
-        if custom_llm_provider == "anthropic" and completion_tokens == 1:
+        if custom_llm_provider == "anthropic":
             return 0
         return completion_tokens
 
@@ -1039,10 +1065,13 @@ class ChunkProcessor:
             returned_usage.prompt_tokens = 0
         returned_usage.completion_tokens = (
             completion_tokens
-            or token_counter(
-                model=model,
-                text=completion_output,
-                count_response_tokens=True,  # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
+            or (
+                token_counter(
+                    model=model,
+                    text=completion_output,
+                    count_response_tokens=True,  # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
+                )
+                + (reasoning_tokens or 0)
             )
         )
         returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
@@ -1066,15 +1095,16 @@ class ChunkProcessor:
                 returned_usage.completion_tokens_details = completion_tokens_details
 
         if reasoning_tokens is not None:
+            capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
             if returned_usage.completion_tokens_details is None:
                 returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
-                    reasoning_tokens=reasoning_tokens
+                    reasoning_tokens=capped_reasoning_tokens,
+                    text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens,
                 )
             elif (
                 returned_usage.completion_tokens_details is not None
                 and returned_usage.completion_tokens_details.reasoning_tokens is None
             ):
-                capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
                 returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
                 if returned_usage.completion_tokens_details.text_tokens is None:
                     returned_usage.completion_tokens_details.text_tokens = (

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

@@ -1561,10 +1561,12 @@ class AnthropicMessagesHandler(BaseTranslation):
 
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
+        tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
+            tool_calls=tool_use_fingerprints if stream_ended else (),
             stream_ended=stream_ended,
+            tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,
         )
 
     @classmethod

litellm/llms/anthropic/chat/handler.py

@@ -632,6 +632,7 @@ class ModelResponseIterator:
         self.tool_name_reverse_map: dict[str, str] = tool_name_reverse_map or {}
         # Generate response ID once per stream to match OpenAI-compatible behavior
         self.response_id = _generate_id()
+        self.served_model: str | None = None
 
         # Track if we're currently streaming a response_format tool
         self.is_response_format_tool: bool = False
@@ -1067,6 +1068,9 @@ class ModelResponseIterator:
                 }
                 """
                 message_start_block: Final = MessageStartBlock(**chunk)
+                start_message: Final = message_start_block["message"]
+                if "model" in start_message:
+                    self.served_model = start_message["model"]
                 if "usage" in message_start_block["message"]:
                     usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"])
             elif type_chunk == "error":
@@ -1098,6 +1102,7 @@ class ModelResponseIterator:
                 ],
                 usage=usage,
                 id=self.response_id,
+                model=self.served_model,
             )
 
             return returned_chunk

litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py

@@ -78,6 +78,7 @@ _PROPAGATED_METADATA_KEYS: Final = (
     "user_api_key_end_user_id",
     "user_api_end_user_max_budget",
     "user_api_key_model_max_budget",
+    "user_api_key_team_model_max_budget",
     "user_api_key_user_model_max_budget",
     "user_api_key_end_user_model_max_budget",
     "litellm_call_id",
@@ -395,9 +396,9 @@ async def _check_summary_model_budget(
     ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no
     per-model budget is configured.
 
-    All three scopes are checked because the summary's spend is charged to all
-    three: this file propagates the key, user and end-user budgets into the
-    subrequest's metadata, so enforcing only two of them would let compaction
+    Every scope is checked because the summary's spend is charged to every
+    scope: this file propagates the key, team, user and end-user budgets into the
+    subrequest's metadata, so skipping one of them would let compaction
     increment a counter it can never be refused by.
     """
     if user_api_key_auth is None:
@@ -444,6 +445,26 @@ async def _check_summary_model_budget(
             )
             return False
 
+    team_model_max_budget: Final = user_api_key_auth.team_model_max_budget
+    team_id: Final = user_api_key_auth.team_id
+    if isinstance(team_model_max_budget, dict) and team_model_max_budget and team_id is not None:
+        try:
+            await model_max_budget_limiter.is_team_within_model_budget(
+                team_id=team_id,
+                team_model_max_budget=team_model_max_budget,
+                key_model_max_budget=model_max_budget if isinstance(model_max_budget, dict) else None,
+                model=summary_model,
+            )
+        except litellm.BudgetExceededError:
+            return False
+        except Exception as e:  # noqa: BLE001  # a budget gate denies on any failure, as the other scopes do
+            verbose_logger.warning(
+                "compact_20260112: unexpected error during team model-budget check for summary_model=%s; denying: %s",
+                summary_model,
+                e,
+            )
+            return False
+
     end_user_model_max_budget: Final[dict[str, object] | None] = getattr(
         user_api_key_auth, "end_user_model_max_budget", None
     )

litellm/llms/azure/responses/transformation.py

@@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
         return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params)
 
     def get_stripped_model_name(self, model: str) -> str:
-        # if "responses/" is in the model name, remove it
-        if "responses/" in model:
-            model = model.replace("responses/", "")
-        if "o_series" in model:
-            model = model.replace("o_series/", "")
-        return model
+        return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
 
     def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
         """

litellm/llms/azure_ai/common_utils.py

@@ -9,6 +9,7 @@ from litellm.types.llms.openai import AllMessageValues
 from litellm.types.router import GenericLiteLLMParams
 
 AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
+AZURE_OPENAI_V1_HOST_SUFFIXES: Final = (".services.ai.azure.com", ".openai.azure.com")
 
 
 def is_foundry_model_inference_base(api_base: str) -> bool:
@@ -19,11 +20,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
     return "/openai/deployments" not in parsed.path
 
 
-def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
+def is_azure_openai_v1_host(api_base: str | None) -> bool:
     host: Final = urlparse(api_base).hostname if api_base else None
-    if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
-        return "api-key"
-    return "Authorization"
+    return host is not None and host.endswith(AZURE_OPENAI_V1_HOST_SUFFIXES)
+
+
+def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
+    return "api-key" if is_azure_openai_v1_host(api_base) else "Authorization"
 
 
 def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
@@ -70,6 +73,17 @@ def get_azure_ai_auth_headers(
 AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
 
 
+def azure_ai_supports_native_responses(model: str | None, api_base: str | None) -> bool:
+    resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
+    if resolved_base is not None and not is_azure_openai_v1_host(resolved_base):
+        return False
+    if model is None:
+        return True
+    if "claude" in model.lower():
+        return False
+    return AzureFoundryModelInfo.get_azure_ai_route(model) == "default"
+
+
 class AzureFoundryModelInfo(BaseLLMModelInfo):
     """Model info for Azure AI / Azure Foundry models."""
 

litellm/llms/azure_ai/responses/transformation.py

@@ -0,0 +1,53 @@
+from typing import Final
+
+import httpx
+
+from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
+from litellm.llms.azure_ai.common_utils import (
+    AzureFoundryModelInfo,
+    api_key_header_for_base,
+    get_azure_ai_auth_headers,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+_PROJECT_PATH_PREFIX: Final = ("api", "projects")
+_RESPONSES_PATH: Final = ("openai", "v1", "responses")
+
+
+def _responses_url(api_base: str) -> str:
+    base_url: Final = httpx.URL(api_base)
+    segments: Final = tuple(segment for segment in base_url.path.split("/") if segment)
+    project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else ()
+    return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None))
+
+
+class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
+    @property
+    def custom_llm_provider(self) -> LlmProviders:
+        return LlmProviders.AZURE_AI
+
+    def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
+        params: Final = litellm_params or GenericLiteLLMParams()
+        auth_headers: Final = get_azure_ai_auth_headers(
+            api_key=AzureFoundryModelInfo.get_api_key(params.api_key),
+            litellm_params=params.model_dump(),
+            api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)),
+        )
+        return {  # mutable-ok: the handler updates the returned headers in place per the dict contract
+            **headers,
+            **auth_headers,
+            "Content-Type": "application/json",
+        }
+
+    def supports_native_websocket(self) -> bool:
+        return False
+
+    def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
+        resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
+        if resolved_base is None:
+            raise ValueError(
+                "api_base is required for the Azure AI Foundry Responses API. "
+                "Set the api_base parameter or the AZURE_AI_API_BASE environment variable."
+            )
+        return _responses_url(resolved_base)

litellm/llms/base_llm/guardrail_translation/base_translation.py

@@ -40,11 +40,15 @@ class StreamingScanKey:
     """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
     compare equal when the round would scan the same content again; ``stream_ended``
     stays out of the comparison and only says whether the handler is on its
-    end-of-stream path, where an empty payload is still scanned today."""
+    end-of-stream path, where an empty payload is still scanned today.
+    ``tool_calls_in_flight`` also stays out of the comparison: it flags that tool
+    calls have streamed which this round cannot scan yet, so a buffered window
+    holding them must stay withheld until the end-of-stream scan covers them."""
 
     texts: tuple[str, ...]
     tool_calls: tuple[str, ...] = ()
     stream_ended: bool = field(default=False, compare=False)
+    tool_calls_in_flight: bool = field(default=False, compare=False)
 
     @property
     def has_nothing_to_scan(self) -> bool:

litellm/llms/bedrock/chat/converse_transformation.py

@@ -53,6 +53,7 @@ from litellm.types.llms.openai import (
     AllMessageValues,
     ChatCompletionAnnotation,
     ChatCompletionAssistantMessage,
+    ChatCompletionAssistantToolCall,
     ChatCompletionRedactedThinkingBlock,
     ChatCompletionResponseMessage,
     ChatCompletionSystemMessage,
@@ -205,6 +206,84 @@ class AmazonConverseConfig(BaseConfig):
 
         return messages_copy
 
+    @staticmethod
+    def _has_orphaned_tool_blocks(messages: list[AllMessageValues]) -> bool:
+        return any(
+            (m.get("role") == "assistant" and m.get("tool_calls")) or m.get("role") in ("tool", "function")
+            for m in messages
+        )
+
+    @staticmethod
+    def _neutralize_orphaned_tool_blocks(
+        messages: list[AllMessageValues], optional_params: dict
+    ) -> list[AllMessageValues]:
+        if optional_params.get("tools") or not AmazonConverseConfig._has_orphaned_tool_blocks(messages):
+            return messages
+
+        from litellm.litellm_core_utils.prompt_templates.common_utils import (
+            convert_content_list_to_str,
+        )
+
+        def _tool_call_text(tool_call: ChatCompletionAssistantToolCall) -> str:
+            function = tool_call.get("function") or {}
+            name = function.get("name") or "unknown_tool"
+            arguments = function.get("arguments") or ""
+            call_id = tool_call.get("id")
+            label = f"tool call {call_id}" if call_id else "tool call"
+            return f"[{label}: {name}({arguments})]"
+
+        def _result_text(message: AllMessageValues) -> str:
+            rendered = convert_content_list_to_str(message).strip()
+            return rendered or "<non-text tool result omitted>"
+
+        guardrail_active: Final = "guardrailConfig" in optional_params
+
+        def _rewrite(message: AllMessageValues) -> AllMessageValues:
+            role = message.get("role")
+            tool_calls = message.get("tool_calls")
+            if role == "assistant" and tool_calls:
+                base_text: Final = convert_content_list_to_str(message)
+                call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
+                text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
+                return ChatCompletionAssistantMessage(role="assistant", content=text)
+            if role in ("tool", "function"):
+                tool_call_id = message.get("tool_call_id")
+                name = message.get("name")
+                label = f"tool result for {tool_call_id or name or 'unknown'}"
+                result_text: Final = f"[{label}: {_result_text(message)}]"
+                # Tool results are externally controlled, so guard them wherever they
+                # land in history; _convert_consecutive_user_messages_to_guarded_text
+                # only covers the trailing user turn.
+                content: Final = [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
+                return ChatCompletionUserMessage(role="user", content=content)
+            return message
+
+        verbose_logger.warning(
+            "litellm.bedrock: request has tool blocks in message history but no "
+            "`tools=` param; neutralizing orphaned tool blocks to text so Bedrock "
+            "accepts the request without a toolConfig. Non-text tool-result "
+            "payloads are dropped. Pass `tools=` to preserve structured tool calling."
+        )
+        return [_rewrite(message) for message in messages]
+
+    @staticmethod
+    def _handle_orphaned_tool_blocks(messages: list[AllMessageValues], optional_params: dict) -> list[AllMessageValues]:
+        if litellm.bedrock_neutralize_orphaned_tool_blocks:
+            return AmazonConverseConfig._neutralize_orphaned_tool_blocks(messages, optional_params)
+
+        if "tools" in optional_params or not has_tool_call_blocks(messages):
+            return messages
+
+        if litellm.modify_params:
+            optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
+            return messages
+
+        raise litellm.utils.UnsupportedParamsError(
+            message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
+            model="",
+            llm_provider="bedrock",
+        )
+
     @classmethod
     def get_config(cls):
         return {
@@ -1609,20 +1688,6 @@ class AmazonConverseConfig(BaseConfig):
         drop_params: bool = False,
         litellm_params: Mapping[str, object] | None = None,
     ) -> CommonRequestObject:
-        ## VALIDATE REQUEST
-        """
-        Bedrock doesn't support tool calling without `tools=` param specified.
-        """
-        if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
-            if litellm.modify_params:
-                optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
-            else:
-                raise litellm.UnsupportedParamsError(
-                    message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
-                    model="",
-                    llm_provider="bedrock",
-                )
-
         # Drop thinking param if thinking is enabled but thinking_blocks are missing
         # This prevents the error: "Expected thinking or redacted_thinking, but found tool_use"
         #
@@ -1735,7 +1800,9 @@ class AmazonConverseConfig(BaseConfig):
         messages, system_content_blocks = self._transform_system_message(messages, model=model)
 
         # Convert last user message to guarded_text if guardrailConfig is present
-        messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
+        messages = self._convert_consecutive_user_messages_to_guarded_text(
+            self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
+        )
         ## TRANSFORMATION ##
 
         _data: Final[CommonRequestObject] = self._transform_request_helper(
@@ -1796,7 +1863,9 @@ class AmazonConverseConfig(BaseConfig):
         messages, system_content_blocks = self._transform_system_message(messages, model=model)
 
         # Convert last user message to guarded_text if guardrailConfig is present
-        messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
+        messages = self._convert_consecutive_user_messages_to_guarded_text(
+            self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
+        )
 
         _data: Final[CommonRequestObject] = self._transform_request_helper(
             model=model,
@@ -1902,7 +1971,7 @@ class AmazonConverseConfig(BaseConfig):
             return None
         tokens_5m: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m")
         tokens_1h: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h")
-        if tokens_5m + tokens_1h != usage.get("cacheWriteInputTokens", 0):
+        if tokens_5m + tokens_1h != AmazonConverseConfig._cache_write_count(usage):
             return None
         return CacheCreationTokenDetails(
             ephemeral_5m_input_tokens=tokens_5m,
@@ -1933,6 +2002,15 @@ class AmazonConverseConfig(BaseConfig):
                 return int(value)
         return 0
 
+    @staticmethod
+    def _cache_read_count(usage_object: Mapping[str, object]) -> int:
+        """Converse reports ``cacheReadInputTokens``; InvokeModel reports ``cacheReadInputTokenCount``."""
+        return AmazonConverseConfig._usage_count(usage_object, "cacheReadInputTokens", "cacheReadInputTokenCount")
+
+    @staticmethod
+    def _cache_write_count(usage_object: Mapping[str, object]) -> int:
+        return AmazonConverseConfig._usage_count(usage_object, "cacheWriteInputTokens", "cacheWriteInputTokenCount")
+
     def usage_from_batch_output(self, usage_object: Mapping[str, object]) -> Usage:
         """Read a Converse-shaped usage block out of a batch output line.
 
@@ -1942,8 +2020,8 @@ class AmazonConverseConfig(BaseConfig):
         """
         input_tokens: Final = self._usage_count(usage_object, "inputTokens")
         output_tokens: Final = self._usage_count(usage_object, "outputTokens")
-        cache_read: Final = self._usage_count(usage_object, "cacheReadInputTokens", "cacheReadInputTokenCount")
-        cache_write: Final = self._usage_count(usage_object, "cacheWriteInputTokens", "cacheWriteInputTokenCount")
+        cache_read: Final = self._cache_read_count(usage_object)
+        cache_write: Final = self._cache_write_count(usage_object)
         return self.transform_usage(
             ConverseTokenUsageBlock(
                 inputTokens=input_tokens,
@@ -1963,19 +2041,12 @@ class AmazonConverseConfig(BaseConfig):
         thinking_ran: bool = False,
         provider_reasoning_tokens: int | None = None,
     ) -> Usage:
-        input_tokens = usage["inputTokens"]
+        raw_input_tokens: Final = usage["inputTokens"]
         output_tokens: Final = usage["outputTokens"]
-        total_tokens: Final = usage["totalTokens"]
-        cache_creation_input_tokens: int = 0
-        cache_read_input_tokens: int = 0
-
-        raw_input_tokens: Final = input_tokens  # capture before inflation
-        if "cacheReadInputTokens" in usage:
-            cache_read_input_tokens = usage["cacheReadInputTokens"]
-            input_tokens += cache_read_input_tokens
-        if "cacheWriteInputTokens" in usage:
-            cache_creation_input_tokens = usage["cacheWriteInputTokens"]
-            input_tokens += cache_creation_input_tokens
+        cache_read_input_tokens: Final = self._cache_read_count(usage)
+        cache_creation_input_tokens: Final = self._cache_write_count(usage)
+        input_tokens: Final = raw_input_tokens + cache_read_input_tokens + cache_creation_input_tokens
+        total_tokens: Final = usage.get("totalTokens", input_tokens + output_tokens)
 
         prompt_tokens_details: Final = PromptTokensDetailsWrapper(
             cached_tokens=cache_read_input_tokens,

litellm/llms/bedrock/chat/invoke_handler.py

@@ -3,6 +3,7 @@ from collections.abc import AsyncIterator, Iterator
 from typing import Final, cast
 
 import httpx
+from pydantic import TypeAdapter
 
 import litellm
 from litellm import verbose_logger
@@ -51,6 +52,15 @@ bedrock_tool_name_mappings: Final[InMemoryCache] = InMemoryCache(max_size_in_mem
 from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
 
 converse_config: Final = AmazonConverseConfig()
+NOVA_INVOKE_STREAM_EVENT_TYPES: Final = (
+    "messageStart",
+    "contentBlockStart",
+    "contentBlockDelta",
+    "contentBlockStop",
+    "messageStop",
+    "metadata",
+)
+NOVA_INVOKE_STREAM_EVENT_PAYLOAD: Final = TypeAdapter(dict[str, object])
 
 
 class AmazonCohereChatConfig:
@@ -601,14 +611,12 @@ class AWSEventStreamDecoder:
             if thinking_blocks:
                 self._thinking_ran = True
 
-            carries_message_content: Final = any(
-                key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason", "trace")
+            trace: Final = chunk_data.get("trace")
+            carries_message_content: Final = bool(trace) or any(
+                key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason")
             )
 
-            model_response_provider_specific_fields: Final = {}
-            if "trace" in chunk_data:
-                trace: Final = chunk_data.get("trace")
-                model_response_provider_specific_fields["trace"] = trace
+            model_response_provider_specific_fields: Final = {"trace": trace} if trace else {}
             response: Final = ModelResponseStream(
                 choices=[
                     StreamingChoices(
@@ -654,10 +662,10 @@ class AWSEventStreamDecoder:
         ):
             return self.converse_chunk_parser(chunk_data=chunk_data)
         ######### /bedrock/invoke nova mappings ###############
-        elif "contentBlockDelta" in chunk_data:
-            # when using /bedrock/invoke/nova, the chunk_data is nested under "contentBlockDelta"
-            _chunk_data: Final = chunk_data.get("contentBlockDelta", {})
-            return self.converse_chunk_parser(chunk_data=_chunk_data)
+        elif nova_event_type := next((key for key in NOVA_INVOKE_STREAM_EVENT_TYPES if key in chunk_data), None):
+            return self.converse_chunk_parser(
+                chunk_data=NOVA_INVOKE_STREAM_EVENT_PAYLOAD.validate_python(chunk_data[nova_event_type])
+            )
         ######## bedrock.mistral mappings ###############
         elif "outputs" in chunk_data:
             if len(chunk_data["outputs"]) == 1 and chunk_data["outputs"][0].get("text", None) is not None:

litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py

@@ -6,12 +6,21 @@ Inherits from `AmazonConverseConfig`
 Nova + Invoke API Tutorial: https://docs.aws.amazon.com/nova/latest/userguide/using-invoke-api.html
 """
 
-from typing import TYPE_CHECKING, Final
+from collections.abc import Callable, Mapping, Sequence
+from functools import reduce
+from typing import TYPE_CHECKING, Final, TypeVar
 
 import httpx
+from pydantic import TypeAdapter, ValidationError
 
 from litellm.litellm_core_utils.litellm_logging import Logging
-from litellm.types.llms.bedrock import BedrockInvokeNovaRequest
+from litellm.types.llms.bedrock import (
+    BedrockInvokeNovaRequest,
+    CachePointBlock,
+    ContentBlock,
+    MessageBlock,
+    SystemContentBlock,
+)
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import ModelResponse
 
@@ -21,6 +30,50 @@ from .base_invoke_transformation import AmazonInvokeConfig
 if TYPE_CHECKING:
     import tiktoken
 
+_CachePointCarrier = TypeVar("_CachePointCarrier", SystemContentBlock, ContentBlock)
+_INJECTION_POINTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
+
+
+def _without_tool_config_injection_points(optional_params: Mapping[str, object]) -> dict[str, object]:
+    """InvokeModel has no tool caching, and a ``tool_config`` point the Converse transform
+    placed would credit the gateway for a cachePoint this request cannot carry.
+    """
+    raw_points: Final = optional_params.get("cache_control_injection_points")
+    if raw_points is None:
+        return dict(optional_params)
+    try:
+        points = _INJECTION_POINTS.validate_python(raw_points)
+    except ValidationError:
+        return dict(optional_params)
+    return {
+        **optional_params,
+        "cache_control_injection_points": [point for point in points if point.get("location") != "tool_config"],
+    }
+
+
+def _system_block_with_cache_point(block: SystemContentBlock, cache_point: CachePointBlock) -> SystemContentBlock:
+    return {**block, "cachePoint": cache_point}
+
+
+def _content_block_with_cache_point(block: ContentBlock, cache_point: CachePointBlock) -> ContentBlock:
+    return {**block, "cachePoint": cache_point}
+
+
+def _inline_block_cache_points(
+    blocks: Sequence[_CachePointCarrier],
+    with_cache_point: Callable[[_CachePointCarrier, CachePointBlock], _CachePointCarrier],
+) -> list[_CachePointCarrier]:
+    def attach(inlined: tuple[_CachePointCarrier, ...], block: _CachePointCarrier) -> tuple[_CachePointCarrier, ...]:
+        cache_point: Final = block.get("cachePoint")
+        if cache_point is None or len(block) != 1:
+            return (*inlined, block)
+        anchor: Final = next((index for index in reversed(range(len(inlined))) if "text" in inlined[index]), None)
+        if anchor is None:
+            return inlined
+        return (*inlined[:anchor], with_cache_point(inlined[anchor], cache_point), *inlined[anchor + 1 :])
+
+    return list(reduce(attach, blocks, ()))
+
 
 class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
     """
@@ -46,7 +99,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
         self,
         model: str,
         messages: list[AllMessageValues],
-        optional_params: dict,
+        optional_params: dict[str, object],
         litellm_params: dict,
         headers: dict,
     ) -> dict:
@@ -54,11 +107,13 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
             self,
             model=model,
             messages=messages,
-            optional_params=optional_params,
+            optional_params=_without_tool_config_injection_points(optional_params),
             litellm_params=litellm_params,
             headers=headers,
         )
-        _bedrock_invoke_nova_request: Final = BedrockInvokeNovaRequest(**_transformed_nova_request)
+        _bedrock_invoke_nova_request: Final = self._inline_cache_points(
+            BedrockInvokeNovaRequest(**_transformed_nova_request)
+        )
         self._remove_empty_system_messages(_bedrock_invoke_nova_request)
         bedrock_invoke_nova_request: Final = self._filter_allowed_fields(_bedrock_invoke_nova_request)
         return bedrock_invoke_nova_request
@@ -92,6 +147,24 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
             json_mode,
         )
 
+    @staticmethod
+    def _inline_cache_points(request: BedrockInvokeNovaRequest) -> BedrockInvokeNovaRequest:
+        """InvokeModel takes ``cachePoint`` as a key of the text block it caches: it rejects the
+        standalone ``{"cachePoint": ...}`` blocks Converse accepts and the key on image, toolUse,
+        and toolResult blocks, so a point behind one of those moves back to the last text block.
+        """
+        return {
+            **request,
+            "system": _inline_block_cache_points(request.get("system", []), _system_block_with_cache_point),
+            "messages": [
+                MessageBlock(
+                    role=message["role"],
+                    content=_inline_block_cache_points(message["content"], _content_block_with_cache_point),
+                )
+                for message in request.get("messages", [])
+            ],
+        }
+
     def _filter_allowed_fields(self, bedrock_invoke_nova_request: BedrockInvokeNovaRequest) -> dict:
         """
         Filter out fields that are not allowed in the `BedrockInvokeNovaRequest` dataclass.

litellm/llms/bedrock/vector_stores/transformation.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
@@ -14,6 +15,7 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import (
     BedrockKBResponse,
     BedrockKBRetrievalConfiguration,
     BedrockKBRetrievalQuery,
+    BedrockKBUserContext,
 )
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.vector_stores import (
@@ -242,10 +244,29 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
             retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters
         if retrieval_config:
             request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config)
+        user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params)
+        if user_context is not None:
+            request_body["userContext"] = user_context
 
         litellm_logging_obj.model_call_details["query"] = query
         return url, request_body
 
+    @staticmethod
+    def _user_context(
+        extra_body: Mapping[str, object] | None, litellm_params: Mapping[str, object]
+    ) -> BedrockKBUserContext | None:
+        sources: Final = tuple(source for source in (extra_body, litellm_params) if isinstance(source, Mapping))
+        found: Final = next(
+            (
+                source[key]
+                for source in sources
+                for key in ("userContext", "user_context")
+                if source.get(key) is not None
+            ),
+            None,
+        )
+        return None if found is None else cast(BedrockKBUserContext, found)
+
     def sign_request(
         self,
         headers: dict,

litellm/llms/bedrock_mantle/chat/transformation.py

@@ -22,6 +22,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
     BEDROCK_MANTLE_DEFAULT_REGION,
     BedrockMantleAuthMixin,
 )
+from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.router import GenericLiteLLMParams
@@ -108,13 +109,22 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
 
     def get_supported_openai_params(self, model: str) -> list:
         base_params: Final = super().get_supported_openai_params(model)
+        extra_params: Final = tuple(
+            param
+            for param, supported in (
+                ("verbosity", is_gpt_reasoning_series_name(model)),
+                ("reasoning_effort", self._supports_reasoning(model)),
+            )
+            if supported and param not in base_params
+        )
+        return [*base_params, *extra_params]  # mutable-ok: fresh list required by the inherited signature
+
+    def _supports_reasoning(self, model: str) -> bool:
         try:
-            if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
-                if "reasoning_effort" not in base_params:
-                    base_params.append("reasoning_effort")
+            return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
         except Exception as e:
             verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e)
-        return base_params
+            return False
 
     def get_model_response_iterator(
         self,

litellm/llms/custom_httpx/http_handler.py

@@ -7,6 +7,7 @@ import ssl
 import sys
 import threading
 import time
+import weakref
 from collections.abc import AsyncIterable, Callable, Iterable, Mapping
 from http.cookiejar import CookieJar, DefaultCookiePolicy
 from io import BytesIO
@@ -185,6 +186,33 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool:
     return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER
 
 
+def _drop_streaming_anchor(_handler: object) -> None:
+    """Release a handler anchored to a streaming response. See ``_anchor_handler_to``.
+
+    The work is the reference held until this point, so there is nothing to do here.
+    """
+
+
+def _anchor_handler_to(response: httpx.Response, handler: object) -> None:
+    """Keep the handler alive for as long as a streaming response can still read.
+
+    A body still arriving reads through the handler's connection pool, and closing
+    the client tears that pool down. The refcount ``_handler_may_close_client``
+    reads cannot see that body: the reference graph runs response -> stream ->
+    connection and stops there, so a client carrying one looks exactly like an
+    unreferenced client, and the finalizer closes it mid-body.
+
+    ``weakref.finalize`` holds the handler in its own registry rather than on the
+    response, which matters twice. The handler stays out of the response's
+    reference cycle, so it is finalized by refcount once the anchor drops and can
+    still schedule an async close, instead of being finalized inside a cyclic
+    collection that reaps its aiohttp session in the same pass. And a handler
+    serving several streams collects only once every one of them is done, because
+    each anchor holds it separately.
+    """
+    weakref.finalize(response, _drop_streaming_anchor, handler)
+
+
 def blocked_cookie_jar() -> CookieJar:
     """A jar that stores no response cookie and sends none, for httpx clients.
 
@@ -778,6 +806,8 @@ class AsyncHTTPHandler:
                 content=request_content,
             )
             response: Final = await self.client.send(req, stream=stream)
+            if stream:
+                _anchor_handler_to(response, self)
             response.raise_for_status()
             return response
         except (httpx.RemoteProtocolError, httpx.ConnectError):
@@ -982,6 +1012,8 @@ class AsyncHTTPHandler:
                 content=request_content,
             )
             response: Final = await self.client.send(req, stream=stream)
+            if stream:
+                _anchor_handler_to(response, self)
             response.raise_for_status()
             return response
         except (httpx.RemoteProtocolError, httpx.ConnectError):
@@ -1451,6 +1483,8 @@ class HTTPHandler:
                     content=request_content,
                 )
             response: Final = self.client.send(req, stream=stream)
+            if stream:
+                _anchor_handler_to(response, self)
             response.raise_for_status()
             return response
         except httpx.TimeoutException:
@@ -1501,6 +1535,8 @@ class HTTPHandler:
                     content=request_content,
                 )
             response: Final = self.client.send(req, stream=stream)
+            if stream:
+                _anchor_handler_to(response, self)
             response.raise_for_status()
             return response
         except httpx.TimeoutException:
@@ -1551,6 +1587,8 @@ class HTTPHandler:
                     content=request_content,
                 )
             response: Final = self.client.send(req, stream=stream)
+            if stream:
+                _anchor_handler_to(response, self)
             return response
         except httpx.TimeoutException:
             raise litellm.Timeout(
@@ -1600,6 +1638,8 @@ class HTTPHandler:
                     content=request_content,
                 )
             response: Final = self.client.send(req, stream=stream)
+            if stream:
+                _anchor_handler_to(response, self)
             response.raise_for_status()
             return response
         except httpx.TimeoutException:

litellm/llms/dashscope/chat/transformation.py

@@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
 
 
 class DashScopeChatConfig(OpenAIGPTConfig):
+    def get_supported_openai_params(self, model: str) -> list[str]:  # mutable-ok: base class contract returns a list
+        return [  # mutable-ok: base class contract returns a list
+            *super().get_supported_openai_params(model=model),
+            "reasoning_effort",
+        ]
+
     def remove_cache_control_flag_from_messages_and_tools(
         self,
         model: str,

litellm/llms/fireworks_ai/chat/transformation.py

@@ -49,6 +49,17 @@ if TYPE_CHECKING:
     import tiktoken
 
 
+def _map_reasoning_effort(value: object) -> object:
+    effort: Final[object] = cast(Mapping[str, object], value).get("effort") if isinstance(value, Mapping) else value
+    if effort is True:
+        return "medium"
+    if effort is False:
+        return "none"
+    if effort == "auto":
+        return None
+    return effort
+
+
 def _extract_fireworks_hidden_params(payload: dict) -> dict:
     """
     Collect Fireworks-specific response fields (perf_metrics, prompt_token_ids,
@@ -327,12 +338,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
             elif param == "max_completion_tokens":
                 optional_params["max_tokens"] = value
             elif param == "reasoning_effort":
-                if value is True:
-                    optional_params["reasoning_effort"] = "medium"
-                elif value is False:
-                    optional_params["reasoning_effort"] = "none"
-                elif value != "auto":
-                    optional_params["reasoning_effort"] = value
+                effort = _map_reasoning_effort(value)
+                if effort is not None:
+                    optional_params["reasoning_effort"] = effort
             elif param in supported_openai_params:
                 if value is not None:
                     optional_params[param] = value

litellm/llms/openai/chat/gpt_5_transformation.py

@@ -4,9 +4,9 @@ from typing import Final
 
 import litellm
 from litellm.utils import (
-    _is_explicitly_disabled_factory,
     _supports_factory,
     declared_value_factory,
+    is_explicitly_disabled_factory,
 )
 
 from .gpt_transformation import OpenAIGPTConfig
@@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
 
         Use this for opt-out checks where unknown models should be allowed through.
         """
-        return _is_explicitly_disabled_factory(
+        return is_explicitly_disabled_factory(
             model=cls._model_map_lookup_name(model),
             custom_llm_provider=None,
             key=f"supports_{level}_reasoning_effort",

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

@@ -792,10 +792,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
     def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
         chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
         stream_ended: Final = self._first_choice_has_finished(responses_so_far)
+        tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far)
         return StreamingScanKey(
             texts=tuple(self._combine_streaming_texts(chunks).values()),
-            tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
+            tool_calls=tool_call_fingerprints if stream_ended else (),
             stream_ended=stream_ended,
+            tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended,
         )
 
     @staticmethod
@@ -804,7 +806,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             stream_item_fingerprint(tool_call)
             for chunk in responses_so_far
             for choice in _stream_chunk_choices(chunk)
-            for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls")
+            for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta"))
         )
 
     @staticmethod
@@ -1342,6 +1344,12 @@ def _stream_chunk_choices(item: object) -> Sequence[object]:
     return ()
 
 
+def _streamed_delta_tool_calls(delta: object) -> tuple[object, ...]:
+    function_call: Final = stream_item_field(delta, "function_call")
+    legacy: Final = () if function_call is None else (function_call,)
+    return stream_item_items(delta, "tool_calls") + legacy
+
+
 def _blocked_stream_identity(
     exc: "ModifyResponseException", responses_so_far: Sequence[object]
 ) -> tuple[str, int, str]:

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -1175,11 +1175,22 @@ class OpenAIResponsesHandler(BaseTranslation):
         last_event_type: Final = stream_item_field(last_event, "type")
         if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
             return None
-        if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
+        if last_event_type in _TERMINAL_ENVELOPE_EVENT_TYPES:
             return self._completed_response_scan_key(stream_item_field(last_event, "response"))
         return StreamingScanKey(
             texts=(self.get_streaming_string_so_far(responses_so_far),),
-            stream_ended=self._check_streaming_has_ended(responses_so_far),
+            tool_calls_in_flight=self._has_streamed_tool_call_events(responses_so_far),
+        )
+
+    @staticmethod
+    def _has_streamed_tool_call_events(responses_so_far: Sequence[object]) -> bool:
+        return any(
+            stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
+            or (
+                stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
+                and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES
+            )
+            for event in responses_so_far
         )
 
     @staticmethod

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -79,6 +79,7 @@ from litellm.utils import (
     CustomStreamWrapper,
     ModelResponse,
     is_base64_encoded,
+    is_explicitly_disabled_factory,
     supports_reasoning,
 )
 
@@ -866,6 +867,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         else:
             raise _unsupported_reasoning_effort(reasoning_effort)
 
+    @staticmethod
+    def _supports_minimal_thinking_level(model: str) -> bool:
+        lowered: Final = model.lower()
+        is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered
+        return is_gemini3flash and not is_explicitly_disabled_factory(
+            model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort"
+        )
+
     @staticmethod
     def _map_reasoning_effort_to_thinking_level(
         reasoning_effort: str,
@@ -880,13 +889,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         Returns:
             GeminiThinkingConfig with thinkingLevel and includeThoughts
         """
-        # Check if this is gemini-3-flash which supports MINIMAL thinking level
-        # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview,
-        # gemini-3.5-flash, and any future 3.x-flash variants.
         is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower())
+        supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
         is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower())
         if reasoning_effort == "minimal":
-            if is_gemini3flash:
+            if supports_minimal:
                 return {"thinkingLevel": "minimal", "includeThoughts": True}
             else:
                 return {"thinkingLevel": "low", "includeThoughts": True}
@@ -899,18 +906,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                 return {"thinkingLevel": "high", "includeThoughts": True}
         elif reasoning_effort == "high":
             return {"thinkingLevel": "high", "includeThoughts": True}
-        elif reasoning_effort == "disable":
-            # Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others
-            if is_gemini3flash:
-                return {"thinkingLevel": "minimal", "includeThoughts": False}
-            else:
-                return {"thinkingLevel": "low", "includeThoughts": False}
-        elif reasoning_effort == "none":
-            # For gemini-3-flash-preview, use "minimal" instead of "low"
-            if is_gemini3flash:
-                return {"thinkingLevel": "minimal", "includeThoughts": False}
-            else:
-                return {"thinkingLevel": "low", "includeThoughts": False}
+        elif reasoning_effort in ("disable", "none"):
+            return {
+                "thinkingLevel": "minimal" if supports_minimal else "low",
+                "includeThoughts": False,
+            }
         else:
             raise _unsupported_reasoning_effort(reasoning_effort)
 
@@ -977,8 +977,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                     params["includeThoughts"] = True
                     # Follow provider defaults unless explicitly opted into legacy behavior.
                     if litellm.enable_gemini_default_thinking_level_low is True:
-                        is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower()
-                        params["thinkingLevel"] = "minimal" if is_gemini3flash else "low"
+                        params["thinkingLevel"] = (
+                            "minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low"
+                        )
             else:
                 # Thinking disabled
                 params["includeThoughts"] = False

litellm/models/verification_token.py

@@ -18,6 +18,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
     key_name: str | None = None
     key_alias: str | None = None
     spend: float = 0.0
+    total_spend: float = 0.0
     max_budget: float | None = None
     expires: str | datetime | None = None
     models: list = []
@@ -69,6 +70,7 @@ class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken):
     """Audit record for deleted keys; mirrors the token plus deletion metadata."""
 
     id: str | None = None
+    organization_id: str | None = None
     deleted_at: datetime | None = None
     deleted_by: str | None = None
     deleted_by_api_key: str | None = None

litellm/ocr/input.py

@@ -1,112 +0,0 @@
-from collections.abc import Mapping
-from os import PathLike
-from typing import Final, Literal, Protocol, cast  # noqa: TID251  # native callables are validated when loaded
-
-from typing_extensions import NotRequired, ReadOnly, TypedDict
-
-from litellm.rust_bridge.bindings import NativeBinding
-from litellm.rust_bridge.configuration import rust_ocr_enabled
-
-
-class FileReader(Protocol):
-    def read(self) -> bytes | str: ...
-
-
-class FileDocument(TypedDict):
-    type: ReadOnly[Literal["file"]]
-    file: ReadOnly[bytes | PathLike[str] | FileReader]
-    mime_type: ReadOnly[NotRequired[str]]
-
-
-class NativeFileDocument(Protocol):
-    def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ...
-
-
-class NativeUploadDocument(Protocol):
-    def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ...
-
-
-class NativeMimeType(Protocol):
-    def __call__(self, file_name: str) -> str: ...
-
-
-_FILE_DOCUMENT: Final = NativeBinding(
-    "_ocr_file_document",
-    validate=lambda value: (
-        cast(  # cast-ok: native export owns the callable signature
-            NativeFileDocument, value
-        )
-        if callable(value)
-        else None
-    ),
-)
-_UPLOAD_DOCUMENT: Final = NativeBinding(
-    "_ocr_upload_document",
-    validate=lambda value: (
-        cast(  # cast-ok: native export owns the callable signature
-            NativeUploadDocument, value
-        )
-        if callable(value)
-        else None
-    ),
-)
-_MAX_FILE_BYTES: Final = NativeBinding(
-    "_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None
-)
-_MIME_TYPE: Final = NativeBinding(
-    "_ocr_mime_type",
-    validate=lambda value: (
-        cast(  # cast-ok: native export owns the callable signature
-            NativeMimeType, value
-        )
-        if callable(value)
-        else None
-    ),
-)
-_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
-
-
-def get_mime_type(file_path: str) -> str:
-    native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None
-    if native is None:
-        from litellm.ocr import legacy
-
-        return legacy.get_mime_type(file_path)
-    return native(file_path)
-
-
-def get_max_file_bytes() -> int:
-    limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None
-    if limit is None:
-        return _PYTHON_MAX_FILE_BYTES
-    return limit
-
-
-def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]:
-    native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None
-    if native is None:
-        from litellm.ocr import legacy
-
-        return legacy.convert_file_document_to_url_document(document)
-    return native(document)
-
-
-def convert_upload_to_url_document(
-    file_content: bytes, filename: str | None, content_type: str | None
-) -> dict[str, str]:
-    native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None
-    if native is None:
-        from litellm.ocr import legacy
-
-        if len(file_content) > _PYTHON_MAX_FILE_BYTES:
-            raise ValueError("OCR file exceeds the size limit")
-        content_mime: Final = content_type.split(";")[0].strip() if content_type else None
-        mime_type: Final = (
-            legacy.get_mime_type(filename)
-            if filename and (not content_mime or content_mime == "application/octet-stream")
-            else content_mime or "application/octet-stream"
-        )
-        return legacy.convert_file_document_to_url_document(
-            {"type": "file", "file": file_content, "mime_type": mime_type}
-        )
-    return native(file_content, filename, content_type)

litellm/ocr/legacy.py

@@ -11,7 +11,7 @@ from collections.abc import Coroutine, Mapping
 from dataclasses import dataclass
 from io import IOBase
 from types import MappingProxyType
-from typing import Final, cast  # noqa: TID251  # adapters preserve the legacy untyped contracts
+from typing import Final, Protocol, cast  # noqa: TID251  # adapters preserve the legacy untyped contracts
 
 import httpx
 
@@ -26,7 +26,6 @@ from litellm.llms.base_llm.ocr.transformation import (
     parse_ocr_request_format,
 )
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from litellm.ocr.input import FileReader
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import CustomPricingLiteLLMParams
 from litellm.utils import ProviderConfigManager, client
@@ -34,6 +33,10 @@ from litellm.utils import ProviderConfigManager, client
 base_llm_http_handler: Final = BaseLLMHTTPHandler()
 
 
+class FileReader(Protocol):
+    def read(self) -> bytes | str: ...
+
+
 @dataclass(frozen=True, slots=True)
 class _PreparedOCRRequest:
     model: str

litellm/ocr/main.py

@@ -5,7 +5,7 @@ import httpx
 
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
 from litellm.ocr import legacy
-from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
+from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type
 from litellm.rust_bridge.bindings import native_exception_types
 from litellm.rust_bridge.configuration import rust_ocr_enabled
 from litellm.rust_bridge.ocr import LiteLLMOcrRequest

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
     UpstreamCredentialProvider,
 )
 from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
+    prepare_mcp_client,
     raise_public,
     raise_token_exchange_challenge,
     raise_user_oauth_challenge,
@@ -2804,6 +2805,8 @@ class MCPServerManager:
                         headers=headers,
                         server_label=server.name or server.server_name or server.alias or server.server_id,
                         relays_upstream_auth=server.is_client_forwarded_token,
+                        auth_type=server.auth_type,
+                        upstream_token_header=server.upstream_token_header,
                     )
                     tool_func.__name__ = prefixed_tool_name
                     tool_func.__doc__ = description
@@ -4259,15 +4262,20 @@ class MCPServerManager:
                     user_api_key_auth=user_api_key_auth,
                     extra_headers=extra_headers,
                 )
-                return MCPClient(
-                    server_url=server_url,
-                    transport_type=transport,
-                    auth_type=resolved_server.auth_type,
-                    timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
-                    extra_headers=extra_headers,
-                    resolved_auth=resolved_auth,
-                    sampling_callback=sampling_cb,
-                    elicitation_callback=elicitation_cb,
+                return await prepare_mcp_client(
+                    resolved_server,
+                    MCPClient(
+                        server_url=server_url,
+                        transport_type=transport,
+                        auth_type=resolved_server.auth_type,
+                        timeout=(
+                            resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT
+                        ),
+                        extra_headers=extra_headers,
+                        resolved_auth=resolved_auth,
+                        sampling_callback=sampling_cb,
+                        elicitation_callback=elicitation_cb,
+                    ),
                 )
 
             # Create SigV4 auth if configured
@@ -4297,17 +4305,20 @@ class MCPServerManager:
                 else AuthResolution.no_auth
             )
             record_auth_resolution(server.server_id, legacy_source)
-            return MCPClient(
-                server_url=server_url,
-                transport_type=transport,
-                auth_type=resolved_server.auth_type,
-                auth_value=auth_value,
-                auth_header_name=auth_header_name,
-                timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
-                extra_headers=extra_headers,
-                aws_auth=aws_auth,
-                sampling_callback=sampling_cb,
-                elicitation_callback=elicitation_cb,
+            return await prepare_mcp_client(
+                resolved_server,
+                MCPClient(
+                    server_url=server_url,
+                    transport_type=transport,
+                    auth_type=resolved_server.auth_type,
+                    auth_value=auth_value,
+                    auth_header_name=auth_header_name,
+                    timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
+                    extra_headers=extra_headers,
+                    aws_auth=aws_auth,
+                    sampling_callback=sampling_cb,
+                    elicitation_callback=elicitation_cb,
+                ),
             )
 
     async def _get_tools_from_server(

litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py

@@ -54,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.proxy._experimental.mcp_server.tool_registry import (
     global_mcp_tool_registry,
 )
-from litellm.types.mcp import credential_redirect_hook, custom_credential_slot
+from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot
 
 
 class _OpenAPIJSONSchema(TypedDict, total=False):
@@ -471,6 +471,8 @@ def create_tool_function(
     headers: dict[str, str] | None = None,
     server_label: str | None = None,
     relays_upstream_auth: bool = False,
+    auth_type: MCPAuthType = None,
+    upstream_token_header: str | None = None,
 ):
     """Create a tool function for an OpenAPI operation.
 
@@ -503,6 +505,18 @@ def create_tool_function(
         by using **kwargs instead of named parameters.
         """
         effective_headers: Final = _merge_openapi_tool_request_headers(headers)
+        if auth_type is not None:
+            from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
+                raise_public,
+                validate_static_credential,
+            )
+            from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
+
+            match validate_static_credential(auth_type, effective_headers, upstream_token_header, headers or ()):
+                case Error(error):
+                    raise_public(error)
+                case Ok():
+                    pass
 
         # Build URL from base_url and path
         url = base_url + path

litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py

@@ -13,15 +13,17 @@ from __future__ import annotations
 
 import base64
 import os
+from collections.abc import Iterable, Mapping
 from typing import TYPE_CHECKING, Final, Literal, NoReturn
 
 from fastapi import HTTPException
 from pydantic import SecretStr
 from typing_extensions import assert_never
 
-from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
+from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
 from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
 from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
+from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
 from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
     DEFAULT_CREDENTIAL_HEADER,
     ApiKeyConfig,
@@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
     Subject,
     TokenExchangeConfig,
 )
-from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth
+from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport
 
 if TYPE_CHECKING:
     from litellm.proxy._types import UserAPIKeyAuth
@@ -79,7 +81,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
 
     BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
     like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
-    to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).
+    to v1 for its static schemes. Declared OBO always stays with the exchange arm.
 
     Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
     an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
@@ -90,8 +92,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
     modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
     oauth2 and SigV4 return None and stay on v1.
     """
-    if server.is_byok:
-        return None  # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
+    if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange:
+        return None  # per-user BYOK source not migrated yet -> defer to v1
     resource: Final = server.url or server.server_id
     auth_type: Final = server.auth_type
     match auth_type:
@@ -165,21 +167,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
     )
 
 
-def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
-    """Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.
-
-    An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
-    ``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
-    the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
-    gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
-    nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
-    (``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
-    normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
-    forwarded only when the operator set it; a missing one is omitted, not derived.
-    """
+def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec:
+    """Keep declared OBO owned by the resolver, including incomplete client configuration."""
     endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
-    if not server.client_id or not server.client_secret:
-        return None
     profile: Final[Literal["rfc8693", "entra_obo"]] = (
         "entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
     )
@@ -193,7 +183,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
             token_exchange_endpoint=endpoint,
             audience=server.audience,
             client_id=server.client_id,
-            client_secret=SecretStr(server.client_secret),
+            client_secret=SecretStr(server.client_secret) if server.client_secret else None,
             token_endpoint_auth_method=server.token_endpoint_auth_method,
             scopes=tuple(server.scopes or ()),
         ),
@@ -397,3 +387,74 @@ def raise_token_exchange_challenge(
         detail="Unauthorized",
         headers={"WWW-Authenticate": www_authenticate},
     )
+
+
+_STATIC_MODES: Final = frozenset(
+    (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization)
+)
+
+
+def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool:
+    if not value:
+        return False
+    if auth_type == MCPAuth.api_key and name != "authorization":
+        return True
+    if value.lower() in ("bearer", "basic", "token", "apikey"):
+        return False
+    if auth_type == MCPAuth.api_key:
+        api_scheme: Final = value.split(None, 1)[0]
+        if api_scheme.lower() in ("bearer", "token", "apikey"):
+            api_credential: Final = strip_auth_scheme(value, api_scheme).strip()
+            return api_credential.lower() != api_scheme.lower()
+    if auth_type in (MCPAuth.bearer_token, MCPAuth.token):
+        scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token"
+        credential: Final = strip_auth_scheme(value, scheme).strip()
+        return bool(credential) and credential.lower() != scheme.lower()
+    if auth_type == MCPAuth.basic:
+        parts: Final = value.split(None, 1)
+        if len(parts) != 2 or parts[0].lower() != "basic":
+            return False
+        try:
+            decoded: Final = base64.b64decode(parts[1], validate=True).strip()
+            return b":" in decoded
+        except ValueError:
+            return False
+    return True
+
+
+def validate_static_credential(
+    auth_type: MCPAuthType,
+    headers: Mapping[str, str],
+    upstream_token_header: str | None = None,
+    static_header_names: Iterable[str] = (),
+) -> Result[None, CredError]:
+    if auth_type not in _STATIC_MODES:
+        return Ok(None)
+    default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization"
+    admin_chosen_slots: Final = tuple(static_header_names) if auth_type == MCPAuth.api_key else ()
+    slots: Final = frozenset(
+        name.lower()
+        for name in (
+            upstream_token_header or default_slot,
+            default_slot,
+            "Authorization",
+            *admin_chosen_slots,
+        )
+    )
+    values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
+    if any(_usable_credential_value(auth_type, name, value) for name, value in values):
+        return Ok(None)
+    return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
+
+
+async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
+    if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
+        return client
+    request: Final = await client.prepare_request_auth()
+    match validate_static_credential(
+        server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
+    ):
+        case Error(error):
+            raise_public(error)
+        case Ok():
+            return client

litellm/proxy/_types.py

@@ -845,6 +845,9 @@ class LiteLLMRoutes(enum.Enum):
     )
 
     self_managed_routes = [
+        # update_team resolves proxy/org/team admin itself and filters team admins
+        # through the team_admin_editable_team_fields setting
+        "/team/update",
         "/team/member_add",
         "/team/member_delete",
         "/management/v1/teams/{team_id}/members/bulk_delete",
@@ -2005,6 +2008,13 @@ RouterSettingsDict = Annotated[
 class NewTeamRequest(TeamBase):
     router_settings: RouterSettingsDict | None = None
     model_aliases: dict | None = None
+    model_max_budget: GenericBudgetConfigType | None = Field(
+        default=None,
+        description=(
+            "Max budget per model for every key on the team, overridable per key "
+            "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})"
+        ),
+    )
     tags: list | None = None
     guardrails: list[str] | None = None
     policies: list[str] | None = None
@@ -2106,6 +2116,13 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
     access_group_ids: list[str] | None = None
     budget_limits: list[BudgetLimitEntry] | None = None  # multiple concurrent budget windows
     default_team_member_models: list[str] | None = None  # default allowed_models seeded onto new team members
+    model_max_budget: GenericBudgetConfigType | None = Field(
+        default=None,
+        description=(
+            "Max budget per model for every key on the team, overridable per key "
+            "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})"
+        ),
+    )
 
 
 class PatchTeamRequest(UpdateTeamRequest):
@@ -3033,6 +3050,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
     team_tpd_limit: int | None = None
     team_max_budget: float | None = None
     team_soft_budget: float | None = None
+    team_model_max_budget: dict[str, object] | None = None
     team_models: list = []
     team_blocked: bool = False
     soft_budget: float | None = None
@@ -3711,6 +3729,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
             "AWS_ACCESS_KEY_ID",
             "AWS_SECRET_ACCESS_KEY",
             "AWS_REGION_NAME",
+            "S3_LOG_PROMPTS_ONLY",
         ],
     )
 
@@ -4452,6 +4471,29 @@ class TeamInfoMember(Member):
     user_alias: str | None = None
 
 
+class TeamEditUnrestricted(BaseModel):
+    kind: Literal["unrestricted"] = "unrestricted"
+
+
+class TeamEditAsTeamAdmin(BaseModel):
+    kind: Literal["team_admin"] = "team_admin"
+    editable_fields: tuple[str, ...]
+
+
+class TeamEditAsTeamAdminDisabled(BaseModel):
+    kind: Literal["team_admin_disabled"] = "team_admin_disabled"
+
+
+class TeamEditNone(BaseModel):
+    kind: Literal["none"] = "none"
+
+
+TeamEditAccess = Annotated[
+    TeamEditUnrestricted | TeamEditAsTeamAdmin | TeamEditAsTeamAdminDisabled | TeamEditNone,
+    Field(discriminator="kind"),
+]
+
+
 class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
     members_with_roles: tuple[TeamInfoMember, ...] = ()
     team_member_budget_table: LiteLLM_BudgetTableFull | None = None
@@ -4463,6 +4505,8 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
     # Parent org's model ceiling, reported only to callers who can manage the team.
     # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
     organization_models: list[str] | None = None
+    model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None
+    caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
 
 
 class TeamInfoResponseObject(TypedDict):

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from fastapi.responses import JSONResponse
 
 import litellm
-from litellm._logging import verbose_proxy_logger
 from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
 from litellm.integrations.custom_guardrail import ModifyResponseException
 from litellm.llms.anthropic.experimental_pass_through.context_management import (
@@ -22,13 +21,16 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_request_processing import (
     ProxyBaseLLMRequestProcessing,
     create_response,
+    log_llm_api_exception,
     proxy_exception_from_http_exception,
+    resolve_litellm_call_id,
 )
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
 from litellm.proxy.common_utils.openai_error_payload import (
     error_status_code,
     openai_error_param,
     openai_error_type,
+    with_litellm_call_id,
 )
 from litellm.types.utils import TokenCountResponse
 
@@ -218,10 +220,12 @@ async def anthropic_response(
         await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data
         )
-        verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
+        log_llm_api_exception(e, base_llm_response_processor.litellm_call_id)
 
         if isinstance(e, ProxyException):
-            return _anthropic_error_json_response(e, request)
+            return _anthropic_error_json_response(
+                with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request
+            )
 
         # Extract model_id from request metadata (same as success path)
         litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@@ -231,7 +235,7 @@ async def anthropic_response(
         # Get headers
         headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
             user_api_key_dict=user_api_key_dict,
-            call_id=data.get("litellm_call_id", ""),
+            call_id=base_llm_response_processor.litellm_call_id,
             model_id=model_id,
             version=version,
             response_cost=0,
@@ -288,6 +292,7 @@ async def count_tokens(
     """
     from litellm.proxy.proxy_server import token_counter as internal_token_counter
 
+    litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
     try:
         request_data: Final = await _read_request_body(request=request)
         data: Final[dict] = {**request_data}
@@ -339,7 +344,7 @@ async def count_tokens(
             detail=detail,
         )
     except Exception as e:
-        verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e)
+        log_llm_api_exception(e, litellm_call_id)
         raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})
 
 

litellm/proxy/auth/auth_checks.py

@@ -856,6 +856,16 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
 )
 
 
+def route_skips_budget_checks(route: str) -> bool:
+    return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
+        route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)
+    )
+
+
+def request_skips_budget_checks(route: str, model: str | list[str] | None, llm_router: Router | None) -> bool:
+    return route_skips_budget_checks(route=route) or _is_model_cost_zero(model=model, llm_router=llm_router)
+
+
 async def common_checks(
     request_body: dict,
     team_object: LiteLLM_TeamTable | None,
@@ -903,10 +913,7 @@ async def common_checks(
         team_id=valid_token.team_id if valid_token is not None else None,
     )
 
-    skip_all_budget_checks: Final = skip_budget_checks or (
-        route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES
-        and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
-    )
+    skip_all_budget_checks: Final = skip_budget_checks or route_skips_budget_checks(route=route)
 
     membership_user_id: Final = (
         valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
@@ -2102,7 +2109,7 @@ async def _fetch_uncached_tags(
 
 @log_db_metrics
 async def get_tag_objects_batch(
-    tag_names: list[str],
+    tag_names: Sequence[str],
     prisma_client: PrismaClient | None,
     user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Span | None = None,
@@ -5861,15 +5868,25 @@ async def _tag_max_budget_check(
     """
     from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
 
-    if prisma_client is None:
+    await tag_max_budget_check_for_tags(
+        tags=get_tags_from_request_body(request_body=request_body),
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+        valid_token=valid_token,
+    )
+
+
+async def tag_max_budget_check_for_tags(
+    tags: Sequence[str],
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    proxy_logging_obj: ProxyLogging,
+    valid_token: UserAPIKeyAuth | None,
+) -> None:
+    if prisma_client is None or not tags:
         return
 
-    # Get tags from request metadata
-    tags: Final = get_tags_from_request_body(request_body=request_body)
-    if not tags:
-        return
-
-    # Batch fetch all tags in one go
     tag_objects: Final = await get_tag_objects_batch(
         tag_names=tags,
         prisma_client=prisma_client,

litellm/proxy/auth/fallback_budget.py

@@ -0,0 +1,166 @@
+"""
+Enforce the caller's budget against router fallback targets.
+
+Budget is checked once, during auth, against the *requested* model group. A zero-cost group takes
+`_is_model_cost_zero`'s bypass and waives every budget check; the router then picks a fallback
+target after auth, inside `run_async_fallback`, and nothing re-checks budget on the group that
+actually bills. So a free model with a paid fallback spends without a gate.
+
+This predicate is injected into the router to re-check budget for each fallback target before it is
+attempted, mirroring `fallback_model_access.py`. It deliberately leaves the primary attempt alone:
+a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default;
+set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour.
+
+Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation
+path before it can be: team, team-member, end-user, org, global and per-model budgets, whose
+auth-path functions enforce rather than report (they raise), so reusing them would fire threshold
+alerts and take spend reservations for a target that is then skipped; and the key's rolling
+`budget_limits` windows, whose accumulated spend lives only in per-window counters
+(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the
+fallback path rather than reusing state auth already loaded.
+
+Two known limitations of that narrow scope, both shared with `fallback_model_access.py`:
+
+* This reads the spend counter, it does not reserve against it. Requests already in flight all
+  observe the same pre-billing figure, so a cap can be crossed by roughly the number of concurrent
+  fallbacks times their cost. Auth-time enforcement avoids this by pre-filling the counter through
+  `reserve_budget_for_request`, which the zero-cost bypass skips. Turning the soft cap into a hard
+  one means reserving per fallback attempt and reconciling on completion.
+* A request that reaches the router without `metadata["user_api_key_auth"]` is not restricted.
+  Only `add_litellm_data_to_request` populates that key, so endpoints that assemble metadata by
+  hand (for example `/queue/chat/completions`) fall through as unauthenticated.
+"""
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from pydantic import BaseModel, ValidationError
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.auth_checks import (
+    _is_model_cost_zero,  # pyright: ignore[reportPrivateUsage]  # the zero-cost predicate the auth-time budget checks use; no public equivalent
+)
+from litellm.router import Router
+
+
+class _RequestMetadata(BaseModel):
+    user_api_key_auth: UserAPIKeyAuth | None = None
+
+
+class _FallbackBudgetSettings(BaseModel):
+    enforce_fallback_budget: bool = True
+
+
+def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None:
+    try:
+        return _RequestMetadata.model_validate(metadata).user_api_key_auth
+    except ValidationError:
+        return None
+
+
+def _user_api_key_auth_from_request(request_kwargs: Mapping[str, object]) -> UserAPIKeyAuth | None:
+    return next(
+        (
+            token
+            for field in ("metadata", "litellm_metadata")
+            if (token := _token_in_metadata(request_kwargs.get(field))) is not None
+        ),
+        None,
+    )
+
+
+def _enforced_by_general_settings() -> bool:
+    from litellm.proxy.proxy_server import general_settings
+
+    return _FallbackBudgetSettings.model_validate(general_settings).enforce_fallback_budget
+
+
+def _applies_user_budget_to_team_keys() -> bool:
+    from litellm.proxy.proxy_server import general_settings
+
+    return general_settings.get("apply_user_budget_to_team_keys") is True
+
+
+async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: float) -> float:
+    """
+    Read a spend counter the same way the auth-time budget checks do.
+
+    `max_budget` is not advisory: it makes `get_current_spend` re-check the counter against the
+    authoritative recorded spend before admitting. A counter restored from an older Redis snapshot
+    reads as a hit rather than a clean miss, so without this the reseed path never runs and a
+    stale-low counter would keep admitting paid fallbacks past the cap.
+    """
+    from litellm.proxy.proxy_server import get_current_spend
+
+    return await get_current_spend(
+        counter_key=counter_key,
+        fallback_spend=fallback_spend,
+        max_budget=max_budget,
+    )
+
+
+async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool:
+    """
+    True when the key and the user behind it can still pay for `model`.
+
+    A zero-cost fallback target is always allowed: refusing it would deny a request on spend some
+    other model accrued, which is the same reasoning behind the auth-time bypass.
+    """
+    if _is_model_cost_zero(model=model, llm_router=llm_router):
+        return True
+
+    key_budget: Final = valid_token.max_budget
+    if key_budget is not None and valid_token.token is not None:
+        key_spend: Final = await _counter_spend(
+            counter_key=f"spend:key:{valid_token.token}",
+            fallback_spend=valid_token.spend or 0.0,
+            max_budget=key_budget,
+        )
+        if key_spend >= key_budget:
+            return False
+
+    # Mirrors `_PROXY_MaxBudgetLimiter`: a team key does not carry the key owner's personal budget
+    # unless the proxy opts in, so the personal cap must not gate the fallback either.
+    user_budget: Final = valid_token.user_max_budget
+    if (
+        user_budget is not None
+        and valid_token.user_id is not None
+        and (valid_token.team_id is None or _applies_user_budget_to_team_keys())
+    ):
+        user_spend: Final = await _counter_spend(
+            counter_key=f"spend:user:{valid_token.user_id}",
+            fallback_spend=valid_token.user_spend or 0.0,
+            max_budget=user_budget,
+        )
+        if user_spend >= user_budget:
+            return False
+
+    return True
+
+
+@dataclass(frozen=True, slots=True)
+class RouterFallbackBudgetCheck:
+    """
+    `FallbackBudgetCheck` for the proxy's router: while `is_enforced()` is true, a paid fallback
+    target is attempted only when the caller is still within budget. Requests that carry no key
+    (for example internal health checks) are not restricted.
+    """
+
+    is_enforced: Callable[[], bool]
+
+    async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
+        if not self.is_enforced():
+            return True
+        valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
+        if valid_token is None:
+            return True
+        try:
+            return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router)
+        except Exception as e:  # noqa: BLE001  # fail closed: a spend lookup failure must not bill the caller
+            verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
+            return False
+
+
+router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings)

litellm/proxy/auth/team_grants.py

@@ -59,6 +59,7 @@ class TeamGrants(TypedDict, total=False):
     team_tpd_limit: ReadOnly[int | None]
     team_max_budget: ReadOnly[float | None]
     team_soft_budget: ReadOnly[float | None]
+    team_model_max_budget: ReadOnly[dict[str, object] | None]
     team_spend: ReadOnly[float | None]
     team_models: ReadOnly[Sequence[str]]
     team_blocked: ReadOnly[bool]
@@ -101,6 +102,7 @@ def team_grants(
         team_tpd_limit=team_object.tpd_limit,
         team_max_budget=team_object.max_budget,
         team_soft_budget=team_object.soft_budget,
+        team_model_max_budget=team_object.model_max_budget,
         team_spend=team_object.spend,
         team_models=tuple(team_object.models),
         team_blocked=team_object.blocked,

litellm/proxy/auth/user_api_key_auth.py

@@ -304,6 +304,16 @@ class _UserModelBudgetLimiter(Protocol):
     ) -> bool: ...
 
 
+class _TeamModelBudgetLimiter(Protocol):
+    async def is_team_within_model_budget(
+        self,
+        team_id: str,
+        team_model_max_budget: Mapping[str, object],
+        key_model_max_budget: Mapping[str, object] | None,
+        model: str,
+    ) -> bool: ...
+
+
 class _TokenTeamModels(Protocol):
     @property
     def team_models(self) -> list[str]: ...
@@ -374,6 +384,25 @@ async def _check_user_model_budget(
         )
 
 
+async def _check_team_model_budget(
+    valid_token: UserAPIKeyAuth,
+    model_max_budget_limiter: _TeamModelBudgetLimiter,
+    models: list[str],
+) -> None:
+    """Enforce the team's `model_max_budget` for every requested model the key does not override."""
+    team_model_max_budget: Final = valid_token.team_model_max_budget
+    if valid_token.team_id is None or not team_model_max_budget:
+        return
+    key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget
+    for model_name in models:
+        await model_max_budget_limiter.is_team_within_model_budget(
+            team_id=valid_token.team_id,
+            team_model_max_budget=team_model_max_budget,
+            key_model_max_budget=key_model_max_budget,
+            model=model_name,
+        )
+
+
 async def _check_key_model_budget_with_fallback(
     valid_token: UserAPIKeyAuth,
     model_max_budget_limiter: _KeyModelBudgetLimiter,
@@ -2376,6 +2405,7 @@ async def _user_api_key_auth_builder(
                         team_id=valid_token.team_id,
                         max_budget=valid_token.team_max_budget,
                         soft_budget=valid_token.team_soft_budget,
+                        model_max_budget=valid_token.team_model_max_budget,
                         spend=valid_token.team_spend,
                         tpm_limit=valid_token.team_tpm_limit,
                         rpm_limit=valid_token.team_rpm_limit,
@@ -2530,6 +2560,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
         team_id=valid_token.team_id,
         max_budget=valid_token.team_max_budget,
         soft_budget=valid_token.team_soft_budget,
+        model_max_budget=valid_token.team_model_max_budget,
         spend=valid_token.team_spend,
         tpm_limit=valid_token.team_tpm_limit,
         rpm_limit=valid_token.team_rpm_limit,
@@ -2571,6 +2602,13 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc
     return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
 
 
+def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, object]) -> bool:
+    return master_key is None and not any(
+        general_settings.get(flag, False)
+        for flag in ("enable_jwt_auth", "enable_oauth2_auth", "enable_oauth2_proxy_auth")
+    )
+
+
 @tracer.wrap()
 async def _run_centralized_common_checks(
     user_api_key_auth_obj: UserAPIKeyAuth,
@@ -2599,6 +2637,7 @@ async def _run_centralized_common_checks(
         litellm_proxy_admin_name,
         llm_router,
         master_key,
+        model_max_budget_limiter,
         prisma_client,
         proxy_logging_obj,
         user_api_key_cache,
@@ -2630,11 +2669,7 @@ async def _run_centralized_common_checks(
     # Running common_checks would block every admin route on these
     # deployments where that was previously not the contract. If any
     # authn is enabled (JWT, OAuth2, OAuth2-proxy), authz must run.
-    if master_key is None and not (
-        general_settings.get("enable_jwt_auth", False)
-        or general_settings.get("enable_oauth2_auth", False)
-        or general_settings.get("enable_oauth2_proxy_auth", False)
-    ):
+    if is_no_auth_dev_mode(master_key, general_settings):
         return
 
     if user_custom_auth is not None and not general_settings.get("custom_auth_run_common_checks", False):
@@ -2871,6 +2906,21 @@ async def _run_centralized_common_checks(
     finally:
         release_spend_counter_batch()
 
+    if not skip_budget_checks:
+        await _check_team_model_budget(
+            valid_token=user_api_key_auth_obj,
+            model_max_budget_limiter=model_max_budget_limiter,
+            models=_get_model_names_for_budget_checks(
+                model=_get_model_from_request_context(
+                    request_data=request_data,
+                    route=route,
+                    request=request,
+                    llm_router=llm_router,
+                    team_id=user_api_key_auth_obj.team_id,
+                )
+            ),
+        )
+
     await _reserve_budget_after_common_checks(
         user_api_key_auth_obj=user_api_key_auth_obj,
         request=request,

litellm/proxy/batches_endpoints/endpoints.py

@@ -7,6 +7,7 @@
 import asyncio
 import os
 from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Any, Final, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
@@ -17,7 +18,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
-from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_request_processing import (
+    ProxyBaseLLMRequestProcessing,
+    log_llm_api_exception,
+    request_litellm_call_id,
+)
 from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
 from litellm.proxy.common_utils.openai_endpoint_utils import (
@@ -383,8 +388,9 @@ async def create_batch(
         await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
         )
-        verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
-        raise handle_exception_on_proxy(e)
+        litellm_call_id: Final = request_litellm_call_id(data)
+        log_llm_api_exception(e, litellm_call_id)
+        raise handle_exception_on_proxy(e, litellm_call_id)
 
 
 @router.get(
@@ -674,8 +680,9 @@ async def retrieve_batch(
         await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
         )
-        verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
-        raise handle_exception_on_proxy(e)
+        litellm_call_id: Final = request_litellm_call_id(data)
+        log_llm_api_exception(e, litellm_call_id)
+        raise handle_exception_on_proxy(e, litellm_call_id)
 
 
 @router.get(
@@ -725,6 +732,7 @@ async def list_batches(
     )
 
     verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit)
+    data: Mapping[str, object] = MappingProxyType({})
     try:
         if llm_router is None:
             raise HTTPException(
@@ -854,10 +862,11 @@ async def list_batches(
         await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict,
             original_exception=e,
-            request_data={"after": after, "limit": limit},
+            request_data={**data, "after": after, "limit": limit},
         )
-        verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
-        raise handle_exception_on_proxy(e)
+        litellm_call_id: Final = request_litellm_call_id(data)
+        log_llm_api_exception(e, litellm_call_id)
+        raise handle_exception_on_proxy(e, litellm_call_id)
 
 
 @router.post(
@@ -1079,8 +1088,9 @@ async def cancel_batch(
         await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
         )
-        verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
-        raise handle_exception_on_proxy(e)
+        litellm_call_id: Final = request_litellm_call_id(data)
+        log_llm_api_exception(e, litellm_call_id)
+        raise handle_exception_on_proxy(e, litellm_call_id)
 
 
 ######################################################################

litellm/proxy/common_request_processing.py

@@ -7,14 +7,25 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Literal,
+    NamedTuple,
+    Protocol,
+    TypeAlias,
+    TypeVar,
+    overload,
+    runtime_checkable,
+)
 
 import anyio
 import httpx
 import orjson
 from fastapi import HTTPException, Request, status
 from fastapi.responses import JSONResponse, Response, StreamingResponse
-from pydantic import ValidationError
+from pydantic import TypeAdapter, ValidationError
 from starlette.types import Receive, Scope, Send
 
 import litellm
@@ -34,7 +45,11 @@ from litellm.constants import (
     UNSAFE_PROXY_RESPONSE_HEADERS,
 )
 from litellm.integrations.custom_guardrail import CustomGuardrail
-from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error
+from litellm.litellm_core_utils.core_helpers import (
+    get_or_create_metadata_bucket,
+    independent_snapshot,
+    is_expected_client_error,
+)
 from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
 from litellm.litellm_core_utils.get_supported_openai_params import (
     get_supported_openai_params,
@@ -49,14 +64,21 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.streaming_handler import (
     backfill_missing_cache_usage_fields,
 )
-from litellm.proxy._types import ProxyException, UserAPIKeyAuth
-from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
-from litellm.proxy.auth.auth_utils import check_response_size_is_safe
+from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
+from litellm.proxy.auth.auth_checks import (
+    can_key_call_resolved_model,
+    request_skips_budget_checks,
+    tag_max_budget_check_for_tags,
+)
+from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route
 from litellm.proxy.common_utils.callback_utils import (
     get_logging_caching_headers,
     get_remaining_tokens_and_requests_from_request_data,
 )
-from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
+from litellm.proxy.common_utils.http_parsing_utils import (
+    get_client_requested_model,
+    get_tags_from_request_body,
+)
 from litellm.proxy.common_utils.openai_error_payload import (
     attribute_of,
     error_status_code,
@@ -643,6 +665,48 @@ async def _resolve_per_request_model_group_alias(
     return target
 
 
+_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
+
+
+def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
+    try:
+        return _REQUEST_MODEL.validate_python(data.get("model"), strict=True)
+    except ValidationError:
+        return None
+
+
+async def _enforce_guardrail_added_tag_budgets(
+    data: Mapping[str, object],
+    tags_before_guardrails: frozenset[str],
+    route: str,
+    llm_router: Router | None,
+    user_api_key_dict: UserAPIKeyAuth,
+    proxy_logging_obj: ProxyLogging,
+) -> None:
+    added_tags: Final = tuple(
+        tag for tag in get_tags_from_request_body(request_body=data) if tag not in tags_before_guardrails
+    )
+    if not added_tags or request_skips_budget_checks(route=route, model=_request_model(data), llm_router=llm_router):
+        return
+    from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
+
+    try:
+        await tag_max_budget_check_for_tags(
+            tags=added_tags,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+            valid_token=user_api_key_dict,
+        )
+    except litellm.BudgetExceededError as e:
+        raise ProxyException(
+            message=e.message,
+            type=ProxyErrorTypes.budget_exceeded,
+            param=None,
+            code=e.status_code,
+        ) from e
+
+
 async def _parse_event_data_for_error(event_line: str | bytes) -> int | None:
     """Parses an event line and returns an error code if present, else None."""
     event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line
@@ -1452,7 +1516,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
 _CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
 
 
-def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
+@runtime_checkable
+class _CarriesLitellmCallId(Protocol):
+    litellm_call_id: str | None
+
+
+def request_litellm_call_id(data: Mapping[str, object]) -> str | None:
+    logging_obj: Final = data.get("litellm_logging_obj")
+    logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None
+    call_id: Final = logged_id or data.get("litellm_call_id")
+    return call_id if isinstance(call_id, str) else None
+
+
+def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
     if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL:
         verbose_proxy_logger.info(
             "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, "
@@ -1531,6 +1607,11 @@ def _timing_values(
 class ProxyBaseLLMRequestProcessing:
     def __init__(self, data: dict):
         self.data = data
+        self._tags_before_guardrails: frozenset[str] | None = None
+
+    @property
+    def litellm_call_id(self) -> str | None:
+        return request_litellm_call_id(self.data)
 
     @staticmethod
     def _merge_passthrough_streaming_headers(
@@ -2020,11 +2101,21 @@ class ProxyBaseLLMRequestProcessing:
         # to run below.
         await _arm_auto_router_compression(data=self.data, llm_router=llm_router)
 
+        if self._tags_before_guardrails is None:
+            self._tags_before_guardrails = frozenset(get_tags_from_request_body(request_body=self.data))
         self.data = await proxy_logging_obj.pre_call_hook(
             user_api_key_dict=user_api_key_dict,
             data=self.data,
             call_type=route_type,
         )
+        await _enforce_guardrail_added_tag_budgets(
+            data=self.data,
+            tags_before_guardrails=self._tags_before_guardrails,
+            route=get_request_route(request=request),
+            llm_router=llm_router,
+            user_api_key_dict=user_api_key_dict,
+            proxy_logging_obj=proxy_logging_obj,
+        )
         if route_type == "aget_responses":
             attach_post_call_pipelines_to_retrieval(
                 data=self.data,
@@ -2062,6 +2153,13 @@ class ProxyBaseLLMRequestProcessing:
     ) -> tuple[dict, LiteLLMLoggingObj]:
         from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
 
+        configured_fallbacks: Final = (
+            self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
+            if llm_router is not None and not self.data.get("disable_fallbacks")
+            else None
+        )
+        pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
+
         try:
             return await self.common_processing_pre_call_logic(
                 request=request,
@@ -2080,14 +2178,19 @@ class ProxyBaseLLMRequestProcessing:
                 llm_router=llm_router,
             )
         except ProxyRateLimitError as original_exc:
-            original_model: Final = self.data.get("model")
-            if not original_model or not llm_router or self.data.get("disable_fallbacks"):
+            rate_limited_data: Final = self.data
+            original_model: Final = rate_limited_data.get("model")
+            if (
+                pristine is None
+                or not configured_fallbacks
+                or rate_limited_data.get("disable_fallbacks")
+                or not isinstance(original_model, str)
+            ):
                 raise
 
             fallback_models: Final = self._resolve_fallback_models(
                 model=original_model,
-                llm_router=llm_router,
-                user_api_key_dict=user_api_key_dict,
+                fallbacks=configured_fallbacks,
             )
             if not fallback_models:
                 raise
@@ -2102,6 +2205,7 @@ class ProxyBaseLLMRequestProcessing:
                 for fallback_model in fallback_models:
                     if fallback_model == original_model:
                         continue
+                    self.data = independent_snapshot(pristine)
                     self.data["model"] = fallback_model
                     try:
                         return await self.common_processing_pre_call_logic(
@@ -2123,39 +2227,30 @@ class ProxyBaseLLMRequestProcessing:
                     except ProxyRateLimitError:
                         continue
             except BaseException:
-                self.data["model"] = original_model
+                self.data = rate_limited_data
                 raise
 
-            self.data["model"] = original_model
+            self.data = rate_limited_data
             raise original_exc
 
-    def _resolve_fallback_models(
-        self,
-        model: str,
-        llm_router: Router,
-        user_api_key_dict: UserAPIKeyAuth,
-    ) -> list | None:
-        from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
-
-        fallbacks = None
-
+    @staticmethod
+    def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None:
         key_router_settings: Final = user_api_key_dict.router_settings
-        if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings:
-            fallbacks = key_router_settings["fallbacks"]
+        key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None
+        fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks
+        return fallbacks if isinstance(fallbacks, list) and fallbacks else None
 
-        if fallbacks is None:
-            fallbacks = llm_router.fallbacks
-
-        if not fallbacks:
-            return None
+    @staticmethod
+    def _resolve_fallback_models(model: str, fallbacks: list) -> list | None:
+        from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
 
         fallback_model_group, generic_fallback_idx = get_fallback_model_group(
             fallbacks=fallbacks,
             model_group=model,
         )
-        if fallback_model_group is None and generic_fallback_idx is not None:
-            fallback_model_group = fallbacks[generic_fallback_idx]["*"]
-        return fallback_model_group
+        if fallback_model_group is not None:
+            return fallback_model_group
+        return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None
 
     @staticmethod
     def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
@@ -3429,11 +3524,7 @@ class ProxyBaseLLMRequestProcessing:
         version: str | None = None,
     ):
         """Raises ProxyException (OpenAI API compatible) if an exception is raised"""
-        logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None)
-        _log_llm_api_exception(
-            e,
-            (logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"),
-        )
+        log_llm_api_exception(e, self.litellm_call_id)
         # Allow callbacks to transform the error response
         transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict,
@@ -3463,9 +3554,7 @@ class ProxyBaseLLMRequestProcessing:
 
         custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
             user_api_key_dict=user_api_key_dict,
-            call_id=(
-                _litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id")
-            ),
+            call_id=self.litellm_call_id,
             model_id=model_id,
             version=version,
             response_cost=0,

litellm/proxy/common_utils/openai_error_payload.py

@@ -9,6 +9,9 @@ from typing import Final
 from fastapi import status
 
 from litellm.constants import STRINGIFIED_NONE
+from litellm.proxy._types import ProxyException
+
+LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id"
 
 _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
     {
@@ -52,3 +55,23 @@ def openai_error_param(exc: object) -> str | None:
     serializes as JSON ``null``."""
     carried: Final = attribute_of(exc, "param")
     return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
+
+
+def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None:  # mutable-ok: ProxyException.headers
+    if litellm_call_id is None:
+        return None
+    return {LITELLM_CALL_ID_HEADER: litellm_call_id}  # mutable-ok: ProxyException mutates its headers dict
+
+
+def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException:
+    """The same error object, answering with ``x-litellm-call-id`` when it was raised without one."""
+    if litellm_call_id is not None:
+        exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
+    return exc
+
+
+def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]:
+    """``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name."""
+    if headers is None:
+        return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id})
+    return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers})

litellm/proxy/common_utils/proxy_rate_limit_error.py

@@ -11,7 +11,7 @@ exception types:
   an upstream LLM provider returns 429.
 * :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks
   such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
-  ``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
+  ``batch_rate_limiter``, ``max_iterations_limiter``,
   etc.
 * :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
   429) — raised by some provider transports.

litellm/proxy/db/create_views.py

@@ -78,6 +78,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None:
                 v.*,
                 t.spend AS team_spend,
                 t.max_budget AS team_max_budget,
+                t.model_max_budget AS team_model_max_budget,
                 t.tpm_limit AS team_tpm_limit,
                 t.rpm_limit AS team_rpm_limit,
                 t.tpd_limit AS team_tpd_limit,