Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_post_call_policy_pipeline

.circleci/config.yml

@@ -1440,6 +1440,7 @@ jobs:
             TEST_FILES=$(printf "%s\n" \
               tests/local_testing/test_dual_cache.py \
               tests/local_testing/test_redis_batch_optimizations.py \
+              tests/local_testing/test_redis_increment_with_floor.py \
               tests/local_testing/test_router_utils.py)
             echo "$TEST_FILES" | circleci tests run \
               --verbose \

.github/workflows/test-rust.yml

@@ -117,6 +117,9 @@ jobs:
 
       - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
 
+      - name: Run pytest tests/test_litellm_rust with the compiled extension
+        run: make test-rust-extension
+
       - run: >-
           uv build --wheel --out-dir panic-dist
           --config-setting "maturin.build-args=--features panic-test,extension-module"

Makefile

@@ -4,6 +4,7 @@
 .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
 	test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
 	test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
+	test-rust-extension \
 	info lint lint-inner lint-dev lint-checks format \
 	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 \
@@ -54,6 +55,7 @@ help:
 	@echo "  make test-proxy-unit-b  - Run proxy_unit_tests (p-z, ~28 files)"
 	@echo "  make test-integration   - Run integration tests"
 	@echo "  make test-unit-helm     - Run helm unit tests"
+	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests"
 	@echo ""
 	@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
 	@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
@@ -289,6 +291,17 @@ pre-commit:
 	@$(MAKE) check
 
 # Testing targets
+test-rust-extension:
+	@temporary=$$(mktemp -d) && \
+	trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
+	$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
+	set -- "$$temporary"/wheels/*.whl && \
+	[ "$$#" -eq 1 ] && \
+	UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
+	$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
+	LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
+	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
+
 test: install-test-deps
 	$(UV_RUN) pytest tests/
 

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -8,14 +8,10 @@ import tempfile
 import time
 from dataclasses import dataclass, replace
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Final, Optional
 
 from litellm_proxy_extras import prisma_toolchain
 from litellm_proxy_extras._logging import logger
-from litellm_proxy_extras.replica_identity import (
-    REPLICA_IDENTITY_FULL_ENV_VAR,
-    apply_replica_identity_full,
-)
 from litellm_proxy_extras.prisma_toolchain import (
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
     PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
     prisma_command_timeout,
     prisma_migrate_deploy_timeout,
 )
+from litellm_proxy_extras.replica_identity import (
+    REPLICA_IDENTITY_FULL_ENV_VAR,
+    apply_replica_identity_full,
+)
+
+if TYPE_CHECKING:
+    import psycopg
+    import psycopg.sql
 
 
 def str_to_bool(value: Optional[str]) -> bool:
@@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
 _MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
 
 _MIGRATION_DEADLOCK_MARKER = "deadlock detected"
+INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
+_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
+_INVALID_LITELLM_INDEXES_SQL: Final = (
+    "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
+    "FROM pg_index i "
+    "JOIN pg_class c ON c.oid = i.indexrelid "
+    "JOIN pg_class t ON t.oid = i.indrelid "
+    "JOIN pg_namespace n ON n.oid = t.relnamespace "
+    "WHERE NOT i.indisvalid "
+    "  AND c.relkind = 'i' "
+    "  AND n.nspname = %s "
+    "  AND t.relname LIKE %s "
+    "  AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
+    "ORDER BY c.relname"
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _InvalidIndex:
+    schema: str
+    name: str
+    table_size: str
 
 MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
 
@@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
     def _strip_prisma_query_params(url: str) -> str:
         """Remove Prisma-specific query params (connection_limit, pool_timeout,
         schema, etc.) from DATABASE_URL so psycopg can parse it."""
-        from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
+        from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
 
         parsed = urlparse(url)
         if not parsed.query:
@@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
             "target_session_attrs",
         }
         kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
-        return urlunparse(parsed._replace(query=urlencode(kept)))
+        return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
 
     @staticmethod
     def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
@@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
             ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
         )
 
+    @staticmethod
+    def _invalid_litellm_indexes(
+        conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
+    ) -> tuple[_InvalidIndex, ...]:
+        rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
+        return tuple(_InvalidIndex(*row) for row in rows)
+
+    @staticmethod
+    def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
+        from psycopg import sql
+
+        target: Final = sql.Identifier(index.schema, index.name)
+        if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
+            return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
+        return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
+
+    @staticmethod
+    def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
+        import psycopg
+
+        statement, action = ProxyExtrasDBManager._index_repair(index)
+        try:
+            conn.execute(statement)
+        except psycopg.Error as e:
+            logger.warning(
+                "Could not repair invalid index %s.%s, will retry on the next startup. "
+                "If this keeps happening, run `%s` by hand as the index owner. Error: %s",
+                index.schema,
+                index.name,
+                statement.as_string(conn),
+                e,
+            )
+            return
+        logger.info("%s invalid index %s.%s", action, index.schema, index.name)
+
+    @staticmethod
+    def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
+        """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
+        INVALID (a migration deadlock between replicas is the usual cause; the
+        retried migration skips them because of IF NOT EXISTS). Never raises:
+        returns True when no invalid index remains, False when the repair was
+        skipped or failed and will be retried on the next startup. Looks in the
+        schema DATABASE_URL names, the only URL Prisma migrates through, but
+        connects over DIRECT_URL when set: the session settings, the advisory
+        lock and REINDEX CONCURRENTLY all need one server session, which a
+        transaction pooler does not give."""
+        prisma_url: Final = os.getenv("DATABASE_URL")
+        if not prisma_url:
+            return False
+
+        try:
+            import psycopg
+            from psycopg import sql
+        except ImportError:
+            logger.warning(
+                "psycopg is not installed; skipping the invalid index check. "
+                "Install the litellm[extra_proxy] extra, which includes psycopg."
+            )
+            return False
+
+        schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public"
+        cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url)
+        try:
+            with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
+                conn.execute("SET statement_timeout = 0")
+                conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
+                found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
+                if not found:
+                    return True
+                logger.warning(
+                    "Found %d invalid index(es) left by an interrupted CREATE INDEX "
+                    "CONCURRENTLY, rebuilding: %s",
+                    len(found),
+                    ", ".join(f"{index.name} (table size {index.table_size})" for index in found),
+                )
+                lock_row: Final = conn.execute(
+                    "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
+                ).fetchone()
+                if lock_row is None or not lock_row[0]:
+                    logger.info("Another replica is already rebuilding the invalid indexes, skipping")
+                    return False
+                for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
+                    ProxyExtrasDBManager._repair_index(conn, index)
+                remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
+        except psycopg.Error as e:
+            logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
+            return False
+        return not remaining
+
     @staticmethod
     def _setup_database_v2(use_migrate: bool) -> bool:
         """
@@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
             use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
         )
         if migrated:
+            ProxyExtrasDBManager.repair_invalid_indexes()
             ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
         return migrated
 

litellm/batches/main.py

@@ -319,6 +319,7 @@ def create_batch(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 create_batch_data=_create_batch_request,
+                custom_endpoint=optional_params.get("custom_endpoint"),
             )
         else:
             raise litellm.exceptions.BadRequestError(

litellm/caching/redis_cache.py

@@ -20,6 +20,8 @@ from contextvars import ContextVar
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import print_verbose, verbose_logger
 from litellm.constants import (
@@ -80,11 +82,29 @@ class _AsyncRedisCommands(Protocol):
 
     def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
 
+    def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
+
 
 _BREAKER_GUARD_FRAME_NAMES: Final = frozenset(
     {"<lambda>", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"}
 )
 
+_INCREMENT_WITH_FLOOR_LUA: Final = (
+    "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) "
+    "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end "
+    "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
+    "return count"
+)
+
+_LUA_COUNT: Final = TypeAdapter(int)
+_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
+
+
+def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]:
+    return _OPTIONAL_COUNTS.validate_python(
+        tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values)
+    )
+
 
 def _get_call_stack_info(num_frames: int = 2) -> str:
     """
@@ -736,6 +756,43 @@ class RedisCache(BaseCache):
             )
             raise e
 
+    @_redis_circuit_breaker_guard_sync
+    def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
+        """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call.
+
+        A counter whose key expired while a request was still in flight would otherwise be
+        recreated negative by that request's decrement. Clamping inside the same call is what
+        keeps it safe: a separate corrective write could land after another pod's increment and
+        erase it.
+
+        The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was
+        created rather than ``ttl`` after it was last touched. Refreshing it on every touch
+        would keep a count a dead worker never decremented alive for as long as the group
+        takes traffic. Returns the resulting count.
+        """
+        namespaced_key: Final = self.check_and_fix_namespace(key=key)
+        count: Final[object] = self.redis_client.eval(  # pyright: ignore[reportAttributeAccessIssue]  # stubs omit eval
+            _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl
+        )
+        return _LUA_COUNT.validate_python(count)
+
+    @_redis_circuit_breaker_guard_sync
+    def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
+        """Read integer counters for ``key_list``, in order, raising when Redis cannot answer.
+
+        ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller
+        cannot tell apart from "every counter is unset". A caller that has to fall back to its
+        own numbers when Redis is unreachable needs the failure, not a dict of zeros.
+        """
+        namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
+        return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys))
+
+    @_redis_circuit_breaker_guard
+    async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
+        """Async twin of ``batch_get_counts``, raising on failure the same way."""
+        namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
+        return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys))
+
     @_redis_circuit_breaker_guard
     async def async_scan_iter(self, pattern: str, count: int = 100) -> list:
         start_time: Final = time.time()
@@ -1241,6 +1298,14 @@ class RedisCache(BaseCache):
             result = result.decode()
         return float(result)
 
+    @_redis_circuit_breaker_guard
+    async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int:
+        """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees."""
+        _redis_client: Final = self._async_commands()
+        namespaced_key: Final = self.check_and_fix_namespace(key=key)
+        count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
+        return _LUA_COUNT.validate_python(count)
+
     async def flush_cache_buffer(self):
         print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
         await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

litellm/constants.py

@@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
 DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
+budget_reservation_disabled_info_emitted = False
 DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
 DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
 SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
@@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096))
 DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3))
 DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1))
 DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5))
+DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float(
+    os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1")
+)
 DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
 DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
 DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))

litellm/exceptions.py

@@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError):
         num_retries: int | None = None,
         headers: dict | None = None,
         exception_status_code: int | None = None,
+        response: httpx.Response | None = None,
     ):
         request: Final = httpx.Request(
             method="POST",
@@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError):
         self.max_retries = max_retries
         self.num_retries = num_retries
         self.headers = headers
+        if response is not None:
+            self.response = response
 
     # custom function to convert to str
     def __str__(self):

litellm/integrations/anthropic_cache_control_hook.py

@@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
 
+from pydantic import TypeAdapter, ValidationError
+
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.custom_prompt_management import CustomPromptManagement
@@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     with_prompt_cache_breakpoint,
 )
+from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
 from litellm.types.integrations.anthropic_cache_control_hook import (
     GATEWAY_INJECTED_CACHE_METADATA_KEY,
     GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
 )
 OPENAI_API_HOST: Final = "api.openai.com"
 OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
+_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
+_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
 
 AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
 
 
+def _validated_object_mapping(value: object) -> dict[object, object] | None:
+    try:
+        return _OBJECT_MAPPING_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _validated_object_list(value: object) -> list[object] | None:
+    try:
+        return _OBJECT_LIST_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
 def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
     model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
     if model_map_flag is not None:
@@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_
 
 
 class AnthropicCacheControlHook(CustomPromptManagement):
+    @staticmethod
+    def _request_value(request_kwargs: object, key: str) -> object:
+        request_mapping: Final = _validated_object_mapping(request_kwargs)
+        if request_mapping is None:
+            return None
+        return request_mapping.get(key)
+
+    @staticmethod
+    def _request_user_agent(request_kwargs: object) -> str | None:
+        proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request")
+        proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request)
+        if proxy_server_request_mapping is None:
+            return None
+        headers: Final = proxy_server_request_mapping.get("headers")
+        headers_mapping: Final = _validated_object_mapping(headers)
+        if headers_mapping is None:
+            return None
+        user_agent: Final = next(
+            (value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"),
+            None,
+        )
+        return user_agent if isinstance(user_agent, str) else None
+
+    @staticmethod
+    def _request_system(request_kwargs: object) -> str | list[object] | None:
+        system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system")
+        if isinstance(system, str):
+            return system
+        return _validated_object_list(system)
+
     def get_chat_completion_prompt(
         self,
         model: str,
@@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         points: Sequence[CacheControlInjectionPoint],
         messages: list[AllMessageValues],
         tools: list[object] | None,
+        cache_control: object,
         model: str,
         custom_llm_provider: str | None,
         api_base: object,
         prompt_cache_options: object,
     ) -> Sequence[Mapping[str, object]] | None:
-        if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
+        if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
             return None
         return AnthropicCacheControlHook._stamped_with_dialect(
             points, model, custom_llm_provider, api_base, prompt_cache_options
@@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         messages: list[AllMessageValues],
         system: str | list | None,
         tools: list | None,
+        cache_control: object = None,
     ) -> bool:
         """Whether configured injection points must yield to client-set cache_control.
 
@@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         """
         if all(point.get("_litellm_judged") for point in points):
             return False
-        return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
+        return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
 
     @staticmethod
     def _request_has_cache_control(
         messages: list[AllMessageValues],
         system: str | list | None,
         tools: list | None = None,
+        cache_control: object = None,
     ) -> bool:
         """Return True if the request already carries any client-supplied cache_control.
 
@@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         carry the mark either at the top level (Anthropic shape) or nested under
         ``function`` (OpenAI shape); the Anthropic chat transform accepts both.
         """
+        if cache_control is not None:
+            return True
         if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
             return True
         if tools is not None:
@@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         custom_llm_provider: str | None,
         tools: list | None = None,
         enable_prompt_caching: bool | None = None,
+        cache_control: object = None,
+        request_kwargs: object = None,
     ) -> list[CacheControlInjectionPoint]:
         """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
 
@@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         if not supports_prompt_caching(model=model, custom_llm_provider=provider):
             return []
 
-        if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
+        if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
+            return []
+
+        if is_claude_code_one_shot_subagent_request(
+            messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs)
+        ):
             return []
 
         control: Final = AnthropicCacheControlHook._default_control()
@@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         models: Iterable[str],
         tools: list[AllToolParamValues] | None = None,
         enable_prompt_caching: bool | None = None,
+        request_kwargs: object = None,
     ) -> list[AllMessageValues]:
         """Return the messages auto prompt caching will send, default breakpoints included.
 
@@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 for candidate in (
                     AnthropicCacheControlHook.get_default_injection_points(
                         messages=messages,
-                        system=None,
                         model=model,
                         custom_llm_provider=None,
                         tools=tools,
                         enable_prompt_caching=enable_prompt_caching,
+                        system=AnthropicCacheControlHook._request_system(request_kwargs),
+                        cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"),
+                        request_kwargs=request_kwargs,
                     )
                     for model in models
                 )
@@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 non_default_params["cache_control_injection_points"],
                 messages,
                 tools,
+                non_default_params.get("cache_control"),
                 model,
                 custom_llm_provider,
                 api_base,
@@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
             custom_llm_provider=custom_llm_provider,
             tools=tools,
             enable_prompt_caching=enable_prompt_caching,
+            cache_control=non_default_params.get("cache_control"),
+            request_kwargs=non_default_params,
         )
         if points:
             non_default_params["cache_control_injection_points"] = points
@@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         enable_prompt_caching: Final = cast(  # cast-ok: kwargs is untyped; key stamped as bool by the proxy
             bool | None, kwargs.pop("enable_prompt_caching", None)
         )
+        cache_control: Final = kwargs.get("cache_control")
         configured: Final = cast(  # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
             list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
         )
-        if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
+        if configured and AnthropicCacheControlHook._should_stand_down(
+            configured, typed_messages, system, tools, cache_control
+        ):
             return messages, system
         injection_points: list[CacheControlInjectionPoint] = configured or []
         if not injection_points and model is not None:
@@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 model=model,
                 custom_llm_provider=custom_llm_provider,
                 enable_prompt_caching=enable_prompt_caching,
+                cache_control=cache_control,
+                request_kwargs=kwargs,
             )
         if not injection_points:
             return messages, system

litellm/integrations/custom_guardrail.py

@@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger):
         event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None,
         supported_event_hooks: list[GuardrailEventHooks],
     ) -> None:
+        allowed_hooks: Final = frozenset(supported_event_hooks) | (
+            frozenset((GuardrailEventHooks.logging_only,))
+            if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
+            else frozenset()
+        )
+
         def _validate_event_hook_list_is_in_supported_event_hooks(
             event_hook: list[GuardrailEventHooks] | list[str],
             supported_event_hooks: list[GuardrailEventHooks],
@@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger):
             for hook in event_hook:
                 if isinstance(hook, str):
                     hook = GuardrailEventHooks(hook)
-                if hook not in supported_event_hooks:
+                if hook not in allowed_hooks:
                     raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}")
 
         if event_hook is None:
@@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger):
                 default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default]
                 _validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks)
         elif isinstance(event_hook, GuardrailEventHooks):
-            if event_hook not in supported_event_hooks:
+            if event_hook not in allowed_hooks:
                 raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}")
 
     @staticmethod
@@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger):
     def uses_apply_guardrail_interface(self) -> bool:
         return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
 
-    def _deployment_pre_call_target(self) -> "CustomLogger":
+    def _deployment_hook_target(self) -> "CustomLogger":
         if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
             return self
         try:
@@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger):
 
         # CHECK IF GUARDRAIL REJECTS THE REQUEST
         if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
-            target: Final = self._deployment_pre_call_target()
+            target: Final = self._deployment_hook_target()
             if target is not self:
                 kwargs["guardrail_to_apply"] = self
             result: Final = await target.async_pre_call_hook(
@@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger):
             return None
 
         # CHECK IF GUARDRAIL REJECTS THE REQUEST
-        result: Final = await self.async_post_call_success_hook(
+        target: Final = self._deployment_hook_target()
+        hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
+        result: Final = await target.async_post_call_success_hook(
             user_api_key_dict=UserAPIKeyAuth(
                 user_id=request_data.get("user_api_key_user_id"),
                 team_id=request_data.get("user_api_key_team_id"),
@@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger):
                 api_key=request_data.get("user_api_key_hash"),
                 request_route=request_data.get("user_api_key_request_route"),
             ),
-            data=request_data,
+            data=hook_request_data,
             response=response,
         )
 

litellm/litellm_core_utils/credential_accessor.py

@@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem
 
 
 class CredentialAccessor:
+    @staticmethod
+    def find_credential(credential_name: str) -> CredentialItem | None:
+        return next(
+            (credential for credential in litellm.credential_list if credential.credential_name == credential_name),
+            None,
+        )
+
     @staticmethod
     def get_credential_values(credential_name: str) -> dict:
         """Safe accessor for credentials."""
 
-        if not litellm.credential_list:
-            return {}
-        for credential in litellm.credential_list:
-            if credential.credential_name == credential_name:
-                return credential.credential_values.copy()
-        return {}
+        credential: Final = CredentialAccessor.find_credential(credential_name)
+        return {} if credential is None else credential.credential_values.copy()
 
     @staticmethod
     def upsert_credentials(credentials: list[CredentialItem]):

litellm/litellm_core_utils/exception_mapping_utils.py

@@ -860,6 +860,7 @@ def _map_bedrock_exception(
                 message=mantle_context_window_message,
                 model=model,
                 llm_provider=custom_llm_provider,
+                response=getattr(original_exception, "response", None),
             )
     if (
         "too many tokens" in error_str
@@ -873,6 +874,7 @@ def _map_bedrock_exception(
             message=f"BedrockException: Context Window Error - {error_str}",
             model=model,
             llm_provider="bedrock",
+            response=getattr(original_exception, "response", None),
         )
     elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str:
         raise BadRequestError(
@@ -924,12 +926,14 @@ def _map_bedrock_exception(
             message=f"BedrockException: Timeout Error - {error_str}",
             model=model,
             llm_provider="bedrock",
+            response=getattr(original_exception, "response", None),
         )
     elif "Could not process image" in error_str:
         raise litellm.InternalServerError(
             message=f"BedrockException - {error_str}",
             model=model,
             llm_provider="bedrock",
+            response=getattr(original_exception, "response", None),
         )
     elif hasattr(original_exception, "status_code"):
         if original_exception.status_code == 500:
@@ -937,10 +941,7 @@ def _map_bedrock_exception(
                 message=f"BedrockException - {original_exception.message}",
                 llm_provider="bedrock",
                 model=model,
-                response=httpx.Response(
-                    status_code=500,
-                    request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
-                ),
+                response=getattr(original_exception, "response", None),
             )
         elif original_exception.status_code == 401:
             raise AuthenticationError(
@@ -969,6 +970,7 @@ def _map_bedrock_exception(
                 model=model,
                 llm_provider=custom_llm_provider,
                 litellm_debug_info=extra_information,
+                response=getattr(original_exception, "response", None),
             )
         elif original_exception.status_code == 422:
             raise BadRequestError(
@@ -1001,6 +1003,7 @@ def _map_bedrock_exception(
                 llm_provider=custom_llm_provider,
                 litellm_debug_info=extra_information,
                 exception_status_code=original_exception.status_code,
+                response=getattr(original_exception, "response", None),
             )
 
 

litellm/litellm_core_utils/llm_response_utils/get_headers.py

@@ -1,7 +1,8 @@
+from collections.abc import Mapping
 from typing import Final
 
 
-def get_response_headers(_response_headers: dict | None = None) -> dict:
+def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict:
     """
 
     Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header}
@@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict:
     return {**llm_provider_headers, **openai_headers}
 
 
-def _get_llm_provider_headers(response_headers: dict) -> dict:
+def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict:
     """
     Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider
 

litellm/llms/anthropic/common_utils.py

@@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
 _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
 _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
 _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
+_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
+_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
+_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
+
+
+def is_claude_code_user_agent(user_agent: str) -> bool:
+    return user_agent.startswith("claude-cli/")
+
+
+def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
+    try:
+        return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _validated_claude_code_list(value: object) -> list[object] | None:
+    try:
+        return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None:
+    stripped: Final = text.strip()
+    if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX):
+        return None
+    fields: Final = tuple(
+        field
+        for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";")
+        if (field := raw_field.strip())
+    )
+    if not fields or any("=" not in field for field in fields):
+        return None
+    parsed_fields: Final = tuple(
+        (parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),)
+    )
+    if any(not key or not value for key, value in parsed_fields):
+        return None
+    return parsed_fields
+
+
+def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
+    if isinstance(system, str):
+        return (system,)
+    blocks: Final = _validated_claude_code_list(system)
+    if blocks is None:
+        return None
+    block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
+    if any(block is None for block in block_mappings):
+        return None
+    text_values: Final = tuple(
+        block.get("text") for block in block_mappings if block is not None and block.get("type") == "text"
+    )
+    if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values):
+        return None
+    meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip())
+    return meaningful_text or None
+
+
+def _is_claude_code_subagent_billing_system(system: object) -> bool:
+    billing_texts: Final = _claude_code_billing_texts(system)
+    if billing_texts is None:
+        return False
+    billing_fields: Final = tuple(
+        fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None
+    )
+    if len(billing_fields) != len(billing_texts):
+        return False
+    subagent_values: Final = tuple(
+        value for fields in billing_fields for key, value in fields if key == "cc_is_subagent"
+    )
+    return subagent_values == ("true",)
+
+
+def is_claude_code_one_shot_subagent_request(
+    messages: list[AllMessageValues],
+    system: object,
+    tools: object,
+    user_agent: str | None,
+) -> bool:
+    only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None
+    return (
+        user_agent is not None
+        and is_claude_code_user_agent(user_agent)
+        and not tools
+        and only_message is not None
+        and only_message.get("role") == "user"
+        and _is_claude_code_subagent_billing_system(system)
+    )
 
 
 def _strip_bedrock_id_suffixes(model: str) -> str:

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

@@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp(
         LiteLLM_Proxy_MCP_Handler,
     )
 
-    mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+    mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
     if not mcp_references:
         return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(

litellm/llms/bedrock/chat/agentcore/transformation.py

@@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=str(response.read()))
+            raise BedrockError(
+                status_code=response.status_code,
+                message=str(response.read()),
+                headers=response.headers,
+                response=response,
+            )
 
         # LOGGING
         logging_obj.post_call(
@@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
                 raise BedrockError(
                     status_code=response.status_code,
                     message=f"AgentCore: Failed to read/parse JSON response body: {e}",
+                    headers=response.headers,
                 )
             parsed: Final = self._parse_json_response(response_json)
 
@@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=str(await response.aread()))
+            raise BedrockError(
+                status_code=response.status_code,
+                message=str(await response.aread()),
+                headers=response.headers,
+                response=response,
+            )
 
         # LOGGING
         logging_obj.post_call(
@@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
                 raise BedrockError(
                     status_code=response.status_code,
                     message=f"AgentCore: Failed to read/parse JSON response body: {e}",
+                    headers=response.headers,
                 )
             parsed: Final = self._parse_json_response(response_json)
 
@@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error processing response: {e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
     def validate_environment(
@@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         return headers
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
 
     def should_fake_stream(
         self,

litellm/llms/bedrock/chat/converse_handler.py

@@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse
 from litellm.utils import CustomStreamWrapper
 
 from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
-from ..common_utils import BedrockError, _get_all_bedrock_regions
+from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
 from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
 
 
@@ -66,7 +66,12 @@ def make_sync_call(
     )
 
     if response.status_code != 200:
-        raise BedrockError(status_code=response.status_code, message=str(response.read()))
+        raise BedrockError(
+            status_code=response.status_code,
+            message=str(response.read()),
+            headers=response.headers,
+            response=response,
+        )
 
     if fake_stream:
         model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@@ -247,7 +252,12 @@ class BedrockConverseLLM(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=error_response_text(err.response),
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -594,7 +604,12 @@ class BedrockConverseLLM(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=error_response_text(err.response),
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 

litellm/llms/bedrock/chat/converse_transformation.py

@@ -2255,6 +2255,7 @@ class AmazonConverseConfig(BaseConfig):
             raise BedrockError(
                 message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues",
                 status_code=422,
+                headers=response.headers,
             )
 
         """

litellm/llms/bedrock/chat/invoke_agent/transformation.py

@@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error processing response: {e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
     def validate_environment(
@@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
         return headers
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
 
     def should_fake_stream(
         self,

litellm/llms/bedrock/chat/invoke_handler.py

@@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk
 from ..common_utils import (
     BedrockError,
     build_bedrock_stream_error,
+    error_response_text,
     get_bedrock_response_stream_shape,
     get_bedrock_tool_name,
 )
@@ -184,7 +185,12 @@ async def make_call(
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=response.text)
+            raise BedrockError(
+                status_code=response.status_code,
+                message=error_response_text(response),
+                headers=response.headers,
+                response=response,
+            )
 
         if fake_stream:
             model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@@ -228,9 +234,16 @@ async def make_call(
         )
 
         return completion_stream, response.headers
+    except BedrockError:
+        raise
     except httpx.HTTPStatusError as err:
         error_code: Final = err.response.status_code
-        raise BedrockError(status_code=error_code, message=err.response.text)
+        raise BedrockError(
+            status_code=error_code,
+            message=error_response_text(err.response),
+            headers=err.response.headers,
+            response=err.response,
+        )
     except httpx.TimeoutException:
         raise BedrockError(status_code=408, message="Timeout error occurred.")
     except Exception as e:
@@ -270,7 +283,12 @@ def make_sync_call(
         )
 
         if response.status_code != 200:
-            raise BedrockError(status_code=response.status_code, message=response.text)
+            raise BedrockError(
+                status_code=response.status_code,
+                message=error_response_text(response),
+                headers=response.headers,
+                response=response,
+            )
 
         if fake_stream:
             model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@@ -314,9 +332,16 @@ def make_sync_call(
         )
 
         return completion_stream, response.headers
+    except BedrockError:
+        raise
     except httpx.HTTPStatusError as err:
         error_code: Final = err.response.status_code
-        raise BedrockError(status_code=error_code, message=err.response.text)
+        raise BedrockError(
+            status_code=error_code,
+            message=error_response_text(err.response),
+            headers=err.response.headers,
+            response=err.response,
+        )
     except httpx.TimeoutException:
         raise BedrockError(status_code=408, message="Timeout error occurred.")
     except Exception as e:

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

@@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
         """Return the appropriate error class for Bedrock."""
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)

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

@@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
         """Return the appropriate error class for Bedrock."""
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)

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

@@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
             raise BedrockError(
                 message=f"Error parsing response: {raw_response.text}, error: {e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
         verbose_logger.debug(
@@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
             raise BedrockError(
                 message=f"Error setting response content: {e}. Response: {completion_response}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
         # Calculate usage from headers

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

@@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
         try:
             completion_response: Final = raw_response.json()
         except Exception:
-            raise BedrockError(message=raw_response.text, status_code=raw_response.status_code)
+            raise BedrockError(
+                message=raw_response.text,
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
         verbose_logger.debug(
             "bedrock invoke response % s",
             json.dumps(completion_response, indent=4, default=str),
@@ -363,6 +367,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error processing={raw_response.text}, Received error={e}",
                 status_code=422,
+                headers=raw_response.headers,
             )
 
         try:
@@ -384,6 +389,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
             raise BedrockError(
                 message=f"Error parsing received text={outputText}.\nError-{e}",
                 status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
 
         ## CALCULATING USAGE - bedrock returns usage in the headers
@@ -431,7 +437,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
         return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)
 
     def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
-        return BedrockError(status_code=status_code, message=error_message)
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
 
     @track_llm_api_timing()
     async def get_async_custom_stream_wrapper(

litellm/llms/bedrock/claude_platform/common_utils.py

@@ -1,7 +1,10 @@
 from typing import Final
 
+import httpx
+
 import litellm
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.secret_managers.main import get_secret_str
 
 CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic"
@@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str:
 
 
 class BedrockClaudePlatformMixin(BaseAWSLLM):
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     @staticmethod
     def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None:
         workspace_id = (

litellm/llms/bedrock/common_utils.py

@@ -33,8 +33,53 @@ if TYPE_CHECKING:
     from litellm.types.llms.openai import AllMessageValues
 
 
+_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"
+
+
+def error_response_text(response: httpx.Response) -> str:
+    try:
+        return response.text
+    except httpx.ResponseNotRead:
+        return response.reason_phrase
+
+
+def _synthesize_error_response(
+    *, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None
+) -> tuple[httpx.Request, httpx.Response]:
+    error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL)
+    safe_headers: Final = (
+        headers
+        if isinstance(headers, httpx.Headers)
+        else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes)))
+    )
+    return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request)
+
+
 class BedrockError(BaseLLMException):
-    pass
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        headers: dict[str, object] | httpx.Headers | None = None,
+        request: httpx.Request | None = None,
+        response: httpx.Response | None = None,
+        body: dict[str, object] | None = None,
+        status_code_is_synthesized: bool = False,
+    ) -> None:
+        error_request, error_response = (
+            _synthesize_error_response(status_code=status_code, headers=headers, request=request)
+            if response is None and headers
+            else (request, response)
+        )
+        super().__init__(
+            status_code=status_code,
+            message=message,
+            headers=headers,
+            request=error_request,
+            response=error_response,
+            body=body,
+            status_code_is_synthesized=status_code_is_synthesized,
+        )
 
 
 _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (

litellm/llms/bedrock/count_tokens/handler.py

@@ -102,6 +102,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
                 raise BedrockError(
                     status_code=response.status_code,
                     message=error_text,
+                    headers=response.headers,
+                    response=response,
                 )
 
             bedrock_response: Final = response.json()
@@ -124,6 +126,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
             raise BedrockError(
                 status_code=e.response.status_code,
                 message=e.response.text,
+                headers=e.response.headers,
+                response=e.response,
             )
         except Exception as e:
             verbose_logger.error("Error in CountTokens handler: %s", e)

litellm/llms/bedrock/embed/embedding.py

@@ -132,7 +132,12 @@ class BedrockEmbedding(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -161,7 +166,12 @@ class BedrockEmbedding(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 

litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py

@@ -20,6 +20,7 @@ import httpx
 
 from litellm._logging import verbose_logger
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.images.main import ImageEditOptionalRequestParams
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import FileTypes, ImageObject, ImageResponse
@@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
         """
         return _supports_nova_canvas_image_edit_from_model_cost(model or "")
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_supported_openai_params(self, model: str) -> list:
         return [
             "n",

litellm/llms/bedrock/image_edit/handler.py

@@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 

litellm/llms/bedrock/image_edit/stability_transformation.py

@@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final
 import httpx
 
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.images.main import ImageEditOptionalRequestParams
 from litellm.types.llms.stability import (
     OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO,
@@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
                 return True
         return False
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_supported_openai_params(self, model: str) -> list:
         """
         Return list of OpenAI params supported by Bedrock Stability.

litellm/llms/bedrock/image_generation/image_handler.py

@@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
         ### FORMAT RESPONSE TO OPENAI FORMAT ###
@@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 

litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py

@@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
     AmazonInvokeConfig,
 )
 from litellm.llms.bedrock.common_utils import (
+    BedrockError,
     apply_bedrock_invoke_structured_output,
     ensure_bedrock_anthropic_messages_tool_names,
     get_anthropic_beta_from_headers,
@@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig(
 
     BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys())
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def __init__(self, **kwargs):
         BaseAnthropicMessagesConfig.__init__(self, **kwargs)
         AmazonInvokeConfig.__init__(self, **kwargs)

litellm/llms/bedrock/passthrough/transformation.py

@@ -2,13 +2,14 @@ import json
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final, Optional, cast
 
+import httpx
 from httpx import Response
 
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
 
 from ..base_aws_llm import BaseAWSLLM
-from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo
+from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo
 
 if TYPE_CHECKING:
     from httpx import URL
@@ -18,6 +19,14 @@ if TYPE_CHECKING:
 
 
 class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig):
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
         return "stream" in endpoint
 

litellm/llms/bedrock/realtime/transformation.py

@@ -9,12 +9,14 @@ import json
 import uuid as uuid_lib
 from typing import Final, cast
 
+import httpx
 from pydantic import BaseModel
 
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm
 from litellm.types.llms.openai import (
     OpenAIRealtimeContentPartDone,
@@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
         self._cumulative_usage = BedrockUsageEvent()
         self._reported_usage = BedrockUsageEvent()
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict:
         """Validate environment - no special validation needed for Bedrock."""
         return headers

litellm/llms/bedrock/rerank/handler.py

@@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 
@@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM):
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
             error_code: Final = err.response.status_code
-            raise BedrockError(status_code=error_code, message=err.response.text)
+            raise BedrockError(
+                status_code=error_code,
+                message=err.response.text,
+                headers=err.response.headers,
+                response=err.response,
+            )
         except httpx.TimeoutException:
             raise BedrockError(status_code=408, message="Timeout error occurred.")
 

litellm/llms/bedrock/search/transformation.py

@@ -39,7 +39,6 @@ from typing import Final
 import httpx
 
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.search.transformation import (
     BaseSearchConfig,
     SearchResponse,
@@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
             raise BedrockError(
                 status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
                 message=f"AgentCore gateway MCP error: {error}",
+                headers=raw_response.headers,
             )
 
         # A failed tools/call is reported in-band, as HTTP 200 with result.isError
@@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
             raise BedrockError(
                 status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
                 message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}",
+                headers=raw_response.headers,
             )
 
         text_items: Final = tuple(
@@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
         raise BedrockError(
             status_code=502,
             message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}",
+            headers=raw_response.headers,
         )
 
     def get_error_class(
@@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
         status_code: int,
         headers: dict,  # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict
     ) -> Exception:
-        return BaseLLMException(
+        return BedrockError(
             status_code=status_code,
             message=error_message,
             headers=headers,

litellm/llms/bedrock/vector_stores/transformation.py

@@ -8,6 +8,7 @@ from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.types.integrations.rag.bedrock_knowledgebase import (
     BedrockKBContent,
     BedrockKBResponse,
@@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
         BaseVectorStoreConfig.__init__(self)
         BaseAWSLLM.__init__(self)
 
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | httpx.Headers,  # mutable-ok: base passes response headers as a dict
+    ) -> BedrockError:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials:
         return {}
 

litellm/llms/bedrock_mantle/chat/transformation.py

@@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the
 from collections.abc import AsyncIterator, Iterator
 from typing import Any, Final
 
+import httpx
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.router import GenericLiteLLMParams
 
+from ...base_llm.chat.transformation import BaseLLMException
+from ...bedrock.common_utils import BedrockError
 from ...openai_like.chat.transformation import OpenAILikeChatConfig
 from ..common_utils import mantle_base_segment
 
@@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
     def get_config(cls):
         return super().get_config()
 
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
+    ) -> BaseLLMException:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def _get_openai_compatible_provider_info(
         self,
         api_base: str | None,

litellm/llms/bedrock_mantle/responses/transformation.py

@@ -19,11 +19,14 @@ import json
 from collections.abc import Mapping
 from typing import Any, Final
 
+import httpx
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
+from litellm.llms.bedrock.common_utils import BedrockError
 from litellm.llms.bedrock_mantle.common_utils import (
     MANTLE_HOST_RE,
     BedrockMantleAuthMixin,
@@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
     def custom_llm_provider(self) -> LlmProviders:
         return LlmProviders.BEDROCK_MANTLE
 
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
+    ) -> BaseLLMException:
+        return BedrockError(status_code=status_code, message=error_message, headers=headers)
+
     def get_complete_url(
         self,
         api_base: str | None,

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler:
         except Exception as e:
             raise self._handle_error(e=e, provider_config=provider_config)
 
+        logging_obj.post_call(
+            api_key=api_key,
+            original_response=response.text,
+            additional_args={"complete_input_dict": data},
+        )
+
         return self._transform_ocr_response(
             provider_config=provider_config,
             model=model,
@@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler:
         except Exception as e:
             raise self._handle_error(e=e, provider_config=provider_config)
 
+        logging_obj.post_call(
+            api_key=api_key,
+            original_response=response.text,
+            additional_args={"complete_input_dict": data},
+        )
+
         # Use async response transform for async operations
         return await provider_config.async_transform_ocr_response(
             model=model,

litellm/llms/vertex_ai/batches/handler.py

@@ -1,6 +1,7 @@
 import json
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Sequence
 from typing import TYPE_CHECKING, Final, Protocol
+from urllib.parse import urlparse
 
 import httpx
 from typing_extensions import ReadOnly, TypedDict
@@ -12,11 +13,13 @@ from litellm.litellm_core_utils.url_utils import (
     safe_get,
 )
 from litellm.llms.custom_httpx.http_handler import (
+    HTTPHandler,
     _get_httpx_client,
     get_async_httpx_client,
 )
 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.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path
 from litellm.types.llms.openai import CreateBatchRequest
 from litellm.types.llms.vertex_ai import (
     VERTEX_CREDENTIALS_TYPES,
@@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict):
     response: ReadOnly[httpx.Response]
 
 
+class _VertexEndpointDeployedModel(TypedDict, total=False):
+    model: ReadOnly[str]
+
+
+class _VertexEndpointResponse(TypedDict, total=False):
+    deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]]
+
+
+class _VertexEndpointPayloadView(TypedDict):
+    """Holds one decoded GET endpoints/<id> response so the payload reads back typed."""
+
+    payload: ReadOnly[_VertexEndpointResponse]
+
+
 def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse:
     return response.json()
 
@@ -78,7 +95,17 @@ class VertexAIBatchPrediction(VertexLLM):
         vertex_location: str | None,
         timeout: float | httpx.Timeout,
         max_retries: int | None,
+        custom_endpoint: bool | None = None,
     ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]:
+        if custom_endpoint:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
+                    "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
+                    "use a publisher model or fine-tuned Gemini endpoint deployment instead."
+                ),
+            )
         sync_handler: Final = _get_httpx_client()
 
         access_token, project_id = self._ensure_access_token(
@@ -87,6 +114,26 @@ class VertexAIBatchPrediction(VertexLLM):
             custom_llm_provider="vertex_ai",
         )
 
+        headers: Final = {
+            "Content-Type": "application/json; charset=utf-8",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+        transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
+            VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
+                request=create_batch_data,
+                vertex_project=vertex_project or project_id,
+                vertex_location=vertex_location or "us-central1",
+            )
+        )
+        vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
+            vertex_batch_request=transformed_batch_request,
+            headers=headers,
+            sync_handler=sync_handler,
+            api_base=api_base,
+            vertex_location=vertex_location or "us-central1",
+        )
+
         default_api_base: Final = self.create_vertex_batch_url(
             vertex_location=vertex_location or "us-central1",
             vertex_project=vertex_project or project_id,
@@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
-        headers: Final = {
-            "Content-Type": "application/json; charset=utf-8",
-            "Authorization": f"Bearer {access_token}",
-        }
-
-        vertex_batch_request: Final[VertexAIBatchPredictionJob] = (
-            VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
-                request=create_batch_data
-            )
-        )
-
         if _is_async is True:
             return self._async_create_batch(
                 vertex_batch_request=vertex_batch_request,
@@ -142,6 +178,77 @@ class VertexAIBatchPrediction(VertexLLM):
         )
         return vertex_batch_response
 
+    @staticmethod
+    def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str:
+        """
+        Builds the GET url for resolving an endpoint resource (`projects/../endpoints/<id>`).
+
+        A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the
+        version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a
+        mount prefix in front of the full default path. The `:operation` suffix convention from
+        `_check_custom_proxy` does not apply to a plain resource GET.
+        """
+        default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        if not api_base:
+            return default_endpoint_url
+        api_base_path: Final = urlparse(api_base).path.rstrip("/")
+        if api_base_path in ("/v1", "/v1beta1"):
+            return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url)
+        return api_base.rstrip("/") + urlparse(default_endpoint_url).path
+
+    def _resolve_fine_tuned_endpoint_model(
+        self,
+        vertex_batch_request: VertexAIBatchPredictionJob,
+        headers: dict[str, str],  # mutable-ok: HTTPHandler.get only accepts dict headers
+        sync_handler: HTTPHandler,
+        api_base: str | None,
+        vertex_location: str,
+    ) -> VertexAIBatchPredictionJob:
+        """
+        A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only
+        accepts Model resources, so swap the endpoint resource for its deployed tuned model
+        (`projects/../locations/../models/<id>`) read from GET endpoints/<id>.
+        """
+        model: Final = vertex_batch_request.get("model", "")
+        if "/endpoints/" not in model:
+            return vertex_batch_request
+
+        endpoint_url: Final = self._build_endpoint_resolution_url(
+            api_base=api_base,
+            model=model,
+            vertex_location=vertex_location,
+        )
+        # ``api_base`` can come from caller-supplied request kwargs, so wrap the
+        # fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata
+        # targets before the bearer token leaves the process (mirrors retrieve_batch).
+        fetched: Final[_FetchedResponseView] = {
+            "response": safe_get(
+                sync_handler,
+                endpoint_url,
+                headers=headers,
+            )
+        }
+        response: Final = fetched["response"]
+        if response.status_code != 200:
+            raise VertexAIError(
+                status_code=response.status_code,
+                message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}",
+            )
+
+        payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()}
+        deployed_models: Final = payload_view["payload"].get("deployedModels") or ()
+        deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else ""
+        if not deployed_model:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model "
+                    "resource to run batch predictions against"
+                ),
+            )
+        resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model}
+        return resolved_request
+
     async def _async_create_batch(
         self,
         vertex_batch_request: VertexAIBatchPredictionJob,

litellm/llms/vertex_ai/batches/transformation.py

@@ -22,6 +22,8 @@ class VertexAIBatchTransformation:
     def transform_openai_batch_request_to_vertex_ai_batch_request(
         cls,
         request: CreateBatchRequest,
+        vertex_project: str | None = None,
+        vertex_location: str | None = None,
     ) -> VertexAIBatchPredictionJob:
         """
         Transforms OpenAI Batch requests to Vertex AI Batch requests
@@ -31,7 +33,11 @@ class VertexAIBatchTransformation:
         if input_file_id is None:
             raise ValueError("input_file_id is required, but not provided")
         input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl")
-        model: Final[str] = cls._get_model_from_gcs_file(input_file_id)
+        model: Final[str] = cls._get_batch_job_model(
+            input_file_id=input_file_id,
+            vertex_project=vertex_project,
+            vertex_location=vertex_location,
+        )
         output_config: Final[OutputConfig] = OutputConfig(
             predictionsFormat="jsonl",
             gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)),
@@ -188,6 +194,33 @@ class VertexAIBatchTransformation:
         path_parts: Final = input_file_id.rsplit("/", 1)
         return path_parts[0]
 
+    @classmethod
+    def _get_batch_job_model(
+        cls,
+        input_file_id: str,
+        vertex_project: str | None,
+        vertex_location: str | None,
+    ) -> str:
+        """
+        Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or
+        the full `projects/../locations/../endpoints/<id>` resource name for a fine-tuned endpoint.
+
+        The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource
+        to its deployed tuned model (`projects/../locations/../models/<id>`) before sending the job.
+        """
+        parsed_model: Final = cls._get_model_from_gcs_file(input_file_id)
+        if not parsed_model.startswith("endpoints/"):
+            return parsed_model
+        if not vertex_project:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require "
+                    "`vertex_project` to build the endpoint resource name"
+                ),
+            )
+        return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}"
+
     @classmethod
     def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str:
         """
@@ -202,6 +235,9 @@ 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"
 
+        Fine-tuned Gemini endpoints are stored as `endpoints/<numeric id>` in the uri and returned
+        in that form.
+
         Raises a 400 `VertexAIError` when the uri carries no parseable model path.
         """
         model: Final = cls._parse_model_from_gcs_file(gcs_file_uri)
@@ -210,11 +246,13 @@ class VertexAIBatchTransformation:
                 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. "
+                    f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' or "
+                    "'endpoints/<numeric endpoint id>' 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>"
+                    "gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file> "
+                    "(or gs://<bucket>/<prefix>/endpoints/<numeric endpoint id>/<file> for fine-tuned models)"
                 ),
             )
         return model
@@ -222,18 +260,26 @@ class VertexAIBatchTransformation:
     @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.
+        Returns the `publishers/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
+        gcs uri, or None if the uri does not contain one.
+
+        A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is
+        used, so a user-configured bucket prefix that happens to contain `endpoints/<digits>` cannot
+        override the model path LiteLLM appended after it.
         """
-        _, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
-        if not separator:
-            return None
+        unquoted_uri: Final = unquote(gcs_file_uri)
+        _, separator, model_path = unquoted_uri.partition("publishers/")
+        if separator:
+            parts: Final = model_path.split("/")
+            if len(parts) >= 3 and parts[1] == "models" and parts[2]:
+                return f"publishers/{'/'.join(parts[:3])}"
 
-        parts: Final = model_path.split("/")
-        if len(parts) < 3 or parts[1] != "models" or not parts[2]:
-            return None
+        _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/")
+        endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else ""
+        if endpoint_id.isdigit():
+            return f"endpoints/{endpoint_id}"
 
-        return f"publishers/{'/'.join(parts[:3])}"
+        return None
 
     @classmethod
     def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:

litellm/llms/vertex_ai/common_utils.py

@@ -370,6 +370,19 @@ def get_vertex_base_model_name(model: str) -> str:
     return model
 
 
+def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None:
+    """
+    Fine-tuned Gemini deployments are addressed by a numeric endpoint id,
+    configured as `vertex_ai/<id>` or `vertex_ai/gemini/<id>`.
+
+    Returns the endpoint id, or None when `model` is a regular publisher model.
+    Mirrors the online chat path in `_get_vertex_url`, which sends numeric
+    models to `endpoints/{id}` instead of `publishers/google/models/{model}`.
+    """
+    candidate: Final = model.split("/")[-1] if "gemini/" in model else model
+    return candidate if candidate.isdigit() else None
+
+
 def validate_vertex_location(vertex_location: str | None) -> str:
     """
     Validate a Vertex AI location before interpolating it into a request host or

litellm/llms/vertex_ai/files/transformation.py

@@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import (
 )
 from litellm.llms.vertex_ai.common_utils import (
     _convert_vertex_datetime_to_openai_datetime,
+    get_vertex_ai_fine_tuned_endpoint_id,
 )
 from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@@ -707,20 +708,39 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
     def _get_gcs_object_name_from_batch_jsonl(
         self,
         openai_jsonl_content: list[dict[str, Any]],
+        deployment_model: str | None = None,
     ) -> str:
         """
         Gets a unique GCS object name for the VertexAI batch prediction job
 
         named as: litellm-vertex-{model}-{uuid}
+
+        The stored model path decides which Vertex model the batch job later executes against, so
+        `deployment_model` (the deployment's own configured model) wins over the user-supplied
+        JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no
+        deployment config.
+
+        Fine-tuned Gemini deployments (numeric endpoint ids) are stored under
+        `endpoints/<id>` so the batch transformation can round-trip them into a
+        `projects/../locations/../endpoints/<id>` batch job model instead of a
+        nonexistent publisher model.
         """
-        _model = openai_jsonl_content[0].get("body", {}).get("model", "")
-        if "publishers/google/models" not in _model:
-            _model = f"publishers/google/models/{_model}"
-        safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model")
+        raw_model: Final = (
+            deployment_model.removeprefix("vertex_ai/")
+            if deployment_model
+            else openai_jsonl_content[0].get("body", {}).get("model", "")
+        )
+        endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
+        model_path: Final = (
+            f"endpoints/{endpoint_id}"
+            if endpoint_id is not None
+            else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}")
+        )
+        safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model")
         object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}"
         return object_name
 
-    def get_object_name(self, file_data: FileTypes, purpose: str) -> str:
+    def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str:
         """
         Get the object name for the request.
 
@@ -728,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         upload is never materialized just to derive the GCS object name.
         """
         if purpose == "batch":
-            ## 1. If jsonl, derive the object name from the first entry's model
+            ## 1. If jsonl, derive the object name from the deployment model (or the first entry's)
             first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None)
             if first_entry is not None:
-                return self._get_gcs_object_name_from_batch_jsonl([first_entry])
+                return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model)
 
         ## 2. If not jsonl, store under a server-generated managed object name
         filename, _ = extract_file_metadata(file_data)
@@ -761,6 +781,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         """
         Get the complete url for the request
         """
+        if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"):
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
+                    "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
+                    "remove this deployment from the batch request (e.g. `target_model_names`) or "
+                    "use a publisher model / fine-tuned Gemini endpoint instead."
+                ),
+            )
         bucket_name = self._get_configured_bucket_name(litellm_params)
         bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name)
         file_data: Final = data.get("file")
@@ -769,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
             raise ValueError("file is required")
         if purpose is None:
             raise ValueError("purpose is required")
-        object_name = self.get_object_name(file_data, purpose)
+        configured_model: Final = litellm_params.get("model")
+        object_name = self.get_object_name(
+            file_data,
+            purpose,
+            deployment_model=configured_model if isinstance(configured_model, str) else None,
+        )
         if object_prefix:
             object_name = f"{object_prefix}/{object_name}"
         encoded_object_name: Final = encode_gcs_object_name_for_url(object_name)

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -3108,15 +3108,17 @@ class MCPRequestHandler:
     @staticmethod
     async def _get_allowed_mcp_servers_for_agent(
         user_api_key_auth: UserAPIKeyAuth | None = None,
-        agent_object_permission=None,
+        agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
     ) -> list[str]:
         """
         Get allowed MCP servers for an agent (from the agent's object_permission).
 
-        Returns the MCP servers from the agent's object_permission.
-        If agent has no object_permission, returns [] (no extra restriction). An entitlement the
-        agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the
-        resolver denies.
+        Returns the agent's direct servers, the servers in its access groups, and the servers reached
+        through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no
+        object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that
+        cannot be read, or a declared toolset that resolves to no grants, raises
+        ``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the
+        agent as unrestricted.
 
         Args:
             user_api_key_auth: User auth with agent_id
@@ -3126,31 +3128,30 @@ class MCPRequestHandler:
         if not user_api_key_auth or not user_api_key_auth.agent_id:
             return []
 
-        obj_perm = agent_object_permission
-        if obj_perm is None:
-            obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        obj_perm: Final = (
+            agent_object_permission
+            if agent_object_permission is not None
+            else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        )
         if obj_perm is None:
             return []
 
         try:
-            direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or []
-            if isinstance(direct_mcp_servers, str):
-                direct_mcp_servers = []
-            mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or []
-            if isinstance(mcp_access_groups, str):
-                mcp_access_groups = []
-
-            # Permission entries may be server_ids OR names/aliases — expand to ids.
             from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
                 global_mcp_server_manager,
             )
 
-            expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers))
-
-            access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups)
-            all_servers: Final = expanded_direct_servers + access_group_servers
-            return list(set(all_servers))
+            expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(
+                obj_perm.mcp_servers or []
+            )
+            access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(
+                obj_perm.mcp_access_groups or []
+            )
+            toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm)
+            return list({*expanded_direct_servers, *access_group_servers, *toolset_grants})
         except Exception as e:
+            if isinstance(e, UnloadableEntitlementError):
+                raise
             verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
             return []
 
@@ -3158,13 +3159,15 @@ class MCPRequestHandler:
     async def _get_agent_tool_permissions_for_server(
         server_id: str,
         user_api_key_auth: UserAPIKeyAuth | None = None,
-        agent_object_permission=None,
+        agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
     ) -> list[str] | None:
         """
-        Get allowed tool names for a server from the agent's object_permission.
-        Returns None if agent has no tool restrictions for this server. An entitlement the agent
-        LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the
-        tool resolver turns into deny-all for the server rather than an unrestricted tool list.
+        Get allowed tool names for a server from the agent's object_permission: the union of its
+        direct tool permissions and the tools its toolsets grant on that server, mirroring the key and
+        team levels. Returns None if agent has no tool restrictions for this server. An entitlement the
+        agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises
+        ``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the
+        server rather than an unrestricted tool list.
 
         Args:
             server_id: Server ID to check permissions for
@@ -3175,24 +3178,30 @@ class MCPRequestHandler:
         if not user_api_key_auth or not user_api_key_auth.agent_id:
             return None
 
-        obj_perm = agent_object_permission
-        if obj_perm is None:
-            obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        obj_perm: Final = (
+            agent_object_permission
+            if agent_object_permission is not None
+            else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
+        )
         if obj_perm is None:
             return None
 
         try:
-            mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None)
-            if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict):
-                return None
-            # Dict keys may be server_ids OR names/aliases; normalize before lookup.
             from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
                 global_mcp_server_manager,
             )
 
-            tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id)
-            return list(tools) if tools else None
+            direct_tools: Final = (
+                global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
+                if obj_perm.mcp_tool_permissions
+                else None
+            )
+            toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id)
+            agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
+            return list(agent_tools) if agent_tools else None
         except Exception as e:
+            if isinstance(e, UnloadableEntitlementError):
+                raise
             verbose_logger.warning("Failed to get agent tool permissions for server: %s", e)
             return None
 

litellm/proxy/_experimental/mcp_server/db.py

@@ -4,7 +4,7 @@ import hashlib
 import json
 from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
 
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
@@ -125,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False):
     server_id: str
 
 
+OAuthGrantState = Literal["valid", "refreshable", "absent"]
+
+
 class _OAuthTokenRefreshResponse(TypedDict, total=False):
     access_token: str
     refresh_token: str
@@ -1465,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in
         return False
 
 
+def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState:
+    """Classify local grant readiness without attempting a refresh or checking upstream revocation."""
+    if not cred or not cred.get("access_token"):
+        return "absent"
+    if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
+        return "valid"
+    return "refreshable" if cred.get("refresh_token") else "absent"
+
+
 async def get_user_oauth_credential(
     prisma_client: PrismaClient,
     user_id: str,
@@ -1727,12 +1739,11 @@ async def resolve_valid_user_oauth_token(
     dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh
     actually happens, so the valid-token path never requires a DB handle.
     """
-    if not cred or not cred.get("access_token"):
+    grant: Final = oauth_grant_state(cred)
+    if cred is None or grant == "absent":
         return None
-    if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
+    if grant == "valid":
         return cred
-    if not cred.get("refresh_token"):
-        return None
     if prisma_client is None:
         from litellm.proxy.utils import get_prisma_client_or_throw
 

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import (
     render_token_fault,
 )
 from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
+    VendorCredentialState,
     aggregate_authorize,
     aggregate_token,
     complete_connect_flow,
+    describe_connect_flow,
     introspect_gateway_token,
     is_gateway_dcr_client_id,
     is_proxy_api_resource,
@@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC
     return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302)
 
 
-async def _bridge_authorize_access_denial(
-    litellm_user_id: str,
-    mcp_server: MCPServer,
-    redirect_uri: str,
-    state: str,
-) -> RedirectResponse | None:
-    """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.
-
-    Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the
-    same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting
-    session can actually list and call the server's tools. Without this gate the flow completes, the
-    client shows connected, and every tool request fail-closes to an empty list with nothing telling
-    the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or
-    deactivated user denies like a missing grant, fail closed.
-    """
+async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
     from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
         MCPRequestHandler,
     )
@@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial(
     )
 
     try:
-        admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id)
+        admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id)
     except HTTPException as exc:
         if exc.status_code >= 500:
             raise
-        return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
-    allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
-    if mcp_server.server_id in allowed_server_ids:
+        return False
+    return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
+
+
+async def _bridge_authorize_access_denial(
+    litellm_user_id: str,
+    mcp_server: MCPServer,
+    redirect_uri: str,
+    state: str,
+) -> RedirectResponse | None:
+    """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
+    if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
         return None
     return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
 
@@ -1910,6 +1907,38 @@ async def token_endpoint(
     )
 
 
+async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState:
+    """Whether the gateway itself can see a live vendor credential for this user and server.
+
+    The one reading of "authorized" the connect page displays and the finish step enforces, so
+    the button a user sees and the grant they get cannot disagree. A read fault is neither, and
+    fails the scoped grant closed."""
+    from litellm.proxy._experimental.mcp_server.db import (  # noqa: PLC0415  # circular import at module load
+        get_user_oauth_credential,
+        oauth_grant_state,
+    )
+    from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # circular import at module load
+
+    if prisma_client is None:
+        return "unavailable"
+    try:
+        credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
+    except Exception:  # noqa: BLE001  # a credential-read fault must fail the scoped grant closed
+        return "unavailable"
+    return "absent" if oauth_grant_state(credential) == "absent" else "present"
+
+
+@router.get("/authorize/flow")
+async def authorize_flow(request: Request, flow: str) -> Response:
+    return await describe_connect_flow(
+        request=request,
+        flow_handle=flow,
+        session_user_id=_session_cookie_user_id(request),
+        lookup_vendor_credential=_vendor_credential_state,
+        lookup_server_reachability=_user_can_reach_mcp_server,
+    )
+
+
 @router.post("/authorize/complete")
 async def authorize_complete(
     request: Request,
@@ -1934,6 +1963,8 @@ async def authorize_complete(
         delivery=delivery,
         team_id=team_id,
         decision=decision,
+        lookup_vendor_credential=_vendor_credential_state,
+        lookup_server_reachability=_user_can_reach_mcp_server,
     )
 
 

litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py

@@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code"
 
 ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"]
 ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]]
-"""Injected live-user revalidation (the token endpoint's mirror of admission):
-``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is
-a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else
-fails the grant closed."""
+VendorCredentialState = Literal["present", "absent", "unavailable"]
+"""The per-user vendor credential read has three outcomes: present, absent, or unavailable."""
 
 _DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry"
 _DB_FAULTED_DESCRIPTION: Final = (
@@ -195,6 +193,16 @@ class ConsentTeam(BaseModel):
     team_alias: str | None = None
 
 
+class LookupVendorCredential(Protocol):
+    """Injected read of a user's vendor credential for one server."""
+
+    def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ...
+
+
+class LookupServerReachability(Protocol):
+    def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ...
+
+
 class LookupConsentTeams(Protocol):
     """Injected lookup of the teams a signed-in user may bind a proxy-API credential to."""
 
@@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
     return "unresolvable"
 
 
+async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
+    return "unavailable"
+
+
+async def _unreachable_server(user_id: str, server_id: str) -> bool:
+    return False
+
+
 class GatewayDcrClient(BaseModel):
     """The registration record sealed into a gateway DCR ``client_id``.
 
@@ -449,7 +465,10 @@ def aggregate_authorize(
 
     A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the
     flow to that one server: the scope is sealed into the flow, carried into the code, and
-    bound into the session token, while the connect page interlude runs exactly as before.
+    bound into the session token. The connect URL carries only the flow handle; the page
+    learns the client origin, the scoped server, and whether its vendor OAuth is done from
+    :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry
+    steers which server the page authorizes or names on the confirmation.
 
     Validation failures respond directly with 400 and never redirect: per RFC 6749
     section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and
@@ -474,10 +493,7 @@ def aggregate_authorize(
         resource_server_id=scoped_server.server_id if scoped_server is not None else None,
         audience=None,
     )
-    connect_url: Final = _append_query_params(
-        f"{base_url}/ui/connect",
-        (("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))),
-    )
+    connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
     response: Final = RedirectResponse(connect_url, status_code=303)
     _set_flow_cookie(response, request, handle, flow)
     return response
@@ -684,6 +700,99 @@ def _origin_only(url: str) -> str:
     return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else ""
 
 
+def _open_flow_for(
+    request: Request, flow_handle: str, session_user_id: str | None, now: datetime
+) -> _ConnectFlow | Response:
+    sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
+    if sealed_flow is None:
+        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
+    flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
+    if flow is None or now.timestamp() >= flow.exp:
+        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
+    if session_user_id is None:
+        return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
+    if session_user_id != flow.user_id:
+        return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
+    return flow
+
+
+async def _flow_target(
+    flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability
+) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]:
+    if flow.resource_server_id is None:
+        return "unscoped", None
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (  # noqa: PLC0415  # import cycle
+        MCPServerManager,
+        global_mcp_server_manager,
+    )
+
+    server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id)
+    if (
+        server is None
+        or not server.is_gateway_managed_oauth2
+        or not await lookup_server_reachability(flow.user_id, server.server_id)
+    ):
+        return "stale", None
+    state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive"
+    return state, server
+
+
+class ConnectFlowDescription(TypedDict):
+    """What the connect page is allowed to know about one in-flight flow."""
+
+    state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]]
+    client_origin: ReadOnly[str]
+    server_id: ReadOnly[str | None]
+    server_name: ReadOnly[str | None]
+    connected: ReadOnly[bool | None]
+
+
+async def _describe_opened_flow(
+    flow: _ConnectFlow,
+    lookup_vendor_credential: LookupVendorCredential,
+    lookup_server_reachability: LookupServerReachability,
+) -> ConnectFlowDescription | Response:
+    state, server = await _flow_target(flow, lookup_server_reachability)
+    if state == "interactive" and server is not None:
+        credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id)
+        if credential == "unavailable":
+            return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
+        interactive_description: Final[ConnectFlowDescription] = {
+            "state": state,
+            "client_origin": _origin_only(flow.redirect_uri),
+            "server_id": server.server_id,
+            "server_name": server.server_name or server.alias or server.name,
+            "connected": credential == "present",
+        }
+        return interactive_description
+    described: Final[ConnectFlowDescription] = {
+        "state": state,
+        "client_origin": _origin_only(flow.redirect_uri),
+        "server_id": None if server is None else server.server_id,
+        "server_name": None if server is None else (server.server_name or server.alias or server.name),
+        "connected": state == "m2m" or None,
+    }
+    return described
+
+
+async def describe_connect_flow(
+    request: Request,
+    flow_handle: str,
+    session_user_id: str | None,
+    lookup_vendor_credential: LookupVendorCredential,
+    lookup_server_reachability: LookupServerReachability,
+) -> Response:
+    opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc))
+    if isinstance(opened, Response):
+        return opened
+    described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
+    return (
+        described
+        if isinstance(described, Response)
+        else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS)
+    )
+
+
 async def complete_connect_flow(
     request: Request,
     flow_handle: str,
@@ -692,56 +801,34 @@ async def complete_connect_flow(
     delivery: str | None = None,
     team_id: str | None = None,
     decision: str | None = None,
+    lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential,
+    lookup_server_reachability: LookupServerReachability = _unreachable_server,
 ) -> Response:
-    """The deliberate finish step of the connect flow: mint the gateway authorization
-    code and send the browser back to the client.
+    """Mint the code only after a deliberate POST by the sealed user.
 
-    Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly
-    per-flow cookie plus an exact match between the signed-in user and the user sealed
-    into the flow: a link crafted by another party dies here with ``access_denied``
-    instead of minting a code for the victim's identity. The flow is single-use (an atomic
-    claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
-
-    ``delivery`` chooses how the code reaches the client. Default (absent or
-    ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"``
-    renders the callback URL on a page instead, for a client whose redirect URI is a
-    loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box,
-    container): the 303 would dereference the browser machine's loopback and the code
-    would never arrive, so the user carries it over by pasting the URL into the client or
-    fetching it from the client machine's terminal. Manual delivery is honored only for
-    loopback redirect URIs; a routable redirect URI works from any browser by
-    construction, so those flows always redirect. The user who sees the page is exactly
-    the user the 303 would have carried the code to, and the same user already sees the
-    code today in the dead redirect's address bar, so the page exposes the code to no new
-    party. Unknown ``delivery`` values are rejected rather than defaulted: a client that
-    asked for manual delivery and got a dead redirect instead would silently lose its
-    code.
-
-    ``decision`` and ``team_id`` come from the native-client consent page. ``"deny"``
-    burns the flow and sends the client ``error=access_denied`` so it stops waiting;
-    ``team_id`` is sealed into the code only for proxy-API flows, where it picks which of
-    the user's teams the minted credential is attributed to.
+    A scoped flow additionally requires its sealed server to have a live vendor credential
+    before a code can be minted. The check happens before the single-use claim, so a
+    premature submit can be retried after authorization; denial deliberately bypasses it.
     """
     if delivery not in (None, "redirect", "manual"):
         return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'")
     if decision not in (None, "approve", "deny"):
         return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'")
-    sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
-    if sealed_flow is None:
-        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
-    flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
-    if flow is None:
-        return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
     now: Final = datetime.now(timezone.utc)
-    if now.timestamp() >= flow.exp:
-        return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection")
-    if session_user_id is None:
-        return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
-    if session_user_id != flow.user_id:
-        return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
+    opened: Final = _open_flow_for(request, flow_handle, session_user_id, now)
+    if isinstance(opened, Response):
+        return opened
+    if decision != "deny":
+        described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
+        if isinstance(described, Response):
+            return described
+        if described["state"] == "stale":
+            return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available")
+        if described["connected"] is False:
+            return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing")
     flow_refusal: Final = _claim_refusal(
         await _SingleUseGuard(cache).claim(
-            f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
+            f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
         ),
         replayed=_oauth_error(
             400, "invalid_request", "this connect flow was already completed; restart the connection"
@@ -750,7 +837,7 @@ async def complete_connect_flow(
     if flow_refusal is not None:
         return flow_refusal
     response: Final = (
-        _denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now)
+        _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now)
     )
     path, secure = _cookie_path_and_secure(request)
     response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")

litellm/proxy/_experimental/mcp_server/server.py

@@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict
 from starlette.requests import Request as StarletteRequest
 from starlette.responses import JSONResponse
 from starlette.types import Message, Receive, Scope, Send
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_logger
 from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
@@ -816,6 +817,11 @@ if MCP_AVAILABLE:
                 }
             }
             return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta})
+        except HTTPException as e:
+            from mcp.shared.exceptions import McpError
+            from mcp.types import INVALID_REQUEST, ErrorData
+
+            raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e
         except Exception as e:
             verbose_logger.exception("Error in list_tools endpoint: %s", e)
             # Return empty list instead of failing completely
@@ -1095,6 +1101,7 @@ if MCP_AVAILABLE:
                     mcp_server_auth_headers=mcp_server_auth_headers,
                     oauth2_headers=oauth2_headers,
                     raw_headers=raw_headers,
+                    client_ip=_client_ip,
                     host_progress_callback=host_progress_callback,
                     **data,  # for logging
                 )
@@ -1128,7 +1135,7 @@ if MCP_AVAILABLE:
             except HTTPException as e:
                 verbose_logger.error("HTTPException in MCP tool call: %s", e)
                 return CallToolResult(
-                    content=[TextContent(text=f"Error: {e.detail}", type="text")],
+                    content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")],
                     isError=True,
                 )
             except MCPUpstreamAuthError as e:
@@ -1392,7 +1399,7 @@ if MCP_AVAILABLE:
     ########################################################
 
     async def _get_allowed_mcp_servers_from_mcp_server_names(
-        mcp_servers: list[str] | None,
+        mcp_servers: Sequence[str] | None,
         allowed_mcp_servers: list[MCPServer],
     ) -> list[MCPServer]:
         """
@@ -1413,13 +1420,10 @@ if MCP_AVAILABLE:
                 server_name_matched = False
 
                 for server in allowed_mcp_servers:
-                    if server:
-                        match_list = [s.lower() for s in iter_known_server_prefixes(server) if s]
-
-                        if server_or_group.lower() in match_list:
-                            filtered_server[server.server_id] = server
-                            server_name_matched = True
-                            break
+                    if server and _server_answers_to(server, server_or_group):
+                        filtered_server[server.server_id] = server
+                        server_name_matched = True
+                        break
 
                 if not server_name_matched:
                     try:
@@ -1449,6 +1453,72 @@ if MCP_AVAILABLE:
 
         return allowed_mcp_servers
 
+    def _http_detail_message(detail: object) -> str:
+        return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
+
+    def _server_answers_to(server: MCPServer, name: str) -> bool:
+        requested: Final = name.lower()
+        return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
+
+    class _McpDeniedDetail(TypedDict):
+        error: ReadOnly[str]
+
+    async def raise_denied_scoped_mcp_access(
+        requested_names: Sequence[str],
+        user_api_key_auth: UserAPIKeyAuth | None,
+        client_ip: str | None = None,
+    ) -> None:
+        """A scoped request (``/mcp/<name>`` path or ``x-mcp-servers`` header) resolved to zero
+        allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy
+        server with no tools. Unknown, unauthorized, and access-group names all share one generic
+        error so scoping cannot probe which servers exist; the agent variant fires only when the
+        same request resolves once the agent binding is stripped, proving the binding caused the veto."""
+        agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None
+        if user_api_key_auth is not None and agent_id:
+            resolved_without_agent: Final = await _get_allowed_mcp_servers(
+                user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})),
+                mcp_servers=requested_names,
+                client_ip=client_ip,
+            )
+
+            def _resolved_to_server(name: str) -> bool:
+                return any(_server_answers_to(server, name) for server in resolved_without_agent)
+
+            vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None)
+            if vetoed_server is not None:
+                agent_denial: Final[_McpDeniedDetail] = {
+                    "error": (
+                        f"MCP server '{vetoed_server}' is not available to this key: the key is bound to "
+                        f"agent '{agent_id}', whose MCP grants do not include this server. Add the server "
+                        f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or "
+                        f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
+                    )
+                }
+                raise HTTPException(status_code=403, detail=agent_denial)
+            vetoed_group: Final = next(
+                (
+                    name
+                    for name in requested_names
+                    if not _resolved_to_server(name)
+                    and any(name in (server.access_groups or ()) for server in resolved_without_agent)
+                ),
+                None,
+            )
+            if vetoed_group is not None:
+                group_denial: Final[_McpDeniedDetail] = {
+                    "error": (
+                        f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to "
+                        f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the "
+                        f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or "
+                        f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
+                    )
+                }
+                raise HTTPException(status_code=403, detail=group_denial)
+        generic_denial: Final[_McpDeniedDetail] = {
+            "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}"
+        }
+        raise HTTPException(status_code=403, detail=generic_denial)
+
     def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
         """
         Check if a tool name matches any name in the filter list.
@@ -1541,7 +1611,7 @@ if MCP_AVAILABLE:
 
     async def _get_allowed_mcp_servers(
         user_api_key_auth: UserAPIKeyAuth | None,
-        mcp_servers: list[str] | None,
+        mcp_servers: Sequence[str] | None,
         client_ip: str | None = None,
     ) -> list[MCPServer]:
         """Return allowed MCP servers for a request after applying filters.
@@ -1977,6 +2047,12 @@ if MCP_AVAILABLE:
                 mcp_servers=mcp_servers,
                 client_ip=client_ip,
             )
+            if mcp_servers and not allowed_mcp_servers:
+                await raise_denied_scoped_mcp_access(
+                    requested_names=mcp_servers,
+                    user_api_key_auth=user_api_key_auth,
+                    client_ip=client_ip,
+                )
 
             # Pre-fetch OAuth credentials only when at least one server uses OAuth2,
             # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
@@ -2404,6 +2480,8 @@ if MCP_AVAILABLE:
             )
             verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
             return listing
+        except HTTPException:
+            raise
         except Exception as e:
             verbose_logger.exception("Error getting tools from managed MCP servers: %s", e)
             # Continue with an empty listing instead of failing completely
@@ -3086,6 +3164,7 @@ if MCP_AVAILABLE:
         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,
+        client_ip: str | None = None,
         **kwargs: Any,
     ) -> CallToolResult:
         """
@@ -3116,6 +3195,12 @@ if MCP_AVAILABLE:
                 mcp_servers=mcp_servers,
                 allowed_mcp_servers=allowed_mcp_servers,
             )
+            if mcp_servers and not allowed_mcp_servers:
+                await raise_denied_scoped_mcp_access(
+                    requested_names=mcp_servers,
+                    user_api_key_auth=user_api_key_auth,
+                    client_ip=client_ip,
+                )
             if not allowed_mcp_servers:
                 raise HTTPException(
                     status_code=403,

litellm/proxy/_experimental/mcp_server/tool_search.py

@@ -366,6 +366,7 @@ async def handle_mcp_tool_call(
     from litellm.proxy._experimental.mcp_server.server import (
         _get_allowed_mcp_servers,
         execute_mcp_tool,
+        raise_denied_scoped_mcp_access,
     )
 
     allowed_mcp_servers: Final = await _get_allowed_mcp_servers(
@@ -373,6 +374,12 @@ async def handle_mcp_tool_call(
         mcp_servers=mcp_servers,
         client_ip=client_ip,
     )
+    if mcp_servers and not allowed_mcp_servers:
+        await raise_denied_scoped_mcp_access(
+            requested_names=mcp_servers,
+            user_api_key_auth=user_api_key_dict,
+            client_ip=client_ip,
+        )
 
     # Reject before dispatch when the key has no accessible servers; otherwise an
     # unprefixed local tool name would fall through to the local registry in

litellm/proxy/_lazy_openapi_snapshot.json

@@ -2634,6 +2634,20 @@
               ],
               "title": "Mcp Tool Permissions"
             },
+            "mcp_toolsets": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Mcp Toolsets"
+            },
             "models": {
               "anyOf": [
                 {
@@ -19872,6 +19886,46 @@
           ]
         }
       },
+      "/authorize/flow": {
+        "get": {
+          "operationId": "authorize_flow_authorize_flow_get",
+          "parameters": [
+            {
+              "in": "query",
+              "name": "flow",
+              "required": true,
+              "schema": {
+                "title": "Flow",
+                "type": "string"
+              }
+            }
+          ],
+          "responses": {
+            "200": {
+              "content": {
+                "application/json": {
+                  "schema": {}
+                }
+              },
+              "description": "Successful Response"
+            },
+            "422": {
+              "content": {
+                "application/json": {
+                  "schema": {
+                    "$ref": "#/components/schemas/HTTPValidationError"
+                  }
+                }
+              },
+              "description": "Validation Error"
+            }
+          },
+          "summary": "Authorize Flow",
+          "tags": [
+            "mcp_discoverable"
+          ]
+        }
+      },
       "/callback": {
         "get": {
           "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n  forwarded back to the validated client ``redirect_uri`` with the\n  original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n  RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n  ``redirect_uri``, the error params are propagated back to the client so\n  its OAuth library can surface them. Otherwise we render an HTML error\n  page so the user is not left on an opaque 422 / blank screen.",

litellm/proxy/_types.py

@@ -2833,7 +2833,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
             "Enable only if your deployment is experiencing phantom "
             "BudgetExceededError responses caused by leaked reservations "
             "(see GitHub issue #27639). "
-            "A proxy-level WARNING is logged on every request while this flag "
+            "An INFO notice is logged once per worker at config load while this flag "
             "is active as a reminder that hard enforcement is relaxed."
         ),
     )

litellm/proxy/auth/auth_utils.py

@@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status
 from pydantic import PositiveInt, TypeAdapter, ValidationError
 
 import litellm
-from litellm import Router, provider_list
+from litellm import Router, constants, provider_list
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
     BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,
@@ -1390,6 +1390,24 @@ def warn_once_if_custom_auth_skips_common_checks(
     _custom_auth_common_checks_warning_emitted = True
 
 
+def log_once_if_budget_reservation_disabled(
+    *,
+    disabled: bool,
+    logger: Logger = verbose_proxy_logger,
+) -> None:
+    if constants.budget_reservation_disabled_info_emitted or not disabled:
+        return
+    logger.info(
+        "disable_budget_reservation is enabled: skipping optimistic budget "
+        "reservation. Budget enforcement is read-time only. Concurrent "
+        "requests can each pass the spend check before their cost is recorded, "
+        "so a configured budget may be briefly exceeded under high concurrency. "
+        "Set disable_budget_reservation to False or remove it to restore "
+        "hard per-request budget enforcement."
+    )
+    constants.budget_reservation_disabled_info_emitted = True  # rebind-ok: process-wide one-shot sentinel
+
+
 def is_pass_through_provider_route(route: str) -> bool:
     PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [
         "vertex-ai",

litellm/proxy/auth/user_api_key_auth.py

@@ -2706,14 +2706,6 @@ async def _reserve_budget_after_common_checks(
     if skip_budget_checks:
         return
     if general_settings.get("disable_budget_reservation") is True:
-        verbose_proxy_logger.warning(
-            "disable_budget_reservation is enabled: skipping optimistic budget "
-            "reservation. Budget enforcement is read-time only — concurrent "
-            "requests can each pass the spend check before their cost is recorded, "
-            "so a configured budget may be briefly exceeded under high concurrency. "
-            "Set disable_budget_reservation to False or remove it to restore "
-            "hard per-request budget enforcement."
-        )
         return
 
     from litellm.proxy.spend_tracking.budget_reservation import (

litellm/proxy/common_request_processing.py

@@ -3508,9 +3508,13 @@ class ProxyBaseLLMRequestProcessing:
             error_body: Final = await http_status_error.response.aread()
             error_text: Final = error_body.decode("utf-8")
 
+            error_headers: Final = {  # mutable-ok: HTTPException takes a plain header dict
+                k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items()
+            }
             raise HTTPException(
                 status_code=http_status_error.response.status_code,
                 detail={"error": error_text},
+                headers=error_headers,
             )
         error_msg: Final = f"{e}"
         # Check for AttributeError in the exception chain.

litellm/proxy/common_utils/registry_read_through.py

@@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool:
         )
         return proxy_server.llm_router is not None
     async with proxy_server.MODEL_RECONCILE_LOCK:
+        await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client)
         proxy_server.proxy_config._add_deployment(db_models=rows)
         proxy_server.llm_model_list = router.get_model_list()
     return True

litellm/proxy/db/check_migration.py

@@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]:
 
 def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
     """Checks for differences between current database and Prisma schema.
+
+    Never raises: a diff that cannot be produced, because the runner is missing,
+    because the command failed, or because it outlived its budget, is reported as
+    "no diff" so boot continues.
+
     Returns:
         A tuple containing:
         - A boolean indicating if differences were found (True) or not (False).
-        - A string with the diff output or error message.
-    Raises:
-        subprocess.CalledProcessError: If the Prisma command fails.
-        Exception: For any other errors during execution.
+        - The SQL commands that would close the diff, empty when there is none.
     """
-    verbose_logger.debug("Checking for Prisma schema diff...")
     try:
-        result: Final = subprocess.run(
+        from litellm_proxy_extras.prisma_toolchain import (
+            PRISMA_COMMAND_TIMEOUT_ENV_VAR,
+            prisma_command_timeout,
+            run_prisma,
+        )
+    except ImportError as e:
+        print(  # noqa: T201  # boot-time operator output, same channel as this helper's other messages
+            f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}"
+        )
+        return False, []
+
+    verbose_logger.debug("Checking for Prisma schema diff...")
+    timeout: Final = prisma_command_timeout()
+    try:
+        result: Final = run_prisma(
             [
                 "prisma",
                 "migrate",
@@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
                 "./schema.prisma",
                 "--script",
             ],
-            capture_output=True,
-            text=True,
-            check=True,
+            timeout=timeout,
+            env=os.environ.copy(),
         )
 
-        # return True, "Migration diff generated successfully."
         sql_commands: Final = extract_sql_commands(result.stdout)
 
         if sql_commands:
@@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
             return True, sql_commands
         else:
             return False, []
+    except subprocess.TimeoutExpired:
+        print(  # noqa: T201  # boot-time operator output, same channel as this helper's other messages
+            f"Timed out after {timeout}s generating the migration diff. "
+            f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer."
+        )
+        return False, []
     except subprocess.CalledProcessError as e:
         error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}"
         print(error_message)  # noqa: T201

litellm/proxy/db/prisma_client.py

@@ -937,9 +937,17 @@ class PrismaManager:
                         use_v2_resolver=use_v2_resolver,
                     )
                 else:
+                    try:
+                        from litellm_proxy_extras.prisma_toolchain import (
+                            prisma_command_timeout,
+                            run_prisma,
+                        )
+                    except ImportError as e:
+                        verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
+                        return False
+
                     PrismaManager._raise_if_partitioned_spend_logs()
-                    # Use prisma db push with increased timeout
-                    subprocess.run(
+                    run_prisma(
                         [
                             "prisma",
                             "db",
@@ -947,13 +955,15 @@ class PrismaManager:
                             "--accept-data-loss",
                             "--skip-generate",
                         ],
-                        timeout=60,
-                        check=True,
+                        timeout=prisma_command_timeout(),
+                        env=os.environ.copy(),
+                        stdout=None,
+                        stderr=None,
                     )
                     PrismaManager._apply_replica_identity_full_if_requested()
                     return True
-            except subprocess.TimeoutExpired:
-                verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1)
+            except subprocess.TimeoutExpired as e:
+                verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout)
                 time.sleep(random.randrange(5, 15))
             except subprocess.CalledProcessError as e:
                 attempts_left = 3 - attempt

litellm/proxy/guardrails/guardrail_hooks/tool_permission.py

@@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         for _tool_call, is_allowed, _rule_id, message in checked:
             if not is_allowed and message is not None:
-                verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
+                verbose_proxy_logger.info("Tool Permission Guardrail: %s", message)
                 if self.on_disallowed_action == "block":
                     raise GuardrailRaisedException(
                         guardrail_name=self.guardrail_name, message=message, blocked_content=True
@@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
         new_tools: Final = self._collect_request_tools(data)
         if not new_tools:
-            verbose_proxy_logger.warning(
+            verbose_proxy_logger.debug(
                 "Tool Permission Guardrail: not running guardrail. No tools or functions in data"
             )
             return data
@@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
             is_allowed, _, message = self._check_tool_permission(tool_name, tool_type)
 
             if not is_allowed and message is not None:
-                verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
+                verbose_proxy_logger.info("Tool Permission Guardrail: %s", message)
                 if self.on_disallowed_action == "block":
                     raise HTTPException(
                         status_code=400,

litellm/proxy/hooks/mcp_semantic_filter/hook.py

@@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference.
 Reduces context window size and improves tool selection accuracy.
 """
 
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Final, Optional
 
 from fastapi import HTTPException
@@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger):
         )
 
         # Parse to separate MCP tools from other tools
-        mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+        mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
         if not mcp_tools:
             return []
@@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger):
         return [name for name in names if name]
 
     @staticmethod
-    def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]:
+    async def _narrow_mcp_references(
+        tools: Sequence[Mapping[str, object]],
+        selected_tool_names: list[str],
+        served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None,
+    ) -> list[object]:
         """
         Restrict each litellm_proxy MCP reference to the semantically selected tools.
 
@@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger):
             LiteLLM_Proxy_MCP_Handler,
         )
 
+        via_gateway: Final = await (
+            LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, served_names)
+            if served_names is not None
+            else LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools)
+        )
         return [
-            (
-                {**tool, "allowed_tools": selected_tool_names}
-                if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool])
-                else tool
-            )
-            for tool in tools
+            {**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool
+            for tool, routed in zip(tools, via_gateway, strict=True)
         ]
 
     def _is_mcp_tool(self, tool: object) -> bool:
@@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger):
                 filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools)
 
                 selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools)
-                narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names)
+                narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names)
                 data["tools"] = narrowed_tools
                 self._emit_filter_metadata_safe(
                     data=data,

litellm/proxy/litellm_pre_call_utils.py

@@ -789,12 +789,6 @@ def apply_missing_session_id_policy(
             )
 
 
-def is_claude_code_user_agent(user_agent: str) -> bool:
-    """Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
-    extensions and the Agent SDK run through the same CLI and share that prefix."""
-    return user_agent.startswith("claude-cli/")
-
-
 def is_codex_user_agent(user_agent: str) -> bool:
     """Codex builds its user agent as ``<originator>/<version> ...`` and ships
     several first-party originators: ``codex-tui``, ``codex_cli_rs``,
@@ -811,6 +805,8 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c
     requests routed to providers that reject them. An explicit drop_params
     from the caller or in the operator's ``litellm_settings`` always wins
     over this default."""
+    from litellm.llms.anthropic.common_utils import is_claude_code_user_agent
+
     if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)):
         return False
     if "drop_params" in data:

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -946,7 +946,14 @@ async def handle_bedrock_count_tokens(
     except BedrockError as e:
         # Convert BedrockError to HTTPException for FastAPI
         verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e)
-        raise HTTPException(status_code=e.status_code, detail={"error": e.message})
+        from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers
+
+        provider_headers: Final = getattr(getattr(e, "response", None), "headers", None)
+        raise HTTPException(
+            status_code=e.status_code,
+            detail={"error": e.message},
+            headers=get_response_headers(provider_headers) if provider_headers else None,
+        )
     except HTTPException:
         # Re-raise HTTP exceptions as-is
         raise

litellm/proxy/policy_engine/pipeline_executor.py

@@ -20,7 +20,11 @@ from litellm.integrations.custom_guardrail import (
     ModifyResponseException,
 )
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.litellm_core_utils.core_helpers import independent_snapshot
+from litellm.litellm_core_utils.core_helpers import (
+    get_metadata_variable_name_from_kwargs,
+    get_or_create_metadata_bucket,
+    independent_snapshot,
+)
 from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
     UnifiedLLMGuardrails,
@@ -30,7 +34,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
     PipelineStep,
     PipelineStepResult,
 )
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.types.utils import GenericGuardrailAPIInputs, StandardLoggingGuardrailInformation
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@@ -258,6 +262,7 @@ class PipelineExecutor:
                 return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
 
             if action == "block":
+                _carry_working_guardrail_information(working_data=working_data, request_data=data)
                 return PipelineExecutionResult(
                     terminal_action="block",
                     step_results=step_results,
@@ -267,6 +272,7 @@ class PipelineExecutor:
                 )
 
             if action == "modify_response":
+                _carry_working_guardrail_information(working_data=working_data, request_data=data)
                 return PipelineExecutionResult(
                     terminal_action="modify_response",
                     step_results=step_results,
@@ -357,16 +363,17 @@ class PipelineExecutor:
             verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail)
             return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
 
+        hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot)
+        snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input))
+
+        # Use unified_guardrail path if callback implements apply_guardrail
+        target: CustomLogger = callback
+        use_unified: Final = PipelineExecutor.supports_unified_execution(callback)
+        if use_unified and streaming_chunks is None:
+            hook_input["guardrail_to_apply"] = callback
+            target = UnifiedLLMGuardrails()
+
         try:
-            hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot)
-
-            # Use unified_guardrail path if callback implements apply_guardrail
-            target: CustomLogger = callback
-            use_unified: Final = PipelineExecutor.supports_unified_execution(callback)
-            if use_unified and streaming_chunks is None:
-                hook_input["guardrail_to_apply"] = callback
-                target = UnifiedLLMGuardrails()
-
             if mode == "pre_call":
                 response = await target.async_pre_call_hook(
                     user_api_key_dict=user_api_key_dict,
@@ -430,6 +437,12 @@ class PipelineExecutor:
             else:
                 verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e)
                 return ("error", None, str(e), e)
+        finally:
+            if hook_input is not data:
+                _append_guardrail_information(
+                    request_data=data,
+                    entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:],
+                )
 
     @staticmethod
     def supports_unified_execution(callback: CustomGuardrail) -> bool:
@@ -486,6 +499,40 @@ def _restore_request_guardrails(
     return {**working_data, "metadata": stripped}  # mutable-ok: request dict
 
 
+_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
+
+
+def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]:
+    bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source))
+    recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None
+    return recorded if isinstance(recorded, list) else []
+
+
+def _append_guardrail_information(
+    request_data: dict[str, object],  # mutable-ok: same request-payload shape as execute_steps' data
+    entries: Sequence[StandardLoggingGuardrailInformation],
+) -> None:
+    if not entries:
+        return
+    _, request_bucket = get_or_create_metadata_bucket(request_data)
+    existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY)
+    if isinstance(existing, list):
+        existing.extend(entries)
+        return
+    request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries)
+
+
+def _carry_working_guardrail_information(
+    working_data: Mapping[str, object],
+    request_data: dict[str, object],  # mutable-ok: same request-payload shape as execute_steps' data
+) -> None:
+    recorded: Final = _recorded_guardrail_information(working_data)
+    existing: Final = _recorded_guardrail_information(request_data)
+    if recorded is existing:
+        return
+    _append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing])
+
+
 def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str:
     """
     Map pipeline step outcome to the configured action.

litellm/proxy/proxy_server.py

@@ -306,6 +306,7 @@ from litellm.proxy.auth.auth_checks import (
 from litellm.proxy.auth.auth_utils import (
     check_response_size_is_safe,
     is_request_body_safe,
+    log_once_if_budget_reservation_disabled,
     warn_once_if_custom_auth_skips_common_checks,
 )
 from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
@@ -5653,6 +5654,10 @@ class ProxyConfig:
                 run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)),
             )
 
+            log_once_if_budget_reservation_disabled(
+                disabled=general_settings.get("disable_budget_reservation") is True,
+            )
+
             custom_key_generate: Final = general_settings.get("custom_key_generate", None)
             if custom_key_generate is not None:
                 user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path)
@@ -7104,11 +7109,10 @@ class ProxyConfig:
                 ],
             )
 
-            # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set)
-            if self._should_load_db_object(object_type="models"):
-                new_models: Final = await self._get_models_from_db(prisma_client=prisma_client)
-
-                # update llm router
+            load_models: Final = self._should_load_db_object(object_type="models")
+            new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None
+            await self.get_credentials(prisma_client=prisma_client)
+            if load_models:
                 still_desired_ids = await self._update_llm_router(
                     new_models=new_models, proxy_logging_obj=proxy_logging_obj
                 )
@@ -7148,12 +7152,9 @@ class ProxyConfig:
         async def _resync_config_from_db() -> None:
             await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
 
-        async def _resync_credentials_from_db() -> None:
-            await self.get_credentials(prisma_client=prisma_client)
-
         subscriber: Final = ConfigSyncSubscriber(
             redis_cache=redis_cache,
-            resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db),
+            resync_callbacks=(_resync_config_from_db,),
         )
         self.config_sync_subscriber = subscriber
         subscriber.start()
@@ -8008,7 +8009,7 @@ class ProxyConfig:
 
     async def get_credentials(self, prisma_client: PrismaClient):
         try:
-            credentials = await CredentialsRepository(prisma_client).find_all()
+            credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all()
             credentials = [self.decrypt_credentials(cred) for cred in credentials]
             await self.delete_credentials(credentials)  # delete credentials that are not in the all-up list
             CredentialAccessor.upsert_credentials(credentials)  # upsert credentials that are in the all-up list
@@ -9592,19 +9593,6 @@ class ProxyStartupEvent:
         )
 
         if store_model_in_db is True:
-            ### GET STORED CREDENTIALS ###
-            scheduler.add_job(
-                proxy_config.get_credentials,
-                "interval",
-                seconds=config_reload_interval_seconds,
-                # REMOVED jitter parameter - major cause of memory leak
-                args=[prisma_client],
-                id="get_credentials_job",
-                replace_existing=True,
-                misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
-            )
-            await proxy_config.get_credentials(prisma_client=prisma_client)
-
             # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum
             # Frequent polling was causing excessive memory allocations
             scheduler.add_job(
@@ -9618,7 +9606,7 @@ class ProxyStartupEvent:
                 misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
             )
 
-            # this will load all existing models on proxy startup
+            # this will load all existing credentials and models on proxy startup
             await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
 
             proxy_config.start_config_sync_subscriber(

litellm/proxy/spend_tracking/savings.py

@@ -299,6 +299,8 @@ def compute_autorouter_savings(
     selected_info: ModelInfo | None = None,
     baseline_info: ModelInfo | None = None,
     cost_breakdown: Mapping[str, object] | None = None,
+    baseline_deployment_id: str | None = None,
+    selected_deployment_id: str | None = None,
 ) -> float:
     """Net dollars the router saved, or cost, by serving this request on ``selected_model``.
 
@@ -334,11 +336,12 @@ def compute_autorouter_savings(
     selected: Final = _resolve_model(selected_model, selected_provider)
     if baseline is None or selected is None:
         return 0.0
-    # Same model is only the same cost when it is also the same deployment. Two
-    # deployments of one model can carry different negotiated rates, and routing from
-    # the dear one to the cheap one is a real saving that short-circuiting on the model
-    # name alone reports as zero.
-    if baseline == selected:
+    same_target: Final = (
+        baseline_deployment_id == selected_deployment_id
+        if baseline_deployment_id and selected_deployment_id
+        else baseline == selected
+    )
+    if same_target:
         return 0.0
     basis: Final = _pricing_basis(cost_breakdown)
     effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline)
@@ -517,6 +520,8 @@ def autorouter_savings_for_request(
         selected_info=_effective_model_info(router_instance, model_id, model or ""),
         baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""),
         cost_breakdown=cost_breakdown,
+        baseline_deployment_id=baseline_id,
+        selected_deployment_id=model_id,
     )
     classifier_cost: Final = classifier_cost_from_decision(decision)
     return gross if classifier_cost is None else gross - classifier_cost

litellm/responses/main.py

@@ -177,7 +177,7 @@ async def aresponses_api_with_mcp(
     (
         mcp_tools_with_litellm_proxy,
         other_tools,
-    ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+    ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
     # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform)
     # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata)
@@ -236,6 +236,7 @@ async def aresponses_api_with_mcp(
         "timeout": timeout,
         "custom_llm_provider": custom_llm_provider,
         **kwargs,
+        "_skip_mcp_handler": True,
     }
 
     # Handle MCP streaming if requested
@@ -898,13 +899,14 @@ def _responses_try_dispatch_mcp_gateway(
     custom_llm_provider: str | None,
     kwargs: dict[str, object],
     _is_async: bool,
+    skip_mcp_handler: bool,
 ) -> Any | None:
     """Return a response when MCP gateway handles the call; otherwise None."""
     from litellm.responses.mcp.litellm_proxy_mcp_handler import (
         LiteLLM_Proxy_MCP_Handler,
     )
 
-    if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
+    if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
         return None
     mcp_call_kwargs: Final = {
         "input": input,
@@ -1074,6 +1076,7 @@ def responses(
         litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
         litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
         _is_async: Final = kwargs.pop("aresponses", False) is True
+        skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
         use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs)
 
         client_headers: Final = kwargs.get("headers")
@@ -1168,6 +1171,7 @@ def responses(
             custom_llm_provider=custom_llm_provider,
             kwargs=kwargs,
             _is_async=_is_async,
+            skip_mcp_handler=skip_mcp_handler,
         )
         if _mcp_dispatch is not None:
             return _mcp_dispatch

litellm/responses/mcp/chat_completions_handler.py

@@ -106,7 +106,7 @@ async def acompletion_with_mcp(
     (
         mcp_tools_with_litellm_proxy,
         other_tools,
-    ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
+    ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
 
     if not mcp_tools_with_litellm_proxy:
         # No MCP tools, proceed with regular completion
@@ -114,6 +114,7 @@ async def acompletion_with_mcp(
             model=model,
             messages=messages,
             tools=tools,
+            _skip_mcp_handler=True,
             **kwargs,
         )
 

litellm/responses/mcp/litellm_proxy_mcp_handler.py

@@ -1,6 +1,6 @@
 import re
 import traceback
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload
 
@@ -11,6 +11,7 @@ from litellm._logging import verbose_logger
 from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy._experimental.mcp_server.utils import (
+    iter_known_server_prefixes,
     logging_safe_mcp_headers,
     split_server_prefix_from_name,
     strip_known_server_prefix,
@@ -23,6 +24,7 @@ from litellm.types.llms.openai import (
     ResponsesAPIStreamingResponse,
 )
 from litellm.types.llms.openai import ToolParam as ResponsesToolParam
+from litellm.types.mcp_server.mcp_server_manager import MCPServer
 from litellm.types.utils import (
     CallTypes,
     ChatCompletionMessageCustomToolCall,
@@ -45,6 +47,7 @@ else:
 
 # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling
 ToolParam: TypeAlias = Mapping[str, object]
+SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]]
 
 
 class MCPToolResult(TypedDict):
@@ -56,14 +59,74 @@ class MCPToolResult(TypedDict):
 LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
 LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
 
-# Matches any URL whose path ends with /mcp/<server_name> — covers both root-path
-# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments.
-# A false-positive match (e.g. an external URL that happens to end with /mcp/<name>) results
-# in a "server not found" error from the internal gateway, not a silent failure or data leak,
-# so this broad pattern is intentional and preferred over anchoring to localhost only.
 _PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$")
 
 
+def _mcp_server_url(tool: ToolParam) -> str | None:
+    if not isinstance(tool, dict) or tool.get("type") != "mcp":
+        return None
+    server_url: Final = tool.get("server_url")
+    return server_url if isinstance(server_url, str) else None
+
+
+def _names_gateway_explicitly(tool: ToolParam) -> bool:
+    return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL)
+
+
+def _proxy_path_mcp_name(tool: ToolParam) -> str | None:
+    server_url: Final = _mcp_server_url(tool)
+    match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)
+    return None if match is None else match.group(1)
+
+
+def _registered_mcp_servers() -> Collection[MCPServer]:
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        global_mcp_server_manager,
+    )
+
+    return global_mcp_server_manager.get_registry().values()
+
+
+def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool:
+    requested: Final = name.lower()
+    return any(
+        requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name))
+        or name in (server.access_groups or ())
+        for server in servers
+    )
+
+
+async def _toolset_exists(name: str) -> bool:
+    try:
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            global_mcp_server_manager,
+        )
+        from litellm.proxy.proxy_server import prisma_client
+
+        if prisma_client is None:
+            return False
+        return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None
+    except Exception as e:
+        verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e)
+        return False
+
+
+async def _gateway_served_names(
+    names: Collection[str],
+    servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers,
+    toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists,
+) -> frozenset[str]:
+    registered: Final = tuple(servers()) if names else ()
+    return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)])
+
+
+async def _served_mcp_path_names(
+    tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]]
+) -> frozenset[str]:
+    names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None)
+    return await served_names(names) if names else frozenset[str]()
+
+
 class LiteLLM_Proxy_MCP_Handler:
     """
     Helper class with static methods for MCP integration with Responses API.
@@ -87,57 +150,41 @@ class LiteLLM_Proxy_MCP_Handler:
 
     @staticmethod
     def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool:
-        """
-        Returns True if any MCP tool should be handled via the litellm proxy MCP gateway.
-        This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/<name>.
-        """
-        if tools:
-            for tool in tools:
-                if isinstance(tool, dict) and tool.get("type") == "mcp":
-                    server_url = tool.get("server_url", "")
-                    if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
-                        return True
-                    if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url):
-                        return True
-        return False
+        """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in
+        /mcp/<name>. `_split_mcp_tools` then settles which of the latter the gateway actually serves."""
+        return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ())
 
     @staticmethod
-    def _parse_mcp_tools(
+    def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools:
+        items: Final = tuple(tools or ())
+        gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)]
+        other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)]
+        return gateway_tools, other_tools
+
+    @staticmethod
+    async def _split_mcp_tools(
         tools: Iterable[Mapping[str, object]] | None,
-    ) -> tuple[list[ToolParam], list[Any]]:
-        """
-        Parse tools and separate MCP tools with litellm_proxy from other tools.
+        served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
+    ) -> SplitTools:
+        items: Final = tuple(tools or ())
+        served: Final = await _served_mcp_path_names(items, served_names)
+        return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
+            [
+                {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"}
+                if (name := _proxy_path_mcp_name(tool)) in served
+                else tool
+                for tool in items
+            ]
+        )
 
-        Returns:
-            Tuple of (mcp_tools_with_litellm_proxy, other_tools)
-        """
-        mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = []
-        other_tools: Final[list[Any]] = []
-
-        if tools:
-            for tool in tools:
-                if isinstance(tool, dict) and tool.get("type") == "mcp":
-                    server_url = tool.get("server_url", "")
-                    if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
-                        mcp_tools_with_litellm_proxy.append(tool)
-                    elif isinstance(server_url, str):
-                        # Also intercept URLs like http://localhost:4000/mcp/atlassian_test
-                        # by rewriting them to the internal litellm_proxy format.
-                        m = _PROXY_MCP_PATH_RE.match(server_url)
-                        if m:
-                            rewritten = {
-                                **tool,
-                                "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}",
-                            }
-                            mcp_tools_with_litellm_proxy.append(rewritten)
-                        else:
-                            other_tools.append(tool)
-                    else:
-                        other_tools.append(tool)
-                else:
-                    other_tools.append(tool)
-
-        return mcp_tools_with_litellm_proxy, other_tools
+    @staticmethod
+    async def routes_through_gateway(
+        tools: Iterable[Mapping[str, object]] | None,
+        served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
+    ) -> tuple[bool, ...]:
+        items: Final = tuple(tools or ())
+        served: Final = await _served_mcp_path_names(items, served_names)
+        return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items)
 
     @staticmethod
     async def _apply_toolset_permissions(

litellm/router.py

@@ -1248,7 +1248,7 @@ class Router:
                 selector = LeastBusyLoggingHandler(router_cache=self.cache)
                 if register_callbacks:
                     if isinstance(litellm.input_callback, list):
-                        litellm.input_callback.append(selector)
+                        litellm.logging_callback_manager.add_litellm_input_callback(selector)
                     else:
                         litellm.input_callback = [selector]
             case RoutingStrategy.USAGE_BASED_ROUTING.value:
@@ -4214,10 +4214,12 @@ class Router:
                 }
             )
         litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object)
-        prompt_management_deployment: Final = self.get_available_deployment(
+        specific_deployment: Final = kwargs.pop("specific_deployment", None)
+        prompt_management_deployment: Final = await self.async_get_available_deployment(
             model=model,
-            messages=[{"role": "user", "content": "prompt"}],
-            specific_deployment=kwargs.pop("specific_deployment", None),
+            messages=cast(list[dict[str, str]], messages),  # cast-ok: selection reads messages structurally
+            specific_deployment=specific_deployment,
+            request_kwargs=kwargs,
         )
 
         self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs)

litellm/router_strategy/complexity_router/README.md

@@ -195,6 +195,30 @@ model_list:
         session_affinity_ttl_seconds: 300
 ```
 
+## Custom dimensions
+
+Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal
+
+```yaml
+custom_dimensions:
+  - name: internalFrameworks
+    weight: 0.9
+    keywords: [orbitmesh, fluxgate]
+  - name: sqlMigration
+    weight: 0.7
+    patterns: ['\b(create|alter|drop)\s{1,4}table\b']
+```
+
+Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request
+
+Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one
+
+Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke
+
+Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules
+
+The existing heuristic-v1 tuning quota covers custom dimensions and their weights: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML or the model API; this change adds no dashboard editor
+
 ## Usage
 
 Once configured, use the model name like any other:

litellm/router_strategy/complexity_router/complexity_router.py

@@ -67,6 +67,7 @@ from litellm.types.utils import (
 from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section
 from .config import (
     CALIBRATION_EXAMPLES_HEADING,
+    CUSTOM_PATTERN_SCAN_CHARS,
     DEFAULT_CLASSIFICATION_RUBRIC,
     DEFAULT_CODE_KEYWORDS,
     DEFAULT_ESCALATION_KEYWORDS,
@@ -1119,6 +1120,10 @@ class ComplexityRouter(CustomLogger):
             self.config.custom_technical_keywords,
         )
         self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
+        self._custom_dimensions = tuple(
+            (dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns))
+            for dimension in self.config.custom_dimensions
+        )
         if self.config.has_custom_tiers:
             self.escalation_keywords: tuple[str, ...] = ()
         elif self.config.escalation_keywords is not None:
@@ -1320,6 +1325,17 @@ class ComplexityRouter(CustomLogger):
         score: Final = score_high if match_count >= high_threshold else score_low
         return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
 
+    def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]:
+        if not self._custom_dimensions:
+            return ()
+        scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS]
+        return tuple(
+            (DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight)
+            for dimension, patterns in self._custom_dimensions
+            if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords)
+            or any(pattern.search(scanned) is not None for pattern in patterns)
+        )
+
     def _score_multi_step(self, text: str) -> DimensionScore:
         """Score based on multi-step patterns."""
         hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text))
@@ -1415,12 +1431,13 @@ class ComplexityRouter(CustomLogger):
             self._score_question_complexity(prompt),
         ]
 
-        # Collect signals
-        signals: Final = [d.signal for d in dimensions if d.signal is not None]
+        custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text)
+        signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None]
 
-        # Compute weighted score
         weights: Final = self.config.dimension_weights
-        weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions)
+        weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum(
+            dimension.score * weight for dimension, weight in custom_dimensions
+        )
 
         boundaries: Final = self._effective_tier_boundaries()
         clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score()

litellm/router_strategy/complexity_router/config.py

@@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas
 All values are configurable via proxy config.yaml.
 """
 
-from collections.abc import Mapping
+import math
+import re
+import warnings
+from collections.abc import Iterable, Mapping
 from enum import Enum
 from types import MappingProxyType
-from typing import Annotated, Final, Literal
+from typing import Annotated, Final, Literal, NamedTuple
 
 from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
 
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", DeprecationWarning)
+    import sre_constants
+    import sre_parse
+
 from litellm.types.llms.openai import REASONING_EFFORT
 from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
 
@@ -569,6 +577,117 @@ class ClassifierLLMConfig(BaseModel):
         return self
 
 
+MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64
+MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048
+MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192
+MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16
+CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048
+
+_ATOM_OPCODES: Final = frozenset(
+    {sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY}
+)
+_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT})
+
+
+class _PatternCost(NamedTuple):
+    paths: int
+    steps: int
+
+
+def _atom_steps(node: object) -> int:
+    if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN:
+        return 1 + len(node[1])
+    return 1
+
+
+def _repeat_cost(argument: object) -> _PatternCost | str:
+    if not isinstance(argument, tuple) or len(argument) != 3:
+        return "unsupported repeat structure"
+    low, high, body = argument
+    if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES:
+        return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}"
+    choices: Final = high - low + 1
+    return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices)
+
+
+def _node_cost(node: object, depth: int) -> _PatternCost | str:
+    if not isinstance(node, tuple) or len(node) != 2:
+        return "unsupported regex structure"
+    opcode, argument = node
+    if opcode in _ATOM_OPCODES or opcode is sre_constants.AT:
+        return _PatternCost(1, _atom_steps(node))
+    if opcode is sre_constants.SUBPATTERN:
+        return _sequence_cost(argument[-1], depth + 1)
+    if opcode is sre_constants.BRANCH:
+        costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1])
+        refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
+        if refused is not None:
+            return refused
+        return _PatternCost(
+            sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)),
+            len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)),
+        )
+    if opcode in _REPEAT_OPCODES:
+        return _repeat_cost(argument)
+    return "contains an unsupported regex construct"
+
+
+def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str:
+    if depth > MAX_CUSTOM_PATTERN_DEPTH:
+        return "nests deeper than 16 levels"
+    costs: Final = tuple(_node_cost(node, depth) for node in nodes)
+    refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
+    if refused is not None:
+        return refused
+    valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost))
+    # Choices multiply across a sequence; every continuation can execute once per preceding path.
+    total: Final = _PatternCost(
+        math.prod(cost.paths for cost in valid),
+        1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)),
+    )
+    if total.steps > MAX_CUSTOM_PATTERN_WORK:
+        return "exceeds the per-pattern regex work budget"
+    return total
+
+
+def custom_pattern_work(pattern: str) -> int | str:
+    try:
+        re.compile(pattern, re.IGNORECASE)
+        parsed: Final = sre_parse.parse(pattern, re.IGNORECASE)
+    except (re.error, RecursionError, OverflowError):
+        return "is not a valid regex"
+    cost: Final = _sequence_cost(tuple(parsed), 0)
+    return cost if isinstance(cost, str) else cost.steps
+
+
+class CustomDimension(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$")
+    weight: float = Field(gt=0, le=1, allow_inf_nan=False)
+    keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
+    patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
+
+    @model_validator(mode="after")
+    def _validate_matchers(self) -> "CustomDimension":
+        matchers: Final = (*self.keywords, *self.patterns)
+        if not matchers or any(not matcher.strip() for matcher in matchers):
+            raise ValueError("custom dimensions require nonblank keywords and/or patterns")
+        if len(matchers) > 32 or sum(map(len, matchers)) > 4096:
+            raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each")
+        costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns)
+        rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str))
+        if rejected:
+            raise ValueError("custom dimension " + "; ".join(rejected))
+        return self
+
+    def pattern_work(self) -> int:
+        """Combined work estimate of the validated patterns."""
+        return sum(
+            work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int)
+        )
+
+
 class ComplexityRouterConfig(BaseModel):
     """Configuration for the ComplexityRouter."""
 
@@ -671,6 +790,19 @@ class ComplexityRouterConfig(BaseModel):
         description="Weights for each scoring dimension",
     )
 
+    custom_dimensions: tuple[CustomDimension, ...] = Field(
+        default=(),
+        max_length=16,
+        description=(
+            "Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
+            "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. "
+            "Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, "
+            "backreferences and lookarounds are rejected. Conservative work limits include alternation paths, "
+            "repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. "
+            "Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota."
+        ),
+    )
+
     # Keyword lists (overridable)
     code_keywords: list[str] | None = Field(
         default=None,
@@ -1245,6 +1377,27 @@ class ComplexityRouterConfig(BaseModel):
             )
         return self
 
+    @model_validator(mode="after")
+    def _validate_custom_dimensions(self) -> "ComplexityRouterConfig":
+        if not self.custom_dimensions:
+            return self
+        if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"):
+            raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid")
+        names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions)
+        reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS)
+        weighted: Final = frozenset(name.casefold() for name in self.dimension_weights)
+        if len(frozenset(names)) != len(names) or frozenset(names) & reserved:
+            raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions")
+        if frozenset(names) & weighted:
+            raise ValueError("custom dimension weights must be inline, not in dimension_weights")
+        work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions)
+        if work > MAX_CUSTOM_DIMENSIONS_WORK:
+            raise ValueError(
+                f"custom_dimensions regex work estimate is {work}; the limit across the router is "
+                f"{MAX_CUSTOM_DIMENSIONS_WORK}"
+            )
+        return self
+
     @field_validator("heuristic_first_max_tier", mode="before")
     @classmethod
     def _coerce_heuristic_first_max_tier(cls, value: object) -> object:

litellm/router_strategy/least_busy.py

@@ -1,17 +1,103 @@
-#### What this does ####
-#   identifies least busy deployment
-#   How is this achieved?
-#   - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"}
-#   - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic}
-#   - use litellm.success + failure callbacks to log when a request completed
-#   - in get_available_deployment, for a given model group name -> pick based on traffic
-
-import random
+from collections.abc import Mapping, Sequence
 from typing import Final
 
+from pydantic import TypeAdapter, ValidationError
+from typing_extensions import ReadOnly, TypedDict
+
+from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
 
+IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
+
+
+class _ModelInfo(TypedDict, total=False):
+    id: ReadOnly[str | int | None]
+
+
+class _Metadata(TypedDict, total=False):
+    model_group: ReadOnly[str | None]
+
+
+class _LitellmParams(TypedDict, total=False):
+    metadata: ReadOnly[_Metadata | None]
+    model_info: ReadOnly[_ModelInfo | None]
+
+
+class _CallKwargs(TypedDict, total=False):
+    litellm_params: ReadOnly[_LitellmParams | None]
+
+
+class _DeploymentModelInfo(TypedDict):
+    id: ReadOnly[str | int]
+
+
+class _Deployment(TypedDict):
+    model_info: ReadOnly[_DeploymentModelInfo]
+
+
+_CALL_KWARGS: Final = TypeAdapter(_CallKwargs)
+_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
+_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None)
+
+
+def _request_count_key(model_group: str, deployment_id: str) -> str:
+    return f"{model_group}_request_count:{deployment_id}"
+
+
+def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
+    try:
+        call: Final = _CALL_KWARGS.validate_python(kwargs)
+    except ValidationError:
+        return None
+    litellm_params: Final = call.get("litellm_params")
+    metadata: Final = litellm_params.get("metadata") if litellm_params else None
+    model_info: Final = litellm_params.get("model_info") if litellm_params else None
+    model_group: Final = metadata.get("model_group") if metadata else None
+    deployment_id: Final = model_info.get("id") if model_info else None
+    if model_group is None or deployment_id is None:
+        return None
+    return model_group, str(deployment_id)
+
+
+def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
+    return tuple(
+        _request_count_key(model_group, str(deployment["model_info"]["id"]))
+        for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
+    )
+
+
+def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]:
+    return tuple(0 if value is None else int(value) for value in values)
+
+
+def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]:
+    values: Final = _MEMORY_COUNTS.validate_python(raw)
+    if values is None or len(values) != len(keys):
+        return (0,) * len(keys)
+    return _as_counts(values)
+
+
+def _least_busy(
+    healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...]
+) -> Mapping[str, object] | None:
+    if not healthy_deployments:
+        return None
+    return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])]
+
+
+def _warn_unreadable(model_group: str, error: Exception) -> None:
+    verbose_router_logger.warning(
+        "least-busy routing could not read the shared in-flight counts for %s, "
+        "falling back to this worker's own counts: %s",
+        model_group,
+        error,
+    )
+
+
+def _warn_unwritable(key: str, error: Exception) -> None:
+    verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error)
+
 
 class LeastBusyLoggingHandler(CustomLogger):
     test_flag: bool = False
@@ -20,195 +106,101 @@ class LeastBusyLoggingHandler(CustomLogger):
 
     def __init__(self, router_cache: DualCache):
         self.router_cache = router_cache
+        self.router_cache_id = str(id(router_cache))
 
-    def log_pre_api_call(self, model, messages, kwargs):
-        """
-        Log when a model is being used.
+    def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
+        self._increment(kwargs, 1)
 
-        Caching based on model group.
-        """
-        try:
-            if kwargs["litellm_params"].get("metadata") is None:
-                pass
-            else:
-                model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
-                id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
-                if model_group is None or id is None:
-                    return
-                elif isinstance(id, int):
-                    id = str(id)
+    def log_success_event(
+        self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
+    ) -> None:
+        self._increment(kwargs, -1)
+        if self.test_flag:
+            self.logged_success += 1
 
-                request_count_api_key: Final = f"{model_group}_request_count"
-                # update cache
-                request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
-                request_count_dict[id] = request_count_dict.get(id, 0) + 1
+    def log_failure_event(
+        self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
+    ) -> None:
+        self._increment(kwargs, -1)
+        if self.test_flag:
+            self.logged_failure += 1
 
-                self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict)
-        except Exception:
-            pass
+    async def async_log_success_event(
+        self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
+    ) -> None:
+        await self._async_increment(kwargs, -1)
+        if self.test_flag:
+            self.logged_success += 1
 
-    def log_success_event(self, kwargs, response_obj, start_time, end_time):
-        try:
-            if kwargs["litellm_params"].get("metadata") is None:
-                pass
-            else:
-                model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
-
-                id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
-                if model_group is None or id is None:
-                    return
-                elif isinstance(id, int):
-                    id = str(id)
-
-                request_count_api_key: Final = f"{model_group}_request_count"
-                # decrement count in cache
-                request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
-                request_count_value: Final[int | None] = request_count_dict.get(id, 0)
-                if request_count_value is None:
-                    return
-                request_count_dict[id] = request_count_value - 1
-                self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict)
-
-                ### TESTING ###
-                if self.test_flag:
-                    self.logged_success += 1
-        except Exception:
-            pass
-
-    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
-        try:
-            if kwargs["litellm_params"].get("metadata") is None:
-                pass
-            else:
-                model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
-                id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
-                if model_group is None or id is None:
-                    return
-                elif isinstance(id, int):
-                    id = str(id)
-
-                request_count_api_key: Final = f"{model_group}_request_count"
-                # decrement count in cache
-                request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
-                request_count_value: Final[int | None] = request_count_dict.get(id, 0)
-                if request_count_value is None:
-                    return
-                request_count_dict[id] = request_count_value - 1
-                self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict)
-
-                ### TESTING ###
-                if self.test_flag:
-                    self.logged_failure += 1
-        except Exception:
-            pass
-
-    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
-        try:
-            if kwargs["litellm_params"].get("metadata") is None:
-                pass
-            else:
-                model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
-
-                id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
-                if model_group is None or id is None:
-                    return
-                elif isinstance(id, int):
-                    id = str(id)
-
-                request_count_api_key: Final = f"{model_group}_request_count"
-                # decrement count in cache
-                request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {}
-                request_count_value: Final[int | None] = request_count_dict.get(id, 0)
-                if request_count_value is None:
-                    return
-                request_count_dict[id] = request_count_value - 1
-                await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict)
-
-                ### TESTING ###
-                if self.test_flag:
-                    self.logged_success += 1
-        except Exception:
-            pass
-
-    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
-        try:
-            if kwargs["litellm_params"].get("metadata") is None:
-                pass
-            else:
-                model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
-                id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
-                if model_group is None or id is None:
-                    return
-                elif isinstance(id, int):
-                    id = str(id)
-
-                request_count_api_key: Final = f"{model_group}_request_count"
-                # decrement count in cache
-                request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {}
-                request_count_value: Final[int | None] = request_count_dict.get(id, 0)
-                if request_count_value is None:
-                    return
-                request_count_dict[id] = request_count_value - 1
-                await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict)
-
-                ### TESTING ###
-                if self.test_flag:
-                    self.logged_failure += 1
-        except Exception:
-            pass
-
-    def _get_available_deployments(
-        self,
-        healthy_deployments: list,
-        all_deployments: dict,
-    ):
-        """
-        Helper to get deployments using least busy strategy
-        """
-        for d in healthy_deployments:
-            ## if healthy deployment not yet used
-            if d["model_info"]["id"] not in all_deployments:
-                all_deployments[d["model_info"]["id"]] = 0
-        # map deployment to id
-        # pick least busy deployment
-        min_traffic = float("inf")
-        min_deployment = None
-        for k, v in all_deployments.items():
-            if v < min_traffic:
-                min_traffic = v
-                min_deployment = k
-        if min_deployment is not None:
-            ## check if min deployment is a string, if so, cast it to int
-            for m in healthy_deployments:
-                if m["model_info"]["id"] == min_deployment:
-                    return m
-            min_deployment = random.choice(healthy_deployments)
-        else:
-            min_deployment = random.choice(healthy_deployments)
-        return min_deployment
+    async def async_log_failure_event(
+        self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
+    ) -> None:
+        await self._async_increment(kwargs, -1)
+        if self.test_flag:
+            self.logged_failure += 1
 
     def get_available_deployments(
-        self,
-        model_group: str,
-        healthy_deployments: list,
-    ):
-        """
-        Sync helper to get deployments using least busy strategy
-        """
-        request_count_api_key: Final = f"{model_group}_request_count"
-        all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
-        return self._get_available_deployments(
-            healthy_deployments=healthy_deployments,
-            all_deployments=all_deployments,
-        )
+        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
+    ) -> Mapping[str, object] | None:
+        keys: Final = _request_count_keys(model_group, healthy_deployments)
+        redis_cache: Final = self.router_cache.redis_cache
+        if redis_cache is not None:
+            try:
+                shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys)))
+            except Exception as e:
+                _warn_unreadable(model_group, e)
+            else:
+                return _least_busy(healthy_deployments, shared)
+        local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
+        return _least_busy(healthy_deployments, local)
 
-    async def async_get_available_deployments(self, model_group: str, healthy_deployments: list):
-        """
-        Async helper to get deployments using least busy strategy
-        """
-        request_count_api_key: Final = f"{model_group}_request_count"
-        all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {}
-        return self._get_available_deployments(
-            healthy_deployments=healthy_deployments,
-            all_deployments=all_deployments,
-        )
+    async def async_get_available_deployments(
+        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
+    ) -> Mapping[str, object] | None:
+        keys: Final = _request_count_keys(model_group, healthy_deployments)
+        redis_cache: Final = self.router_cache.redis_cache
+        if redis_cache is not None:
+            try:
+                shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys)))
+            except Exception as e:
+                _warn_unreadable(model_group, e)
+            else:
+                return _least_busy(healthy_deployments, shared)
+        local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys)
+        return _least_busy(healthy_deployments, local)
+
+    def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
+        ref: Final = _deployment_ref(kwargs)
+        if ref is None:
+            return
+        key: Final = _request_count_key(*ref)
+        redis_cache: Final = self.router_cache.redis_cache
+        try:
+            local: Final = self.router_cache.increment_cache(
+                key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
+            )
+            if local < 0:
+                self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
+            if redis_cache is None:
+                return
+            redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
+        except Exception as e:
+            _warn_unwritable(key, e)
+
+    async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None:
+        ref: Final = _deployment_ref(kwargs)
+        if ref is None:
+            return
+        key: Final = _request_count_key(*ref)
+        redis_cache: Final = self.router_cache.redis_cache
+        try:
+            local: Final = await self.router_cache.async_increment_cache(
+                key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
+            )
+            if local is not None and local < 0:
+                await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
+            if redis_cache is None:
+                return
+            await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
+        except Exception as e:
+            _warn_unwritable(key, e)

litellm/router_strategy/lowest_latency.py

@@ -32,6 +32,12 @@ def _average_latency(samples: Sequence[float]) -> float:
     return sum(samples) / len(samples)
 
 
+def _ttft_seconds(elapsed: timedelta | float) -> float:
+    if isinstance(elapsed, timedelta):
+        return elapsed.total_seconds()
+    return float(elapsed)
+
+
 class LowestLatencyLoggingHandler(CustomLogger):
     test_flag: bool = False
     logged_success: int = 0
@@ -86,14 +92,13 @@ class LowestLatencyLoggingHandler(CustomLogger):
                     # breaks JSON serialization when the router cache syncs to
                     # Redis (issue #33169)
                     response_ms = response_ms.total_seconds()
-                time_to_first_token_response_time = None
+                time_to_first_token: float | None = None
 
                 if kwargs.get("stream", None) is not None and kwargs["stream"] is True:
                     # only log ttft for streaming request
-                    time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time
+                    time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time)
 
                 final_value: float = response_ms
-                time_to_first_token: float | None = None
                 total_tokens = 0
 
                 if isinstance(response_obj, ModelResponse):
@@ -111,13 +116,6 @@ class LowestLatencyLoggingHandler(CustomLogger):
                         else:
                             final_value = response_seconds
 
-                        if time_to_first_token_response_time is not None:
-                            if isinstance(time_to_first_token_response_time, timedelta):
-                                ttft_seconds = time_to_first_token_response_time.total_seconds()
-                            else:
-                                ttft_seconds = time_to_first_token_response_time
-                            time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens)
-
                 # ------------
                 # Update usage
                 # ------------
@@ -138,14 +136,14 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 ## Time to first token
                 if time_to_first_token is not None:
                     if (
-                        len(request_count_dict[id].get("time_to_first_token", []))
+                        len(request_count_dict[id].get("time_to_first_token_seconds", []))
                         < self.routing_args.max_latency_list_size
                     ):
-                        request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token)
+                        request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token)
                     else:
-                        request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][
-                            1:
-                        ] + [time_to_first_token]
+                        request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][
+                            "time_to_first_token_seconds"
+                        ][1:] + [time_to_first_token]
 
                 if precise_minute not in request_count_dict[id]:
                     request_count_dict[id][precise_minute] = {}
@@ -252,7 +250,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
                     {model_group}_map: {
                         id: {
                             "latency": [..]
-                            "time_to_first_token": [..]
+                            "time_to_first_token_seconds": [..]
                             f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
                         }
                     }
@@ -273,14 +271,13 @@ class LowestLatencyLoggingHandler(CustomLogger):
                     # breaks JSON serialization when the router cache syncs to
                     # Redis (issue #33169)
                     response_ms = response_ms.total_seconds()
-                time_to_first_token_response_time = None
+                time_to_first_token: float | None = None
                 if kwargs.get("stream", None) is not None and kwargs["stream"] is True:
                     # only log ttft for streaming request
-                    time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time
+                    time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time)
 
                 final_value: float = response_ms
                 total_tokens = 0
-                time_to_first_token: float | None = None
 
                 if isinstance(response_obj, ModelResponse):
                     _usage: Final = getattr(response_obj, "usage", None)
@@ -296,13 +293,6 @@ class LowestLatencyLoggingHandler(CustomLogger):
                             final_value = float(normalized_value)
                         else:
                             final_value = response_seconds
-
-                        if time_to_first_token_response_time is not None:
-                            if isinstance(time_to_first_token_response_time, timedelta):
-                                ttft_seconds = time_to_first_token_response_time.total_seconds()
-                            else:
-                                ttft_seconds = time_to_first_token_response_time
-                            time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens)
                 # ------------
                 # Update usage
                 # ------------
@@ -328,14 +318,14 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 ## Time to first token
                 if time_to_first_token is not None:
                     if (
-                        len(request_count_dict[id].get("time_to_first_token", []))
+                        len(request_count_dict[id].get("time_to_first_token_seconds", []))
                         < self.routing_args.max_latency_list_size
                     ):
-                        request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token)
+                        request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token)
                     else:
-                        request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][
-                            1:
-                        ] + [time_to_first_token]
+                        request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][
+                            "time_to_first_token_seconds"
+                        ][1:] + [time_to_first_token]
 
                 if precise_minute not in request_count_dict[id]:
                     request_count_dict[id][precise_minute] = {}
@@ -433,7 +423,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 or float("inf")
             )
             item_latency = item_map.get("latency", [])
-            item_ttft_latency = item_map.get("time_to_first_token", [])
+            item_ttft_latency = item_map.get("time_to_first_token_seconds", [])
             item_rpm = item_map.get(precise_minute, {}).get("rpm", 0)
             item_tpm = item_map.get(precise_minute, {}).get("tpm", 0)
 

litellm/router_utils/auto_router_tuning_baseline.py

@@ -20,6 +20,7 @@ HEURISTIC_V1_TUNING_FIELDS: Final = (
     "reasoning_override_min_score",
     "token_thresholds",
     "dimension_weights",
+    "custom_dimensions",
     "code_keywords",
     "reasoning_keywords",
     "technical_keywords",

litellm/router_utils/cooldown_cache.py

@@ -12,6 +12,7 @@ from typing_extensions import TypedDict
 from litellm import verbose_logger
 from litellm.caching.caching import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS
 from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
 
 if TYPE_CHECKING:
@@ -36,10 +37,19 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0
 
 
 class CooldownCache:
-    def __init__(self, cache: DualCache, default_cooldown_time: float):
+    def __init__(
+        self,
+        cache: DualCache,
+        default_cooldown_time: float,
+        redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS,
+    ):
         self.cache = cache
         self.default_cooldown_time = default_cooldown_time
         self.in_memory_cache = InMemoryCache()
+        self._cooldown_store = DualCache(
+            in_memory_cache=self.in_memory_cache,
+            default_redis_batch_cache_expiry=redis_read_interval_seconds,
+        )
         # Initialize the masker with custom settings for exception strings
         self.exception_masker = SensitiveDataMasker(
             visible_prefix=50,  # Show first 50 characters
@@ -48,6 +58,21 @@ class CooldownCache:
             mask_short_values=False,  # Truncate long messages only; keep short ones readable
         )
 
+    @property
+    def cooldown_store(self) -> DualCache:
+        """
+        The cache cooldown entries live in, with the router's Redis attached on first use.
+
+        It is kept separate from the router-wide cache so that a key missing from memory is
+        re-read from Redis every `redis_read_interval_seconds` rather than on the router
+        cache's much longer batch interval, which is what lets a sibling replica see a
+        cooldown another replica wrote, and so that unrelated router keys cannot evict a
+        cooldown from the in-memory tier before it expires. Redis is attached lazily because
+        the router builds its cooldown cache before it wires up the shared Redis client.
+        """
+        self._cooldown_store.attach_redis_cache(self.cache.redis_cache)
+        return self._cooldown_store
+
     def _common_add_cooldown_logic(
         self, model_id: str, original_exception, exception_status, cooldown_time: float
     ) -> tuple[str, CooldownCacheValue]:
@@ -93,7 +118,7 @@ class CooldownCache:
             )
 
             # Set the cache with a TTL equal to the cooldown time
-            self.cache.set_cache(
+            self.cooldown_store.set_cache(
                 value=cooldown_data,
                 key=cooldown_key,
                 ttl=_cooldown_time,
@@ -122,13 +147,13 @@ class CooldownCache:
         cooldown_cache_value: Final = CooldownCacheValue(**result)  # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code
         remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time
         if remaining <= 0:
-            self.cache.in_memory_cache.delete_cache(key)
+            self.in_memory_cache.delete_cache(key)
             return None
-        current_expiry: Final = self.cache.in_memory_cache.ttl_dict.get(key)
+        current_expiry: Final = self.in_memory_cache.ttl_dict.get(key)
         if current_expiry is not None and current_expiry > current_time + remaining + 5:
             corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS)
-            self.cache.in_memory_cache.delete_cache(key)
-            self.cache.in_memory_cache.set_cache(key, result, ttl=corrected_ttl)
+            self.in_memory_cache.delete_cache(key)
+            self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl)
         return cooldown_cache_value
 
     async def async_get_active_cooldowns(
@@ -137,12 +162,7 @@ class CooldownCache:
         # Generate the keys for the deployments
         keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
 
-        # Retrieve the values for the keys using mget
-        ## more likely to be none if no models ratelimited. So just check redis every 1s
-        ## each redis call adds ~100ms latency.
-
-        ## check in memory cache first
-        results: Final = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span)
+        results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span)
         active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = []
 
         if results is None or all(v is None for v in results):
@@ -164,7 +184,7 @@ class CooldownCache:
         # Generate the keys for the deployments
         keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
         # Retrieve the values for the keys using mget
-        results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
+        results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
 
         active_cooldowns: Final = []
         current_time: Final = time.time()
@@ -184,7 +204,7 @@ class CooldownCache:
         keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
 
         # Retrieve the values for the keys using mget
-        results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
+        results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
 
         min_cooldown_time: float | None = None
         # Process the results

litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py

@@ -90,6 +90,7 @@ class PromptCachingDeploymentCheck(CustomLogger):
                 enable_prompt_caching=(
                     request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None
                 ),
+                request_kwargs=request_kwargs,
             )
 
             model_id_dict: Final = await prompt_cache.async_get_model_id(

litellm/types/agents.py

@@ -1,9 +1,9 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, Literal
 
 from pydantic import BaseModel, PrivateAttr, StrictInt
-from typing_extensions import Required, TypedDict
+from typing_extensions import ReadOnly, Required, TypedDict
 
 from litellm.types.llms.base import LiteLLMPydanticObjectBase
 
@@ -172,6 +172,7 @@ class AugmentedAgentCard(AgentCard):
 class AgentObjectPermission(TypedDict, total=False):
     mcp_servers: list[str] | None
     mcp_access_groups: list[str] | None
+    mcp_toolsets: ReadOnly[Sequence[str] | None]
     mcp_tool_permissions: dict[str, list[str]] | None
     models: list[str] | None
     agents: list[str] | None

litellm/utils.py

@@ -672,11 +672,19 @@ def load_credentials_from_list(kwargs: dict):
     CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor")
 
     credential_name: Final = kwargs.get("litellm_credential_name")
-    if credential_name and litellm.credential_list:
-        credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name)
-        for key, value in credential_accessor.items():
-            if key not in kwargs:
-                kwargs[key] = value
+    if not credential_name:
+        return
+    credential: Final = CredentialAccessor.find_credential(credential_name)
+    if credential is None:
+        verbose_logger.warning(
+            "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it",
+            credential_name,
+            len(litellm.credential_list),
+        )
+        return
+    for key, value in credential.credential_values.items():
+        if key not in kwargs:
+            kwargs[key] = value
 
 
 def get_dynamic_callbacks(

pyproject.toml

@@ -340,6 +340,7 @@ markers = [
     "asyncio: mark test as an asyncio test",
     "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
     "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests",
+    "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension",
 ]
 filterwarnings = [
     # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests)

ruff-strict-budget.json

@@ -1,6 +1,6 @@
 {
   "ANN001": {
-    "limit": 2956
+    "limit": 2918
   },
   "ANN002": {
     "limit": 71
@@ -9,10 +9,10 @@
     "limit": 806
   },
   "ANN201": {
-    "limit": 1979
+    "limit": 1965
   },
   "ANN202": {
-    "limit": 831
+    "limit": 829
   },
   "ANN204": {
     "limit": 683
@@ -57,7 +57,7 @@
     "limit": 3
   },
   "BLE001": {
-    "limit": 2916
+    "limit": 2914
   },
   "C401": {
     "limit": 8
@@ -189,7 +189,7 @@
     "limit": 0
   },
   "S110": {
-    "limit": 217
+    "limit": 207
   },
   "S112": {
     "limit": 22

tests/llm_translation/test_bedrock_embedding.py

@@ -1,14 +1,12 @@
 import json
 import os
-from datetime import datetime
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import Mock, patch
 import pytest
 import base64
-import httpx
 
 
 import litellm
-from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
 
 titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}
 
@@ -394,8 +392,6 @@ def test_bedrock_embedding_uses_correct_region_when_specified():
             os.environ["AWS_REGION_NAME"] = original_region_name
         else:
             os.environ.pop("AWS_REGION_NAME", None)
-
-
 def test_bedrock_embedding_region_bug_reproduction():
     """
     Reproduces the bug where aws_region_name is ignored when passed explicitly.
@@ -458,13 +454,3 @@ def test_bedrock_embedding_region_bug_reproduction():
             os.environ["AWS_REGION_NAME"] = original_region_name
         else:
             os.environ.pop("AWS_REGION_NAME", None)
-
-
-def test_bedrock_titan_g1_text_02_model_info():
-    """Test that amazon.titan-embed-g1-text-02 has correct pricing metadata"""
-    model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02")
-    assert model_info is not None, "Model info should not be None"
-    assert model_info["litellm_provider"] == "bedrock"
-    assert model_info["mode"] == "embedding"
-    assert model_info["input_cost_per_token"] == 1e-07
-    assert model_info["max_input_tokens"] == 8192

tests/llm_translation/test_bedrock_embedding_pricing.py

@@ -1,34 +0,0 @@
-"""
-Tests for AWS Bedrock embedding model pricing in the model cost map.
-
-Regression test for the Amazon Titan Text Embeddings V2 commercial price,
-which was previously set 10x too high (2e-07 instead of 2e-08).
-AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens
-(= $0.00002 per 1K tokens = 2e-08 per token).
-"""
-
-import importlib
-
-
-class TestBedrockEmbeddingPricing:
-    """Test suite for Bedrock embedding model pricing in the cost map."""
-
-    def test_titan_embed_v2_commercial_input_cost(self, monkeypatch):
-        """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08)."""
-        # Scope the local-cost-map flag to this test only, so it does not leak
-        # into sibling tests. monkeypatch restores the environment on teardown.
-        monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-
-        import litellm.litellm_core_utils.get_model_cost_map
-        import litellm
-
-        # Reload so the cost map is re-read from the local file with the flag set.
-        importlib.reload(litellm.litellm_core_utils.get_model_cost_map)
-        importlib.reload(litellm)
-
-        model = litellm.model_cost["amazon.titan-embed-text-v2:0"]
-
-        assert model["input_cost_per_token"] == 2e-08
-        assert model["output_cost_per_token"] == 0.0
-        assert model["litellm_provider"] == "bedrock"
-        assert model["mode"] == "embedding"

tests/llm_translation/test_bedrock_govcloud.py

@@ -40,38 +40,6 @@ class TestBedrockGovCloudSupport:
         assert "us-gov-east-1" in all_regions
         assert "us-gov-west-1" in all_regions
 
-    def test_govcloud_models_in_model_cost(self):
-        """Test that GovCloud models are present in model cost configuration"""
-        from litellm import model_cost
-
-        # Test Claude models in GovCloud
-        assert (
-            "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-            in model_cost
-        )
-        assert (
-            "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-            in model_cost
-        )
-        assert (
-            "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
-        )
-        assert (
-            "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
-        )
-        assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
-        assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
-
-        # Test Llama models in GovCloud
-        assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost
-        assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost
-        assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost
-        assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost
-
-        # Test Titan models in GovCloud
-        assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost
-        assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost
-
     def test_govcloud_model_routing(self):
         """Test that GovCloud models are routed correctly"""
         # Test Claude model routing
@@ -148,135 +116,6 @@ class TestBedrockGovCloudSupport:
         assert not any("us-gov-east-1" in model for model in litellm.bedrock_models)
         assert not any("us-gov-west-1" in model for model in litellm.bedrock_models)
 
-    def test_govcloud_model_cost_properties(self):
-        """Test that GovCloud models have proper cost configuration"""
-        from litellm import model_cost
-
-        # Check a specific GovCloud model has all required properties
-        govcloud_model = model_cost[
-            "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-        ]
-
-        assert "max_tokens" in govcloud_model
-        assert "max_input_tokens" in govcloud_model
-        assert "max_output_tokens" in govcloud_model
-        assert "input_cost_per_token" in govcloud_model
-        assert "output_cost_per_token" in govcloud_model
-        assert govcloud_model["litellm_provider"] == "bedrock"
-        assert govcloud_model["mode"] == "chat"
-
-    def test_govcloud_model_pricing_verification(self):
-        """Test that GovCloud models have correct pricing that differs from base models"""
-        from litellm import model_cost
-
-        # Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id
-        base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
-        gov_east_model = (
-            "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-        )
-        gov_west_model = (
-            "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
-        )
-
-        # Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok)
-        base_pricing = model_cost[base_model]
-        assert base_pricing["input_cost_per_token"] == 1.1e-06
-        assert base_pricing["output_cost_per_token"] == 5.5e-06
-
-        # Verify GovCloud models have different (higher) pricing
-        gov_east_pricing = model_cost[gov_east_model]
-        gov_west_pricing = model_cost[gov_west_model]
-
-        # GovCloud models should have ~20% higher pricing than base models
-        assert gov_east_pricing["input_cost_per_token"] == 1.2e-06
-        assert gov_east_pricing["output_cost_per_token"] == 6e-06
-        assert gov_west_pricing["input_cost_per_token"] == 1.2e-06
-        assert gov_west_pricing["output_cost_per_token"] == 6e-06
-
-        # Verify the pricing difference is approximately 20%
-        assert (
-            abs(
-                gov_east_pricing["input_cost_per_token"]
-                / base_pricing["input_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-        assert (
-            abs(
-                gov_east_pricing["output_cost_per_token"]
-                / base_pricing["output_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-        assert (
-            abs(
-                gov_west_pricing["input_cost_per_token"]
-                / base_pricing["input_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-        assert (
-            abs(
-                gov_west_pricing["output_cost_per_token"]
-                / base_pricing["output_cost_per_token"]
-                - 1.2
-            )
-            < 0.15
-        )
-
-        # Test Claude 3 Haiku pricing
-        base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0"
-        gov_east_haiku_model = (
-            "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0"
-        )
-        gov_west_haiku_model = (
-            "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0"
-        )
-
-        # Verify base Haiku model pricing
-        base_haiku_pricing = model_cost[base_haiku_model]
-        assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07  # 0.00000025
-        assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06  # 0.00000125
-
-        # Verify GovCloud Haiku models have different (higher) pricing
-        gov_east_haiku_pricing = model_cost[gov_east_haiku_model]
-        gov_west_haiku_pricing = model_cost[gov_west_haiku_model]
-
-        # GovCloud Haiku models should have 20% higher pricing than base models
-        assert (
-            gov_east_haiku_pricing["input_cost_per_token"] == 3e-07
-        )  # 0.0000003 (20% higher)
-        assert (
-            gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06
-        )  # 0.0000015 (20% higher)
-        assert (
-            gov_west_haiku_pricing["input_cost_per_token"] == 3e-07
-        )  # 0.0000003 (20% higher)
-        assert (
-            gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06
-        )  # 0.0000015 (20% higher)
-
-        # Verify the pricing difference is exactly 20%
-        assert (
-            gov_east_haiku_pricing["input_cost_per_token"]
-            == base_haiku_pricing["input_cost_per_token"] * 1.2
-        )
-        assert (
-            gov_east_haiku_pricing["output_cost_per_token"]
-            == base_haiku_pricing["output_cost_per_token"] * 1.2
-        )
-        assert (
-            gov_west_haiku_pricing["input_cost_per_token"]
-            == base_haiku_pricing["input_cost_per_token"] * 1.2
-        )
-        assert (
-            gov_west_haiku_pricing["output_cost_per_token"]
-            == base_haiku_pricing["output_cost_per_token"] * 1.2
-        )
-
     @patch("litellm.completion")
     def test_govcloud_completion_cost_calculation(self, mock_completion):
         """Test that completion requests use correct pricing for GovCloud models"""

tests/llm_translation/test_crusoe.py

@@ -4,7 +4,6 @@ Tests for Crusoe provider integration
 import os
 from unittest import mock
 
-import litellm
 
 CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1"
 
@@ -71,38 +70,3 @@ def test_get_llm_provider_crusoe():
     )
     assert model == "meta-llama/Llama-3.3-70B-Instruct"
     assert provider == "crusoe"
-
-
-def test_crusoe_models_configuration():
-    """Test that Crusoe models are configured correctly"""
-    from litellm import get_model_info
-
-    original_model_cost = litellm.model_cost
-    original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP")
-    try:
-        os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-        litellm.model_cost = litellm.get_model_cost_map(url="")
-
-        crusoe_models = [
-            "crusoe/meta-llama/Llama-3.3-70B-Instruct",
-            "crusoe/deepseek-ai/DeepSeek-R1-0528",
-            "crusoe/deepseek-ai/DeepSeek-V3-0324",
-            "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507",
-            "crusoe/moonshotai/Kimi-K2-Thinking",
-            "crusoe/openai/gpt-oss-120b",
-            "crusoe/google/gemma-3-12b-it",
-        ]
-
-        for model in crusoe_models:
-            model_info = get_model_info(model)
-            assert model_info is not None, f"Model info not found for {model}"
-            assert model_info.get("litellm_provider") == "crusoe", (
-                f"{model} should have crusoe as provider"
-            )
-            assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
-    finally:
-        litellm.model_cost = original_model_cost
-        if original_env is None:
-            os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None)
-        else:
-            os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env

tests/llm_translation/test_hyperbolic.py

@@ -1,8 +1,3 @@
-import os
-from datetime import datetime
-from unittest.mock import MagicMock
-
-import pytest
 
 
 import litellm
@@ -69,34 +64,6 @@ def test_hyperbolic_in_provider_lists():
     assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints
 
 
-def test_hyperbolic_models_configuration():
-    """Test that Hyperbolic models are properly configured"""
-    import json
-
-    # Load model configuration directly from the JSON file
-    json_path = os.path.join(
-        os.path.dirname(__file__), "../../model_prices_and_context_window.json"
-    )
-    with open(json_path, "r") as f:
-        model_data = json.load(f)
-
-    # Test a few key models
-    test_models = [
-        "hyperbolic/deepseek-ai/DeepSeek-V3",
-        "hyperbolic/Qwen/Qwen2.5-Coder-32B-Instruct",
-        "hyperbolic/deepseek-ai/DeepSeek-R1",
-    ]
-
-    for model in test_models:
-        assert model in model_data
-        model_info = model_data[model]
-        assert model_info["litellm_provider"] == "hyperbolic"
-        assert model_info["mode"] == "chat"
-        assert "max_tokens" in model_info
-        assert "input_cost_per_token" in model_info
-        assert "output_cost_per_token" in model_info
-
-
 def test_hyperbolic_supported_params():
     """Test that supported OpenAI parameters are correctly configured"""
     from litellm.llms.hyperbolic.chat.transformation import HyperbolicChatConfig

tests/llm_translation/test_lambda_ai.py

@@ -8,7 +8,6 @@ from unittest import mock
 import pytest
 
 import litellm
-from litellm import completion
 from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig
 
 
@@ -103,48 +102,6 @@ async def test_lambda_ai_completion_call():
             raise
 
 
-def test_lambda_ai_models_configuration():
-    """Test that Lambda AI models are configured correctly"""
-    from litellm import get_model_info
-
-    # Reload model cost map to pick up local changes
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-
-    # Clear and repopulate lambda_ai_models list after reloading model_cost
-    litellm.lambda_ai_models = set()
-    litellm.add_known_models()
-
-    # Some Lambda AI models to test
-    lambda_ai_models = [
-        "lambda_ai/deepseek-llama3.3-70b",
-        "lambda_ai/hermes3-8b",
-        "lambda_ai/llama3.1-8b-instruct",
-        "lambda_ai/llama3.2-11b-vision-instruct",
-        "lambda_ai/qwen25-coder-32b-instruct",
-    ]
-
-    for model in lambda_ai_models:
-        model_info = get_model_info(model)
-        assert model_info is not None, f"Model info not found for {model}"
-        assert (
-            model_info.get("litellm_provider") == "lambda_ai"
-        ), f"{model} should have lambda_ai as provider"
-        assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
-        assert (
-            model_info.get("supports_function_calling") is True
-        ), f"{model} should support function calling"
-        assert (
-            model_info.get("supports_system_messages") is True
-        ), f"{model} should support system messages"
-
-        # Check vision support for vision models
-        if "vision" in model:
-            assert (
-                model_info.get("supports_vision") is True
-            ), f"{model} should support vision"
-
-
 def test_lambda_ai_model_list_populated():
     """Test that lambda_ai_models list is populated correctly"""
     # Ensure we're using local model cost map and repopulate models

tests/llm_translation/test_morph.py

@@ -68,24 +68,6 @@ def test_morph_in_provider_lists():
     )
 
 
-def test_morph_model_info():
-    """Test that morph models have correct configuration."""
-    import litellm
-
-    model_info = litellm.get_model_info("morph/morph-v3-large")
-
-    assert model_info["litellm_provider"] == "morph"
-    assert model_info["mode"] == "chat"
-    assert model_info["max_tokens"] == 16000
-    assert model_info["max_input_tokens"] == 16000
-    assert model_info["max_output_tokens"] == 16000
-    assert model_info["input_cost_per_token"] == 9e-07  # $0.9/1M tokens
-    assert model_info["output_cost_per_token"] == 1.9e-06  # $1.9/1M tokens
-    assert model_info["supports_function_calling"] is False
-    assert model_info["supports_vision"] is False
-    assert model_info["supports_system_messages"] is True
-
-
 def test_morph_supported_params():
     """Test that MorphChatConfig returns correct supported parameters."""
     config = MorphChatConfig()

tests/llm_translation/test_openai_o1.py

@@ -1,15 +1,11 @@
-import json
 import os
-from datetime import datetime
-from unittest.mock import AsyncMock, patch, MagicMock
+from unittest.mock import patch
 
 
-
-import httpx
 import pytest
 
 import litellm
-from litellm import Choices, Message, ModelResponse
+from litellm import ModelResponse
 from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest
 
 
@@ -74,7 +70,6 @@ async def test_o1_handle_tool_calling_optional_params(
     - max_tokens is translated to 'max_completion_tokens'
     - role 'system' is translated to 'user'
     """
-    from openai import AsyncOpenAI
     from litellm.utils import ProviderConfigManager
     from litellm.types.utils import LlmProviders
 
@@ -186,15 +181,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest):
         pass
 
 
-def test_o1_supports_vision():
-    """Test that o1 supports vision"""
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    for k, v in litellm.model_cost.items():
-        if k.startswith("o1") and v.get("litellm_provider") == "openai":
-            assert v.get("supports_vision") is True, f"{k} does not support vision"
-
-
 def test_o3_reasoning_effort():
     resp = litellm.completion(
         model="o3-mini",

tests/llm_translation/test_v0.py

@@ -8,7 +8,6 @@ from unittest import mock
 import pytest
 
 import litellm
-from litellm import completion
 from litellm.llms.v0.chat.transformation import V0ChatConfig
 
 
@@ -111,33 +110,3 @@ def test_v0_supported_params():
     ]
 
     assert set(supported_params) == set(expected_params)
-
-
-def test_v0_models_configuration():
-    """Test that v0 models are configured correctly"""
-    from litellm import get_model_info
-
-    # Reload model cost map to pick up local changes
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-
-    # All v0 models
-    v0_models = ["v0/v0-1.0-md", "v0/v0-1.5-md", "v0/v0-1.5-lg"]
-
-    for model in v0_models:
-        model_info = get_model_info(model)
-        assert model_info is not None, f"Model info not found for {model}"
-        # All v0 models support vision (multimodal)
-        assert (
-            model_info.get("supports_vision") is True
-        ), f"{model} should support vision"
-        assert (
-            model_info.get("litellm_provider") == "v0"
-        ), f"{model} should have v0 as provider"
-        assert model_info.get("mode") == "chat", f"{model} should be in chat mode"
-        assert (
-            model_info.get("supports_function_calling") is True
-        ), f"{model} should support function calling"
-        assert (
-            model_info.get("supports_system_messages") is True
-        ), f"{model} should support system messages"

tests/local_testing/test_get_model_info.py

@@ -1,8 +1,6 @@
 # What is this?
 ## Unit testing for the 'get_model_info()' function
 import os
-import traceback
-import json
 
 
 from typing import List, Dict, Any
@@ -11,7 +9,7 @@ import pytest
 
 import litellm
 from litellm import get_model_info
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 
 def test_get_model_info_simple_model_name():
@@ -49,34 +47,6 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch):
     assert model_info["input_cost_per_token"] == 0.0
 
 
-def test_get_model_info_shows_correct_supports_vision():
-    info = litellm.get_model_info("gemini/gemini-2.0-flash")
-    print("info", info)
-    assert info["supports_vision"] is True
-
-
-def test_get_model_info_shows_assistant_prefill():
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    info = litellm.get_model_info("deepseek/deepseek-chat")
-    print("info", info)
-    assert info.get("supports_assistant_prefill") is True
-
-
-def test_get_model_info_shows_supports_prompt_caching():
-    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
-    litellm.model_cost = litellm.get_model_cost_map(url="")
-    info = litellm.get_model_info("deepseek/deepseek-chat")
-    print("info", info)
-    assert info.get("supports_prompt_caching") is True
-
-
-def test_get_model_info_finetuned_models():
-    info = litellm.get_model_info("ft:gpt-3.5-turbo:my-org:custom_suffix:id")
-    print("info", info)
-    assert info["input_cost_per_token"] == 0.000003
-
-
 def test_get_model_info_gemini_pro():
     info = litellm.get_model_info("gemini-2.0-flash")
     print("info", info)
@@ -219,7 +189,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch):
 def test_get_model_info_custom_provider():
     # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server:
     import litellm
-    from litellm import CustomLLM, completion, get_llm_provider
+    from litellm import CustomLLM, completion
 
     class MyCustomLLM(CustomLLM):
         def completion(self, *args, **kwargs) -> litellm.ModelResponse:

tests/local_testing/test_least_busy_routing.py

@@ -33,8 +33,8 @@ def test_model_added():
         }
     }
     least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs)
-    request_count_api_key = f"gpt-3.5-turbo_request_count"
-    assert test_cache.get_cache(key=request_count_api_key) is not None
+    request_count_api_key = "gpt-3.5-turbo_request_count:1234"
+    assert test_cache.get_cache(key=request_count_api_key) == 1
 
 
 def test_get_available_deployments():
@@ -52,8 +52,8 @@ def test_get_available_deployments():
         }
     }
     least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs)
-    request_count_api_key = f"{model_group}_request_count"
-    assert test_cache.get_cache(key=request_count_api_key) is not None
+    request_count_api_key = f"{model_group}_request_count:1234"
+    assert test_cache.get_cache(key=request_count_api_key) == 1
 
 
 # test_get_available_deployments()
@@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test):
     router.leastbusy_logger.test_flag = True
 
     model_group = "azure-model"
-    request_count_dict = {1: 10, 2: 54, 3: 100}
-    cache_key = f"{model_group}_request_count"
+    request_count_dict = {"1": 10, "2": 54, "3": 100}
+    cache_keys = {
+        deployment_id: f"{model_group}_request_count:{deployment_id}"
+        for deployment_id in request_count_dict
+    }
     if async_test is True:
-        await router.cache.async_set_cache(key=cache_key, value=request_count_dict)
+        for deployment_id, count in request_count_dict.items():
+            await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count)
         deployment = await router.async_get_available_deployment(
             model=model_group, messages=None, request_kwargs={}
         )
     else:
-        router.cache.set_cache(key=cache_key, value=request_count_dict)
+        for deployment_id, count in request_count_dict.items():
+            router.cache.set_cache(key=cache_keys[deployment_id], value=count)
         deployment = router.get_available_deployment(model=model_group, messages=None)
     print(f"deployment: {deployment}")
     assert deployment["model_info"]["id"] == "1"
@@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test):
         messages=[{"role": "user", "content": "Hey, how's it going?"}],
     )
 
-    return_dict = router.cache.get_cache(key=cache_key)
-
     # wait 2 seconds
     time.sleep(2)
 
+    return_dict = {
+        deployment_id: router.cache.get_cache(key=cache_key)
+        for deployment_id, cache_key in cache_keys.items()
+    }
+
     assert router.leastbusy_logger.logged_success == 1
-    assert return_dict[1] == 10
-    assert return_dict[2] == 54
-    assert return_dict[3] == 100
+    assert return_dict["1"] == 10
+    assert return_dict["2"] == 54
+    assert return_dict["3"] == 100
 
 
 ## Test with Real calls ##
@@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming():
     await asyncio.sleep(random.uniform(0, 2))
     await router.atext_completion(model=model, prompt=prompt, stream=True)
 
-    cache_key = f"{model}_request_count"
     ## check if calls equally distributed
-    cache_dict = router.cache.get_cache(key=cache_key)
+    cache_dict = {
+        deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}")
+        for deployment_id in ("1", "2", "3")
+    }
     for k, v in cache_dict.items():
         assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}"
 
@@ -259,8 +269,10 @@ async def test_router_completion_streaming():
     await asyncio.sleep(random.uniform(0, 2))
     await router.acompletion(model=model, messages=messages, stream=True)
 
-    cache_key = f"{model}_request_count"
     ## check if calls equally distributed
-    cache_dict = router.cache.get_cache(key=cache_key)
+    cache_dict = {
+        deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}")
+        for deployment_id in ("1", "2", "3")
+    }
     for k, v in cache_dict.items():
         assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}"

tests/local_testing/test_lowest_latency_routing.py

@@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async():
         ), f"Oldest latency {oldest_latency} should have been discarded"
 
 
-def test_ttft_list_trimming_discards_oldest_entry():
-    """
-    The time_to_first_token list trims the oldest entry when full, matching
-    the behavior of the latency list.
-    """
-    max_size = 3
-    test_cache = DualCache()
-    lowest_latency_logger = LowestLatencyLoggingHandler(
-        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
-    )
-
-    model_group = "gpt-3.5-turbo"
-    deployment_id = "test-deployment"
-
-    ttft_values = []
-    for i in range(max_size + 1):
-        start_time = time.time()
-        expected_ttft = float(i + 1) * 0.1  # 0.1, 0.2, 0.3, 0.4
-        completion_start_time = start_time + expected_ttft
-        end_time = start_time + float(i + 1)
-        ttft_values.append(expected_ttft)
-
-        kwargs = {
-            "litellm_params": {
-                "metadata": {
-                    "model_group": model_group,
-                    "deployment": "azure/gpt-4.1-mini",
-                },
-                "model_info": {"id": deployment_id},
-            },
-            "stream": True,
-            "completion_start_time": completion_start_time,
-        }
-        # TTFT is only recorded when response_obj is a ModelResponse.
-        response_obj = litellm.ModelResponse(
-            usage=litellm.Usage(completion_tokens=1, total_tokens=1)
-        )
-
-        lowest_latency_logger.log_success_event(
-            response_obj=response_obj,
-            kwargs=kwargs,
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-    latency_key = f"{model_group}_map"
-    cached_data = test_cache.get_cache(key=latency_key)
-    ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
-
-    assert (
-        len(ttft_list) == max_size
-    ), f"Expected {max_size} entries, got {len(ttft_list)}"
-
-    newest_ttft = ttft_values[-1]
-    oldest_ttft = ttft_values[0]
-    tolerance = 0.05
-
-    assert (
-        abs(ttft_list[-1] - newest_ttft) < tolerance
-    ), f"Newest TTFT {newest_ttft} should be at end of list"
-
-    for ttft in ttft_list:
-        assert (
-            abs(ttft - oldest_ttft) > tolerance
-        ), f"Oldest TTFT {oldest_ttft} should have been discarded"
-
-
 @pytest.mark.asyncio
 async def test_timeout_penalty_discards_oldest_entry():
     """
@@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims():
         assert (
             abs(latency_list[i] - expected) < tolerance
         ), f"At index {i}, expected ~{expected}, got {latency_list[i]}"
-
-
-@pytest.mark.asyncio
-async def test_ttft_list_trimming_discards_oldest_entry_async():
-    """
-    Async counterpart: the time_to_first_token list trims the oldest entry
-    when full. Exercises the async_log_success_event TTFT path, which only
-    runs when response_obj is a ModelResponse and the call is marked as
-    streaming with a completion_start_time.
-    """
-    max_size = 3
-    test_cache = DualCache()
-    lowest_latency_logger = LowestLatencyLoggingHandler(
-        router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
-    )
-
-    model_group = "gpt-3.5-turbo"
-    deployment_id = "test-deployment"
-
-    ttft_values = []
-    for i in range(max_size + 1):
-        start_time = time.time()
-        expected_ttft = float(i + 1) * 0.1  # 0.1, 0.2, 0.3, 0.4
-        completion_start_time = start_time + expected_ttft
-        end_time = start_time + float(i + 1)
-        ttft_values.append(expected_ttft)
-
-        kwargs = {
-            "litellm_params": {
-                "metadata": {
-                    "model_group": model_group,
-                    "deployment": "azure/gpt-4.1-mini",
-                },
-                "model_info": {"id": deployment_id},
-            },
-            "stream": True,
-            "completion_start_time": completion_start_time,
-        }
-        response_obj = litellm.ModelResponse(
-            usage=litellm.Usage(completion_tokens=1, total_tokens=1)
-        )
-
-        await lowest_latency_logger.async_log_success_event(
-            response_obj=response_obj,
-            kwargs=kwargs,
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-    latency_key = f"{model_group}_map"
-    cached_data = await test_cache.async_get_cache(key=latency_key)
-    ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
-
-    assert (
-        len(ttft_list) == max_size
-    ), f"Expected {max_size} entries, got {len(ttft_list)}"
-
-    newest_ttft = ttft_values[-1]
-    oldest_ttft = ttft_values[0]
-    tolerance = 0.05
-
-    assert (
-        abs(ttft_list[-1] - newest_ttft) < tolerance
-    ), f"Newest TTFT {newest_ttft} should be at end of list"
-
-    for ttft in ttft_list:
-        assert (
-            abs(ttft - oldest_ttft) > tolerance
-        ), f"Oldest TTFT {oldest_ttft} should have been discarded"

tests/local_testing/test_redis_increment_with_floor.py

@@ -0,0 +1,80 @@
+"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the
+create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so
+these are the only tests that fail when the script itself is wrong."""
+
+import os
+import uuid
+from typing import Final
+
+import pytest
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from litellm.caching.redis_cache import RedisCache
+
+TTL: Final = 600
+
+
+@pytest.fixture
+def counter():
+    cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT"))
+    key: Final = f"lit7039-{uuid.uuid4()}"
+    yield cache, key, cache.check_and_fix_namespace(key=key)
+    cache.delete_cache(key)
+
+
+def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter):
+    cache, key, _ = counter
+
+    assert cache.increment_with_floor(key, 3, TTL) == 3
+    assert cache.increment_with_floor(key, 2, TTL) == 5
+    assert cache.batch_get_counts([key]) == (5,)
+
+
+def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter):
+    """A worker whose counter expired mid-request decrements a key that is no longer there.
+    Without the clamp that deployment reads negative, and least-busy pins every later request
+    on it until the count climbs back to zero."""
+    cache, key, _ = counter
+
+    assert cache.increment_with_floor(key, 1, TTL) == 1
+    assert cache.increment_with_floor(key, -5, TTL) == 0
+    assert cache.batch_get_counts([key]) == (0,)
+
+
+def test_traffic_never_pushes_a_counters_expiry_back_out(counter):
+    """The TTL is what releases a count whose worker died mid-request. Rewriting it on every
+    touch would keep that stuck count alive for as long as the group takes traffic."""
+    cache, key, namespaced_key = counter
+
+    cache.increment_with_floor(key, 1, TTL)
+    assert cache.redis_client.ttl(namespaced_key) > TTL - 60
+
+    cache.redis_client.expire(namespaced_key, 30)
+    cache.increment_with_floor(key, 1, TTL)
+
+    assert cache.redis_client.ttl(namespaced_key) <= 30
+
+
+def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter):
+    cache, key, namespaced_key = counter
+
+    cache.increment_with_floor(key, 1, TTL)
+    cache.redis_client.expire(namespaced_key, 30)
+
+    assert cache.increment_with_floor(key, -5, TTL) == 0
+    assert cache.redis_client.ttl(namespaced_key) <= 30
+
+
+@pytest.mark.asyncio
+async def test_the_async_counter_behaves_the_same_way(counter):
+    cache, key, namespaced_key = counter
+
+    assert await cache.async_increment_with_floor(key, 2, TTL) == 2
+    assert await cache.async_batch_get_counts([key]) == (2,)
+
+    cache.redis_client.expire(namespaced_key, 30)
+
+    assert await cache.async_increment_with_floor(key, -9, TTL) == 0
+    assert cache.redis_client.ttl(namespaced_key) <= 30

tests/proxy_migration_tests/test_invalid_index_repair.py

@@ -0,0 +1,264 @@
+import os
+import threading
+import uuid
+from collections.abc import Iterator, Mapping
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager
+
+psycopg = pytest.importorskip("psycopg")
+
+pytestmark = pytest.mark.timeout(120)
+
+requires_db: Final = pytest.mark.skipif(
+    "DATABASE_URL" not in os.environ,
+    reason="requires a postgres database (DATABASE_URL)",
+)
+
+HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable"
+HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx"
+HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC'
+LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable"
+LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx"
+PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable"
+PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx"
+
+
+def _base_url() -> str:
+    return os.environ["DATABASE_URL"].split("?")[0]
+
+
+def _index_validity(schema: str) -> Mapping[str, bool]:
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        rows = conn.execute(
+            "SELECT c.relname, i.indisvalid FROM pg_index i "
+            "JOIN pg_class c ON c.oid = i.indexrelid "
+            "JOIN pg_namespace n ON n.oid = c.relnamespace "
+            "WHERE n.nspname = %s",
+            (schema,),
+        ).fetchall()
+    return MappingProxyType(dict(rows))
+
+
+def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None:
+    """Abort a CONCURRENTLY build while it waits on an older snapshot, the same
+    spot the deadlock loser dies at, so it leaves its index INVALID."""
+    with psycopg.connect(_base_url()) as pin:
+        pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
+        pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
+        with psycopg.connect(_base_url(), autocommit=True) as builder:
+            builder.execute("SET statement_timeout = '1s'")
+            with pytest.raises(psycopg.errors.QueryCanceled):
+                builder.execute(statement)
+
+
+def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None:
+    _interrupt_concurrent_build(
+        schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})'
+    )
+
+
+def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None:
+    _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"')
+
+
+@pytest.fixture
+def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
+    schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        conn.execute(f'CREATE SCHEMA "{schema}"')
+        conn.execute(
+            f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)'
+        )
+        conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)')
+
+    monkeypatch.delenv("DIRECT_URL", raising=False)
+    monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}")
+    yield schema
+
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        conn.execute(f'DROP SCHEMA "{schema}" CASCADE')
+
+
+@pytest.fixture
+def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
+    """A brand-new database, what a first deploy sees. A scratch schema would
+    not do: the migrations guard on pg_constraint by name across every schema,
+    so a LiteLLM schema already pushed into public makes them skip and then
+    fail, which is exactly what CI's database looks like."""
+    admin_url: Final = _base_url()
+    name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
+    with psycopg.connect(admin_url, autocommit=True) as conn:
+        conn.execute(f'CREATE DATABASE "{name}"')
+
+    monkeypatch.delenv("DIRECT_URL", raising=False)
+    monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}")
+    yield "public"
+
+    with psycopg.connect(admin_url, autocommit=True) as conn:
+        conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)')
+
+
+@requires_db
+def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None:
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+    _leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id")
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False}
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False}
+
+
+@requires_db
+def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None:
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+    _leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX)
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"')
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"')
+    before: Final = _index_validity(scratch_schema)
+    assert len(before) == 4
+    assert set(before.values()) == {False}
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+@requires_db
+def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None:
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})')
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+@requires_db
+def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None:
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)')
+        conn.execute(
+            f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" '
+            f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)'
+        )
+        conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)')
+    assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+
+    assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
+
+
+@requires_db
+def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None:
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+
+    with psycopg.connect(_base_url(), autocommit=True) as other_replica:
+        other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,))
+        assert ProxyExtrasDBManager.repair_invalid_indexes() is False
+        assert _index_validity(scratch_schema) == {HEALTH_INDEX: False}
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+@requires_db
+def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None:
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+
+    with psycopg.connect(_base_url()) as pin:
+        pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
+        pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"')
+        assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False
+        blocked: Final = _index_validity(scratch_schema)
+        assert blocked[HEALTH_INDEX] is False
+        assert [name for name in blocked if name.endswith("_ccnew")]
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None:
+    with psycopg.connect(_base_url()) as pin:
+        pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
+        pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
+        pinned.set()
+        pin.execute("SELECT pg_sleep(%s)", (seconds,))
+
+
+@requires_db
+def test_repair_outlives_a_statement_timeout_passed_through_database_url_options(
+    scratch_schema: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+    monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000")
+    pinned: Final = threading.Event()
+    holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0))
+    holder.start()
+    pinned.wait()
+    try:
+        assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+    finally:
+        holder.join()
+
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+@requires_db
+def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None:
+    table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}"
+    index: Final = f"{table}_id_idx"
+    monkeypatch.setenv("DATABASE_URL", _base_url())
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        conn.execute(f'CREATE TABLE public."{table}" (id TEXT)')
+    try:
+        _leave_invalid_index("public", table, index, "id")
+        assert _index_validity("public")[index] is False
+
+        assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+
+        assert _index_validity("public")[index] is True
+    finally:
+        with psycopg.connect(_base_url(), autocommit=True) as conn:
+            conn.execute(f'DROP TABLE public."{table}"')
+
+
+def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("DIRECT_URL", raising=False)
+    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever")
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is False
+
+
+@requires_db
+def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None:
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+    with pytest.MonkeyPatch.context() as env:
+        env.setenv("DIRECT_URL", f"{_base_url()}?schema=public")
+        env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}")
+        assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+def _invalidate_deployed_index(schema: str) -> None:
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"')
+    _leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+
+
+@requires_db
+@pytest.mark.timeout(300)
+@pytest.mark.parametrize("use_v2_resolver", [True, False])
+def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None:
+    assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
+    _invalidate_deployed_index(fresh_database)
+    assert _index_validity(fresh_database)[HEALTH_INDEX] is False
+
+    assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
+
+    assert _index_validity(fresh_database)[HEALTH_INDEX] is True

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -1760,6 +1760,35 @@ class TestUnmanagedVertexRouting:
         )
         router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash")
 
+    def test_flag_on_routes_fine_tuned_endpoint_to_vertex_deployment(self):
+        """A fine-tuned Gemini batch stores `endpoints/<id>` in the gs:// path; the bare model
+        (the endpoint id) must round-trip to the deployment configured as
+        `vertex_ai/gemini/<id>` (LIT-6899)."""
+        endpoint_id = "7768560373388541952"
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = None
+        router.get_model_list.return_value = [
+            {
+                "model_name": "gemini-2.5-flash-dts-usc1",
+                "litellm_params": {
+                    "model": f"vertex_ai/gemini/{endpoint_id}",
+                    "custom_llm_provider": "vertex_ai",
+                },
+                "model_info": {"id": "deploy-ft"},
+            },
+        ]
+        instance = self._instance(track_unmanaged=True, router=router)
+        job = self._job(
+            file_object=_unmanaged_vertex_file_object(
+                input_file_id=f"gs://bucket/litellm-vertex-files/endpoints/{endpoint_id}/abc.jsonl"
+            )
+        )
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(job, MagicMock())
+
+        assert result == ("deploy-ft", "8823717160934178816")
+
     def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self):
         """Flag on, but the only deployment for the model group is a non-vertex_ai
         provider: must not be selected, even though the model group name matches."""

tests/router_unit_tests/test_router_cooldown_per_deployment.py

@@ -246,12 +246,12 @@ class TestCooldownCacheTTLCorrection:
             "timestamp": time.time() - 120.0,
             "cooldown_time": 60.0,
         }
-        cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600)
+        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
 
         active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
 
         assert active == [], "Expired cooldown entry must not appear in active cooldowns"
-        assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
+        assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
 
     def test_active_entry_is_returned(self):
         """
@@ -267,7 +267,7 @@ class TestCooldownCacheTTLCorrection:
             "timestamp": time.time(),
             "cooldown_time": 60.0,
         }
-        cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60)
+        cc.in_memory_cache.set_cache(key, active_value, ttl=60)
 
         active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
 
@@ -290,14 +290,14 @@ class TestCooldownCacheTTLCorrection:
             "timestamp": time.time() - (60.0 - remaining),
             "cooldown_time": 60.0,
         }
-        cc.cache.in_memory_cache.set_cache(key, value, ttl=600)
+        cc.in_memory_cache.set_cache(key, value, ttl=600)
 
-        before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
+        before_expiry = cc.in_memory_cache.ttl_dict.get(key)
         assert before_expiry is not None
 
         cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
 
-        after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
+        after_expiry = cc.in_memory_cache.ttl_dict.get(key)
         assert after_expiry is not None
         corrected_remaining = after_expiry - time.time()
         assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
@@ -318,12 +318,12 @@ class TestCooldownCacheTTLCorrection:
             "timestamp": time.time() - 120.0,
             "cooldown_time": 60.0,
         }
-        cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600)
+        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
 
         active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
 
         assert active == [], "Expired entry must not appear in async active cooldowns"
-        assert cc.cache.in_memory_cache.get_cache(key) is None
+        assert cc.in_memory_cache.get_cache(key) is None
 
 
 class TestFallbackDeploymentCooldown:

tests/test_litellm/batches/test_main.py

@@ -158,6 +158,14 @@ def test_create__vertex_ai_dispatch(seams):
     _assert_only(seams.vertex.create_batch, seams, "create_batch")
 
 
+def test_create__vertex_ai_forwards_custom_endpoint(seams):
+    """The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher
+    must forward the flag for the handler to act on."""
+    bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
+
+    assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True
+
+
 def test_create__provider_config_routes_to_base_http_handler(seams):
     """model + a provider batches config (bedrock-style) routes to the generic
     base_llm_http_handler, NOT the per-provider instance."""

tests/test_litellm/caching/test_redis_cache.py

@@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re
     assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {}
 
 
+def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache):
+    """A caller that must fall back when Redis is unreachable needs the failure, not zeros.
+
+    The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell
+    apart from "every counter is unset". Least-busy routing read that as an idle deployment and
+    kept sending traffic to it instead of falling back to this worker's own in-flight counts.
+    """
+    assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {}
+
+    with pytest.raises(OSError, match="redis unavailable"):
+        sync_batch_redis_cache.batch_get_counts(["lit7039"])
+
+
+@pytest.mark.asyncio
+async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None):
+    """Async twin: the async batch read hides the same failure behind an empty dict."""
+    failing_client = AsyncMock()
+    failing_client.mget.side_effect = OSError("redis unavailable")
+    with patch(  # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
+        "litellm._redis.get_redis_client", return_value=MagicMock()
+    ):
+        cache = RedisCache(host="127.0.0.1", port=6379)
+
+    with patch.object(cache, "init_async_client", return_value=failing_client):
+        assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {}
+
+        with pytest.raises(OSError, match="redis unavailable"):
+            await cache.async_batch_get_counts(["lit7039"])
+
+
+@pytest.mark.parametrize("stored", [b"3", "3"])
+def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None):
+    """Counters come back positionally, so an unset key has to stay a hole rather than shift the
+    rest of the row onto the wrong deployments, and a count has to survive whether the client
+    hands it back as bytes or as text."""
+    with patch(  # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
+        "litellm._redis.get_redis_client", return_value=MagicMock()
+    ):
+        cache = RedisCache(host="127.0.0.1", port=6379)
+    cache.redis_client.mget.return_value = [stored, None, b"0"]
+
+    assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0)
+
+
 @pytest.fixture
 def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]:
     service_logger = ServiceLogging(mock_testing=True)