Merge pull request #36286 from BerriAI/litellm_internal_staging

chore(ci): promote internal staging to main

.github/ci-coverage-allowlist.yml

@@ -136,13 +136,6 @@ test_paths:
       - tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
 
 dockerfiles:
-  - reason: >-
-      The componentized images the microservices chart deploys are built by no job; wiring both into
-      the scan workflow costs a full image build each and is deferred to a change that prices the
-      whole set
-    paths:
-      - backend/Dockerfile
-      - gateway/Dockerfile
   - reason: >-
       The dashboard container is a static Next.js export served by nginx, and the dashboard build
       and lint workflows already exercise that output, so building the image adds no signal about it

.github/pull_request_template.md

@@ -13,6 +13,33 @@ How it solves it:
 - <blah>
 - ...
 
+## User Flow
+
+<!-- Two ordered lists, Before and After, walking the same end user through the same task, written strictly from that user's seat
+     Read the linked issue, ticket, or customer thread first so the flow reflects the real application and the routes its users actually hit; don't invent a generic scenario
+     Lead each list with one plain sentence saying where the flow fails (Before) or succeeds (After), then number the steps
+     Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
+     No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong
+     Keep the two lists step-for-step identical until they diverge, so the changed step is obvious
+     If the bug had a security or authorization consequence, end each list with what another user could or could no longer do
+     Regenerate this section whenever new commits change the PR's behavior, so it never describes an older revision
+
+Example:
+
+Before: a developer whose app streams chat completions gets no token counts back, so their cost dashboard reads zero
+
+1. They send POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
+2. The last SSE chunk arrives with `"usage": null`, so their app records 0 prompt and 0 completion tokens
+3. They open https://litellm-domain/ui/?page=logs and see the request logged at $0 spend
+
+After: the same request comes back with real token counts, so the dashboard shows real spend
+
+1. The proxy admin sets `always_include_stream_usage: true` and restarts the proxy
+2. The developer sends the same POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
+3. The last SSE chunk now carries a `usage` object with real prompt and completion token counts
+4. https://litellm-domain/ui/?page=logs shows that request at non-zero spend
+-->
+
 ## Relevant issues
 
 <!-- e.g., "Fixes #000" -->

.github/workflows/check-schema-sync.yml

@@ -10,6 +10,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   check-sync:
     name: Verify schema.prisma copies match root

.github/workflows/check-ui-api-types.yml

@@ -2,18 +2,19 @@ name: Check UI API Types Sync
 
 on:
   pull_request:
-    paths:
-      - "litellm/proxy/**"
-      - "litellm/types/**"
-      - "ui/litellm-dashboard/src/lib/http/schema.d.ts"
-      - "ui/litellm-dashboard/scripts/gen-api-types.mjs"
-      - "ui/litellm-dashboard/package.json"
-      - "ui/litellm-dashboard/package-lock.json"
-      - ".github/workflows/check-ui-api-types.yml"
+    branches:
+      - main
+      - litellm_internal_staging
+      - litellm_oss_staging
+      - "litellm_**"
 
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
 jobs:
   check-sync:
     name: Verify schema.d.ts matches the proxy OpenAPI spec
@@ -24,18 +25,39 @@ jobs:
         uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
         with:
           persist-credentials: false
+          fetch-depth: 2
+
+      - name: Detect changes that can affect the generated types
+        id: changes
+        run: |
+          set -euo pipefail
+          if ! base="$(git rev-parse --verify --quiet HEAD^2 >/dev/null && git rev-parse HEAD^1)"; then
+            echo "Not a pull request merge commit, running the full check."
+            echo "relevant=true" >> "$GITHUB_OUTPUT"
+            exit 0
+          fi
+          files="$(git diff --name-only "$base" HEAD)"
+          if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)' <<< "$files"; then
+            echo "relevant=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "No proxy, types or generator changes in this pull request, nothing to verify."
+            echo "relevant=false" >> "$GITHUB_OUTPUT"
+          fi
 
       - name: Set up Python
+        if: steps.changes.outputs.relevant == 'true'
         uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: "3.12"
 
       - name: Set up uv
+        if: steps.changes.outputs.relevant == 'true'
         uses: ./.github/actions/setup-uv-with-retries
         with:
           version: "0.10.9"
 
       - name: Cache uv dependencies
+        if: steps.changes.outputs.relevant == 'true'
         uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
         with:
           path: |
@@ -46,14 +68,17 @@ jobs:
             ${{ runner.os }}-uv-
 
       - name: Install backend dependencies
+        if: steps.changes.outputs.relevant == 'true'
         run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
 
       - name: Generate Prisma client
+        if: steps.changes.outputs.relevant == 'true'
         env:
           PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
         run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
 
       - name: Set up Node.js
+        if: steps.changes.outputs.relevant == 'true'
         uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
         with:
           node-version-file: ui/litellm-dashboard/.nvmrc
@@ -61,16 +86,19 @@ jobs:
           cache-dependency-path: ui/litellm-dashboard/package-lock.json
 
       - name: Install dashboard dependencies
+        if: steps.changes.outputs.relevant == 'true'
         working-directory: ui/litellm-dashboard
         run: npm ci
 
       - name: Regenerate types from the live spec
+        if: steps.changes.outputs.relevant == 'true'
         working-directory: ui/litellm-dashboard
         env:
           LITELLM_PYTHON: "uv run --no-sync python"
         run: npm run gen:api
 
       - name: Fail if types are stale
+        if: steps.changes.outputs.relevant == 'true'
         run: |
           if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
             echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec."

.github/workflows/conventional-commits.yml

@@ -14,6 +14,10 @@ on:
 permissions:
   pull-requests: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   lint-pr-title:
     name: Validate PR title

.github/workflows/guard-fork-dependencies.yml

@@ -15,6 +15,10 @@ on:
 
 permissions: {}
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   guard:
     name: Block fork dependency changes

.github/workflows/helm_unit_test.yml

@@ -9,6 +9,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   unit-test:
     runs-on: ubuntu-latest

.github/workflows/image-scan.yml

@@ -12,6 +12,11 @@ on:
       - docker/Dockerfile.non_root
       - migrations/Dockerfile
       - migrations/run.py
+      - gateway/Dockerfile
+      - gateway/main.py
+      - backend/Dockerfile
+      - backend/main.py
+      - docker/component_entrypoint.sh
       - litellm-proxy-extras/**
       - tests/proxy_migration_tests/**
       - uv.lock
@@ -147,3 +152,63 @@ jobs:
         run: |
           python -m pip install "pytest==9.0.3"
           python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
+
+  gateway-image:
+    name: gateway-image
+    runs-on: ubuntu-latest
+    if: >-
+      github.event_name != 'pull_request' ||
+      github.event.pull_request.head.repo.full_name == github.repository
+    timeout-minutes: 30
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - name: Build gateway image
+        run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} .
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - name: Verify the gateway serves offline as a non-root uid
+        env:
+          LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }}
+          LITELLM_COMPONENT_PORT: "4000"
+        run: |
+          python -m pip install "pytest==9.0.3"
+          python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
+
+  backend-image:
+    name: backend-image
+    runs-on: ubuntu-latest
+    if: >-
+      github.event_name != 'pull_request' ||
+      github.event.pull_request.head.repo.full_name == github.repository
+    timeout-minutes: 30
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - name: Build backend image
+        run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} .
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - name: Verify the backend serves offline as a non-root uid
+        env:
+          LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }}
+          LITELLM_COMPONENT_PORT: "4001"
+        run: |
+          python -m pip install "pytest==9.0.3"
+          python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v

.github/workflows/test-linting.yml

@@ -11,6 +11,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   lint:
     runs-on: ubuntu-latest

.github/workflows/test-litellm-ui-build.yml

@@ -10,6 +10,10 @@ on:
       - litellm_oss_staging
       - "litellm_**"
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   build-ui:
     runs-on: ubuntu-latest

.github/workflows/test-litellm-ui-lint.yml

@@ -10,6 +10,10 @@ on:
       - litellm_oss_staging
       - "litellm_**"
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   frontend-lint:
     runs-on: ubuntu-latest

.github/workflows/test-mcp.yml

@@ -11,6 +11,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   test:
     runs-on: ubuntu-latest

.github/workflows/test-model-map.yaml

@@ -11,6 +11,10 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   validate-model-prices-json:
     runs-on: ubuntu-latest

.github/workflows/test-unit-proxy-endpoints.yml

@@ -38,12 +38,15 @@ jobs:
         tests/test_litellm/proxy/google_endpoints
         tests/test_litellm/proxy/openai_files_endpoint
         tests/test_litellm/proxy/batches_endpoints
+        tests/test_litellm/proxy/fine_tuning_endpoints
+        tests/test_litellm/proxy/vector_store_files_endpoints
         tests/test_litellm/proxy/video_endpoints
         tests/test_litellm/proxy/response_api_endpoints
         tests/test_litellm/proxy/image_endpoints
         tests/test_litellm/proxy/vector_store_endpoints
         tests/test_litellm/proxy/agent_endpoints
         tests/test_litellm/proxy/a2a
+        tests/test_litellm/proxy/credential_endpoints
         tests/test_litellm/proxy/discovery_endpoints
         tests/test_litellm/proxy/health_endpoints
         tests/test_litellm/proxy/shutdown

CLAUDE.md

@@ -9,7 +9,7 @@ Don't assume that the existing code is correct or the right way of doing things
 - easy to maintain/change
 - modern
 
-In that order of importance
+In descending order of importance
 
 When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
 
@@ -41,7 +41,7 @@ Python max line length is 120, not 88
 
 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
 
-`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
+`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
 
 If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
 

Makefile

@@ -8,7 +8,7 @@
 	lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
 	lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
 	install-dev install-proxy-dev install-test-deps install-hooks \
-	install-helm-unittest check-circular-imports check-import-safety pre-commit \
+	install-helm-unittest check-circular-imports check-import-safety check pre-commit \
 	lint-install lint-fetch-base bootstrap
 
 # Default target
@@ -22,7 +22,8 @@ help:
 	@echo "  make install-test-deps  - Install the full local test environment"
 	@echo "  make install-helm-unittest - Install helm unittest plugin"
 	@echo "  make install-hooks      - Install git hooks (Conventional Commits + Branches)"
-	@echo "  make pre-commit         - Run CI-equivalent lint on staged files (run before committing)"
+	@echo "  make check              - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
+	@echo "  make pre-commit         - Legacy alias for make check"
 	@echo "  make format             - Apply ruff format code formatting"
 	@echo "  make format-check       - Check ruff format code formatting (matches CI)"
 	@echo "  make lint               - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
@@ -236,13 +237,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline
 # Faster linting for local development (only checks changed code)
 lint-dev: lint-format-changed check-circular-imports check-import-safety
 
-# Run the gating CI checks against your staged files right before committing. Mirrors
+# Run the gating CI checks against your changes. Scopes to staged files when anything
+# is staged (warning about changed files left unstaged); with nothing staged it falls
+# back to the working tree's diff against the merge base with the base branch, so a
+# fresh merge commit or an unstaged working tree still gets checked. Mirrors
 # test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
-# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage.
+# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope.
 # Not auto-installed as a git hook so it never slows an unrelated human commit.
-pre-commit: bootstrap
+check: bootstrap
 	./scripts/pre_commit_lint.sh
 
+pre-commit:
+	@echo "make pre-commit is a legacy alias; use make check" >&2
+	@$(MAKE) check
+
 # Testing targets
 test: install-test-deps
 	$(UV_RUN) pytest tests/

basedpyright-code-budget.json

@@ -1,9 +1,9 @@
 {
   "reportAny": {
-    "limit": 28842
+    "limit": 27731
   },
   "reportArgumentType": {
-    "limit": 2634
+    "limit": 2626
   },
   "reportAssignmentType": {
     "limit": 329
@@ -12,7 +12,7 @@
     "limit": 514
   },
   "reportCallIssue": {
-    "limit": 117
+    "limit": 116
   },
   "reportConstantRedefinition": {
     "limit": 40
@@ -24,7 +24,7 @@
     "limit": 19
   },
   "reportExplicitAny": {
-    "limit": 9103
+    "limit": 8807
   },
   "reportFunctionMemberAccess": {
     "limit": 7
@@ -54,10 +54,10 @@
     "limit": 0
   },
   "reportMissingParameterType": {
-    "limit": 5843
+    "limit": 5835
   },
   "reportMissingTypeArgument": {
-    "limit": 15816
+    "limit": 15790
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -72,7 +72,7 @@
     "limit": 0
   },
   "reportOptionalMemberAccess": {
-    "limit": 1078
+    "limit": 1077
   },
   "reportOptionalOperand": {
     "limit": 0
@@ -84,34 +84,34 @@
     "limit": 56
   },
   "reportPrivateUsage": {
-    "limit": 1825
+    "limit": 1824
   },
   "reportRedeclaration": {
     "limit": 8
   },
   "reportReturnType": {
-    "limit": 218
+    "limit": 217
   },
   "reportTypedDictNotRequiredAccess": {
-    "limit": 27
+    "limit": 26
   },
   "reportUndefinedVariable": {
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 45110
+    "limit": 45063
   },
   "reportUnknownLambdaType": {
     "limit": 113
   },
   "reportUnknownMemberType": {
-    "limit": 39838
+    "limit": 39773
   },
   "reportUnknownParameterType": {
-    "limit": 20237
+    "limit": 20207
   },
   "reportUnknownVariableType": {
-    "limit": 31383
+    "limit": 31281
   },
   "reportUnnecessaryCast": {
     "limit": 122
@@ -123,7 +123,7 @@
     "limit": 5
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 864
+    "limit": 862
   },
   "reportUntypedBaseClass": {
     "limit": 0

enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py

@@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.proxy._types import *
+from litellm.proxy.auth.auth_checks import delete_cached_project_object
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
 from litellm.proxy.management_helpers.utils import (
@@ -514,6 +515,7 @@ async def update_project(
         litellm_proxy_admin_name,
         premium_user,
         prisma_client,
+        user_api_key_cache,
     )
 
     try:
@@ -672,6 +674,11 @@ async def update_project(
             include={"litellm_budget_table": True, "object_permission": True},
         )
 
+        await delete_cached_project_object(
+            project_id=data.project_id,
+            user_api_key_cache=user_api_key_cache,
+        )
+
         return updated_project
     except Exception as e:
         verbose_proxy_logger.exception(
@@ -710,7 +717,7 @@ async def delete_project(
     }'
     ```
     """
-    from litellm.proxy.proxy_server import premium_user, prisma_client
+    from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
 
     try:
         if not premium_user:
@@ -773,6 +780,11 @@ async def delete_project(
                 prisma_models.LiteLLM_ProjectTable | None
             ) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id})
 
+            await delete_cached_project_object(
+                project_id=project_id,
+                user_api_key_cache=user_api_key_cache,
+            )
+
             deleted_projects.append(deleted_project)
 
         return deleted_projects

litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql

@@ -0,0 +1 @@
+ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}';

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
   total_tokens          BigInt   @default(0)
   spend                 Float    @default(0)
   saved_spend           Float    @default(0)
+  tier_turns            Json     @default("{}")
 
   @@id([api_key, session_id, router_name])
   @@index([last_turn_at], map: "idx_autorouter_session_last_turn")

litellm/a2a_protocol/main.py

@@ -13,7 +13,6 @@ import asyncio
 import datetime
 import uuid
 from collections.abc import AsyncIterator, Coroutine
-from http.cookiejar import DefaultCookiePolicy
 from typing import TYPE_CHECKING, Any, Final, Optional, cast
 
 import litellm
@@ -80,8 +79,6 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
 # Use our custom resolver instead of the default A2A SDK resolver
 A2ACardResolver: Final = LiteLLMA2ACardResolver
 
-_BLOCK_ALL_COOKIES: Final = DefaultCookiePolicy(allowed_domains=())
-
 
 def _set_usage_on_logging_obj(
     kwargs: dict[str, Any],
@@ -770,7 +767,6 @@ async def create_a2a_client(
         params={"timeout": timeout},
     )
     httpx_client: Final = _async_handler.client
-    httpx_client.cookies.jar.set_policy(_BLOCK_ALL_COOKIES)
     if extra_headers:
         verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
 

litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py

@@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s
 """
 
 import asyncio
-from collections.abc import AsyncIterator
-from typing import Any, Final, cast
+from collections.abc import AsyncIterator, Mapping, Sequence
+from typing import Any, Final, Protocol, cast, runtime_checkable
 from uuid import uuid4
 
+from pydantic import TypeAdapter
+
 from litellm._logging import verbose_logger
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     get_async_httpx_client,
 )
 
+_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
+_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
+_LIST_ADAPTER: Final = TypeAdapter(list[object])
+_TEXT_ADAPTER: Final = TypeAdapter(str)
+
+
+@runtime_checkable
+class _SupportsModelDump(Protocol):
+    def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ...
+
+
+@runtime_checkable
+class _SupportsPydanticDict(Protocol):
+    def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ...
+
 
 class PydanticAITransformation:
     """
@@ -28,7 +45,7 @@ class PydanticAITransformation:
     """
 
     @staticmethod
-    def _remove_none_values(obj: Any) -> Any:
+    def _remove_none_values(obj: object) -> object:
         """
         Recursively remove None values from a dict/list structure.
 
@@ -42,14 +59,18 @@ class PydanticAITransformation:
             Cleaned object with None values removed
         """
         if isinstance(obj, dict):
-            return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None}
+            typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj)
+            return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None}
         elif isinstance(obj, list):
-            return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None]
+            typed_list: Final = _LIST_ADAPTER.validate_python(obj)
+            return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None]
         else:
             return obj
 
     @staticmethod
-    def _params_to_dict(params: Any) -> dict[str, Any]:
+    def _params_to_dict(
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
+    ) -> Mapping[str, object]:
         """
         Convert params to a dict, handling Pydantic models.
 
@@ -59,10 +80,10 @@ class PydanticAITransformation:
         Returns:
             Dict representation of params
         """
-        if hasattr(params, "model_dump"):
+        if isinstance(params, _SupportsModelDump):
             # Pydantic v2 model
             return params.model_dump(mode="python", exclude_none=True)
-        elif hasattr(params, "dict"):
+        elif isinstance(params, _SupportsPydanticDict):
             # Pydantic v1 model
             return params.dict(exclude_none=True)
         elif isinstance(params, dict):
@@ -75,12 +96,12 @@ class PydanticAITransformation:
     async def _poll_for_completion(
         client: AsyncHTTPHandler,
         endpoint: str,
-        task_id: str,
+        task_id: object,
         request_id: str,
         max_attempts: int = 30,
         poll_interval: float = 0.5,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Poll for task completion using tasks/get method.
 
@@ -112,10 +133,10 @@ class PydanticAITransformation:
                 },
             )
             response.raise_for_status()
-            poll_data = response.json()
+            poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
 
-            result = poll_data.get("result", {})
-            status = result.get("status", {})
+            result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {}))
+            status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
             state = status.get("state", "")
 
             verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state)
@@ -133,10 +154,10 @@ class PydanticAITransformation:
     async def _send_and_poll_raw(
         api_base: str,
         request_id: str,
-        params: Any,
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
         timeout: float = 60.0,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Send a request to Pydantic AI agent and return the raw task response.
 
@@ -153,14 +174,16 @@ class PydanticAITransformation:
             Raw Pydantic AI task response (with history/artifacts)
         """
         # Convert params to dict if it's a Pydantic model
-        params_dict = PydanticAITransformation._params_to_dict(params)
-
         # Remove None values - FastA2A doesn't accept null for optional fields
-        params_dict = PydanticAITransformation._remove_none_values(params_dict)
+        params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(
+            PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params))
+        )
 
         # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI
         if "message" in params_dict:
-            params_dict["message"]["kind"] = "message"
+            message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"])
+            message_value["kind"] = "message"
+            params_dict["message"] = message_value
 
         # Build A2A JSON-RPC request using message/send method for FastA2A compatibility
         a2a_request: Final = {
@@ -189,11 +212,11 @@ class PydanticAITransformation:
             },
         )
         response.raise_for_status()
-        response_data = response.json()
+        response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
 
         # Check if task is already completed
-        result: Final = response_data.get("result", {})
-        status: Final = result.get("status", {})
+        result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
+        status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
         state: Final = status.get("state", "")
 
         if state != "completed":
@@ -217,10 +240,10 @@ class PydanticAITransformation:
     async def send_non_streaming_request(
         api_base: str,
         request_id: str,
-        params: Any,
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
         timeout: float = 60.0,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
 
@@ -253,10 +276,10 @@ class PydanticAITransformation:
     async def send_and_get_raw_response(
         api_base: str,
         request_id: str,
-        params: Any,
+        params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
         timeout: float = 60.0,
         agent_extra_headers: dict[str, str] | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Send a request to Pydantic AI agent and return the raw task response.
 
@@ -282,9 +305,9 @@ class PydanticAITransformation:
 
     @staticmethod
     def _transform_to_a2a_response(
-        response_data: dict[str, Any],
+        response_data: Mapping[str, object],
         request_id: str,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Transform Pydantic AI task response to standard A2A non-streaming format.
 
@@ -328,7 +351,7 @@ class PydanticAITransformation:
         }
 
     @staticmethod
-    def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]:
+    def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]:
         """
         Extract response text from completed task response.
 
@@ -342,52 +365,53 @@ class PydanticAITransformation:
         Returns:
             Tuple of (full_text, message_id, parts)
         """
-        result: Final = response_data.get("result", {})
+        result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
 
         # Try to extract from artifacts first (preferred for results)
         artifacts: Final = result.get("artifacts", [])
         if artifacts:
-            for artifact in artifacts:
-                parts = artifact.get("parts", [])
+            for artifact in _LIST_ADAPTER.validate_python(artifacts):
+                parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", []))
                 for part in parts:
-                    if part.get("kind") == "text":
-                        text = part.get("text", "")
+                    if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
+                        text = part_dict.get("text", "")
                         if text:
                             return text, str(uuid4()), parts
 
         # Fall back to history - get the last agent message
-        history: Final = result.get("history", [])
+        history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
         for msg in reversed(history):
-            if msg.get("role") == "agent":
-                parts = msg.get("parts", [])
-                message_id = msg.get("messageId", str(uuid4()))
+            if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent":
+                parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", []))
+                message_id = msg_dict.get("messageId", str(uuid4()))
                 full_text = ""
                 for part in parts:
-                    if part.get("kind") == "text":
-                        full_text += part.get("text", "")
+                    if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
+                        full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
                 if full_text:
                     return full_text, message_id, parts
 
         # Fall back to message field (original format)
         message: Final = result.get("message", {})
         if message:
-            parts = message.get("parts", [])
-            message_id = message.get("messageId", str(uuid4()))
+            message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
+            parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", []))
+            message_id = message_dict.get("messageId", str(uuid4()))
             full_text = ""
             for part in parts:
-                if part.get("kind") == "text":
-                    full_text += part.get("text", "")
+                if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
+                    full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
             return full_text, message_id, parts
 
         return "", str(uuid4()), []
 
     @staticmethod
     async def fake_streaming_from_response(
-        response_data: dict[str, Any],
+        response_data: Mapping[str, object],
         request_id: str,
         chunk_size: int = 50,
         delay_ms: int = 10,
-    ) -> AsyncIterator[dict[str, Any]]:
+    ) -> AsyncIterator[dict[str, object]]:
         """
         Convert a non-streaming A2A response into fake streaming chunks.
 
@@ -410,12 +434,12 @@ class PydanticAITransformation:
         full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data)
 
         # Extract input message from raw response for history
-        result: Final = response_data.get("result", {})
-        history: Final = result.get("history", [])
-        input_message = {}
+        result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
+        history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
+        input_message = _STR_KEY_DICT_ADAPTER.validate_python({})
         for msg in history:
-            if msg.get("role") == "user":
-                input_message = msg
+            if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user":
+                input_message = msg_dict
                 break
 
         # Generate IDs for streaming events
@@ -426,45 +450,49 @@ class PydanticAITransformation:
 
         # 1. Emit initial task event (kind: "task", status: "submitted")
         # Format matches A2ACompletionBridgeTransformation.create_task_event
-        task_event: Final = {
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "contextId": context_id,
-                "history": [
-                    {
-                        "contextId": context_id,
-                        "kind": "message",
-                        "messageId": input_message_id,
-                        "parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
-                        "role": "user",
-                        "taskId": task_id,
-                    }
-                ],
-                "id": task_id,
-                "kind": "task",
-                "status": {
-                    "state": "submitted",
+        task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "contextId": context_id,
+                    "history": [
+                        {
+                            "contextId": context_id,
+                            "kind": "message",
+                            "messageId": input_message_id,
+                            "parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
+                            "role": "user",
+                            "taskId": task_id,
+                        }
+                    ],
+                    "id": task_id,
+                    "kind": "task",
+                    "status": {
+                        "state": "submitted",
+                    },
                 },
-            },
-        }
+            }
+        )
         yield task_event
 
         # 2. Emit status update (kind: "status-update", status: "working")
         # Format matches A2ACompletionBridgeTransformation.create_status_update_event
-        working_event: Final = {
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "contextId": context_id,
-                "final": False,
-                "kind": "status-update",
-                "status": {
-                    "state": "working",
+        working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "contextId": context_id,
+                    "final": False,
+                    "kind": "status-update",
+                    "status": {
+                        "state": "working",
+                    },
+                    "taskId": task_id,
                 },
-                "taskId": task_id,
-            },
-        }
+            }
+        )
         yield working_event
 
         # Small delay to simulate processing
@@ -473,29 +501,32 @@ class PydanticAITransformation:
         # 3. Emit artifact update chunks (kind: "artifact-update")
         # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event
         if full_text:
+            full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text)
             # Split text into chunks
-            for i in range(0, len(full_text), chunk_size):
-                chunk_text = full_text[i : i + chunk_size]
-                is_last_chunk = (i + chunk_size) >= len(full_text)
+            for i in range(0, len(full_text_str), chunk_size):
+                chunk_text = full_text_str[i : i + chunk_size]
+                is_last_chunk = (i + chunk_size) >= len(full_text_str)
 
-                artifact_event = {
-                    "jsonrpc": "2.0",
-                    "id": request_id,
-                    "result": {
-                        "contextId": context_id,
-                        "kind": "artifact-update",
-                        "taskId": task_id,
-                        "artifact": {
-                            "artifactId": artifact_id,
-                            "parts": [
-                                {
-                                    "kind": "text",
-                                    "text": chunk_text,
-                                }
-                            ],
+                artifact_event = _STR_KEY_DICT_ADAPTER.validate_python(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": request_id,
+                        "result": {
+                            "contextId": context_id,
+                            "kind": "artifact-update",
+                            "taskId": task_id,
+                            "artifact": {
+                                "artifactId": artifact_id,
+                                "parts": [
+                                    {
+                                        "kind": "text",
+                                        "text": chunk_text,
+                                    }
+                                ],
+                            },
                         },
-                    },
-                }
+                    }
+                )
                 yield artifact_event
 
                 # Add delay between chunks (except for last chunk)
@@ -503,19 +534,21 @@ class PydanticAITransformation:
                     await asyncio.sleep(delay_ms / 1000.0)
 
         # 4. Emit final status update (kind: "status-update", status: "completed", final: true)
-        completed_event: Final = {
-            "jsonrpc": "2.0",
-            "id": request_id,
-            "result": {
-                "contextId": context_id,
-                "final": True,
-                "kind": "status-update",
-                "status": {
-                    "state": "completed",
+        completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": {
+                    "contextId": context_id,
+                    "final": True,
+                    "kind": "status-update",
+                    "status": {
+                        "state": "completed",
+                    },
+                    "taskId": task_id,
                 },
-                "taskId": task_id,
-            },
-        }
+            }
+        )
         yield completed_event
 
         verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id)

litellm/constants.py

@@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_non
 DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
 AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
 ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
+ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
 DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
 DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
@@ -279,6 +280,7 @@ TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECO
 GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
     os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
 )
+BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
 # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
 # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
 MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))
@@ -1320,6 +1322,7 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
 LITELLM_METADATA_FIELD: Final = "litellm_metadata"
 OLD_LITELLM_METADATA_FIELD: Final = "metadata"
 RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
+SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
 INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
 LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
 LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (

litellm/integrations/azure_sentinel/azure_sentinel.py

@@ -16,7 +16,10 @@ import asyncio
 import os
 import time
 import traceback
+from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Final
+from urllib.parse import urlparse
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
@@ -27,6 +30,16 @@ from litellm.llms.custom_httpx.http_handler import (
 )
 from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
 
+DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com"
+DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default"
+
+MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE,
+        "login.microsoftonline.us": "https://monitor.azure.us/.default",
+    }
+)
+
 
 class AzureSentinelLogger(CustomBatchLogger):
     """
@@ -42,6 +55,7 @@ class AzureSentinelLogger(CustomBatchLogger):
         client_id: str | None = None,
         client_secret: str | None = None,
         audit_stream_name: str | None = None,
+        authority_host: str | None = None,
         **kwargs,
     ):
         """
@@ -62,6 +76,10 @@ class AzureSentinelLogger(CustomBatchLogger):
                 If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var.
             audit_stream_name (str, optional): Stream name from DCR for audit logs.
                 If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name.
+            authority_host (str, optional): Microsoft Entra authority host that issues the OAuth2 token,
+                e.g. "https://login.microsoftonline.us" for Azure Government. If not provided, will use
+                AZURE_SENTINEL_AUTHORITY_HOST or AZURE_AUTHORITY_HOST env vars, or default to the Azure
+                Public Cloud authority. The Azure Monitor audience is derived from it.
         """
         self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
 
@@ -76,6 +94,12 @@ class AzureSentinelLogger(CustomBatchLogger):
         resolved_client_secret: Final = (
             client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET")
         )
+        resolved_authority_host: Final = self._normalize_authority_host(
+            authority_host
+            or os.getenv("AZURE_SENTINEL_AUTHORITY_HOST")
+            or os.getenv("AZURE_AUTHORITY_HOST")
+            or DEFAULT_AZURE_AUTHORITY_HOST
+        )
 
         if not resolved_dcr_immutable_id:
             raise ValueError(
@@ -119,7 +143,8 @@ class AzureSentinelLogger(CustomBatchLogger):
         )
 
         # OAuth2 scope for Azure Monitor
-        self.oauth_scope = "https://monitor.azure.com/.default"
+        self.authority_host = resolved_authority_host
+        self.oauth_scope = self._resolve_oauth_scope(authority_host=resolved_authority_host)
         self.oauth_token: str | None = None
         self.oauth_token_expires_at: float | None = None
 
@@ -129,6 +154,26 @@ class AzureSentinelLogger(CustomBatchLogger):
         self.log_queue: list[StandardLoggingPayload] = []
         self.audit_log_queue: list[StandardAuditLogPayload] = []
 
+    @staticmethod
+    def _normalize_authority_host(authority_host: str) -> str:
+        """
+        Normalize an authority host into an absolute URL with no trailing slash.
+
+        Accepts the scheme-qualified form litellm documents ("https://login.microsoftonline.us")
+        and the bare-host form the azure-identity AzureAuthorityHosts constants use.
+        """
+        stripped: Final = authority_host.strip().rstrip("/")
+        return stripped if "://" in stripped else f"https://{stripped}"
+
+    @staticmethod
+    def _resolve_oauth_scope(authority_host: str) -> str:
+        """
+        Map an authority host to the Azure Monitor Logs Ingestion audience for the same cloud,
+        falling back to the Azure Public Cloud audience for an unrecognized host.
+        """
+        host: Final = urlparse(authority_host).hostname or ""
+        return MONITOR_SCOPE_BY_AUTHORITY_HOST.get(host, DEFAULT_AZURE_MONITOR_SCOPE)
+
     @staticmethod
     def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str:
         return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01"
@@ -150,7 +195,7 @@ class AzureSentinelLogger(CustomBatchLogger):
         assert self.client_id is not None, "client_id is required"
         assert self.client_secret is not None, "client_secret is required"
 
-        token_url: Final = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
+        token_url: Final = f"{self.authority_host}/{self.tenant_id}/oauth2/v2.0/token"
 
         token_data: Final = {
             "client_id": self.client_id,

litellm/integrations/otel/__init__.py

@@ -57,6 +57,7 @@ from litellm.integrations.otel.model.semconv import (
     Metric,
     Network,
     NetworkTransport,
+    RpcSystem,
     Server,
     resolve_operation,
     resolve_provider,
@@ -102,6 +103,7 @@ __all__ = [
     "ProxyRequestSpanData",
     "RequestContext",
     "RequestIdentity",
+    "RpcSystem",
     "Server",
     "ServerInfo",
     "ServiceSpanData",

litellm/integrations/otel/logger.py

@@ -33,6 +33,7 @@ from litellm.integrations.otel.model.payloads import (
     is_mcp_list_tools,
     is_mcp_tool_call,
 )
+from litellm.integrations.otel.model.semconv import Error
 from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service
 from litellm.integrations.otel.model.utils import to_ns
 from litellm.integrations.otel.plumbing.context import (
@@ -634,18 +635,23 @@ class OpenTelemetryV2(CustomLogger):
         """Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a
         failure that dies before any LLM-call span exists (malformed body, auth /
         validation rejection). Called from the proxy's global exception handler via
-        ``_close_dangling_otel_server_span``. The instrumentor still owns the span's
-        status and lifecycle, so this only decorates it — never sets status, never
-        ends it — and emits no exception event, matching v1's SERVER-span behavior
-        and avoiding a duplicate of the event ``async_post_call_failure_hook`` or
-        the ``auth`` phase span already records."""
+        ``_close_dangling_otel_server_span``, which swallows the exception into a
+        ``JSONResponse`` so the instrumentor never sees it and leaves the span
+        ``UNSET``; the status is set here instead (v1 did the same from the handler)
+        so a failed request reads as failed and not merely as a span carrying an
+        error message. The instrumentor still owns the span's lifecycle, so this
+        never ends it. The exception event is recorded only when nothing stamped
+        this span already — ``async_post_call_failure_hook`` and the ``auth`` phase
+        span record their own, and a second event would duplicate it — while the
+        attributes are always restamped so ``error.code`` stays pinned to the real
+        response status."""
         if span is None or not is_recordable_span(span):
             return
+        already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ())
         stamp_error(
             span,
             _span_error_from_exception(exception, status_code=status_code),
-            record_event=False,
-            set_status=False,
+            record_event=not already_stamped,
         )
 
     async def async_post_call_failure_hook(

litellm/integrations/otel/mappers/genai.py

@@ -31,7 +31,9 @@ from litellm.integrations.otel.model.semconv import (
     MCP,
     Error,
     GenAI,
+    JsonRpc,
     LiteLLM,
+    RpcSystem,
     Server,
 )
 from litellm.integrations.otel.model.spans import db_system
@@ -94,11 +96,14 @@ class GenAIMapper:
 
     _MCP_ATTRS: dict[str, Callable[[MCPToolCallSpanData], AttrValue | None]] = {
         GenAI.OPERATION_NAME: lambda d: d.operation.value,
+        JsonRpc.SYSTEM: lambda d: RpcSystem.JSONRPC.value if d.server_address and d.server_port else None,
         MCP.METHOD_NAME: lambda d: d.method,
         MCP.SESSION_ID: lambda d: d.session_id,
         GenAI.TOOL_NAME: lambda d: d.tool_name or None,
         GenAI.TOOL_CALL_ARGUMENTS: lambda d: d.arguments_json,
         GenAI.TOOL_CALL_RESULT: lambda d: d.result_json,
+        Server.ADDRESS: lambda d: d.server_address,
+        Server.PORT: lambda d: d.server_port,
         LiteLLM.MCP_SERVER_NAME: lambda d: d.server_name,
         LiteLLM.CALL_ID: lambda d: d.identity.call_id or None,
         f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost,

litellm/integrations/otel/model/payloads.py

@@ -364,6 +364,34 @@ class LLMCallSpanData:
 # --- the MCP tool-call model ------------------------------------------------- #
 
 
+def _upstream_address_port(resource: str | None) -> tuple[str | None, int | None]:
+    """Split a redacted MCP server origin into ``server.address`` / ``server.port``.
+
+    ``mcp_server_resource`` is a scheme + host + port origin with userinfo, path,
+    query and fragment already stripped. The port falls back to the scheme default
+    when the origin omits it, because a consumer that keys a downstream dependency
+    off the address renders a missing port as ``0``.
+
+    The origin is rebuilt without its IPv6 brackets upstream, so reading the port can
+    raise on an address the host check still admits: a zone-scoped ``fe80::1%25eth0``
+    leaves a truthy hostname of ``fe80`` behind. Both halves are read inside the guard
+    so an unparseable origin yields no address rather than propagating out of span
+    construction, matching how the redactor guards the same split.
+    """
+    if not resource:
+        return None, None
+    try:
+        parsed: Final = urlsplit(resource)
+        hostname: Final = parsed.hostname
+        port: Final = parsed.port
+    except ValueError:
+        return None, None
+    if not hostname:
+        return None, None
+    default_port: Final = 443 if parsed.scheme == "https" else 80 if parsed.scheme == "http" else None
+    return hostname, port or default_port
+
+
 @dataclass(frozen=True)
 class MCPToolCallSpanData:
     """One MCP ``tools/call`` execution, parsed from a closed request's payload.
@@ -378,6 +406,8 @@ class MCPToolCallSpanData:
     method: str
     tool_name: str
     server_name: str | None
+    server_address: str | None
+    server_port: int | None
     session_id: str | None
     arguments_json: str | None
     result_json: str | None
@@ -390,11 +420,14 @@ class MCPToolCallSpanData:
         cls, payload: StandardLoggingPayload, capture_content: bool = False
     ) -> MCPToolCallSpanData:
         meta: Final = _mcp_tool_call_metadata(cast(Mapping[str, object], payload))
+        address, port = _upstream_address_port(as_str(meta.get("mcp_server_resource")) or None)
         return cls(
             operation=resolve_operation(as_str(payload.get("call_type"))),
             method=MCPMethod.TOOLS_CALL.value,
             tool_name=as_str(meta.get("name")) or "",
             server_name=as_str(meta.get("mcp_server_name")),
+            server_address=address,
+            server_port=port,
             session_id=as_str(meta.get("mcp_session_id")),
             arguments_json=(
                 _json_or_none(meta.get("arguments")) if capture_content and meta.get("arguments") is not None else None

litellm/integrations/otel/model/semconv.py

@@ -130,11 +130,25 @@ class JsonRpc:
     """JSON-RPC keys carried on MCP spans. The error/status code lives in the
     ``rpc.*`` namespace per semconv, not ``jsonrpc.*``."""
 
+    SYSTEM: Final = "rpc.system"
     REQUEST_ID: Final = "jsonrpc.request.id"
     PROTOCOL_VERSION: Final = "jsonrpc.protocol.version"
     RESPONSE_STATUS_CODE: Final = "rpc.response.status_code"
 
 
+class RpcSystem(str, Enum):
+    """Well-known values for ``rpc.system``. MCP frames every message as JSON-RPC 2.0.
+
+    Naming the system also classifies the span: a CLIENT span carrying none of the
+    ``rpc.*``/``http.*``/``db.*``/``messaging.*`` families records no span type or
+    subtype in backends that derive those from the attribute family. It is emitted
+    only alongside ``server.address``/``server.port``, since a backend that reads it
+    as a downstream dependency names that dependency from the server address.
+    """
+
+    JSONRPC = "jsonrpc"
+
+
 class NetworkTransport(str, Enum):
     """Well-known values for ``network.transport``."""
 

litellm/integrations/websearch_interception/handler.py

@@ -9,7 +9,7 @@ server-side using litellm router's search tools.
 import asyncio
 import math
 import uuid
-from collections.abc import AsyncIterator, Mapping
+from collections.abc import AsyncIterator, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, cast
 
 import litellm
@@ -37,10 +37,17 @@ from litellm.types.integrations.custom_logger import (
     AgenticLoopRequestPatch,
 )
 from litellm.types.integrations.websearch_interception import (
+    AnthropicSearchQuery,
+    AnthropicServerToolUseBlock,
     WebSearchInterceptionConfig,
 )
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import CallTypes, LlmProviders
+from litellm.types.utils import (
+    AgenticLoopParams,
+    CallTypes,
+    LlmProviders,
+    StandardLoggingUserAPIKeyMetadata,
+)
 from litellm.utils import ProviderConfigManager
 
 if TYPE_CHECKING:
@@ -263,7 +270,7 @@ class WebSearchInterceptionLogger(CustomLogger):
             return None
 
         # Check if request has tools with native web_search
-        tools: Final = kwargs.get("tools")
+        tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools")
         if not tools:
             return None
 
@@ -312,7 +319,9 @@ class WebSearchInterceptionLogger(CustomLogger):
 
         return kwargs
 
-    def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None:
+    def _convert_responses_tools(
+        self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]]
+    ) -> dict[str, object] | None:
         """Convert Responses API web search tools to the LiteLLM standard function tool."""
         if not any(is_web_search_tool_responses(tool) for tool in tools):
             return None
@@ -377,7 +386,7 @@ class WebSearchInterceptionLogger(CustomLogger):
         )
 
     @staticmethod
-    def _tool_name(tool: dict[str, Any]) -> str | None:
+    def _tool_name(tool: Mapping[str, object]) -> object:
         """Effective tool name, handling OpenAI ``function`` wrapper shape."""
         fn: Final = tool.get("function")
         if tool.get("type") == "function" and isinstance(fn, dict):
@@ -833,22 +842,48 @@ class WebSearchInterceptionLogger(CustomLogger):
     def _build_native_result_blocks(
         tool_calls: list[dict],
         structured_results: list[SearchResponse | None],
-    ) -> list[dict[str, object]]:
-        """Build one ``web_search_tool_result`` block per tool_call."""
-        blocks: Final[list[dict[str, object]]] = []
-        for i, tool_call in enumerate(tool_calls):
-            tool_use_id = tool_call.get("id") or ""
-            structured = structured_results[i] if i < len(structured_results) else None
-            blocks.append(
-                WebSearchTransformation.build_web_search_tool_result_block(
-                    tool_use_id=tool_use_id,
-                    search_response=structured,
-                )
+    ) -> tuple[Mapping[str, object], ...]:
+        """
+        Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call.
+
+        The pair is what Anthropic's spec requires: a bare result block, or one
+        keyed by the model's ``toolu_...`` id instead of a ``srvtoolu_...`` one,
+        is rejected on replay ("String should match pattern '^srvtoolu_'") and
+        leaves native clients without a search to attach the sources to.
+        """
+        return tuple(
+            block
+            for i, tool_call in enumerate(tool_calls)
+            for block in WebSearchInterceptionLogger._native_result_pair(
+                query=WebSearchInterceptionLogger._tool_call_query(tool_call),
+                search_response=structured_results[i] if i < len(structured_results) else None,
             )
-        return blocks
+        )
 
     @staticmethod
-    def _inject_native_blocks(response: Any, native_blocks: list[dict[str, object]]) -> Any:
+    def _tool_call_query(tool_call: Mapping[str, object]) -> str:
+        tool_input: Final = tool_call.get("input")
+        if not isinstance(tool_input, Mapping):
+            return ""
+        query: Final = tool_input.get("query")
+        return query if isinstance(query, str) else ""
+
+    @staticmethod
+    def _native_result_pair(
+        query: str,
+        search_response: SearchResponse | None,
+    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+        tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
+        return (
+            AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(),
+            WebSearchTransformation.build_web_search_tool_result_block(
+                tool_use_id=tool_use_id,
+                search_response=search_response,
+            ),
+        )
+
+    @staticmethod
+    def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
         """Prepend native blocks to response content, dict or object form."""
         if not native_blocks:
             return response
@@ -1243,7 +1278,7 @@ class WebSearchInterceptionLogger(CustomLogger):
         kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs)
 
         if logging_obj is not None:
-            agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
+            agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {})
             full_model_name = agentic_params.get("model", model)
         verbose_logger.debug(
             "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]",
@@ -1288,6 +1323,7 @@ class WebSearchInterceptionLogger(CustomLogger):
             search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
             search_provider: str | None = None
             search_litellm_params: dict[str, Any] = {}
+            search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
             if search_tool is not None:
                 await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
                 search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
@@ -1304,12 +1340,30 @@ class WebSearchInterceptionLogger(CustomLogger):
             verbose_logger.debug(
                 "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider
             )
+            user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs)
+            search_metadata: Final = (
+                None
+                if user_api_key_auth is None
+                else self._build_search_request_metadata(
+                    user_api_key_auth=user_api_key_auth,
+                    search_tool_name=search_tool_name,
+                )
+            )
             search_kwargs: Final = {
                 key: value
                 for key, value in search_litellm_params.items()
                 if key != "search_provider" and value is not None
             }
-            result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
+            result: Final = (
+                await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
+                if search_metadata is None
+                else await litellm.asearch(
+                    query=query,
+                    search_provider=search_provider,
+                    litellm_metadata=search_metadata,
+                    **search_kwargs,
+                )
+            )
 
             # Format using transformation function
             search_result_text: Final = WebSearchTransformation.format_search_response(result)
@@ -1366,6 +1420,35 @@ class WebSearchInterceptionLogger(CustomLogger):
                 team_object=team_object,
             )
 
+    @staticmethod
+    def _build_search_request_metadata(
+        user_api_key_auth: "UserAPIKeyAuth",
+        search_tool_name: str | None,
+    ) -> Mapping[str, object]:
+        """
+        Spend-tracking metadata for the intercepted search, so its provider cost is logged
+        and billed against the key/user/team that made the originating LLM request instead
+        of being dropped by the proxy's spend hook for lack of an owner.
+        """
+        from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+
+        user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = (
+            LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth)
+        )
+        return {  # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches
+            **user_api_key_metadata,
+            "model_group": search_tool_name,
+            "user_api_key": user_api_key_auth.api_key,
+            "user_api_key_auth": user_api_key_auth,
+        }
+
+    @staticmethod
+    def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None:
+        if search_tool is None:
+            return None
+        search_tool_name: Final = search_tool.get("search_tool_name")
+        return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None
+
     @staticmethod
     def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None":
         if not kwargs:

litellm/integrations/websearch_interception/transformation.py

@@ -412,6 +412,15 @@ class WebSearchTransformation:
         block that should accompany the model's text reply when the original
         request used a native ``web_search_*`` tool.
 
+        The spec'd shape carries page text only in ``encrypted_content``, an
+        opaque server-issued blob that we cannot mint. Emitting the four spec
+        fields alone would drop the snippet entirely, leaving the client (and
+        the model, on any replayed follow-up turn) with URLs and titles but no
+        evidence to answer from, forcing a fetch per result. So the snippet is
+        carried in an additive ``snippet`` key alongside the spec fields.
+        ``encrypted_content`` stays empty rather than holding plaintext, which
+        would assert encryption semantics that do not hold.
+
         Spec reference:
         https://docs.anthropic.com/en/api/web-search-tool
 
@@ -438,6 +447,7 @@ class WebSearchTransformation:
                         "title": title,
                         "page_age": page_age,
                         "encrypted_content": "",
+                        "snippet": getattr(r, "snippet", "") or "",
                     }
                 )
         return {

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -6,7 +6,7 @@ import io
 import json
 import mimetypes
 import re
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from os import PathLike
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
@@ -1775,3 +1775,24 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]:
         idx = end_idx
 
     return results
+
+
+def text_completion_prompt_to_messages(prompt: object) -> tuple[AllMessageValues, ...]:
+    """
+    Wrap an OpenAI ``/v1/completions`` ``prompt`` into Chat Completion messages.
+
+    Mirrors what ``litellm.text_completion`` does on the real-time path: a
+    string becomes a single user message, and a list of strings becomes one
+    user message per element. Pre-tokenized prompts (``list[int]`` /
+    ``list[list[int]]``) are only meaningful for the OpenAI-family text
+    endpoints, so they are rejected here rather than silently forwarded, as is
+    an empty prompt, which every chat-shaped provider rejects downstream.
+    """
+    prompt_type_name: Final = type(prompt).__name__
+    if isinstance(prompt, str) and prompt:
+        return (ChatCompletionUserMessage(role="user", content=prompt),)
+    entries: Final = cast("Sequence[object]", prompt) if isinstance(prompt, Sequence) else ()
+    string_entries: Final = tuple(entry for entry in entries if isinstance(entry, str) and entry)
+    if string_entries and len(string_entries) == len(entries):
+        return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in string_entries)
+    raise ValueError(f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {prompt_type_name}.")

litellm/litellm_core_utils/realtime_streaming.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Protocol, cast
 
 import litellm
@@ -20,11 +21,24 @@ from .litellm_logging import Logging as LiteLLMLogging
 if TYPE_CHECKING:
     from websockets.asyncio.client import ClientConnection
 
+    from litellm.types.guardrails import GuardrailEventHooks
+
     CLIENT_CONNECTION_CLASS = ClientConnection
 else:
     CLIENT_CONNECTION_CLASS = Any
 
 
+class _ClientWebSocketExceptions(Protocol):
+    ConnectionClosed: type[Exception]
+
+
+class _ClientWebSocket(Protocol):
+    exceptions: _ClientWebSocketExceptions
+
+    async def send_text(self, data: str) -> None: ...
+    async def receive_text(self) -> str: ...
+
+
 class RealtimeEventNormalizer(Protocol):
     def should_drop(self, event: object) -> bool: ...
     def normalize(self, event: dict) -> dict: ...
@@ -48,13 +62,13 @@ class RealTimeStreaming:
         logging_obj: LiteLLMLogging,
         provider_config: BaseRealtimeConfig | None = None,
         model: str = "",
-        user_api_key_dict: Any | None = None,
+        user_api_key_dict: object | None = None,
         request_data: dict | None = None,
         backend_uses_beta_protocol: bool | None = None,
         force_transcription_model: str | None = None,
         event_normalizer: RealtimeEventNormalizer | None = None,
     ):
-        self.websocket = websocket
+        self.websocket: _ClientWebSocket = websocket
         self.backend_ws = backend_ws
         self.logging_obj = logging_obj
         self.messages: list[OpenAIRealtimeEvents] = []
@@ -127,7 +141,7 @@ class RealTimeStreaming:
         ]
     )
     _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"])
-    _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = {
+    _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = {
         "pcm16": {"type": "audio/pcm", "rate": 24000},
         "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
         "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000},
@@ -281,6 +295,7 @@ class RealTimeStreaming:
             if event_obj.get("type") != "response.done":
                 return
             response: Final = cast(dict[str, Any], event_obj.get("response", {}))
+            item: Mapping[str, object]
             for item in response.get("output", []):
                 if item.get("type") == "function_call":
                     self.tool_calls.append(
@@ -384,7 +399,7 @@ class RealTimeStreaming:
             return message
 
         try:
-            message_obj: Final = json.loads(message)
+            message_obj: Final[Mapping[str, object]] = json.loads(message)
         except (json.JSONDecodeError, TypeError):
             return message
 
@@ -487,7 +502,7 @@ class RealTimeStreaming:
         if self._backend_setup_complete and not self._flushing_pending_messages_until_setup:
             return False
         try:
-            msg_obj: Final = json.loads(message)
+            msg_obj: Final[Mapping[str, object]] = json.loads(message)
         except (json.JSONDecodeError, TypeError):
             return False
         return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
@@ -555,7 +570,7 @@ class RealTimeStreaming:
     def _event_to_client_json(self, event: dict) -> str:
         return json.dumps(self._normalize_event_for_ga_client(event))
 
-    async def _send_event_to_client(self, event: Any, event_str: str) -> bool:
+    async def _send_event_to_client(self, event: object, event_str: str) -> bool:
         if self._should_drop_event_from_client(event):
             return False
         if isinstance(event, dict):
@@ -595,12 +610,12 @@ class RealTimeStreaming:
 
     def _make_disable_auto_response_message(self) -> str:
         """Return a session.update that disables VAD auto-response."""
-        turn_detection: Final[dict[str, Any]] = {
+        turn_detection: Final[dict[str, str | bool]] = {
             "type": "server_vad",
             "create_response": False,
         }
         if self._backend_uses_beta_protocol:
-            session: dict[str, Any] = {"turn_detection": turn_detection}
+            session: dict[str, object] = {"turn_detection": turn_detection}
         else:
             session = {
                 "type": "realtime",
@@ -654,7 +669,7 @@ class RealTimeStreaming:
 
     def _has_realtime_guardrails_for_event_hooks(
         self,
-        event_hooks: list[Any],
+        event_hooks: Sequence["GuardrailEventHooks"],
     ) -> bool:
         """Return True if any callback would run for one of ``event_hooks``."""
         from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -699,7 +714,7 @@ class RealTimeStreaming:
         transcript: str,
         item_id: str | None = None,
         pre_block_backend_message: str | None = None,
-        event_hooks: list[Any] | None = None,
+        event_hooks: Sequence["GuardrailEventHooks"] | None = None,
     ) -> bool:
         """
         Run registered guardrails on realtime text (transcript, user message, tool output).
@@ -753,7 +768,7 @@ class RealTimeStreaming:
                     raise
                 # Extract the human-readable error from the detail dict (HTTPException)
                 # or fall back to str(e) for plain ValueError.
-                detail = getattr(e, "detail", None)
+                detail: object | None = getattr(e, "detail", None)
                 if isinstance(detail, dict):
                     safe_msg = detail.get("error") or str(e)
                 elif detail is not None:
@@ -826,7 +841,7 @@ class RealTimeStreaming:
                 return True
         return False
 
-    async def _handle_provider_config_message(self, raw_response) -> None:
+    async def _handle_provider_config_message(self, raw_response: str) -> None:
         """Process a backend message when a provider_config is set (transformed path)."""
         returned_object: Final = self.provider_config.transform_realtime_response(
             raw_response,
@@ -910,7 +925,7 @@ class RealTimeStreaming:
             await self._send_event_to_client(event, event_str)
 
     @staticmethod
-    def _parse_backend_event(raw_response: str) -> dict | None:
+    def _parse_backend_event(raw_response: str) -> dict[str, object] | None:
         """Parse a backend frame once. Returns None for non-JSON or non-object frames."""
         try:
             event: Final = json.loads(raw_response)
@@ -1020,7 +1035,7 @@ class RealTimeStreaming:
         objects and any test doubles that expose a .scope dict.
         """
         try:
-            headers: Final = websocket.scope.get("headers", [])
+            headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", [])
             for name, value in headers:
                 if isinstance(name, bytes):
                     name = name.decode("latin-1")
@@ -1071,9 +1086,9 @@ class RealTimeStreaming:
                     session["output_modalities"] = ["text"]
 
         # 3-7. Lift flat audio fields into the nested audio object
-        audio: Final[dict[str, Any]] = {}
-        inp: Final[dict[str, Any]] = {}
-        out: Final[dict[str, Any]] = {}
+        audio: Final[dict[str, object]] = {}
+        inp: Final[dict[str, object]] = {}
+        out: Final[dict[str, object]] = {}
 
         # voice → audio.output.voice
         if "voice" in session:
@@ -1190,7 +1205,7 @@ class RealTimeStreaming:
                             # model; check them with the same guardrail used for
                             # user text so an attacker cannot smuggle blocked
                             # content into a function_call_output.
-                            output = item.get("output", "")
+                            output: object = item.get("output", "")
                             output_text = output if isinstance(output, str) else json.dumps(output)
                             if output_text:
                                 # Build the sanitized function_call_output up
@@ -1241,7 +1256,7 @@ class RealTimeStreaming:
                                     # interaction turn.
                                     continue
                         elif item.get("role") == "user":
-                            content_list = item.get("content", [])
+                            content_list: Sequence[object] = item.get("content", [])
                             texts = [
                                 c.get("text", "")
                                 for c in content_list
@@ -1280,7 +1295,7 @@ class RealTimeStreaming:
                         and not self._guardrail_turn_detection_update_sent
                         and self._has_audio_transcription_guardrails()
                     ):
-                        session = msg_obj.setdefault("session", {})
+                        session: object = msg_obj.setdefault("session", {})
                         if isinstance(session, dict):
                             existing_td = session.get("turn_detection")
                             if not isinstance(existing_td, dict):

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -3,7 +3,7 @@ import time
 from collections.abc import Iterator, Mapping, Sequence
 from itertools import groupby
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Union, cast
+from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast
 
 from litellm._logging import verbose_logger
 from litellm.types.llms.openai import (
@@ -30,6 +30,7 @@ from litellm.types.utils import (
 from litellm.utils import print_verbose, token_counter
 
 if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging
     from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
         UsagePerChunk,
     )
@@ -39,6 +40,60 @@ if TYPE_CHECKING:
     )
 
 
+class _ThinkingBlockFragment(TypedDict, total=False):
+    type: str | None
+    data: str | None
+    thinking: str | None
+    signature: str | None
+
+
+class _ThinkingDelta(TypedDict, total=False):
+    thinking_blocks: Sequence[_ThinkingBlockFragment]
+
+
+class _ThinkingChoice(TypedDict, total=False):
+    delta: _ThinkingDelta
+
+
+class _ThinkingChunk(TypedDict):
+    choices: Sequence[_ThinkingChoice]
+
+
+class _ContentChoice(TypedDict, total=False):
+    delta: Mapping[str, str | None]
+
+
+class _ContentChunk(TypedDict):
+    choices: Sequence[_ContentChoice]
+
+
+class _AudioDelta(TypedDict, total=False):
+    audio: ChatCompletionAudioDelta | None
+
+
+class _AudioChoice(TypedDict, total=False):
+    delta: _AudioDelta
+
+
+class _AudioChunk(TypedDict):
+    choices: Sequence[_AudioChoice]
+
+
+class _UsageBearingChunk(TypedDict, total=False):
+    usage: Usage | None
+    _hidden_params: Mapping[str, str]
+
+
+class _UsageSummary(TypedDict):
+    prompt_tokens: int | None
+    completion_tokens: int | None
+    cache_creation_input_tokens: int | None
+    cache_read_input_tokens: int | None
+    completion_tokens_details: CompletionTokensDetails | None
+    prompt_tokens_details: PromptTokensDetailsWrapper | None
+    cost: float | None
+
+
 def capture_cache_creation_token_details(
     prompt_tokens_details: PromptTokensDetailsWrapper | None,
     current: CacheCreationTokenDetails | None,
@@ -78,7 +133,7 @@ class ChunkProcessor:
             return []
 
         first_chunk: Final = chunks[0]
-        first_hidden_params: dict[str, Any] = {}
+        first_hidden_params: dict[str, object] = {}
         if isinstance(first_chunk, dict):
             candidate = first_chunk.get("_hidden_params", {})
             if isinstance(candidate, dict):
@@ -115,8 +170,8 @@ class ChunkProcessor:
     @staticmethod
     def apply_provider_assembled_streaming_metadata(
         response: ModelResponse,
-        chunks: list[Any],
-        logging_obj: Any | None = None,
+        chunks: list[object],
+        logging_obj: "Logging | None" = None,
     ) -> None:
         if not chunks:
             return
@@ -456,7 +511,7 @@ class ChunkProcessor:
         )
 
     def get_combined_content(
-        self, chunks: list[dict[str, Any]], delta_key: str = "content"
+        self, chunks: Sequence["_ContentChunk"], delta_key: str = "content"
     ) -> ChatCompletionAssistantContentValue:
         content_list: Final[list[str]] = []
         for chunk in chunks:
@@ -475,7 +530,7 @@ class ChunkProcessor:
         return combined_content
 
     def get_combined_thinking_content(
-        self, chunks: list[dict[str, Any]]
+        self, chunks: Sequence["_ThinkingChunk"]
     ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None:
         from litellm.types.llms.openai import (
             ChatCompletionRedactedThinkingBlock,
@@ -532,10 +587,10 @@ class ChunkProcessor:
             return thinking_blocks
         return None
 
-    def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue:
+    def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue:
         return self.get_combined_content(chunks, delta_key="reasoning_content")
 
-    def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse:
+    def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse:
         base64_data_list: Final[list[str]] = []
         transcript_list: Final[list[str]] = []
         expires_at: int | None = None
@@ -544,7 +599,7 @@ class ChunkProcessor:
         for chunk in chunks:
             choices = chunk["choices"]
             for choice in choices:
-                delta = choice.get("delta") or {}
+                delta: _AudioDelta = choice.get("delta") or {}
                 audio: ChatCompletionAudioDelta | None = delta.get("audio")
                 if audio is not None:
                     for k, v in audio.items():
@@ -565,7 +620,7 @@ class ChunkProcessor:
             id=id,
         )
 
-    def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict:
+    def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary":
         prompt_tokens = 0
         completion_tokens = 0
         ## anthropic prompt caching information ##
@@ -623,8 +678,8 @@ class ChunkProcessor:
         return reasoning_tokens
 
     @staticmethod
-    def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
-        usage_chunk: Usage | dict[str, Any] | None = None
+    def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None:
+        usage_chunk: Usage | None = None
         if hasattr(chunk, "usage") and chunk.usage is not None:
             usage_chunk = chunk.usage
         elif "usage" in chunk:
@@ -640,7 +695,7 @@ class ChunkProcessor:
 
     def _calculate_usage_per_chunk(
         self,
-        chunks: list[dict[str, Any] | ModelResponse],
+        chunks: Sequence["_UsageBearingChunk | ModelResponse"],
     ) -> "UsagePerChunk":
         from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
             UsagePerChunk,
@@ -721,13 +776,7 @@ class ChunkProcessor:
                         "web_search_requests",
                     )
 
-                prompt_tokens_details = (
-                    cast(
-                        PromptTokensDetailsWrapper | None,
-                        usage_chunk_dict["prompt_tokens_details"],
-                    )
-                    or prompt_tokens_details
-                )
+                prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details
 
                 cache_creation_token_details = capture_cache_creation_token_details(
                     prompt_tokens_details, cache_creation_token_details
@@ -758,7 +807,7 @@ class ChunkProcessor:
 
     @staticmethod
     def _reset_anthropic_cursor_completion_tokens(
-        chunks: list[dict[str, Any] | ModelResponse],
+        chunks: Sequence["_UsageBearingChunk | ModelResponse"],
         completion_tokens: int,
         completion_usage_updates: int,
     ) -> int:
@@ -797,7 +846,7 @@ class ChunkProcessor:
 
     def calculate_usage(
         self,
-        chunks: list[dict[str, Any] | ModelResponse],
+        chunks: Sequence["_UsageBearingChunk | ModelResponse"],
         model: str,
         completion_output: str,
         messages: list | None = None,
@@ -851,8 +900,8 @@ class ChunkProcessor:
             setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens)  # for anthropic
         if completion_tokens_details is not None:
             if isinstance(completion_tokens_details, CompletionTokensDetails):
-                returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
-                    **completion_tokens_details.model_dump()
+                returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate(
+                    completion_tokens_details.model_dump()
                 )
             else:
                 returned_usage.completion_tokens_details = completion_tokens_details

litellm/llms/anthropic/common_utils.py

@@ -4,9 +4,12 @@ This file contains common utils for anthropic calls.
 
 import copy
 import re
-from typing import Any, Final
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import Any, Final, Literal
 
 import httpx
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 
 import litellm
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -1057,6 +1060,152 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any
     return out
 
 
+class _ReplayedSearchQuery(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    query: str = ""
+
+
+class _ReplayedWebSearchResult(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    type: Literal["web_search_result"]
+    url: str = ""
+    title: str = ""
+    snippet: str = ""
+    encrypted_content: str = ""
+
+
+class _ReplayedWebSearchToolResult(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    type: Literal["web_search_tool_result"]
+    tool_use_id: str
+    content: tuple[_ReplayedWebSearchResult, ...]
+
+
+class _ReplayedServerToolUse(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    type: Literal["server_tool_use"]
+    id: str
+    input: _ReplayedSearchQuery = _ReplayedSearchQuery()
+
+
+class _TextBlock(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+_WEB_SEARCH_TOOL_RESULT_ADAPTER: Final = TypeAdapter(_ReplayedWebSearchToolResult)
+_SERVER_TOOL_USE_ADAPTER: Final = TypeAdapter(_ReplayedServerToolUse)
+
+
+def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchToolResult | None:
+    """
+    The parsed block when it is a ``web_search_tool_result`` carrying no
+    ``encrypted_content``, else None for anything Anthropic itself issued.
+
+    An empty ``content`` list is flattenable too. It is what the interceptor emits
+    when a search legitimately returns nothing and when a search raises, and it
+    carries neither evidence to preserve nor an ``encrypted_content`` to respect,
+    so leaving it in place only buys the 400 this whole function exists to avoid.
+    """
+    try:
+        parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block)
+    except ValidationError:
+        return None
+    if any(result.encrypted_content for result in parsed.content):
+        return None
+    return parsed
+
+
+def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
+    try:
+        return _SERVER_TOOL_USE_ADAPTER.validate_python(block)
+    except ValidationError:
+        return None
+
+
+def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str:
+    header: Final = f"Web search results for '{query}':" if query else "Web search results:"
+    if not results:
+        return f"{header}\n\nNo results were returned."
+    body: Final = "\n\n".join(
+        "\n".join(
+            line
+            for line in (
+                f"Title: {result.title}" if result.title else "",
+                f"URL: {result.url}" if result.url else "",
+                f"Snippet: {result.snippet}" if result.snippet else "",
+            )
+            if line
+        )
+        for result in results
+    )
+    return f"{header}\n\n{body}" if body else header
+
+
+def _rewrite_replayed_web_search_block(
+    block: object,
+    flattenable: Mapping[str, _ReplayedWebSearchToolResult],
+    queries: Mapping[str, str],
+) -> object | None:
+    parsed_result: Final = _flattenable_web_search_tool_result(block)
+    if parsed_result is not None:
+        return _TextBlock(
+            text=_render_web_search_results(queries.get(parsed_result.tool_use_id, ""), parsed_result.content)
+        ).model_dump()
+    parsed_use: Final = _replayed_server_tool_use(block)
+    if parsed_use is not None and parsed_use.id in flattenable:
+        return None
+    return block
+
+
+def _flatten_web_search_results_in_message(message: object) -> object:
+    if not isinstance(message, Mapping) or not isinstance(message.get("content"), Sequence):
+        return message
+    content: Final = message["content"]
+    if isinstance(content, str):
+        return message
+    flattenable: Final = MappingProxyType(
+        {
+            parsed.tool_use_id: parsed
+            for parsed in (_flattenable_web_search_tool_result(block) for block in content)
+            if parsed is not None
+        }
+    )
+    if not flattenable:
+        return message
+    queries: Final = MappingProxyType(
+        {
+            parsed.id: parsed.input.query
+            for parsed in (_replayed_server_tool_use(block) for block in content)
+            if parsed is not None
+        }
+    )
+    rewritten: Final = tuple(_rewrite_replayed_web_search_block(block, flattenable, queries) for block in content)
+    return {**message, "content": [b for b in rewritten if b is not None]}  # mutable-ok: JSON wire format
+
+
+def flatten_unencrypted_web_search_results_in_anthropic_messages(  # mutable-ok: as sibling sanitizers
+    messages: list[Any],
+) -> list[Any]:
+    """
+    Return a new message list with replayed ``web_search_tool_result`` blocks that
+    carry no ``encrypted_content`` rewritten into plain ``text`` blocks holding the
+    same title / url / snippet evidence.
+
+    ``encrypted_content`` is an opaque blob only Anthropic's own search backend can
+    mint, so blocks synthesized by LiteLLM (websearch interception against a search
+    provider) are rejected with ``Invalid encrypted_content in search_result block``
+    when a native client loops them back as history. Flattening them keeps the
+    evidence in the conversation instead of 400ing the follow-up turn, and leaves
+    genuine Anthropic-issued blocks untouched.
+    """
+    return [_flatten_web_search_results_in_message(m) for m in messages]  # mutable-ok: JSON wire format
+
+
 def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
     openai_headers: Final = {}
     if "anthropic-ratelimit-requests-limit" in headers:

litellm/llms/anthropic/experimental_pass_through/adapters/handler.py

@@ -1,8 +1,9 @@
-from collections.abc import AsyncIterator, Coroutine, Iterator
+from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
 from typing import (
     TYPE_CHECKING,
     Any,
     Final,
+    TypeAlias,
     cast,
 )
 
@@ -33,8 +34,12 @@ if TYPE_CHECKING:
 # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
 ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"})
 
+_AnthropicMessages: TypeAlias = "list[dict[str, object]]"
+_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None"
+_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None"
 
-def _messages_have_compaction_block(messages: list[dict]) -> bool:
+
+def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool:
     """Return True when any message carries a ``compaction`` content block."""
     for msg in messages:
         content = msg.get("content")
@@ -54,8 +59,10 @@ def _proxy_router_fallback() -> "Router | None":
     return _proxy_router
 
 
-def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
-    """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise.
+def _extract_proxy_litellm_metadata(
+    kwargs: Mapping[str, object],
+) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]":
+    """Return ``(kwargs["litellm_metadata"], its user_api_key_auth)`` when it's a dict; ``(None, None)`` otherwise.
 
     The proxy attaches its auth/spend-attribution fields (``user_api_key``,
     ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth``
@@ -68,18 +75,19 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] |
     """
     litellm_metadata: Final = kwargs.get("litellm_metadata")
     if not isinstance(litellm_metadata, dict):
-        return None
-    return litellm_metadata
+        return None, None
+    user_api_key_auth: Final[UserAPIKeyAuth | None] = litellm_metadata.get("user_api_key_auth")
+    return litellm_metadata, user_api_key_auth
 
 
 async def _prepare_context_managed_request(
     *,
     model: str,
-    messages: list[dict],
-    tools: list[dict] | None,
-    system: Any | None,
-    context_management_spec: Any,
-    litellm_metadata: dict | None,
+    messages: _AnthropicMessages,
+    tools: list[dict[str, object]] | None,
+    system: _AnthropicSystem,
+    context_management_spec: _ContextManagementSpec,
+    litellm_metadata: dict[str, object] | None,
     additional_drop_params: list[str] | None,
     llm_router: "Router | None",
     user_api_key_auth: "UserAPIKeyAuth | None" = None,
@@ -102,11 +110,11 @@ async def _prepare_context_managed_request(
 
     if polyfill_will_run:
         history_result: PolyfillResult | None = None
-        working_messages: list[dict] = messages
-        working_system: Any | None = system
+        working_messages: _AnthropicMessages = messages
+        working_system: _AnthropicSystem = system
     else:
         history_result = apply_client_compaction_block_history(
-            messages=cast(list[dict[str, Any]], messages),
+            messages=messages,
             system=system,
         )
         working_messages = history_result.messages if history_result is not None else messages
@@ -136,7 +144,7 @@ async def _prepare_context_managed_request(
     # to non-Anthropic backends that would reject them.
     if polyfill_will_run and history_result is None:
         history_result = apply_client_compaction_block_history(
-            messages=cast(list[dict[str, Any]], messages),
+            messages=messages,
             system=system,
         )
     return history_result
@@ -144,7 +152,7 @@ async def _prepare_context_managed_request(
 
 def _polyfill_will_run(
     *,
-    context_management_spec: Any,
+    context_management_spec: _ContextManagementSpec,
     additional_drop_params: list[str] | None,
 ) -> bool:
     """Return True when ``compact_20260112`` will run via the polyfill dispatcher.
@@ -171,7 +179,7 @@ def _polyfill_will_run(
 
 def _spec_has_non_compact_edits(
     *,
-    context_management_spec: Any,
+    context_management_spec: _ContextManagementSpec,
     additional_drop_params: list[str] | None,
 ) -> bool:
     """Return True when the spec includes edits other than ``compact_20260112``.
@@ -209,9 +217,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N
 
 def _normalize_spec_edits(
     *,
-    context_management_spec: Any,
+    context_management_spec: _ContextManagementSpec,
     additional_drop_params: list[str] | None,
-) -> list[dict[str, Any]] | None:
+) -> list[dict[str, object]] | None:
     """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run.
 
     Delegates spec-shape normalization to the dispatcher's ``_normalize_spec``
@@ -236,11 +244,11 @@ def _normalize_spec_edits(
 async def _run_polyfill_if_enabled(
     *,
     model: str,
-    messages: list[dict],
-    tools: list[dict] | None,
-    system: Any | None,
-    context_management_spec: Any,
-    litellm_metadata: dict | None,
+    messages: _AnthropicMessages,
+    tools: list[dict[str, object]] | None,
+    system: _AnthropicSystem,
+    context_management_spec: _ContextManagementSpec,
+    litellm_metadata: dict[str, object] | None,
     additional_drop_params: list[str] | None,
     llm_router: "Router | None",
     user_api_key_auth: "UserAPIKeyAuth | None" = None,
@@ -306,7 +314,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     def _route_openai_thinking_to_responses_api_if_needed(
         completion_kwargs: dict[str, Any],
         *,
-        thinking: dict[str, Any] | None,
+        thinking: Mapping[str, object] | None,
     ) -> None:
         """
         When users call `litellm.anthropic.messages.*` with a non-Anthropic model and
@@ -407,12 +415,12 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     def _prepare_completion_kwargs(
         *,
         max_tokens: int,
-        messages: list[dict],
+        messages: _AnthropicMessages,
         model: str,
         metadata: dict | None = None,
         stop_sequences: list[str] | None = None,
         stream: bool | None = False,
-        system: str | list[dict[str, Any]] | None = None,
+        system: _AnthropicSystem = None,
         temperature: float | None = None,
         thinking: dict | None = None,
         tool_choice: dict | None = None,
@@ -420,7 +428,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         top_k: int | None = None,
         top_p: float | None = None,
         output_format: dict | None = None,
-        extra_kwargs: dict[str, Any] | None = None,
+        extra_kwargs: Mapping[str, object] | None = None,
     ) -> tuple[dict[str, Any], dict[str, str]]:
         """Prepare kwargs for litellm.completion/acompletion.
 
@@ -433,7 +441,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
             Logging as LiteLLMLoggingObject,
         )
 
-        request_data: Final = {
+        request_data: Final[dict[str, object]] = {
             "model": model,
             "messages": messages,
             "max_tokens": max_tokens,
@@ -528,7 +536,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     @staticmethod
     async def async_anthropic_messages_handler(
         max_tokens: int,
-        messages: list[dict],
+        messages: _AnthropicMessages,
         model: str,
         metadata: dict | None = None,
         stop_sequences: list[str] | None = None,
@@ -537,7 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         temperature: float | None = None,
         thinking: dict | None = None,
         tool_choice: dict | None = None,
-        tools: list[dict] | None = None,
+        tools: list[dict[str, object]] | None = None,
         top_k: int | None = None,
         top_p: float | None = None,
         output_format: dict | None = None,
@@ -551,10 +559,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
             requested_router if requested_router is not None else _proxy_router_fallback()
         )
 
-        proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs)
-        user_api_key_auth: Final[UserAPIKeyAuth | None] = (
-            proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None
-        )
+        proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
 
         polyfill_result: Final = await _prepare_context_managed_request(
             model=model,
@@ -618,7 +623,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
     @staticmethod
     def anthropic_messages_handler(
         max_tokens: int,
-        messages: list[dict],
+        messages: _AnthropicMessages,
         model: str,
         metadata: dict | None = None,
         stop_sequences: list[str] | None = None,
@@ -627,7 +632,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         temperature: float | None = None,
         thinking: dict | None = None,
         tool_choice: dict | None = None,
-        tools: list[dict] | None = None,
+        tools: list[dict[str, object]] | None = None,
         top_k: int | None = None,
         top_p: float | None = None,
         output_format: dict | None = None,
@@ -688,10 +693,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
         if context_management is None and not _messages_have_compaction_block(messages):
             polyfill_result: PolyfillResult | None = None
         else:
-            proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs)
-            user_api_key_auth: Final[UserAPIKeyAuth | None] = (
-                proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None
-            )
+            proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
             polyfill_result = run_async_function(
                 _prepare_context_managed_request,
                 model=model,

litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py

@@ -348,6 +348,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
             merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits))
         return self._augment_message_delta_usage(merged_chunk)
 
+    def _handle_choiceless_chunk(self, chunk: "ModelResponseStream") -> bool:
+        """Consume an OpenAI-compatible chunk that carries no ``choices``.
+
+        ``choices`` is legitimately empty on metadata-only chunks; the final
+        usage chunk emitted when ``stream_options.include_usage`` is set is the
+        common case (vLLM and other OpenAI-compatible servers do this). Such a
+        chunk carries no content-block information, so the caller must not run
+        the content-block state machine over it.
+
+        Returns True when a merged ``message_delta`` was queued (usage folded
+        into the held stop-reason chunk); False when the chunk should be
+        skipped entirely.
+        """
+        if self.holding_stop_reason_chunk is not None and _optional_attr(chunk, "usage") is not None:
+            self.chunk_queue.append(self._merge_usage_into_held_stop_reason_chunk(chunk))
+            self.queued_usage_chunk = True
+            self.holding_stop_reason_chunk = None
+            return True
+        return False
+
     def _ensure_context_management_attached(self, message_delta_chunk: MessageBlockDelta) -> MessageBlockDelta:
         """Attach ``context_management`` to a ``message_delta`` chunk if
         ``self.applied_edits`` is non-empty and the chunk does not already
@@ -509,6 +529,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
                 if chunk == "None" or chunk is None:
                     raise Exception
 
+                if not getattr(chunk, "choices", None):
+                    if self._handle_choiceless_chunk(chunk):
+                        return self.chunk_queue.popleft()
+                    continue
+
                 should_start_new_block = self._should_start_new_content_block(chunk)
                 is_opening_first_block = self.sent_content_block_start is False
                 if is_opening_first_block and self._is_blank_delta(chunk):
@@ -732,6 +757,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
                 if chunk == "None" or chunk is None:
                     raise Exception
 
+                if not getattr(chunk, "choices", None):
+                    if self._handle_choiceless_chunk(chunk):
+                        return self.chunk_queue.popleft()
+                    continue
+
                 should_start_new_block = self._should_start_new_content_block(chunk)
                 is_opening_first_block = self.sent_content_block_start is False
                 if is_opening_first_block and self._is_blank_delta(chunk):

litellm/llms/anthropic/experimental_pass_through/messages/handler.py

@@ -14,6 +14,7 @@ from typing import Any, Final, cast
 import litellm
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.anthropic.common_utils import (
+    flatten_unencrypted_web_search_results_in_anthropic_messages,
     sanitize_tool_use_ids_in_anthropic_messages,
     strip_empty_text_blocks_from_anthropic_messages,
 )
@@ -222,6 +223,7 @@ async def anthropic_messages(
     # Replay of cross-provider tool history (e.g. kimi -> Anthropic) may carry
     # ids like ``functions.Bash:0`` that violate Anthropic's id pattern.
     messages = sanitize_tool_use_ids_in_anthropic_messages(messages)
+    messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages)
 
     from litellm.integrations.anthropic_cache_control_hook import (
         AnthropicCacheControlHook,
@@ -413,6 +415,7 @@ def anthropic_messages_handler(
     if not kwargs.pop("_litellm_messages_presanitized", False):
         messages = strip_empty_text_blocks_from_anthropic_messages(messages)
         messages = sanitize_tool_use_ids_in_anthropic_messages(messages)
+        messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages)
 
     from litellm.integrations.anthropic_cache_control_hook import (
         AnthropicCacheControlHook,

litellm/llms/azure/assistants.py

@@ -1,8 +1,9 @@
 from collections.abc import Coroutine, Iterable
-from typing import Any, Final, Literal
+from typing import Any, Final, Literal, TypedDict
 
 import httpx
 from openai import AsyncAzureOpenAI, AzureOpenAI
+from openai.types.shared_params.metadata import Metadata
 from typing_extensions import overload
 
 from ...types.llms.openai import (
@@ -22,6 +23,16 @@ from ...types.llms.openai import (
 from .common_utils import BaseAzureLLM
 
 
+class _RunThreadStreamData(TypedDict):
+    thread_id: str
+    assistant_id: str
+    additional_instructions: str | None
+    instructions: str | None
+    metadata: Metadata | None
+    model: str | None
+    tools: Iterable[AssistantToolParam] | None
+
+
 class AzureAssistantsAPI(BaseAzureLLM):
     def __init__(self) -> None:
         super().__init__()
@@ -212,9 +223,9 @@ class AzureAssistantsAPI(BaseAzureLLM):
         response_obj: OpenAIMessage | None = None
         if getattr(thread_message, "status", None) is None:
             thread_message.status = "completed"
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         else:
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         return response_obj
 
     # fmt: off
@@ -301,9 +312,9 @@ class AzureAssistantsAPI(BaseAzureLLM):
         response_obj: OpenAIMessage | None = None
         if getattr(thread_message, "status", None) is None:
             thread_message.status = "completed"
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         else:
-            response_obj = OpenAIMessage(**thread_message.dict())
+            response_obj = OpenAIMessage.model_validate(thread_message.dict())
         return response_obj
 
     async def async_get_messages(
@@ -443,7 +454,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         message_thread: Final = await openai_client.beta.threads.create(**data)
 
-        return Thread(**message_thread.dict())
+        return Thread.model_validate(message_thread.dict())
 
     # fmt: off
 
@@ -539,7 +550,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         message_thread: Final = azure_openai_client.beta.threads.create(**data)
 
-        return Thread(**message_thread.dict())
+        return Thread.model_validate(message_thread.dict())
 
     async def async_get_thread(
         self,
@@ -566,7 +577,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id)
 
-        return Thread(**response.dict())
+        return Thread.model_validate(response.dict())
 
     # fmt: off
 
@@ -642,7 +653,7 @@ class AzureAssistantsAPI(BaseAzureLLM):
 
         response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id)
 
-        return Thread(**response.dict())
+        return Thread.model_validate(response.dict())
 
     # def delete_thread(self):
     #     pass
@@ -730,7 +741,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
         event_handler: AssistantEventHandler | None,
         litellm_params: dict | None = None,
     ) -> AssistantStreamManager[AssistantEventHandler]:
-        data: Final[dict[str, Any]] = {
+        stream_fn: Final = client.beta.threads.runs.stream
+        base_data: Final[_RunThreadStreamData] = {
             "thread_id": thread_id,
             "assistant_id": assistant_id,
             "additional_instructions": additional_instructions,
@@ -740,8 +752,8 @@ class AzureAssistantsAPI(BaseAzureLLM):
             "tools": tools,
         }
         if event_handler is not None:
-            data["event_handler"] = event_handler
-        return client.beta.threads.runs.stream(**data)
+            return stream_fn(**base_data, event_handler=event_handler)
+        return stream_fn(**base_data)
 
     # fmt: off
 

litellm/llms/bedrock/batches/transformation.py

@@ -12,7 +12,6 @@ from litellm.litellm_core_utils.cloud_storage_security import (
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
-from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.bedrock import (
     BedrockCreateBatchRequest,
     BedrockCreateBatchResponse,
@@ -29,7 +28,7 @@ from litellm.types.llms.openai import (
 from litellm.types.utils import LiteLLMBatch, LlmProviders
 
 from ..base_aws_llm import BaseAWSLLM
-from ..common_utils import CommonBatchFilesUtils
+from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id
 
 # Bedrock batch input files are uploaded as
 # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see
@@ -200,7 +199,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
         )
 
         # Add optional KMS encryption key ID if provided
-        s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID")
+        s3_encryption_key_id = resolve_s3_encryption_key_id(
+            litellm_params=litellm_params,
+            optional_params=optional_params,
+        )
         if s3_encryption_key_id:
             s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id
 

litellm/llms/bedrock/common_utils.py

@@ -26,7 +26,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import (
 )
 from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
-from litellm.secret_managers.main import get_secret
+from litellm.secret_managers.main import get_secret, get_secret_str
 
 if TYPE_CHECKING:
     from litellm.types.llms.openai import AllMessageValues
@@ -1304,6 +1304,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> list[str]:
     return []
 
 
+def resolve_s3_encryption_key_id(
+    litellm_params: Mapping[str, Any],
+    optional_params: Mapping[str, Any] | None = None,
+) -> str | None:
+    """
+    Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects.
+
+    Precedence: `s3_encryption_key_id` in litellm_params, then optional_params
+    (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var.
+    """
+    candidates: Final = tuple(
+        source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None
+    )
+    explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None)
+    return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID")
+
+
 class CommonBatchFilesUtils:
     """
     Common utilities for Bedrock batch and file operations.

litellm/llms/bedrock/files/transformation.py

@@ -2,7 +2,9 @@ import base64
 import json
 import os
 import time
-from collections.abc import Mapping, MutableMapping
+from collections.abc import Iterable, Mapping, MutableMapping
+from functools import cache
+from itertools import chain
 from types import MappingProxyType
 from typing import Any, Final
 from urllib.parse import unquote
@@ -10,7 +12,7 @@ from urllib.parse import unquote
 import httpx
 from httpx import Headers, Response
 from openai.types.file_deleted import FileDeleted
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, TypeAdapter
 
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
@@ -26,12 +28,16 @@ from litellm.litellm_core_utils.cloud_storage_security import (
     split_configured_cloud_bucket_name,
     validate_managed_cloud_file_id,
 )
-from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    extract_file_data,
+    text_completion_prompt_to_messages,
+)
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.files.transformation import (
     BaseFilesConfig,
     LiteLLMLoggingObj,
 )
+from litellm.types.llms.bedrock import BedrockBatchRecordKind
 from litellm.types.llms.openai import (
     AllMessageValues,
     CreateFileRequest,
@@ -41,12 +47,14 @@ from litellm.types.llms.openai import (
     OpenAICreateFileRequestOptionalParams,
     OpenAIFileObject,
     PathLike,
+    ResponseInputParam,
+    ResponsesAPIOptionalRequestParams,
 )
 from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums
 from litellm.utils import get_llm_provider
 
 from ..base_aws_llm import BaseAWSLLM
-from ..common_utils import BedrockError
+from ..common_utils import BedrockError, resolve_s3_encryption_key_id
 
 # litellm_params key used to hand the SigV4-signed GET headers from
 # `transform_file_content_request` to `validate_environment` (the only hook
@@ -55,6 +63,26 @@ from ..common_utils import BedrockError
 S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers"
 
 
+def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]:
+    return MappingProxyType(dict(items))
+
+
+# JSONL batch records are untyped json, so the `/v1/responses` fields are
+# validated into their concrete Responses API types before being handed to the
+# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't
+# define, which is what the bridge would ignore anyway. Built on first use
+# rather than at import: `ResponseInputParam` is a deep union and only batch
+# files carrying `/v1/responses` records need it.
+@cache
+def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]:
+    return TypeAdapter(str | ResponseInputParam)
+
+
+@cache
+def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]:
+    return TypeAdapter(ResponsesAPIOptionalRequestParams)
+
+
 class _BedrockS3RequestParams(BaseModel):
     """Typed view of the credential/region params the S3 GetObject path reads."""
 
@@ -303,41 +331,55 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
     # example; add others here as they adopt the same schema.
     CONVERSE_INVOKE_PROVIDERS = ("nova",)
 
-    # OpenAI batch URL that signals an embedding request. Per OpenAI Batch API
-    # spec, every JSONL record carries a `url` field; we use it as the
-    # authoritative signal to route the line to the embedding code path
-    # instead of inferring from the presence of `input` vs `messages`.
+    # OpenAI batch URLs that select which request shape a JSONL line carries.
+    # Per the OpenAI Batch API spec every record carries a `url`, so we use it
+    # as the authoritative routing signal instead of inferring from the
+    # presence of `input` vs `prompt` vs `messages`.
     OPENAI_EMBEDDINGS_URL = "/v1/embeddings"
+    OPENAI_TEXT_COMPLETIONS_URL = "/v1/completions"
+    OPENAI_RESPONSES_URL = "/v1/responses"
 
     @staticmethod
-    def _is_embedding_record(openai_jsonl_record: dict[str, Any]) -> bool:
+    def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind:
         """
-        Decide whether an OpenAI batch JSONL line is an embedding request.
+        Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries.
 
-        Precedence (strict - any explicit `url` short-circuits):
-          1. `url == "/v1/embeddings"` -> embedding. Authoritative per the
-             OpenAI Batch API spec.
-          2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT
-             embedding. We trust the caller's explicit signal even if the
-             body would otherwise suggest embedding; misrouting a chat
-             record into the embedding transformer would corrupt the
-             modelInput, while a chat-shaped body sent to the chat path
-             either succeeds or fails cleanly inside that transformer.
-          3. `url` missing/empty -> fall back to body shape. Requires
-             `input` present AND `messages` absent so a malformed record
-             carrying both keys routes to the chat path (safer default:
-             Anthropic transforms ignore unknown top-level keys, whereas
-             the embedding transformer would silently drop the messages).
+        Precedence (strict - any recognized `url` short-circuits):
+          1. A `url` matching a supported endpoint wins. Authoritative per the
+             OpenAI Batch API spec, which requires it on every record.
+          2. Any other non-empty `url` -> chat. We trust the caller's explicit
+             signal rather than re-deriving it from the body, and an
+             unexpectedly-shaped body fails cleanly inside the chat
+             transformer instead of being silently misrouted.
+          3. `url` missing/empty -> fall back to body shape. `messages` wins
+             over the other keys so a malformed record carrying several of
+             them keeps its conversation instead of having it dropped, and a
+             bare `input` stays an embedding for backwards compatibility
+             (that ambiguity with `/v1/responses` is only resolvable from
+             `url`).
         """
-        url: Final = openai_jsonl_record.get("url")
-        if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL:
-            return True
-        if url:
-            return False
-        body: Final = openai_jsonl_record.get("body", {})
-        if not isinstance(body, dict):
-            return False
-        return "input" in body and "messages" not in body
+        match openai_jsonl_record.get("url"):
+            case BedrockFilesConfig.OPENAI_EMBEDDINGS_URL:
+                return BedrockBatchRecordKind.EMBEDDING
+            case BedrockFilesConfig.OPENAI_TEXT_COMPLETIONS_URL:
+                return BedrockBatchRecordKind.TEXT_COMPLETION
+            case BedrockFilesConfig.OPENAI_RESPONSES_URL:
+                return BedrockBatchRecordKind.RESPONSES
+            case None | "":
+                pass
+            case _:
+                return BedrockBatchRecordKind.CHAT
+
+        body: Final = openai_jsonl_record.get("body")
+        if not isinstance(body, Mapping):
+            return BedrockBatchRecordKind.CHAT
+        if "messages" in body:
+            return BedrockBatchRecordKind.CHAT
+        if "prompt" in body:
+            return BedrockBatchRecordKind.TEXT_COMPLETION
+        if "input" in body:
+            return BedrockBatchRecordKind.EMBEDDING
+        return BedrockBatchRecordKind.CHAT
 
     # Identifier for the Bedrock Titan v2 InvokeModel body schema as stored
     # in `model_prices_and_context_window.json`. Centralized so future
@@ -544,9 +586,84 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
         )
         return dict(titan_config._transform_request(input=input_text, inference_params=inference_params))
 
+    @staticmethod
+    def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]:
+        """
+        Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body.
+
+        Bedrock batch `modelInput` is the model's InvokeModel/Converse body, and
+        no Bedrock batch model takes a bare `prompt`, so the wrapping that
+        `litellm.text_completion` does in real time has to happen here too.
+        """
+        prompt: Final = openai_request_body.get("prompt")
+        if prompt is None:
+            raise ValueError(
+                "Batch record for /v1/completions is missing required `prompt` field: "
+                f"model={openai_request_body.get('model', '')}"
+            )
+        return _frozen_mapping(
+            chain(
+                ((key, value) for key, value in openai_request_body.items() if key != "prompt"),
+                (("messages", text_completion_prompt_to_messages(prompt)),),
+            )
+        )
+
+    @staticmethod
+    def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]:
+        """
+        Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body.
+
+        Delegates to the same Responses-to-Chat bridge the real-time path uses
+        for providers without a native Responses API (which is every Bedrock
+        model), so `input`, `instructions`, `max_output_tokens` and the tool
+        params translate identically in batch and real time. The bridge always
+        emits a `tools` key; an empty one is dropped rather than shipped as an
+        empty array inside `modelInput`.
+        """
+        from litellm.responses.litellm_completion_transformation.transformation import (
+            LiteLLMCompletionResponsesConfig,
+        )
+
+        responses_input: Final = openai_request_body.get("input")
+        if responses_input is None:
+            raise ValueError(
+                "Batch record for /v1/responses is missing required `input` field: "
+                f"model={openai_request_body.get('model', '')}"
+            )
+        chat_body: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
+            model=openai_request_body.get("model", ""),
+            input=_responses_input_adapter().validate_python(responses_input),
+            responses_api_request=_responses_request_adapter().validate_python(
+                _frozen_mapping(
+                    (key, value) for key, value in openai_request_body.items() if key not in ("model", "input")
+                )
+            ),
+            metadata=openai_request_body.get("metadata"),
+        )
+        return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value)
+
+    @staticmethod
+    def _transform_batch_body_to_chat_body(
+        openai_request_body: Mapping[str, Any],
+        record_kind: BedrockBatchRecordKind,
+    ) -> Mapping[str, Any]:
+        """
+        Normalize a non-embedding batch body to the Chat Completions shape the
+        per-provider Bedrock transformations expect.
+        """
+        match record_kind:
+            case BedrockBatchRecordKind.TEXT_COMPLETION:
+                return BedrockFilesConfig._transform_text_completion_body_to_chat_body(openai_request_body)
+            case BedrockBatchRecordKind.RESPONSES:
+                return BedrockFilesConfig._transform_responses_body_to_chat_body(openai_request_body)
+            case BedrockBatchRecordKind.CHAT:
+                return openai_request_body
+            case BedrockBatchRecordKind.EMBEDDING:
+                raise ValueError("Embedding batch records do not have a chat-completion equivalent")
+
     def _map_openai_to_bedrock_params(
         self,
-        openai_request_body: dict[str, Any],
+        openai_request_body: Mapping[str, Any],
         provider: str | None = None,
     ) -> dict[str, Any]:
         """
@@ -658,14 +775,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             provider = self.get_bedrock_invoke_provider(model)
 
             # Route to the embedding transformer when the OpenAI batch line
-            # targets /v1/embeddings; otherwise fall back to the existing
-            # chat-completion path. We branch here (rather than inside
+            # targets /v1/embeddings; every other endpoint shape is normalized
+            # to chat completions first. We branch here (rather than inside
             # `_map_openai_to_bedrock_params`) so the chat helper keeps its
             # narrow contract and the embedding helper can evolve independently.
-            if self._is_embedding_record(_openai_jsonl_content):
+            record_kind = self._classify_batch_record(_openai_jsonl_content)
+            if record_kind is BedrockBatchRecordKind.EMBEDDING:
                 model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body)
             else:
-                model_input = self._map_openai_to_bedrock_params(openai_request_body=openai_body, provider=provider)
+                model_input = self._map_openai_to_bedrock_params(
+                    openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind),
+                    provider=provider,
+                )
 
             # Create Bedrock batch record
             record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}")
@@ -733,6 +854,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             content=file_content,
             api_base=api_base,
             optional_params=optional_params,
+            s3_encryption_key_id=resolve_s3_encryption_key_id(
+                litellm_params=litellm_params,
+                optional_params=optional_params,
+            ),
         )
 
         litellm_params["upload_url"] = api_base
@@ -750,6 +875,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
         content: str,
         api_base: str,
         optional_params: dict,
+        s3_encryption_key_id: str | None = None,
     ) -> tuple[dict, str]:
         """
         Sign S3 PUT request using the same proven logic as S3Logger.
@@ -782,12 +908,25 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
         content_hash: Final = hashlib.sha256(content.encode("utf-8")).hexdigest()
 
         # Prepare headers with required S3 headers (same as s3_v2.py)
-        request_headers: Final = {
-            "Content-Type": "application/json",  # JSONL files are JSON content
-            "x-amz-content-sha256": content_hash,  # REQUIRED by S3
-            "Content-Language": "en",
-            "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
-        }
+        sse_headers: Final = (
+            MappingProxyType(
+                {
+                    "x-amz-server-side-encryption": "aws:kms",
+                    "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id,
+                }
+            )
+            if s3_encryption_key_id
+            else MappingProxyType({})
+        )
+        request_headers: Final = MappingProxyType(
+            {
+                "Content-Type": "application/json",  # JSONL files are JSON content
+                "x-amz-content-sha256": content_hash,  # REQUIRED by S3
+                "Content-Language": "en",
+                "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
+                **sse_headers,
+            }
+        )
 
         # Use requests.Request to prepare the request (same pattern as s3_v2.py)
         req: Final = requests.Request("PUT", api_base, data=content, headers=request_headers)

litellm/llms/custom_httpx/http_handler.py

@@ -8,11 +8,12 @@ import sys
 import threading
 import time
 from collections.abc import Callable, Mapping
+from http.cookiejar import CookieJar, DefaultCookiePolicy
 from typing import TYPE_CHECKING, Any, Final, Optional
 
 import certifi
 import httpx
-from aiohttp import ClientSession, TCPConnector
+from aiohttp import ClientSession, DummyCookieJar, TCPConnector
 from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport
 from httpx._types import RequestFiles
 
@@ -144,6 +145,15 @@ 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 blocked_cookie_jar() -> CookieJar:
+    """A jar that stores no response cookie and sends none, for httpx clients.
+
+    LiteLLM's outbound clients are pooled and shared by every caller, so a cookie one
+    upstream sets would be replayed to every other upstream on a matching domain.
+    """
+    return CookieJar(policy=DefaultCookiePolicy(allowed_domains=()))
+
+
 _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS: Final = 5.0
 _STREAMING_ERROR_BODY_READ_EXECUTOR: Final = concurrent.futures.ThreadPoolExecutor(
     max_workers=50,
@@ -587,6 +597,7 @@ class AsyncHTTPHandler:
             verify=ssl_config,
             cert=cert,
             headers=default_headers,
+            cookies=blocked_cookie_jar(),
             follow_redirects=True,
         )
 
@@ -1063,6 +1074,7 @@ class AsyncHTTPHandler:
         def session_factory() -> ClientSession:
             return ClientSession(
                 connector=TCPConnector(**transport_connector_kwargs),
+                cookie_jar=DummyCookieJar(),
                 trust_env=trust_env,
             )
 
@@ -1132,6 +1144,7 @@ class HTTPHandler:
             verify=ssl_config,
             cert=cert,
             headers=default_headers,
+            cookies=blocked_cookie_jar(),
             follow_redirects=True,
         )
 

litellm/llms/vertex_ai/batches/handler.py

@@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
     _get_httpx_client,
     get_async_httpx_client,
 )
-from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
+from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
 from litellm.types.llms.openai import CreateBatchRequest
 from litellm.types.llms.vertex_ai import (
@@ -98,9 +98,6 @@ class VertexAIBatchPrediction(VertexLLM):
             data=json.dumps(vertex_batch_request),
         )
 
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
-
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
             response=_json_response
@@ -130,8 +127,6 @@ class VertexAIBatchPrediction(VertexLLM):
                 error_body[:1000],
             )
             raise
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
 
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -243,7 +238,9 @@ class VertexAIBatchPrediction(VertexLLM):
         )
 
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -293,7 +290,9 @@ class VertexAIBatchPrediction(VertexLLM):
             headers=headers,
         )
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -366,7 +365,9 @@ class VertexAIBatchPrediction(VertexLLM):
         )
 
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response: Final = (
@@ -391,7 +392,9 @@ class VertexAIBatchPrediction(VertexLLM):
             params=params,
         )
         if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
+            raise VertexAIError(
+                status_code=response.status_code, message=f"Error: {response.status_code} {response.text}"
+            )
 
         _json_response: Final = response.json()
         vertex_batch_response: Final = (
@@ -461,7 +464,7 @@ class VertexAIBatchPrediction(VertexLLM):
 
         sync_handler: Final = _get_httpx_client()
         try:
-            response: Final = sync_handler.post(
+            sync_handler.post(
                 url=api_base,
                 headers=headers,
                 data=json.dumps({}),
@@ -475,9 +478,6 @@ class VertexAIBatchPrediction(VertexLLM):
             )
             raise
 
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
-
         # HTTPHandler.get() does not accept a timeout parameter
         retrieve_response: Final = sync_handler.get(
             url=retrieve_api_base,
@@ -489,7 +489,10 @@ class VertexAIBatchPrediction(VertexLLM):
                 retrieve_response.status_code,
                 retrieve_response.text[:1000],
             )
-            raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}")
+            raise VertexAIError(
+                status_code=retrieve_response.status_code,
+                message=f"Error: {retrieve_response.status_code} {retrieve_response.text}",
+            )
 
         _json_response: Final = retrieve_response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
@@ -508,7 +511,7 @@ class VertexAIBatchPrediction(VertexLLM):
             llm_provider=litellm.LlmProviders.VERTEX_AI,
         )
         try:
-            response: Final = await client.post(
+            await client.post(
                 url=api_base,
                 headers=headers,
                 data=json.dumps({}),
@@ -521,8 +524,6 @@ class VertexAIBatchPrediction(VertexLLM):
                 e.response.text[:1000],
             )
             raise
-        if response.status_code != 200:
-            raise Exception(f"Error: {response.status_code} {response.text}")
 
         # AsyncHTTPHandler.get() does not accept a timeout parameter
         retrieve_response: Final = await client.get(
@@ -535,7 +536,10 @@ class VertexAIBatchPrediction(VertexLLM):
                 retrieve_response.status_code,
                 retrieve_response.text[:1000],
             )
-            raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}")
+            raise VertexAIError(
+                status_code=retrieve_response.status_code,
+                message=f"Error: {retrieve_response.status_code} {retrieve_response.text}",
+            )
 
         _json_response: Final = retrieve_response.json()
         vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(

litellm/llms/vertex_ai/batches/transformation.py

@@ -1,7 +1,9 @@
 from typing import Any, Final
+from urllib.parse import unquote
 
 from litellm._uuid import uuid
 from litellm.llms.vertex_ai.common_utils import (
+    VertexAIError,
     _convert_vertex_datetime_to_openai_datetime,
 )
 from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
@@ -199,16 +201,40 @@ class VertexAIBatchTransformation:
 
         gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8
         returns: "publishers/google/models/gemini-1.5-flash-001"
+
+        Raises a 400 `VertexAIError` when the uri carries no parseable model path.
         """
-        from urllib.parse import unquote
-
-        decoded_uri: Final = unquote(gcs_file_uri)
-
-        model_path: Final = decoded_uri.split("publishers/")[1]
-        parts: Final = model_path.split("/")
-        model: Final = f"publishers/{'/'.join(parts[:3])}"
+        model: Final = cls._parse_model_from_gcs_file(gcs_file_uri)
+        if model is None:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch creation requires the model to be part of `input_file_id`, but "
+                    f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' path segment. "
+                    "Either upload the input file through LiteLLM (POST /v1/files with "
+                    "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or "
+                    "pass a uri of the form "
+                    "gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file>"
+                ),
+            )
         return model
 
+    @classmethod
+    def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None:
+        """
+        Returns the `publishers/<publisher>/models/<model>` path from a gcs uri, or None if the uri
+        does not contain one.
+        """
+        _, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
+        if not separator:
+            return None
+
+        parts: Final = model_path.split("/")
+        if len(parts) < 3 or parts[1] != "models" or not parts[2]:
+            return None
+
+        return f"publishers/{'/'.join(parts[:3])}"
+
     @classmethod
     def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:
         """
@@ -216,7 +242,11 @@ class VertexAIBatchTransformation:
         LiteLLM-managed unified file id) with a `publishers/` model path that
         `_get_model_from_gcs_file` can parse.
         """
-        return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id
+        return (
+            input_file_id is not None
+            and input_file_id.startswith("gs://")
+            and cls._parse_model_from_gcs_file(input_file_id) is not None
+        )
 
     @classmethod
     def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str:

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -109,7 +109,7 @@ if MCP_AVAILABLE:
     ############ MCP Server REST API Routes #################
     async def _safe_fire_mcp_tool_call_logging(
         logging_obj: Any | None,
-        result: Any,
+        result: "CallToolResult",
         start_time: datetime,
         end_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,

litellm/proxy/_experimental/mcp_server/server.py

@@ -13,9 +13,9 @@ import time
 import traceback
 import types
 import uuid
-from collections.abc import AsyncIterator, Callable, Mapping
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 import httpx
 from fastapi import FastAPI, HTTPException
@@ -145,7 +145,7 @@ try:
     )
 
     # Robust auth lookup keyed by session_object.
-    _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary()
+    _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary()
 except ImportError as e:
     verbose_logger.debug("MCP module not found: %s", e)
     MCP_AVAILABLE = False
@@ -493,14 +493,14 @@ if MCP_AVAILABLE:
     def _gateway_create_initialization_options(
         self,
         notification_options: NotificationOptions | None = None,
-        experimental_capabilities: dict[str, dict[str, Any]] | None = None,
+        experimental_capabilities: dict[str, dict[str, object]] | None = None,
     ) -> InitializationOptions:
         opts: Final = Server.create_initialization_options(
             self,
             notification_options=notification_options,
             experimental_capabilities=experimental_capabilities or {},
         )
-        updates: Final[dict[str, Any]] = {}
+        updates: Final[dict[str, str]] = {}
         merged: Final = _mcp_gateway_initialize_instructions.get()
         if merged is not None:
             updates["instructions"] = merged
@@ -549,6 +549,17 @@ if MCP_AVAILABLE:
     _stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
     _stateful_session_active_request_counts: Final[dict[str, int]] = {}
 
+    class _TerminableTransport(Protocol):
+        async def terminate(self) -> None: ...
+
+    class _TransportRegistry(Protocol):
+        def __contains__(self, session_id: object, /) -> bool: ...
+
+        def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ...
+
+    def _stateful_server_instances() -> _TransportRegistry:
+        return getattr(session_manager_stateful, "_server_instances", {})
+
     def _remove_stateful_session_tracking(session_id: str) -> None:
         _stateful_session_auth_contexts.pop(session_id, None)
         _stateful_session_auth_context_last_seen.pop(session_id, None)
@@ -578,8 +589,8 @@ if MCP_AVAILABLE:
     ) -> None:
         """Terminate expired stateful sessions and drop their auth contexts."""
         now = time.monotonic() if now is None else now
-        server_instances: Final = getattr(session_manager_stateful, "_server_instances", {})
-        expired_session_ids: Final = []
+        server_instances: Final = _stateful_server_instances()
+        expired_session_ids: Final[list[str]] = []
         for session_id, last_seen in _stateful_session_auth_context_last_seen.items():
             if _stateful_session_active_request_counts.get(session_id, 0) > 0:
                 continue
@@ -619,7 +630,7 @@ if MCP_AVAILABLE:
         session may proceed, or ``False`` when the caller is already at the cap
         with every session in flight (the new ``initialize`` should be rejected).
         """
-        server_instances: Final = getattr(session_manager_stateful, "_server_instances", {})
+        server_instances: Final = _stateful_server_instances()
 
         def _owned_live_session_ids() -> list[str]:
             return [
@@ -778,7 +789,7 @@ if MCP_AVAILABLE:
                     get_virtual_tool_definitions,
                 )
 
-                return [Tool(**d) for d in get_virtual_tool_definitions()]
+                return [Tool.model_validate(d) for d in get_virtual_tool_definitions()]
 
             # Get mcp_servers from context variable
             verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools")
@@ -847,7 +858,7 @@ if MCP_AVAILABLE:
 
     async def _build_virtual_call_logging_obj(
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         user_api_key_auth: UserAPIKeyAuth,
     ) -> LiteLLMLoggingObj | None:
         """Run the pre-call pipeline (guardrails + logging setup) for a virtual
@@ -885,7 +896,7 @@ if MCP_AVAILABLE:
 
     async def _dispatch_virtual_mcp_tool(
         name: str,
-        arguments: dict[str, Any] | None,
+        arguments: dict[str, object] | None,
         user_api_key_auth: UserAPIKeyAuth | None,
         client_ip: str | None,
         mcp_servers: list[str] | None = None,
@@ -957,7 +968,7 @@ if MCP_AVAILABLE:
         )
 
     @server.call_tool()
-    async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult:
+    async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult:
         """
         Call a specific tool with the provided arguments
         Args:
@@ -1621,7 +1632,7 @@ if MCP_AVAILABLE:
     async def _get_user_oauth_extra_headers_from_db(
         server: MCPServer,
         user_api_key_auth: UserAPIKeyAuth | None,
-        prefetched_creds: dict[str, dict[str, Any]] | None = None,
+        prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None,
     ) -> dict[str, str] | None:
         """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None.
 
@@ -1646,7 +1657,7 @@ if MCP_AVAILABLE:
 
         Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops.
         """
-        user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
+        user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
         if not user_id:
             return {}
         try:
@@ -1871,7 +1882,7 @@ if MCP_AVAILABLE:
 
         list_tools_start_time: Final = datetime.now()
         litellm_logging_obj: LiteLLMLoggingObj | None = None
-        list_tools_request_data: dict[str, Any] = {}
+        list_tools_request_data: dict[str, object] = {}
 
         if log_list_tools_to_spendlogs:
             # This is intentionally minimal: only async_success_handler / post_call_failure_hook
@@ -1879,7 +1890,7 @@ if MCP_AVAILABLE:
             list_tools_call_id: Final = str(uuid.uuid4())
             # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool)
             effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers)
-            spend_logs_metadata: Final[dict[str, Any]] = {
+            spend_logs_metadata: Final[dict[str, object]] = {
                 "mcp_operation": "list_tools",
             }
             if isinstance(list_tools_log_source, str):
@@ -2615,7 +2626,7 @@ if MCP_AVAILABLE:
 
     async def execute_mcp_tool(
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         allowed_mcp_servers: list[MCPServer],
         start_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,
@@ -2882,7 +2893,7 @@ if MCP_AVAILABLE:
                 _request_auth_header.reset(_auth_token)
                 _request_extra_headers.reset(_extra_token)
                 _request_resolved_auth_headers.reset(_resolved_token)
-            response = CallToolResult(content=cast(Any, local_content), isError=False)
+            response = CallToolResult(content=local_content, isError=False)
 
         # Try managed MCP server tool (the name is bare; the prefix boundary was
         # already resolved above against this server's registered prefixes)
@@ -2956,7 +2967,7 @@ if MCP_AVAILABLE:
                     arguments = hook_result["arguments"]  # pyright: ignore[reportAny]  # hook returns untyped args
 
             local_content = await _handle_local_mcp_tool(original_tool_name, arguments)
-            response = CallToolResult(content=cast(Any, local_content), isError=False)
+            response = CallToolResult(content=local_content, isError=False)
 
         return await _run_post_mcp_call_guardrails(
             result=response,
@@ -3003,7 +3014,7 @@ if MCP_AVAILABLE:
 
     async def _fire_mcp_tool_call_logging(
         logging_obj: LiteLLMLoggingObj,
-        result: Any,
+        result: CallToolResult,
         start_time: datetime,
         end_time: datetime,
         user_api_key_auth: UserAPIKeyAuth | None = None,
@@ -3070,7 +3081,7 @@ if MCP_AVAILABLE:
     @client
     async def call_mcp_tool(
         name: str,
-        arguments: dict[str, Any] | None = None,
+        arguments: dict[str, object] | None = None,
         user_api_key_auth: UserAPIKeyAuth | None = None,
         mcp_auth_header: str | None = None,
         mcp_servers: list[str] | None = None,
@@ -3161,7 +3172,7 @@ if MCP_AVAILABLE:
 
     async def mcp_get_prompt(
         name: str,
-        arguments: dict[str, Any] | None = None,
+        arguments: dict[str, object] | None = None,
         user_api_key_auth: UserAPIKeyAuth | None = None,
         mcp_auth_header: str | None = None,
         mcp_servers: list[str] | None = None,
@@ -3262,7 +3273,7 @@ if MCP_AVAILABLE:
 
     def _get_standard_logging_mcp_tool_call(
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         server_name: str | None,
         session_id: str | None = None,
     ) -> StandardLoggingMCPToolCall:
@@ -3291,13 +3302,13 @@ if MCP_AVAILABLE:
     async def _handle_managed_mcp_tool(
         server_name: str,
         name: str,
-        arguments: dict[str, Any],
+        arguments: dict[str, object],
         user_api_key_auth: UserAPIKeyAuth | None = None,
         mcp_auth_header: str | None = None,
         mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
         oauth2_headers: dict[str, str] | None = None,
         raw_headers: dict[str, str] | None = None,
-        litellm_logging_obj: Any | None = None,
+        litellm_logging_obj: LiteLLMLoggingObj | None = None,
         host_progress_callback: Callable | None = None,
     ) -> CallToolResult:
         """Handle tool execution for managed server tools"""
@@ -3320,7 +3331,7 @@ if MCP_AVAILABLE:
         return call_tool_result
 
     async def _handle_local_mcp_tool(
-        name: str, arguments: dict[str, Any]
+        name: str, arguments: dict[str, object]
     ) -> list[TextContent | ImageContent | EmbeddedResource]:
         """
         Handle tool execution for local registry tools
@@ -3426,7 +3437,8 @@ if MCP_AVAILABLE:
         Extract mcp-session-id from ASGI scope headers.
         Returns None if not present.
         """
-        for header_name, header_value in scope.get("headers", []):
+        scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", [])
+        for header_name, header_value in scope_headers:
             name = header_name if isinstance(header_name, bytes) else header_name.encode()
             if name.lower() == b"mcp-session-id":
                 return header_value.decode() if isinstance(header_value, bytes) else str(header_value)
@@ -3528,7 +3540,7 @@ if MCP_AVAILABLE:
             if message.get("type") != "http.request":
                 break
 
-            body = message.get("body", b"") or b""
+            body: bytes = message.get("body", b"") or b""
             if body:
                 # Only retain up to the remaining peek budget for sniffing.
                 # The full ``message`` is already in memory (delivered by
@@ -3571,9 +3583,9 @@ if MCP_AVAILABLE:
         Fixes https://github.com/BerriAI/litellm/issues/20992
         """
         _mcp_session_header: Final = b"mcp-session-id"
-        _headers: Final = scope.get("headers", [])
+        _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", [])
 
-        def _normalize_header_name(header_name: Any) -> bytes | None:
+        def _normalize_header_name(header_name: object) -> bytes | None:
             if isinstance(header_name, bytes):
                 return header_name.lower()
             if isinstance(header_name, str):
@@ -3902,7 +3914,8 @@ if MCP_AVAILABLE:
 
     def _get_authorization_header_from_scope(scope: Scope) -> str | None:
         """First ``Authorization`` header value in the ASGI scope, or None."""
-        for key, value in scope.get("headers", []):
+        scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+        for key, value in scope_headers:
             if key.lower() == b"authorization":
                 return value.decode("latin-1")
         return None
@@ -3921,7 +3934,8 @@ if MCP_AVAILABLE:
         ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream
         would leak the proxy key to a third-party MCP server.
         """
-        has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", []))
+        scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+        has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers)
         if not has_litellm_key_header:
             return None
         return _get_authorization_header_from_scope(scope)
@@ -4115,7 +4129,7 @@ if MCP_AVAILABLE:
     async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
         """Handle MCP requests through StreamableHTTP."""
         try:
-            path: Final = scope.get("path", "")
+            path: Final[str] = scope.get("path", "")
             (
                 user_api_key_auth,
                 mcp_auth_header,
@@ -4135,7 +4149,8 @@ if MCP_AVAILABLE:
             )
 
             # Strip any client-supplied x-mcp-toolset-id to prevent forgery.
-            scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"]
+            scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+            scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"]
 
             # Apply toolset scope if set server-side via ContextVar (set by
             # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py).
@@ -4436,7 +4451,7 @@ if MCP_AVAILABLE:
     async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None:
         """Handle MCP requests through SSE."""
         try:
-            path: Final = scope.get("path", "")
+            path: Final[str] = scope.get("path", "")
             (
                 user_api_key_auth,
                 mcp_auth_header,
@@ -4456,7 +4471,8 @@ if MCP_AVAILABLE:
             )
 
             # Strip any client-supplied x-mcp-toolset-id to prevent forgery.
-            scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"]
+            scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", [])
+            scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"]
 
             # Apply toolset scope if set server-side via ContextVar so the
             # downstream probe list matches the fully-authorized server set
@@ -4680,7 +4696,8 @@ if MCP_AVAILABLE:
     ) -> Send:
         async def wrapped_send(message: Message) -> None:
             if message.get("type") == "http.response.start":
-                for key, value in message.get("headers", []):
+                response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", [])
+                for key, value in response_headers:
                     header_name = key if isinstance(key, bytes) else str(key).encode()
                     if header_name.lower() == b"mcp-session-id":
                         session_id = value.decode() if isinstance(value, bytes) else str(value)

litellm/proxy/_types.py

@@ -2488,6 +2488,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
             "is active as a reminder that hard enforcement is relaxed."
         ),
     )
+    apply_user_budget_to_team_keys: bool | None = Field(
+        None,
+        description=(
+            "If True, a user's personal max_budget is enforced on every request they "
+            "make, including requests made with a team-scoped key. Defaults to False, "
+            "where a team-scoped key is governed only by the team and team-member "
+            "budgets and the key owner's personal max_budget does not apply "
+            "(see GitHub issue #12905)."
+        ),
+    )
     user_url_validation: bool | None = Field(
         None,
         description=(
@@ -3884,12 +3894,19 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse):
 ##########################################
 
 
+class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase):
+    access_group_id: str
+    access_group_name: str
+    models: tuple[str, ...]
+
+
 class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
     team_member_budget_table: LiteLLM_BudgetTableFull | None = None
     # Resources inherited from access groups (separate from direct assignments)
     access_group_models: list[str] | None = None
     access_group_mcp_server_ids: list[str] | None = None
     access_group_agent_ids: list[str] | None = None
+    access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
 
 
 class TeamInfoResponseObject(TypedDict):
@@ -4540,10 +4557,10 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase):
 
     user_role: (
         Literal[
-            LitellmUserRoles.INTERNAL_USER,
-            LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
             LitellmUserRoles.PROXY_ADMIN,
             LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
+            LitellmUserRoles.INTERNAL_USER,
+            LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
         ]
         | None
     ) = Field(

litellm/proxy/agent_endpoints/agent_registry.py

@@ -1,8 +1,10 @@
+import asyncio
 import hashlib
 import json
 from collections.abc import Iterator, Mapping, Sequence
 from datetime import datetime, timezone
-from typing import Any, Final, Protocol, TypedDict
+from types import MappingProxyType
+from typing import Any, Final, NamedTuple, Protocol, TypedDict
 
 import litellm
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@@ -10,7 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import (
     handle_update_object_permission_common,
 )
 from litellm.proxy.utils import PrismaClient
-from litellm.repositories.table_repositories import AgentsRepository
+from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository
 from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest
 
 
@@ -86,10 +88,32 @@ def agents_table(prisma_client: PrismaClient) -> AgentTableClient:
     return table
 
 
+class ObjectPermissionGrantRecord(Protocol):
+    object_permission_id: str
+    agents: list[str] | None
+
+
+class ObjectPermissionTableClient(Protocol):
+    async def find_many(self, where: Mapping[str, object]) -> Sequence[ObjectPermissionGrantRecord]: ...
+
+    async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
+
+
+def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient:
+    table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table
+    return table
+
+
+class GrantMigrationResult(NamedTuple):
+    rewritten: int
+    missed: int
+
+
 class AgentRegistry:
     def __init__(self):
         self.agent_list: list[AgentResponse] = []
         self.config_agents: tuple[AgentConfig, ...] = ()
+        self.config_agent_legacy_ids: Mapping[str, str] = MappingProxyType({})
 
     def reset_agent_list(self):
         self.agent_list = []
@@ -100,23 +124,33 @@ class AgentRegistry:
     def deregister_agent(self, agent_name: str):
         self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name]
 
-    def get_agent_list(self, agent_names: Sequence[str] | None = None):
+    def get_agent_list(self, agent_names: Sequence[str] | None = None) -> tuple[AgentResponse, ...]:
         if agent_names is not None:
-            return [agent for agent in self.agent_list if agent.agent_name in agent_names]
-        return self.agent_list
+            return tuple(agent for agent in self.agent_list if agent.agent_name in agent_names)
+        return tuple(self.agent_list)
 
-    def get_public_agent_list(self) -> list[AgentResponse]:
-        public_agent_list: Final[list[AgentResponse]] = []
-        if litellm.public_agent_groups is None:
-            return public_agent_list
-        for agent in self.agent_list:
-            if agent.agent_id in litellm.public_agent_groups:
-                public_agent_list.append(agent)
-        return public_agent_list
+    def get_public_agent_list(self) -> tuple[AgentResponse, ...]:
+        public_agent_groups: Final = litellm.public_agent_groups
+        if public_agent_groups is None:
+            return ()
+        return tuple(
+            agent for agent in self.agent_list if not self.ids_for_agent(agent.agent_id).isdisjoint(public_agent_groups)
+        )
 
     def _create_agent_id(self, agent_config: AgentConfig) -> str:
+        return hashlib.sha256(agent_config["agent_name"].encode()).hexdigest()
+
+    def _create_legacy_agent_id(self, agent_config: AgentConfig) -> str:
         return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest()
 
+    def ids_for_agent(self, agent_id: str) -> frozenset[str]:
+        return frozenset(
+            {agent_id, *(legacy for legacy, stable in self.config_agent_legacy_ids.items() if stable == agent_id)}
+        )
+
+    def stable_agent_id(self, agent_id: str) -> str:
+        return self.config_agent_legacy_ids.get(agent_id, agent_id)
+
     def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None):
         """
         Register the agents declared in config.yaml and remember them for later rebuilds.
@@ -131,12 +165,20 @@ class AgentRegistry:
         if agent_config is None:
             return
 
-        self.config_agents = tuple(agent_config)
-
         for agent_config_item in agent_config:
             if not isinstance(agent_config_item, dict):
                 raise ValueError("agent_config must be a list of dictionaries")
 
+        self.config_agents = tuple(agent_config)
+        self.config_agent_legacy_ids = MappingProxyType(
+            {
+                self._create_legacy_agent_id(agent_config_item): self._create_agent_id(agent_config_item)
+                for agent_config_item in agent_config
+                if agent_config_item.get("agent_name") and agent_config_item.get("agent_card_params")
+            }
+        )
+
+        for agent_config_item in agent_config:
             agent_name = agent_config_item.get("agent_name")
             agent_card_params = agent_config_item.get("agent_card_params")
             if not all([agent_name, agent_card_params]):
@@ -180,6 +222,45 @@ class AgentRegistry:
         self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents)
         return self.agent_list
 
+    async def migrate_legacy_grant_ids(self, table: ObjectPermissionTableClient) -> GrantMigrationResult:
+        """
+        Rewrite object_permission.agents rows holding a legacy full-entry hash to the
+        stable name-derived id.
+
+        Only the running proxy can do this: the legacy hash is computed from the
+        resolved config entry (secrets included), so no SQL migration can know it.
+        Persisting the stable id here is what keeps a grant alive across a later
+        secret rotation, which re-mints the legacy hash and would otherwise orphan
+        the stored value. Idempotent; runs of it after the first find no rows.
+
+        Each write is a compare-and-swap against the agents array read above, so a
+        grant edited concurrently is left untouched; the runtime alias keeps covering
+        it and the next boot retries the rewrite.
+        """
+        legacy_ids: Final = tuple(legacy for legacy, stable in self.config_agent_legacy_ids.items() if legacy != stable)
+        if not legacy_ids:
+            return GrantMigrationResult(rewritten=0, missed=0)
+        rows: Final = await table.find_many(where={"agents": {"has_some": legacy_ids}})
+        updates: Final = tuple(
+            (
+                row.object_permission_id,
+                tuple(row.agents or ()),
+                tuple(dict.fromkeys(self.stable_agent_id(agent_id) for agent_id in row.agents or ())),
+            )
+            for row in rows
+        )
+        counts: Final = await asyncio.gather(
+            *(
+                table.update_many(
+                    where={"object_permission_id": object_permission_id, "agents": {"equals": snapshot_agents}},
+                    data={"agents": translated_agents},
+                )
+                for object_permission_id, snapshot_agents, translated_agents in updates
+            )
+        )
+        rewritten: Final = sum(counts)
+        return GrantMigrationResult(rewritten=rewritten, missed=len(updates) - rewritten)
+
     ###########################################################
     ########### DB management helpers for agents ###########
     ############################################################
@@ -492,6 +573,14 @@ class AgentRegistry:
                 if agent.agent_id == agent_id:
                     return agent
 
+            translated_id: Final = self.config_agent_legacy_ids.get(agent_id)
+            if translated_id is None:
+                return None
+
+            for agent in self.agent_list:
+                if agent.agent_id == translated_id:
+                    return agent
+
             return None
         except Exception as e:
             raise Exception(f"Error getting agent from DB: {e}")

litellm/proxy/agent_endpoints/auth/agent_permission_handler.py

@@ -5,7 +5,8 @@ Handles agent permission checking for keys and teams using object_permission_id.
 Follows the same pattern as MCP permission handling.
 """
 
-from typing import Final
+from dataclasses import dataclass
+from typing import Final, TypeAlias
 
 from litellm._logging import verbose_logger
 from litellm.proxy._types import (
@@ -17,6 +18,27 @@ from litellm.proxy._types import (
 from litellm.repositories.table_repositories import AgentsRepository
 
 
+@dataclass(frozen=True, slots=True)
+class UnrestrictedAgentAccess:
+    """No agent grant exists on the key or its team, so every agent is reachable."""
+
+
+@dataclass(frozen=True, slots=True)
+class RestrictedAgentAccess:
+    """Only ``agent_ids`` are reachable. An empty set denies every agent."""
+
+    agent_ids: frozenset[str]
+
+
+AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess
+
+
+def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]:
+    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+
+    return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids)
+
+
 class AgentRequestHandler:
     """
     Class to handle agent permission checking, including:
@@ -32,37 +54,32 @@ class AgentRequestHandler:
     """
 
     @staticmethod
-    async def get_allowed_agents(
+    async def resolve_agent_access(
         user_api_key_auth: UserAPIKeyAuth | None = None,
-    ) -> list[str]:
+    ) -> AgentAccess:
         """
-        Get list of allowed agent IDs for the given user/key based on permissions.
+        Resolve the agents the given user/key may reach.
 
-        Returns:
-            List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all).
+        ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team
+        carries any grant. Grants that intersect to nothing stay restricted, so
+        narrowing a caller can never widen what it reaches.
         """
         try:
-            allowed_agents: list[str] = []
-            allowed_agents_for_key: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth)
-            allowed_agents_for_team: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth)
+            key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth)
+            team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth)
 
-            # If team has agent restrictions, handle inheritance and intersection logic
-            if len(allowed_agents_for_team) > 0:
-                if len(allowed_agents_for_key) > 0:
-                    # Key has its own agent permissions - use intersection with team permissions
-                    for agent_id in allowed_agents_for_key:
-                        if agent_id in allowed_agents_for_team:
-                            allowed_agents.append(agent_id)
-                else:
-                    # Key has no agent permissions - inherit from team
-                    allowed_agents = allowed_agents_for_team
-            else:
-                allowed_agents = allowed_agents_for_key
-
-            return list(set(allowed_agents))
+            match (key_access, team_access):
+                case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()):
+                    return UnrestrictedAgentAccess()
+                case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)):
+                    return RestrictedAgentAccess(_to_stable_ids(team_ids))
+                case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()):
+                    return RestrictedAgentAccess(_to_stable_ids(key_ids))
+                case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)):
+                    return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids))
         except Exception as e:
             verbose_logger.warning("Failed to get allowed agents: %s", e)
-            return []
+            return UnrestrictedAgentAccess()
 
     @staticmethod
     async def is_agent_allowed(
@@ -79,13 +96,14 @@ class AgentRequestHandler:
         Returns:
             bool: True if agent is allowed, False otherwise
         """
-        allowed_agents: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth)
+        from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
 
-        # Empty list means no restrictions - allow all
-        if len(allowed_agents) == 0:
-            return True
-
-        return agent_id in allowed_agents
+        match await AgentRequestHandler.resolve_agent_access(user_api_key_auth):
+            case UnrestrictedAgentAccess():
+                return True
+            case RestrictedAgentAccess(allowed_agent_ids):
+                stable_id: Final = global_agent_registry.stable_agent_id(agent_id)
+                return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids)
 
     @staticmethod
     def _get_key_object_permission(
@@ -139,55 +157,58 @@ class AgentRequestHandler:
     @staticmethod
     async def _get_allowed_agents_for_key(
         user_api_key_auth: UserAPIKeyAuth | None = None,
-    ) -> list[str]:
+    ) -> AgentAccess:
         """
         Get allowed agents for a key.
 
         1. First checks native key-level agent permissions (object_permission)
         2. Also includes agents from key's access_group_ids (unified access groups)
 
+        A key that declares agents or access groups is restricted even when those
+        declarations resolve to nothing, so an emptied or deleted access group denies
+        rather than opening the key up. Lookup failures still propagate to the caller,
+        which keeps them fail-open.
+
         Note: object_permission is already loaded by get_key_object() in main auth flow.
         """
         if user_api_key_auth is None:
-            return []
+            return UnrestrictedAgentAccess()
 
         try:
-            all_agents: list[str] = []
-
             # 1. Get agents from object_permission (native permissions)
             key_object_permission: Final = AgentRequestHandler._get_key_object_permission(user_api_key_auth)
-            if key_object_permission is not None:
-                # Get direct agents
-                direct_agents: Final = key_object_permission.agents or []
-
-                # Get agents from access groups
-                access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups(
-                    key_object_permission.agent_access_groups or []
-                )
-
-                all_agents = direct_agents + access_group_agents
-
+            direct_agents: Final = tuple(
+                key_object_permission.agents or () if key_object_permission is not None else ()
+            )
+            declared_access_groups: Final = tuple(
+                key_object_permission.agent_access_groups or () if key_object_permission is not None else ()
+            )
             # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups)
-            key_access_group_ids: Final = user_api_key_auth.access_group_ids or []
-            if key_access_group_ids:
-                from litellm.proxy.auth.auth_checks import (
-                    _get_agent_ids_from_access_groups,
-                )
+            key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ())
 
-                unified_agents: Final = await _get_agent_ids_from_access_groups(
-                    access_group_ids=key_access_group_ids,
-                )
-                all_agents.extend(unified_agents)
+            if not direct_agents and not declared_access_groups and not key_access_group_ids:
+                return UnrestrictedAgentAccess()
 
-            return list(set(all_agents))
+            access_group_agents: Final = (
+                tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups)))
+                if declared_access_groups
+                else ()
+            )
+            unified_agents: Final = (
+                tuple(await AgentRequestHandler._get_unified_access_group_agents(list(key_access_group_ids)))
+                if key_access_group_ids
+                else ()
+            )
+
+            return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents))
         except Exception as e:
             verbose_logger.warning("Failed to get allowed agents for key: %s", e)
-            return []
+            return UnrestrictedAgentAccess()
 
     @staticmethod
     async def _get_allowed_agents_for_team(
         user_api_key_auth: UserAPIKeyAuth | None = None,
-    ) -> list[str]:
+    ) -> AgentAccess:
         """
         Get allowed agents for a team.
 
@@ -195,12 +216,13 @@ class AgentRequestHandler:
         2. Also includes agents from team's access_group_ids (unified access groups)
 
         Fetches the team object once and reuses it for both permission sources.
+        Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`.
         """
         if user_api_key_auth is None:
-            return []
+            return UnrestrictedAgentAccess()
 
         if user_api_key_auth.team_id is None:
-            return []
+            return UnrestrictedAgentAccess()
 
         try:
             from litellm.proxy.auth.auth_checks import get_team_object
@@ -211,7 +233,7 @@ class AgentRequestHandler:
             )
 
             if not prisma_client:
-                return []
+                return UnrestrictedAgentAccess()
 
             # Fetch the team object once for both permission sources
             team_obj: Final = await get_team_object(
@@ -223,42 +245,38 @@ class AgentRequestHandler:
             )
 
             if team_obj is None:
-                return []
-
-            all_agents: list[str] = []
+                return UnrestrictedAgentAccess()
 
             # 1. Get agents from object_permission (native permissions)
             object_permissions: Final = team_obj.object_permission
-            if object_permissions is not None:
-                # Get direct agents
-                direct_agents: Final = object_permissions.agents or []
-
-                # Get agents from access groups
-                access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups(
-                    object_permissions.agent_access_groups or []
-                )
-
-                all_agents = direct_agents + access_group_agents
-
+            direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ())
+            declared_access_groups: Final = tuple(
+                object_permissions.agent_access_groups or () if object_permissions is not None else ()
+            )
             # 2. Also include agents from team's access_group_ids (unified access groups)
-            team_access_group_ids: Final = team_obj.access_group_ids or []
-            if team_access_group_ids:
-                from litellm.proxy.auth.auth_checks import (
-                    _get_agent_ids_from_access_groups,
-                )
+            team_access_group_ids: Final = tuple(team_obj.access_group_ids or ())
 
-                unified_agents: Final = await _get_agent_ids_from_access_groups(
-                    access_group_ids=team_access_group_ids,
-                )
-                all_agents.extend(unified_agents)
+            if not direct_agents and not declared_access_groups and not team_access_group_ids:
+                return UnrestrictedAgentAccess()
 
-            return list(set(all_agents))
+            access_group_agents: Final = (
+                tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups)))
+                if declared_access_groups
+                else ()
+            )
+            unified_agents: Final = (
+                tuple(await AgentRequestHandler._get_unified_access_group_agents(list(team_access_group_ids)))
+                if team_access_group_ids
+                else ()
+            )
+
+            return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents))
         except Exception as e:
             # litellm-dashboard is the default UI team and will never have agents;
             # skip noisy warnings for it.
             if user_api_key_auth.team_id != UI_TEAM_ID:
                 verbose_logger.warning("Failed to get allowed agents for team: %s", e)
-            return []
+            return UnrestrictedAgentAccess()
 
     @staticmethod
     def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]:
@@ -277,18 +295,26 @@ class AgentRequestHandler:
     async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]:
         """
         Helper to get agent_ids from DB agents that match any of the given access groups.
+
+        Query failures propagate so the caller can tell "this group is empty" (deny)
+        apart from "the lookup failed" (fail-open).
         """
-        agent_ids: Final[set[str]] = set()
-        if access_groups and prisma_client is not None:
-            try:
-                agents: Final = await AgentsRepository(prisma_client).table.find_many(
-                    where={"agent_access_groups": {"hasSome": access_groups}}
-                )
-                for agent in agents:
-                    agent_ids.add(agent.agent_id)
-            except Exception as e:
-                verbose_logger.debug("Error getting agents from access groups: %s", e)
-        return agent_ids
+        if not access_groups or prisma_client is None:
+            return set()
+
+        agents: Final = await AgentsRepository(prisma_client).table.find_many(
+            where={"agent_access_groups": {"hasSome": access_groups}}
+        )
+        return {agent.agent_id for agent in agents}
+
+    @staticmethod
+    async def _get_unified_access_group_agents(access_group_ids: list[str]) -> list[str]:
+        """
+        Resolve unified access group ids to agent IDs.
+        """
+        from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups
+
+        return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids)
 
     @staticmethod
     async def _get_agents_from_access_groups(
@@ -300,20 +326,17 @@ class AgentRequestHandler:
         from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
         from litellm.proxy.proxy_server import prisma_client
 
-        try:
-            # Use the helper for config-loaded agents
-            agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups(
-                global_agent_registry.agent_list, access_groups
-            )
+        # Use the helper for config-loaded agents
+        config_agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups(
+            global_agent_registry.agent_list, access_groups
+        )
 
-            # Use the helper for DB agents
-            db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups(prisma_client, access_groups)
-            agent_ids.update(db_agent_ids)
+        # Use the helper for DB agents
+        db_agent_ids: Final = await AgentRequestHandler._get_db_agent_ids_for_access_groups(
+            prisma_client, access_groups
+        )
 
-            return list(agent_ids)
-        except Exception as e:
-            verbose_logger.warning("Failed to get agents from access groups: %s", e)
-            return []
+        return list(config_agent_ids | db_agent_ids)
 
     @staticmethod
     async def get_agent_access_groups(

litellm/proxy/agent_endpoints/endpoints.py

@@ -101,7 +101,11 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client)
     foreign key. Mirrors how spend is joined into the agent response so the UI
     never has to cross-reference a full key dump client-side. Only non-secret
     fields are exposed (alias, masked key_name, hashed token)."""
-    agent_ids: Final = [agent.agent_id for agent in agents]
+    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+
+    agent_ids: Final = tuple(
+        alias_id for agent in agents for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
+    )
     if not agent_ids:
         return
     key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many(
@@ -117,7 +121,12 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client)
             )
         )
     for agent in agents:
-        agent.keys = keys_by_agent.get(agent.agent_id)
+        matched_keys = [
+            key_summary
+            for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
+            for key_summary in keys_by_agent.get(alias_id) or ()
+        ]
+        agent.keys = matched_keys or None
 
 
 def _redact_sensitive_agent_fields(
@@ -239,10 +248,12 @@ async def get_agents(
     from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
     from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
         AgentRequestHandler,
+        RestrictedAgentAccess,
+        UnrestrictedAgentAccess,
     )
 
     try:
-        returned_agents: list[AgentResponse] = []
+        returned_agents: Sequence[AgentResponse] = ()
 
         # Admin users get all agents
         if (
@@ -252,37 +263,45 @@ async def get_agents(
             returned_agents = global_agent_registry.get_agent_list()
         else:
             # Get allowed agents from object_permission (key/team level)
-            allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict)
+            agent_access: Final = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict)
+            all_agents: Final = global_agent_registry.get_agent_list()
 
-            # If no restrictions (empty list), return all agents
-            if len(allowed_agent_ids) == 0:
-                returned_agents = global_agent_registry.get_agent_list()
-            else:
-                # Filter agents by allowed IDs
-                all_agents: Final = global_agent_registry.get_agent_list()
-                returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids]
+            match agent_access:
+                case UnrestrictedAgentAccess():
+                    returned_agents = all_agents
+                case RestrictedAgentAccess(allowed_agent_ids):
+                    returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids]
 
         # Fetch current spend from DB for all returned agents
         from litellm.proxy.proxy_server import prisma_client
 
         if prisma_client is not None:
-            agent_ids: Final = [agent.agent_id for agent in returned_agents]
+            agent_ids: Final = tuple(
+                alias_id
+                for agent in returned_agents
+                for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
+            )
             if agent_ids:
                 db_agents: Final = await agents_table(prisma_client).find_many(
                     where={"agent_id": {"in": agent_ids}},
                 )
                 spend_map: Final = {a.agent_id: a.spend for a in db_agents}
                 for agent in returned_agents:
-                    if agent.agent_id in spend_map:
-                        agent.spend = spend_map[agent.agent_id]
+                    matched_spends = tuple(
+                        spend_map[alias_id]
+                        for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
+                        if alias_id in spend_map
+                    )
+                    if matched_spends:
+                        agent.spend = sum(matched_spends)
                 await _attach_keys_to_agents(returned_agents, prisma_client)
 
         # add is_public field to each agent - we do it this way, to allow setting config agents as public
         for agent in returned_agents:
             if agent.litellm_params is None:
                 agent.litellm_params = {}
-            agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and (
-                agent.agent_id in litellm.public_agent_groups
+            agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and not (
+                global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
             )
 
         # Redact sensitive fields for non-admin users
@@ -863,7 +882,7 @@ async def make_agent_public(
         if litellm.public_agent_groups is None:
             litellm.public_agent_groups = []
         # handle duplicates
-        if agent.agent_id in litellm.public_agent_groups:
+        if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups):
             raise HTTPException(
                 status_code=400,
                 detail=f"Agent with name {agent.agent_name} already in public agent groups",
@@ -1043,27 +1062,29 @@ async def get_agent_daily_activity(
     # intersect their explicit `agent_ids` filter with the same allowlist.
     from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
         AgentRequestHandler,
+        RestrictedAgentAccess,
+        UnrestrictedAgentAccess,
     )
     from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
 
     where_condition: Final[dict[str, object]] = {}
     if not _user_has_admin_view(user_api_key_dict):
-        permitted_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict)
-        # `get_allowed_agents` returns an empty list when the caller's key
-        # and team carry no agent restrictions. For activity scoping that's
-        # not "see everything" — fall back to the agents the caller
-        # created so they cannot enumerate other tenants' agents.
+        permitted_agent_ids: list[str] = []
+        # An unrestricted caller is not "see everything" for activity scoping. Fall
+        # back to the agents the caller created so they cannot enumerate other
+        # tenants' agents.
         # Guard against `user_id is None`: a literal None in Prisma
         # `where={"created_by": None}` resolves to ``created_by IS NULL``
         # and would expose every ownerless agent's rows.
-        if not permitted_agent_ids:
-            if user_api_key_dict.user_id is None:
-                permitted_agent_ids = []
-            else:
-                owned_records: Final = await agents_table(prisma_client).find_many(
-                    where={"created_by": user_api_key_dict.user_id}
-                )
-                permitted_agent_ids = [a.agent_id for a in owned_records]
+        match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict):
+            case RestrictedAgentAccess(allowed_agent_ids):
+                permitted_agent_ids = list(allowed_agent_ids)
+            case UnrestrictedAgentAccess():
+                if user_api_key_dict.user_id is not None:
+                    owned_records: Final = await agents_table(prisma_client).find_many(
+                        where={"created_by": user_api_key_dict.user_id}
+                    )
+                    permitted_agent_ids = [a.agent_id for a in owned_records]
 
         if agent_ids_list:
             permitted_agent_id_set: Final = set(permitted_agent_ids)

litellm/proxy/agent_endpoints/model_list_helpers.py

@@ -4,8 +4,6 @@ Helper functions for appending A2A agents to model lists.
 Used by proxy model endpoints to make agents appear in UI alongside models.
 """
 
-from typing import Final
-
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.proxy.management_endpoints.model_management_endpoints import (
@@ -27,20 +25,23 @@ async def append_agents_to_model_group(
         from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
         from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
             AgentRequestHandler,
+            RestrictedAgentAccess,
         )
 
-        allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict)
-
-        for agent_id in allowed_agent_ids:
-            agent = global_agent_registry.get_agent_by_id(agent_id)
-            if agent is not None:
-                model_groups.append(
-                    ModelGroupInfoProxy(
-                        model_group=f"a2a/{agent.agent_name}",
-                        mode="chat",
-                        providers=["a2a"],
-                    )
-                )
+        match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict):
+            case RestrictedAgentAccess(allowed_agent_ids):
+                for agent_id in allowed_agent_ids:
+                    agent = global_agent_registry.get_agent_by_id(agent_id)
+                    if agent is not None:
+                        model_groups.append(
+                            ModelGroupInfoProxy(
+                                model_group=f"a2a/{agent.agent_name}",
+                                mode="chat",
+                                providers=["a2a"],
+                            )
+                        )
+            case _:
+                pass
     except Exception as e:
         verbose_proxy_logger.debug("Error appending agents to model_group/info: %s", e)
 
@@ -61,30 +62,33 @@ async def append_agents_to_model_info(
         from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
         from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
             AgentRequestHandler,
+            RestrictedAgentAccess,
         )
 
-        allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict)
-
-        for agent_id in allowed_agent_ids:
-            agent = global_agent_registry.get_agent_by_id(agent_id)
-            if agent is not None:
-                models.append(
-                    {
-                        "model_name": f"a2a/{agent.agent_name}",
-                        "litellm_params": {
-                            "model": f"a2a/{agent.agent_name}",
-                            "custom_llm_provider": "a2a",
-                        },
-                        "model_info": {
-                            "id": agent.agent_id,
-                            "mode": "chat",
-                            "db_model": True,
-                            "created_by": agent.created_by,
-                            "created_at": agent.created_at,
-                            "updated_at": agent.updated_at,
-                        },
-                    }
-                )
+        match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict):
+            case RestrictedAgentAccess(allowed_agent_ids):
+                for agent_id in allowed_agent_ids:
+                    agent = global_agent_registry.get_agent_by_id(agent_id)
+                    if agent is not None:
+                        models.append(
+                            {
+                                "model_name": f"a2a/{agent.agent_name}",
+                                "litellm_params": {
+                                    "model": f"a2a/{agent.agent_name}",
+                                    "custom_llm_provider": "a2a",
+                                },
+                                "model_info": {
+                                    "id": agent.agent_id,
+                                    "mode": "chat",
+                                    "db_model": True,
+                                    "created_by": agent.created_by,
+                                    "created_at": agent.created_at,
+                                    "updated_at": agent.updated_at,
+                                },
+                            }
+                        )
+            case _:
+                pass
     except Exception as e:
         verbose_proxy_logger.debug("Error appending agents to v2/model/info: %s", e)
 

litellm/proxy/auth/auth_checks.py

@@ -13,6 +13,7 @@ import asyncio
 import math
 import re
 import time
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
 
 from fastapi import HTTPException, Request, status
@@ -648,28 +649,29 @@ async def common_checks(
             )
 
         async def _user_max_budget_check() -> None:
-            # 4.1 personal budget, if personal key
-            if (
-                (team_object is None or team_object.team_id is None)
-                and user_object is not None
-                and user_object.max_budget is not None
-            ):
-                from litellm.proxy.proxy_server import get_current_spend
+            # 4.1 personal budget
+            if user_object is None or user_object.max_budget is None:
+                return
+            is_team_key: Final = team_object is not None and team_object.team_id is not None
+            if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True:
+                return
 
-                user_budget: Final = user_object.max_budget
-                user_spend: Final = await get_current_spend(
-                    counter_key=f"spend:user:{user_object.user_id}",
-                    fallback_spend=user_object.spend or 0.0,
+            from litellm.proxy.proxy_server import get_current_spend
+
+            user_budget: Final = user_object.max_budget
+            user_spend: Final = await get_current_spend(
+                counter_key=f"spend:user:{user_object.user_id}",
+                fallback_spend=user_object.spend or 0.0,
+                max_budget=user_budget,
+            )
+            if math.isfinite(user_budget) and user_spend >= user_budget:
+                raise litellm.BudgetExceededError(
+                    current_cost=user_spend,
                     max_budget=user_budget,
+                    message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
+                    entity_type=Litellm_EntityType.USER.value,
+                    entity_id=user_object.user_id,
                 )
-                if math.isfinite(user_budget) and user_spend >= user_budget:
-                    raise litellm.BudgetExceededError(
-                        current_cost=user_spend,
-                        max_budget=user_budget,
-                        message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
-                        entity_type=Litellm_EntityType.USER.value,
-                        entity_id=user_object.user_id,
-                    )
 
         # Each scope reads a distinct counter key with no cross-scope ordering
         # dependency, so the per-scope Redis-first reads run concurrently instead
@@ -2833,7 +2835,7 @@ async def get_org_object(
 
 
 async def _get_resources_from_access_groups(
-    access_group_ids: list[str],
+    access_group_ids: Sequence[str],
     resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
     prisma_client: PrismaClient | None = None,
     user_api_key_cache: UserApiKeyCache | None = None,
@@ -2892,7 +2894,7 @@ async def _get_resources_from_access_groups(
 
 
 async def _get_models_from_access_groups(
-    access_group_ids: list[str],
+    access_group_ids: Sequence[str],
     prisma_client: PrismaClient | None = None,
     user_api_key_cache: UserApiKeyCache | None = None,
     proxy_logging_obj: ProxyLogging | None = None,
@@ -4262,6 +4264,10 @@ async def _project_soft_budget_check(
             )
 
 
+def _project_cache_key(project_id: str) -> str:
+    return f"project_id:{project_id}"
+
+
 async def get_project_object(
     project_id: str,
     prisma_client: PrismaClient | None,
@@ -4279,7 +4285,7 @@ async def get_project_object(
         return None
 
     # Check cache first
-    cache_key: Final = f"project_id:{project_id}"
+    cache_key: Final = _project_cache_key(project_id)
     deserialized_project: Final = await user_api_key_cache.async_get_cache(
         key=cache_key,
         model_type=LiteLLM_ProjectTableCachedObj,
@@ -4310,6 +4316,32 @@ async def get_project_object(
     return project_obj
 
 
+async def delete_cached_project_object(
+    project_id: str,
+    user_api_key_cache: UserApiKeyCache,
+) -> None:
+    """
+    Every endpoint that mutates litellm_projecttable must call this: get_project_object
+    serves auth cache-first with no freshness check, so without invalidation a stale
+    project (e.g. a pre-update empty model allowlist) keeps being enforced until the
+    TTL expires (LIT-3803). Best-effort on both steps: the DB write has already
+    committed, so a cache backend error must not fail the endpoint; the stale entry
+    then expires via TTL.
+    """
+    from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
+
+    cache_key: Final = _project_cache_key(project_id)
+    try:
+        await user_api_key_cache.async_delete_cache(key=cache_key)
+    except Exception as e:  # noqa: BLE001  # best-effort eviction: any cache backend error must not fail the mutation
+        verbose_proxy_logger.warning(
+            "Failed to evict cached project entry %s; a stale project may be served until its TTL expires: %s",
+            cache_key,
+            e,
+        )
+    await publish_auth_cache_invalidation(cache_key=cache_key)
+
+
 async def _organization_max_budget_check(
     valid_token: UserAPIKeyAuth | None,
     team_object: LiteLLM_TeamTable | None,

litellm/proxy/auth/model_checks.py

@@ -1,6 +1,7 @@
 # What is this?
 ## Common checks for /v1/models and `/model/info`
 import copy
+from collections.abc import Sequence
 from typing import Any, Final
 
 import litellm
@@ -178,8 +179,8 @@ def get_team_models(
 
 
 def get_complete_model_list(
-    key_models: list[str],
-    team_models: list[str],
+    key_models: Sequence[str],
+    team_models: Sequence[str],
     proxy_model_list: list[str],
     user_model: str | None,
     infer_model_from_keys: bool | None,
@@ -203,7 +204,7 @@ def get_complete_model_list(
 
     def append_unique(models):
         for model in models:
-            if model not in unique_models:
+            if model not in unique_models and model != SpecialModelNames.no_default_models.value:
                 unique_models.append(model)
 
     if key_models:

litellm/proxy/auth/user_api_key_auth.py

@@ -1044,6 +1044,22 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
     request.state.parent_otel_span = parent_otel_span
 
 
+async def _read_request_body_deferring_parse_failure(
+    request: Request,
+) -> tuple[dict, ProxyException | None]:
+    """Parse the body, returning a parse failure instead of raising it.
+
+    A body that fails to parse is still a request from a known caller, so auth
+    must run (resolving identity onto the request's trace) before the 400 goes
+    out; the caller re-raises the returned exception once identity is seeded.
+    """
+    try:
+        parsed_body: Final = await _read_request_body(request=request)
+    except ProxyException as parse_exception:
+        return {}, parse_exception  # mutable-ok: request_data is a plain dict across the whole auth path
+    return populate_request_with_path_params(request_data=parsed_body, request=request), None
+
+
 async def _user_api_key_auth_builder(
     request: Request,
     api_key: str,
@@ -2470,6 +2486,7 @@ async def _reserve_budget_after_common_checks(
         proxy_logging_obj=proxy_logging_obj,
         end_user_id=end_user_id,
         end_user_object=end_user_object,
+        apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
         fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
     )
 
@@ -2515,6 +2532,72 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) ->
     )
 
 
+async def _authorize_authenticated_request(
+    user_api_key_auth_obj: UserAPIKeyAuth,
+    request: Request,
+    request_data: dict,
+    route: str,
+    api_key: str,
+) -> UserAPIKeyAuth | None:
+    """Authorize an already-authenticated request: disabled-route check, the single
+    ``common_checks`` gate (which also reserves budget), and end-user fallback
+    resolution. Returns the auth object the exception handler recovered when a check
+    failed but the request may proceed anyway, else ``None``.
+    """
+    ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
+    RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
+
+    # Single authorization point. Builder paths MUST NOT call common_checks.
+    # Route through the same exception handler the builder uses so
+    # authorization failures (ProxyException, or plain Exception from
+    # admin-only-route / model-access / budget checks) surface as
+    # ProxyException consistently with pre-refactor behavior.
+    try:
+        await _run_centralized_common_checks(
+            user_api_key_auth_obj=user_api_key_auth_obj,
+            request=request,
+            request_data=request_data,
+            route=route,
+        )
+    except Exception as e:
+        return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
+            e=e,
+            request=request,
+            request_data=request_data,
+            route=route,
+            parent_otel_span=user_api_key_auth_obj.parent_otel_span,
+            api_key=api_key,
+            resolved_identity=user_api_key_auth_obj,
+        )
+
+    # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
+    # paths (no master key, /user/auth route, JWT short-circuits) that bypass
+    # the end-user resolution block. If those paths produced an auth obj
+    # without an ``end_user_id`` set, fall back to extracting from the request
+    # body so spend logs are still attributed correctly. Validation honours
+    # ``litellm.validate_end_user_id_in_db``.
+    if user_api_key_auth_obj.end_user_id is None:
+        from litellm.proxy.proxy_server import (
+            prisma_client,
+            proxy_logging_obj,
+            user_api_key_cache,
+        )
+
+        raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
+        if raw_end_user_id is not None:
+            resolved_end_user_id: Final = await resolve_and_validate_end_user_id(
+                raw_end_user_id=raw_end_user_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=user_api_key_auth_obj.parent_otel_span,
+                proxy_logging_obj=proxy_logging_obj,
+                route=route,
+            )
+            if resolved_end_user_id is not None:
+                user_api_key_auth_obj.end_user_id = resolved_end_user_id
+    return None
+
+
 @tracer.wrap()
 async def user_api_key_auth(
     request: Request,
@@ -2535,8 +2618,7 @@ async def user_api_key_auth(
     # close, and the trace never reaches the backend.
     _ensure_parent_otel_span_on_request_state(request)
 
-    request_data = await _read_request_body(request=request)
-    request_data = populate_request_with_path_params(request_data=request_data, request=request)
+    request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request)
     route: Final[str] = get_request_route(request=request)
     ## CHECK IF ROUTE IS ALLOWED
 
@@ -2544,69 +2626,41 @@ async def user_api_key_auth(
     # triggers (key/user/team object reads) nest under it instead of flattening
     # onto the server span. No-op when OTel V2 isn't active.
     with phase_span(f"auth {route}"):
-        user_api_key_auth_obj: Final = await _user_api_key_auth_builder(
-            request=request,
-            api_key=api_key,
-            azure_api_key_header=azure_api_key_header,
-            anthropic_api_key_header=anthropic_api_key_header,
-            google_ai_studio_api_key_header=google_ai_studio_api_key_header,
-            azure_apim_header=azure_apim_header,
-            request_data=request_data,
-            custom_litellm_key_header=custom_litellm_key_header,
-        )
+        try:
+            user_api_key_auth_obj: Final = await _user_api_key_auth_builder(
+                request=request,
+                api_key=api_key,
+                azure_api_key_header=azure_api_key_header,
+                anthropic_api_key_header=anthropic_api_key_header,
+                google_ai_studio_api_key_header=google_ai_studio_api_key_header,
+                azure_apim_header=azure_apim_header,
+                request_data=request_data,
+                custom_litellm_key_header=custom_litellm_key_header,
+            )
+        except Exception:
+            # The body was read first, so a caller who sent both a malformed body and
+            # a rejected key used to get the 400; the response is unchanged, and the
+            # auth failure is still recorded on the trace by the handler that ran.
+            if body_parse_exception is not None:
+                raise body_parse_exception
+            raise
         user_api_key_auth_obj.budget_reservation = None
 
-        ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
-        RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
-
-        # Single authorization point. Builder paths MUST NOT call common_checks.
-        # Route through the same exception handler the builder uses so
-        # authorization failures (ProxyException, or plain Exception from
-        # admin-only-route / model-access / budget checks) surface as
-        # ProxyException consistently with pre-refactor behavior.
-        try:
-            await _run_centralized_common_checks(
+        # A body that never parsed is authenticated (so the trace carries identity
+        # and this ``auth`` span) but not authorized: there is no model to check it
+        # against, and budget reservation would increment live spend counters that
+        # only the endpoint's post-call path releases; the endpoint never runs, since
+        # the parse failure is raised below.
+        if body_parse_exception is None:
+            recovered_auth_obj: Final = await _authorize_authenticated_request(
                 user_api_key_auth_obj=user_api_key_auth_obj,
                 request=request,
                 request_data=request_data,
                 route=route,
-            )
-        except Exception as e:
-            return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
-                e=e,
-                request=request,
-                request_data=request_data,
-                route=route,
-                parent_otel_span=user_api_key_auth_obj.parent_otel_span,
                 api_key=api_key,
-                resolved_identity=user_api_key_auth_obj,
             )
-
-        # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
-        # paths (no master key, /user/auth route, JWT short-circuits) that bypass
-        # the end-user resolution block. If those paths produced an auth obj
-        # without an ``end_user_id`` set, fall back to extracting from the request
-        # body so spend logs are still attributed correctly. Validation honours
-        # ``litellm.validate_end_user_id_in_db``.
-        if user_api_key_auth_obj.end_user_id is None:
-            from litellm.proxy.proxy_server import (
-                prisma_client,
-                proxy_logging_obj,
-                user_api_key_cache,
-            )
-
-            raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
-            if raw_end_user_id is not None:
-                resolved_end_user_id: Final = await resolve_and_validate_end_user_id(
-                    raw_end_user_id=raw_end_user_id,
-                    prisma_client=prisma_client,
-                    user_api_key_cache=user_api_key_cache,
-                    parent_otel_span=user_api_key_auth_obj.parent_otel_span,
-                    proxy_logging_obj=proxy_logging_obj,
-                    route=route,
-                )
-                if resolved_end_user_id is not None:
-                    user_api_key_auth_obj.end_user_id = resolved_end_user_id
+            if recovered_auth_obj is not None:
+                return recovered_auth_obj
 
     # Identity is now resolved. Seed it AFTER the auth span closes so the Baggage
     # persists on the request task (detaching the span's context token inside the
@@ -2618,6 +2672,9 @@ async def user_api_key_auth(
     )
     user_api_key_auth_obj.request_route = normalize_request_route(route)
 
+    if body_parse_exception is not None:
+        raise body_parse_exception
+
     # Resolve caller identity once, here at the seam, into a single per-request
     # Principal projected off the key object the builder already fetched (no
     # second lookup). Downstream consumers read identity off this instead of

litellm/proxy/batches_endpoints/endpoints.py

@@ -36,6 +36,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
     get_original_file_id,
     prepare_data_with_credentials,
     update_batch_in_database,
+    validate_managed_id_requirement,
 )
 from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
 from litellm.repositories.table_repositories import ManagedFileRepository
@@ -176,6 +177,12 @@ async def create_batch(
             }
 
         input_file_id: Final = _create_batch_data.get("input_file_id", None)
+        await validate_managed_id_requirement(
+            resource_id=input_file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         unified_file_id: str | Literal[False] = False
 
         model_from_file_id = None
@@ -392,6 +399,12 @@ async def retrieve_batch(
 
     data: dict = {}
     try:
+        await validate_managed_id_requirement(
+            resource_id=batch_id,
+            resource_kind="batch",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         model_from_id: Final = decode_model_from_file_id(batch_id)
         _retrieve_batch_request: Final = RetrieveBatchRequest(
             batch_id=batch_id,
@@ -840,6 +853,13 @@ async def cancel_batch(
 
     data: dict = {}
     try:
+        await validate_managed_id_requirement(
+            resource_id=batch_id,
+            resource_kind="batch",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         # Check for encoded batch ID with model info
         model_from_id: Final = decode_model_from_file_id(batch_id)
 

litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py

@@ -0,0 +1,153 @@
+import asyncio
+import json
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Final
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy.common_utils.config_sync_pubsub import (
+    _ConfigSyncPubSub,
+    _pubsub_capable_client,
+    coordination_redis_cache,
+)
+
+if TYPE_CHECKING:
+    from litellm.caching.redis_cache import RedisCache
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
+_POLL_TIMEOUT_SECONDS: Final = 1.0
+_BACKOFF_INITIAL_SECONDS: Final = 5.0
+_BACKOFF_MAX_SECONDS: Final = 60.0
+
+
+def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
+    if redis_cache.namespace is None:
+        return AUTH_CACHE_INVALIDATION_CHANNEL
+    return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}"
+
+
+@dataclass(frozen=True, slots=True)
+class _CacheInvalidationMessage:
+    cache_key: str
+
+
+def _cache_invalidation_message_json(cache_key: str) -> str:
+    return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key)))
+
+
+def _cache_key_from_message_data(data: object) -> str | None:
+    if isinstance(data, bytes):
+        data = data.decode("utf-8", errors="replace")
+    if not isinstance(data, str):
+        return None
+    try:
+        parsed: Final = json.loads(data)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(parsed, dict):
+        return None
+    cache_key: Final = parsed.get("cache_key")
+    return cache_key if isinstance(cache_key, str) else None
+
+
+async def publish_auth_cache_invalidation(cache_key: str) -> None:
+    """
+    Best-effort broadcast so every worker drops its local in-memory copy of a
+    mutated management object; without this, only the handling worker and Redis
+    are evicted and other workers keep serving the stale object until its TTL.
+    """
+    redis_cache: Final = coordination_redis_cache()
+    if redis_cache is None:
+        return
+    try:
+        client: Final = _pubsub_capable_client(redis_cache)
+        if client is None:
+            verbose_proxy_logger.debug(
+                "auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support",
+                cache_key,
+            )
+            return
+        await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key))
+    except Exception as e:  # noqa: BLE001  # best-effort publish; mutations must never fail on redis errors
+        verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
+
+
+class AuthCacheInvalidationSubscriber:
+    __slots__ = ("_redis_cache", "_task", "_user_api_key_cache")
+
+    def __init__(
+        self,
+        redis_cache: "RedisCache",
+        user_api_key_cache: "UserApiKeyCache",
+    ) -> None:
+        self._redis_cache = redis_cache
+        self._user_api_key_cache = user_api_key_cache
+        self._task: asyncio.Task[None] | None = None
+
+    def start(self) -> None:
+        if self._task is not None:
+            return
+        self._task = asyncio.create_task(self._run())
+
+    async def stop(self) -> None:
+        task: Final = self._task
+        if task is None:
+            return
+        self._task = None
+        _ = task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+    async def _run(self) -> None:
+        backoff_seconds = _BACKOFF_INITIAL_SECONDS  # rebind-ok: exponential backoff accumulator across reconnects
+        while True:
+            try:
+                client = _pubsub_capable_client(self._redis_cache)  # rebind-ok: re-resolved on every reconnect
+                if client is None:
+                    verbose_proxy_logger.warning(
+                        "auth cache invalidation subscriber disabled: cluster redis client has no pub/sub support; "
+                        "cross-worker eviction falls back to the local cache TTL"
+                    )
+                    return
+                pubsub = client.pubsub()  # rebind-ok: fresh pubsub per reconnect
+                try:
+                    await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache))
+                    backoff_seconds = _BACKOFF_INITIAL_SECONDS  # rebind-ok: reset after successful subscribe
+                    await self._consume(pubsub)
+                finally:
+                    await self._close_pubsub(pubsub)
+            except asyncio.CancelledError:
+                raise
+            except Exception as e:  # noqa: BLE001  # any redis failure falls through to backoff and reconnect
+                verbose_proxy_logger.warning(
+                    "auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs",
+                    e,
+                    backoff_seconds,
+                )
+                await asyncio.sleep(backoff_seconds)
+                backoff_seconds = min(backoff_seconds * 2, _BACKOFF_MAX_SECONDS)  # rebind-ok: backoff accumulator
+
+    async def _consume(self, pubsub: _ConfigSyncPubSub) -> None:
+        while True:
+            message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS)
+            if message is None:
+                continue
+            self._apply_message(message)
+
+    def _apply_message(self, message: object) -> None:
+        data: Final = message.get("data") if isinstance(message, dict) else None
+        cache_key: Final = _cache_key_from_message_data(data)
+        if cache_key is None:
+            return
+        in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
+        if in_memory_cache is not None:
+            in_memory_cache.delete_cache(cache_key)
+
+    @staticmethod
+    async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None:
+        try:
+            await pubsub.aclose()
+        except Exception as e:  # noqa: BLE001  # best-effort close of a possibly-broken connection
+            verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", e)

litellm/proxy/common_utils/callback_utils.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional
 import litellm
 from litellm import get_secret
 from litellm._logging import verbose_proxy_logger
-from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
+from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import (
     get_metadata_variable_name_from_kwargs,
@@ -425,6 +425,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
         "_guardrail_pipelines",
         "_pipeline_managed_guardrails",
         PRE_CALL_EXECUTED_GUARDRAILS_KEY,
+        SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
         "disable_global_guardrails",
         "disable_global_guardrail",
         "opted_out_global_guardrails",

litellm/proxy/credential_endpoints/endpoints.py

@@ -346,4 +346,4 @@ async def update_credential(
 
         return {"success": True, "message": "Credential updated successfully"}
     except Exception as e:
-        return handle_exception_on_proxy(e)
+        raise handle_exception_on_proxy(e)

litellm/proxy/db/autorouter_session_rollup.py

@@ -33,6 +33,53 @@ if TYPE_CHECKING:
 CACHE_TTL_5M_SECONDS: Final = 300
 CACHE_TTL_1H_SECONDS: Final = 3600
 
+AUTOROUTER_BENCHMARKS_SQL: Final = """
+WITH windowed AS (
+    SELECT * FROM "LiteLLM_AutoRouterSession"
+    WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
+),
+tier_maps AS (
+    SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
+    FROM (
+        SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns
+        FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv
+        GROUP BY router_name, router_type, kv.key
+    ) per_tier
+    GROUP BY router_name, router_type
+)
+SELECT
+    agg.*,
+    COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns
+FROM (
+SELECT
+    router_name,
+    router_type,
+    COUNT(*)::int AS sessions,
+    COALESCE(SUM(turns), 0)::int AS turns,
+    COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns,
+    COALESCE(SUM(covered_turns), 0)::int AS covered_turns,
+    COALESCE(SUM(cache_hits), 0)::int AS cache_hits,
+    COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns,
+    COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits,
+    COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns,
+    COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits,
+    COALESCE(SUM(return_turns), 0)::int AS return_turns,
+    COALESCE(SUM(return_hits), 0)::int AS return_hits,
+    COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses,
+    COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses,
+    COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns,
+    COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns,
+    COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
+    COALESCE(SUM(spend), 0)::float8 AS spend,
+    COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
+    COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
+FROM windowed
+GROUP BY router_name, router_type
+) agg
+LEFT JOIN tier_maps USING (router_name, router_type)
+ORDER BY agg.spend DESC
+"""
+
 
 @dataclass(frozen=True, slots=True)
 class AutoRouterTurnTransaction:
@@ -49,6 +96,7 @@ class AutoRouterTurnTransaction:
     cache_hit: bool
     cache_ttl_seconds: int | None
     cache_touched: bool
+    tier: str | None = None
 
 
 class TurnCacheFacts(NamedTuple):
@@ -152,11 +200,13 @@ def build_autorouter_turn_transaction(
         return None
     usage_object_raw: Final = metadata.get("usage_object")
     cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
+    tier_raw: Final = routing_decision.get("tier")
     return AutoRouterTurnTransaction(
         api_key=api_key,
         session_id=_bounded_session_id(session_id),
         router_name=router_name,
         router_type=str(routing_decision.get("router_type") or "unknown"),
+        tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None,
         model=model,
         turn_at=turn_at,
         total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
@@ -184,6 +234,8 @@ _COVERED: Final = _p("covered")
 _CACHE_HIT: Final = _p("cache_hit")
 _CACHE_TTL: Final = _p("cache_ttl_seconds")
 _TOUCHED: Final = _p("cache_touched")
+_TIER: Final = f"{_p('tier')}::text"
+_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)"
 
 _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at"
 _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}"
@@ -201,7 +253,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
     last_model, models, turns, unordered_turns, covered_turns, cache_hits,
     same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
     return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
-    ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend
+    ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns
 )
 VALUES (
     {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
@@ -211,7 +263,8 @@ VALUES (
     0, 0, 0, 0,
     (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
     (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
-    {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8
+    {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
+    {_TIER_DELTA}
 )
 ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
     turns = t.turns + 1,
@@ -242,6 +295,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
                 ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END)
     )),
     last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END),
+    tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")}
+        THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1)
+        ELSE t.tier_turns END),
     first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at),
     last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at)
 """

litellm/proxy/db/prisma_client.py

@@ -486,6 +486,38 @@ class PrismaWrapper:
         os.environ[self._db_url_env_var] = _db_url
         return _db_url
 
+    @property
+    def engine_generation(self) -> int:
+        """How many query-engine replacements have completed on this wrapper.
+
+        Bumped under `_reconnection_lock` only after a replacement engine has
+        connected, so a change across an await proves a *successful* planned
+        replacement happened in between — a replacement that failed (a real
+        outage) leaves it untouched.
+        """
+        return self._engine_generation
+
+    async def _reconnection_settled(self) -> None:
+        async with self._reconnection_lock:
+            pass
+
+    async def wait_for_planned_engine_replacement(self, timeout_seconds: float) -> None:
+        """Wait, bounded, for an in-flight planned engine replacement to finish.
+
+        Both replacement paths (`recreate_prisma_client` and
+        `_safe_refresh_token`) hold `_reconnection_lock` across their whole
+        kill/connect window, so re-acquiring it means the replacement has
+        settled one way or the other. Gives up silently on timeout: a caller
+        that stopped waiting must treat the replacement as not completed and
+        consult `engine_generation` rather than assume success.
+        """
+        if timeout_seconds <= 0 or not self._reconnection_lock.locked():
+            return
+        try:
+            await asyncio.wait_for(self._reconnection_settled(), timeout=timeout_seconds)
+        except asyncio.TimeoutError:
+            return
+
     async def recreate_prisma_client(
         self,
         new_db_url: str,

litellm/proxy/fine_tuning_endpoints/endpoints.py

@@ -17,6 +17,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.openai_files_endpoints.common_utils import (
     _is_base64_encoded_unified_file_id,
+    validate_managed_id_requirement,
 )
 from litellm.proxy.utils import handle_exception_on_proxy
 from litellm.types.utils import LiteLLMFineTuningJob
@@ -134,6 +135,18 @@ async def create_fine_tuning_job(
         ## CHECK IF MANAGED FILE ID
         unified_file_id: str | Literal[False] = False
         training_file: Final = fine_tuning_request.training_file
+        await validate_managed_id_requirement(
+            resource_id=training_file,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+        await validate_managed_id_requirement(
+            resource_id=fine_tuning_request.validation_file,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         response: LiteLLMFineTuningJob | None = None
         if training_file:
             unified_file_id = _is_base64_encoded_unified_file_id(training_file)
@@ -246,6 +259,12 @@ async def retrieve_fine_tuning_job(
     try:
         if premium_user is not True:
             raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}")
+        await validate_managed_id_requirement(
+            resource_id=fine_tuning_job_id,
+            resource_kind="fine-tuning job",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         # Include original request and headers in the data
         base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
         (
@@ -513,6 +532,12 @@ async def cancel_fine_tuning_job(
     try:
         if premium_user is not True:
             raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}")
+        await validate_managed_id_requirement(
+            resource_id=fine_tuning_job_id,
+            resource_kind="fine-tuning job",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
         # Include original request and headers in the data
         base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
         (

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -6,10 +6,10 @@ import concurrent.futures
 import inspect
 import json
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timezone
 from types import UnionType
-from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin
 from urllib.parse import urlparse
 
 from fastapi import APIRouter, Depends, HTTPException, Request
@@ -54,8 +54,8 @@ from litellm.types.guardrails import (
 if TYPE_CHECKING:
     from types import CodeType
 
-    from prisma.actions import LiteLLM_GuardrailsTableActions
     from prisma.models import LiteLLM_GuardrailsTable
+    from pydantic.fields import FieldInfo
 
     from litellm.proxy.utils import PrismaClient
 
@@ -65,24 +65,44 @@ router: Final = APIRouter()
 GUARDRAIL_REGISTRY: Final = GuardrailRegistry()
 
 
-def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]":
-    table: Final[LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]] = GuardrailsRepository(prisma_client).table
+class _GuardrailsTableActions(Protocol):
+    async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ...
+
+    async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
+
+    async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
+
+    async def find_many(
+        self, where: Mapping[str, object], order: Mapping[str, str]
+    ) -> "Sequence[LiteLLM_GuardrailsTable]": ...
+
+    async def update(
+        self, where: Mapping[str, object], data: Mapping[str, object]
+    ) -> "LiteLLM_GuardrailsTable | None": ...
+
+
+def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]:
+    return mapping
+
+
+def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions:
+    table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table
     return table
 
 
 async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable":
-    row: Final[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.create(data=data)
+    row: Final = await _guardrails_table(prisma_client).create(data=data)
     return row
 
 
 async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None:
-    await GuardrailsRepository(prisma_client).table.delete(where=where)
+    await _guardrails_table(prisma_client).delete(where=where)
 
 
 async def _find_team_guardrail_rows(
     prisma_client: "PrismaClient", where: Mapping[str, object]
 ) -> "Sequence[LiteLLM_GuardrailsTable]":
-    rows: Final[Sequence[LiteLLM_GuardrailsTable]] = await GuardrailsRepository(prisma_client).table.find_many(
+    rows: Final = await _guardrails_table(prisma_client).find_many(
         where=where,
         order={"created_at": "desc"},
     )
@@ -499,10 +519,12 @@ async def update_guardrail(
         if existing_guardrail is None:
             raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found")
 
-        result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
-            guardrail_id=guardrail_id,
-            guardrail=request.guardrail,
-            prisma_client=prisma_client,
+        result: Final = _as_str_object_mapping(
+            await GUARDRAIL_REGISTRY.update_guardrail_in_db(
+                guardrail_id=guardrail_id,
+                guardrail=request.guardrail,
+                prisma_client=prisma_client,
+            )
         )
 
         guardrail_name: Final = result.get("guardrail_name", "Unknown")
@@ -613,7 +635,7 @@ class RegisterGuardrailRequest(BaseModel):
     """Request body for POST /guardrails/register. Follows Generic Guardrail API config."""
 
     guardrail_name: str
-    litellm_params: dict[str, Any]  # guardrail, mode, api_base required; api_key, headers, etc. optional
+    litellm_params: dict[str, object]  # guardrail, mode, api_base required; api_key, headers, etc. optional
     guardrail_info: dict[str, object] | None = None
     team_id: str | None = None
 
@@ -1172,12 +1194,14 @@ async def patch_guardrail(
         )
 
         # Update litellm_params if default_on is provided or pii_entities_config is provided
-        litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {})))
+        existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {})))
+        litellm_params = LitellmParams(**existing_litellm_params)
         if request.litellm_params is not None:
             requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True)
             litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True)
             litellm_params_dict.update(requested_litellm_params)
-            litellm_params = LitellmParams(**litellm_params_dict)
+            merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict)
+            litellm_params = LitellmParams(**merged_litellm_params)
 
         # Update guardrail_info if provided
         guardrail_info: Final = (
@@ -1193,10 +1217,12 @@ async def patch_guardrail(
             litellm_params=litellm_params,
             guardrail_info=guardrail_info,
         )
-        result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
-            guardrail_id=guardrail_id,
-            guardrail=guardrail,
-            prisma_client=prisma_client,
+        result: Final = _as_str_object_mapping(
+            await GUARDRAIL_REGISTRY.update_guardrail_in_db(
+                guardrail_id=guardrail_id,
+                guardrail=guardrail,
+                prisma_client=prisma_client,
+            )
         )
 
         guardrail_name = result.get("guardrail_name", "Unknown")
@@ -1552,31 +1578,46 @@ async def validate_blocked_words_file(request: dict[str, str]):
         return {"valid": False, "error": f"Validation error: {e}"}
 
 
-def _get_field_type_from_annotation(field_annotation: Any) -> str:
+def _dunder_origin(annotation: object) -> object:
+    origin: Final[object] = getattr(annotation, "__origin__", None)
+    return origin
+
+
+def _dunder_name(annotation: object) -> object:
+    name: Final[object] = getattr(annotation, "__name__", None)
+    return name
+
+
+def _dunder_args(annotation: object) -> tuple[object, ...]:
+    args: Final[tuple[object, ...]] = getattr(annotation, "__args__", ())
+    return args
+
+
+def _get_field_type_from_annotation(field_annotation: object) -> str:
     """
     Convert a Python type annotation to a UI-friendly type string
     """
     # Handle Union types (like Optional[T])
     if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType:
         # For Optional[T], get the non-None type
-        args: Final = get_args(field_annotation)
+        args: Final[tuple[object, ...]] = get_args(field_annotation)
         non_none_args: Final = [arg for arg in args if arg is not type(None)]
         if non_none_args:
             field_annotation = non_none_args[0]
 
     # Handle List types
-    if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list:
+    if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is list:
         return "array"
 
     # Handle Dict types
-    if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict:
+    if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is dict:
         return "dict"
 
     # Handle Literal types
     if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"):
         # Check for Literal types (Python 3.8+)
-        origin: Final = field_annotation.__origin__
-        if hasattr(origin, "__name__") and origin.__name__ == "Literal":
+        origin: Final = _dunder_origin(field_annotation)
+        if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal":
             return "select"  # For dropdown/select inputs
 
     # Handle basic types
@@ -1595,66 +1636,66 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str:
     return "string"
 
 
-def _extract_literal_values(annotation: Any) -> list[str]:
+def _extract_literal_values(annotation: object) -> Sequence[object]:
     """
     Extract literal values from a Literal type annotation
     """
     if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"):
-        origin: Final = annotation.__origin__
-        if hasattr(origin, "__name__") and origin.__name__ == "Literal":
-            return list(annotation.__args__)
+        origin: Final = _dunder_origin(annotation)
+        if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal":
+            return list(_dunder_args(annotation))
     return []
 
 
-def _get_dict_key_options(field_annotation: Any) -> list[str] | None:
+def _get_dict_key_options(field_annotation: object) -> Sequence[object] | None:
     """
     Extract key options from Dict[Literal[...], T] types
     """
     if (
         hasattr(field_annotation, "__origin__")
-        and field_annotation.__origin__ is dict
+        and _dunder_origin(field_annotation) is dict
         and hasattr(field_annotation, "__args__")
     ):
-        args: Final = field_annotation.__args__
+        args: Final = _dunder_args(field_annotation)
         if len(args) >= 2:
             key_type: Final = args[0]
             return _extract_literal_values(key_type)
     return None
 
 
-def _get_dict_value_type(field_annotation: Any) -> str:
+def _get_dict_value_type(field_annotation: object) -> str:
     """
     Get the value type from Dict[K, V] types
     """
     if (
         hasattr(field_annotation, "__origin__")
-        and field_annotation.__origin__ is dict
+        and _dunder_origin(field_annotation) is dict
         and hasattr(field_annotation, "__args__")
     ):
-        args: Final = field_annotation.__args__
+        args: Final = _dunder_args(field_annotation)
         if len(args) >= 2:
             value_type: Final = args[1]
             return _get_field_type_from_annotation(value_type)
     return "string"
 
 
-def _get_list_element_options(field_annotation: Any) -> list[str] | None:
+def _get_list_element_options(field_annotation: object) -> Sequence[object] | None:
     """
     Extract element options from List[Literal[...]] types
     """
     if (
         hasattr(field_annotation, "__origin__")
-        and field_annotation.__origin__ is list
+        and _dunder_origin(field_annotation) is list
         and hasattr(field_annotation, "__args__")
     ):
-        args: Final = field_annotation.__args__
+        args: Final = _dunder_args(field_annotation)
         if len(args) >= 1:
             element_type: Final = args[0]
             return _extract_literal_values(element_type)
     return None
 
 
-def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool:
+def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool:
     """Check if optional_params field should be skipped (not meaningfully overridden)."""
     if field_name != "optional_params":
         return False
@@ -1664,12 +1705,12 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool
 
     # Check if the annotation is still a generic TypeVar (not specialized)
     if isinstance(field_annotation, TypeVar) or (
-        hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar
+        hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is TypeVar
     ):
         return True
 
     # Also skip if it's a generic type that wasn't specialized
-    if hasattr(field_annotation, "__name__") and field_annotation.__name__ in (
+    if hasattr(field_annotation, "__name__") and _dunder_name(field_annotation) in (
         "T",
         "TypeVar",
     ):
@@ -1677,18 +1718,18 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool
 
     # Handle Optional[T] where T is still a TypeVar
     if hasattr(field_annotation, "__args__"):
-        non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)]
+        non_none_args: Final = [arg for arg in _dunder_args(field_annotation) if arg is not type(None)]
         if non_none_args and isinstance(non_none_args[0], TypeVar):
             return True
 
     return False
 
 
-def _unwrap_optional_type(field_annotation: Any) -> Any:
+def _unwrap_optional_type(field_annotation: object) -> object:
     """Unwrap Optional types to get the actual type."""
     if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType:
         # For Optional[BaseModel], get the non-None type
-        args: Final = get_args(field_annotation)
+        args: Final[tuple[object, ...]] = get_args(field_annotation)
         non_none_args: Final = [arg for arg in args if arg is not type(None)]
         if non_none_args:
             return non_none_args[0]
@@ -1696,20 +1737,20 @@ def _unwrap_optional_type(field_annotation: Any) -> Any:
 
 
 def _build_field_dict(
-    field: Any,
-    field_annotation: Any,
+    field: "FieldInfo",
+    field_annotation: object,
     description: str,
     required: bool,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """Build field dictionary for non-nested fields."""
     # Determine the field type from annotation
     field_type = _get_field_type_from_annotation(field_annotation)
 
     # Check for custom UI type override
-    field_json_schema_extra: Final = getattr(field, "json_schema_extra", {})
+    field_json_schema_extra: Final[Mapping[str, object]] = getattr(field, "json_schema_extra", {})
     if field_json_schema_extra and "ui_type" in field_json_schema_extra:
         ui_type: Final = field_json_schema_extra["ui_type"]
-        field_type = ui_type.value if hasattr(ui_type, "value") else ui_type
+        field_type = getattr(ui_type, "value", ui_type)
     elif field_json_schema_extra and "type" in field_json_schema_extra:
         field_type = field_json_schema_extra["type"]
 
@@ -1748,8 +1789,9 @@ def _build_field_dict(
             field_dict["options"] = literal_options
 
     # Add default value if it exists
-    if field.default is not None and field.default is not ...:
-        field_dict["default_value"] = field.default
+    field_default: Final[object] = getattr(field, "default", None)
+    if field_default is not None and field_default is not ...:
+        field_dict["default_value"] = field_default
 
     # Copy min, max, step from json_schema_extra for number/percentage inputs
     if field_json_schema_extra:
@@ -1763,7 +1805,7 @@ def _build_field_dict(
 def _extract_fields_recursive(
     model: type[BaseModel],
     depth: int = 0,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     # Check if we've exceeded the maximum recursion depth
     if depth > DEFAULT_MAX_RECURSE_DEPTH:
         raise HTTPException(
@@ -1817,7 +1859,7 @@ def _extract_fields_recursive(
     return fields
 
 
-def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]:
+def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]:
     """
     Get the fields from a Pydantic model as a nested dictionary structure
     """
@@ -2141,7 +2183,26 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type:
     return "response" if input_type == "response" else "request"
 
 
-def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None:
+class _GuardrailLoggingObj(Protocol):
+    call_type: str
+    model_call_details: dict[str, object]
+
+    @property
+    def update_messages(self) -> "Callable[..., object]": ...
+
+    @property
+    def async_success_handler(self) -> "Callable[..., Awaitable[object]]": ...
+
+    @property
+    def success_handler(self) -> "Callable[..., object]": ...
+
+
+class _GuardrailProxyLogging(Protocol):
+    @property
+    def post_call_success_hook(self) -> "Callable[..., Awaitable[object]]": ...
+
+
+def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj, request: ApplyGuardrailRequest) -> None:
     """Configure the logging object so Langfuse/OTEL extract input and output correctly."""
     litellm_logging_obj.call_type = "pass_through_endpoint"
     litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint"
@@ -2151,8 +2212,8 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGua
 
 
 async def _emit_guardrail_success_logs(
-    proxy_logging_obj: Any,
-    litellm_logging_obj: Any,
+    proxy_logging_obj: _GuardrailProxyLogging,
+    litellm_logging_obj: _GuardrailLoggingObj | None,
     data: dict,
     user_api_key_dict: UserAPIKeyAuth,
     response: ApplyGuardrailResponse,

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -9,11 +9,14 @@ import os
 import sys
 
 sys.path.insert(0, os.path.abspath("../.."))  # Adds the parent directory to the system path
+import asyncio
 import copy
 import json
+import re
 import sys
-from collections.abc import AsyncGenerator, Mapping
+from collections.abc import AsyncGenerator, Mapping, Sequence
 from datetime import datetime, timezone
+from itertools import accumulate, groupby
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
 
 import httpx
@@ -23,6 +26,7 @@ from pydantic import TypeAdapter, ValidationError
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import DualCache
+from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
 from litellm.exceptions import ModifyResponseException
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
@@ -46,6 +50,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
     BedrockGuardrailOutput,
     BedrockGuardrailQualifier,
     BedrockGuardrailResponse,
+    BedrockGuardrailUsage,
     BedrockRequest,
     BedrockTextContent,
 )
@@ -53,6 +58,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
     from botocore.awsrequest import AWSPreparedRequest
+    from botocore.credentials import Credentials
 
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
@@ -71,6 +77,17 @@ from litellm.types.utils import (
 
 GUARDRAIL_NAME: Final = "bedrock"
 _BEDROCK_DYNAMIC_BODY_DENYLIST: Final = frozenset({"content", "source"})
+_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = (
+    "text unit",
+    "maximum input size",
+    "content size",
+    "too long",
+    "too large",
+    "exceeds the maximum",
+)
+_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3
+_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5
+_BEDROCK_WHITESPACE: Final = re.compile(r"\s")
 # Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required).
 _BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke"
 # InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with
@@ -118,6 +135,29 @@ class GuardrailMessageFilterResult(NamedTuple):
     target_indices: list[int] | None
 
 
+class BedrockContentChunkResult(NamedTuple):
+    """One chunk's ApplyGuardrail response, paired with enough bookkeeping to
+    reconstruct global masked-output positions once every chunk is back.
+
+    `content` is the exact content items this chunk was called with -- needed
+    so an all-clear chunk (empty `outputs`) can still contribute one unmasked
+    placeholder per item it covers, keeping every later chunk's masked text
+    aligned to its original global position. `fragment_group_size` is 1 for an
+    ordinary chunk, and otherwise the total number of consecutive chunk results
+    that together make up ONE original content item's own text (split because a
+    list of length 1 could not be bisected by list length). All of them must be
+    concatenated back into that one item's masked output rather than treated as
+    separate items. It is a count rather than a boolean because one item can be
+    bisected more than once: two levels of splitting produce four fragments for
+    a single item, not two, and grouping them in fixed pairs would emit two
+    outputs for one message and shift every later message's masked text.
+    """
+
+    response: BedrockGuardrailResponse
+    content: tuple[BedrockContentItem, ...]
+    fragment_group_size: int
+
+
 class ApplyGuardrailMessageSelection(NamedTuple):
     """Messages selected for an apply_guardrail scan + write-back metadata."""
 
@@ -168,12 +208,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         content_filter_threshold: float | None = 0.5,
         prompt_attack_threshold: float | None = 0.5,
         pii_confidence_threshold: float | None = 0.5,
+        chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
         **kwargs,
     ):
         self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
         self.guardrailIdentifier = guardrailIdentifier
         self.guardrailVersion = guardrailVersion
         self.guardrail_provider = "bedrock"
+        self.chunk_budget_chars = chunk_budget_chars
         self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
 
         # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
@@ -759,12 +801,35 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         request_data: dict | None = None,
         logging_event_type: GuardrailEventHooks | None = None,
     ) -> BedrockGuardrailResponse:
+        """Scan `messages`/`response` with ApplyGuardrail, chunking if it is too large.
+
+        Content is bin-packed into budget-sized batches and each batch posted
+        sequentially, every batch independently falling back to bisection if AWS
+        rejects it. The per-batch responses are merged so callers cannot tell whether
+        chunking happened.
+
+        Content using contextual grounding opts out of chunking entirely: grounding is
+        scored holistically against the whole reference source, so bisecting it would
+        fragment that evaluation and yield misleading scores. Such a request keeps the
+        old behavior of surfacing a too-large error rather than being split.
+
+        `logging_event_type` drives what UI and spend logs report. It is distinct from
+        Bedrock's `source`, which is INPUT vs OUTPUT for the API body and must not be
+        confused with the proxy hook (pre_call / during_call / post_call); when omitted,
+        the legacy source-derived mapping is kept for backward compatibility.
+
+        A guardrail *block* is logged where it happens, in
+        `_post_apply_guardrail_content`, because chunking stops immediately and there is
+        no later merged response to log instead. Everything else that fails out of the
+        chunking flow (an unrecoverable too-large error, a non-size validation error,
+        exhausted throttle retries) is a genuine end-to-end failure of this one logical
+        guardrail call and is logged exactly once here.
+        """
         start_time: Final = datetime.now(timezone.utc)
         credentials, aws_region_name = self._load_credentials()
         bedrock_request_data: Final[dict] = dict(
             self.convert_to_bedrock_format(source=source, messages=messages, response=response)
         )
-        bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse()
         api_key: str | None = None
         if request_data:
             dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
@@ -778,6 +843,257 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             if request_data.get("api_key") is not None:
                 api_key = request_data["api_key"]
 
+        event_type: Final = (
+            logging_event_type
+            if logging_event_type is not None
+            else (GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call)
+        )
+
+        content: Final[tuple[BedrockContentItem, ...]] = tuple(bedrock_request_data.get("content") or ())
+        allow_chunking: Final = not self._content_uses_contextual_grounding(content)
+
+        try:
+            responses: Final = await self._apply_guardrail_content_with_chunking(
+                content=content,
+                base_request_data=bedrock_request_data,
+                credentials=credentials,
+                aws_region_name=aws_region_name,
+                api_key=api_key,
+                request_data=request_data,
+                event_type=event_type,
+                start_time=start_time,
+                allow_chunking=allow_chunking,
+            )
+        except HTTPException as exc:
+            if not isinstance(exc.detail, dict):
+                self._log_apply_guardrail_failure(
+                    detail=exc.detail,
+                    request_data=request_data,
+                    event_type=event_type,
+                    start_time=start_time,
+                )
+            raise
+        merged_response: Final = self._merge_bedrock_guardrail_responses(responses)
+        self._log_apply_guardrail_success(
+            merged_response=merged_response,
+            request_data=request_data,
+            event_type=event_type,
+            start_time=start_time,
+        )
+        return merged_response
+
+    async def _apply_guardrail_content_with_chunking(
+        self,
+        content: Sequence[BedrockContentItem],
+        base_request_data: Mapping[str, Any],
+        credentials: "Credentials",
+        aws_region_name: str,
+        api_key: str | None,
+        request_data: dict | None,  # mutable-ok: proxy request body dict, mutated by the logging helper
+        event_type: GuardrailEventHooks,
+        start_time: "datetime",
+        allow_chunking: bool,
+    ) -> tuple[BedrockContentChunkResult, ...]:
+        """Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large.
+
+        Tries `content` as a single call first. AWS's per-request "maximum input
+        size in text units" quota is account/region/policy-dependent and cannot be
+        predicted ahead of time, so it is only ever discovered reactively: on an
+        error whose message indicates the input was too large (a ThrottlingException
+        in practice, a ValidationException per the docs -- see
+        ``_is_input_too_large_error``), the content is re-sent in smaller pieces.
+
+        Probing with the whole payload first is what keeps a request AWS would have
+        accepted at exactly one call. Packing into fixed batches up front instead
+        would split conversations AWS was happy to take whole, multiplying billed
+        calls and guardrail latency on traffic that never had a size problem, and
+        no fixed budget can avoid that because the real cap is unknown here.
+
+        Once a rejection proves the payload is over the cap, a multi-item payload is
+        re-sent as ``chunk_budget_chars``-sized batches rather than bisected: that
+        reaches a working size in one step instead of paying an O(log n) ladder of
+        rejected calls. Bisection remains the fallback for anything bin-packing
+        cannot make smaller, which is what makes the recursion terminate: a batch
+        already inside the budget packs back to itself, so it falls through to the
+        split below. A single oversized
+        content item (one very long message) is split by its own text instead of
+        by list length, since a list of length 1 has no items left to bisect --
+        the resulting fragments all carry a ``fragment_group_size`` so the merge
+        step can recombine them into the one content item they came from, rather
+        than treating each fragment as its own item when reconstructing positions
+        for masking. That count covers however many fragments the item ended up
+        split into, not just two, since it can be bisected repeatedly: the
+        outermost single-item split stamps the total leaf count on every leaf
+        below it, overwriting any smaller count an inner split had set. A real
+        guardrail block on any (sub-)chunk raises immediately
+        -- callers must not lose that signal by continuing to post the remaining
+        chunks.
+        """
+        try:
+            response: Final = await self._post_apply_guardrail_content_with_retry(
+                content=content,
+                base_request_data=base_request_data,
+                credentials=credentials,
+                aws_region_name=aws_region_name,
+                api_key=api_key,
+                request_data=request_data,
+                event_type=event_type,
+                start_time=start_time,
+            )
+            return (
+                BedrockContentChunkResult(
+                    response=response,
+                    content=tuple(content),
+                    fragment_group_size=1,
+                ),
+            )
+        except HTTPException as exc:
+            if allow_chunking and self._is_input_too_large_error(exc.detail):
+                batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars)
+                if len(batches) > 1:
+                    verbose_proxy_logger.warning(
+                        "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; "
+                        "re-sending as %d batches of at most %d characters",
+                        len(content),
+                        len(batches),
+                        self.chunk_budget_chars,
+                    )
+                    batch_results: Final = [  # mutable-ok: await needs a list comprehension; frozen to a tuple below
+                        await self._apply_guardrail_content_with_chunking(
+                            content=batch,
+                            base_request_data=base_request_data,
+                            credentials=credentials,
+                            aws_region_name=aws_region_name,
+                            api_key=api_key,
+                            request_data=request_data,
+                            event_type=event_type,
+                            start_time=start_time,
+                            allow_chunking=allow_chunking,
+                        )
+                        for batch in batches
+                    ]
+                    return tuple(result for results in batch_results for result in results)
+                split_content: Final = self._split_bedrock_content(content)
+                if split_content is None:
+                    raise
+                first_half, second_half = split_content
+                is_single_item_text_split: Final = len(content) == 1
+                verbose_proxy_logger.warning(
+                    "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; "
+                    "splitting into %d + %d and retrying each",
+                    len(content),
+                    len(first_half),
+                    len(second_half),
+                )
+                first_results: Final = await self._apply_guardrail_content_with_chunking(
+                    content=first_half,
+                    base_request_data=base_request_data,
+                    credentials=credentials,
+                    aws_region_name=aws_region_name,
+                    api_key=api_key,
+                    request_data=request_data,
+                    event_type=event_type,
+                    start_time=start_time,
+                    allow_chunking=allow_chunking,
+                )
+                second_results: Final = await self._apply_guardrail_content_with_chunking(
+                    content=second_half,
+                    base_request_data=base_request_data,
+                    credentials=credentials,
+                    aws_region_name=aws_region_name,
+                    api_key=api_key,
+                    request_data=request_data,
+                    event_type=event_type,
+                    start_time=start_time,
+                    allow_chunking=allow_chunking,
+                )
+                combined_results: Final = tuple(first_results) + tuple(second_results)
+                if is_single_item_text_split:
+                    return tuple(
+                        result._replace(fragment_group_size=len(combined_results)) for result in combined_results
+                    )
+                return combined_results
+            raise
+
+    async def _post_apply_guardrail_content_with_retry(
+        self,
+        content: Sequence[BedrockContentItem],
+        base_request_data: Mapping[str, Any],
+        credentials: "Credentials",
+        aws_region_name: str,
+        api_key: str | None,
+        request_data: dict | None,  # mutable-ok: proxy request body dict, mutated by the logging helper
+        event_type: GuardrailEventHooks,
+        start_time: "datetime",
+    ) -> BedrockGuardrailResponse:
+        """Post one ApplyGuardrail call for `content`, retrying with exponential
+        backoff on AWS ThrottlingException (HTTP 429).
+
+        Chunking already trades one oversized call for several smaller ones, so
+        retries here are capped low -- they must not multiply per-request latency
+        by an order of magnitude when the account's per-second text-unit quota is
+        the binding constraint rather than the per-request size quota.
+
+        A too-large rejection is deliberately excluded from the retry. AWS reports
+        it as a ThrottlingException (429), not only as a ValidationException, but
+        unlike a genuine throttle it is not transient: re-posting the same
+        oversized content can never succeed. Retrying it would burn every backoff
+        sleep and every (billed) attempt before the caller's bisection gets a
+        chance to split the content, at every level of the recursion.
+        """
+        for attempt in range(_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + 1):
+            try:
+                return await self._post_apply_guardrail_content(
+                    content=content,
+                    base_request_data=base_request_data,
+                    credentials=credentials,
+                    aws_region_name=aws_region_name,
+                    api_key=api_key,
+                    request_data=request_data,
+                    event_type=event_type,
+                    start_time=start_time,
+                )
+            except HTTPException as exc:
+                if (
+                    exc.status_code != 429
+                    or self._is_input_too_large_error(exc.detail)
+                    or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES
+                ):
+                    raise
+                await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt))
+        raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted")
+
+    async def _post_apply_guardrail_content(
+        self,
+        content: Sequence[BedrockContentItem],
+        base_request_data: Mapping[str, Any],
+        credentials: "Credentials",
+        aws_region_name: str,
+        api_key: str | None,
+        request_data: dict | None,  # mutable-ok: proxy request body dict, mutated by the logging helper
+        event_type: GuardrailEventHooks,
+        start_time: "datetime",
+    ) -> BedrockGuardrailResponse:
+        """Make exactly one signed ApplyGuardrail HTTP call for `content` and
+        parse the result. Raises HTTPException on a guardrail block or any
+        non-200 response (including 429, handled by the retry wrapper above).
+
+        AWS also reports some failures inside a 200 body, tagging ``Output.__type``
+        with an Exception marker. Those deliberately do NOT raise: the request proceeds,
+        matching the behaviour of this code before chunking existed. The marker survives
+        the merge, so the one consolidated log entry still records
+        ``guardrail_failed_to_respond`` rather than a success. Making that path fail
+        closed is a separate change, tracked apart from this PR, and belongs behind the
+        existing ``unreachable_fallback`` setting rather than a hardcoded status.
+
+        A block is logged here rather than by the caller: it ends the whole chunking
+        flow immediately, with no further chunks attempted, so there is no later
+        merged response for the caller to log instead.
+        """
+        bedrock_request_data: Final = {  # mutable-ok: outbound JSON request body
+            **base_request_data,
+            "content": content,
+        }  # mutable-ok: outbound JSON request body
         prepared_request: Final = self._prepare_request(
             credentials=credentials,
             data=bedrock_request_data,
@@ -792,42 +1108,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             prepared_request.headers,
         )
 
-        # UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API
-        # body, which must not be confused with the proxy hook (pre_call / during_call /
-        # post_call). When omitted, keep legacy mapping for backward compatibility.
-        if logging_event_type is not None:
-            event_type = logging_event_type
-        else:
-            event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call
-
         httpx_response: Final = await self._sign_and_post(
             prepared_request=prepared_request,
             request_data=request_data,
             event_type=event_type,
             start_time=start_time,
+            log_transport_failure=False,
         )
 
-        #########################################################
-        # Add guardrail information to request trace
-        #########################################################
-        _json_response: Final = httpx_response.json()
-        tracing_detail: Final = self._build_tracing_detail(_json_response)
-
-        # Raw Bedrock JSON is passed here; match/regex redaction runs once inside
-        # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
-        self.add_standard_logging_guardrail_information_to_request_data(
-            guardrail_provider=self.guardrail_provider,
-            guardrail_json_response=_json_response,
-            request_data=request_data or {},
-            guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
-            start_time=start_time.timestamp(),
-            end_time=datetime.now(timezone.utc).timestamp(),
-            duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
-            event_type=event_type,
-            tracing_detail=tracing_detail or None,
-        )
-        #########################################################
         if httpx_response.status_code == 200:
+            _json_response: Final = httpx_response.json()
             # check if the response was flagged
             verbose_proxy_logger.debug(
                 "Bedrock AI response : %s",
@@ -835,19 +1125,462 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             )
             bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
             if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
+                self._log_apply_guardrail_attempt(
+                    httpx_response=httpx_response,
+                    json_response=_json_response,
+                    request_data=request_data,
+                    event_type=event_type,
+                    start_time=start_time,
+                )
                 raise self._get_http_exception_for_blocked_guardrail(
                     bedrock_guardrail_response, request_data=request_data
                 )
-        else:
-            status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
-            verbose_proxy_logger.error(
-                "Bedrock AI: error in response. Status code: %s, response: %s",
-                httpx_response.status_code,
-                httpx_response.text,
-            )
-            raise HTTPException(status_code=status_code, detail=detail_message)
+            return bedrock_guardrail_response
 
-        return bedrock_guardrail_response
+        status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
+        verbose_proxy_logger.error(
+            "Bedrock AI: error in response. Status code: %s, response: %s",
+            httpx_response.status_code,
+            httpx_response.text,
+        )
+        raise HTTPException(status_code=status_code, detail=detail_message)
+
+    def _log_apply_guardrail_attempt(
+        self,
+        httpx_response: httpx.Response,
+        json_response: dict,  # mutable-ok: raw AWS JSON payload
+        request_data: dict | None,  # mutable-ok: proxy request body dict, mutated by the logging helper
+        event_type: GuardrailEventHooks,
+        start_time: "datetime",
+    ) -> None:
+        """Log a single ApplyGuardrail HTTP attempt as-is (its own status,
+        derived from its own response). Used only for the blocked-content
+        case, which ends the whole chunking flow immediately."""
+        tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response))
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_provider=self.guardrail_provider,
+            guardrail_json_response=json_response,
+            request_data=request_data or {},  # mutable-ok: logging helper requires a dict
+            guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
+            start_time=start_time.timestamp(),
+            end_time=datetime.now(timezone.utc).timestamp(),
+            duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
+            event_type=event_type,
+            tracing_detail=tracing_detail or None,
+        )
+
+    def _log_apply_guardrail_success(
+        self,
+        merged_response: BedrockGuardrailResponse,
+        request_data: dict | None,  # mutable-ok: proxy request body dict, mutated by the logging helper
+        event_type: GuardrailEventHooks,
+        start_time: "datetime",
+    ) -> None:
+        """Log one logical ApplyGuardrail call -- possibly several chunk calls
+        under the hood -- using its final merged response, so a chunked
+        request produces exactly one telemetry entry, the same as an
+        unchunked one would.
+
+        AWS can report a failure inside an HTTP 200 body by tagging
+        ``Output.__type`` with an exception marker. That marker survives the merge,
+        so the status is derived from the merged response rather than assumed to be
+        a success, which is what the pre-chunking code reported for that shape."""
+        tracing_detail: Final = self._build_tracing_detail(merged_response)
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_provider=self.guardrail_provider,
+            guardrail_json_response=dict(merged_response),  # mutable-ok: logging helper requires a dict
+            request_data=request_data or {},  # mutable-ok: logging helper requires a dict
+            guardrail_status=(
+                "guardrail_failed_to_respond"
+                if "Exception" in str((merged_response.get("Output") or {}).get("__type", ""))
+                else "success"
+            ),
+            start_time=start_time.timestamp(),
+            end_time=datetime.now(timezone.utc).timestamp(),
+            duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
+            event_type=event_type,
+            tracing_detail=tracing_detail or None,
+        )
+
+    def _log_apply_guardrail_failure(
+        self,
+        detail: object,
+        request_data: dict | None,  # mutable-ok: proxy request body dict, mutated by the logging helper
+        event_type: GuardrailEventHooks,
+        start_time: "datetime",
+    ) -> None:
+        """Log one logical ApplyGuardrail call that failed end-to-end (an
+        unrecoverable too-large error, a non-size validation error, or
+        exhausted throttle retries) as a single failure, rather than logging
+        every failed attempt chunking made along the way."""
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_provider=self.guardrail_provider,
+            guardrail_json_response={"error": str(detail)},  # mutable-ok: logging helper requires a dict
+            request_data=request_data or {},  # mutable-ok: logging helper requires a dict
+            guardrail_status="guardrail_failed_to_respond",
+            start_time=start_time.timestamp(),
+            end_time=datetime.now(timezone.utc).timestamp(),
+            duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
+            event_type=event_type,
+        )
+
+    @staticmethod
+    def _content_uses_contextual_grounding(content: Sequence[BedrockContentItem]) -> bool:
+        """True if any content item carries a contextual-grounding qualifier
+        (``grounding_source``, ``query``, or the ``guard_content`` the response
+        itself is tagged with once grounding is present)."""
+        for item in content:
+            if (item.get("text") or {}).get("qualifiers"):  # mutable-ok: read-only empty fallback
+                return True
+        return False
+
+    @staticmethod
+    def _bin_pack_bedrock_content(
+        content: Sequence[BedrockContentItem],
+        budget: int,
+    ) -> tuple[tuple[BedrockContentItem, ...], ...]:
+        """Pack whole content items, in order, into batches whose combined text
+        length stays within `budget`, in a single pass that carries the running
+        total rather than re-summing the open batch per item.
+
+        This is the fast-path half of the hybrid chunking strategy: bin-packing
+        at a conservative fixed budget keeps the common case at O(n / budget)
+        ApplyGuardrail calls instead of the O(log n) round trips pure reactive
+        bisection pays on every oversized request. An item whose own text
+        already exceeds `budget` is not split here -- it becomes its own
+        (still oversized) batch and is sent as-is; if AWS rejects that batch as
+        too large, `_apply_guardrail_content_with_chunking`'s existing
+        recursive-bisection fallback takes over for that batch only.
+
+        `budget` comes from the guardrail's ``chunk_budget_chars`` setting and
+        defaults to 25,000, matching ApplyGuardrail's default quota of 25 text
+        units (roughly 1,000 characters each) per second. Packing to that size and
+        posting sequentially is what keeps chunking from tripping the rate quota
+        and trading a size error for a throttle. Accounts with raised quotas can
+        configure a larger budget to spend fewer calls.
+
+        The budget is not a correctness dependency either way. AWS's effective cap
+        varies by account, region, and policy, is not a fixed character count, and
+        cannot be read from config, so any batch it still rejects falls back to
+        bisection, which self-corrects however wrong the value was. An over-large
+        budget therefore costs one extra probe-and-bisect round trip rather than
+        failing the request.
+        """
+        if not content:
+            return (tuple(content),)
+
+        lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content)
+
+        def assign(carried: tuple[int, int], length: int) -> tuple[int, int]:
+            batch_index, used = carried
+            if used + length <= budget:
+                return batch_index, used + length
+            return batch_index + 1, length
+
+        batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:])
+        return tuple(
+            tuple(item for _, item in group)
+            for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0])
+        )
+
+    @staticmethod
+    def _split_bedrock_content(
+        content: Sequence[BedrockContentItem],
+    ) -> tuple[tuple[BedrockContentItem, ...], tuple[BedrockContentItem, ...]] | None:
+        """Bisect `content` into two roughly-equal, non-empty halves.
+
+        When `content` already holds more than one item, it is split by list
+        length. When it holds exactly one item, that item's own text is split
+        instead (a list of length 1 has no items left to bisect, but one very
+        long message is still a single content item) -- at the whitespace
+        character nearest the midpoint rather than a raw character index, so
+        the cut never lands inside a word/token. This is a plain, lossless
+        cut with no overlap: concatenating the two fragments in order always
+        reproduces the original text exactly, so merging back at
+        ``_merge_logical_unit_outputs`` needs no reconciliation step.
+
+        Known, accepted limitation: whitespace splitting only guards against
+        *accidentally* severing a single token (one denied word, one PII
+        pattern) across the cut. It does not, and cannot without an overlap
+        window, stop a *multi-word* denied phrase deliberately positioned to
+        straddle the boundary -- each fragment can scan clean on its own and
+        still reassemble into the flagged phrase. AWS's own guidance on this
+        API acknowledges the same gap for input chunking ("a critical piece of
+        text could span two (or more) chunks if not carefully divided") with
+        no documented resolution, and overlap-and-reconcile was evaluated and
+        rejected for this PR: AWS's masking output has no documented
+        length-preservation guarantee, so reconciling an overlap region against
+        masked text is not sound in general. Out of scope for this PR.
+
+        Returns None when there is nothing left to split -- a single item
+        whose text is too short to halve into two non-empty pieces -- so the
+        caller can give up and propagate the original too-large error instead
+        of recursing forever.
+        """
+        if len(content) > 1:
+            midpoint: Final = max(1, len(content) // 2)
+            return tuple(content[:midpoint]), tuple(content[midpoint:])
+
+        text_content: Final = content[0].get("text") or BedrockTextContent()
+        text: Final = text_content.get("text") or ""
+        if len(text) < 2:
+            return None
+        split_at: Final = BedrockGuardrail._nearest_whitespace_split_index(text)
+        qualifiers: Final = text_content.get("qualifiers")
+
+        def fragment(piece: str) -> BedrockContentItem:
+            block: Final = (
+                BedrockTextContent(text=piece, qualifiers=qualifiers) if qualifiers else BedrockTextContent(text=piece)
+            )
+            return BedrockContentItem(text=block)
+
+        return (fragment(text[:split_at]),), (fragment(text[split_at:]),)
+
+    @staticmethod
+    def _nearest_whitespace_split_index(text: str) -> int:
+        """Return the index nearest `text`'s midpoint that falls on a whitespace
+        boundary, so splitting `text[:i]` / `text[i:]` there never severs a word.
+
+        Any Unicode whitespace counts, not just an ASCII space. Matching only `" "`
+        would leave the boundary unguarded for exactly the payloads that get large
+        enough to need splitting: JSON lines, source code, logs and transcripts are
+        newline or tab delimited, so a deny-listed word sitting at the midpoint of
+        one would be cut in half, scan clean on both fragments, and reassemble
+        intact.
+
+        The returned index always leaves both sides non-empty, which is what makes
+        the caller's recursion terminate. A boundary that would put the split at 0
+        or at ``len(text)`` is discarded: it would hand back a fragment identical to
+        the text just rejected as too large, AWS would reject that again, and each
+        retry would re-split it into the same unchanged fragment until the stack ran
+        out. The dangerous shape is a text whose only space at or after the midpoint
+        is its final character.
+
+        Falls back to the raw midpoint when no usable whitespace boundary exists, either
+        because `text` has none at all (a single giant token) or because the only
+        candidates were degenerate. That is still a correct, lossless split, just no
+        longer guaranteed word-safe for those cases. `text` must be at least two
+        characters, which `_split_bedrock_content` guarantees, so the midpoint itself
+        is never degenerate.
+        """
+        midpoint: Final = len(text) // 2
+        before: Final = max((found.end() for found in _BEDROCK_WHITESPACE.finditer(text, 0, midpoint)), default=None)
+        after_match: Final = _BEDROCK_WHITESPACE.search(text, midpoint)
+        candidates: Final = sorted(
+            (split for split in (before, after_match.end() if after_match else None) if split is not None),
+            key=lambda split: abs(split - midpoint),
+        )
+        return next((split for split in candidates if 0 < split < len(text)), midpoint)
+
+    @staticmethod
+    def _is_input_too_large_error(detail: object) -> bool:
+        """True if `detail` is an AWS error message for input exceeding the
+        per-request text-unit quota.
+
+        Matched on the message rather than the status code on purpose: AWS is not
+        consistent about which error it raises for this. Observed against a live
+        guardrail with an active content-filter policy, an oversized request comes
+        back as a *ThrottlingException* (429) reading ``Input text size (3273 text
+        units) exceeds the maximum allowed (1000 text units) for the content filter
+        policy (Classic tier)``, while the documented failure mode is a
+        ValidationException (400). Keying off the message covers both.
+
+        A guardrail *block* is also raised as an HTTPException with status 400,
+        but its ``detail`` is always a dict (built by
+        ``_get_http_exception_for_blocked_guardrail``); a non-200 API error's
+        ``detail`` is always the plain string returned by
+        ``_parse_bedrock_guardrail_error_response``. Checking ``isinstance(detail,
+        str)`` is therefore sufficient to never mistake a real block for a
+        too-large error.
+        """
+        if not isinstance(detail, str):
+            return False
+        lowered: Final = detail.lower()
+        return any(substring in lowered for substring in _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS)
+
+    @staticmethod
+    def _merge_bedrock_guardrail_responses(
+        chunk_results: Sequence[BedrockContentChunkResult],
+    ) -> BedrockGuardrailResponse:
+        """Merge the per-chunk ApplyGuardrail responses of a chunked request into
+        one, so a caller cannot tell whether chunking happened.
+
+        Only ever called with responses that all passed (a block raises
+        immediately from ``_apply_guardrail_content_with_chunking`` and is never
+        added to this list). ``action`` is only set on the merged response when
+        at least one chunk's raw response included it, and left absent otherwise
+        -- mirroring a real single-call response and matching what
+        ``_build_tracing_detail`` treats as "Bedrock didn't report an action".
+
+        Fields this merge has no opinion on (``actionReason``, ``guardrailCoverage``,
+        ``blockedResponse``, anything AWS adds later) are carried over from the chunk
+        responses rather than dropped, so the response and the logged telemetry keep
+        the shape a single unchunked call returned. The merged keys below win.
+
+        Per AWS's documented ApplyGuardrail contract, a single call's ``outputs``
+        is positionally parallel to the ``content`` items *of that call*: an
+        entry per item when anything in the call was masked, or an empty list
+        when nothing in the whole call was masked. Downstream masking
+        (``_apply_masking_to_messages``) walks the merged ``outputs`` by a single
+        running index across the *original, unchunked* message list, so a later
+        chunk's masked text must land at the same global position it would have
+        if chunking had never happened. Naively concatenating each chunk's
+        ``outputs`` breaks that whenever a chunk had nothing masked (its empty
+        list would otherwise silently swallow its items' slots, shifting every
+        later chunk's masked text left onto the wrong message). So every
+        item -- masked or not -- always contributes exactly one entry here,
+        falling back to that item's own original (unmasked) text when its
+        chunk returned no output for it; a wholly-untouched result is then
+        collapsed back to an empty ``outputs`` list to match a real single-call
+        no-op response. A chunk that returns a nonzero output count not equal
+        to its item count is passed through as-is instead of guessed at, since
+        AWS's docs don't cover partial masking within one multi-item call.
+        """
+        logical_units: Final = BedrockGuardrail._group_fragment_units(chunk_results)
+        per_unit_outputs: Final = tuple(BedrockGuardrail._merge_logical_unit_outputs(unit) for unit in logical_units)
+        merged_outputs: Final = [  # mutable-ok: logged payload; redaction only traverses dict/list
+            output for outputs, _ in per_unit_outputs for output in outputs
+        ]
+        any_masked: Final = any(masked for _, masked in per_unit_outputs)
+
+        actions: Final = tuple(
+            chunk_result.response.get("action")
+            for chunk_result in chunk_results
+            if isinstance(chunk_result.response.get("action"), str)
+        )
+        merged_action: Final = (
+            "GUARDRAIL_INTERVENED" if "GUARDRAIL_INTERVENED" in actions else (actions[-1] if actions else None)
+        )
+        merged_assessments: Final = [  # mutable-ok: logged payload; redaction only traverses dict/list
+            assessment
+            for chunk_result in chunk_results
+            for assessment in (chunk_result.response.get("assessments") or [])  # mutable-ok: logged payload
+        ]
+        any_usage_reported: Final = any(chunk_result.response.get("usage") for chunk_result in chunk_results)
+
+        merged: Final[BedrockGuardrailResponse] = cast(  # cast-ok: TypedDict assembled from a comprehension
+            BedrockGuardrailResponse,
+            {  # mutable-ok: builds the TypedDict payload
+                key: value for chunk_result in chunk_results for key, value in chunk_result.response.items()
+            },
+        )
+        if merged_action is not None:
+            merged["action"] = merged_action
+        if merged_outputs and any_masked:
+            merged["outputs"] = merged_outputs
+            merged["output"] = merged_outputs
+        if merged_assessments:
+            merged["assessments"] = merged_assessments
+        if any_usage_reported:
+            merged["usage"] = BedrockGuardrail._sum_bedrock_guardrail_usage(chunk_results)
+        return merged
+
+    @staticmethod
+    def _sum_bedrock_guardrail_usage(
+        chunk_results: Sequence[BedrockContentChunkResult],
+    ) -> BedrockGuardrailUsage:
+        """Sum each chunk's ``usage`` counters field-by-field into one totals dict.
+
+        Keys are taken from the responses rather than from a fixed list, so a counter
+        this code does not know about (AWS has added several) is still summed and
+        reported instead of being silently dropped to zero."""
+        chunk_usages: Final = tuple(
+            chunk_result.response.get("usage") or {}  # mutable-ok: read-only empty fallback
+            for chunk_result in chunk_results
+        )
+        return cast(  # cast-ok: TypedDict assembled from a comprehension
+            BedrockGuardrailUsage,
+            {  # mutable-ok: builds the TypedDict payload
+                key: sum(usage.get(key) or 0 for usage in chunk_usages)
+                for key in dict.fromkeys(key for usage in chunk_usages for key in usage)
+            },
+        )
+
+    @staticmethod
+    def _group_fragment_units(
+        chunk_results: Sequence[BedrockContentChunkResult],
+    ) -> tuple[tuple[BedrockContentChunkResult, ...], ...]:
+        """Group consecutive text-fragment chunk results back into the one content
+        item each group came from, leaving every ordinary chunk result as a unit of
+        one.
+
+        The group size is read off the results themselves rather than assumed,
+        because a single content item can be bisected repeatedly: two levels of
+        splitting yield four fragments for one item, not two. Assuming a fixed pair
+        here would emit two outputs for one message and shift every later message's
+        masked text onto the wrong message."""
+
+        def advance(carried: tuple[int, bool], result: BedrockContentChunkResult) -> tuple[int, bool]:
+            remaining, _ = carried
+            if remaining == 0:
+                return max(1, result.fragment_group_size) - 1, True
+            return remaining - 1, False
+
+        starts: Final = tuple(
+            index
+            for index, (_, starts_unit) in enumerate(tuple(accumulate(chunk_results, advance, initial=(0, False)))[1:])
+            if starts_unit
+        )
+        return tuple(tuple(chunk_results[start:end]) for start, end in zip(starts, starts[1:] + (len(chunk_results),)))
+
+    @staticmethod
+    def _merge_logical_unit_outputs(
+        unit: tuple[BedrockContentChunkResult, ...],
+    ) -> tuple[tuple[BedrockGuardrailOutput, ...], bool]:
+        """Reduce one logical unit (a fragment group of any size, or a single chunk
+        result) to the ``BedrockGuardrailOutput`` entries it contributes to the
+        merged response, plus whether any masking actually happened in it.
+
+        Per AWS's documented ApplyGuardrail contract, a single call's
+        ``outputs`` is positionally parallel to the ``content`` items *of that
+        call*: an entry per item when anything in the call was masked, or an
+        empty list when nothing in the whole call was masked. Downstream
+        masking (``_apply_masking_to_messages``) walks the merged ``outputs``
+        by a single running index across the *original, unchunked* message
+        list, so a later chunk's masked text must land at the same global
+        position it would have if chunking had never happened. So every item
+        -- masked or not -- always contributes exactly one entry here, falling
+        back to that item's own original (unmasked) text when its chunk
+        returned no output for it. A chunk that returns a nonzero output count
+        not equal to its item count is passed through as-is instead of guessed
+        at, since AWS's docs don't cover partial masking within one multi-item
+        call.
+
+        A unit holding more than one result is a fragment group: every result in it
+        is one fragment of a single content item's text, so the group collapses to
+        one entry built from each fragment's masked text (or that fragment's own
+        original text where it came back unmasked), concatenated in order. This
+        holds for any group size, not only two.
+        """
+        if len(unit) > 1:
+
+            def fragment_outputs(result: BedrockContentChunkResult) -> tuple[BedrockGuardrailOutput, ...]:
+                return tuple(result.response.get("outputs") or result.response.get("output") or ())
+
+            def fragment_text(result: BedrockContentChunkResult) -> str:
+                source: Final = (result.content[0].get("text") or {}).get(  # mutable-ok: read-only fallback
+                    "text"
+                ) or ""
+                outputs: Final = fragment_outputs(result)
+                masked: Final = outputs[0].get("text") if outputs else None
+                return masked if masked is not None else source
+
+            merged_text: Final = "".join(fragment_text(result) for result in unit)
+            any_masked: Final = any(fragment_outputs(result) for result in unit)
+            return (BedrockGuardrailOutput(text=merged_text),), any_masked
+
+        (chunk_result,) = unit
+        chunk_outputs: Final = chunk_result.response.get("outputs") or chunk_result.response.get("output") or ()
+        if len(chunk_outputs) == len(chunk_result.content):
+            return tuple(chunk_outputs), bool(chunk_outputs)
+        if not chunk_outputs:
+            return tuple(
+                BedrockGuardrailOutput(
+                    text=(item.get("text") or {}).get("text") or ""  # mutable-ok: read-only fallback
+                )
+                for item in chunk_result.content
+            ), False
+        return tuple(chunk_outputs), True
 
     async def _sign_and_post(
         self,
@@ -855,6 +1588,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         request_data: dict | None,
         event_type: GuardrailEventHooks,
         start_time: "datetime",
+        log_transport_failure: bool = True,
     ) -> httpx.Response:
         """POST a signed Bedrock request, logging+raising on network/HTTP errors.
 
@@ -862,6 +1596,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         transport-error handling cannot drift. Returns the raw ``httpx.Response`` on
         success (including non-2xx that httpx did not raise on); the 200-path logging,
         status and tracing stay with each caller because the two APIs report differently.
+
+        ``log_transport_failure=False`` suppresses the ``guardrail_failed_to_respond``
+        entry for a non-200 that is re-raised as an ``HTTPException``, for callers that
+        own consolidated per-request logging. The ApplyGuardrail path needs this:
+        ``AsyncHTTPHandler.post`` calls ``raise_for_status()``, so every non-200 lands
+        in this handler, and one logical request can legitimately produce several of
+        them (a too-large probe, then each rejected bisection level) while still
+        succeeding overall. Logging per attempt would report a recovered request as
+        several failures plus a success.
+
+        The connection-level branch below (timeout, endpoint down) still logs
+        unconditionally: it re-raises the original exception rather than an
+        ``HTTPException``, so no consolidating caller catches it, and suppressing it
+        would drop the only record of the failure.
         """
         try:
             return await self.async_handler.post(
@@ -882,16 +1630,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                         status_code,
                         detail_message,
                     ) = self._parse_bedrock_guardrail_error_response(err_response)
-                    self.add_standard_logging_guardrail_information_to_request_data(
-                        guardrail_provider=self.guardrail_provider,
-                        guardrail_json_response={"error": detail_message},
-                        request_data=request_data or {},
-                        guardrail_status="guardrail_failed_to_respond",
-                        start_time=start_time.timestamp(),
-                        end_time=datetime.now(timezone.utc).timestamp(),
-                        duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
-                        event_type=event_type,
-                    )
+                    if log_transport_failure:
+                        self.add_standard_logging_guardrail_information_to_request_data(
+                            guardrail_provider=self.guardrail_provider,
+                            guardrail_json_response={  # mutable-ok: logging helper requires a dict
+                                "error": detail_message
+                            },
+                            request_data=request_data or {},  # mutable-ok: logging helper requires a dict
+                            guardrail_status="guardrail_failed_to_respond",
+                            start_time=start_time.timestamp(),
+                            end_time=datetime.now(timezone.utc).timestamp(),
+                            duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
+                            event_type=event_type,
+                        )
                     raise HTTPException(status_code=status_code, detail=detail_message) from e
                 except HTTPException:
                     raise
@@ -900,7 +1651,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             self.add_standard_logging_guardrail_information_to_request_data(
                 guardrail_provider=self.guardrail_provider,
                 guardrail_json_response={"error": str(e)},
-                request_data=request_data or {},
+                request_data=request_data or {},  # mutable-ok: logging helper requires a dict
                 guardrail_status="guardrail_failed_to_respond",
                 start_time=start_time.timestamp(),
                 end_time=datetime.now(timezone.utc).timestamp(),
@@ -1027,7 +1778,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             self.add_standard_logging_guardrail_information_to_request_data(
                 guardrail_provider=self.guardrail_provider,
                 guardrail_json_response={"error": detail_message},
-                request_data=request_data or {},
+                request_data=request_data or {},  # mutable-ok: logging helper requires a dict
                 guardrail_status="guardrail_failed_to_respond",
                 start_time=start_time.timestamp(),
                 end_time=datetime.now(timezone.utc).timestamp(),
@@ -1043,7 +1794,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             self.add_standard_logging_guardrail_information_to_request_data(
                 guardrail_provider=self.guardrail_provider,
                 guardrail_json_response={"error": str(e)},
-                request_data=request_data or {},
+                request_data=request_data or {},  # mutable-ok: logging helper requires a dict
                 guardrail_status="guardrail_failed_to_respond",
                 start_time=start_time.timestamp(),
                 end_time=datetime.now(timezone.utc).timestamp(),
@@ -1061,7 +1812,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         self.add_standard_logging_guardrail_information_to_request_data(
             guardrail_provider=self.guardrail_provider,
             guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response),
-            request_data=request_data or {},
+            request_data=request_data or {},  # mutable-ok: logging helper requires a dict
             guardrail_status=self._get_invoke_checks_status(bool(violations)),
             start_time=start_time.timestamp(),
             end_time=datetime.now(timezone.utc).timestamp(),

litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py

@@ -19,7 +19,7 @@ request is sent with the ``X-Cisco-AI-Defense-API-Key`` header.
 
 import json
 import os
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Mapping, Sequence
 from dataclasses import dataclass, replace
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal
@@ -94,13 +94,13 @@ class _CiscoVerdict:
     is_safe: bool | None
     classifications: list[str]
     severity: str | None
-    rules: list[dict[str, Any]]
+    rules: list[dict[str, object]]
     explanation: str | None
     event_id: str | None
     action: str | None = None
     sanitized_text: str | None = None
-    sanitized_messages: list[dict[str, Any]] | None = None
-    sanitized_mcp_arguments: dict[str, Any] | None = None
+    sanitized_messages: list[dict[str, object]] | None = None
+    sanitized_mcp_arguments: dict[str, object] | None = None
 
 
 class CiscoAIDefenseGuardrailMissingSecrets(Exception):
@@ -136,7 +136,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         api_base: str | None = None,
         inspection_type: str | None = None,
         inspect_path: str | None = None,
-        enabled_rules: list[dict[str, Any]] | None = None,
+        enabled_rules: Sequence[object] | None = None,
         integration_profile_id: str | None = None,
         integration_profile_version: str | None = None,
         integration_tenant_id: str | None = None,
@@ -415,7 +415,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
-        response: AsyncIterator[Any],
+        response: AsyncIterator[object],
         request_data: dict,
     ):
         """Buffer and inspect streaming chat output before delivery."""
@@ -437,7 +437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
             self.guardrail_name,
         )
 
-        all_chunks: Final[list[Any]] = []
+        all_chunks: Final[list[object]] = []
         try:
             async for chunk in response:
                 all_chunks.append(chunk)
@@ -497,7 +497,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
                 response_obj=assembled,
             )
         except HTTPException as exc:
-            error_obj: dict[str, Any] = self._http_exception_to_error_obj(exc)
+            error_obj: dict[str, object] = self._http_exception_to_error_obj(exc)
             verbose_proxy_logger.warning(
                 "Cisco AI Defense guardrail (%s): streaming response "
                 "blocked — emitting SSE error event instead of "
@@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
             for chunk in all_chunks:
                 yield chunk
 
-    def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, Any]:
+    def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, object]:
         """Canonical block payload used across all four block paths.
 
         Same dict is the ``HTTPException.detail`` for chat / MCP request
@@ -555,34 +555,34 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
             "event_id": verdict.event_id,
         }
 
-    def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, Any]:
+    def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, object]:
         """Wrap an ``HTTPException`` detail into the SSE ``error`` payload.
 
         For Cisco's own blocks the detail is already the canonical block
         payload, so this is a near-passthrough that just adds ``code``
         / ``guardrail`` defaults for non-Cisco / unstructured details.
         """
-        error_obj: dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)}
+        error_obj: dict[str, object] = {**exc.detail} if isinstance(exc.detail, dict) else {"message": str(exc.detail)}
         error_obj.setdefault("message", error_obj.get("error", "Guardrail block"))
         error_obj.setdefault("code", exc.status_code)
         error_obj.setdefault("guardrail", self.guardrail_name)
         return error_obj
 
     @classmethod
-    def _streaming_content_was_modified(cls, original_chunks: list[Any], assembled: ModelResponse) -> bool:
+    def _streaming_content_was_modified(cls, original_chunks: Sequence[object], assembled: ModelResponse) -> bool:
         """Decide whether redact changed content or tool/function arguments."""
         original_text: Final = cls._extract_streaming_chunk_scan_text(original_chunks)
         assembled_text: Final = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled))
         return original_text != assembled_text
 
     @classmethod
-    def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str:
+    def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str:
         original_text = ""
         argument_text = ""
         for chunk in chunks:
             choices = getattr(chunk, "choices", None) or []
             for c in choices:
-                delta = getattr(c, "delta", None)
+                delta: object | None = getattr(c, "delta", None)
                 if delta is None:
                     continue
                 text = getattr(delta, "content", None)
@@ -595,7 +595,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
                     args = cls._extract_tool_call_arguments(tc)
                     if args:
                         argument_text += args
-                fc = getattr(delta, "function_call", None)
+                fc: object | None = getattr(delta, "function_call", None)
                 if fc is not None:
                     args = cls._extract_function_call_arguments(fc)
                     if args:
@@ -673,7 +673,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         allow, WARNING for intervened/redacted, ERROR is left for
         upstream API failures.
         """
-        fields: Final[dict[str, Any]] = {
+        fields: Final[dict[str, object]] = {
             "guardrail": self.guardrail_name,
             "surface": context.surface,
             "direction": context.direction,
@@ -752,7 +752,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         user_api_key_dict: UserAPIKeyAuth,
         direction: str = "input",
         response_obj: object = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         url: Final = f"{self.api_base}{self.inspect_path}"
         payload: Final = self._build_chat_payload(messages, request_data, user_api_key_dict)
         start_time: Final = datetime.now()
@@ -784,7 +784,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         messages: list[dict[str, str]],
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         return {
             "messages": messages,
             "metadata": self._build_metadata(request_data, user_api_key_dict),
@@ -798,9 +798,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     async def _post_inspection(
         self,
         url: str,
-        payload: dict[str, Any],
+        payload: dict[str, object],
         surface: str,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         headers: Final = self._build_headers()
         verbose_proxy_logger.debug(
             "Cisco AI Defense guardrail: posting %s inspection to %s",
@@ -856,8 +856,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-    ) -> dict[str, Any]:
-        metadata: Final[dict[str, Any]] = {}
+    ) -> dict[str, object]:
+        metadata: Final[dict[str, object]] = {}
 
         user: Final = request_data.get("user") or getattr(user_api_key_dict, "user_id", None)
         if user:
@@ -884,8 +884,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
         return metadata
 
-    def _build_config(self) -> dict[str, Any]:
-        config: Final[dict[str, Any]] = {}
+    def _build_config(self) -> dict[str, object]:
+        config: Final[dict[str, object]] = {}
         if self.enabled_rules:
             config["enabled_rules"] = self.enabled_rules
         if self.integration_profile_id:
@@ -899,7 +899,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return config
 
     @staticmethod
-    def _normalize_rule(rule: object) -> dict[str, Any]:
+    def _normalize_rule(rule: object) -> dict[str, object]:
         """Coerce a user-supplied rule into the wire-shape dict Cisco expects.
 
         Accepts ``str``, ``dict``, and Pydantic model inputs.
@@ -922,7 +922,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
                     rule = dumped
 
         if isinstance(rule, dict):
-            normalized: Final[dict[str, Any]] = {}
+            normalized: Final[dict[str, object]] = {}
             rule_name: Final = rule.get("rule_name")
             if rule_name:
                 normalized["rule_name"] = rule_name
@@ -950,7 +950,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         context: _ScanContext,
         start_time: datetime,
         response_obj: object = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """Parse, log, and (optionally) raise/redact on the Cisco verdict.
 
         ``context.direction`` is ``"input"`` for request scans and ``"output"``
@@ -1119,10 +1119,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     @classmethod
     def _sanitize_response_for_logging(
         cls,
-        inspect_response: dict[str, Any],
+        inspect_response: Mapping[str, object],
         surface: str,
         action: str | None = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """Drop bulky / privacy-sensitive fields, recursing into nested dicts.
 
         MCP verdicts are commonly nested under ``result``, so a
@@ -1138,9 +1138,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return sanitized
 
     @classmethod
-    def _strip_sensitive_keys(cls, d: dict[str, Any]) -> dict[str, Any]:
+    def _strip_sensitive_keys(cls, d: Mapping[str, object]) -> dict[str, object]:
         """Recursively strip privacy-sensitive keys from a verdict dict."""
-        out: Final[dict[str, Any]] = {}
+        out: Final[dict[str, object]] = {}
         for key, value in d.items():
             if key.startswith("_") or key in cls._REDACTED_LOG_KEYS:
                 continue
@@ -1222,8 +1222,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _extract_jsonrpc_error(
-        inspect_response: dict[str, Any],
-    ) -> dict[str, Any] | None:
+        inspect_response: Mapping[str, object],
+    ) -> dict[str, object] | None:
         """Detect a JSON-RPC error envelope inside an HTTP 200 response.
 
         The Cisco Inspect API can return ``{"error": {...}}`` (or nest one
@@ -1270,7 +1270,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _extract_sanitized_text(
-        inspect_response: dict[str, Any],
+        inspect_response: Mapping[str, object],
     ) -> str | None:
         """Pull ``sanitized_text`` (or camelCase variant) off the verdict."""
         for key in ("sanitized_text", "sanitizedText"):
@@ -1287,8 +1287,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _extract_sanitized_messages(
-        inspect_response: dict[str, Any],
-    ) -> list[dict[str, Any]] | None:
+        inspect_response: Mapping[str, object],
+    ) -> list[dict[str, object]] | None:
         """Pull a sanitized OpenAI-format messages array off the verdict.
 
         Cisco can return the rewrite under several keys; we accept any of
@@ -1354,7 +1354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     def _redact_mcp_input(
         request_data: dict,
         sanitized_text: str | None,
-        sanitized_mcp_arguments: dict[str, Any] | None,
+        sanitized_mcp_arguments: dict[str, object] | None,
     ) -> bool:
         """Rewrite MCP request arguments in all locations the proxy reads."""
         if sanitized_mcp_arguments is not None:
@@ -1388,7 +1388,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         request_data: dict,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Rewrite chat request input (``messages`` or ``input``)."""
         if sanitized_messages and self._extract_tool_definition_text(request_data):
@@ -1444,7 +1444,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         cls,
         request_data: dict,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         if sanitized_messages:
             instruction_text: Final = cls._instruction_text_from_messages(sanitized_messages)
@@ -1457,7 +1457,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return False
 
     @classmethod
-    def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str | None:
+    def _instruction_text_from_messages(cls, messages: list[dict[str, object]]) -> str | None:
         for message in messages:
             if not isinstance(message, dict):
                 continue
@@ -1468,7 +1468,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return None
 
     @classmethod
-    def _non_instruction_messages(cls, messages: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None:
+    def _non_instruction_messages(cls, messages: list[dict[str, object]] | None) -> list[dict[str, object]] | None:
         if messages is None:
             return None
         return [
@@ -1499,7 +1499,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         response_obj: object,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Rewrite chat response (``ModelResponse`` or ``ResponsesAPIResponse``)."""
         if response_obj is None:
@@ -1526,7 +1526,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     def _redact_model_response_choices(
         choices: list,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Redact every returned choice, including tool-call/reasoning fields."""
         if sanitized_messages:
@@ -1570,7 +1570,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
     def _redact_text_completion_choices(
         choices: list,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         """Rewrite ``/v1/completions`` text choices after Cisco redaction."""
         replacement = sanitized_text
@@ -1638,7 +1638,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         self,
         output_items: list,
         sanitized_text: str | None,
-        sanitized_messages: list[dict[str, Any]] | None,
+        sanitized_messages: list[dict[str, object]] | None,
     ) -> bool:
         replacement_text: str | None = sanitized_text
         if not replacement_text and sanitized_messages:
@@ -1672,14 +1672,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
 
     @staticmethod
     def _sanitized_messages_to_responses_input(
-        sanitized_messages: list[dict[str, Any]],
-    ) -> list[dict[str, Any]] | None:
+        sanitized_messages: list[dict[str, object]],
+    ) -> list[dict[str, object]] | None:
         """Convert chat-shape sanitized_messages to Responses API ``input``.
 
         Returns ``None`` if nothing usable could be converted, so the
         caller falls back to ``on_flagged_action``.
         """
-        out: Final[list[dict[str, Any]]] = []
+        out: Final[list[dict[str, object]]] = []
         for m in sanitized_messages:
             if not isinstance(m, dict):
                 continue
@@ -1764,7 +1764,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         start_time: datetime | None = None,
         surface: str = "chat",
         direction: str = "input",
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         verbose_proxy_logger.error(
             "Cisco AI Defense guardrail (%s): API communication failed: %s",
             surface,
@@ -2060,7 +2060,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail):
         return getattr(obj, key, None)
 
     @classmethod
-    def _field_list(cls, obj: object, key: str) -> list[Any]:
+    def _field_list(cls, obj: object, key: str) -> list[object]:
         value: Final = cls._field(obj, key)
         return value if isinstance(value, list) else []
 

litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py

@@ -148,6 +148,27 @@ def _restore_protected_messages(
     ]
 
 
+def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, object]:
+    """Build error details for failed /v1/compress responses.
+
+    Adds troubleshooting hints for known deployment-related errors while
+    preserving the upstream status code and response body.
+    """
+    if status_code == 404:
+        return {
+            "status_code": status_code,
+            "body": body,
+            "hint": (
+                "The Headroom compression endpoint returned HTTP 404. "
+                "Verify that the configured Headroom endpoint is correct and that "
+                "the compression endpoint is available. If you are using a "
+                "self-hosted deployment, some deployments require enabling remote "
+                "compression (for example, HEADROOM_COMPRESS_ALLOW_REMOTE=1)."
+            ),
+        }
+    return {"status_code": status_code, "body": body}
+
+
 def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]:
     hashes: Final[list[str]] = []
     for msg in messages:
@@ -417,7 +438,7 @@ class HeadroomGuardrail(CustomGuardrail):
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned an error",
-                    {"status_code": e.response.status_code, "body": e.response.text},
+                    _build_compress_failure_detail(e.response.status_code, e.response.text),
                 ),
                 False,
                 {},
@@ -449,7 +470,7 @@ class HeadroomGuardrail(CustomGuardrail):
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned an error",
-                    {"status_code": response.status_code, "body": response.text},
+                    _build_compress_failure_detail(response.status_code, response.text),
                 ),
                 False,
                 {},

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -8,8 +8,8 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
 
 import copy
 import json
-from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Any, Final
+from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from fastapi import HTTPException
 
@@ -34,6 +34,9 @@ if TYPE_CHECKING:
     # Imported lazily at runtime (inside the streaming hook) to avoid a
     # module-level cyclic import with litellm.integrations.custom_guardrail.
     from litellm.integrations.custom_guardrail import ModifyResponseException
+    from litellm.llms.base_llm.guardrail_translation.base_translation import (
+        BaseTranslation,
+    )
 
 # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error
 A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
@@ -41,12 +44,35 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
 GUARDRAIL_NAME: Final = "unified_llm_guardrails"
 
 
+class _EndpointTranslation(Protocol):
+    @property
+    def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ...
+
+    @property
+    def process_output_response(self) -> "Callable[..., Awaitable[object]]": ...
+
+    @property
+    def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ...
+
+    @property
+    def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
+
+
+def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation:
+    return translation
+
+
+def _chunk_choices(item: object) -> Sequence[object]:
+    choices: Final[Sequence[object]] = getattr(item, "choices", None) or []
+    return choices
+
+
 class _StreamTerminated(Exception):
     """Internal signal that the incremental transform stream has already emitted
     its terminal chunks (block message or in-stream error) and must stop."""
 
 
-def _get_a2a_request_id(responses_so_far: list[Any], request_data: dict) -> str | None:
+def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None:
     """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting."""
     for item in responses_so_far:
         if isinstance(item, dict) and "id" in item:
@@ -138,7 +164,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         except ValueError:
             return data  # handle unmapped call types
 
-        endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        endpoint_translation: Final = _as_endpoint_translation(
+            endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        )
 
         _ensure_litellm_metadata(data, user_api_key_dict)
 
@@ -156,7 +184,7 @@ class UnifiedLLMGuardrails(CustomLogger):
 
     async def async_moderation_hook(
         self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral
-    ) -> Any:
+    ) -> object:
         """
         Runs in parallel to LLM API call
         Runs on only Input
@@ -187,7 +215,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
             return data
 
-        endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        endpoint_translation: Final = _as_endpoint_translation(
+            endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        )
 
         _ensure_litellm_metadata(data, user_api_key_dict)
 
@@ -202,7 +232,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         response,
-    ) -> Any:
+    ) -> object:
         """
         Runs on response from LLM API call
 
@@ -271,7 +301,9 @@ class UnifiedLLMGuardrails(CustomLogger):
             )
             return response
 
-        endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        endpoint_translation: Final = _as_endpoint_translation(
+            endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
+        )
 
         try:
             response = await endpoint_translation.process_output_response(
@@ -299,10 +331,10 @@ class UnifiedLLMGuardrails(CustomLogger):
     async def _handle_streaming_block(
         self,
         exc: "ModifyResponseException",
-        endpoint_translation: Any,
+        endpoint_translation: _EndpointTranslation,
         stream_started: bool,
-        responses_so_far: list[Any],
-    ) -> AsyncGenerator[Any, None]:
+        responses_so_far: Sequence[object],
+    ) -> AsyncGenerator[object, None]:
         """
         Terminate a streamed response cleanly when a guardrail blocks it.
 
@@ -323,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger):
     @staticmethod
     def _resolve_transform_call_type(
         user_api_key_dict: UserAPIKeyAuth,
-        mappings: dict,
+        mappings: Mapping[CallTypes, type["BaseTranslation"]],
     ) -> str | None:
         """Resolve the call type for the incremental_diff path, or None if the
         route is unresolvable / unsupported.
@@ -356,9 +388,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         self,
         exc: HTTPException,
         call_type: str | None,
-        responses_so_far: list[Any],
+        responses_so_far: Sequence[object],
         request_data: dict,
-    ) -> AsyncGenerator[Any, None]:
+    ) -> AsyncGenerator[object, None]:
         """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the
         response has already started, so emit an in-stream JSON-RPC error chunk;
         otherwise re-raise so the proxy can report it.
@@ -387,7 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger):
     def _build_transform_chunk(
         self,
         *,
-        reference_chunk: Any,
+        reference_chunk: object,
         mutated_text_per_choice: dict[int, str],
         emitted_text_per_choice: dict[int, str],
         holdback_per_choice: dict[int, int],
@@ -500,18 +532,18 @@ class UnifiedLLMGuardrails(CustomLogger):
     async def _emit_transform_round(
         self,
         *,
-        endpoint_translation: Any,
+        endpoint_translation: _EndpointTranslation,
         guardrail_to_apply: CustomGuardrail,
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         call_type: str,
-        reference_chunk: Any,
-        responses_so_far: list[Any],
-        responses_yielded: list[Any],
+        reference_chunk: object,
+        responses_so_far: Sequence[object],
+        responses_yielded: list[object],
         emitted_text_per_choice: dict[int, str],
         finish_reason_per_choice: dict[int, str | None],
         is_final: bool,
-    ) -> AsyncGenerator[Any, None]:
+    ) -> AsyncGenerator[object, None]:
         """Run one guardrail processing round and emit the resulting diff chunk.
 
         Raises ``_StreamTerminated`` (after emitting the terminal block message or
@@ -564,14 +596,14 @@ class UnifiedLLMGuardrails(CustomLogger):
         self,
         *,
         guardrail_to_apply: CustomGuardrail,
-        response: Any,
+        response: AsyncIterable[object],
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         call_type: str,
         sampling_rate: int,
         end_of_stream_only: bool,
-        mappings: dict,
-    ) -> AsyncGenerator[Any, None]:
+        mappings: Mapping[CallTypes, type["BaseTranslation"]],
+    ) -> AsyncGenerator[object, None]:
         """Emit guardrail text transformations as new deltas on the stream.
 
         Raw chunks are withheld and accumulated; on each sampled processing round
@@ -580,15 +612,15 @@ class UnifiedLLMGuardrails(CustomLogger):
         synthetic chunk. A BLOCK terminates the stream via the shared block
         handler; an underflow surfaces as an HTTPException.
         """
-        endpoint_translation: Final = mappings[CallTypes(call_type)]()
-        responses_so_far: Final[list[Any]] = []
-        responses_yielded: Final[list[Any]] = []
+        endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
+        responses_so_far: Final[list[object]] = []
+        responses_yielded: Final[list[object]] = []
         emitted_text_per_choice: Final[dict[int, str]] = {}
         finish_reason_per_choice: Final[dict[int, str | None]] = {}
         chunk_counter = 0
-        last_chunk: Any | None = None
+        last_chunk: object | None = None
 
-        def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]:
+        def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]:
             return self._emit_transform_round(
                 endpoint_translation=endpoint_translation,
                 guardrail_to_apply=guardrail_to_apply,
@@ -694,13 +726,13 @@ class UnifiedLLMGuardrails(CustomLogger):
     async def _inspect_full_response_for_block(
         self,
         *,
-        endpoint_translation: Any,
+        endpoint_translation: _EndpointTranslation,
         guardrail_to_apply: CustomGuardrail,
         request_data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-        responses_so_far: list[Any],
-        responses_yielded: list[Any],
-    ) -> AsyncGenerator[Any, None]:
+        responses_so_far: Sequence[object],
+        responses_yielded: Sequence[object],
+    ) -> AsyncGenerator[object, None]:
         """Run the block-only guardrail inspection over the full assembled
         response (text + tool calls) so nothing bypasses the block decision.
 
@@ -734,17 +766,17 @@ class UnifiedLLMGuardrails(CustomLogger):
             raise _StreamTerminated()
 
     @staticmethod
-    def _chunk_has_tool_calls(item: Any) -> bool:
-        for choice in getattr(item, "choices", None) or []:
+    def _chunk_has_tool_calls(item: object) -> bool:
+        for choice in _chunk_choices(item):
             delta = getattr(choice, "delta", None)
             if getattr(delta, "tool_calls", None):
                 return True
         return False
 
     @staticmethod
-    def _chunk_carries_text(item: Any) -> bool:
+    def _chunk_carries_text(item: object) -> bool:
         """True if any choice in this chunk has non-empty string ``delta.content``."""
-        for choice in getattr(item, "choices", None) or []:
+        for choice in _chunk_choices(item):
             delta = getattr(choice, "delta", None)
             content = getattr(delta, "content", None)
             if isinstance(content, str) and content != "":
@@ -753,7 +785,7 @@ class UnifiedLLMGuardrails(CustomLogger):
 
     @staticmethod
     def _tool_call_passthrough_chunk(
-        item: Any,
+        item: object,
         finish_reason_per_choice: "dict[int, str | None] | None" = None,
     ) -> ModelResponseStream:
         """Copy of a chunk carrying tool calls with all text content stripped.
@@ -772,7 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger):
         redaction purpose.
         """
         synthetic_choices: Final[list[StreamingChoices]] = []
-        for choice in getattr(item, "choices", None) or []:
+        for choice in _chunk_choices(item):
             delta = getattr(choice, "delta", None)
             idx = getattr(choice, "index", 0) or 0
             original_finish = getattr(choice, "finish_reason", None)
@@ -801,15 +833,15 @@ class UnifiedLLMGuardrails(CustomLogger):
         )
 
     @staticmethod
-    def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None:
-        for choice in getattr(item, "choices", None) or []:
+    def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None:
+        for choice in _chunk_choices(item):
             finish_reason = getattr(choice, "finish_reason", None)
             if finish_reason is not None:
                 finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason
 
     @staticmethod
-    def _chunk_has_finish_reason(item: Any) -> bool:
-        choices: Final = getattr(item, "choices", None) or []
+    def _chunk_has_finish_reason(item: object) -> bool:
+        choices: Final = _chunk_choices(item)
         return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
 
     async def async_post_call_streaming_iterator_hook(
@@ -845,22 +877,22 @@ class UnifiedLLMGuardrails(CustomLogger):
         # Get streaming configuration. Resolution order (later wins): default
         # < guardrail attribute < guardrail_config dict < this callback's
         # optional_params.
-        def _streaming_flag(name: str, default: Any) -> Any:
+        def _streaming_flag(name: str, default: object) -> Any:
             value = default
             if guardrail_to_apply is not None:
                 value = getattr(guardrail_to_apply, name, value)
-                config: Final = getattr(guardrail_to_apply, "guardrail_config", {})
+                config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {})
                 if isinstance(config, dict):
                     value = config.get(name, value)
             return self.optional_params.get(name, value)
 
-        sampling_rate: Final = _streaming_flag("streaming_sampling_rate", 5)
+        sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
         # Only apply the guardrail at end of stream (not per chunk).
-        end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False)
+        end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False)
         # "block_only" (default) drops guardrail text rewrites on the streaming
         # path; "incremental_diff" emits them as synthetic deltas (see
         # _run_incremental_transform_stream).
-        streaming_transform_mode: Final = _streaming_flag("streaming_transform_mode", "block_only")
+        streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only")
         # Withhold every chunk until end-of-stream moderation passes, then
         # release the original chunks (clean) or only the block message
         # (blocked) -- moderating the whole response *before* any content
@@ -868,7 +900,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         # release the original chunks are replayed as-is, so a
         # content-rewriting guardrail (e.g. PII masking) would leak
         # unredacted content. Guarded below via mask_response_content.
-        buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default)
+        buffer_until_moderated: bool = _streaming_flag(
+            "streaming_buffer_until_moderated", buffer_until_moderated_default
+        )
 
         if (
             buffer_until_moderated
@@ -939,9 +973,9 @@ class UnifiedLLMGuardrails(CustomLogger):
         # Infer call type from first chunk
         call_type = None
         chunk_counter = 0
-        responses_so_far: Final[list[Any]] = []
-        responses_yielded: Final[list[Any]] = []
-        pending_end_of_stream_items: Final[list[Any]] = []
+        responses_so_far: Final[list[object]] = []
+        responses_yielded: Final[list[object]] = []
+        pending_end_of_stream_items: Final[list[object]] = []
         # Whether any real response chunk has been forwarded to the client.
         # Drives how a block terminates the stream: continue the in-progress
         # message (True) vs emit a standalone block message (False, buffered).

litellm/proxy/guardrails/guardrail_initializers.py

@@ -20,6 +20,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
         content_filter_threshold=litellm_params.content_filter_threshold,
         prompt_attack_threshold=litellm_params.prompt_attack_threshold,
         pii_confidence_threshold=litellm_params.pii_confidence_threshold,
+        chunk_budget_chars=litellm_params.chunk_budget_chars,
         default_on=litellm_params.default_on,
         disable_exception_on_block=litellm_params.disable_exception_on_block,
         mask_request_content=litellm_params.mask_request_content,

litellm/proxy/hooks/litellm_skills/main.py

@@ -26,7 +26,8 @@ Usage:
 
 import base64
 import json
-from typing import Any, Final
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -36,7 +37,10 @@ from litellm.llms.litellm_proxy.skills.prompt_injection import (
     SkillPromptInjectionHandler,
 )
 from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
-from litellm.types.utils import CallTypes, CallTypesLiteral
+from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes
+
+if TYPE_CHECKING:
+    from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor
 
 
 class SkillsInjectionHook(CustomLogger):
@@ -99,7 +103,7 @@ class SkillsInjectionHook(CustomLogger):
         verbose_proxy_logger.debug("SkillsInjectionHook: Processing %s skills", len(skills))
 
         litellm_skills: Final[list[LiteLLM_SkillsTable]] = []
-        anthropic_skills: Final[list[dict[str, Any]]] = []
+        anthropic_skills: Final[list[dict[str, object]]] = []
 
         # Separate skills by prefix
         for skill in skills:
@@ -324,9 +328,9 @@ class SkillsInjectionHook(CustomLogger):
     async def async_post_call_success_deployment_hook(
         self,
         request_data: dict,
-        response: Any,
+        response: LLMResponseTypes,
         call_type: CallTypes | None,
-    ) -> Any | None:
+    ) -> LLMResponseTypes | None:
         """
         Post-call hook to handle automatic code execution.
 
@@ -372,7 +376,7 @@ class SkillsInjectionHook(CustomLogger):
         # Check if any tool call needs execution (litellm_code_execution or skill tool)
         has_executable_tool = False
         for tc in tool_calls:
-            tool_name = tc.get("name", "")
+            tool_name: str = tc.get("name", "")
             # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx)
             if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX):
                 has_executable_tool = True
@@ -441,7 +445,7 @@ class SkillsInjectionHook(CustomLogger):
         data: dict,
         response: Any,
         skill_files: dict[str, bytes],
-    ) -> Any:
+    ) -> LLMResponseTypes | None:
         """
         Execute the code execution loop for messages API (Anthropic format).
 
@@ -466,7 +470,7 @@ class SkillsInjectionHook(CustomLogger):
         max_tokens: Final = data.get("max_tokens", 4096)
 
         executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
-        generated_files: Final[list[dict[str, Any]]] = []
+        generated_files: Final[list[dict[str, object]]] = []
         current_response = response
 
         for iteration in range(self.max_iterations):
@@ -511,9 +515,9 @@ class SkillsInjectionHook(CustomLogger):
             # Process tool calls
             tool_results = []
             for tc in tool_calls:
-                tool_name = tc.get("name", "")
+                tool_name: str = tc.get("name", "")
                 tool_id = tc.get("id", "")
-                tool_input = tc.get("input", {})
+                tool_input: Mapping[str, str] = tc.get("input", {})
 
                 # Execute if it's litellm_code_execution OR a skill tool
                 if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value:
@@ -561,8 +565,8 @@ class SkillsInjectionHook(CustomLogger):
         self,
         code: str,
         skill_files: dict[str, bytes],
-        executor: Any,
-        generated_files: list[dict[str, Any]],
+        executor: "SkillsSandboxExecutor",
+        generated_files: list[dict[str, object]],
     ) -> str:
         """Execute code in sandbox and return result string."""
         try:
@@ -574,7 +578,8 @@ class SkillsInjectionHook(CustomLogger):
 
             # Collect generated files
             if exec_result.get("files"):
-                for f in exec_result["files"]:
+                files: Final[Sequence[Mapping[str, str]]] = exec_result["files"]
+                for f in files:
                     generated_files.append(
                         {
                             "name": f["name"],
@@ -595,10 +600,10 @@ class SkillsInjectionHook(CustomLogger):
     async def _execute_skill_tool(
         self,
         tool_name: str,
-        tool_input: dict[str, Any],
+        tool_input: Mapping[str, str],
         skill_files: dict[str, bytes],
-        executor: Any,
-        generated_files: list[dict[str, Any]],
+        executor: "SkillsSandboxExecutor",
+        generated_files: list[dict[str, object]],
     ) -> str:
         """Execute a skill tool by generating and running code based on skill content."""
         # Generate code based on available skill modules
@@ -670,7 +675,7 @@ print('No executable skill module found')
         data: dict,
         response: Any,
         skill_files: dict[str, bytes],
-    ) -> Any:
+    ) -> LLMResponseTypes:
         """
         Execute the code execution loop until model gives final response.
 
@@ -704,7 +709,7 @@ print('No executable skill module found')
         kwargs: Final = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS}
 
         executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
-        generated_files: Final[list[dict[str, Any]]] = []
+        generated_files: Final[list[dict[str, object]]] = []
         current_response: Any = response
 
         for iteration in range(self.max_iterations):
@@ -713,7 +718,7 @@ print('No executable skill module found')
             stop_reason = current_response.choices[0].finish_reason
 
             # Build assistant message for conversation history
-            assistant_msg_dict: dict[str, Any] = {
+            assistant_msg_dict: dict[str, object] = {
                 "role": "assistant",
                 "content": assistant_message.content,
             }
@@ -781,13 +786,13 @@ print('No executable skill module found')
         self,
         tool_call: Any,
         skill_files: dict[str, bytes],
-        executor: Any,
-        generated_files: list[dict[str, Any]],
+        executor: "SkillsSandboxExecutor",
+        generated_files: list[dict[str, object]],
     ) -> str:
         """Execute a litellm_code_execution tool call and return result string."""
         try:
             args: Final = json.loads(tool_call.function.arguments)
-            code: Final = args.get("code", "")
+            code: Final[str] = args.get("code", "")
 
             verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code))
 
@@ -802,7 +807,8 @@ print('No executable skill module found')
             # Collect generated files
             if exec_result.get("files"):
                 tool_result += "\n\nGenerated files:"
-                for f in exec_result["files"]:
+                files: Final[Sequence[Mapping[str, str]]] = exec_result["files"]
+                for f in files:
                     file_content = base64.b64decode(f["content_base64"])
                     generated_files.append(
                         {
@@ -830,8 +836,8 @@ print('No executable skill module found')
     def _attach_files_to_response(
         self,
         response: Any,
-        generated_files: list[dict[str, Any]],
-    ) -> Any:
+        generated_files: list[dict[str, object]],
+    ) -> LLMResponseTypes:
         """
         Attach generated files to the response object.
 
@@ -841,11 +847,13 @@ print('No executable skill module found')
         if not generated_files:
             return response
 
+        raw_response: Final = response
+
         # Handle dict response (Anthropic/messages API format)
         if isinstance(response, dict):
             response["_litellm_generated_files"] = generated_files
             verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to dict response", len(generated_files))
-            return response
+            return raw_response
 
         # Handle object response (OpenAI format)
         try:

litellm/proxy/hooks/max_budget_limiter.py

@@ -32,9 +32,12 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
             if max_budget is None or user_id is None:
                 return
 
-            # Personal budget applies only to non-team requests, matching
-            # the explicit team-key exemption in common_checks section 4.1.
-            if user_api_key_dict.team_id is not None:
+            from litellm.proxy.proxy_server import general_settings
+
+            if (
+                user_api_key_dict.team_id is not None
+                and general_settings.get("apply_user_budget_to_team_keys") is not True
+            ):
                 return
 
             # The reservation path admits at the strict-`<` boundary and

litellm/proxy/litellm_pre_call_utils.py

@@ -18,6 +18,7 @@ from litellm.constants import (
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
     LITELLM_PROXY_MASTER_KEY_ALIAS,
     PRE_CALL_EXECUTED_GUARDRAILS_KEY,
+    SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
 )
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
 from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
@@ -226,6 +227,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
     "applied_policies",
     "policy_sources",
     "routing_decision",
+    SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
     "standard_logging_object",
     "proxy_server_request",

litellm/proxy/management_endpoints/auto_router_endpoints.py

@@ -4,8 +4,9 @@ AUTO ROUTER MANAGEMENT ENDPOINTS
 POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config
 """
 
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from datetime import datetime, timedelta, timezone
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Annotated, Final
 
 from pydantic import BaseModel, TypeAdapter
@@ -25,6 +26,7 @@ from litellm.proxy.auth.auth_checks import (
     can_key_call_resolved_model,
 )
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL
 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
 from litellm.repositories.team_repository import TeamRepository
 from litellm.router_strategy.complexity_router import ComplexityRouter
@@ -260,6 +262,7 @@ async def preview_auto_router_routing(
 class _SessionAggRow(BaseModel):
     router_name: str
     router_type: str
+    tier_turns: Mapping[str, int]
     sessions: int
     turns: int
     unordered_turns: int
@@ -283,35 +286,6 @@ class _SessionAggRow(BaseModel):
 
 _SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow])
 
-_BENCHMARKS_SQL: Final = """
-SELECT
-    router_name,
-    router_type,
-    COUNT(*)::int AS sessions,
-    COALESCE(SUM(turns), 0)::int AS turns,
-    COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns,
-    COALESCE(SUM(covered_turns), 0)::int AS covered_turns,
-    COALESCE(SUM(cache_hits), 0)::int AS cache_hits,
-    COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns,
-    COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits,
-    COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns,
-    COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits,
-    COALESCE(SUM(return_turns), 0)::int AS return_turns,
-    COALESCE(SUM(return_hits), 0)::int AS return_hits,
-    COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses,
-    COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses,
-    COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns,
-    COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns,
-    COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
-    COALESCE(SUM(spend), 0)::float8 AS spend,
-    COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
-    COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
-FROM "LiteLLM_AutoRouterSession"
-WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
-GROUP BY router_name, router_type
-ORDER BY SUM(spend) DESC
-"""
-
 
 def _parse_benchmark_day(value: str) -> datetime:
     try:
@@ -366,6 +340,7 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
     return _SessionAggRow(
         router_name="",
         router_type="",
+        tier_turns=MappingProxyType({}),
         sessions=sum(row.sessions for row in rows),
         turns=sum(row.turns for row in rows),
         unordered_turns=sum(row.unordered_turns for row in rows),
@@ -434,7 +409,7 @@ async def get_auto_router_benchmarks(
         raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
 
     raw_rows: Final = await prisma_client.db.query_raw(
-        _BENCHMARKS_SQL,
+        AUTOROUTER_BENCHMARKS_SQL,
         start_day.isoformat(),
         (end_day + timedelta(days=1)).isoformat(),
     )
@@ -443,6 +418,7 @@ async def get_auto_router_benchmarks(
         AutoRouterBenchmarkGroup(
             router_name=row.router_name,
             router_type=row.router_type,
+            tier_turns=row.tier_turns,
             **_benchmark_totals(row).model_dump(),
         )
         for row in rows

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -17,7 +17,7 @@ import json
 import traceback
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timezone
-from typing import Any, Final, cast
+from typing import Any, Final, Literal, cast
 
 import fastapi
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -18,7 +18,7 @@ import os
 import re
 import secrets
 import traceback
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timedelta, timezone
 from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast
 
@@ -171,8 +171,12 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
 
     async def count(self, *, where: Mapping[str, object] | None = None) -> int: ...
 
+    async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ...
+
     async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ...
 
+    async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ...
+
     async def update(
         self,
         *,
@@ -181,6 +185,10 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
     ) -> _PrismaRowT | None: ...
 
 
+class _TxTables(Protocol):
+    litellm_proxymodeltable: _PrismaTableActions[object]
+
+
 def _prisma_table(
     repository: BaseRepository[_RepositoryModelT],
 ) -> _PrismaTableActions[_RepositoryModelT]:
@@ -1650,9 +1658,12 @@ async def generate_key_fn(
                 detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
             )
 
-        if user_custom_key_generate is not None:
-            if inspect.iscoroutinefunction(user_custom_key_generate):
-                result: Final = await user_custom_key_generate(data)
+        custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
+            user_custom_key_generate
+        )
+        if custom_key_generate_hook is not None:
+            if inspect.iscoroutinefunction(custom_key_generate_hook):
+                result: Final = await custom_key_generate_hook(data)
             else:
                 raise ValueError("user_custom_key_generate must be a coroutine")
             decision: Final = result.get("decision", True)
@@ -1847,9 +1858,10 @@ async def generate_service_account_key_fn(
 
     verbose_proxy_logger.debug("entered /key/generate")
 
-    if user_custom_key_generate is not None:
-        if inspect.iscoroutinefunction(user_custom_key_generate):
-            result: Final = await user_custom_key_generate(data)
+    custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate
+    if custom_key_generate_hook is not None:
+        if inspect.iscoroutinefunction(custom_key_generate_hook):
+            result: Final = await custom_key_generate_hook(data)
         else:
             raise ValueError("user_custom_key_generate must be a coroutine")
         decision: Final = result.get("decision", True)
@@ -1918,7 +1930,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
             )
         casted_metadata[reserved_field] = existing_value
 
-    data_json: Final = data.model_dump(exclude_unset=True, exclude_none=True)
+    data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True)
 
     try:
         for k, v in data_json.items():
@@ -2179,7 +2191,7 @@ async def _process_single_key_update(
     llm_router: Router | None,
     user_custom_key_update: Callable | None = None,
     existing_key_row: LiteLLM_VerificationToken | None = None,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """
     Process a single key update with all validations and checks.
 
@@ -2722,9 +2734,10 @@ async def update_key_fn(
         )
 
         # Custom key update hook
-        if user_custom_key_update is not None:
-            if inspect.iscoroutinefunction(user_custom_key_update):
-                result: Final = await user_custom_key_update(data)
+        custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update
+        if custom_key_update_hook is not None:
+            if inspect.iscoroutinefunction(custom_key_update_hook):
+                result: Final = await custom_key_update_hook(data)
             else:
                 raise ValueError("user_custom_key_update must be a coroutine")
             decision: Final = result.get("decision", True)
@@ -3528,11 +3541,34 @@ async def info_key_fn(
 ):
     """
     Retrieve information about a key.
+
     Parameters:
-        key: Optional[str] = Query parameter representing the key in the request
-        user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key
+    - key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash.
+      Defaults to the key in the Authorization header.
+
     Returns:
-        Dict containing the key and its associated information
+    - key: str - The key that was looked up, echoed back as it was passed in
+    - info: dict - The key's row, minus the hashed token
+        - key_alias: str | None - User-friendly key alias
+        - spend: float - Amount spent by the key. When budget_duration is set this covers only the
+          current budget window, not the key's lifetime
+        - max_budget: float | None - Max budget for the key, enforced against spend
+        - budget_duration: str | None - Budget reset period ("30d", "1h", etc.)
+        - budget_reset_at: datetime | None - When the current budget window ends and spend is next
+          reset to 0, not when it was last reset. Reset times snap to standard boundaries in the
+          configured timezone (30d and 1mo land on the 1st of the month, 7d on Monday, 1h on the
+          hour), so subtracting budget_duration from it does not give the window's start
+        - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
+        - model_max_budget_usage: dict | None - Current-window spend per model, present only when
+          the key has per-model budgets
+        - models: list - Model_name's the key is allowed to call
+        - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits
+        - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"}
+        - blocked: bool | None - Whether the key is blocked
+        - expires: datetime | None - When the key stops authenticating requests
+        - last_active: datetime | None - When the key was last used
+        - object_permission: dict | None - Resolved vector store / MCP permissions when the key has
+          an object_permission_id
 
     Example Curl:
     ```
@@ -4066,10 +4102,11 @@ async def delete_verification_tokens(
     failed_tokens: list = []
     try:
         if prisma_client:
-            tokens = [_hash_token_if_needed(token=key) for key in tokens]
-            _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository(
-                prisma_client
-            ).table.find_many(where={"token": {"in": tokens}})
+            hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens]
+            tokens = hashed_tokens
+            _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table(
+                VerificationTokenRepository(prisma_client)
+            ).find_many(where={"token": {"in": hashed_tokens}})
 
             if len(_keys_being_deleted) == 0:
                 raise HTTPException(
@@ -4268,7 +4305,7 @@ async def _rotate_master_key(
     if models:
         decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models)
         verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models))
-        new_models: Final = []
+        new_models: Final[list[dict[str, object]]] = []
         for model in decrypted_models:
             new_model = await _add_model_to_db(
                 model_params=Deployment(**model),
@@ -4283,7 +4320,8 @@ async def _rotate_master_key(
                 _dumped["model_info"] = prisma.Json(_dumped["model_info"])
                 new_models.append(_dumped)
         verbose_proxy_logger.debug("Resetting proxy model table")
-        async with prisma_client.db.tx() as tx:
+        async with prisma_client.db.tx() as tx_ctx:
+            tx: Final[_TxTables] = tx_ctx
             await tx.litellm_proxymodeltable.delete_many()
             verbose_proxy_logger.debug("Creating %s models", len(new_models))
             await tx.litellm_proxymodeltable.create_many(
@@ -4607,7 +4645,7 @@ async def _execute_virtual_key_regeneration(
             _validate_key_alias_format(key_alias=new_key_alias)
         verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
     update_data.update(non_default_values)
-    update_data = prisma_client.jsonify_object(data=update_data)
+    jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
 
     # If grace period set, insert deprecated key so old key remains valid
     await _insert_deprecated_key(
@@ -4619,9 +4657,9 @@ async def _execute_virtual_key_regeneration(
 
     updated_token: Final = await VerificationTokenRepository(prisma_client).table.update(
         where={"token": hashed_api_key},
-        data=update_data,
+        data=jsonified_update_data,
     )
-    updated_token_dict: Final = dict(updated_token) if updated_token is not None else {}
+    updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {}
     updated_token_dict["key"] = new_token
     updated_token_dict["token_id"] = updated_token_dict.pop("token")
 
@@ -5566,7 +5604,7 @@ async def key_aliases(
         where_sql: Final = " AND ".join(where_parts)
 
         count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}'
-        count_rows: Final = await prisma_client.db.query_raw(count_sql, *query_params)
+        count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params)
         total_count: Final = int(count_rows[0]["count"]) if count_rows else 0
 
         aliases_params: Final = query_params + [size, (page - 1) * size]
@@ -5579,7 +5617,7 @@ async def key_aliases(
             f" ORDER BY key_alias ASC"
             f" LIMIT ${limit_idx} OFFSET ${offset_idx}"
         )
-        alias_rows: Final = await prisma_client.db.query_raw(aliases_sql, *aliases_params)
+        alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params)
         aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")]
 
         total_pages: Final = -(-total_count // size) if total_count > 0 else 0
@@ -5672,7 +5710,7 @@ def _build_key_filter_conditions(
     agent_id: str | None = None,
     use_substring_matching: bool = False,
     expires_filter: str | None = None,
-) -> dict[str, str | dict[str, Any] | list[dict[str, Any]]]:
+) -> Mapping[str, object]:
     """Build filter conditions for key listing.
 
     Visibility rules:
@@ -5684,14 +5722,14 @@ def _build_key_filter_conditions(
       so former members cannot see service accounts they created after leaving.
     """
     # Prepare filter conditions
-    where: dict[str, str | dict[str, Any] | list[dict[str, Any]]] = {}
+    where: dict[str, object] = {}
     where.update(_get_condition_to_filter_out_ui_session_tokens())
 
     # Build the OR conditions for user's keys and admin team keys
-    or_conditions: Final[list[dict[str, Any]]] = []
+    or_conditions: Final[list[dict[str, object]]] = []
 
     # Base conditions for user's own keys
-    user_condition: Final[dict[str, Any]] = {}
+    user_condition: Final[dict[str, object]] = {}
     if user_id and isinstance(user_id, str):
         if use_substring_matching:
             user_condition["user_id"] = {
@@ -5761,7 +5799,7 @@ def _build_key_filter_conditions(
 
     # Apply team_id, project_id and access_group_id as global AND filters so they
     # narrow results across all visibility conditions (own keys, team keys, etc.)
-    global_filters: tuple[dict[str, Any], ...] = (
+    global_filters: Final[tuple[dict[str, object], ...]] = (
         *(
             (
                 {"key_alias": {"contains": key_alias, "mode": "insensitive"}}
@@ -5782,7 +5820,7 @@ def _build_key_filter_conditions(
             else ()
         ),
     )
-    combined_where = {"AND": [where, *global_filters]} if global_filters else where
+    combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where
     verbose_proxy_logger.debug("Filter conditions: %s", combined_where)
     return combined_where
 
@@ -5963,7 +6001,7 @@ async def _list_key_helper(
     )
 
 
-def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, Any]:
+def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]:
     """
     Condition to filter out UI session tokens
     """
@@ -6372,7 +6410,7 @@ async def _can_user_query_key_info(
 async def test_key_logging(
     user_api_key_dict: UserAPIKeyAuth,
     request: Request,
-    key_logging: list[dict[str, Any]],
+    key_logging: Sequence[Mapping[str, str]],
 ) -> LoggingCallbackStatus:
     """
     Test the key-based logging

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update.
 import asyncio
 import datetime
 import json
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Mapping, Sequence
 from json import JSONDecodeError
-from typing import Any, Final, Literal, cast
+from typing import Final, Literal, Protocol, cast
 
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
@@ -78,9 +78,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import
 from litellm.types.router import (
     SPECIAL_MODEL_INFO_PARAMS,
     Deployment,
-    DeploymentTypedDict,
     GenericLiteLLMParams,
-    LiteLLMParamsTypedDict,
     updateDeployment,
 )
 from litellm.utils import get_utc_datetime
@@ -104,10 +102,80 @@ class UpdatePublicModelGroupsRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
+class _ProxyModelRow(Protocol):
+    model_id: str
+    model_name: str
+    model_info: Mapping[str, object] | None
+
+    def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
+
+
+class _ProxyModelTable(Protocol):
+    def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
+
+    def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
+
+    def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ...
+
+    def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
+
+    def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ...
+
+
+class _TxModelTables(Protocol):
+    litellm_proxymodeltable: _ProxyModelTable
+
+
+class _TeamRow(Protocol):
+    models: Sequence[str]
+
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _TeamTable(Protocol):
+    def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ...
+
+    def update(
+        self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool]
+    ) -> Awaitable[LiteLLM_TeamTable]: ...
+
+
+class _TeamIdRef(Protocol):
+    team_id: str
+
+
+class _ModelAliasRow(Protocol):
+    id: int
+    model_aliases: dict[str, str]
+    team: _TeamIdRef | None
+
+
+class _ModelAliasTable(Protocol):
+    def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ...
+
+    def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ...
+
+
+def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable:
+    return ModelRepository(prisma_client).table
+
+
+def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable:
+    return TeamRepository(prisma_client).table
+
+
+def _db_team_table(prisma_client: PrismaClient) -> _TeamTable:
+    return prisma_client.db.litellm_teamtable
+
+
+def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable:
+    return ModelTableRepository(prisma_client).table
+
+
 async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None:
     db_model: Final = cast(
         BaseModel | None,
-        await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}),
+        await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}),
     )
 
     if not db_model:
@@ -166,14 +234,9 @@ def _raise_on_strategy_router_write_violation(
 
 
 def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel:
-    merged_deployment_dict: Final = DeploymentTypedDict(
-        model_name=db_model.model_name,
-        litellm_params=LiteLLMParamsTypedDict(**db_model.litellm_params.model_dump(exclude_none=True)),
-        model_info=db_model.model_info.model_dump(exclude_none=True),
-    )
-    # update model name
-    if updated_patch.model_name:
-        merged_deployment_dict["model_name"] = updated_patch.model_name
+    merged_model_name: Final = updated_patch.model_name or db_model.model_name
+    merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True)
+    merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True)
 
     # update litellm params
     if updated_patch.litellm_params:
@@ -182,13 +245,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
             k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items()
         }
 
-        merged_deployment_dict["litellm_params"].update(encrypted_params)
+        merged_litellm_params.update(encrypted_params)
 
     # update model info
     if updated_patch.model_info:
-        if "model_info" not in merged_deployment_dict:
-            merged_deployment_dict["model_info"] = {}
-        merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True))
+        merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True))
 
     # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI
     # passes through (which today re-sends the OLD pricing on every save) cannot
@@ -202,29 +263,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
     if updated_patch.litellm_params:
         for field in updated_patch.litellm_params.model_fields_set:
             if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None:
-                merged_deployment_dict["litellm_params"].pop(field, None)
-                merged_deployment_dict.get("model_info", {}).pop(field, None)
+                merged_litellm_params.pop(field, None)
+                merged_model_info.pop(field, None)
     if updated_patch.model_info:
         for field in updated_patch.model_info.model_fields_set:
             if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None:
-                merged_deployment_dict["model_info"].pop(field, None)
-                merged_deployment_dict.get("litellm_params", {}).pop(field, None)
+                merged_model_info.pop(field, None)
+                merged_litellm_params.pop(field, None)
 
     # convert to prisma compatible format
 
-    prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel()
-    if "model_name" in merged_deployment_dict:
-        prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"]
+    for key, value in merged_model_info.items():
+        if isinstance(value, datetime.datetime):
+            merged_model_info[key] = value.isoformat()
 
-    if "litellm_params" in merged_deployment_dict:
-        prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"])
-
-    if "model_info" in merged_deployment_dict:
-        model_info: Final = merged_deployment_dict["model_info"]
-        for key, value in model_info.items():
-            if isinstance(value, datetime.datetime):
-                model_info[key] = value.isoformat()
-        prisma_compatible_model_dict["model_info"] = json.dumps(model_info)
+    prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel(
+        model_name=merged_model_name,
+        litellm_params=json.dumps(merged_litellm_params),
+        model_info=json.dumps(merged_model_info),
+    )
 
     if updated_patch.blocked is not None:
         prisma_compatible_model_dict["blocked"] = updated_patch.blocked
@@ -338,7 +395,7 @@ async def patch_model(
         update_data["updated_at"] = cast(str, get_utc_datetime())
 
         # Perform partial update
-        updated_model: Final = await ModelRepository(prisma_client).table.update(
+        updated_model: Final = await _proxy_model_table(prisma_client).update(
             where={"model_id": model_id},
             data=update_data,
         )
@@ -769,8 +826,8 @@ async def _setup_new_team_model_assignment(
 
 
 async def _get_team_deployments(
-    team_id: str, prisma_client: PrismaClient, table: Any | None = None
-) -> list[LiteLLM_ProxyModelTable]:
+    team_id: str, prisma_client: PrismaClient, table: _ProxyModelTable | None = None
+) -> Sequence[_ProxyModelRow]:
     """
     Fetch all deployments for a given team_id from the database.
 
@@ -785,7 +842,7 @@ async def _get_team_deployments(
     existing transaction.
     """
     prefix: Final = f"model_name_{team_id}_"
-    table = table or ModelRepository(prisma_client).table
+    table = table or _proxy_model_table(prisma_client)
     response: Final = await table.find_many(
         where={
             "model_name": {"startswith": prefix},
@@ -806,7 +863,7 @@ async def _get_team_deployments(
 async def delete_team_models(
     team_ids: list[str],
     prisma_client: PrismaClient,
-    llm_router: Any | None,
+    llm_router: Router | None,
 ) -> list[str]:
     """
     Delete every BYOK model owned by the given teams, from the DB and the router.
@@ -820,7 +877,8 @@ async def delete_team_models(
     Returns the model_ids that were deleted.
     """
     deleted_model_ids: Final[list[str]] = []
-    async with prisma_client.db.tx() as tx:
+    async with prisma_client.db.tx() as tx_ctx:
+        tx: Final[_TxModelTables] = tx_ctx
         for team_id in team_ids:
             rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
             model_ids = [row.model_id for row in rows]
@@ -920,11 +978,11 @@ async def _remove_unbacked_team_models(
     if not names_to_remove:
         return
 
-    existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id})
+    existing_team_row: Final = await _db_team_table(prisma_client).find_unique(where={"team_id": team_id})
     if existing_team_row is None:
         return
 
-    updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update(
+    updated_team_row: Final[LiteLLM_TeamTable] = await _db_team_table(prisma_client).update(
         where={"team_id": team_id},
         data={"models": [model for model in existing_team_row.models if model not in names_to_remove]},
         include={"object_permission": True},
@@ -953,7 +1011,7 @@ async def _update_existing_team_model_assignment(
     """
 
     def _get_team_public_model_name(
-        model_info: dict | str | None,
+        model_info: object,
     ) -> str | None:
         parsed: Final = model_info_as_mapping(model_info)
         if parsed is None:
@@ -1062,7 +1120,7 @@ class ModelManagementAuthChecks:
                 detail={"error": CommonProxyErrors.not_premium_user.value},
             )
 
-        _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique(
+        _existing_team_row: Final = await _repo_team_table(prisma_client).find_unique(
             where={"team_id": model_params.model_info.team_id}
         )
 
@@ -1091,7 +1149,7 @@ class ModelManagementAuthChecks:
     ) -> Literal[True]:
         ## Check team model auth
         if model_params.model_info is not None and model_params.model_info.team_id is not None:
-            team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique(
+            team_obj_row: Final = await _repo_team_table(prisma_client).find_unique(
                 where={"team_id": model_params.model_info.team_id}
             )
             if team_obj_row is None:
@@ -1192,7 +1250,7 @@ async def delete_model(
             - store keys separately
             """
             # encrypt litellm params #
-            result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id})
+            result: Final = await _proxy_model_table(prisma_client).delete(where={"model_id": model_info.id})
 
             if result is None:
                 raise HTTPException(
@@ -1265,9 +1323,9 @@ async def delete_team_model_alias(
     Returns:
     - List of team id + model alias pairs that were removed
     """
-    team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True})
+    team_model_aliases: Final = await _model_alias_table(prisma_client).find_many(include={"team": True})
     tasks: Final = []
-    removed_model_aliases: Final = []
+    removed_model_aliases: Final[list[tuple[str, str]]] = []
     for team_model_alias in team_model_aliases:
         model_aliases = team_model_alias.model_aliases  # {"alias": "public model name"}
         id = team_model_alias.id
@@ -1278,7 +1336,7 @@ async def delete_team_model_alias(
                 removed_model_aliases.append((team_model_alias.team.team_id, key))
             del model_aliases[key]
             tasks.append(
-                ModelTableRepository(prisma_client).table.update(
+                _model_alias_table(prisma_client).update(
                     where={"id": id},
                     data={"model_aliases": json.dumps(model_aliases)},
                 )
@@ -1492,7 +1550,7 @@ async def update_model(
                 },
             )
 
-        _model_id = None
+        _model_id: str | None = None
         _model_info: Final = getattr(model_params, "model_info", None)
         if _model_info is None:
             raise Exception("model_info not provided")
@@ -1551,11 +1609,11 @@ async def update_model(
                 else:
                     pass
 
-            _data: Final[dict] = {
+            _data: Final[dict[str, str]] = {
                 "litellm_params": json.dumps(merged_dictionary),
                 "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
             }
-            model_response: Final = await ModelRepository(prisma_client).table.update(
+            model_response: Final = await _proxy_model_table(prisma_client).update(
                 where={"model_id": _model_id},
                 data=_data,
             )

litellm/proxy/management_endpoints/team_endpoints.py

@@ -15,11 +15,11 @@ import math
 import traceback
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timezone
-from typing import Annotated, Final, Protocol, TypeVar, cast
+from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast
 
 import fastapi
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
-from pydantic import BaseModel
+from pydantic import BaseModel, JsonValue
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
     UI_TEAM_ID,
     BlockTeamRequest,
+    BudgetNewRequest,
     CommonProxyErrors,
     DeleteTeamRequest,
     LiteLLM_AccessGroupTable,
@@ -57,6 +58,7 @@ from litellm.proxy._types import (
     SpecialManagementEndpointEnums,
     SpecialModelNames,
     SpecialProxyStrings,
+    TeamAccessGroupModelGrant,
     TeamAddMemberResponse,
     TeamInfoResponseObject,
     TeamInfoResponseObjectTeamTable,
@@ -155,6 +157,15 @@ router: Final = APIRouter()
 _DbRecordT = TypeVar("_DbRecordT")
 
 
+class _TeamIdKeyCount(TypedDict):
+    team_id: int
+
+
+class _TeamIdGroupRow(TypedDict):
+    team_id: str
+    _count: _TeamIdKeyCount
+
+
 class _PrismaTableActions(Protocol[_DbRecordT]):
     async def find_unique(
         self,
@@ -219,59 +230,127 @@ class _PrismaTableActions(Protocol[_DbRecordT]):
         where: Mapping[str, object] | None = None,
     ) -> int: ...
 
+    async def group_by(
+        self,
+        by: Sequence[str],
+        where: Mapping[str, object] | None = None,
+        count: Mapping[str, bool] | None = None,
+    ) -> Sequence[_TeamIdGroupRow]: ...
+
+
+class _HasTableActions(Protocol[_DbRecordT]):
+    @property
+    def table(self) -> "_PrismaTableActions[_DbRecordT]": ...
+
+
+def _typed_table(
+    repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT]
+) -> "_PrismaTableActions[_DbRecordT]":
+    return repo.table
+
+
+def _as_object(value: object) -> object:
+    return value
+
+
+def _nullable(value: _DbRecordT | None) -> _DbRecordT | None:
+    return value
+
+
+class _UserIdRow(Protocol):
+    @property
+    def user_id(self) -> str | None: ...
+
+
+class _HasUserIdTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_UserIdRow]": ...
+
+
+def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]":
+    return repo.table
+
+
+class _RawTeamRow(Protocol):
+    @property
+    def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ...
+
+
+class _HasRawTeamTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_RawTeamRow]": ...
+
+
+def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]":
+    return repo.table
+
+
+class _BudgetWriteCall(Protocol):
+    async def __call__(
+        self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth
+    ) -> LiteLLM_BudgetTableFull: ...
+
+
+def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall":
+    return fn
+
+
+class _TeamFindManyArgs(TypedDict, total=False):
+    take: int
+    skip: int
+    order: Mapping[str, str]
+    cursor: Mapping[str, object]
+
+
+class _TeamUiViewFilters(TypedDict, total=False):
+    team_id: Mapping[str, str]
+    team_alias: Mapping[str, str]
+
+
+class _TeamIdInFilter(TypedDict, total=False):
+    team_id: Mapping[str, Sequence[str]]
+
 
 def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
-    team_table: Final[_PrismaTableActions[LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
-    return team_table
+    return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
 
 
 def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]":
-    membership_table: Final[_PrismaTableActions[LiteLLM_TeamMembership]] = TeamMembershipRepository(prisma_client).table
-    return membership_table
+    return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership)
 
 
 def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]":
-    user_table: Final[_PrismaTableActions[LiteLLM_UserTable]] = UserRepository(prisma_client).table
-    return user_table
+    return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable)
 
 
 def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]":
-    model_table: Final[_PrismaTableActions[LiteLLM_ModelTable]] = ModelTableRepository(prisma_client).table
-    return model_table
+    return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable)
 
 
 def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]":
-    org_table: Final[_PrismaTableActions[LiteLLM_OrganizationTable]] = OrganizationRepository(prisma_client).table
-    return org_table
+    return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable)
 
 
 def _org_membership_db(
     prisma_client: PrismaClient | None,
 ) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]":
-    org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository(
-        prisma_client
-    ).table
-    return org_membership_table
+    return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable)
 
 
 def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]":
-    budget_table: Final[_PrismaTableActions[LiteLLM_BudgetTableFull]] = BudgetRepository(prisma_client).table
-    return budget_table
+    return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull)
 
 
 def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]":
-    deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table
-    return deleted_team_table
+    return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable)
 
 
 def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]":
-    access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table
-    return access_group_table
+    return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable)
 
 
 def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]":
-    tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table
-    return tokens_table
+    return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken)
 
 
 def _sanitize_for_log(value: object) -> str:
@@ -407,7 +486,7 @@ class TeamMemberBudgetHandler:
         if team_member_budget_duration is not None:
             budget_request.budget_duration = team_member_budget_duration
 
-        team_member_budget_table: Final = await new_budget(
+        team_member_budget_table: Final = await _as_budget_write(new_budget)(
             budget_obj=budget_request,
             user_api_key_dict=user_api_key_dict,
         )
@@ -455,7 +534,7 @@ class TeamMemberBudgetHandler:
             if team_member_budget_duration is not None:
                 budget_request.budget_duration = team_member_budget_duration
 
-            budget_row: Final = await update_budget(
+            budget_row: Final = await _as_budget_write(update_budget)(
                 budget_obj=budget_request,
                 user_api_key_dict=user_api_key_dict,
             )
@@ -570,7 +649,7 @@ class TeamMemberBudgetHandler:
                 )
 
         if missing:
-            await TeamMembershipRepository(prisma_client).table.create_many(
+            await _team_membership_db(prisma_client).create_many(
                 data=missing,
                 skip_duplicates=True,  # safety net against concurrent races
             )
@@ -1406,9 +1485,10 @@ async def new_team(
             complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"])
 
         complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict)
+        team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict
 
-        team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create(
-            data=complete_team_data_dict,
+        team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create(
+            data=team_creation_data,
             include={"litellm_model_table": True},
         )
 
@@ -1855,7 +1935,7 @@ async def update_team(
                 detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
             )
 
-        existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id})
+        existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
 
         if existing_team_row is None:
             raise HTTPException(
@@ -1883,7 +1963,7 @@ async def update_team(
                     )
 
         if data.max_budget is not None:
-            existing_soft_budget: Final = getattr(existing_team_row, "soft_budget", None)
+            existing_soft_budget: Final[object] = _as_object(getattr(existing_team_row, "soft_budget", None))
             soft_budget_to_check: Final = data.soft_budget if data.soft_budget is not None else existing_soft_budget
             if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)):
                 if data.max_budget <= soft_budget_to_check:
@@ -1942,7 +2022,7 @@ async def update_team(
             data.organization_id = None
 
         # check org team limits - if updating team that belongs to an org
-        org_id_to_check: Final = (
+        org_id_to_check: Final[object] = _as_object(
             data.organization_id if data.organization_id is not None else existing_team_row.organization_id
         )
         if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None:
@@ -1975,7 +2055,7 @@ async def update_team(
             TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"])
 
         if "metadata" in updated_kv:
-            stored_metadata: Final = (
+            stored_metadata: Final[Mapping[str, JsonValue] | None] = (
                 {  # mutable-ok: the validator payload's isinstance guard requires a plain dict
                     key: value
                     for key, value in existing_team_row.metadata.items()
@@ -2078,16 +2158,19 @@ async def update_team(
             updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"])
 
         updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
-        team_row: Final[LiteLLM_TeamTable | None] = await TeamRepository(prisma_client).table.update(
-            where={"team_id": data.team_id},
-            data=updated_kv,
-            # `object_permission` is included so `_refresh_cached_team`
-            # doesn't write a cached team with the relation nulled out —
-            # see team_model_add for the full rationale.
-            include={
-                "litellm_model_table": True,
-                "object_permission": True,
-            },
+        team_update_data: Final[Mapping[str, object]] = updated_kv
+        team_row: Final[LiteLLM_TeamTable | None] = _nullable(
+            await _team_db(prisma_client).update(
+                where={"team_id": data.team_id},
+                data=team_update_data,
+                # `object_permission` is included so `_refresh_cached_team`
+                # doesn't write a cached team with the relation nulled out —
+                # see team_model_add for the full rationale.
+                include={
+                    "litellm_model_table": True,
+                    "object_permission": True,
+                },
+            )
         )
 
         if team_row is None or team_row.team_id is None:
@@ -2602,7 +2685,7 @@ async def _resolve_existing_member_user_ids(
     if not requested_user_ids:
         return frozenset()
 
-    found: Final = await UserRepository(prisma_client).table.find_many(
+    found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many(
         where={  # mutable-ok: Prisma query filters are dict-shaped
             "user_id": {  # mutable-ok: Prisma query filters are dict-shaped
                 "in": sorted(requested_user_ids)
@@ -3097,7 +3180,9 @@ async def team_member_delete(
         key_val["user_id"] = data.user_id
     elif data.user_email is not None:
         key_val["user_email"] = data.user_email
-    existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val)
+    existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many(
+        where=key_val
+    )
 
     if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0):
         for existing_user in existing_user_rows:
@@ -3105,7 +3190,7 @@ async def team_member_delete(
             if data.team_id in existing_user.teams:
                 team_list = existing_user.teams
                 team_list.remove(data.team_id)
-                await UserRepository(prisma_client).table.update(
+                await _user_db(prisma_client).update(
                     where={
                         "user_id": existing_user.user_id,
                     },
@@ -3113,7 +3198,7 @@ async def team_member_delete(
                 )
 
     # Also clean up any existing team membership rows for this user and team
-    user_ids_to_delete: Final = set()
+    user_ids_to_delete: Final = set[str]()
     if data.user_id is not None:
         user_ids_to_delete.add(data.user_id)
     if existing_user_rows is not None and isinstance(existing_user_rows, list):
@@ -3122,9 +3207,7 @@ async def team_member_delete(
                 user_ids_to_delete.add(existing_user.user_id)
 
     for _uid in user_ids_to_delete:
-        await TeamMembershipRepository(prisma_client).table.delete_many(
-            where={"team_id": data.team_id, "user_id": _uid}
-        )
+        await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid})
 
     ## DELETE KEYS CREATED BY USER FOR THIS TEAM
     if user_ids_to_delete:
@@ -3133,9 +3216,7 @@ async def team_member_delete(
         )
 
         # Fetch keys before deletion to persist them
-        keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository(
-            prisma_client
-        ).table.find_many(
+        keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
             where={
                 "user_id": {"in": list(user_ids_to_delete)},
                 "team_id": data.team_id,
@@ -3150,7 +3231,7 @@ async def team_member_delete(
                 litellm_changed_by=None,
             )
 
-        await VerificationTokenRepository(prisma_client).table.delete_many(
+        await _tokens_db(prisma_client).delete_many(
             where={
                 "user_id": {"in": list(user_ids_to_delete)},
                 "team_id": data.team_id,
@@ -3310,7 +3391,7 @@ async def team_member_update(
 
     ### upsert new budget
     budget_patch: Final = _build_member_budget_patch(data)
-    async with prisma_client.db.tx() as tx:
+    async with prisma_client.tx() as tx:
         await _upsert_budget_and_membership(
             tx=tx,
             team_id=data.team_id,
@@ -3653,7 +3734,7 @@ async def delete_team(
         _persist_deleted_verification_tokens,
     )
 
-    keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
+    keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many(
         where={"team_id": {"in": data.team_ids}}
     )
 
@@ -3829,7 +3910,7 @@ async def _add_team_member_budget_table(
 ) -> TeamInfoResponseObjectTeamTable:
     try:
         team_budget: Final = await _budget_db(prisma_client).find_unique(where={"budget_id": team_member_budget_id})
-        team_info_response_object.team_member_budget_table = team_budget
+        return team_info_response_object.model_copy(update={"team_member_budget_table": team_budget})
     except Exception:
         verbose_proxy_logger.info(
             "Team member budget table not found, passed team_member_budget_id=%s", team_member_budget_id
@@ -3838,21 +3919,34 @@ async def _add_team_member_budget_table(
     return team_info_response_object
 
 
-async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None:
-    """Populate access_group_models / mcp_server_ids / agent_ids on the team
-    info response by resolving inherited resources from its access groups."""
+async def _resolve_team_access_group_resources(
+    _team_info: TeamInfoResponseObjectTeamTable,
+) -> TeamInfoResponseObjectTeamTable:
+    """Return a copy of the team info with access_group_models / mcp_server_ids /
+    agent_ids / details resolved from its access groups."""
     if not _team_info.access_group_ids:
-        return
+        return _team_info
     ag_lookup: Final = await _batch_resolve_access_group_resources(_team_info.access_group_ids)
-    models, mcp_ids, agent_ids = set(), set(), set()
-    for ag_id in _team_info.access_group_ids:
-        if ag_id in ag_lookup:
-            models.update(ag_lookup[ag_id]["models"])
-            mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"])
-            agent_ids.update(ag_lookup[ag_id]["agent_ids"])
-    _team_info.access_group_models = list(models)
-    _team_info.access_group_mcp_server_ids = list(mcp_ids)
-    _team_info.access_group_agent_ids = list(agent_ids)
+    resolved_groups: Final = tuple(
+        ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup
+    )
+    return _team_info.model_copy(
+        update={
+            "access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}),
+            "access_group_mcp_server_ids": list(
+                {s for group in resolved_groups for s in (group.access_mcp_server_ids or [])}
+            ),
+            "access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}),
+            "access_group_details": tuple(
+                TeamAccessGroupModelGrant(
+                    access_group_id=group.access_group_id,
+                    access_group_name=group.access_group_name,
+                    models=tuple(group.access_model_names or ()),
+                )
+                for group in resolved_groups
+            ),
+        }
+    )
 
 
 @router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)])
@@ -3958,11 +4052,11 @@ async def team_info(
             )
 
         # Resolve resources inherited from access groups
-        await _resolve_team_access_group_resources(_team_info)
+        resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info)
 
         response_object: Final = TeamInfoResponseObject(
             team_id=team_id,
-            team_info=_team_info,
+            team_info=resolved_team_info,
             keys=keys,
             team_memberships=returned_tm,
         )
@@ -4391,32 +4485,21 @@ async def _build_team_list_where_conditions(
 
 async def _batch_resolve_access_group_resources(
     all_access_group_ids: list[str],
-) -> dict[str, dict[str, list[str]]]:
+) -> dict[str, LiteLLM_AccessGroupTable]:
     """
-    Batch-fetch access groups in a single DB query and return a per-group
-    resource map.
-
-    Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}.
-    Missing/invalid groups are silently omitted.
+    Batch-fetch access groups in a single DB query and return them keyed by
+    access_group_id. Missing/invalid groups are silently omitted.
     """
     from litellm.proxy.proxy_server import prisma_client as _prisma_client
 
     if not all_access_group_ids or _prisma_client is None:
         return {}
 
-    unique_ids: Final = list(set(all_access_group_ids))
+    unique_ids: Final = tuple(frozenset(all_access_group_ids))
     rows: Final = await _access_group_db(_prisma_client).find_many(
         where={"access_group_id": {"in": unique_ids}},
     )
-
-    result: Final[dict[str, dict[str, list[str]]]] = {}
-    for row in rows:
-        result[row.access_group_id] = {
-            "models": list(row.access_model_names or []),
-            "mcp_server_ids": list(row.access_mcp_server_ids or []),
-            "agent_ids": list(row.access_agent_ids or []),
-        }
-    return result
+    return {row.access_group_id: row for row in rows}
 
 
 def _convert_teams_to_response_models(
@@ -4466,7 +4549,7 @@ async def _get_keys_count_by_team(
     if not page_team_ids:
         return {}
 
-    grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by(
+    grouped: Final = await _tokens_db(prisma_client).group_by(
         by=["team_id"],
         where={"team_id": {"in": page_team_ids}},
         count={"team_id": True},
@@ -4710,15 +4793,18 @@ async def list_team_v2(
             all_ag_ids: Final = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])]
             ag_lookup: Final = await _batch_resolve_access_group_resources(all_ag_ids)
             for team_item in team_items_with_ag:
-                models, mcp_ids, agent_ids = set(), set(), set()
-                for ag_id in team_item.access_group_ids or []:
-                    if ag_id in ag_lookup:
-                        models.update(ag_lookup[ag_id]["models"])
-                        mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"])
-                        agent_ids.update(ag_lookup[ag_id]["agent_ids"])
-                team_item.access_group_models = list(models)
-                team_item.access_group_mcp_server_ids = list(mcp_ids)
-                team_item.access_group_agent_ids = list(agent_ids)
+                team_groups = tuple(
+                    ag_lookup[ag_id] for ag_id in (team_item.access_group_ids or []) if ag_id in ag_lookup
+                )
+                team_item.access_group_models = list(
+                    {m for group in team_groups for m in (group.access_model_names or [])}
+                )
+                team_item.access_group_mcp_server_ids = list(
+                    {s for group in team_groups for s in (group.access_mcp_server_ids or [])}
+                )
+                team_item.access_group_agent_ids = list(
+                    {a for group in team_groups for a in (group.access_agent_ids or [])}
+                )
 
     return {
         "teams": team_list,
@@ -4780,7 +4866,7 @@ async def _authorize_and_filter_teams(
 
     if allowed_org_ids is not None:
         # Org admin: query DB for teams in their orgs
-        org_teams: Final = await TeamRepository(prisma_client).table.find_many(
+        org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
             where={"organization_id": {"in": allowed_org_ids}},
             include={"litellm_model_table": True},
         )
@@ -4794,7 +4880,9 @@ async def _authorize_and_filter_teams(
         ]
     elif user_id:
         # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays)
-        response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})
+        response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many(
+            include={"litellm_model_table": True}
+        )
         return [
             team
             for team in response
@@ -4802,7 +4890,7 @@ async def _authorize_and_filter_teams(
         ]
     else:
         # Proxy admin: all teams
-        return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}))
+        return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True}))
 
 
 @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)])
@@ -4854,7 +4942,7 @@ async def list_team(
                 _team_memberships.append(tm)
 
         # add all keys that belong to the team
-        keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id})
+        keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id})
 
         try:
             returned_responses.append(
@@ -4905,7 +4993,7 @@ async def get_paginated_teams(
         total_count: Final = await _team_db(prisma_client).count()
 
         # Get paginated teams
-        teams: Final = await TeamRepository(prisma_client).table.find_many(
+        teams: Final = await _team_db(prisma_client).find_many(
             skip=skip,
             take=page_size,
             order={"team_alias": "asc"},  # Sort by team_alias
@@ -4955,7 +5043,7 @@ async def ui_view_teams(
         skip: Final = (page - 1) * page_size
 
         # Build where conditions based on provided parameters
-        where_conditions: Final = {}
+        where_conditions: Final[_TeamUiViewFilters] = {}
 
         if team_id:
             where_conditions["team_id"] = {
@@ -4970,7 +5058,7 @@ async def ui_view_teams(
             }
 
         # Query users with pagination and filters
-        teams: Final = await TeamRepository(prisma_client).table.find_many(
+        teams: Final = await _team_db(prisma_client).find_many(
             where=where_conditions,
             skip=skip,
             take=page_size,
@@ -5160,13 +5248,13 @@ async def team_model_delete(
         )
 
     # Get current models list
-    current_models: Final = team_obj.models or []
+    current_models: Final[Sequence[str]] = team_obj.models or []
 
     # Remove specified models
     updated_models: Final = [m for m in current_models if m not in data.models]
 
     # Update team. See team_model_add for the rationale on `include`.
-    updated_team: Final = await TeamRepository(prisma_client).table.update(
+    updated_team: Final = await _team_db(prisma_client).update(
         where={"team_id": data.team_id},
         data={"models": updated_models},
         include={"object_permission": True},
@@ -5419,7 +5507,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi
     BATCH_SIZE: Final = 500
 
     while True:
-        find_args: dict = {
+        find_args: _TeamFindManyArgs = {
             "take": BATCH_SIZE,
             "order": {"team_id": "asc"},
         }
@@ -5427,7 +5515,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi
             find_args["cursor"] = {"team_id": cursor}
             find_args["skip"] = 1
 
-        teams = await TeamRepository(prisma_client).table.find_many(**find_args)
+        teams = await _team_db(prisma_client).find_many(**find_args)
 
         if not teams:
             break
@@ -5522,11 +5610,11 @@ async def get_team_daily_activity(
                     )
 
     ## Fetch team aliases and check team admin status
-    where_condition: Final = {}
+    where_condition: Final[_TeamIdInFilter] = {}
     if team_ids_list:
         where_condition["team_id"] = {"in": list(team_ids_list)}
-    team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition)
-    team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases}
+    team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition)
+    team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases}
 
     # Check if user is team admin or has /team/daily/activity permission
     # If not, filter by user's API keys.

litellm/proxy/management_endpoints/ui_sso.py

@@ -16,9 +16,22 @@ import json
 import os
 import re
 import secrets
+from collections.abc import Mapping, Sequence
 from copy import deepcopy
 from html import escape
-from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Literal,
+    NoReturn,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 from urllib.parse import parse_qs, urlencode, urlparse
 
 if TYPE_CHECKING:
@@ -155,6 +168,102 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset(
     }
 )
 
+_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
+
+
+class _PrismaTableActions(Protocol[_DbRecordT]):
+    async def find_unique(
+        self,
+        where: Mapping[str, object],
+    ) -> _DbRecordT | None: ...
+
+    async def find_first(
+        self,
+        where: Mapping[str, object] | None = None,
+    ) -> _DbRecordT | None: ...
+
+    async def find_many(
+        self,
+        where: Mapping[str, object] | None = None,
+    ) -> Sequence[_DbRecordT]: ...
+
+    async def update(
+        self,
+        where: Mapping[str, object],
+        data: Mapping[str, object],
+    ) -> _DbRecordT: ...
+
+    async def update_many(
+        self,
+        where: Mapping[str, object],
+        data: Mapping[str, object],
+    ) -> int: ...
+
+
+class _UserMetadataRow(Protocol):
+    @property
+    def metadata(self) -> Mapping[str, object] | None: ...
+
+
+class _HasUserMetadataTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ...
+
+
+def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]":
+    return repo.table
+
+
+class _SsoConfigRow(Protocol):
+    @property
+    def sso_settings(self) -> Mapping[str, object] | None: ...
+
+
+class _HasSsoConfigTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ...
+
+
+def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]":
+    return repo.table
+
+
+class _TeamDetailRow(Protocol):
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _HasTeamDetailTable(Protocol):
+    @property
+    def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ...
+
+
+def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]":
+    return repo.table
+
+
+class _CustomSsoCall(Protocol):
+    async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ...
+
+
+class _ServicePrincipalAssignment(Protocol):
+    def get(self, key: str) -> str: ...
+
+
+class _ServicePrincipalPage(Protocol):
+    @overload
+    def get(
+        self,
+        key: Literal["value"],
+        default: Sequence["_ServicePrincipalAssignment"],
+    ) -> Sequence["_ServicePrincipalAssignment"]: ...
+
+    @overload
+    def get(self, key: Literal["@odata.nextLink"]) -> str | None: ...
+
+
+def _as_object(value: object) -> object:
+    return value
+
 
 def _hash_cli_sso_secret(secret: str) -> str:
     return hashlib.sha256(secret.encode("utf-8")).hexdigest()
@@ -256,7 +365,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict:
         flow = cache.get_cache(key=cache_key)
     if isinstance(flow, str):
         try:
-            flow = json.loads(flow)
+            flow = _as_object(json.loads(flow))
         except ValueError:
             flow = None
     if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
@@ -421,7 +530,7 @@ def _flatten_cli_sso_metadata_for_poll(
 
 def build_cli_sso_attribution_metadata(
     result: CustomOpenID | OpenID | dict,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """
     Build allowlisted, non-secret scalar attribution metadata from an SSO result.
 
@@ -432,7 +541,7 @@ def build_cli_sso_attribution_metadata(
     if not claim_map:
         return {}
 
-    metadata: Final[dict[str, Any]] = {}
+    metadata: Final[dict[str, object]] = {}
     for source_claim, dest_key in claim_map:
         if not _is_safe_cli_sso_metadata_dest_key(dest_key):
             verbose_proxy_logger.debug("Skipping unsafe CLI SSO metadata destination key: %s", dest_key)
@@ -474,14 +583,14 @@ def _merge_cli_sso_attribution_metadata(
 async def _persist_cli_sso_user_metadata(
     prisma_client: PrismaClient,
     user_id: str,
-    attribution_metadata: dict[str, Any],
+    attribution_metadata: dict[str, object],
 ) -> None:
     if not attribution_metadata:
         return
 
     try:
-        user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
-        existing_metadata: dict[str, Any] = {}
+        user_row: Final = await _user_meta_db(UserRepository(prisma_client)).find_unique(where={"user_id": user_id})
+        existing_metadata: dict[str, object] = {}
         if user_row is not None:
             row_metadata: Final = user_row.metadata
             if isinstance(row_metadata, dict):
@@ -491,7 +600,7 @@ async def _persist_cli_sso_user_metadata(
             existing_metadata=existing_metadata,
             attribution_metadata=attribution_metadata,
         )
-        await UserRepository(prisma_client).table.update_many(
+        await _user_meta_db(UserRepository(prisma_client)).update_many(
             where={"user_id": user_id},
             data={"metadata": merged_metadata},
         )
@@ -1104,7 +1213,7 @@ def generic_response_convertor(
                 )
 
     # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified
-    extra_fields: dict[str, Any] | None = None
+    extra_fields: dict[str, object] | None = None
     if generic_user_extra_attributes:
         extra_fields = {}
         for attr_name in generic_user_extra_attributes.split(","):
@@ -1193,7 +1302,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]:
 
         prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy")
 
-        sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
+        sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique(
+            where={"id": "sso_config"}
+        )
 
         if sso_db_record and sso_db_record.sso_settings:
             sso_settings_dict: Final = dict(sso_db_record.sso_settings)
@@ -1225,7 +1336,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]:
 
         prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy")
 
-        sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
+        sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique(
+            where={"id": "sso_config"}
+        )
 
         if sso_db_record and sso_db_record.sso_settings:
             sso_settings_dict: Final = dict(sso_db_record.sso_settings)
@@ -1273,7 +1386,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]:
     return role_mappings
 
 
-def _parse_generic_sso_headers() -> dict:
+def _parse_generic_sso_headers() -> dict[str, str]:
     """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict."""
     raw: Final = os.getenv("GENERIC_SSO_HEADERS", None)
     if raw is None:
@@ -1677,7 +1790,7 @@ def _build_sso_user_update_data(
     result: Union["CustomOpenID", OpenID, dict] | None,
     user_email: str | None,
     user_id: str | None,
-) -> dict:
+) -> dict[str, object]:
     """
     Build the update data dictionary for SSO user upsert.
 
@@ -1689,7 +1802,7 @@ def _build_sso_user_update_data(
     Returns:
         dict: Update data containing user_email and optionally user_role if valid
     """
-    update_data: Final[dict] = {"user_email": normalize_email(user_email)}
+    update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)}
 
     # Get SSO role from result and include if valid
     sso_role: Final = getattr(result, "user_role", None)
@@ -1740,7 +1853,7 @@ async def _sync_user_role_from_jwt_role_map(
 
     # Update existing DB record if role differs
     if user_info is not None and user_info.user_role != mapped_role.value:
-        await UserRepository(prisma_client).table.update(
+        await _user_meta_db(UserRepository(prisma_client)).update(
             where={"user_id": user_info.user_id},
             data={"user_role": mapped_role.value},
         )
@@ -1796,7 +1909,7 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism
             return user_role
 
         if prisma_client:
-            await UserRepository(prisma_client).table.update(
+            await _user_meta_db(UserRepository(prisma_client)).update(
                 where={"user_id": user_id},
                 data={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
             )
@@ -2016,10 +2129,11 @@ async def _build_cli_sso_user_defined_values(
 ) -> SSOUserDefinedValues | None:
     from litellm.proxy.proxy_server import user_custom_sso
 
+    custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso
     user_id: Final = parsed_openid_result.get("user_id")
-    if user_custom_sso is not None:
-        if inspect.iscoroutinefunction(user_custom_sso):
-            return await user_custom_sso(result)
+    if custom_sso_handler is not None:
+        if inspect.iscoroutinefunction(custom_sso_handler):
+            return await custom_sso_handler(result)
         raise ValueError("user_custom_sso must be a coroutine function")
     if user_id is None:
         return None
@@ -2035,12 +2149,14 @@ async def _build_cli_sso_user_defined_values(
 
 async def _fetch_cli_sso_team_details(
     prisma_client: PrismaClient,
-    teams: list[str],
-) -> list[dict[str, Any]]:
-    team_details: Final[list[dict[str, Any]]] = []
+    teams: Sequence[str],
+) -> list[dict[str, object]]:
+    team_details: Final[list[dict[str, object]]] = []
     try:
         if teams:
-            prisma_teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}})
+            prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many(
+                where={"team_id": {"in": teams}}
+            )
             for team_row in prisma_teams:
                 team_dict = team_row.model_dump()
                 team_details.append(
@@ -2257,12 +2373,12 @@ async def cli_poll_key(
                 verbose_proxy_logger.info("Returning teams list for user %s to select from: %s", user_id, user_teams)
                 # Best-effort construction of team_details if it wasn't
                 # already cached for some reason.
-                team_details_response: list[dict[str, Any]] | None = None
+                team_details_response: list[dict[str, object]] | None = None
                 if isinstance(user_team_details, list) and user_team_details:
                     team_details_response = user_team_details
                 elif user_teams:
                     team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams]
-                poll_response: dict[str, Any] = {
+                poll_response: dict[str, object] = {
                     "status": "ready",
                     "user_id": user_id,
                     "teams": user_teams,
@@ -2997,7 +3113,9 @@ class SSOAuthenticationHandler:
                     user_id=user_id,
                 )
 
-                await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data)
+                await _user_meta_db(UserRepository(prisma_client)).update_many(
+                    where={"user_id": user_id}, data=update_data
+                )
             else:
                 verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB")
                 # user not in DB, insert User into LiteLLM DB
@@ -3089,7 +3207,9 @@ class SSOAuthenticationHandler:
                 code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             )
         try:
-            team_obj: Final = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id})
+            team_obj: Final = await _team_detail_db(TeamRepository(prisma_client)).find_first(
+                where={"team_id": litellm_team_id}
+            )
             verbose_proxy_logger.debug("Team object: %s", team_obj)
 
             # only create a new team if it doesn't exist
@@ -3278,9 +3398,10 @@ class SSOAuthenticationHandler:
         # But if it is, we want their models preferences
         user_defined_values: SSOUserDefinedValues | None = None
 
-        if user_custom_sso is not None:
-            if inspect.iscoroutinefunction(user_custom_sso):
-                user_defined_values = await user_custom_sso(result)
+        custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso
+        if custom_sso_handler is not None:
+            if inspect.iscoroutinefunction(custom_sso_handler):
+                user_defined_values = await custom_sso_handler(result)
             else:
                 raise ValueError("user_custom_sso must be a coroutine function")
         elif user_id is not None:
@@ -3448,7 +3569,7 @@ class SSOAuthenticationHandler:
             dict: Token exchange parameters
         """
         # Prepare token exchange parameters (may add code_verifier: str later)
-        token_params: Final[dict[str, Any]] = {"include_client_id": generic_include_client_id}
+        token_params: Final[dict[str, object]] = {"include_client_id": generic_include_client_id}
 
         # Retrieve PKCE code_verifier if PKCE was used in authorization.
         # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip
@@ -3663,7 +3784,7 @@ class SSOAuthenticationHandler:
         access_token string. Raises ProxyException on any validation failure.
         """
         try:
-            token_response_raw: Final = response.json()
+            token_response_raw: Final[object] = _as_object(response.json())
         except Exception as json_err:
             verbose_proxy_logger.error(
                 "Failed to parse token response as JSON: %s. Body: %s",
@@ -4253,7 +4374,7 @@ class MicrosoftSSOHandler:
 
         while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES:
             response = await async_client.get(next_link, headers=headers)
-            response_json = response.json()
+            response_json: _ServicePrincipalPage = response.json()
             verbose_proxy_logger.debug("Response from service principal app role assigned to: %s", response_json)
 
             for _object in response_json.get("value", []):

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1,9 +1,10 @@
 import base64
 import mimetypes
 import re
+from collections.abc import Mapping
 from dataclasses import dataclass, field
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Final, Literal, Optional
+from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable
 
 from litellm.repositories.table_repositories import (
     ManagedFileRepository,
@@ -16,10 +17,26 @@ if TYPE_CHECKING:
     from prisma.models import LiteLLM_ManagedObjectTable
 
     from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.utils import PrismaClient
     from litellm.router import Router
     from litellm.types.utils import LiteLLMBatch
 
 
+@runtime_checkable
+class ManagedResourceAccessChecker(Protocol):
+    async def can_user_call_unified_file_id(
+        self,
+        unified_file_id: str,
+        user_api_key_dict: "UserAPIKeyAuth",
+    ) -> bool: ...
+
+    async def can_user_call_unified_object_id(
+        self,
+        unified_object_id: str,
+        user_api_key_dict: "UserAPIKeyAuth",
+    ) -> bool: ...
+
+
 def _is_base64_encoded_unified_file_id(b64_uid: str) -> str | Literal[False]:
     # Ensure b64_uid is a string and not a mock object
     if not isinstance(b64_uid, str):
@@ -879,6 +896,65 @@ def validate_managed_files_requirement(
         )
 
 
+async def validate_managed_id_requirement(
+    resource_id: str | None,
+    resource_kind: Literal["file", "batch", "fine-tuning job"],
+    user_api_key_dict: "UserAPIKeyAuth",
+    managed_files_obj: object | None,
+) -> None:
+    """
+    Enforce proxy-level managed resources on every route that accepts a provider-issued id
+    when ``litellm.require_managed_files`` is enabled, and authenticate managed ids against
+    the caller's stored ownership record.
+
+    Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the
+    provider under shared credentials without any tenant check; knowing another tenant's provider
+    id would be enough to read, reuse, or destroy the object behind it.
+
+    Raises:
+        HTTPException: 400 for a raw id, 403 for an inaccessible managed id, or 500 when
+            ownership validation is unavailable.
+    """
+    from fastapi import HTTPException
+
+    import litellm
+
+    if litellm.require_managed_files is not True:
+        return
+
+    if not resource_id:
+        return
+
+    if not _is_base64_encoded_unified_file_id(resource_id):
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in "
+                f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the "
+                f"{resource_kind} was created."
+            ),
+        )
+
+    if not isinstance(managed_files_obj, ManagedResourceAccessChecker):
+        raise HTTPException(
+            status_code=500,
+            detail="Managed resource ownership validation is unavailable.",
+        )
+
+    can_access: Final = (
+        await managed_files_obj.can_user_call_unified_file_id(resource_id, user_api_key_dict)
+        if resource_kind == "file"
+        else await managed_files_obj.can_user_call_unified_object_id(resource_id, user_api_key_dict)
+    )
+    if can_access:
+        return
+
+    raise HTTPException(
+        status_code=403,
+        detail=f"The caller does not have access to this managed {resource_kind} id.",
+    )
+
+
 def _extract_model_param(request: "Request", request_body: dict) -> str | None:
     """
     Extract model parameter from request.
@@ -1002,6 +1078,34 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None:
             pass
 
 
+async def map_raw_file_ids_to_unified(
+    raw_file_ids: frozenset[str], prisma_client: "PrismaClient | None"
+) -> Mapping[str, str]:
+    if not raw_file_ids or not prisma_client:
+        return MappingProxyType({})
+    managed_files: Final = await ManagedFileRepository(prisma_client).table.find_many(
+        where={"flat_model_file_ids": {"hasSome": sorted(raw_file_ids)}}  # mutable-ok: prisma where is a plain dict
+    )
+    return MappingProxyType(
+        {
+            raw_id: managed_file.unified_file_id
+            for managed_file in managed_files
+            for raw_id in managed_file.flat_model_file_ids
+            if raw_id in raw_file_ids
+        }
+    )
+
+
+def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mapping[str, str]) -> None:
+    for file_attr, raw_id in (
+        ("input_file_id", getattr(response, "input_file_id", None)),
+        ("output_file_id", getattr(response, "output_file_id", None)),
+        ("error_file_id", getattr(response, "error_file_id", None)),
+    ):
+        if isinstance(raw_id, str) and raw_id in unified_id_by_raw_id:
+            setattr(response, file_attr, unified_id_by_raw_id[raw_id])
+
+
 async def ensure_batch_response_managed_file_ids(
     response,
     managed_files_obj,

litellm/proxy/openai_files_endpoints/files_endpoints.py

@@ -50,6 +50,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
     handle_model_based_routing,
     prepare_data_with_credentials,
     validate_managed_files_requirement,
+    validate_managed_id_requirement,
 )
 from litellm.proxy.utils import ProxyLogging, is_known_model
 from litellm.repositories.table_repositories import ManagedFileRepository
@@ -612,6 +613,13 @@ async def get_file_content(
 
     data: dict = {"file_id": file_id}
     try:
+        await validate_managed_id_requirement(
+            resource_id=file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         # Include original request and headers in the data
         base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
         (
@@ -908,6 +916,13 @@ async def get_file(
 
     data: dict = {"file_id": file_id}
     try:
+        await validate_managed_id_requirement(
+            resource_id=file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         custom_llm_provider: Final = (
             provider
             or get_custom_llm_provider_from_request_headers(request=request)
@@ -1098,6 +1113,13 @@ async def delete_file(
 
     data: dict = {"file_id": file_id}
     try:
+        await validate_managed_id_requirement(
+            resource_id=file_id,
+            resource_kind="file",
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        )
+
         custom_llm_provider: Final = (
             provider
             or get_custom_llm_provider_from_request_headers(request=request)

litellm/proxy/proxy_server.py

@@ -17,7 +17,7 @@ import traceback
 import warnings
 from collections.abc import AsyncGenerator, Callable, Mapping
 from datetime import datetime, timedelta, timezone
-from types import UnionType
+from types import MappingProxyType, UnionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -297,6 +297,9 @@ from litellm.proxy.common_request_processing import (
     _should_return_raw_model_name,
     create_response,
 )
+from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
+    AuthCacheInvalidationSubscriber,
+)
 from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
 from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber
 from litellm.proxy.common_utils.debug_utils import init_verbose_loggers
@@ -868,7 +871,7 @@ async def proxy_shutdown_event():
 async def _initialize_shared_aiohttp_session():
     """Initialize shared aiohttp session for connection reuse with connection limits."""
     try:
-        from aiohttp import ClientSession, TCPConnector
+        from aiohttp import ClientSession, DummyCookieJar, TCPConnector
 
         from litellm.llms.custom_httpx.http_handler import (
             _build_aiohttp_keepalive_socket_factory,
@@ -889,7 +892,7 @@ async def _initialize_shared_aiohttp_session():
             connector_kwargs["socket_factory"] = socket_factory
 
         connector: Final = TCPConnector(**connector_kwargs)
-        session: Final = ClientSession(connector=connector)
+        session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar())
 
         verbose_proxy_logger.info(
             "SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)",
@@ -1013,6 +1016,37 @@ async def proxy_startup_event(app: FastAPI):
 
         asyncio.create_task(_run_pw_migration())
 
+        async def _run_agent_grant_id_migration() -> None:
+            from litellm.proxy.agent_endpoints.agent_registry import (
+                global_agent_registry,
+                object_permission_table,
+            )
+
+            for attempt in range(3):
+                try:
+                    result = await global_agent_registry.migrate_legacy_grant_ids(
+                        table=object_permission_table(prisma_client)
+                    )
+                    if result.rewritten:
+                        verbose_proxy_logger.info(
+                            "Rewrote %s object_permission rows from legacy config agent ids", result.rewritten
+                        )
+                    if result.missed == 0:
+                        return
+                    verbose_proxy_logger.warning(
+                        "Legacy agent grant id migration attempt %s/3 left %s rows unmigrated",
+                        attempt + 1,
+                        result.missed,
+                    )
+                except Exception as e:  # noqa: BLE001  # startup task must survive any DB error and retry
+                    verbose_proxy_logger.warning(
+                        "Legacy agent grant id migration attempt %s/3 failed: %s", attempt + 1, e
+                    )
+                if attempt < 2:
+                    await asyncio.sleep(5)
+
+        asyncio.create_task(_run_agent_grant_id_migration())
+
     ## A coordination_redis block saved from the admin UI lives in the database,
     ## which is only reachable once the prisma client exists. Apply it here, before
     ## the coordination Redis is published to its consumers below.
@@ -1193,6 +1227,8 @@ async def proxy_startup_event(app: FastAPI):
 
     await proxy_config.stop_config_sync_subscriber()
 
+    await proxy_config.stop_auth_cache_invalidation_subscriber()
+
     await proxy_shutdown_event()
 
 
@@ -3904,6 +3940,7 @@ class ProxyConfig:
         self._last_hashicorp_vault_config: dict[str, Any] | None = None
         self.worker_registry: list[WorkerRegistryEntry] = []
         self.config_sync_subscriber: ConfigSyncSubscriber | None = None
+        self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None
         from litellm.litellm_core_utils.get_model_cost_map import (
             get_model_cost_map_loaded_at,
         )
@@ -6090,6 +6127,15 @@ class ProxyConfig:
             else:
                 general_settings["disable_auto_add_proxy_admin_to_teams"] = value if value is None else bool(value)
 
+        if "apply_user_budget_to_team_keys" in _general_settings and (
+            "apply_user_budget_to_team_keys" not in self._yaml_general_settings_keys
+        ):
+            db_value: Final = _general_settings["apply_user_budget_to_team_keys"]
+            if isinstance(db_value, str):
+                general_settings["apply_user_budget_to_team_keys"] = db_value.lower() == "true"
+            else:
+                general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value)
+
         ## STORE MODEL IN DB ##
         if "store_model_in_db" in _general_settings:
             value = _general_settings["store_model_in_db"]
@@ -6391,6 +6437,30 @@ class ProxyConfig:
         except Exception as e:
             verbose_proxy_logger.error("Error stopping config sync subscriber: %s", e)
 
+    def start_auth_cache_invalidation_subscriber(
+        self,
+        redis_cache: RedisCache | None,
+        user_api_key_cache: UserApiKeyCache,
+    ) -> None:
+        if redis_cache is None or self.auth_cache_invalidation_subscriber is not None:
+            return
+        subscriber: Final = AuthCacheInvalidationSubscriber(
+            redis_cache=redis_cache,
+            user_api_key_cache=user_api_key_cache,
+        )
+        self.auth_cache_invalidation_subscriber = subscriber
+        subscriber.start()
+
+    async def stop_auth_cache_invalidation_subscriber(self) -> None:
+        subscriber: Final = self.auth_cache_invalidation_subscriber
+        if subscriber is None:
+            return
+        self.auth_cache_invalidation_subscriber = None
+        try:
+            await subscriber.stop()
+        except Exception as e:  # noqa: BLE001  # best-effort: a failing stop must not break proxy shutdown
+            verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e)
+
     async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient):
         """
         Use this to read non-llm objects from the db and initialize them
@@ -8330,6 +8400,11 @@ class ProxyStartupEvent:
             misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
         )
 
+        proxy_config.start_auth_cache_invalidation_subscriber(
+            redis_cache=redis_usage_cache,
+            user_api_key_cache=user_api_key_cache,
+        )
+
         if store_model_in_db is True:
             ### GET STORED CREDENTIALS ###
             scheduler.add_job(
@@ -14947,6 +15022,29 @@ Keep it more precise, to prevent overwrite other values unintentially
 
 _PLUGIN_KEY_REDACTED: Final = "***"
 
+_GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "max_parallel_requests": "Integer",
+        "global_max_parallel_requests": "Integer",
+        "max_request_size_mb": "Integer",
+        "max_response_size_mb": "Integer",
+        "proxy_config_reload_interval_seconds": "Integer",
+        "pass_through_endpoints": "PydanticModel",
+        "store_model_in_db": "Boolean",
+        "store_prompts_in_spend_logs": "Boolean",
+        "maximum_spend_logs_retention_period": "String",
+        "mcp_internal_ip_ranges": "List",
+        "mcp_trusted_proxy_ranges": "List",
+        "mcp_xff_num_trusted_hops": "Integer",
+        "always_include_stream_usage": "Boolean",
+        "forward_client_headers_to_llm_api": "Boolean",
+        "mcp_required_fields": "List",
+        "cancel_on_disconnect": "Boolean",
+        "disable_auto_add_proxy_admin_to_teams": "Boolean",
+        "apply_user_budget_to_team_keys": "Boolean",
+    }
+)
+
 
 def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object:
     """Restore real plugin_key values the client never sees.
@@ -15445,25 +15543,7 @@ async def get_config_list(
     else:
         db_general_settings_dict = {}
 
-    allowed_args: Final = {
-        "max_parallel_requests": {"type": "Integer"},
-        "global_max_parallel_requests": {"type": "Integer"},
-        "max_request_size_mb": {"type": "Integer"},
-        "max_response_size_mb": {"type": "Integer"},
-        "proxy_config_reload_interval_seconds": {"type": "Integer"},
-        "pass_through_endpoints": {"type": "PydanticModel"},
-        "store_model_in_db": {"type": "Boolean"},
-        "store_prompts_in_spend_logs": {"type": "Boolean"},
-        "maximum_spend_logs_retention_period": {"type": "String"},
-        "mcp_internal_ip_ranges": {"type": "List"},
-        "mcp_trusted_proxy_ranges": {"type": "List"},
-        "mcp_xff_num_trusted_hops": {"type": "Integer"},
-        "always_include_stream_usage": {"type": "Boolean"},
-        "forward_client_headers_to_llm_api": {"type": "Boolean"},
-        "mcp_required_fields": {"type": "List"},
-        "cancel_on_disconnect": {"type": "Boolean"},
-        "disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"},
-    }
+    allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES
 
     return_val: Final = []
 
@@ -15471,7 +15551,7 @@ async def get_config_list(
         if field_name in allowed_args:
             ## HANDLE TYPED DICT
 
-            typed_dict_type = allowed_args[field_name]["type"]
+            typed_dict_type = allowed_args[field_name]
 
             if typed_dict_type == "PydanticModel":
                 if field_name == "pass_through_endpoints":
@@ -15513,7 +15593,7 @@ async def get_config_list(
 
                     _response_obj = ConfigList(
                         field_name=field_name,
-                        field_type=allowed_args[field_name]["type"],
+                        field_type=allowed_args[field_name],
                         field_description=field_info.description or "",
                         field_value=_redact_general_setting_value(
                             field_name,
@@ -15541,7 +15621,7 @@ async def get_config_list(
 
                 _response_obj = ConfigList(
                     field_name=field_name,
-                    field_type=allowed_args[field_name]["type"],
+                    field_type=allowed_args[field_name],
                     field_description=field_info.description or "",
                     field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin),
                     stored_in_db=_stored_in_db,

litellm/proxy/public_endpoints/public_endpoints.py

@@ -220,7 +220,7 @@ async def get_agents(request: Request):
             "url": get_custom_url(str(request.base_url), route=f"a2a/{agent.agent_id}"),
         }
         for agent in agents
-        if agent.agent_id in litellm.public_agent_groups
+        if not global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
     ]
 
 

litellm/proxy/schema.prisma

@@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession {
   total_tokens          BigInt   @default(0)
   spend                 Float    @default(0)
   saved_spend           Float    @default(0)
+  tier_turns            Json     @default("{}")
 
   @@id([api_key, session_id, router_name])
   @@index([last_turn_at], map: "idx_autorouter_session_last_turn")

litellm/proxy/spend_tracking/budget_reservation.py

@@ -156,6 +156,7 @@ async def reserve_budget_for_request(
     proxy_logging_obj: ProxyLogging,
     end_user_id: str | None = None,
     end_user_object: Any | None = None,
+    apply_user_budget_to_team_keys: bool = False,
     fail_closed_budget_enforcement: bool = False,
 ) -> dict | None:
     if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
@@ -175,6 +176,7 @@ async def reserve_budget_for_request(
         proxy_logging_obj=proxy_logging_obj,
         end_user_id=end_user_id,
         end_user_object=end_user_object,
+        apply_user_budget_to_team_keys=apply_user_budget_to_team_keys,
     )
     if not counters:
         return None
@@ -332,6 +334,7 @@ async def _get_budget_counters(
     proxy_logging_obj: ProxyLogging,
     end_user_id: str | None = None,
     end_user_object: Any | None = None,
+    apply_user_budget_to_team_keys: bool = False,
 ) -> list[_BudgetCounter]:
     counters: Final[list[_BudgetCounter]] = []
 
@@ -380,8 +383,9 @@ async def _get_budget_counters(
             )
         )
 
+    is_team_key: Final = team_object is not None and team_object.team_id is not None
     if (
-        (team_object is None or team_object.team_id is None)
+        (not is_team_key or apply_user_budget_to_team_keys)
         and user_object is not None
         and user_object.user_id is not None
         and user_object.max_budget is not None

litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py

@@ -4,11 +4,11 @@ import json
 import os
 from collections import Counter
 from collections.abc import Mapping
-from typing import Any, Final
+from typing import Any, Final, Protocol, TypeVar
 from urllib.parse import urlparse
 
 from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
-from pydantic import ConfigDict, ValidationError, create_model
+from pydantic import ConfigDict, JsonValue, ValidationError, create_model
 from pydantic.fields import FieldInfo
 
 import litellm
@@ -36,6 +36,73 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
 
 router: Final = APIRouter()
 
+_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
+
+
+class _PrismaTableActions(Protocol[_DbRecordT]):
+    async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ...
+
+    async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ...
+
+    async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ...
+
+
+class _SsoSettingsMappingRow(Protocol):
+    @property
+    def sso_settings(self) -> Mapping[str, object] | None: ...
+
+
+class _HasSsoSettingsMappingTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ...
+
+
+def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]:
+    return repo.table
+
+
+class _StoredSsoSettingsRow(Protocol):
+    @property
+    def sso_settings(self) -> object: ...
+
+
+class _HasStoredSsoSettingsTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ...
+
+
+def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]:
+    return repo.table
+
+
+class _UiSettingsRow(Protocol):
+    @property
+    def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ...
+
+
+class _HasUiSettingsTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_UiSettingsRow]: ...
+
+
+def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]:
+    return repo.table
+
+
+class _ConfigParamRow(Protocol):
+    @property
+    def param_value(self) -> str | Mapping[str, object] | None: ...
+
+
+class _HasConfigParamTable(Protocol):
+    @property
+    def table(self) -> _PrismaTableActions[_ConfigParamRow]: ...
+
+
+def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]:
+    return repo.table
+
+
 # Maps each UIThemeConfig field to the env var the UI branding path reads it
 # from. /update/ui_theme_settings writes both the stored ui_theme_config and
 # these env vars, so /get/ui_theme_settings resolves the same env vars to
@@ -54,7 +121,7 @@ def _is_public_http_url(value: str | None) -> bool:
     return parsed.scheme in ("http", "https") and bool(parsed.netloc)
 
 
-def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None:
+def _resolve_ui_theme_field(stored_values: Mapping[str, object], field_name: str) -> str | None:
     """Resolve one UI theme field to the value the branding path actually uses.
 
     The stored ui_theme_config wins; a field absent or blank there falls back to
@@ -263,7 +330,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
 # include generics like ``Optional[int]`` / ``List[str]`` that are not
 # instances of ``type`` — so tightening this to ``type`` would reject
 # valid inputs.
-_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[Any, FieldInfo]]] = {}
+_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[object, FieldInfo]]] = {}
 
 # Settings OSS knows about as enterprise-gated. If a caller sends one of
 # these keys and no extension package has registered it, the PATCH
@@ -275,7 +342,7 @@ _ENTERPRISE_ONLY_UI_SETTINGS: Final[set[str]] = {"enable_projects_ui"}
 _EFFECTIVE_UI_SETTINGS_CLASS: type[UISettings] | None = None
 
 
-def register_extra_ui_setting(name: str, annotation: Any, field: FieldInfo) -> None:
+def register_extra_ui_setting(name: str, annotation: object, field: FieldInfo) -> None:
     """Register an additional UI settings field contributed by an extension package.
 
     ``field`` must be a ``FieldInfo`` instance — construct it directly
@@ -470,7 +537,7 @@ async def delete_allowed_ip(
 
 async def _get_settings_with_schema(
     settings_key: str,
-    settings_class: Any,
+    settings_class: type[BaseModel],
     config: dict,
 ) -> dict:
     """
@@ -842,7 +909,9 @@ async def get_sso_settings():
     # Resolve the effective SSO config: the stored row wins, else the process
     # environment, else each field's default. Unlike the legacy read path this
     # does not write os.environ; a GET has no business mutating the environment.
-    sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
+    sso_db_record: Final = await _sso_settings_mapping_db(SSOConfigRepository(prisma_client)).find_unique(
+        where={"id": "sso_config"}
+    )
     sso_db_settings: Final = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None
     resolved: Final = resolve_sso_config(sso_db_settings, os.environ)
 
@@ -914,8 +983,10 @@ async def update_sso_settings(
     # before-snapshot has the same shape as after_value, and rely on
     # create_config_audit_log's secret-name redaction to mask the
     # *_client_secret fields before the audit row is written.
-    existing_sso_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"})
-    before_sso_data: dict[str, Any] | None = None
+    existing_sso_record: Final = await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).find_unique(
+        where={"id": "sso_config"}
+    )
+    before_sso_data: dict[str, JsonValue] | None = None
     if existing_sso_record and existing_sso_record.sso_settings:
         stored = existing_sso_record.sso_settings
         if isinstance(stored, str):
@@ -948,7 +1019,7 @@ async def update_sso_settings(
     encrypted_sso_data: Final = proxy_config._encrypt_env_variables(environment_variables=sso_data)
 
     # Save to dedicated SSO table
-    await SSOConfigRepository(prisma_client).table.upsert(
+    await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).upsert(
         where={"id": "sso_config"},
         data={
             "create": {
@@ -974,7 +1045,7 @@ async def update_sso_settings(
 
     # Remove SSO-related env vars from config.environment_variables
     try:
-        env_var_entry: Final = await ConfigRepository(prisma_client).table.find_unique(
+        env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique(
             where={"param_name": "environment_variables"}
         )
 
@@ -982,7 +1053,7 @@ async def update_sso_settings(
         if env_var_entry is not None:
             if env_var_entry.param_value is not None:
                 if isinstance(env_var_entry.param_value, str):
-                    environment_variables = json.loads(env_var_entry.param_value)
+                    environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value)
                 else:
                     environment_variables = dict(env_var_entry.param_value)
             else:
@@ -993,7 +1064,7 @@ async def update_sso_settings(
                 key: value for key, value in environment_variables.items() if key not in env_vars_to_remove
             }
 
-            await ConfigRepository(prisma_client).table.update(
+            await _config_param_db(ConfigRepository(prisma_client)).update(
                 where={"param_name": "environment_variables"},
                 data={
                     "param_value": json.dumps(filtered_env_vars, default=str),
@@ -1239,8 +1310,10 @@ async def get_ui_settings_cached() -> dict[str, Any]:
     if prisma_client is None:
         return {}
 
-    db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"})
-    ui_settings: dict[str, Any] = {}
+    db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
+        where={"id": "ui_settings"}
+    )
+    ui_settings: dict[str, JsonValue] = {}
     if db_record and db_record.ui_settings:
         raw: Final = db_record.ui_settings
         ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw)
@@ -1272,9 +1345,11 @@ async def get_ui_settings():
             detail={"error": "Database not connected. Please connect a database."},
         )
 
-    ui_settings: dict[str, Any] = {}
+    ui_settings: Mapping[str, JsonValue] = {}
 
-    db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"})
+    db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
+        where={"id": "ui_settings"}
+    )
 
     if db_record and db_record.ui_settings:
         ui_settings_json: Final = db_record.ui_settings
@@ -1300,7 +1375,7 @@ async def get_ui_settings():
     await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL)
 
     # Build config-like object for schema helper
-    config: Final[dict[str, Any]] = {"litellm_settings": {"ui_settings": ui_settings}}
+    config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}}
 
     return await _get_settings_with_schema(
         settings_key="ui_settings",
@@ -1315,7 +1390,7 @@ async def get_ui_settings():
     dependencies=[Depends(user_api_key_auth)],
 )
 async def update_ui_settings(
-    settings_body: dict[str, Any] = Body(...),
+    settings_body: dict[str, object] = Body(...),
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
@@ -1352,7 +1427,7 @@ async def update_ui_settings(
         raise HTTPException(status_code=422, detail=e.errors())
 
     # Only include fields the caller actually sent (not Pydantic defaults).
-    settings_dict: Final = settings.model_dump(exclude_unset=True)
+    settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True)
 
     # Reject enterprise-only settings up front so the caller gets a clear
     # signal instead of a silent drop.
@@ -1373,15 +1448,17 @@ async def update_ui_settings(
 
     # Merge with existing persisted settings so a partial PATCH doesn't
     # overwrite fields the caller didn't send.
-    existing: dict = {}
-    db_existing: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"})
+    existing: dict[str, JsonValue] = {}
+    db_existing: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
+        where={"id": "ui_settings"}
+    )
     if db_existing and db_existing.ui_settings:
         raw: Final = db_existing.ui_settings
         existing = json.loads(raw) if isinstance(raw, str) else dict(raw)
 
     ui_settings: Final = {**existing, **incoming}
 
-    await UISettingsRepository(prisma_client).table.upsert(
+    await _ui_settings_db(UISettingsRepository(prisma_client)).upsert(
         where={"id": "ui_settings"},
         data={
             "create": {

litellm/proxy/utils.py

@@ -165,6 +165,7 @@ if TYPE_CHECKING:
     from prisma.client import TransactionManager
 
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.models.team import LiteLLM_TeamTableCachedObj
     from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
     from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
 
@@ -3000,6 +3001,13 @@ class PrismaClient:
     ] = []  # mutable-ok: drained queue, mirrors tool_usage_transactions
     _autorouter_turn_transactions_lock = asyncio.Lock()
 
+    # How long a health probe failure waits for an in-flight planned engine
+    # replacement to settle before deciding whether to report itself. Generous
+    # against a replacement that takes well under a second, and far short of the
+    # reconnect budget an outage-hung `connect()` runs under, so a real outage
+    # is never waited out.
+    PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS: ClassVar[float] = 5.0
+
     def __init__(
         self,
         database_url: str,
@@ -4970,6 +4978,101 @@ class PrismaClient:
                 else:
                     verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e)
 
+    def _probe_target_wrapper(self) -> PrismaWrapper:
+        """The Prisma wrapper a `SELECT 1` health probe actually reaches.
+
+        `health_check()` issues `query_raw`, which `RoutingPrismaWrapper` sends
+        to the reader unless the reader is degraded. The writer's engine state
+        therefore says nothing about a probe that failed against the reader, so
+        the gate has to follow the same routing rule the probe did.
+        """
+        if isinstance(self.db, RoutingPrismaWrapper):
+            return self.db.writer if self.db.reader_unavailable else self.db.reader
+        return self.db
+
+    async def _run_health_probe(self, wrapper: PrismaWrapper) -> object:
+        """Issue the `SELECT 1` a health check is made of, against `wrapper`.
+
+        Takes the wrapper rather than re-reading `self.db`, because routing is
+        re-resolved on every attribute access: a reader that recovers between
+        the caller picking its target and the query going out would send the
+        probe to a different engine than the one whose generation the caller is
+        about to check, and attribute the failure to the wrong replacement.
+        """
+        sql_query: Final = "SELECT 1"
+        response: Final = await wrapper.query_raw(sql_query)
+        return response
+
+    async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool:
+        try:
+            await self._run_health_probe(wrapper)
+        except Exception as probe_error:  # noqa: BLE001  # any failure means the database is not answering
+            verbose_proxy_logger.debug("Prisma health_check() confirmation probe failed: %s", probe_error)
+            return False
+        return True
+
+    async def _planned_engine_replacement_absorbed(
+        self,
+        e: Exception,
+        wrapper: PrismaWrapper,
+        generation_before: int,
+    ) -> bool:
+        """True iff `e` is a connection-class probe failure that a completed
+        planned query-engine replacement explains.
+
+        Planned replacements (RDS IAM token refresh, guarded reconnect) kill the
+        running query engine and spawn a new one. A `SELECT 1` probe that races
+        that sub-second window fails with a transport error against the engine's
+        local HTTP port even though nothing is wrong with the database, and
+        reporting it drives a false-positive `db_exceptions` alert on every
+        replacement.
+
+        Two things must both hold, because neither is sufficient alone. The
+        engine generation must have moved, which says a replacement completed
+        rather than merely being attempted: reconnect attempts during a real
+        outage hold the same lock for tens of seconds, so gating on an in-flight
+        replacement would swallow most of an outage's alerts. And a fresh probe
+        must succeed, because `Prisma.connect()` polls the query engine's own
+        `/status` endpoint rather than round-tripping to the database, so a
+        future engine that binds before it validates its connection pool would
+        let the generation advance with the database still unreachable.
+
+        Waiting for an in-flight replacement to settle is what makes the
+        generation check meaningful, since the generation has not moved yet at
+        the instant the probe fails. The wait is generous against a replacement
+        that takes well under a second and short enough that an outage-hung
+        reconnect is not waited out; a replacement that has not settled by then
+        reports rather than stays silent.
+        """
+        if not PrismaDBExceptionHandler.is_database_connection_error(e):
+            return False
+        await wrapper.wait_for_planned_engine_replacement(self.PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS)
+        if wrapper.engine_generation == generation_before:
+            return False
+        return await self._probe_answers_now(wrapper)
+
+    async def _report_health_check_failure(
+        self,
+        e: Exception,
+        duration: float,
+        traceback_str: str,
+        wrapper: PrismaWrapper,
+        generation_before: int,
+    ) -> None:
+        if await self._planned_engine_replacement_absorbed(e, wrapper, generation_before):
+            verbose_proxy_logger.info(
+                "Prisma health_check() connection error raced a planned query-engine replacement; "
+                "not reporting it as a DB exception: %s",
+                e,
+            )
+            return
+        await self.proxy_logging_obj.failure_handler(
+            original_exception=e,
+            duration=duration,
+            call_type="health_check",
+            traceback_str=traceback_str,
+        )
+
     @backoff.on_exception(
         backoff.expo,
         Exception,
@@ -4982,13 +5085,10 @@ class PrismaClient:
         Health check endpoint for the prisma client
         """
         start_time: Final = time.time()
+        probe_wrapper: Final = self._probe_target_wrapper()
+        generation_before: Final = probe_wrapper.engine_generation
         try:
-            sql_query: Final = "SELECT 1"
-
-            # Execute the raw query
-            # The asterisk before `user_id_list` unpacks the list into separate arguments
-            response: Final = await self.db.query_raw(sql_query)
-            return response
+            return await self._run_health_probe(probe_wrapper)
         except Exception as e:
             import traceback
 
@@ -4998,11 +5098,12 @@ class PrismaClient:
             end_time: Final = time.time()
             _duration: Final = end_time - start_time
             asyncio.create_task(
-                self.proxy_logging_obj.failure_handler(
-                    original_exception=e,
+                self._report_health_check_failure(
+                    e=e,
                     duration=_duration,
-                    call_type="health_check",
                     traceback_str=error_traceback,
+                    wrapper=probe_wrapper,
+                    generation_before=generation_before,
                 )
             )
             raise e
@@ -6319,6 +6420,74 @@ def construct_database_url_from_env_vars() -> str | None:
     return None
 
 
+async def _get_validated_team_object(
+    user_api_key_dict: "UserAPIKeyAuth",
+    team_id: str,
+    prisma_client: "PrismaClient",
+    user_api_key_cache: "UserApiKeyCache",
+    proxy_logging_obj: "ProxyLogging",
+) -> "LiteLLM_TeamTableCachedObj":
+    from litellm.proxy.auth.auth_checks import get_team_object
+    from litellm.proxy.management_endpoints.team_endpoints import validate_membership
+
+    team_object: Final = await get_team_object(
+        team_id=team_id,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object)
+    return team_object
+
+
+async def _get_team_object_for_access_groups(
+    team_id: str | None,
+    prisma_client: Optional["PrismaClient"],
+    user_api_key_cache: Optional["UserApiKeyCache"],
+    proxy_logging_obj: Optional["ProxyLogging"],
+) -> Optional["LiteLLM_TeamTableCachedObj"]:
+    from litellm.proxy.auth.auth_checks import get_team_object
+
+    if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None:
+        return None
+    try:
+        return await get_team_object(
+            team_id=team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except HTTPException:
+        verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id)
+        return None
+
+
+async def _get_access_group_models(
+    user_api_key_dict: "UserAPIKeyAuth",
+    team_object: Optional["LiteLLM_TeamTableCachedObj"],
+    prisma_client: Optional["PrismaClient"],
+    user_api_key_cache: Optional["UserApiKeyCache"],
+    proxy_logging_obj: Optional["ProxyLogging"],
+) -> tuple[str, ...]:
+    from litellm.proxy.auth.auth_checks import (
+        _get_models_from_access_groups,
+        get_authorized_resources_from_key_access_groups,
+    )
+
+    team_group_models: Final = await _get_models_from_access_groups(
+        access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (),
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    key_group_models: Final = await get_authorized_resources_from_key_access_groups(
+        valid_token=user_api_key_dict,
+        team_object=team_object,
+        resource_field="access_model_names",
+    )
+    return tuple(dict.fromkeys((*team_group_models, *key_group_models)))
+
+
 async def get_available_models_for_user(
     user_api_key_dict: "UserAPIKeyAuth",
     llm_router: Optional["Router"],
@@ -6350,13 +6519,11 @@ async def get_available_models_for_user(
     Returns:
         List of model names available to the user
     """
-    from litellm.proxy.auth.auth_checks import get_team_object
     from litellm.proxy.auth.model_checks import (
         get_complete_model_list,
         get_key_models,
         get_team_models,
     )
-    from litellm.proxy.management_endpoints.team_endpoints import validate_membership
 
     # Get proxy model list and access groups
     if llm_router is None:
@@ -6366,31 +6533,33 @@ async def get_available_models_for_user(
         proxy_model_list = llm_router.get_model_names()
         model_access_groups = llm_router.get_model_access_groups()
 
-    # Get key models
-    key_models = get_key_models(
-        user_api_key_dict=user_api_key_dict,
-        proxy_model_list=proxy_model_list,
-        model_access_groups=model_access_groups,
-        include_model_access_groups=include_model_access_groups,
-    )
-
-    # Get team models
-    team_models: list[str] = user_api_key_dict.team_models
-
-    # If specific team_id is provided, validate and get team models
-    if team_id and prisma_client and proxy_logging_obj and user_api_key_cache:
-        key_models = []
-        team_object: Final = await get_team_object(
+    requested_team_object: Final = (
+        await _get_validated_team_object(
+            user_api_key_dict=user_api_key_dict,
             team_id=team_id,
             prisma_client=prisma_client,
             user_api_key_cache=user_api_key_cache,
             proxy_logging_obj=proxy_logging_obj,
         )
-        await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object)
-        team_models = team_object.models
+        if team_id and prisma_client and proxy_logging_obj and user_api_key_cache
+        else None
+    )
 
-    team_models = get_team_models(
-        team_models=team_models,
+    key_models: Final[Sequence[str]] = (
+        ()
+        if requested_team_object is not None
+        else get_key_models(
+            user_api_key_dict=user_api_key_dict,
+            proxy_model_list=proxy_model_list,
+            model_access_groups=model_access_groups,
+            include_model_access_groups=include_model_access_groups,
+        )
+    )
+
+    team_models: Final = get_team_models(
+        team_models=(
+            requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models
+        ),
         proxy_model_list=proxy_model_list,
         model_access_groups=model_access_groups,
         include_model_access_groups=include_model_access_groups,
@@ -6398,10 +6567,31 @@ async def get_available_models_for_user(
 
     effective_team_id: Final = team_id or user_api_key_dict.team_id
 
+    access_group_models: Final = (
+        await _get_access_group_models(
+            user_api_key_dict=user_api_key_dict,
+            team_object=requested_team_object
+            or await _get_team_object_for_access_groups(
+                team_id=effective_team_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            ),
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+        if key_models or team_models
+        else ()
+    )
+
+    granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models
+    granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models
+
     # Get complete model list
     all_models: Final = get_complete_model_list(
-        key_models=key_models,
-        team_models=team_models,
+        key_models=granted_key_models,
+        team_models=granted_team_models,
         proxy_model_list=proxy_model_list,
         user_model=user_model,
         infer_model_from_keys=general_settings.get("infer_model_from_keys", False),

litellm/proxy/vector_store_endpoints/management_endpoints.py

@@ -10,10 +10,17 @@ All /vector_store management endpoints
 
 import copy
 import json
-from typing import Any, Final
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from fastapi import APIRouter, Depends, HTTPException
 
+if TYPE_CHECKING:
+    from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow
+
+    from litellm.proxy.utils import PrismaClient
+    from litellm.router import Router
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.in_memory_cache import InMemoryCache
@@ -43,6 +50,25 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
 
 router: Final = APIRouter()
 
+
+class _VectorStoreTableActions(Protocol):
+    async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ...
+
+    async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ...
+
+    async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ...
+
+    async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ...
+
+
+def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions:
+    return ManagedVectorStoresRepository(prisma_client).table
+
+
+def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore:
+    return LiteLLM_ManagedVectorStore(**row.model_dump())
+
+
 _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker()
 
 
@@ -117,22 +143,20 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An
 async def _fetch_and_authorize_vector_store(
     vector_store_id: str,
     user_api_key_dict: UserAPIKeyAuth,
-    prisma_client: Any,
+    prisma_client: "PrismaClient",
 ) -> "LiteLLM_ManagedVectorStore":
     """
     Look up a vector store by id and confirm the caller can access it.
     Raises HTTPException(404) on miss and HTTPException(403) on access
     denial.
     """
-    row: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
-        where={"vector_store_id": vector_store_id}
-    )
+    row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id})
     if row is None:
         raise HTTPException(
             status_code=404,
             detail=f"Vector store with ID {vector_store_id} not found",
         )
-    typed: Final = LiteLLM_ManagedVectorStore(**row.model_dump())
+    typed: Final = _row_to_vector_store(row)
     if not await _check_vector_store_access(typed, user_api_key_dict):
         raise HTTPException(
             status_code=403,
@@ -141,7 +165,7 @@ async def _fetch_and_authorize_vector_store(
     return typed
 
 
-def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, Any] | None:
+def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None:
     """
     Resolve embedding config from router's config-defined models.
 
@@ -177,7 +201,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d
                 litellm_params = deployment.litellm_params
 
                 # Build embedding config from model params
-                embedding_config: dict[str, Any] = {}
+                embedding_config: dict[str, object] = {}
 
                 # Extract api_key
                 api_key = getattr(litellm_params, "api_key", None)
@@ -217,7 +241,9 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d
     return None
 
 
-async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> dict[str, Any] | None:
+async def _resolve_embedding_config_from_db(
+    embedding_model: str, prisma_client: "PrismaClient"
+) -> dict[str, object] | None:
     """
     Resolve embedding config from database model configuration.
 
@@ -307,7 +333,9 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client)
     return None
 
 
-async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> dict[str, Any] | None:
+async def _resolve_embedding_config(
+    embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None
+) -> dict[str, object] | None:
     """
     Resolve embedding config from either router (config-defined) or database models.
 
@@ -388,7 +416,7 @@ async def _check_vector_store_access(
 async def create_vector_store_in_db(
     vector_store_id: str,
     custom_llm_provider: str,
-    prisma_client,
+    prisma_client: "PrismaClient | None",
     vector_store_name: str | None = None,
     vector_store_description: str | None = None,
     vector_store_metadata: dict | None = None,
@@ -417,7 +445,7 @@ async def create_vector_store_in_db(
         raise HTTPException(status_code=500, detail="Database not connected")
 
     # Check if vector store already exists
-    existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
+    existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique(
         where={"vector_store_id": vector_store_id}
     )
     if existing_vector_store is not None:
@@ -427,7 +455,7 @@ async def create_vector_store_in_db(
         )
 
     # Prepare data for database
-    data_to_create: Final[dict[str, Any]] = {
+    data_to_create: Final[dict[str, object]] = {
         "vector_store_id": vector_store_id,
         "custom_llm_provider": custom_llm_provider,
     }
@@ -463,9 +491,9 @@ async def create_vector_store_in_db(
         data_to_create["litellm_params"] = safe_dumps({})
 
     # Create in database
-    _new_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create)
+    _new_vector_store: Final = await _vector_store_table(prisma_client).create(data=data_to_create)
 
-    new_vector_store: Final[LiteLLM_ManagedVectorStore] = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump())
+    new_vector_store: Final[LiteLLM_ManagedVectorStore] = _row_to_vector_store(_new_vector_store)
 
     # Add vector store to registry
     if litellm.vector_store_registry is not None:
@@ -682,12 +710,12 @@ async def delete_vector_store(
         memory_vector_store_exists = False
         vector_store_to_check = None
 
-        existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
+        existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique(
             where={"vector_store_id": data.vector_store_id}
         )
         if existing_vector_store is not None:
             db_vector_store_exists = True
-            vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump())
+            vector_store_to_check = _row_to_vector_store(existing_vector_store)
 
         # Check in-memory registry
         if litellm.vector_store_registry is not None:
@@ -715,9 +743,7 @@ async def delete_vector_store(
 
         # Delete from database if exists
         if db_vector_store_exists:
-            await ManagedVectorStoresRepository(prisma_client).table.delete(
-                where={"vector_store_id": data.vector_store_id}
-            )
+            await _vector_store_table(prisma_client).delete(where={"vector_store_id": data.vector_store_id})
 
         # Delete from in-memory registry if exists
         if memory_vector_store_exists and litellm.vector_store_registry is not None:
@@ -829,7 +855,7 @@ async def update_vector_store(
 
     try:
         update_data: Final = data.model_dump(exclude_unset=True)
-        vector_store_id: Final = update_data.pop("vector_store_id")
+        vector_store_id: Final[str] = update_data.pop("vector_store_id")
 
         # Per-store access control: anyone authenticated who passes the
         # premium-feature gate could otherwise update *any* vector store —
@@ -857,12 +883,12 @@ async def update_vector_store(
             update_data["litellm_params"] = safe_dumps(litellm_params_dict)
 
         # Update in database
-        updated: Final = await ManagedVectorStoresRepository(prisma_client).table.update(
+        updated: Final = await _vector_store_table(prisma_client).update(
             where={"vector_store_id": vector_store_id},
             data=update_data,
         )
 
-        updated_vs: Final = LiteLLM_ManagedVectorStore(**updated.model_dump())
+        updated_vs: Final = _row_to_vector_store(updated)
 
         # Immediately update in-memory registry to keep it in sync
         if litellm.vector_store_registry is not None:

litellm/proxy/vector_store_files_endpoints/endpoints.py

@@ -30,10 +30,12 @@ if TYPE_CHECKING:
 router: Final = APIRouter()
 
 
-def _update_request_data_with_managed_file_id(
+async def _update_request_data_with_managed_file_id(
     data: dict,
     file_id: str,
     request: Request,
+    user_api_key_dict: UserAPIKeyAuth,
+    managed_files_obj: object | None,
     llm_router: Optional["Router"] = None,
 ) -> tuple[dict, str | None]:
     """
@@ -65,6 +67,16 @@ def _update_request_data_with_managed_file_id(
         is_base64_encoded_unified_id,
         parse_unified_id,
     )
+    from litellm.proxy.openai_files_endpoints.common_utils import (
+        validate_managed_id_requirement,
+    )
+
+    await validate_managed_id_requirement(
+        resource_id=file_id,
+        resource_kind="file",
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=managed_files_obj,
+    )
 
     # First, check if this is a unified managed file ID (base64 encoded)
     decoded_id: Final = is_base64_encoded_unified_id(file_id)
@@ -509,8 +521,13 @@ async def vector_store_file_create(
     # Handle managed file IDs if present in request body
     original_managed_file_id = None
     if "file_id" in data:
-        data, original_managed_file_id = _update_request_data_with_managed_file_id(
-            data=data, file_id=data["file_id"], request=request, llm_router=llm_router
+        data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+            data=data,
+            file_id=data["file_id"],
+            request=request,
+            user_api_key_dict=user_api_key_dict,
+            managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+            llm_router=llm_router,
         )
 
     # Then handle managed vector store IDs
@@ -707,8 +724,13 @@ async def vector_store_file_retrieve(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs
@@ -809,8 +831,13 @@ async def vector_store_file_content(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs
@@ -911,8 +938,13 @@ async def vector_store_file_update(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs
@@ -1013,8 +1045,13 @@ async def vector_store_file_delete(
     )
 
     # Handle managed file IDs first
-    data, original_managed_file_id = _update_request_data_with_managed_file_id(
-        data=data, file_id=file_id, request=request, llm_router=llm_router
+    data, original_managed_file_id = await _update_request_data_with_managed_file_id(
+        data=data,
+        file_id=file_id,
+        request=request,
+        user_api_key_dict=user_api_key_dict,
+        managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"),
+        llm_router=llm_router,
     )
 
     # Then handle managed vector store IDs

litellm/repositories/table_repositories.py

@@ -42,6 +42,10 @@ class AgentsRepository(PrismaTableRepository):
     table_name = "litellm_agentstable"
 
 
+class ObjectPermissionRepository(PrismaTableRepository):
+    table_name = "litellm_objectpermissiontable"
+
+
 class GuardrailsRepository(PrismaTableRepository):
     table_name = "litellm_guardrailstable"
 

litellm/responses/litellm_completion_transformation/transformation.py

@@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion  (Chat Completion
 
 import json
 import re
-from collections.abc import Sequence
-from typing import Any, Final, Literal, cast
+from collections.abc import Iterator, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable
 
 from openai.types.chat.chat_completion_named_tool_choice_param import (
     ChatCompletionNamedToolChoiceParam,
@@ -16,6 +16,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import (
 from openai.types.responses import ResponseFunctionToolCall
 from openai.types.responses.response_create_params import ResponseInputParam
 from openai.types.responses.tool_param import FunctionToolParam
+from pydantic import TypeAdapter
 from typing_extensions import TypedDict
 
 from litellm._logging import verbose_logger
@@ -78,9 +79,35 @@ from .custom_tools import (
     unwrap_custom_tool_arguments,
 )
 
+if TYPE_CHECKING:
+    from openai.types.responses.response_apply_patch_tool_call import (
+        ResponseApplyPatchToolCall,
+    )
+
 ########### Initialize Classes used for Responses API  ###########
 TOOL_CALLS_CACHE: Final = InMemoryCache()
 
+_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
+_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
+_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
+_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
+_TEXT_ADAPTER: Final = TypeAdapter(str)
+
+
+@runtime_checkable
+class _SupportsIter(Protocol):
+    def __iter__(self) -> Iterator[object]: ...
+
+
+@runtime_checkable
+class _HasToolCalls(Protocol):
+    tool_calls: object
+
+
+@runtime_checkable
+class _HasId(Protocol):
+    id: object
+
 
 class ChatCompletionSession(TypedDict, total=False):
     messages: list[
@@ -205,7 +232,7 @@ class LiteLLMCompletionResponsesConfig:
         responses_api_request: ResponsesAPIOptionalRequestParams,
         custom_llm_provider: str | None = None,
         stream: bool | None = None,
-        extra_headers: dict[str, Any] | None = None,
+        extra_headers: Mapping[str, object] | None = None,
         **kwargs,
     ) -> dict:
         """
@@ -462,7 +489,9 @@ class LiteLLMCompletionResponsesConfig:
                     if not chat_completion_messages:
                         continue
 
-                    deduped_in_place: list[Any] = []
+                    deduped_in_place: list[
+                        AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage
+                    ] = []
                     for m in chat_completion_messages:
                         role = ""
                         if isinstance(m, dict):
@@ -472,7 +501,7 @@ class LiteLLMCompletionResponsesConfig:
 
                         # Drop assistant tool_calls wrappers if we already have this call_id
                         if role == "assistant":
-                            tool_calls: Any = (
+                            tool_calls: object = (
                                 m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None)
                             )
                             call_id = ""
@@ -534,7 +563,7 @@ class LiteLLMCompletionResponsesConfig:
             call_id = ""
 
             if role == "assistant":
-                tool_calls: Any = None
+                tool_calls: object = None
                 if isinstance(tool_call_message, dict):
                     tool_calls = tool_call_message.get("tool_calls")
                 else:
@@ -578,7 +607,16 @@ class LiteLLMCompletionResponsesConfig:
         return False
 
     @staticmethod
-    def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None:
+    def _find_previous_assistant_idx(
+        messages: Sequence[
+            AllMessageValues
+            | GenericChatCompletionMessage
+            | ChatCompletionResponseMessage
+            | ChatCompletionMessageToolCall
+            | Message
+        ],
+        current_idx: int,
+    ) -> int | None:
         """Find the index of the previous assistant message."""
         for j in range(current_idx - 1, -1, -1):
             if messages[j].get("role") == "assistant":
@@ -586,7 +624,18 @@ class LiteLLMCompletionResponsesConfig:
         return None
 
     @staticmethod
-    def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str:
+    def _recover_tool_call_id_from_assistant(
+        assistant_message: AllMessageValues
+        | GenericChatCompletionMessage
+        | ChatCompletionResponseMessage
+        | ChatCompletionMessageToolCall
+        | Message,
+        message: AllMessageValues
+        | GenericChatCompletionMessage
+        | ChatCompletionResponseMessage
+        | ChatCompletionMessageToolCall
+        | Message,
+    ) -> str:
         """Try to recover empty tool_call_id from assistant message's tool_calls."""
         tool_calls_raw: Final = (
             assistant_message.get("tool_calls")
@@ -594,17 +643,23 @@ class LiteLLMCompletionResponsesConfig:
             else getattr(assistant_message, "tool_calls", None)
         )
         if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0:
-            first_tool_call: Final = tool_calls_raw[0]
+            first_tool_call: Final = _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)[0]
             if isinstance(first_tool_call, dict):
-                tool_call_id_raw = first_tool_call.get("id", "")
+                tool_call_id_raw = _ANY_KEY_DICT_ADAPTER.validate_python(first_tool_call).get("id", "")
                 return str(tool_call_id_raw) if tool_call_id_raw is not None else ""
-            elif hasattr(first_tool_call, "id"):
-                tool_call_id_raw = getattr(first_tool_call, "id", None)
+            elif isinstance(first_tool_call, _HasId):
+                tool_call_id_raw = first_tool_call.id
                 return str(tool_call_id_raw) if tool_call_id_raw is not None else ""
         return ""
 
     @staticmethod
-    def _get_tool_calls_list(assistant_message: Any) -> list[Any]:
+    def _get_tool_calls_list(
+        assistant_message: AllMessageValues
+        | GenericChatCompletionMessage
+        | ChatCompletionResponseMessage
+        | ChatCompletionMessageToolCall
+        | Message,
+    ) -> Sequence[object]:
         """Extract tool_calls as a list from assistant message."""
         tool_calls_raw: Final = (
             assistant_message.get("tool_calls")
@@ -614,18 +669,18 @@ class LiteLLMCompletionResponsesConfig:
         if tool_calls_raw is None:
             return []
         if isinstance(tool_calls_raw, list):
-            return tool_calls_raw
-        if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)):
+            return _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)
+        if isinstance(tool_calls_raw, _SupportsIter) and not isinstance(tool_calls_raw, (str, bytes)):
             return list(tool_calls_raw)
         return []
 
     @staticmethod
-    def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool:
+    def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool:
         """Check if a tool_call with the given ID exists in the list."""
         for tool_call in tool_calls:
-            tool_call_id_to_check: str | None = None
+            tool_call_id_to_check: object = None
             if isinstance(tool_call, dict):
-                tool_call_id_to_check = tool_call.get("id")
+                tool_call_id_to_check = _ANY_KEY_DICT_ADAPTER.validate_python(tool_call).get("id")
             elif hasattr(tool_call, "id"):
                 tool_call_id_to_check = getattr(tool_call, "id", None)
             if tool_call_id_to_check == tool_call_id:
@@ -633,12 +688,13 @@ class LiteLLMCompletionResponsesConfig:
         return False
 
     @staticmethod
-    def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None:
+    def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None:
         """Reconstruct a minimal tool_call definition from tools list."""
         for tool in tools:
             if isinstance(tool, dict):
-                tool_function = tool.get("function") or {}
-                tool_name = tool_function.get("name") or tool.get("name") or ""
+                tool_map = _ANY_KEY_DICT_ADAPTER.validate_python(tool)
+                tool_function = _ANY_KEY_DICT_ADAPTER.validate_python(tool_map.get("function") or {})
+                tool_name = tool_function.get("name") or tool_map.get("name") or ""
                 if tool_name:
                     return {
                         "id": tool_call_id,
@@ -651,7 +707,7 @@ class LiteLLMCompletionResponsesConfig:
         return None
 
     @staticmethod
-    def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any:
+    def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object:
         """
         Safely read a field from dict-like or attribute-based objects.
         """
@@ -659,7 +715,7 @@ class LiteLLMCompletionResponsesConfig:
             return default
 
         if isinstance(obj, dict):
-            return obj.get(key, default)
+            return _ANY_KEY_DICT_ADAPTER.validate_python(obj).get(key, default)
 
         getter: Final = getattr(obj, "get", None)
         if callable(getter):
@@ -672,13 +728,13 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _create_tool_call_chunk(
-        tool_use_definition: dict[str, Any], tool_call_id: str, index: int
+        tool_use_definition: Mapping[object, object], tool_call_id: str, index: int
     ) -> ChatCompletionToolCallChunk:
         """Create a ChatCompletionToolCallChunk from tool_use_definition."""
         function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function")
         function_name_raw: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name")
         function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments")
-        function: Final[dict[str, Any]] = {
+        function: Final[dict[str, object]] = {
             "name": function_name_raw or "",
             "arguments": function_arguments_raw or "{}",
         }
@@ -697,7 +753,7 @@ class LiteLLMCompletionResponsesConfig:
         )
 
     @staticmethod
-    def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None:
+    def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[object, object] | None:
         """
         Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk.
         """
@@ -705,7 +761,7 @@ class LiteLLMCompletionResponsesConfig:
             return None
 
         if isinstance(tool_use_definition, dict):
-            normalized_definition: dict[str, Any] = dict(tool_use_definition)
+            normalized_definition: dict[object, object] = _ANY_KEY_DICT_ADAPTER.validate_python(tool_use_definition)
         else:
             tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id")
             tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type")
@@ -738,7 +794,7 @@ class LiteLLMCompletionResponsesConfig:
         return normalized_definition
 
     @staticmethod
-    def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
+    def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None:
         """Add a tool_call to an assistant message."""
         if isinstance(assistant_message, dict):
             prev_assistant_dict: Final = cast(dict[str, Any], assistant_message)
@@ -747,7 +803,7 @@ class LiteLLMCompletionResponsesConfig:
             tool_calls_list: Final = prev_assistant_dict["tool_calls"]
             if isinstance(tool_calls_list, list):
                 tool_calls_list.append(tool_call_chunk)
-        elif hasattr(assistant_message, "tool_calls"):
+        elif isinstance(assistant_message, _HasToolCalls):
             if assistant_message.tool_calls is None:
                 assistant_message.tool_calls = []
             if isinstance(assistant_message.tool_calls, list):
@@ -762,7 +818,7 @@ class LiteLLMCompletionResponsesConfig:
             | ChatCompletionMessageToolCall
             | Message
         ],
-        tools: list[Any] | None = None,
+        tools: Sequence[object] | None = None,
     ) -> list[
         AllMessageValues
         | GenericChatCompletionMessage
@@ -851,7 +907,7 @@ class LiteLLMCompletionResponsesConfig:
                 tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant)
 
                 if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id):
-                    _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id)
+                    _tool_use_definition: object = TOOL_CALLS_CACHE.get_cache(key=tool_call_id)
 
                     if not _tool_use_definition and tools:
                         _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools(
@@ -908,7 +964,7 @@ class LiteLLMCompletionResponsesConfig:
                 function_call=input_item
             )
         else:
-            content: Final = input_item.get("content")
+            content: Final[object] = input_item.get("content")
             # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content
             # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages
             if content is None:
@@ -923,7 +979,7 @@ class LiteLLMCompletionResponsesConfig:
             ]
 
     @staticmethod
-    def _is_input_item_tool_call_output(input_item: Any) -> bool:
+    def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool:
         """
         Check if the input item is a tool call output
         """
@@ -936,7 +992,7 @@ class LiteLLMCompletionResponsesConfig:
         ]
 
     @staticmethod
-    def _is_input_item_function_call(input_item: Any) -> bool:
+    def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool:
         """
         Check if the input item is a function call or custom tool call.
         Both need to be reconstructed as assistant tool_calls for Chat
@@ -946,7 +1002,7 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_responses_api_tool_call_output_to_chat_completion_message(
-        tool_call_output: dict[str, Any],
+        tool_call_output: Mapping[str, object],
     ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
         """
         ChatCompletionToolMessage is used to indicate the output from a tool call
@@ -958,7 +1014,7 @@ class LiteLLMCompletionResponsesConfig:
             return []
 
         def _normalize_function_call_output_to_tool_content(
-            output: Any,
+            output: object,
         ) -> Any:
             """
             Normalize Responses API function_call_output.output into a shape that downstream
@@ -981,7 +1037,7 @@ class LiteLLMCompletionResponsesConfig:
 
             # Some adapters represent tool output as a list of "input_*" parts
             if isinstance(output, list):
-                normalized_blocks: Final[list[dict[str, Any]]] = []
+                normalized_blocks: Final[list[dict[str, object]]] = []
                 text_acc: Final[list[str]] = []
                 for part in output:
                     if not isinstance(part, dict):
@@ -1082,7 +1138,7 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_responses_api_function_call_to_chat_completion_message(
-        function_call: dict[str, Any],
+        function_call: Mapping[str, str],
     ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]:
         """
         Transform a Responses API function_call into a Chat Completion message with tool calls
@@ -1127,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig:
         return [chat_completion_response_message]
 
     @staticmethod
-    def _resolve_file_id(item: dict[str, Any]) -> str | None:
+    def _resolve_file_id(item: Mapping[str, object]) -> object:
         """
         Return the effective file_id for a Responses API input_file item.
         Explicit file_id takes precedence; file_url is used as fallback so
@@ -1136,7 +1192,7 @@ class LiteLLMCompletionResponsesConfig:
         return item.get("file_id") or item.get("file_url") or None
 
     @staticmethod
-    def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]:
+    def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]:
         """
         Transform a Responses API input_file item to a Chat Completion file item
 
@@ -1146,21 +1202,21 @@ class LiteLLMCompletionResponsesConfig:
         Returns:
             Dictionary with transformed file structure for Chat Completion
         """
-        file_dict: Final[dict[str, Any]] = {}
+        file_dict: Final[dict[str, object]] = {}
         file_id: Final = LiteLLMCompletionResponsesConfig._resolve_file_id(item)
         if file_id:
             file_dict["file_id"] = file_id
         if item.get("file_data"):
             file_dict["file_data"] = item["file_data"]
 
-        new_item: Final[dict[str, Any]] = {"type": "file", "file": file_dict}
+        new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict}
         if "cache_control" in item:
             new_item["cache_control"] = item["cache_control"]
         return new_item
 
     @staticmethod
     def _transform_input_image_item_to_image_item(
-        item: dict[str, Any],
+        item: Mapping[str, str],
     ) -> ChatCompletionImageObject:
         """
         Transform a Responses API input_image item to a Chat Completion image item
@@ -1173,8 +1229,8 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_responses_api_content_to_chat_completion_content(
-        content: Any,
-    ) -> str | list[str | dict[str, Any]]:
+        content: object,
+    ) -> str | list[str | dict[str, object]]:
         """
         Transform a Responses API content into a Chat Completion content
 
@@ -1188,7 +1244,7 @@ class LiteLLMCompletionResponsesConfig:
         elif isinstance(content, str):
             return content
         elif isinstance(content, list):
-            content_list: Final[list[str | dict[str, Any]]] = []
+            content_list: Final[list[str | dict[str, object]]] = []
             for item in content:
                 if isinstance(item, str):
                     content_list.append(item)
@@ -1198,8 +1254,8 @@ class LiteLLMCompletionResponsesConfig:
                             LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item)
                         )
                     elif item.get("type") == "input_image":
-                        image_block = dict(
-                            LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)
+                        image_block = _STR_KEY_DICT_ADAPTER.validate_python(
+                            dict(LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item))
                         )
                         if "cache_control" in item:
                             image_block["cache_control"] = item["cache_control"]
@@ -1209,7 +1265,7 @@ class LiteLLMCompletionResponsesConfig:
                         text_value = item.get("text")
                         if text_value is None:
                             continue
-                        content_block: dict[str, Any] = {
+                        content_block: dict[str, object] = {
                             "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type(
                                 item.get("type") or "text"
                             ),
@@ -1299,7 +1355,7 @@ class LiteLLMCompletionResponsesConfig:
                 parameters = dict(typed_tool.get("parameters", {}) or {})
                 if not parameters or "type" not in parameters:
                     parameters["type"] = "object"
-                chat_completion_tool: dict[str, Any] = {
+                chat_completion_tool: dict[str, object] = {
                     "type": "function",
                     "function": {
                         "name": typed_tool.get("name") or "",
@@ -1340,7 +1396,7 @@ class LiteLLMCompletionResponsesConfig:
     @staticmethod
     def transform_chat_completion_tool_params_to_responses_api_tools(
         chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None,
-    ) -> list[dict[str, Any]]:
+    ) -> list[dict[str, object]]:
         """
         Transform Chat Completion tool params (e.g. from guardrail output) back to
         Responses API request tool format. Inverse of
@@ -1348,7 +1404,7 @@ class LiteLLMCompletionResponsesConfig:
         """
         if chat_completion_tools is None or not chat_completion_tools:
             return []
-        result: Final[list[dict[str, Any]]] = []
+        result: Final[list[dict[str, object]]] = []
         for tool in chat_completion_tools:
             if not isinstance(tool, dict):
                 result.append(tool)
@@ -1358,7 +1414,7 @@ class LiteLLMCompletionResponsesConfig:
                 parameters = dict(fn.get("parameters", {}) or {})
                 if not parameters or "type" not in parameters:
                     parameters["type"] = "object"
-                responses_tool: dict[str, Any] = {
+                responses_tool: dict[str, object] = {
                     "type": "function",
                     "name": fn.get("name") or "",
                     "description": fn.get("description") or "",
@@ -1510,7 +1566,7 @@ class LiteLLMCompletionResponsesConfig:
     def convert_response_function_tool_call_to_chat_completion_tool_call(
         tool_call_item: Any,
         index: int = 0,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format.
 
@@ -1536,7 +1592,7 @@ class LiteLLMCompletionResponsesConfig:
                     else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {})
                 )
 
-        function_dict: Final[dict[str, Any]] = {
+        function_dict: Final[dict[str, object]] = {
             "name": tool_call_item.name,
             "arguments": tool_call_item.arguments,
         }
@@ -1544,7 +1600,7 @@ class LiteLLMCompletionResponsesConfig:
         if provider_specific_fields:
             function_dict["provider_specific_fields"] = provider_specific_fields
 
-        tool_call_dict: Final[dict[str, Any]] = {
+        tool_call_dict: Final[dict[str, object]] = {
             "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item(
                 getattr(tool_call_item, "id", None),
                 getattr(tool_call_item, "call_id", None),
@@ -1561,9 +1617,9 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def convert_apply_patch_tool_call_to_chat_completion_tool_call(
-        tool_call_item: Any,
+        tool_call_item: "ResponseApplyPatchToolCall",
         index: int = 0,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """
         Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format.
 
@@ -1581,7 +1637,7 @@ class LiteLLMCompletionResponsesConfig:
         import json
 
         operation_dict: Final = tool_call_item.operation.model_dump()
-        tool_call_dict: Final[dict[str, Any]] = {
+        tool_call_dict: Final[dict[str, object]] = {
             "id": tool_call_item.call_id,
             "function": {
                 "name": "apply_patch",
@@ -1795,9 +1851,11 @@ class LiteLLMCompletionResponsesConfig:
         if not images:
             return image_generation_items
 
-        for idx, image_item in enumerate(images):
+        for idx, image_item in enumerate(_DICT_ITEMS_LIST_ADAPTER.validate_python(images)):
             # Extract base64 from data URL
-            image_url = image_item.get("image_url", {}).get("url", "")
+            image_url = _TEXT_ADAPTER.validate_python(
+                _ANY_KEY_DICT_ADAPTER.validate_python(image_item.get("image_url", {})).get("url", "")
+            )
             base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url)
 
             if base64_data:
@@ -2048,8 +2106,8 @@ class LiteLLMCompletionResponsesConfig:
 
     @staticmethod
     def _transform_text_format_to_response_format(
-        text_param: dict[str, Any] | Any,
-    ) -> dict[str, Any] | None:
+        text_param: object,
+    ) -> dict[str, object] | None:
         """
         Transform Responses API text.format parameter to Chat Completion response_format parameter.
 

litellm/responses/main.py

@@ -1064,6 +1064,7 @@ def responses(
                 extra_headers=extra_headers,
                 extra_body=extra_body,
                 timeout=timeout if timeout is not None else request_timeout,
+                allowed_openai_params=allowed_openai_params,
                 **kwargs,
             )
 

litellm/responses/streaming_iterator.py

@@ -5,7 +5,7 @@ import json
 import time
 import traceback
 import uuid
-from collections.abc import Mapping
+from collections.abc import Awaitable, Callable, Mapping
 from datetime import datetime
 from functools import lru_cache
 from types import MappingProxyType
@@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi
 from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
 from litellm.types.llms.openai import (
     PART_UNION_TYPES,
+    ResponseAPIUsage,
     ResponsesAPIResponse,
     ResponsesAPIStreamEvents,
     ResponsesAPIStreamingResponse,
@@ -112,7 +113,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(
 
 
 def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]:
-    if isinstance(error_obj, dict):
+    if _is_json_object(error_obj):
         raw_message = error_obj.get("message")
         raw_type = error_obj.get("type")
         raw_code = error_obj.get("code")
@@ -243,7 +244,9 @@ class BaseResponsesAPIStreamingIterator:
                 # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a
                 # truthy child Mock for any attribute, which breaks tests and is wrong on stream.
                 if "response" in parsed_chunk:
-                    response_object: Final = getattr(openai_responses_api_chunk, "response", None)
+                    response_object: Final[ResponsesAPIResponse | None] = getattr(
+                        openai_responses_api_chunk, "response", None
+                    )
                     if response_object is not None:
                         response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
                             responses_api_response=response_object,
@@ -279,7 +282,9 @@ class BaseResponsesAPIStreamingIterator:
                             model_id=_stream_model_id,
                         )
                 elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE:
-                    _part: Final = getattr(openai_responses_api_chunk, "part", None)
+                    _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr(
+                        openai_responses_api_chunk, "part", None
+                    )
                     if _part is not None:
                         if isinstance(_part, dict):
                             ResponsesAPIRequestUtils._encode_container_ids_in_annotations(
@@ -302,7 +307,7 @@ class BaseResponsesAPIStreamingIterator:
                         openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
                         openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
                     ):
-                        item: Final = getattr(openai_responses_api_chunk, "item", None)
+                        item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None)
                         if item:
                             encrypted_content: Final = getattr(item, "encrypted_content", None)
                             if encrypted_content and isinstance(encrypted_content, str):
@@ -324,9 +329,11 @@ class BaseResponsesAPIStreamingIterator:
                     self.completed_response = openai_responses_api_chunk
                     # Add cost to usage object if include_cost_in_streaming_usage is True
                     if litellm.include_cost_in_streaming_usage and self.logging_obj is not None:
-                        response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None)
+                        response_obj: Final[ResponsesAPIResponse | None] = getattr(
+                            openai_responses_api_chunk, "response", None
+                        )
                         if response_obj:
-                            usage_obj: Final[Any | None] = getattr(response_obj, "usage", None)
+                            usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
                             if usage_obj is not None:
                                 try:
                                     cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj)
@@ -414,7 +421,9 @@ class BaseResponsesAPIStreamingIterator:
         async_failure_handler / failure_handler so logging integrations correctly
         record the call as failed.
         """
-        response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None
+        response_obj: Final[ResponsesAPIResponse | None] = (
+            getattr(self.completed_response, "response", None) if self.completed_response else None
+        )
         error_info: Final = getattr(response_obj, "error", None) if response_obj else None
         error_message, error_type, error_code = _error_event_fields(error_info)
         self._record_failed_response_usage(response_obj)
@@ -429,7 +438,7 @@ class BaseResponsesAPIStreamingIterator:
     def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
         if response_obj is None or self.logging_obj is None:
             return
-        usage_obj: Final = getattr(response_obj, "usage", None)
+        usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
         if usage_obj is None:
             return
         try:
@@ -506,7 +515,7 @@ class BaseResponsesAPIStreamingIterator:
             return
 
         request_kwargs = getattr(caching_handler, "request_kwargs", None)
-        if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True:
+        if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True:
             return
         request_kwargs = request_kwargs.copy()
         preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
@@ -606,7 +615,7 @@ class BaseResponsesAPIStreamingIterator:
         if self.completed_response is None:
             return
 
-        request_payload: Final[dict[str, Any]] = {}
+        request_payload: Final[dict[str, object]] = {}
         if isinstance(self.request_data, dict):
             request_payload.update(self.request_data)
         try:
@@ -695,11 +704,15 @@ class BaseResponsesAPIStreamingIterator:
             pass
 
 
-async def call_post_streaming_hooks_for_testing(iterator, chunk):
+async def call_post_streaming_hooks_for_testing(
+    iterator: object, chunk: ResponsesAPIStreamingResponse
+) -> ResponsesAPIStreamingResponse:
     """
     Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped.
     """
-    hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None)
+    hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = (
+        getattr(iterator, "_call_post_streaming_deployment_hook", None)
+    )
     if hook_fn is None:
         return chunk
     return await hook_fn(chunk)
@@ -1019,7 +1032,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
 def _dump_response_object(obj: Any) -> dict[str, Any]:
     if hasattr(obj, "model_dump"):
         return obj.model_dump()
-    if isinstance(obj, dict):
+    if _is_json_object(obj):
         return obj
     return {}
 
@@ -1684,7 +1697,7 @@ class ResponsesWebSocketStreaming:
             return response_str
 
         try:
-            evt_obj: Final = json.loads(response_str)
+            evt_obj: Final[Mapping[str, object]] = json.loads(response_str)
         except (json.JSONDecodeError, TypeError):
             return response_str
 
@@ -1925,7 +1938,7 @@ class ManagedResponsesWebSocketHandler:
 
     @staticmethod
     def _extract_output_messages(
-        completed_event: dict[str, Any],
+        completed_event: dict[str, object],
     ) -> list[dict[str, object]]:
         """
         Convert the output items in a ``response.completed`` event into
@@ -2065,7 +2078,7 @@ class ManagedResponsesWebSocketHandler:
           Flat:   {"type": "response.create", "input": [...], "model": "...", ...}
         """
         nested: Final = msg_obj.get("response")
-        response_params: Final[dict[str, Any]] = (
+        response_params: Final[dict[str, object]] = (
             nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"}
         )
         return {
@@ -2076,7 +2089,7 @@ class ManagedResponsesWebSocketHandler:
 
     def _apply_history(
         self,
-        call_kwargs: dict[str, Any],
+        call_kwargs: dict[str, object],
         previous_response_id: str | None,
         current_messages: list[dict[str, object]],
         prior_history: list[dict[str, object]],
@@ -2129,7 +2142,7 @@ class ManagedResponsesWebSocketHandler:
             return False
         return event_provider == self._connection_provider
 
-    def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None:
+    def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None:
         """Inject connection-level credentials and metadata into call_kwargs."""
         if self.api_key is not None:
             call_kwargs["api_key"] = self.api_key

litellm/router.py

@@ -46,6 +46,7 @@ from litellm.constants import (
     DEFAULT_HEALTH_CHECK_INTERVAL,
     DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
     DEFAULT_MAX_LRU_CACHE_SIZE,
+    SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
 )
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.asyncify import run_async_function
@@ -109,6 +110,7 @@ from litellm.router_utils.common_utils import (
     filter_team_based_models,
     filter_web_search_deployments,
     resolve_model_group_alias,
+    truncate_fallback_error_detail,
 )
 from litellm.router_utils.cooldown_cache import CooldownCache
 from litellm.router_utils.cooldown_handlers import (
@@ -134,6 +136,7 @@ from litellm.router_utils.handle_error import (
 from litellm.router_utils.health_state_cache import DeploymentHealthCache
 from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
     DeploymentAffinityCheck,
+    warn_on_unknown_model_group_affinity_flags,
 )
 from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
     build_io_token_rate_limit_headers,
@@ -342,6 +345,12 @@ def _replay_live_router_model_cost() -> None:
 set_live_deployment_replay(_replay_live_router_model_cost)
 
 
+# Kwargs that log_retry must not copy into a retry breadcrumb. The breadcrumbs reach spend
+# logs and logging callbacks, and these carry either the request payload or router-internal
+# walk state rather than anything that identifies the failed attempt.
+RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(("messages", "original_function", "attempted_targets"))
+
+
 class Router:
     model_names: set = set()
     cache_responses: bool | None = False
@@ -596,6 +605,10 @@ class Router:
         # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``.
         self._zero_cost_cache: dict[str, bool] = {}
 
+        self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
+        self.model_group_affinity_config = model_group_affinity_config
+        warn_on_unknown_model_group_affinity_flags(model_group_affinity_config)
+
         if model_list is not None:
             # set_model_list will build indices automatically
             self.set_model_list(model_list)
@@ -737,7 +750,6 @@ class Router:
             litellm.failure_callback = [self.deployment_callback_on_failure]
         self.routing_strategy_args = routing_strategy_args
         self.provider_budget_config = provider_budget_config
-        self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
         self.router_budget_logger: RouterBudgetLimiting | None = None
         if RouterBudgetLimiting.should_init_router_budget_limiter(
             model_list=model_list, provider_budget_config=self.provider_budget_config
@@ -759,7 +771,6 @@ class Router:
                 )
 
         self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy
-        self.model_group_affinity_config: dict[str, list[str]] | None = model_group_affinity_config
 
         self.allowed_fails_policy: AllowedFailsPolicy | None = None
         if allowed_fails_policy is not None:
@@ -782,21 +793,8 @@ class Router:
         # If model_group_affinity_config is set but no global affinity checks were
         # enabled, we still need the DeploymentAffinityCheck callback (with global
         # flags all False) so per-group config can activate affinity per model group.
-        if self.model_group_affinity_config and not any(
-            isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])
-        ):
-            if self.optional_callbacks is None:
-                self.optional_callbacks = []
-            affinity_callback: Final = DeploymentAffinityCheck(
-                cache=self.cache,
-                ttl_seconds=self.deployment_affinity_ttl_seconds,
-                enable_user_key_affinity=False,
-                enable_responses_api_affinity=False,
-                enable_session_id_affinity=False,
-                model_group_affinity_config=self.model_group_affinity_config,
-            )
-            self.optional_callbacks.append(affinity_callback)
-            litellm.logging_callback_manager.add_litellm_callback(affinity_callback)
+        if self.model_group_affinity_config:
+            self._ensure_deployment_affinity_callback()
 
         if self.alerting_config is not None:
             self._initialize_alerting()
@@ -1655,6 +1653,28 @@ class Router:
             _move_before_deployment_affinity(self.optional_callbacks, ec_callback)
             _move_before_deployment_affinity(litellm.callbacks, ec_callback)
 
+    def _ensure_deployment_affinity_callback(self) -> None:
+        """Register the DeploymentAffinityCheck callback (global flags all False) if absent.
+
+        Needed when nothing enabled a global affinity flag but affinity can still
+        activate per request: per-group `model_group_affinity_config` entries, or the
+        session-affinity marker a complexity router stamps at pre-routing time.
+        """
+        if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])):
+            return
+        if self.optional_callbacks is None:
+            self.optional_callbacks = []
+        affinity_callback: Final = DeploymentAffinityCheck(
+            cache=self.cache,
+            ttl_seconds=self.deployment_affinity_ttl_seconds,
+            enable_user_key_affinity=False,
+            enable_responses_api_affinity=False,
+            enable_session_id_affinity=False,
+            model_group_affinity_config=self.model_group_affinity_config,
+        )
+        self.optional_callbacks.append(affinity_callback)
+        litellm.logging_callback_manager.add_litellm_callback(affinity_callback)
+
     def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None):
         if optional_pre_call_checks is None:
             return
@@ -6361,17 +6381,16 @@ class Router:
                 return response
         except Exception as new_exception:
             parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
-            fallback_failure_exception_str = redact_string(str(new_exception))
+            fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception)))
             cooldown_info: Final = await _async_get_cooldown_deployments_with_debug_info(
                 litellm_router_instance=self,
                 parent_otel_span=parent_otel_span,
             )
             verbose_router_logger.error(
                 "litellm.router.py::async_function_with_fallbacks() - "
-                "Error occurred while trying to do fallbacks - %s\n%s\n"
+                "Error occurred while trying to do fallbacks - %s\n"
                 "Debug Information:\nCooldown Deployments=%s",
                 fallback_failure_exception_str,
-                redact_string(traceback.format_exc()),
                 cooldown_info,
             )
 
@@ -7162,7 +7181,7 @@ class Router:
                 k,
                 v,
             ) in kwargs.items():  # log everything in kwargs except the old previous_models value - prevent nesting
-                if k not in [_metadata_var, "messages", "original_function"]:
+                if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS:
                     previous_model[k] = v
                 elif k == _metadata_var and isinstance(v, dict):
                     previous_model[_metadata_var] = {}
@@ -7677,6 +7696,8 @@ class Router:
             strategy=complexity_router,
             strategy_label="Complexity-router",
         )
+        if complexity_router._uses_deployment_pin:
+            self._ensure_deployment_affinity_callback()
 
     def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
         """True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
@@ -11184,6 +11205,9 @@ class Router:
         router_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs)
         if router_strategy is None:
             self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
+            self._stamp_or_clear_metadata_key(
+                request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None
+            )
             return None
 
         pre_routing_hook_response: Final = await router_strategy.async_pre_routing_hook(
@@ -11197,6 +11221,11 @@ class Router:
             request_kwargs=request_kwargs,
             routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None),
         )
+        self._stamp_or_clear_metadata_key(
+            request_kwargs=request_kwargs,
+            key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
+            value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None),
+        )
 
         # `model` (the alias, e.g. "smart-router") is never the deployment actually
         # called - apply the alias's own litellm_params (besides `model` itself,
@@ -11228,21 +11257,40 @@ class Router:
         to the deployment that actually served the request. Every attempt therefore
         writes or clears, never just writes.
         """
-        if routing_decision is None:
+        Router._stamp_or_clear_metadata_key(
+            request_kwargs=request_kwargs,
+            key="routing_decision",
+            value=(
+                None
+                if routing_decision is None
+                else Router._redact_prompt_text_if_needed(
+                    request_kwargs=request_kwargs, routing_decision=routing_decision
+                )
+            ),
+        )
+
+    @staticmethod
+    def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None:
+        """Write a proxy-internal metadata key for THIS routing attempt, or clear it.
+
+        Fallbacks and retries re-enter the pre-routing hook with the same
+        `request_kwargs`, so every attempt must write or clear, never just write;
+        a value left behind by an earlier attempt would be attributed to this one.
+        `get_or_create_metadata_bucket` is the single owner of "which dict holds
+        proxy-internal metadata": it picks `litellm_metadata` when present (so the
+        value never lands in the `metadata` dict that routes like /v1/messages
+        forward to the provider) and replaces a non-dict value rather than silently
+        skipping the write. Clearing pops from BOTH buckets so a request whose
+        bucket resolution changed between attempts cannot resurrect a stale value.
+        """
+        if value is None:
             for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")):
                 if isinstance(bucket, dict):
-                    bucket.pop("routing_decision", None)
+                    bucket.pop(key, None)
             return
 
-        # `get_or_create_metadata_bucket` is the single owner of "which dict holds
-        # proxy-internal metadata": it picks `litellm_metadata` when present (so the
-        # decision never lands in the `metadata` dict that routes like /v1/messages
-        # forward to the provider) and replaces a non-dict value rather than silently
-        # skipping the write.
         _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs)
-        metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed(
-            request_kwargs=request_kwargs, routing_decision=routing_decision
-        )
+        metadata_bucket[key] = value
 
     @staticmethod
     def _redact_prompt_text_if_needed(

litellm/router_strategy/complexity_router/complexity_router.py

@@ -1651,6 +1651,28 @@ class ComplexityRouter(CustomLogger):
         caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped"
         return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}"
 
+    @property
+    def _uses_tier_pin(self) -> bool:
+        return bool(self.config.session_affinity and not self.config.plugins)
+
+    @property
+    def _uses_deployment_pin(self) -> bool:
+        """session_affinity implies the deployment pin: a session frozen onto one model
+        group but load-balanced across its deployments would still go cache-cold, which
+        is the exact failure both flags exist to prevent."""
+        return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins)
+
+    def _with_session_deployment_affinity(
+        self, response: PreRoutingHookResponse | None
+    ) -> PreRoutingHookResponse | None:
+        if response is None or not self._uses_deployment_pin:
+            return response
+        return response.model_copy(
+            update={  # mutable-ok: model_copy types update as a plain dict
+                "session_affinity_ttl_seconds": self.config.session_affinity_ttl_seconds
+            }
+        )
+
     async def async_pre_routing_hook(
         self,
         model: str,
@@ -1685,7 +1707,7 @@ class ComplexityRouter(CustomLogger):
         resolved_messages: Final = self._resolve_messages(messages, request_kwargs)
         conversation_continuing: Final = _conversation_is_continuing(resolved_messages)
 
-        use_session_affinity: Final = self.config.session_affinity and not self.config.plugins
+        use_session_affinity: Final = self._uses_tier_pin
         session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None
         cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
 
@@ -1724,16 +1746,19 @@ class ComplexityRouter(CustomLogger):
                         "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model
                     )
                     has_original_messages: Final = messages is not None and len(messages) > 0
-                    return PreRoutingHookResponse(
-                        model=routed_model,
-                        messages=messages if has_original_messages else None,
-                        routing_decision=self._build_routing_decision(
-                            routed_model=routed_model,
-                            cause=cause,
-                            escalation_keyword=pin_escalation_keyword,
-                            escalated=escalated,
-                            conversation_continuing=conversation_continuing,
-                        ),
+                    return self._with_session_deployment_affinity(
+                        PreRoutingHookResponse(
+                            model=routed_model,
+                            messages=messages if has_original_messages else None,
+                            routing_decision=self._build_routing_decision(
+                                routed_model=routed_model,
+                                cause=cause,
+                                tier=self._tier_for_model(routed_model),
+                                escalation_keyword=pin_escalation_keyword,
+                                escalated=escalated,
+                                conversation_continuing=conversation_continuing,
+                            ),
+                        )
                     )
 
         response: Final = await self._classify_and_route(
@@ -1751,7 +1776,7 @@ class ComplexityRouter(CustomLogger):
                 value=response.model,
                 ttl=self.config.session_affinity_ttl_seconds,
             )
-        return response
+        return self._with_session_deployment_affinity(response)
 
     async def _classify_and_route(
         self,
@@ -1797,7 +1822,8 @@ class ComplexityRouter(CustomLogger):
 
         if user_message is None:
             verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model")
-            if not self.config.plugins and self.config.default_model:
+            default_model_first: Final = not self.config.plugins and self.config.default_model
+            if default_model_first:
                 # No plugins configured: preserve the pre-existing default_model-first
                 # priority exactly (changing it would be a silent behavior change for
                 # every non-plugin user, not just a security fix).
@@ -1809,12 +1835,14 @@ class ComplexityRouter(CustomLogger):
                 routed_model = await self._pick_model_for_tier(
                     ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
                 )
+            fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM
             return PreRoutingHookResponse(
                 model=routed_model,
                 messages=messages if has_original_messages else None,
                 routing_decision=self._build_routing_decision(
                     routed_model=routed_model,
                     cause="default_fallback",
+                    tier=fallback_tier,
                     conversation_continuing=conversation_continuing,
                 ),
             )

litellm/router_strategy/complexity_router/config.py

@@ -508,13 +508,39 @@ class ComplexityRouterConfig(BaseModel):
             "session's first turn and reuse it for every later turn, skipping re-classification. "
             "Off by default so every turn is classified on its own merits and routed to the cheapest "
             "adequate tier. Set True to keep a multi-turn session on one model, which preserves "
-            "provider prompt caches and avoids cross-model conversation-history errors."
+            "provider prompt caches and avoids cross-model conversation-history errors. Always "
+            "implies the deployment pin regardless of deployment_affinity: the session sticks to "
+            "one deployment of the pinned model, since freezing the model while re-shuffling its "
+            "deployments would still go cache-cold."
+        ),
+    )
+    deployment_affinity: bool = Field(
+        default=True,
+        description=(
+            "When True and a session_id is resolvable on the request, pin the deployment chosen "
+            "inside each routed model group and reuse it whenever the session returns to that "
+            "group, without pinning which group the session routes to. Independent of "
+            "session_affinity, which pins the model group instead (and always carries this "
+            "deployment pin with it): with session_affinity off, "
+            "every turn is still classified on its own merits while a session that escalates to a "
+            "stronger tier and comes back still lands on the deployment it used before, which is "
+            "what keeps a provider prompt cache warm. Pins are held per model group, so switching "
+            "tiers does not disturb the pin left behind in the previous group. On by default "
+            "because re-shuffling a conversation across deployments of the same model discards "
+            "that cache for no benefit; set False to keep every turn load-balanced across the "
+            "group, which is what a deployment set with tight per-deployment rate limits wants. "
+            "Inert when no session_id is resolvable, since there is nothing to key a pin on, and "
+            "suppressed when plugins are configured, for the same reason session_affinity is."
         ),
     )
     session_affinity_ttl_seconds: int = Field(
         default=3600,
         gt=0,
-        description="TTL for the session affinity pin; refreshed on every cache hit",
+        description=(
+            "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the "
+            "session_affinity model pin and the deployment_affinity deployment pin, so it measures "
+            "idle time for the session's routing decisions rather than total session length"
+        ),
     )
 
     plugins: list[RoutingPlugin] | None = Field(

litellm/router_utils/common_utils.py

@@ -7,6 +7,7 @@ if TYPE_CHECKING:
     from litellm.types.llms.openai import OpenAIFileObject
 
 from litellm._logging import verbose_logger
+from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS
 from litellm.exceptions import BadRequestError
 from litellm.types.router import CredentialLiteLLMParams
 
@@ -43,6 +44,22 @@ def resolve_model_group_alias(model_group_alias: object, model: str) -> str | No
     return target
 
 
+def truncate_fallback_error_detail(detail: str) -> str:
+    """
+    Bound a fallback failure detail before it is logged or appended to an exception message.
+
+    Each level of the fallback walk records the failure of the level below it, so an
+    untruncated detail carries every nested failure with it and grows superlinearly with
+    the number of attempted model groups. One deterministic pre-network failure walked
+    through a small fallback graph is enough to turn that into hundreds of megabytes of
+    output on the event-loop thread, which starves the process that produced it.
+    """
+    if len(detail) <= ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS:
+        return detail
+    dropped: Final = len(detail) - ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS
+    return f"{detail[:ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS]}... [truncated {dropped} characters]"
+
+
 def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str:
     """
     Hash of the credential params, used for mapping the file id to the right model

litellm/router_utils/fallback_event_handlers.py

@@ -1,3 +1,6 @@
+import hashlib
+import json
+from dataclasses import dataclass
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Final
 
@@ -19,6 +22,52 @@ else:
     LitellmRouter = Any
 
 
+def fallback_attempt_key(fallback_target: object) -> str | None:
+    """
+    Identity of one fallback attempt, so the same attempt is never made twice per request.
+
+    A bare model group name and a `{"model": name}` entry describe the same attempt. An
+    entry carrying anything else describes a different one and keeps its own identity: a
+    client-side fallback list overrides request params such as `messages`, and the router
+    re-targets the group that just failed by attaching `_target_order` or
+    `_excluded_deployment_ids` to select a different set of deployments inside it. The
+    payload is hashed rather than kept, so a large `messages` override does not make the
+    request hold a second copy of itself.
+
+    Returns None for a shape with no usable identity, which is never skipped.
+    """
+    if isinstance(fallback_target, str):
+        return fallback_target
+    if not isinstance(fallback_target, dict):
+        return None
+    model: Final = fallback_target.get("model")
+    if tuple(fallback_target) == ("model",) and isinstance(model, str):
+        return model
+    serialized: Final = json.dumps(fallback_target, sort_keys=True, default=str)
+    return hashlib.sha256(serialized.encode()).hexdigest()
+
+
+@dataclass(slots=True)
+class AttemptedFallbackTargets:
+    """
+    The fallback attempts a single request has already made.
+
+    One instance is created on the first fallback hop and shared by reference for the rest
+    of the walk, so an attempt made in one branch is not repeated in a sibling branch.
+    Without it the walk enumerates paths rather than attempts: a fallback graph containing
+    a cycle retries one deterministic failure once per path through the cycle, and a
+    client-side fallback list is re-walked at every level of the recursion.
+    """
+
+    keys: frozenset[str] = frozenset()
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.keys
+
+    def record(self, key: str) -> None:
+        self.keys = self.keys | frozenset((key,))
+
+
 def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool:
     """
     Handles wildcard routing scenario
@@ -106,7 +155,14 @@ async def run_async_fallback(
         fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"]
         original_model_group: The original model group. example: "gpt-3.5-turbo"
         original_exception: The original exception.
-        **kwargs: Keyword arguments.
+        **kwargs: Keyword arguments. `attempted_targets` carries the fallback attempts
+            already made for this request, created on the first hop and shared by reference
+            for the rest of the walk. A target already in it is skipped, so neither a
+            fallback graph that loops back on itself nor a client-side fallback list
+            re-walked at each level can repeat an attempt that has already failed. Identity
+            comes from `fallback_attempt_key`, so an entry that overrides request params or
+            re-targets the failed group with a different deployment selection stays distinct
+            from a bare name.
 
     Returns:
         The response from the successful fallback model group.
@@ -120,10 +176,27 @@ async def run_async_fallback(
 
     error_from_fallbacks = original_exception
     fallback_errors = (get_fallback_error_info(original_exception),)
+    # Read out of kwargs and narrowed here rather than declared as a parameter: every caller
+    # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter
+    # would carry an annotation that no call site can actually be checked against.
+    carried_targets: Final = kwargs.get("attempted_targets")
+    attempted: Final = (
+        carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets()
+    )
+    attempted.record(original_model_group)
 
     for mg in fallback_model_group:
         if mg == original_model_group:
             continue
+        attempt_key = fallback_attempt_key(mg)
+        if attempt_key is not None:
+            if attempt_key in attempted:
+                verbose_router_logger.info(
+                    "Skipping fallback to model_group = %s, already attempted for this request",
+                    mask_sensitive_structure(mg),
+                )
+                continue
+            attempted.record(attempt_key)
         try:
             # LOGGING
             kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception)
@@ -138,6 +211,7 @@ async def run_async_fallback(
             fallback_depth = fallback_depth + 1
             kwargs["fallback_depth"] = fallback_depth
             kwargs["max_fallbacks"] = max_fallbacks
+            kwargs["attempted_targets"] = attempted
             if include_fallback_errors:
                 kwargs["include_fallback_errors"] = include_fallback_errors
             response = await litellm_router.async_function_with_fallbacks(*args, **kwargs)

litellm/router_utils/pre_call_checks/deployment_affinity_check.py

@@ -13,12 +13,15 @@ where routing to a consistent deployment is still beneficial.
 """
 
 import hashlib
+import json
+from collections.abc import Mapping, Sequence
 from typing import Any, Final, cast
 
 from typing_extensions import TypedDict
 
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
 from litellm.integrations.custom_logger import CustomLogger, Span
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import AllMessageValues
@@ -29,6 +32,47 @@ class DeploymentAffinityCacheValue(TypedDict):
     model_id: str
 
 
+VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset(
+    {
+        "deployment_affinity",
+        "responses_api_deployment_check",
+        "session_affinity",
+        "encrypted_content_affinity",
+    }
+)
+
+
+def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapping[str, Sequence[str]] | None) -> None:
+    """`model_group_affinity_config` is one Router-level config consumed by two callbacks:
+    DeploymentAffinityCheck acts on three of the flags and EncryptedContentAffinityCheck
+    on the fourth, so typo detection lives here at the schema, not inside either consumer.
+    """
+    if model_group_affinity_config is None:
+        return
+    for group, flags in model_group_affinity_config.items():
+        unknown = set(flags) - VALID_MODEL_GROUP_AFFINITY_FLAGS
+        if unknown:
+            verbose_router_logger.warning(
+                "model_group_affinity_config: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s",
+                unknown,
+                group,
+                VALID_MODEL_GROUP_AFFINITY_FLAGS,
+            )
+
+
+_CLAIM_PIN_SCRIPT: Final = """
+local current = redis.call('GET', KEYS[1])
+if current == false then
+  redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
+  return ARGV[1]
+end
+if current == ARGV[1] then
+  redis.call('EXPIRE', KEYS[1], ARGV[2])
+end
+return current
+"""
+
+
 class DeploymentAffinityCheck(CustomLogger):
     """
     Router deployment affinity callback.
@@ -38,14 +82,6 @@ class DeploymentAffinityCheck(CustomLogger):
     """
 
     CACHE_KEY_PREFIX = "deployment_affinity:v1"
-    VALID_FLAGS = frozenset(
-        {
-            "deployment_affinity",
-            "responses_api_deployment_check",
-            "session_affinity",
-            "encrypted_content_affinity",
-        }
-    )
 
     def __init__(
         self,
@@ -63,15 +99,6 @@ class DeploymentAffinityCheck(CustomLogger):
         self.enable_responses_api_affinity = enable_responses_api_affinity
         self.enable_session_id_affinity = enable_session_id_affinity
         self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {}
-        for group, flags in self.model_group_affinity_config.items():
-            unknown = set(flags) - self.VALID_FLAGS
-            if unknown:
-                verbose_router_logger.warning(
-                    "DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s",
-                    unknown,
-                    group,
-                    self.VALID_FLAGS,
-                )
 
     def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]:
         """
@@ -218,8 +245,13 @@ class DeploymentAffinityCheck(CustomLogger):
         return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}"
 
     @classmethod
-    def get_session_affinity_cache_key(cls, model_group: str, session_id: str) -> str:
-        return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{session_id}"
+    def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str:
+        """Session pins are scoped by the caller's hashed API key so two callers reusing
+        the same client-supplied session_id cannot read or steer each other's pin.
+        `"unscoped"` covers direct Router usage with no authenticated caller, matching
+        the complexity router's own session pin key."""
+        hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped"
+        return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}"
 
     @staticmethod
     def _get_user_key_from_metadata_dict(metadata: dict) -> str | None:
@@ -278,6 +310,97 @@ class DeploymentAffinityCheck(CustomLogger):
                 return session_id
         return None
 
+    @staticmethod
+    def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None:
+        """TTL from the session-affinity marker the Router stamps at pre-routing time
+        when an auto-router routed this request with session_affinity enabled.
+        Marker presence enables session pinning for this request only; anything that
+        is not a positive int is treated as absent."""
+        for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs):
+            ttl = metadata.get(SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY)
+            if isinstance(ttl, int) and not isinstance(ttl, bool) and ttl > 0:
+                return ttl
+        return None
+
+    @staticmethod
+    def _pinned_model_id(stored: object) -> str | None:
+        """Deployment id held by a stored pin, for both the dict shape this writes and the
+        bare string older writers left behind. None when the value is neither."""
+        if isinstance(stored, dict):
+            model_id: Final = stored.get("model_id")
+            return str(model_id) if model_id is not None else None
+        if isinstance(stored, str):
+            return stored
+        return None
+
+    def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None:
+        """The one owner of authoritative local pin writes: a plain set keeps a live
+        key's original expiry (`allow_ttl_override`), so the entry is replaced to make
+        the TTL real. Every local pin write goes through here so the redis-winner sync
+        and the pod-local claim can never disagree about expiry again."""
+        self.cache.in_memory_cache.delete_cache(cache_key)
+        self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds)
+
+    async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None:
+        """First-writer-wins pin write: store `pin_value` only when the key is absent and
+        return the deployment id the key holds afterwards, so a caller learns whether it won
+        by comparing against its own id, and None when the stored value is one no reader can
+        interpret. Concurrent claimers converge on the
+        first write instead of the last. Re-claiming with the stored value refreshes its
+        TTL, the same keepalive the complexity router's model pin documents: an active
+        session must not lose its pin mid-conversation just because it outlives the
+        original write, so `session_affinity_ttl_seconds` bounds idle time, not total
+        session length. On Redis one Lua script does the get-or-set-or-refresh
+        atomically (same registration seam the rate limiters use) and the in-memory
+        tier is synchronized to the winner; without Redis, and whenever Redis is
+        unreachable, the pod-local check-and-set below stands in and is atomic because it
+        runs synchronously on the event loop. Degrading to a pod-local claim rather than
+        propagating the fault is what keeps same-pod stickiness through a Redis blip: the
+        caller only logs this result, so an escaping error would leave the session with no
+        pin at all and reshuffle every turn for the outage, which is worse than losing
+        cross-pod agreement. The redis tier is
+        resolved per call because the proxy attaches it after Router construction
+        (`Router._update_redis_cache`); the compiled script is cached per event loop
+        underneath the registration seam.
+        """
+        redis_cache: Final = self.cache.redis_cache
+        if redis_cache is not None:
+            try:
+                claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
+                raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds)))
+                decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
+                if not isinstance(decoded, str):
+                    return pin_value["model_id"]
+                try:
+                    winner: object = json.loads(decoded)
+                except json.JSONDecodeError:
+                    winner = decoded
+                self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds)
+                return self._pinned_model_id(winner)
+            except Exception as e:  # noqa: BLE001  # any Redis/Lua failure degrades to the pod-local claim, never unpins
+                verbose_router_logger.debug(
+                    "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e
+                )
+
+        return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds)
+
+    def _claim_pin_in_memory(
+        self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int
+    ) -> str | None:
+        """Pod-local half of the claim, used when no Redis tier is attached and as the
+        fallback when the Redis claim fails. Mirrors the Lua script exactly, including
+        the keepalive: re-claiming with the stored value slides the idle window through
+        `_set_local_pin`. Both branches stay synchronous, hence atomic on the event
+        loop."""
+        existing: Final = self.cache.in_memory_cache.get_cache(cache_key)
+        if existing is not None:
+            existing_model_id: Final = self._pinned_model_id(existing)
+            if existing_model_id == pin_value["model_id"]:
+                self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds)
+            return existing_model_id
+        self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds)
+        return pin_value["model_id"]
+
     @staticmethod
     def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
         for deployment in healthy_deployments:
@@ -334,12 +457,21 @@ class DeploymentAffinityCheck(CustomLogger):
         if stable_model_map_key is None:
             return typed_healthy_deployments
 
+        session_affinity_active: Final = (
+            enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None
+        )
+        user_key: Final = (
+            self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
+            if (session_affinity_active or enable_user_key)
+            else None
+        )
+
         # 2) Session-id -> deployment affinity
-        if enable_session_id:
+        if session_affinity_active:
             session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs)
             if session_id is not None:
                 session_cache_key: Final = self.get_session_affinity_cache_key(
-                    model_group=stable_model_map_key, session_id=session_id
+                    model_group=stable_model_map_key, session_id=session_id, user_key=user_key
                 )
                 session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key)
 
@@ -371,7 +503,6 @@ class DeploymentAffinityCheck(CustomLogger):
         if not enable_user_key:
             return typed_healthy_deployments
 
-        user_key: Final = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
         if user_key is None:
             return typed_healthy_deployments
 
@@ -438,18 +569,22 @@ class DeploymentAffinityCheck(CustomLogger):
             enable_session_id,
         ) = self._get_effective_flags(deployment_model_name)
 
-        if not enable_user_key and not enable_session_id:
+        marker_session_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs=kwargs)
+        session_affinity_active: Final = enable_session_id or marker_session_ttl is not None
+
+        if not enable_user_key and not session_affinity_active:
             return None
 
-        user_key = None
-        if enable_user_key:
-            user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
+        user_key: Final = (
+            self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
+            if (enable_user_key or session_affinity_active)
+            else None
+        )
+        session_id: Final = (
+            self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if session_affinity_active else None
+        )
 
-        session_id = None
-        if enable_session_id:
-            session_id = self._get_session_id_from_request_kwargs(request_kwargs=kwargs)
-
-        if user_key is None and session_id is None:
+        if not ((enable_user_key and user_key is not None) or session_id is not None):
             return None
 
         model_info = kwargs.get("model_info")
@@ -473,22 +608,31 @@ class DeploymentAffinityCheck(CustomLogger):
             verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.")
             return None
 
-        if user_key is not None:
+        pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id))
+
+        if enable_user_key and user_key is not None:
             try:
                 cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key)
-                await self.cache.async_set_cache(
-                    cache_key,
-                    DeploymentAffinityCacheValue(model_id=str(model_id)),
-                    ttl=self.ttl_seconds,
-                )
-
-                verbose_router_logger.debug(
-                    "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
-                    deployment_model_name,
-                    model_id,
-                    self.ttl_seconds,
-                    self._shorten_for_logs(user_key),
+                claimed_user_pin: Final = await self._claim_pin(
+                    cache_key=cache_key,
+                    pin_value=pin_value,
+                    ttl_seconds=self.ttl_seconds,
                 )
+                if claimed_user_pin == pin_value["model_id"]:
+                    verbose_router_logger.debug(
+                        "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
+                        deployment_model_name,
+                        model_id,
+                        self.ttl_seconds,
+                        self._shorten_for_logs(user_key),
+                    )
+                else:
+                    verbose_router_logger.debug(
+                        "DeploymentAffinityCheck: affinity pin already claimed model_map_key=%s existing=%s ours=%s",
+                        deployment_model_name,
+                        claimed_user_pin,
+                        model_id,
+                    )
             except Exception as e:
                 # Non-blocking: affinity is a best-effort optimization.
                 verbose_router_logger.debug(
@@ -500,21 +644,31 @@ class DeploymentAffinityCheck(CustomLogger):
         # Also persist Session-ID affinity if enabled and session-id is provided
         if session_id is not None:
             try:
+                session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds
                 session_cache_key: Final = self.get_session_affinity_cache_key(
-                    model_group=deployment_model_name, session_id=session_id
+                    model_group=deployment_model_name, session_id=session_id, user_key=user_key
                 )
-                await self.cache.async_set_cache(
-                    session_cache_key,
-                    DeploymentAffinityCacheValue(model_id=str(model_id)),
-                    ttl=self.ttl_seconds,
-                )
-                verbose_router_logger.debug(
-                    "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
-                    deployment_model_name,
-                    model_id,
-                    self.ttl_seconds,
-                    session_id,
+                claimed_session_pin: Final = await self._claim_pin(
+                    cache_key=session_cache_key,
+                    pin_value=pin_value,
+                    ttl_seconds=session_affinity_ttl,
                 )
+                if claimed_session_pin == pin_value["model_id"]:
+                    verbose_router_logger.debug(
+                        "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
+                        deployment_model_name,
+                        model_id,
+                        session_affinity_ttl,
+                        session_id,
+                    )
+                else:
+                    verbose_router_logger.debug(
+                        "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s",
+                        deployment_model_name,
+                        claimed_session_pin,
+                        model_id,
+                        session_id,
+                    )
             except Exception as e:
                 verbose_router_logger.debug(
                     "DeploymentAffinityCheck: failed to set session affinity cache. model_map_key=%s error=%s",

litellm/types/guardrails.py

@@ -5,6 +5,7 @@ from typing import Any, Final, Literal
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Required, TypedDict
 
+from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
 from litellm.types.proxy.guardrails.guardrail_hooks.akto import (
     AktoConfigModel,
 )
@@ -525,6 +526,15 @@ class BedrockGuardrailConfigModel(BaseModel):
         description="InvokeGuardrailChecks: block when any sensitiveInformation confidenceScore "
         ">= this value (scores are in [0,1]). Set to null to make PII detection detect-only.",
     )
+    chunk_budget_chars: int = Field(
+        default=BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
+        gt=0,
+        description="ApplyGuardrail: batch size, in characters, used to re-send content after AWS "
+        "has rejected a request as too large. Requests AWS accepts are always sent in a single "
+        "call, so this has no effect until a rejection happens. Defaults to 25,000; a batch AWS "
+        "still rejects is bisected automatically, so this value only trades round trips against "
+        "batch size and cannot fail a request on its own.",
+    )
 
 
 class LakeraV2GuardrailConfigModel(BaseModel):

litellm/types/integrations/websearch_interception.py

@@ -2,7 +2,28 @@
 Type definitions for WebSearch Interception integration.
 """
 
-from typing import TypedDict
+from typing import Literal, TypedDict
+
+from pydantic import BaseModel
+
+
+class AnthropicSearchQuery(BaseModel):
+    """``input`` of an Anthropic ``server_tool_use`` block for a web search."""
+
+    query: str
+
+
+class AnthropicServerToolUseBlock(BaseModel):
+    """
+    The ``server_tool_use`` block that must accompany a ``web_search_tool_result``.
+
+    Anthropic requires the pair, with a ``srvtoolu_``-prefixed id shared by both.
+    """
+
+    type: Literal["server_tool_use"] = "server_tool_use"
+    id: str
+    name: Literal["web_search"] = "web_search"
+    input: AnthropicSearchQuery
 
 
 class WebSearchInterceptionConfig(TypedDict, total=False):