Merge pull request #50136 from BerriAI/litellm_backport_1108_picks_lts_1_108_x

fix(backport): backport 21 main fixes plus 2 prerequisites to lts/1.108.x for v1.108.1

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta, timezone
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
+from pydantic import TypeAdapter, ValidationError
 
 from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
     ManagedObjectCreatorAttribution,
@@ -46,6 +47,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool:
 
 CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
 RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
+_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
 
 PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
     "completed",
@@ -500,6 +502,10 @@ class CheckBatchCost:
         llm_provider: str,
         bare_model_name: str,
     ) -> Optional[Tuple[str, str]]:
+        stored_model_id: Final = self._get_stored_model_id(job)
+        if stored_model_id is not None and self._batch_deployment_exists(stored_model_id):
+            return stored_model_id, job.unified_object_id
+
         deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
         if deployment_id is None:
             verbose_proxy_logger.info(
@@ -611,23 +617,36 @@ class CheckBatchCost:
         )
 
     @staticmethod
-    def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
-        from litellm.types.utils import LiteLLMBatch
-
+    def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]:
         file_object: Final = job.file_object
         try:
-            parsed_file_object: Final[object] = (
+            return _FILE_OBJECT_ADAPTER.validate_python(
                 json.loads(file_object) if isinstance(file_object, str) else file_object
             )
-        except (json.JSONDecodeError, ValueError):
+        except (json.JSONDecodeError, ValueError, ValidationError):
             return None
-        if not isinstance(parsed_file_object, dict):
+
+    @classmethod
+    def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]:
+        from litellm.types.utils import LiteLLMBatch
+
+        parsed_file_object: Final = cls._parse_file_object(job)
+        if parsed_file_object is None:
             return None
         try:
             return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
         except Exception:
             return None
 
+    @classmethod
+    def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]:
+        """The deployment that created the batch, written into file_object by the managed
+        files hook at create time. Older rows have no such key and fall back to matching
+        the bare model name against the router."""
+        parsed_file_object: Final = cls._parse_file_object(job)
+        model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None
+        return model_id if isinstance(model_id, str) and model_id else None
+
     async def _track_completed_batch_cost(
         self,
         job: "ManagedObjectRow",

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -19,6 +19,7 @@ from typing import (
     TypedDict,
     Union,
     cast,
+    runtime_checkable,
 )
 from uuid import NAMESPACE_URL, uuid5
 
@@ -366,6 +367,12 @@ class _ManagedResourcePrismaClient(Protocol):
     def db(self) -> _ManagedResourceDatabase: ...
 
 
+@runtime_checkable
+class _WriterRoutedDatabase(Protocol):
+    @property
+    def writer(self) -> _ManagedResourceDatabase: ...
+
+
 class _SchedulerWithJobLookup(Protocol):
     def get_job(self, job_id: str) -> object: ...
 
@@ -398,6 +405,15 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
     return prisma_client.db.litellm_managedobjecttable
 
 
+def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
+    return prisma_client.db
+
+
+def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase:
+    """Reads that feed a write must see the create that just happened; a lagging read replica may not"""
+    return db.writer if isinstance(db, _WriterRoutedDatabase) else db
+
+
 _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
 
@@ -431,6 +447,25 @@ def _without_model_output(
     )
 
 
+def _with_model_id(
+    file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
+    model_id: str | None,
+) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
+    return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id})
+
+
+async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None:
+    row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id})
+    if row is None:
+        return None
+    try:
+        parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object))
+    except (json.JSONDecodeError, ValidationError):
+        return None
+    model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None
+    return model_id if isinstance(model_id, str) and model_id else None
+
+
 def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
     hidden_params: Final = cast(  # cast-ok: _hidden_params is an untyped attribute the upload path sets
         "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@@ -599,9 +634,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         persist_attribution: bool = False,
         create_if_missing: bool = True,
         batch_processed: bool = False,
+        creating_model_id: str | None = None,
     ) -> None:
         """Persist a managed object row, caching it and upserting it in the DB.
 
+        creating_model_id is the router deployment that created the object. It is kept
+        inside the stored file_object JSON as `model_id`, because the raw provider id an
+        unmanaged batch stores as unified_object_id says nothing about which deployment
+        made it, and the cost poller needs that deployment to read the result back. A
+        later observe of the row (create_if_missing=False) keeps the stored value rather
+        than overwriting it with whichever deployment served the retrieve.
+
         persist_attribution is set only by the batch and background response creates,
         the callers that can speak for the creator; it gates the api_key, request_tags, and,
         for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
@@ -622,7 +665,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         It is written only in the upsert create branch.
         """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
-        stored_object: Final = _without_model_output(file_object)
+        managed_database: Final = _managed_database(self.prisma_client)
+        managed_object_table: Final = managed_database.litellm_managedobjecttable
+        preserved_model_id: Final = (
+            creating_model_id
+            if creating_model_id is not None or create_if_missing
+            else await _stored_model_id(
+                _writer_database(managed_database).litellm_managedobjecttable, unified_object_id
+            )
+        )
+        stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id)
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
             model_object_id=model_object_id,
@@ -654,7 +706,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             "status": file_object.status,
             "updated_by": user_api_key_dict.user_id,
         }
-        managed_object_table: Final = _managed_object_table(self.prisma_client)
         if not create_if_missing:
             await managed_object_table.update_many(
                 where={"unified_object_id": unified_object_id},
@@ -2069,6 +2120,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}),
                 persist_attribution=is_batch_create,
                 create_if_missing=is_batch_create,
+                creating_model_id=model_id if is_batch_create else None,
             )
             if not is_batch_create:
                 await self.store_batch_output_file_ownership(

enterprise/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-enterprise"
-version = "0.1.79"
+version = "0.1.79.post1"
 description = "Package for LiteLLM Enterprise features"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.1.79"
+version = "0.1.79.post1"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-enterprise==",

litellm/caching/evicted_client_closer.py

@@ -29,6 +29,9 @@ it does not belong to. Queued clients are therefore bucketed by what it takes to
 close them, and each bucket is ordered by deadline, so a reap walks the entries
 that are due rather than the whole queue.
 
+A bucket whose loop has been collected is dropped, and its clients are left to the
+collector, as they are when the queue is full.
+
 The queue holds its clients weakly, so waiting out a grace window never keeps
 alive anything the collector would have reclaimed first.
 """
@@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first.
 import asyncio
 import contextlib
 import inspect
+import itertools
 import threading
 import time
 import weakref
@@ -53,6 +57,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere"
 _CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
 
 _BucketKey = str | int
+_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary()
+_LOOP_TOKEN_COUNTER: Final = itertools.count()
+_LOOP_TOKEN_LOCK: Final = threading.Lock()
 
 
 @dataclass(frozen=True, slots=True)
@@ -63,12 +70,12 @@ class _PendingClose:
     collector would otherwise have reclaimed first.
 
     ``needs_loop`` is set for a client whose close is a coroutine; those can only
-    be closed from the event loop they were evicted on, recorded in ``loop_id``.
+    be closed from the event loop they were evicted on, recorded in ``loop_token``.
     A client that closes synchronously carries neither constraint.
     """
 
     client_ref: "weakref.ref[object]"
-    loop_id: int | None
+    loop_token: int | None
     needs_loop: bool
     close_after: float
 
@@ -77,16 +84,24 @@ def _bucket_key(pending: _PendingClose) -> _BucketKey:
     """Which reaps can close this entry: any at all, any running a loop, or one loop's."""
     if not pending.needs_loop:
         return _CLOSABLE_ANYWHERE
-    if pending.loop_id is None:
+    if pending.loop_token is None:
         return _CLOSABLE_ON_ANY_LOOP
-    return pending.loop_id
+    return pending.loop_token
 
 
-def _running_loop_id() -> int | None:
+def running_loop_token() -> int | None:
     try:
-        return id(asyncio.get_running_loop())
+        loop: Final = asyncio.get_running_loop()
     except RuntimeError:
         return None
+    try:
+        token: Final[int | None] = _LOOP_TOKENS.get(loop)
+        if token is not None:
+            return token
+        with _LOOP_TOKEN_LOCK:
+            return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER))
+    except TypeError:
+        return id(loop)
 
 
 def _close_function(client: object) -> Callable[[], object] | None:
@@ -167,6 +182,8 @@ class EvictedClientCloser:
         self._clock = clock
         self._owned: weakref.WeakSet[object] = weakref.WeakSet()
         self._buckets: dict[_BucketKey, deque[_PendingClose]] = {}  # mutable-ok: deadline-ordered queues
+        self._dead_loop_tokens: deque[int] = deque()  # mutable-ok: loop finalizers append without taking the queue lock
+        self._watched_loop_tokens: set[int] = set()  # mutable-ok: tracks tokens until their dead buckets are discarded
         self._pending_count = 0
         self._queue_lock = threading.Lock()  # the cache is reachable from every worker thread's loop
         self._close_tasks: set[asyncio.Task[None]] = set()  # mutable-ok: strong refs to running closes
@@ -197,13 +214,19 @@ class EvictedClientCloser:
         close_fn: Final = _close_function(client)
         if close_fn is None:
             return
+        loop_token: Final = running_loop_token()
+        needs_loop: Final = inspect.iscoroutinefunction(close_fn)
         if self._pending_count >= self._max_pending:
-            return
+            self._discard_dead_loops()
+            if self._pending_count >= self._max_pending:
+                return
+        if needs_loop and loop_token is not None:
+            self._watch_loop(loop_token)
         self._enqueue(
             _PendingClose(
                 client_ref=weakref.ref(client),
-                loop_id=_running_loop_id(),
-                needs_loop=inspect.iscoroutinefunction(close_fn),
+                loop_token=loop_token,
+                needs_loop=needs_loop,
                 close_after=self._clock() + self._grace_seconds,
             )
         )
@@ -217,7 +240,7 @@ class EvictedClientCloser:
         if not self._pending_count:
             return
         now: Final = self._clock()
-        for pending in self._take_due(_running_loop_id(), now):
+        for pending in self._take_due(running_loop_token(), now):
             client = pending.client_ref()
             if client is None:
                 continue
@@ -230,6 +253,29 @@ class EvictedClientCloser:
     def pending_count(self) -> int:
         return self._pending_count
 
+    def _watch_loop(self, loop_token: int) -> None:
+        if loop_token in self._watched_loop_tokens:
+            return
+        try:
+            weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token)
+        except TypeError:
+            return
+        self._watched_loop_tokens.add(loop_token)
+
+    def _discard_dead_loops(self) -> None:
+        with self._queue_lock:
+            self._discard_dead_loops_locked()
+
+    def _discard_dead_loops_locked(self) -> None:
+        while self._dead_loop_tokens:
+            self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft())
+
+    def _discard_dead_loop_token_locked(self, loop_token: int) -> None:
+        self._watched_loop_tokens.discard(loop_token)
+        bucket: Final = self._buckets.pop(loop_token, None)
+        if bucket is not None:
+            self._pending_count -= len(bucket)
+
     def _enqueue(self, pending: _PendingClose) -> None:
         """Append to the entry's bucket, dropping any dead entries it queues behind.
 
@@ -238,6 +284,7 @@ class EvictedClientCloser:
         the front rather than having to be searched for.
         """
         with self._queue_lock:
+            self._discard_dead_loops_locked()
             bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
             while bucket and bucket[0].client_ref() is None:
                 bucket.popleft()
@@ -245,9 +292,12 @@ class EvictedClientCloser:
             bucket.append(pending)
             self._pending_count += 1
 
-    def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
-        buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
+    def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]:
+        buckets = (
+            (_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token)
+        )
         with self._queue_lock:
+            self._discard_dead_loops_locked()
             return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
 
     def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:

litellm/caching/llm_caching_handler.py

@@ -2,10 +2,9 @@
 Add the event loop to the cache key, to prevent event loop closed errors.
 """
 
-import asyncio
 from typing import Final
 
-from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
+from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token
 from .in_memory_cache import InMemoryCache
 
 
@@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache):
         self.evicted_client_closer.schedule(evicted)
         self.evicted_client_closer.reap()
 
-    def update_cache_key_with_event_loop(self, key):
+    def update_cache_key_with_event_loop(self, key: str) -> str:
         """
         Add the event loop to the cache key, to prevent event loop closed errors.
         If none, use the key as is.
         """
-        try:
-            event_loop: Final = asyncio.get_running_loop()
-            stringified_event_loop: Final = str(id(event_loop))
-            return f"{key}-{stringified_event_loop}"
-        except RuntimeError:  # handle no current running event loop
-            return key
+        token: Final = running_loop_token()
+        return f"{key}-{token}" if token is not None else key
 
-    def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
+    def _scoped_key(self, key: str, loop_bound: bool) -> str:
+        return self.update_cache_key_with_event_loop(key) if loop_bound else key
+
+    def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **kwargs):
         """``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
         if litellm_owned_client:
             self.evicted_client_closer.mark_owned(value)
-        key = self.update_cache_key_with_event_loop(key)
-        return super().set_cache(key, value, **kwargs)
+        return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs)
 
     async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
         if litellm_owned_client:
@@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache):
         key = self.update_cache_key_with_event_loop(key)
         return await super().async_set_cache(key, value, **kwargs)
 
-    def get_cache(self, key, **kwargs):
-        key = self.update_cache_key_with_event_loop(key)
+    def get_cache(self, key: str, loop_bound: bool = True, **kwargs):
+        scoped_key: Final = self._scoped_key(key, loop_bound)
         self.evicted_client_closer.reap()
 
-        return super().get_cache(key, **kwargs)
+        return super().get_cache(scoped_key, **kwargs)
 
-    async def async_get_cache(self, key, **kwargs):
+    async def async_get_cache(self, key: str, **kwargs):
         key = self.update_cache_key_with_event_loop(key)
 
         return await super().async_get_cache(key, **kwargs)

litellm/constants.py

@@ -446,6 +446,7 @@ REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_sp
 REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer"
 REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
 REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
+REDIS_TOOL_DISCOVERY_BUFFER_KEY: Final = "litellm_tool_discovery_buffer"
 MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)
 REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer"
 REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000
@@ -701,6 +702,7 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
 )
 
 LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED: Final = 499
+LITELLM_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
 
 EMAIL_BUDGET_ALERT_TTL: Final = get_env_int("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)  # 24 hours in seconds
 EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = get_env_float(
@@ -1810,6 +1812,7 @@ LITELLM_UI_SESSION_DURATION: Final = os.getenv("LITELLM_UI_SESSION_DURATION", "2
 
 ########################### DB CRON JOB NAMES ###########################
 DB_SPEND_UPDATE_JOB_NAME: Final = "db_spend_update_job"
+TOOL_DISCOVERY_COMMIT_JOB_NAME: Final = "tool_discovery_commit_job"
 DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job"
 PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics"
 CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data"
@@ -1985,6 +1988,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30
 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000
 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000
+SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS: Final = 30
 SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100
 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated
 # callers from forcing a DB query per request for unknown names, while bounding

litellm/cost_calculator.py

@@ -42,6 +42,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
     generic_cost_per_token,
     get_batch_cost_rates,
     get_billable_input_tokens,
+    get_cost_per_unit,
+    get_modality_token_cost_key,
     get_regional_uplift_multiplier,
     get_service_tier_cost_key,
     get_token_type_cost_breakdown,
@@ -2832,12 +2834,17 @@ def _batch_rate(
     model_info: ModelInfo,
     key: Literal[
         "input_cost_per_audio_token_batches",
+        "cache_read_input_audio_token_cost_batches",
         "input_cost_per_image_token_batches",
         "input_cost_per_video_token_batches",
     ],
     fallback: float,
+    *,
+    prompt_tokens: int = 0,
 ) -> float:
-    rate: Final = model_info.get(key)
+    context_key: Final = get_modality_token_cost_key(model_info, key, prompt_tokens)
+    context_rate: Final = get_cost_per_unit(model_info, context_key, None)
+    rate: Final = model_info.get(key) if context_rate is None else context_rate
     return fallback if rate is None else rate
 
 
@@ -2965,6 +2972,7 @@ def _batch_prompt_cost(
 ) -> float:
     details: Final = parse_prompt_tokens_details(usage)
     cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0
+    cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0
     written_tokens: Final = details["cache_creation_tokens"] if cache_creation_rate is not None else 0
     audio_tokens, image_tokens, video_tokens = (
         details["audio_tokens"],
@@ -2976,10 +2984,13 @@ def _batch_prompt_cost(
     )
     return (
         text_tokens * input_rate
-        + audio_tokens * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate)
+        + audio_tokens
+        * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate, prompt_tokens=usage.prompt_tokens)
         + image_tokens * _batch_rate(model_info, "input_cost_per_image_token_batches", input_rate)
         + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate)
-        + cached_tokens * (cache_read_rate or 0.0)
+        + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0)
+        + cached_audio_tokens
+        * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0)
         + written_tokens * (cache_creation_rate or 0.0)
     )
 

litellm/images/main.py

@@ -29,7 +29,6 @@ from litellm.litellm_core_utils.mock_functions import mock_image_generation
 from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from litellm.llms.custom_llm import CustomLLM
 from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params
 
 #################### Initialize provider clients ####################
@@ -523,10 +522,7 @@ def image_generation(
             )
         elif custom_llm_provider in litellm._custom_providers:  # Assume custom LLM provider
             # Get the Custom Handler
-            custom_handler: CustomLLM | None = None
-            for item in litellm.custom_provider_map:
-                if item["provider"] == custom_llm_provider:
-                    custom_handler = item["custom_handler"]
+            custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider)
 
             if custom_handler is None:
                 raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)
@@ -799,10 +795,7 @@ def image_edit(
 
         # Check for custom provider
         if custom_llm_provider in litellm._custom_providers:
-            custom_handler: CustomLLM | None = None
-            for item in litellm.custom_provider_map:
-                if item["provider"] == custom_llm_provider:
-                    custom_handler = item["custom_handler"]
+            custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider)
 
             if custom_handler is None:
                 raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)

litellm/integrations/agentops/agentops.py

@@ -99,18 +99,14 @@ class AgentOps(OpenTelemetry):
             "Connection": "keep-alive",
         }
 
-        client: Final = get_httpx_client()
-        try:
-            response: Final = client.post(
-                url=auth_endpoint,
-                headers=headers,
-                json={"api_key": api_key},
-                timeout=10,
-            )
+        response: Final = get_httpx_client().post(
+            url=auth_endpoint,
+            headers=headers,
+            json={"api_key": api_key},
+            timeout=10,
+        )
 
-            if response.status_code != 200:
-                raise Exception(f"Failed to fetch auth token: {response.text}")
+        if response.status_code != 200:
+            raise Exception(f"Failed to fetch auth token: {response.text}")
 
-            return response.json()
-        finally:
-            client.close()
+        return response.json()

litellm/litellm_core_utils/hidden_params.py

@@ -77,6 +77,11 @@ def get_hidden_params(obj: object) -> MutableMapping[str, object] | None:
     return _as_hidden_params_mapping(hidden_params)
 
 
+def served_from_cache(obj: object) -> bool:
+    hidden_params: Final = get_hidden_params(obj)
+    return hidden_params is not None and hidden_params.get("cache_hit") is True
+
+
 def set_hidden_params(obj: object, hidden_params: dict[str, object] | HiddenParams) -> None:
     if isinstance(obj, dict):
         obj[_HIDDEN_PARAMS_ATTR] = hidden_params

litellm/litellm_core_utils/litellm_logging.py

@@ -81,14 +81,19 @@ from litellm.litellm_core_utils.core_helpers import (
 )
 from litellm.litellm_core_utils.error_normalization import normalize_error
 from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
-from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_param
+from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, served_from_cache, set_hidden_param
 from litellm.litellm_core_utils.internal_call_metadata import (
     MODEL_ACCESS_GROUP_METADATA_KEY,
     is_unbilled_non_inference_call,
 )
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
+    bedrock_guardrail_trace,
     cost_breakdown_with_guardrail,
     guardrail_information_cost,
+    is_str_keyed_dict,
+    native_bedrock_guardrail_information,
+    object_mapping,
+    with_guardrail_entry,
 )
 from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
     StandardBuiltInToolCostTracking,
@@ -2478,6 +2483,31 @@ class Logging(LiteLLMLoggingBaseClass):
             return
         self.model_call_details["response_headers"] = headers
 
+    def record_native_bedrock_guardrail(self, result: object) -> None:
+        """Bill the guardrail trace a Bedrock response carries when the deployment's ``guardrailConfig`` ran it, by
+        recording it as guardrail information on this call, the same entry the proxy-side Bedrock guardrail writes.
+        A cache hit carries the original response's trace but made no Bedrock call, so it bills nothing."""
+        trace: Final = bedrock_guardrail_trace(result)
+        if trace is None or served_from_cache(result) or self.model_call_details.get("cache_hit") is True:
+            return
+        optional_params: Final = object_mapping(getattr(self, "optional_params", None)) or {}
+        guardrail_config: Final = object_mapping(optional_params.get("guardrailConfig")) or {}
+        identifier: Final = guardrail_config.get("guardrailIdentifier")
+        litellm_params_in: Final = object_mapping(getattr(self, "litellm_params", None)) or {}
+        region: Final = optional_params.get("aws_region_name") or litellm_params_in.get("aws_region_name")
+        entry: Final = native_bedrock_guardrail_information(
+            trace,
+            aws_region_name=region if isinstance(region, str) else None,
+            guardrail_identifier=identifier if isinstance(identifier, str) else None,
+        )
+        litellm_params: Final[object] = self.model_call_details.get("litellm_params")
+        if entry is None or not is_str_keyed_dict(litellm_params):
+            return
+        metadata: Final[object] = litellm_params.get("metadata")
+        bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {}
+        litellm_params["metadata"] = bucket
+        with_guardrail_entry(bucket, entry)
+
     def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None:
         """
         Copy response._hidden_params into litellm_params.metadata['hidden_params'].
@@ -2664,7 +2694,14 @@ class Logging(LiteLLMLoggingBaseClass):
                 start_time = self.start_time
             if end_time is None:
                 end_time = datetime.datetime.now()
-            if self.completion_start_time is None:
+            # A request re-dispatched to a fallback reuses this logging object but runs in
+            # its own window, so a first-token time from the previous attempt cannot belong
+            # to it. Keeping it makes the attempt report a negative TTFT.
+            first_token_seconds: Final = _timestamp_seconds(self.completion_start_time)
+            attempt_start_seconds: Final = _timestamp_seconds(start_time)
+            if first_token_seconds is None or (
+                attempt_start_seconds is not None and first_token_seconds < attempt_start_seconds
+            ):
                 self.completion_start_time = end_time
                 self.model_call_details["completion_start_time"] = self.completion_start_time
 

litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py

@@ -1,15 +1,24 @@
 import math
-from collections.abc import Mapping
+from collections.abc import Iterator, Mapping, Sequence
+from functools import partial
+from itertools import chain
 from typing import Annotated, Final
 
-from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
+from typing_extensions import TypeIs  # noqa: TID251  # TypeIs reaches typing only on 3.13
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.core_helpers import (
+    redact_nested_match_and_regex_keys,  # pyright: ignore[reportUnknownVariableType]  # signature uses bare dict
+)
+from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.llms.base import LiteLLMBaseModel
-from litellm.types.utils import CostBreakdown
+from litellm.types.utils import CostBreakdown, GuardrailStatus, StandardLoggingGuardrailInformation
 
 BEDROCK_GUARDRAIL_PRICING_KEY: Final = "bedrock/guardrails"
+BEDROCK_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace")
+_INTERVENING_ACTIONS: Final = frozenset({"BLOCKED", "ANONYMIZED"})
 
 
 class GuardrailPricing(LiteLLMBaseModel):
@@ -96,6 +105,161 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str
     return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name))
 
 
+_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object])
+_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...])
+
+
+def object_mapping(value: object) -> Mapping[str, object] | None:
+    """``value`` as a str-keyed mapping, None for anything else. Vendor JSON arrives untyped, so this is the one place
+    that turns it into something the rest of the module can read without guessing."""
+    try:
+        return _OBJECT_MAPPING.validate_python(value)
+    except ValidationError:
+        return None
+
+
+def object_items(value: object) -> tuple[object, ...]:
+    if isinstance(value, (str, bytes, Mapping)):
+        return ()
+    try:
+        return _OBJECT_ITEMS.validate_python(value)
+    except ValidationError:
+        return ()
+
+
+def is_str_keyed_dict(
+    value: object,
+) -> TypeIs[dict[str, object]]:  # guard-ok: JSON object keys are str  # mutable-ok: the caller edits this dict in place
+    return isinstance(value, dict)
+
+
+def _trace_field(response: object, field: str) -> Mapping[str, object] | None:
+    source: Final = object_mapping(response)
+    trace: Final = object_mapping(source.get(field) if source is not None else getattr(response, field, None))
+    if trace is None or object_mapping(trace.get("guardrail")) is None:
+        return None
+    return trace
+
+
+def bedrock_guardrail_trace(response: object) -> Mapping[str, object] | None:
+    """The guardrail trace a Bedrock response carries: Converse puts it under ``trace``, Invoke under
+    ``amazon-bedrock-trace``. Returned as is, or None when the response has no guardrail section."""
+    traces: Final = map(partial(_trace_field, response), BEDROCK_TRACE_FIELDS)
+    return next((trace for trace in traces if trace is not None), None)
+
+
+class _Assessment(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="allow")
+    invocation_metrics: Mapping[str, JsonValue] = Field(default_factory=dict, alias="invocationMetrics")
+
+    @property
+    def billable_usage(self) -> Mapping[str, int]:
+        usage: Final = self.invocation_metrics.get("usage")
+        if not isinstance(usage, dict):
+            return {}
+        return {
+            counter: units
+            for counter, units in usage.items()
+            if counter.endswith("Units") and isinstance(units, int) and not isinstance(units, bool)
+        }
+
+
+class _GuardrailTrace(BaseModel):
+    """Both trace shapes: Converse ``inputAssessment`` {id: a} and ``outputAssessments`` {id: [a, ...]}, Invoke
+    ``input`` {id: a} and ``outputs`` [{id: a}, ...]."""
+
+    model_config = ConfigDict(frozen=True, extra="allow")
+    input_assessment: Mapping[str, _Assessment] = Field(default_factory=dict, alias="inputAssessment")
+    input: Mapping[str, _Assessment] = Field(default_factory=dict)
+    output_assessments: Mapping[str, tuple[_Assessment, ...]] = Field(default_factory=dict, alias="outputAssessments")
+    outputs: tuple[Mapping[str, _Assessment], ...] = ()
+
+    @property
+    def assessments(self) -> tuple[tuple[str, _Assessment], ...]:
+        listed: Final = (
+            (gid, a) for gid, items in self.output_assessments.items() for a in items
+        )  # comprehension-ok: id then its list
+        return (
+            *self.input_assessment.items(),
+            *self.input.items(),
+            *listed,
+            *chain.from_iterable(by_id.items() for by_id in self.outputs),
+        )
+
+
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+def _summed_usage(usages: Sequence[Mapping[str, int]]) -> Mapping[str, int]:
+    counters: Final = tuple(dict.fromkeys(chain.from_iterable(usages)))
+    return {counter: sum(usage.get(counter, 0) for usage in usages) for counter in counters}
+
+
+def _actions(node: JsonValue) -> Iterator[str]:
+    if isinstance(node, list):
+        yield from chain.from_iterable(map(_actions, node))
+        return
+    if not isinstance(node, dict):
+        return
+    action: Final = node.get("action")
+    if isinstance(action, str):
+        yield action
+    yield from chain.from_iterable(map(_actions, node.values()))
+
+
+def native_bedrock_guardrail_information(
+    trace: Mapping[str, object], aws_region_name: str | None, guardrail_identifier: str | None
+) -> StandardLoggingGuardrailInformation | None:
+    """The guardrail entry for a trace Bedrock returned because the deployment's ``guardrailConfig`` ran with
+    ``trace: enabled``: every assessment's ``invocationMetrics.usage`` summed and priced like the proxy-side Bedrock
+    guardrail. None when the trace carries no usage, which is what ``trace: disabled`` leaves."""
+    try:
+        guardrail: Final = _JSON_OBJECT.validate_python(trace.get("guardrail"))
+        pairs: Final = _GuardrailTrace.model_validate(guardrail).assessments
+    except ValidationError:
+        return None
+    usage: Final = _summed_usage(tuple(assessment.billable_usage for _, assessment in pairs))
+    if not usage:
+        return None
+    cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage, aws_region_name)
+    name: Final = guardrail_identifier or pairs[0][0]
+    status: Final[GuardrailStatus] = (
+        "guardrail_intervened" if any(action in _INTERVENING_ACTIONS for action in _actions(guardrail)) else "success"
+    )
+    priced: Final[StandardLoggingGuardrailInformation] = (
+        {"guardrail_cost_by_unit": cost_by_unit} if cost_by_unit is not None else {}
+    )
+    return {
+        "guardrail_name": name,
+        "guardrail_id": name,
+        "guardrail_provider": "bedrock",
+        "guardrail_mode": GuardrailEventHooks.during_call,
+        "guardrail_response": redact_nested_match_and_regex_keys(guardrail),
+        "guardrail_status": status,
+        "guardrail_usage": usage,
+        "guardrail_cost": guardrail_cost_total(cost_by_unit),
+        **priced,
+    }
+
+
+def guardrail_entries(raw: object) -> tuple[object, ...]:
+    """The entries under ``standard_logging_guardrail_information``, which older writers stored as one dict."""
+    single: Final = object_mapping(raw)
+    return (single,) if single is not None else object_items(raw)
+
+
+def with_guardrail_entry(
+    metadata: dict[str, object],  # mutable-ok: the proxy's own request metadata, read later for headers and spend
+    entry: StandardLoggingGuardrailInformation,
+) -> None:
+    """Append ``entry`` to the metadata's ``standard_logging_guardrail_information`` once, so the same trace recorded
+    from the response headers path and again from the success handler bills a single time."""
+    existing: Final = guardrail_entries(metadata.get("standard_logging_guardrail_information"))
+    if entry in existing:
+        return
+    metadata["standard_logging_guardrail_information"] = [*existing, entry]  # rebind-ok: appending here is the job
+
+
 AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"
 
 

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -1088,21 +1088,55 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
     )
 
 
+def get_modality_token_cost_key(
+    model_info: ModelInfo,
+    cost_key: str,
+    prompt_tokens: int,
+    *,
+    threshold_is_inclusive: bool = False,
+) -> str:
+    suffix: Final = next((suffix for suffix in _NON_STANDARD_THRESHOLD_SUFFIXES if cost_key.endswith(suffix)), "")
+    base_key: Final = cost_key.removesuffix(suffix)
+    threshold_key: Final = max(
+        (
+            key
+            for key in model_info
+            if key.startswith(f"{base_key}_above_")
+            and key.endswith(("_tokens", f"_tokens{suffix}"))
+            and model_info.get(key) is not None
+            and _prompt_exceeds_threshold(prompt_tokens, _parse_above_token_threshold(key), threshold_is_inclusive)
+        ),
+        key=_parse_above_token_threshold,
+        default=None,
+    )
+    if threshold_key is None:
+        return cost_key
+    return f"{threshold_key.removesuffix(suffix)}{suffix}"
+
+
 def _modality_token_cost(
     model_info: ModelInfo,
     modality_cost_key: str,
     tokens: int,
     prompt_base_cost: float,
     service_tier: str | None,
+    *,
+    prompt_tokens: int = 0,
+    threshold_is_inclusive: bool = False,
 ) -> float:
     """
     Audio, image and video prompt tokens bill at their modality rate when the model has one.
     Without one they are ordinary prompt tokens, so they bill at `prompt_base_cost`, which
     `_get_token_base_cost` has already resolved for the service tier and long-context tier.
     """
-    tiered_key: Final = get_service_tier_cost_key(modality_cost_key, service_tier)
-    cost_key: Final = tiered_key if model_info.get(tiered_key) is not None else modality_cost_key
-    if model_info.get(cost_key) is None:
+    cost_key: Final = get_modality_token_cost_key(
+        model_info,
+        get_service_tier_cost_key(modality_cost_key, service_tier),
+        prompt_tokens,
+        threshold_is_inclusive=threshold_is_inclusive,
+    )
+    standard_key: Final = cost_key.removesuffix(get_service_tier_cost_key("", service_tier))
+    if model_info.get(cost_key) is None and model_info.get(standard_key) is None:
         return float(tokens) * prompt_base_cost
     return calculate_cost_component(model_info, cost_key, tokens)
 
@@ -1115,6 +1149,9 @@ def _calculate_input_cost(
     cache_creation_cost: float,
     cache_creation_cost_above_1hr: float,
     service_tier: str | None = None,
+    *,
+    prompt_tokens: int = 0,
+    threshold_is_inclusive: bool = False,
 ) -> float:
     """
     Calculates the input cost for a given model, prompt tokens, and completion tokens.
@@ -1149,6 +1186,8 @@ def _calculate_input_cost(
             prompt_tokens_details["audio_tokens"],
             prompt_base_cost,
             service_tier,
+            prompt_tokens=prompt_tokens,
+            threshold_is_inclusive=threshold_is_inclusive,
         )
 
     ### IMAGE TOKEN COST
@@ -1467,6 +1506,8 @@ def generic_cost_per_token(
         cache_creation_cost=cache_creation_cost,
         cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
         service_tier=service_tier,
+        prompt_tokens=usage.prompt_tokens,
+        threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
     )
 
     ## CALCULATE OUTPUT COST
@@ -1524,11 +1565,9 @@ def generic_cost_per_token(
 
     ## IMAGE COST
     if not is_text_tokens_total and image_tokens and image_tokens > 0:
-        _output_cost_per_image_token = get_cost_per_unit(resolved_model_info, "output_cost_per_image_token", None)
-        _output_cost_per_image_token = (
-            _output_cost_per_image_token if _output_cost_per_image_token is not None else completion_base_cost
-        )
-        completion_cost += float(image_tokens) * _output_cost_per_image_token
+        image_cost_key: Final = get_service_tier_cost_key("output_cost_per_image_token", service_tier)
+        image_rate: Final = get_cost_per_unit(resolved_model_info, image_cost_key, None)
+        completion_cost += float(image_tokens) * (completion_base_cost if image_rate is None else image_rate)
 
     ## VIDEO COST
     if not is_text_tokens_total and video_tokens and video_tokens > 0:

litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py

@@ -5,6 +5,7 @@ from typing import Final
 from pydantic import TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
+from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
 from litellm.types.utils import (
     CompletionTokensDetailsWrapper,
     PromptTokensDetailsWrapper,
@@ -16,6 +17,8 @@ ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total"
 
 _TEXT_INPUT_RATE: Final = "input_cost_per_token"
 _AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token"
+_CACHE_READ_INPUT_RATE: Final = "cache_read_input_token_cost"
+_CACHE_READ_AUDIO_INPUT_RATE: Final = "cache_read_input_audio_token_cost"
 _TEXT_OUTPUT_RATE: Final = "output_cost_per_token"
 _AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token"
 _RATE_KEY_MARKERS: Final = ("cost", "pricing")
@@ -33,19 +36,38 @@ def _tokens(value: object) -> int:
 
 
 def used_pricing_keys(usage: Usage) -> tuple[str, ...]:
-    prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details)
+    prompt_details: Final = parse_prompt_tokens_details(usage)
+    prompt_audio: Final = prompt_details["audio_tokens"]
     completion_audio: Final = _audio_tokens(usage.completion_tokens_details)
-    prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio
+    prompt_text: Final = max(
+        _tokens(usage.prompt_tokens)
+        - prompt_details["cache_hit_tokens"]
+        - prompt_audio
+        - prompt_details["image_tokens"]
+        - prompt_details["video_tokens"],
+        0,
+    )
+    cached_text: Final = max(prompt_details["cache_hit_tokens"] - prompt_details["cache_hit_audio_tokens"], 0)
     completion_text: Final = _tokens(usage.completion_tokens) - completion_audio
     components: Final = (
         (_TEXT_INPUT_RATE, prompt_text),
         (_AUDIO_INPUT_RATE, prompt_audio),
+        (_CACHE_READ_INPUT_RATE, cached_text),
+        (_CACHE_READ_AUDIO_INPUT_RATE, prompt_details["cache_hit_audio_tokens"]),
         (_TEXT_OUTPUT_RATE, completion_text),
         (_AUDIO_OUTPUT_RATE, completion_audio),
     )
     return tuple(key for key, count in components if count > 0)
 
 
+def _effective_pricing_key(key: str, pricing_entry: Mapping[str, object]) -> str:
+    if key == _CACHE_READ_AUDIO_INPUT_RATE and pricing_entry.get(key) is None:
+        return _CACHE_READ_INPUT_RATE if pricing_entry.get(_CACHE_READ_INPUT_RATE) is not None else _TEXT_INPUT_RATE
+    if key == _CACHE_READ_INPUT_RATE and pricing_entry.get(key) is None:
+        return _TEXT_INPUT_RATE
+    return key
+
+
 def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None:
     try:
         return _NESTED_PRICING.validate_python(value)
@@ -89,7 +111,8 @@ def diagnose_zero_cost(
     pricing_entry: Mapping[str, object],
     calculation_failed: bool,
 ) -> StandardLoggingZeroCostDiagnostic | None:
-    used_keys: Final = used_pricing_keys(usage)
+    effective_keys: Final = tuple(_effective_pricing_key(key, pricing_entry) for key in used_pricing_keys(usage))
+    used_keys: Final = tuple(key for index, key in enumerate(effective_keys) if key not in effective_keys[:index])
     if not used_keys:
         return None
     missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None)

litellm/litellm_core_utils/redact_messages.py

@@ -14,6 +14,8 @@ from collections.abc import Mapping
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any, Final
 
+from pydantic import JsonValue, TypeAdapter, ValidationError
+
 import litellm
 from litellm.constants import REDACTED_BY_LITELLM
 from litellm.integrations.custom_logger import CustomLogger
@@ -72,11 +74,55 @@ def _redacted_streaming_response_copy(streaming_response):
     return redacted_response
 
 
+BEDROCK_GUARDRAIL_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace")
+_MATCHED_TEXT_KEYS: Final = frozenset({"match", "regex"})
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+def _without_matched_text(trace: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+    redacted: Final = dict(copy.deepcopy(trace))
+    stack: Final[list[JsonValue]] = [redacted]  # mutable-ok: iterative walk, recursive_detector bans recursion here
+    while stack:
+        node = stack.pop()
+        if isinstance(node, dict):
+            for key in _MATCHED_TEXT_KEYS & node.keys():
+                node[key] = "[REDACTED]"
+            stack.extend(node.values())
+        elif isinstance(node, list):
+            stack.extend(node)
+    return redacted
+
+
+def _redacted_bedrock_trace(trace: object) -> object:
+    try:
+        return _without_matched_text(_JSON_OBJECT.validate_python(trace))
+    except ValidationError:
+        return trace
+
+
+def _redact_bedrock_guardrail_trace(response: object) -> None:
+    """A Bedrock guardrail trace echoes the matched text of its sensitive-information and word findings back."""
+    mapping: Final = _as_mapping(response)
+    if isinstance(response, dict):
+        traces: Final = {
+            field: _redacted_bedrock_trace(value)
+            for field, value in (mapping or {}).items()
+            if field in BEDROCK_GUARDRAIL_TRACE_FIELDS
+        }
+        response.update(traces)
+        return
+    for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS):
+        if value is not None:
+            setattr(response, field, _redacted_bedrock_trace(value))
+
+
 def _redact_streaming_response(streaming_response):
     if hasattr(streaming_response, "choices"):
         for choice in streaming_response.choices:
             _redact_choice_content(choice)
         redact_vertex_ai_metadata_from_logged_object(streaming_response)
+        if isinstance(streaming_response, litellm.ModelResponse):
+            _redact_bedrock_guardrail_trace(streaming_response)
     elif hasattr(streaming_response, "output"):
         _redact_responses_api_output(streaming_response.output)
         if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None:
@@ -234,6 +280,7 @@ def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str,
             if isinstance(response.get("choices"), list):
                 _redact_model_response_dict_choices(response["choices"], redacted_str)
             redact_vertex_ai_metadata_from_logged_object(response)
+            _redact_bedrock_guardrail_trace(response)
         elif isinstance(response, str):
             standard_logging_object["response"] = redacted_str
         else:
@@ -393,11 +440,13 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
                 for choice in _result.choices:
                     _redact_choice_content(choice)
             redact_vertex_ai_metadata_from_logged_object(_result)
+            _redact_bedrock_guardrail_trace(_result)
         elif isinstance(_result, dict) and "choices" in _result:
             # Handle dict representation of ModelResponse (e.g., from model_dump())
             if _result.get("choices") is not None:
                 _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM)
             redact_vertex_ai_metadata_from_logged_object(_result)
+            _redact_bedrock_guardrail_trace(_result)
         elif isinstance(_result, dict) and "output" in _result:
             if isinstance(_result.get("output"), list):
                 _redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM)

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, Required
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_content_from_model_response,  # pyright: ignore[reportUnknownVariableType]  # response param accepts an untyped dict
 )
@@ -1389,6 +1390,35 @@ def _simple_text_part(choices: Sequence[object]) -> str | None:
     return content if isinstance(content, str) else ""
 
 
+_NO_TRACE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+def _carry_bedrock_guardrail_trace(chunks: Sequence[object], response: ModelResponse) -> None:
+    """Bedrock reports a guardrail trace once per stream, on the final metadata chunk, so the built response carries it
+    under the same field as the non-streaming response (``trace`` on Converse, ``amazon-bedrock-trace`` on Invoke)."""
+    traced: Final = map(_bedrock_trace_fields, reversed(chunks))
+    carried: Final = next((fields for fields in traced if fields), _NO_TRACE_FIELDS)
+    for field, trace in carried.items():
+        setattr(response, field, trace)
+
+
+def _record_bedrock_guardrail_trace(
+    chunks: Sequence[object], response: ModelResponse, logging_obj: Optional["Logging"]
+) -> None:
+    _carry_bedrock_guardrail_trace(chunks, response)
+    if logging_obj is not None:
+        logging_obj.record_native_bedrock_guardrail(response)
+
+
+def _bedrock_trace_fields(chunk: object) -> Mapping[str, object]:
+    as_mapping: Final = object_mapping(chunk)
+    raw: Final = (
+        as_mapping.get("provider_specific_fields") if as_mapping else getattr(chunk, "provider_specific_fields", None)
+    )
+    fields: Final[Mapping[str, object]] = object_mapping(raw) or {}
+    return {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields}
+
+
 def stream_chunk_builder(
     chunks: list,
     messages: Sequence | None = None,
@@ -1465,6 +1495,11 @@ def stream_chunk_builder(
                     )
                     break
 
+            _record_bedrock_guardrail_trace(
+                object_items(chunks),  # pyright: ignore[reportUnknownArgumentType]  # chunks is an untyped list param
+                response,
+                logging_obj,
+            )
             _stamp_streaming_usage_cost(usage, response, logging_obj)
             _set_stream_builder_response_cost(response, logging_obj)
 
@@ -1652,6 +1687,11 @@ def stream_chunk_builder(
         _stamp_streaming_usage_cost(usage, response, logging_obj)
         _set_stream_builder_response_cost(response, logging_obj)
 
+        _record_bedrock_guardrail_trace(
+            object_items(chunks),  # pyright: ignore[reportUnknownArgumentType]  # chunks is an untyped list param
+            response,
+            logging_obj,
+        )
         processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
         return response
     except Exception as e:

litellm/litellm_core_utils/streaming_handler.py

@@ -1431,11 +1431,7 @@ class CustomStreamWrapper:
                     setattr(
                         model_response,
                         "usage",
-                        litellm.Usage(
-                            prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None,
-                            completion_tokens=response_obj["usage"].get("completion_tokens", None) or None,
-                            total_tokens=response_obj["usage"].get("total_tokens", None) or None,
-                        ),
+                        _usage_from_payload(response_obj["usage"]),
                     )
                 elif isinstance(response_obj["usage"], Usage):
                     setattr(
@@ -1452,6 +1448,8 @@ class CustomStreamWrapper:
         return _ProviderChunkParsed(response_obj)
 
     def chunk_creator(self, chunk: Any):
+        if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields:
+            self.logging_obj.record_native_bedrock_guardrail(chunk.provider_specific_fields)
         if self.custom_llm_provider != "cached_response" and isinstance(
             chunk, (ChatCompletionChunk, ModelResponseStream)
         ):
@@ -1572,6 +1570,8 @@ class CustomStreamWrapper:
                 model_response.usage = Usage(**chunk_usage.model_dump())
 
             ## RETURN ARG
+            if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields:
+                self.copy_model_response_level_provider_specific_fields(chunk, model_response)
             result: Final = self.return_processed_chunk_logic(
                 completion_obj=completion_obj,
                 model_response=model_response,
@@ -2478,6 +2478,11 @@ class CustomStreamWrapper:
     _strip_sse_data_from_chunk = strip_sse_data_from_chunk
 
 
+def _usage_from_payload(payload: Mapping[str, Any]) -> Usage:
+    """Build ``Usage`` from a wire-shaped usage dict, forwarding every key."""
+    return litellm.Usage(**payload)
+
+
 def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int:
     for key in keys:
         value = getattr(details, key, None)

litellm/litellm_core_utils/token_counter.py

@@ -51,6 +51,7 @@ from litellm.types.llms.anthropic import (
     AnthropicMessagesTextParam,
     AnthropicMessagesToolResultParam,
     AnthropicMessagesToolUseParam,
+    AnthropicSearchResultParam,
 )
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -64,6 +65,10 @@ from litellm.types.llms.openai import (
 from litellm.types.utils import Message, SelectTokenizerResponse
 
 
+class UnsupportedContentTypeError(ValueError):
+    pass
+
+
 def get_modified_max_tokens(
     model: str,
     base_model: str,
@@ -783,6 +788,7 @@ def _validate_anthropic_content(content: Mapping[str, object]) -> type:
     mapping: Final = {
         "tool_use": AnthropicMessagesToolUseParam,
         "tool_result": AnthropicMessagesToolResultParam,
+        "search_result": AnthropicSearchResultParam,
     }
 
     expected_cls: Final = mapping.get(content_type) if isinstance(content_type, str) else None
@@ -938,7 +944,9 @@ def _count_anthropic_content(
         except Exception as e:
             if default_token_count is not None:
                 return default_token_count
-            raise ValueError(f"Error counting field '{field_name}': {e}")
+            raise (UnsupportedContentTypeError if isinstance(e, UnsupportedContentTypeError) else ValueError)(
+                f"Error counting field '{field_name}': {e}"
+            )
     return tokens
 
 
@@ -952,6 +960,7 @@ def _count_content_list(
         | AnthropicMessagesTextParam
         | AnthropicMessagesImageParam
         | AnthropicMessagesDocumentParam
+        | AnthropicSearchResultParam
     ],
     use_default_image_token_count: bool,
     default_token_count: int | None,
@@ -986,7 +995,7 @@ def _count_content_list(
                     count_function,
                     use_default_image_token_count,
                 )
-            elif c["type"] in ("tool_use", "tool_result"):
+            elif c["type"] in ("tool_use", "tool_result", "search_result"):
                 num_tokens += _count_anthropic_content(
                     c,
                     count_function,
@@ -1012,17 +1021,18 @@ def _count_content_list(
                     num_tokens += count_function(tool_name)
             else:
                 content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__
-                raise ValueError(
+                raise UnsupportedContentTypeError(
                     f"Invalid content item type: {content_type}. "
                     f"Expected str or dict with 'type' field "
                     f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
-                    f"tool_reference)."
+                    f"tool_reference, search_result)."
                 )
         return num_tokens
     except Exception as e:
         if default_token_count is not None:
             return default_token_count
-        raise ValueError(
+        error_type: Final = UnsupportedContentTypeError if isinstance(e, UnsupportedContentTypeError) else ValueError
+        raise error_type(
             f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}"
         )
 

litellm/llms/bedrock/chat/converse_transformation.py

@@ -2813,6 +2813,8 @@ class AmazonConverseConfig(BaseConfig):
         # Add "trace" from Bedrock guardrails - if user has opted in to returning it
         if "trace" in completion_response:
             setattr(model_response, "trace", completion_response["trace"])
+            if logging_obj is not None:
+                logging_obj.record_native_bedrock_guardrail(model_response)
 
         # Add service_tier if present in Bedrock response
         # Map Bedrock serviceTier (object) to OpenAI service_tier (string)

litellm/llms/bedrock/chat/invoke_handler.py

@@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache
 from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
 from litellm.litellm_core_utils.core_helpers import map_finish_reason
 from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping
 from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
 from litellm.litellm_core_utils.optional_imports import ensure_optional_import
 from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data
@@ -21,6 +22,7 @@ from litellm.llms.anthropic.chat.handler import (
 )
 from litellm.llms.anthropic.chat.transformation import JsonModeRequest
 from litellm.llms.base_llm.base_utils import ResponseFormatSchema
+from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -951,8 +953,19 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
             response_format_schema=response_format_schema,
         )
 
-    def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
-        return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
+    def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream:
+        parsed: Final = self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
+        trace: Final = object_mapping(chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD))
+        if trace is None:
+            return parsed
+        return parsed.model_copy(
+            update={
+                "provider_specific_fields": {
+                    **(parsed.provider_specific_fields or {}),
+                    BEDROCK_INVOKE_TRACE_FIELD: trace,
+                }
+            }
+        )
 
 
 class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder):

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

@@ -5,6 +5,7 @@ import httpx
 from pydantic import TypeAdapter
 
 from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping
 from litellm.litellm_core_utils.prompt_templates.factory import (
     convert_to_anthropic_image_obj,
 )
@@ -22,6 +23,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
     AmazonInvokeConfig,
 )
 from litellm.llms.bedrock.common_utils import (
+    BEDROCK_INVOKE_TRACE_FIELD,
     apply_bedrock_invoke_structured_output,
     bedrock_supports_tool_search,
     get_anthropic_beta_from_headers,
@@ -377,7 +379,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
         api_key: str | None = None,
         json_mode: bool | None = None,
     ) -> ModelResponse:
-        return AnthropicConfig.transform_response(
+        response: Final = AnthropicConfig.transform_response(
             self,
             model=model,
             raw_response=raw_response,
@@ -391,3 +393,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
             api_key=api_key,
             json_mode=json_mode,
         )
+        body: Final = object_mapping(raw_response.json()) or {}
+        trace: Final = object_mapping(body.get(BEDROCK_INVOKE_TRACE_FIELD))
+        if trace is not None:
+            setattr(response, BEDROCK_INVOKE_TRACE_FIELD, trace)
+            logging_obj.record_native_bedrock_guardrail(response)
+        return response

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

@@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None:
     return headers.get(name, None)
 
 
-def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
+def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
     try:
         guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config)
     except ValidationError as e:
@@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
         guardrail_headers: Final = (
             ()
             if raw_guardrail_config is None
-            else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
+            else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
         )
         owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
         return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)

litellm/llms/bedrock/common_utils.py

@@ -42,6 +42,8 @@ if TYPE_CHECKING:
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
     from litellm.types.llms.openai import AllMessageValues
 
+BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace"
+
 
 _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"
 _OPENAI_FAMILY_MODEL_RE: Final = re.compile(r"(^|[./])openai\.")

litellm/llms/bedrock/cost_calculation.py

@@ -3,9 +3,13 @@ Helper util for handling bedrock-specific cost calculation
 - e.g.: prompt caching
 """
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Final
 
-from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
+import litellm
+from litellm.litellm_core_utils.llm_cost_calc.utils import (
+    generic_cost_per_token,
+    get_provider_specific_geo_multiplier,
+)
 
 if TYPE_CHECKING:
     from litellm.types.utils import Usage
@@ -17,9 +21,13 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
 
     Follows the same logic as Anthropic's cost per token calculation.
     """
-    return generic_cost_per_token(
+    model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="bedrock")
+    prompt_cost, completion_cost = generic_cost_per_token(
         model=model,
         usage=usage,
         custom_llm_provider="bedrock",
         service_tier=service_tier,
+        model_info=model_info,
     )
+    geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
+    return prompt_cost * geo_multiplier, completion_cost * geo_multiplier

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

@@ -32,6 +32,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import (
 from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
 from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
     AmazonInvokeConfig,
+    bedrock_invoke_guardrail_headers,
 )
 from litellm.llms.bedrock.common_utils import (
     BedrockError,
@@ -114,8 +115,14 @@ class AmazonAnthropicClaudeMessagesConfig(
         api_key: str | None = None,
         api_base: str | None = None,
     ) -> tuple[dict, str | None]:
+        raw_guardrail_config: Final = litellm_params.get("guardrailConfig")  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]  # litellm_params is an untyped dict param
+        guardrail_headers: Final = (
+            ()
+            if raw_guardrail_config is None
+            else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items())  # pyright: ignore[reportUnknownArgumentType]  # read from the untyped dict param
+        )
         owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
-        return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base
+        return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names), api_base  # pyright: ignore[reportUnknownArgumentType]  # headers is an untyped dict param
 
     def sign_request(
         self,

litellm/llms/custom_httpx/aiohttp_transport.py

@@ -6,6 +6,7 @@ import ssl
 import sys
 import typing
 import urllib.request
+import weakref
 from collections.abc import Callable, Generator
 from typing import ClassVar, Final
 
@@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]:
         raise mapped_exc(message) from exc
 
 
+def _release_transport_anchor(_transport: object) -> None:
+    """Release a transport anchored to a response stream. See ``_anchor_transport_to``."""
+
+
+def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None:
+    """Keep the transport alive for as long as the stream can still read.
+
+    The transport finalizer closes the session it owns, and nothing on the path from a
+    live response back to that session references the transport, so a client dropped
+    mid-stream would close the session under the body. ``weakref.finalize`` holds the
+    transport in its own registry rather than on the stream, which keeps the transport
+    out of the response's reference cycle (an exception carrying the response, for one)
+    so it is finalized by refcount once the anchor drops and can still schedule an async
+    close, instead of being collected in the same cyclic pass as the session it closes.
+    """
+    weakref.finalize(stream, _release_transport_anchor, transport)
+
+
 class AiohttpResponseStream(httpx.AsyncByteStream):
     CHUNK_SIZE = 1024 * 16
 
@@ -289,6 +308,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
         cls._background_close_tasks.add(task)
         task.add_done_callback(cls._on_close_task_done)
 
+    def __del__(self) -> None:
+        session: Final = getattr(self, "client", None)
+        if not getattr(self, "_owns_session", False) or not isinstance(session, ClientSession):
+            return
+        with contextlib.suppress(Exception):
+            self._close_recycled_session(session)
+
     def get_valid_client_session(self) -> ClientSession:
         """
         Helper to get a valid ClientSession for the current event loop.
@@ -456,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
                 # Re-raise if it's a different RuntimeError
                 raise
 
+        stream: Final = AiohttpResponseStream(response)
+        _anchor_transport_to(stream, self)
         return httpx.Response(
             status_code=response.status,
             headers=response.headers,
-            stream=AiohttpResponseStream(response),
+            stream=stream,
             request=request,
         )
 

litellm/llms/custom_httpx/http_handler.py

@@ -1123,6 +1123,17 @@ class AsyncHTTPHandler:
     # Mirrors LiteLLMAiohttpTransport._background_close_tasks.
     _finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set()  # mutable-ok: strong refs for pending closes
 
+    @classmethod
+    async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None:
+        try:
+            await client.aclose()
+        except (OSError, RuntimeError) as exc:
+            verbose_logger.debug("Error closing client at finalization: %s", exc)
+        finally:
+            task: Final = asyncio.current_task()
+            if task is not None:
+                cls._finalizer_close_tasks.discard(task)
+
     @classmethod
     def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
         cls._finalizer_close_tasks.discard(task)
@@ -1190,8 +1201,8 @@ class AsyncHTTPHandler:
                 # here is the cross-loop path the transport refuses.
                 self._dispose_wrapped_aiohttp_session()
                 return
-            task: Final = loop.create_task(self._client.aclose())
             cls: Final = type(self)
+            task: Final = loop.create_task(cls._close_client_at_finalization(self._client))
             cls._finalizer_close_tasks.add(task)
             task.add_done_callback(cls._on_finalizer_close_done)
         except Exception:
@@ -1823,7 +1834,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
         cache = LLMClientCache()
         setattr(litellm, "in_memory_llm_clients_cache", cache)
 
-    _cached_client: Final = cache.get_cache(_cache_key_name)
+    _cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False)
     if _cached_client:
         return _cached_client
 
@@ -1839,6 +1850,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
         value=_new_client,
         ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
         litellm_owned_client=True,
+        loop_bound=False,
     )
     return _new_client
 

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2430,6 +2430,7 @@ class BaseLLMHTTPHandler:
                 raw_response=response,
                 logging_obj=logging_obj,
             )
+            logging_obj.record_native_bedrock_guardrail(initial_response)
 
         return await self._finalize_anthropic_messages_response(
             initial_response=initial_response,
@@ -7582,7 +7583,7 @@ class BaseLLMHTTPHandler:
         """
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -7777,7 +7778,7 @@ class BaseLLMHTTPHandler:
         """
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -7957,7 +7958,7 @@ class BaseLLMHTTPHandler:
         """
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -8120,7 +8121,7 @@ class BaseLLMHTTPHandler:
     ):
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -8260,7 +8261,7 @@ class BaseLLMHTTPHandler:
     ):
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -8432,7 +8433,7 @@ class BaseLLMHTTPHandler:
     ):
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -8618,7 +8619,7 @@ class BaseLLMHTTPHandler:
     ):
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -8750,7 +8751,7 @@ class BaseLLMHTTPHandler:
         """
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -8831,7 +8832,7 @@ class BaseLLMHTTPHandler:
         """
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:
@@ -9015,7 +9016,7 @@ class BaseLLMHTTPHandler:
         """
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider),
+                llm_provider=custom_llm_provider,
                 params={"ssl_verify": litellm_params.get("ssl_verify", None)},
             )
         else:

litellm/llms/openai/common_utils.py

@@ -253,6 +253,10 @@ def build_output_token_limit_response(
     return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk))
 
 
+def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool:
+    return client_initialization_params.get("is_async") is not False
+
+
 class BaseOpenAILLM:
     """
     Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings
@@ -267,7 +271,9 @@ class BaseOpenAILLM:
             client_initialization_params=client_initialization_params,
             client_type=client_type,
         )
-        _cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
+        _cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(
+            _cache_key, loop_bound=_is_loop_bound(client_initialization_params)
+        )
         return _cached_client
 
     @staticmethod
@@ -308,6 +314,7 @@ class BaseOpenAILLM:
             value=openai_client,
             ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
             litellm_owned_client=litellm_owned_client,
+            loop_bound=_is_loop_bound(client_initialization_params),
         )
 
     @staticmethod

litellm/llms/openai/openai.py

@@ -11,13 +11,14 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
 
 import openai
-from openai import AsyncOpenAI, OpenAI
+from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
 from openai._base_client import make_request_options
 from openai._constants import RAW_RESPONSE_HEADER
 from openai._legacy_response import LegacyAPIResponse
 from openai._types import RequestOptions
 from openai.types import CreateEmbeddingResponse
 from openai.types.beta.assistant_deleted import AssistantDeleted
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from openai.types.file_deleted import FileDeleted
 from pydantic import BaseModel, TypeAdapter
 from typing_extensions import overload
@@ -352,7 +353,14 @@ _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({})
 _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body"))
 
 
-def _embedding_request_without_sdk_defaults(
+def _raw_chat_response(
+    result: ChatCompletion | Stream[ChatCompletionChunk] | AsyncStream[ChatCompletionChunk],
+) -> LegacyAPIResponse[ChatCompletion]:
+    assert isinstance(result, LegacyAPIResponse)
+    return result
+
+
+def _request_without_sdk_transform(
     data: Mapping[str, object], timeout: float | httpx.Timeout
 ) -> tuple[Mapping[str, object], RequestOptions]:
     body: Final = {k: v for k, v in data.items() if k not in _SDK_OPTION_KEYS}
@@ -515,7 +523,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         """
         start_time: Final = time.time()
         try:
-            raw_response = await openai_aclient.chat.completions.with_raw_response.create(**data, timeout=timeout)
+            body, options = _request_without_sdk_transform(data, timeout)
+            raw_response = _raw_chat_response(
+                await openai_aclient.post(
+                    "/chat/completions",
+                    body=body,
+                    cast_to=ChatCompletion,
+                    options=options,
+                    stream=bool(data.get("stream")),
+                    stream_cls=AsyncStream[ChatCompletionChunk],
+                )
+            )
             end_time = time.time()
 
             if hasattr(raw_response, "headers"):
@@ -556,7 +574,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         """
         raw_response = None
         try:
-            raw_response = openai_client.chat.completions.with_raw_response.create(**data, timeout=timeout)
+            body, options = _request_without_sdk_transform(data, timeout)
+            raw_response = _raw_chat_response(
+                openai_client.post(
+                    "/chat/completions",
+                    body=body,
+                    cast_to=ChatCompletion,
+                    options=options,
+                    stream=bool(data.get("stream")),
+                    stream_cls=Stream[ChatCompletionChunk],
+                )
+            )
 
             if hasattr(raw_response, "headers"):
                 headers = dict(raw_response.headers)
@@ -1255,7 +1283,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         logging_obj: LiteLLMLoggingObj,
     ) -> LegacyAPIResponse[CreateEmbeddingResponse]:
         if "encoding_format" not in data:
-            body, options = _embedding_request_without_sdk_defaults(data, timeout)
+            body, options = _request_without_sdk_transform(data, timeout)
             bypass_response: Final = await openai_aclient.post(
                 "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse
             )
@@ -1272,7 +1300,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         logging_obj: LiteLLMLoggingObj,
     ) -> LegacyAPIResponse[CreateEmbeddingResponse]:
         if "encoding_format" not in data:
-            body, options = _embedding_request_without_sdk_defaults(data, timeout)
+            body, options = _request_without_sdk_transform(data, timeout)
             bypass_response: Final = openai_client.post(
                 "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse
             )

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -40,6 +40,7 @@ from litellm.llms.custom_httpx.http_handler import (
     get_httpx_client,
 )
 from litellm.types.llms.anthropic import AnthropicThinkingParam
+from litellm.types.llms.base import CachedTokensDetails
 from litellm.types.llms.gemini import BidiGenerateContentServerMessage
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -2004,26 +2005,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                 elif modality == "DOCUMENT":
                     cached_text_tokens = (cached_text_tokens or 0) + token_count
 
-        ## Calculate non-cached tokens by subtracting cached from total (per modality)
-        ## This is necessary because promptTokensDetails includes both cached and non-cached tokens
-        ## See: https://github.com/BerriAI/litellm/issues/18750
-        if cached_text_tokens is not None and prompt_text_tokens is not None:
-            # Explicit caching: subtract cached tokens per modality from cacheTokensDetails
-            prompt_text_tokens = prompt_text_tokens - cached_text_tokens
-        elif (
+        if (
             cached_tokens is not None
             and prompt_text_tokens is not None
             and cached_text_tokens is None
             and "cacheTokensDetails" not in usage_metadata
         ):
-            # Implicit caching: only cachedContentTokenCount is provided (no cacheTokensDetails)
-            # Subtract from text tokens since implicit caching is primarily for text content
-            # See: https://github.com/BerriAI/litellm/issues/16341
             prompt_text_tokens = prompt_text_tokens - cached_tokens
-        if cached_audio_tokens is not None and prompt_audio_tokens is not None:
-            prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens
-        if cached_image_tokens is not None and prompt_image_tokens is not None:
-            prompt_image_tokens = prompt_image_tokens - cached_image_tokens
         if cached_video_tokens is not None and prompt_video_tokens is not None:
             prompt_video_tokens = prompt_video_tokens - cached_video_tokens
 
@@ -2044,9 +2032,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         billable_text_tokens: Final = (
             prompt_text_tokens if prompt_text_tokens is None else prompt_text_tokens + billable_tool_use_prompt_tokens
         )
+        cached_tokens_details_fields: Final = MappingProxyType(
+            {
+                key: token_count
+                for key, token_count in (
+                    ("text_tokens", cached_text_tokens),
+                    ("audio_tokens", cached_audio_tokens),
+                    ("image_tokens", cached_image_tokens),
+                )
+                if token_count is not None
+            }
+        )
+        cached_tokens_details: Final = (
+            CachedTokensDetails.model_validate(cached_tokens_details_fields) if cached_tokens_details_fields else None
+        )
 
         prompt_tokens_details = PromptTokensDetailsWrapper(
             cached_tokens=cached_tokens,
+            cached_tokens_details=cached_tokens_details,
             audio_tokens=prompt_audio_tokens,
             text_tokens=billable_text_tokens,
             image_tokens=prompt_image_tokens,

litellm/main.py

@@ -232,7 +232,7 @@ from .llms.codestral.completion.handler import CodestralTextCompletion
 from .llms.cohere.embed import handler as cohere_embed
 from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler
 from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
-from .llms.custom_llm import CustomLLM, custom_chat_llm_router
+from .llms.custom_llm import custom_chat_llm_router
 from .llms.databricks.embed.handler import DatabricksEmbeddingHandler
 from .llms.deprecated_providers import aleph_alpha
 from .llms.gdc.chat.transformation import GDCGeminiConfig
@@ -4954,10 +4954,7 @@ def _complete_custom_providers(
     stream: Final = ctx.stream
     timeout: Final = ctx.timeout
 
-    custom_handler: CustomLLM | None = None
-    for item in litellm.custom_provider_map:
-        if item["provider"] == custom_llm_provider:
-            custom_handler = item["custom_handler"]
+    custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider)
 
     if custom_handler is None:
         raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)
@@ -7366,10 +7363,7 @@ def embedding(
                 litellm_params={},
             )
         elif custom_llm_provider in litellm._custom_providers:
-            custom_handler: CustomLLM | None = None
-            for item in litellm.custom_provider_map:
-                if item["provider"] == custom_llm_provider:
-                    custom_handler = item["custom_handler"]
+            custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider)
 
             if custom_handler is None:
                 raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)

litellm/proxy/_types.py

@@ -5198,6 +5198,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False):
     team_id: str | None  # team that triggered discovery
     key_alias: str | None  # human-readable key alias
     user_agent: str | None  # HTTP User-Agent of the caller
+    call_count: ReadOnly[int]  # invocations folded into this row since the previous flush
 
 
 from litellm.models.managed_files import (  # noqa: E402

litellm/proxy/analytics_endpoints/cache_activity.py

@@ -2,31 +2,22 @@ import asyncio
 import json
 from collections.abc import Sequence
 from datetime import datetime
-from typing import Final, Protocol
+from typing import TYPE_CHECKING, Final
 
 from pydantic import TypeAdapter
 
 from litellm.proxy._types import LiteLLMRoutes
 from litellm.proxy.auth.authorization import ReadScope, read_scope_bind_params
+from litellm.proxy.utils import query_raw_within_statement_timeout
 from litellm.types.llms.base import LiteLLMBaseModel
 
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+
 UNKNOWN_CALL_TYPE: Final = "Unknown"
 INFO_ROUTES_JSON: Final = json.dumps(LiteLLMRoutes.info_routes.value)
 
 
-class _SupportsQueryRaw(Protocol):
-    """The single database operation the cache-activity queries issue."""
-
-    async def query_raw(self, query: str, *args: object) -> Sequence[object]: ...
-
-
-class _SupportsRawQueryDb(Protocol):
-    """A prisma client handle, narrowed to the raw-query surface used here."""
-
-    @property
-    def db(self) -> _SupportsQueryRaw: ...
-
-
 class CacheActivityGroup(LiteLLMBaseModel):
     call_type: str
     api_requests: int
@@ -174,7 +165,7 @@ def compute_totals(groups: Sequence[CacheActivityGroup]) -> CacheActivityTotals:
 
 
 async def get_cache_activity(
-    prisma_client: _SupportsRawQueryDb,
+    prisma_client: "PrismaClient",
     read_scope: ReadScope,
     start_date: datetime,
     end_date: datetime,
@@ -185,14 +176,46 @@ async def get_cache_activity(
     models_json: Final = json.dumps(list(models))
     scope_params: Final = read_scope_bind_params(read_scope)
     group_rows, error_rows, key_alias_rows, model_rows = await asyncio.gather(
-        prisma_client.db.query_raw(
-            GROUPS_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON, *scope_params
+        query_raw_within_statement_timeout(
+            prisma_client,
+            GROUPS_SQL,
+            start_date,
+            end_date,
+            key_aliases_json,
+            models_json,
+            INFO_ROUTES_JSON,
+            *scope_params,
+            what="Cache activity",
         ),
-        prisma_client.db.query_raw(
-            ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON, *scope_params
+        query_raw_within_statement_timeout(
+            prisma_client,
+            ERROR_BREAKDOWN_SQL,
+            start_date,
+            end_date,
+            key_aliases_json,
+            models_json,
+            INFO_ROUTES_JSON,
+            *scope_params,
+            what="Cache activity",
+        ),
+        query_raw_within_statement_timeout(
+            prisma_client,
+            KEY_ALIAS_OPTIONS_SQL,
+            start_date,
+            end_date,
+            INFO_ROUTES_JSON,
+            *scope_params,
+            what="Cache activity",
+        ),
+        query_raw_within_statement_timeout(
+            prisma_client,
+            MODEL_OPTIONS_SQL,
+            start_date,
+            end_date,
+            INFO_ROUTES_JSON,
+            *scope_params,
+            what="Cache activity",
         ),
-        prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON, *scope_params),
-        prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON, *scope_params),
     )
     groups: Final = _groups_adapter.validate_python(group_rows or [])
     return CacheActivityResponse(

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -18,6 +18,7 @@ from litellm.anthropic_interface.exceptions import (
     AnthropicExceptionMapping,
 )
 from litellm.integrations.custom_guardrail import ModifyResponseException
+from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError
 from litellm.litellm_core_utils.tool_fields import first_validation_error
 from litellm.llms.anthropic.pass_through.context_management import (
     AnthropicContextManagementError,
@@ -413,6 +414,10 @@ async def count_tokens(
             messages_format="anthropic",
             **_count_tokens_body_options(data, auto_drop_params),
         )
+    except UnsupportedContentTypeError as e:
+        return _anthropic_error_json_response(
+            ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400), request
+        )
     except HTTPException as e:
         return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)
     except ProxyException as e:

litellm/proxy/common_request_processing.py

@@ -6,6 +6,7 @@ import math
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
 from datetime import datetime
 from functools import lru_cache
+from itertools import chain
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
@@ -38,6 +39,7 @@ from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthr
 from litellm.constants import (
     DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
     DEFAULT_MAX_RECURSE_DEPTH,
+    LITELLM_CLIENT_DISCONNECT_DETAIL,
     LITELLM_DETAILED_TIMING,
     LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
     MAX_LITELLM_CALL_ID_LENGTH,
@@ -68,7 +70,11 @@ from litellm.litellm_core_utils.get_supported_openai_params import (
 )
 from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
+from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
+    guardrail_entries,
+    guardrail_information_cost,
+    object_mapping,
+)
 from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
 )
@@ -298,6 +304,23 @@ async def close_guarded_stream(stream: object) -> None:
             )
 
 
+def _guardrail_entries_for_headers(
+    request_metadata_bucket: object, logging_obj: LiteLLMLoggingObj
+) -> tuple[object, ...]:
+    """Guardrail entries the proxy hooks wrote on the request plus the ones the call recorded on its own logging
+    object (a Bedrock deployment's native guardrail trace), each once: the hooks copy theirs into the logging object."""
+    litellm_params: Final = object_mapping(logging_obj.model_call_details.get("litellm_params")) or {}
+    buckets: Final = (object_mapping(request_metadata_bucket), object_mapping(litellm_params.get("metadata")))
+    entries: Final = tuple(
+        chain.from_iterable(
+            guardrail_entries(metadata.get("standard_logging_guardrail_information"))
+            for metadata in buckets
+            if metadata is not None
+        )
+    )
+    return tuple(entry for index, entry in enumerate(entries) if entry not in entries[:index])
+
+
 def resolve_litellm_call_id(client_call_id: str | None) -> str:
     if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH:
         return client_call_id
@@ -1784,7 +1807,7 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
     return False
 
 
-_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
+_CLIENT_DISCONNECT_DETAIL: Final = LITELLM_CLIENT_DISCONNECT_DETAIL
 
 
 @runtime_checkable
@@ -3186,9 +3209,9 @@ class ProxyBaseLLMRequestProcessing:
             if is_unbilled_non_inference_call_from_params(logging_obj.call_type, logging_obj.litellm_params, response)
             else computed_cost_for_headers
         )
-        _, request_metadata_bucket = get_or_create_metadata_bucket(self.data)
+        request_metadata_bucket: Final[object] = get_or_create_metadata_bucket(self.data)[1]  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]  # self.data is an untyped dict
         guardrail_cost_for_headers: Final = guardrail_information_cost(
-            request_metadata_bucket.get("standard_logging_guardrail_information")
+            _guardrail_entries_for_headers(request_metadata_bucket, logging_obj)  # pyright: ignore[reportUnknownArgumentType]  # bucket comes from the untyped request dict
         )
         response_cost_for_headers: Final = (
             (llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0)

litellm/proxy/common_utils/http_parsing_utils.py

@@ -5,16 +5,19 @@ from types import MappingProxyType, UnionType
 from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin
 
 import orjson
-from fastapi import Request, UploadFile, status
+from fastapi import HTTPException, Request, UploadFile, status
 from pydantic import TypeAdapter
 from starlette._utils import get_route_path
-from starlette.requests import HTTPConnection
+from starlette.datastructures import FormData
+from starlette.requests import ClientDisconnect, HTTPConnection
 from typing_extensions import NotRequired, ReadOnly, Required, assert_never
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
     AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX,
     CLIENT_REQUESTED_MODEL_SCOPE_KEY,
+    LITELLM_CLIENT_DISCONNECT_DETAIL,
+    LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
     MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB,
 )
 from litellm.integrations.otel.runtime import phase_event
@@ -208,6 +211,29 @@ def _parse_form_metadata(raw_metadata: str) -> object:
     return parsed
 
 
+async def _read_form_data(request: Request) -> FormData:
+    try:
+        form_data: Final = await request.form()
+        _mark_body_received(_declared_content_length(request.headers))
+    except ClientDisconnect:
+        raise
+    except Exception as e:
+        # ``request.form()`` raises on malformed multipart (missing
+        # boundary, malformed chunk encoding, …). Surface as 400 so
+        # the auth-time pre-read does not silently cache ``{}`` while
+        # a later raw-body re-read sees the original payload —
+        # banned-param checks must see the same body the handler
+        # acts on.
+        verbose_proxy_logger.error("Invalid form payload: %s", e)
+        raise ProxyException(
+            message=f"Invalid form payload: {e}",
+            type="invalid_request_error",
+            param="request_body",
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+    return form_data
+
+
 async def read_request_body(request: Request | None) -> dict:
     """
     Safely read the request body and parse it as JSON.
@@ -218,10 +244,9 @@ async def read_request_body(request: Request | None) -> dict:
     Returns:
     - dict: Parsed request data as a dictionary or an empty dictionary if parsing fails
     """
+    if request is None:
+        return {}
     try:
-        if request is None:
-            return {}
-
         if is_otlp_trace_request(request):
             return {}
 
@@ -238,24 +263,7 @@ async def read_request_body(request: Request | None) -> dict:
             _mark_body_received(len(binary_body))
             parsed_body = _parse_binary_body(binary_body)
         elif _is_form_content_type(content_type):
-            try:
-                form_data: Final = await request.form()
-                _mark_body_received(_declared_content_length(request.headers))
-            except Exception as e:
-                # ``request.form()`` raises on malformed multipart (missing
-                # boundary, malformed chunk encoding, …). Surface as 400 so
-                # the auth-time pre-read does not silently cache ``{}`` while
-                # a later raw-body re-read sees the original payload —
-                # banned-param checks must see the same body the handler
-                # acts on.
-                verbose_proxy_logger.error("Invalid form payload: %s", e)
-                raise ProxyException(
-                    message=f"Invalid form payload: {e}",
-                    type="invalid_request_error",
-                    param="request_body",
-                    code=status.HTTP_400_BAD_REQUEST,
-                )
-            form_body: Final = dict(form_data)
+            form_body: Final = dict(await _read_form_data(request))
             raw_metadata: Final = form_body.get("metadata")
             parsed_body = (
                 {**form_body, "metadata": _parse_form_metadata(raw_metadata)}
@@ -318,6 +326,17 @@ async def read_request_body(request: Request | None) -> dict:
         # Re-raise ProxyException as-is
         verbose_proxy_logger.error("Invalid JSON payload received: %s", e)
         raise
+    except ClientDisconnect as e:
+        from litellm.proxy.proxy_server import (
+            _close_dangling_otel_server_span,  # pyright: ignore[reportPrivateUsage]  # proxy_server keeps the span-close helper private; the Starlette HTTPException handler does not stamp the OTel server span opened at auth time, so the body reader must before it aborts
+        )
+
+        verbose_proxy_logger.info("Client disconnected before the request body was fully received")
+        disconnected: Final = HTTPException(
+            status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail=LITELLM_CLIENT_DISCONNECT_DETAIL
+        )
+        _close_dangling_otel_server_span(request, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, exc=disconnected)
+        raise disconnected from e
     except Exception as e:
         # Catch unexpected errors to avoid crashes
         verbose_proxy_logger.exception("Unexpected error reading request body - %s", e)

litellm/proxy/db/db_spend_update_writer.py

@@ -37,7 +37,6 @@ from litellm.constants import (
 from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.proxy._types import (
-    DB_CONNECTION_ERROR_TYPES,
     DB_RETRY_SAFE_ERROR_TYPES,
     BaseDailySpendTransaction,
     DailyAgentSpendTransaction,
@@ -74,7 +73,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
     WindowSpendTransaction,
     WindowSpendUpdateQueue,
 )
-from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
+from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe
 from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction
 from litellm.proxy.db.request_error_tracking import request_error_accumulator
 from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout
@@ -231,16 +230,6 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]):
     ) -> None: ...
 
 
-_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"})
-
-
-def _spend_commit_failure_is_requeue_safe(e: Exception) -> bool:
-    if isinstance(e, DB_CONNECTION_ERROR_TYPES):
-        return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES)
-    sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
-    return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES
-
-
 _SpendTableName = Literal[
     "user_list_transactions",
     "end_user_list_transactions",
@@ -278,7 +267,7 @@ def _spend_tables_left_to_send(
         _SPEND_TABLE_COMMIT_ORDER[len(committed)] if len(committed) < len(_SPEND_TABLE_COMMIT_ORDER) else None
     )
     dropped: Final[frozenset[_SpendTableName]] = (
-        frozenset() if in_flight is None or _spend_commit_failure_is_requeue_safe(failure) else frozenset({in_flight})
+        frozenset() if in_flight is None or db_write_failure_is_requeue_safe(failure) else frozenset({in_flight})
     )
     if dropped and in_flight is not None:
         spend_log_error(
@@ -1757,7 +1746,7 @@ class DBSpendUpdateWriter:
                             window_spend_transactions=window_spend_update_transactions,
                         )
                     except Exception as e:
-                        if not _spend_commit_failure_is_requeue_safe(e):
+                        if not db_write_failure_is_requeue_safe(e):
                             uncommitted.pop("window_spend_update_transactions", None)
                             spend_log_error(
                                 "Spend tracking - dropped %d budget window increments: the failed statement may have "
@@ -1923,7 +1912,7 @@ class DBSpendUpdateWriter:
                 window_spend_transactions=window_spend_update_transactions,
             )
         except Exception as e:  # noqa: BLE001  # the increments go back on the queue; the rest of the flush must run
-            if _spend_commit_failure_is_requeue_safe(e):
+            if db_write_failure_is_requeue_safe(e):
                 spend_log_error(
                     "Spend tracking - failed to commit budget window spend updates. "
                     "Re-queued %d window increments for retry on next tick. Error: %s",
@@ -1941,9 +1930,6 @@ class DBSpendUpdateWriter:
                     exc=e,
                 )
 
-        ################## Tool Registry Upserts ##################
-        await self._flush_tool_discovery_queue(prisma_client=prisma_client)
-
     async def commit_daily_tag_spend_to_db(
         self,
         prisma_client: PrismaClient,
@@ -2075,20 +2061,6 @@ class DBSpendUpdateWriter:
             )
             raise
 
-    async def _flush_tool_discovery_queue(
-        self,
-        prisma_client: PrismaClient,
-    ) -> None:
-        """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable."""
-        from litellm.proxy.db.tool_registry_writer import batch_upsert_tools
-
-        try:
-            items: Final = self.tool_discovery_queue.flush()
-            if items:
-                await batch_upsert_tools(prisma_client=prisma_client, items=items)
-        except Exception as e:
-            verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e)
-
     @staticmethod
     async def handle_spend_update_failure(
         e: Exception,
@@ -2100,7 +2072,6 @@ class DBSpendUpdateWriter:
         """Retry a failed spend-update transaction on connection errors, deadlocks or a rollup
         ``lock_timeout`` (55P03), else re-raise. All three roll the transaction back before any
         increment applied, so re-sending the same batch cannot double-count."""
-        from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
         from litellm.proxy.utils import raise_failed_update_spend_exception
 
         is_retryable = (
@@ -2609,7 +2580,7 @@ class DBSpendUpdateWriter:
                                 _mark_daily_spend_commit_started()
                             _mark_daily_spend_commit_finished()
                         except Exception as batch_error:
-                            if _spend_commit_failure_is_requeue_safe(batch_error):
+                            if db_write_failure_is_requeue_safe(batch_error):
                                 spend_log_error(
                                     "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
                                     entity_type,

litellm/proxy/db/db_transaction_queue/redis_update_buffer.py

@@ -9,8 +9,11 @@ import json
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from functools import reduce
+from itertools import chain
 from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
 
+from pydantic import TypeAdapter
+
 from redis.exceptions import RedisError
 
 from litellm._internal_context import with_service_target
@@ -26,6 +29,7 @@ from litellm.constants import (
     REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
     REDIS_SPEND_LOGS_BUFFER_KEY,
     REDIS_SPEND_LOGS_BUFFER_MAX_ROWS,
+    REDIS_TOOL_DISCOVERY_BUFFER_KEY,
     REDIS_UPDATE_BUFFER_KEY,
     REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
 )
@@ -41,12 +45,17 @@ from litellm.proxy._types import (
     DBSpendUpdateTransactions,
     Litellm_EntityType,
     SpendUpdateQueueItem,
+    ToolDiscoveryQueueItem,
 )
 from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
 from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
     DailySpendUpdateQueue,
 )
 from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
+from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
+    ToolDiscoveryQueue,
+    aggregate_tool_discoveries,
+)
 from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
     WindowSpendTransaction,
     WindowSpendUpdateQueue,
@@ -133,6 +142,9 @@ def _merged_entity_transactions(
     )
 
 
+_TOOL_DISCOVERY_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem])
+
+
 class RedisUpdateBuffer:
     """
     Handles buffering database `UPDATE` transactions in Redis before committing them to the database
@@ -551,6 +563,56 @@ class RedisUpdateBuffer:
                 str(e),
             )
 
+    @with_service_target(SPEND_QUEUE_TARGET)
+    async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None:
+        """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable."""
+        if self.redis_cache is None:
+            return
+        items: Final = tool_discovery_queue.flush()
+        if not items:
+            return
+        try:
+            _ = await self.redis_cache.async_rpush_pipeline(
+                rpush_list=[
+                    RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)])
+                ]
+            )
+        except (RedisError, OSError, asyncio.TimeoutError) as e:
+            verbose_proxy_logger.error(
+                "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s",
+                len(items),
+                str(e),
+            )
+            for item in items:
+                tool_discovery_queue.add_update(item)
+
+    @with_service_target(SPEND_QUEUE_TARGET)
+    async def restore_tool_discoveries_to_redis(self, items: Sequence[ToolDiscoveryQueueItem]) -> None:
+        """Put drained tool rows back for the next drain after a write that provably did not apply."""
+        if self.redis_cache is None or not items:
+            return
+        _ = await self.redis_cache.async_rpush_pipeline(
+            rpush_list=[
+                RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(list(items))])
+            ]
+        )
+
+    @with_service_target(SPEND_QUEUE_TARGET)
+    async def get_all_tool_discoveries_from_redis_buffer(self) -> tuple[ToolDiscoveryQueueItem, ...]:
+        """Drain every worker's tool rows and fold them to one row per tool_name."""
+        if self.redis_cache is None:
+            return ()
+        popped: Final = await self.redis_cache.async_lpop_pipeline(
+            lpop_list=[
+                RedisPipelineLpopOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT)
+            ]
+        )
+        entries: Final = popped[0] if popped else None
+        if not entries:
+            return ()
+        batches: Final = tuple(_TOOL_DISCOVERY_ITEMS.validate_json(entry) for entry in entries)
+        return aggregate_tool_discoveries(chain.from_iterable(batches))
+
     @with_service_target(SPEND_QUEUE_TARGET)
     async def store_spend_logs_in_redis(
         self,

litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py

@@ -1,52 +1,57 @@
 """
-In-memory buffer for tool registry upserts.
-
-Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue
-uses set-deduplication: each unique tool_name is only queued once per flush
-cycle (~30s). The seen-set is cleared on every flush so that call_count
-increments in subsequent cycles rather than stopping after the first flush.
+In-memory buffer for tool registry upserts, aggregated by tool_name so one flush carries one row per
+tool with the number of times it was invoked since the previous flush.
 """
 
+from collections.abc import Iterable
 from typing import Final
 
+from pydantic import TypeAdapter
+
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import ToolDiscoveryQueueItem
 
+_ITEM: Final = TypeAdapter(ToolDiscoveryQueueItem)
+
+
+def _merged(existing: ToolDiscoveryQueueItem | None, item: ToolDiscoveryQueueItem) -> ToolDiscoveryQueueItem:
+    """The first sighting keeps its origin and caller fields, later sightings only add to call_count."""
+    increment: Final = item.get("call_count", 1)
+    if existing is None:
+        return _ITEM.validate_python({**item, "call_count": increment})
+    return _ITEM.validate_python({**existing, "call_count": existing.get("call_count", 1) + increment})
+
+
+def aggregate_tool_discoveries(items: Iterable[ToolDiscoveryQueueItem]) -> tuple[ToolDiscoveryQueueItem, ...]:
+    """Collapse items to one per tool_name, summing call_count, in first-seen order."""
+    by_name: Final[dict[str, ToolDiscoveryQueueItem]] = {}  # mutable-ok: single pass fold keyed by tool_name
+    for item in items:
+        if not (tool_name := item.get("tool_name", "")):
+            continue
+        by_name[tool_name] = _merged(by_name.get(tool_name), item)
+    return tuple(by_name.values())
+
 
 class ToolDiscoveryQueue:
-    """
-    In-memory buffer for tool registry upserts.
-
-    Deduplicates by tool_name within each flush cycle: a tool is only queued
-    once per ~30s batch, so call_count increments once per flush cycle the
-    tool appears in (not once per invocation, but not once per pod lifetime
-    either). The seen-set is cleared on flush so subsequent batches can
-    re-count the same tool.
-    """
+    """In-memory buffer for tool registry upserts, one pending row per tool_name with a running call_count."""
 
     def __init__(self) -> None:
-        self._seen_tool_names: set[str] = set()
-        self._pending: list[ToolDiscoveryQueueItem] = []
+        self._pending: dict[str, ToolDiscoveryQueueItem] = {}  # mutable-ok: folded per add_update, cleared on flush
 
     def add_update(self, item: ToolDiscoveryQueueItem) -> None:
-        """Enqueue a tool discovery item if tool_name has not been seen before."""
+        """Count one more invocation of item's tool_name, keeping the first sighting's origin and caller fields."""
         tool_name: Final = item.get("tool_name", "")
         if not tool_name:
             return
-        if tool_name in self._seen_tool_names:
-            verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name)
-            return
-        self._seen_tool_names.add(tool_name)
-        self._pending.append(item)
+        self._pending[tool_name] = _merged(self._pending.get(tool_name), item)
         verbose_proxy_logger.debug(
-            "ToolDiscoveryQueue: queued new tool %s (origin=%s)",
+            "ToolDiscoveryQueue: tool %s call_count=%s (origin=%s)",
             tool_name,
+            self._pending[tool_name].get("call_count"),
             item.get("origin"),
         )
 
     def flush(self) -> list[ToolDiscoveryQueueItem]:
-        """Return and clear all pending items. Resets seen-set so the next
-        flush cycle can re-count the same tools."""
-        items, self._pending = self._pending, []
-        self._seen_tool_names.clear()
+        """Return and clear all pending rows."""
+        items, self._pending = list(self._pending.values()), {}
         return items

litellm/proxy/db/exception_handler.py

@@ -8,6 +8,7 @@ from pydantic import TypeAdapter, ValidationError
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import (
     DB_CONNECTION_ERROR_TYPES,
+    DB_RETRY_SAFE_ERROR_TYPES,
     ProxyErrorTypes,
     ProxyException,
 )
@@ -33,6 +34,8 @@ _CONNECTION_CAPACITY_PHRASES: Final = (
 _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
 _LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
 _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
+_STATEMENT_TIMEOUT_SQLSTATE: Final = "57014"
+_STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout"
 
 
 def _exception_chain(e: BaseException) -> Iterator[BaseException]:
@@ -283,6 +286,17 @@ class PrismaDBExceptionHandler:
             return True
         return _LOCK_TIMEOUT_PHRASE in str(e).lower()
 
+    @staticmethod
+    def is_statement_timeout_error(e: Exception) -> bool:
+        """True iff Postgres cancelled a statement for exceeding ``statement_timeout`` (SQLSTATE 57014)."""
+        import prisma
+
+        if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
+            return False
+        if PrismaDBExceptionHandler.postgres_sqlstate(e) == _STATEMENT_TIMEOUT_SQLSTATE:
+            return True
+        return _STATEMENT_TIMEOUT_PHRASE in str(e).lower()
+
     @staticmethod
     def postgres_sqlstate(e: Exception) -> str | None:
         """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None."""
@@ -637,3 +651,14 @@ async def call_with_db_reconnect_retry(
         # propagate — repeated reconnect-loops are the watchdog's job, not
         # this helper's.
         return await coro_factory()
+
+
+_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"})
+
+
+def db_write_failure_is_requeue_safe(e: Exception) -> bool:
+    """Whether a failed write provably never applied, so its rows can be sent again without double counting."""
+    if isinstance(e, DB_CONNECTION_ERROR_TYPES):
+        return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES)
+    sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
+    return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES

litellm/proxy/db/tool_registry_writer.py

@@ -14,8 +14,17 @@ from typing import TYPE_CHECKING, Final, Protocol
 from pydantic import TypeAdapter
 
 from litellm._logging import verbose_proxy_logger
+from litellm.constants import TOOL_DISCOVERY_COMMIT_JOB_NAME
 from litellm.proxy._types import ToolDiscoveryQueueItem
-from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
+from litellm.proxy.db.db_span import db_span
+from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
+from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
+from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
+    ToolDiscoveryQueue,
+    aggregate_tool_discoveries,
+)
+from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe
+from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert
 from litellm.repositories.object_permission_repository import ObjectPermissionRepository
 from litellm.repositories.prisma_protocols import TableActions
 from litellm.repositories.table_repositories import ToolRepository
@@ -105,58 +114,75 @@ def _row_to_model(row: object) -> LiteLLM_ToolTableRow:
 
 async def batch_upsert_tools(
     prisma_client: "PrismaClient",
-    items: list[ToolDiscoveryQueueItem],
+    items: Sequence[ToolDiscoveryQueueItem],
 ) -> None:
-    """
-    Batch-upsert tool registry rows via Prisma.
+    """Write a batch of discovered tools to LiteLLM_ToolTable in a single statement. Raises on failure so the
+    caller decides whether the rows can be retried."""
+    rows: Final = aggregate_tool_discoveries(items)
+    if not rows:
+        return
+    sql, params = build_tool_bulk_upsert(rows)
+    async with db_span("commit_tool_discoveries", "LiteLLM_ToolTable"):
+        await prisma_client.db.execute_raw(sql, *params)  # pyright: ignore[reportAny]  # untyped prisma client
+    verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(rows))
 
-    On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1.
-    On conflict: increments call_count; preserves existing policies.
-    """
+
+async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discovery_queue: ToolDiscoveryQueue) -> None:
+    """One worker's own rows straight to LiteLLM_ToolTable, the path without the Redis buffer."""
+    items: Final = tool_discovery_queue.flush()
     if not items:
         return
     try:
-        data: Final = [item for item in items if item.get("tool_name")]
-        if not data:
-            return
-        now: Final = datetime.now(timezone.utc)
-        table: Final = _tool_table_actions(prisma_client)
-        for item in data:
-            tool_name = item.get("tool_name", "")
-            origin = item.get("origin") or "user_defined"
-            created_by = item.get("created_by") or "system"
-            key_hash = item.get("key_hash")
-            team_id = item.get("team_id")
-            key_alias = item.get("key_alias")
-            user_agent = item.get("user_agent")
-            await table.upsert(
-                where={"tool_name": tool_name},
-                data={
-                    "create": {
-                        "tool_id": str(uuid.uuid4()),
-                        "tool_name": tool_name,
-                        "origin": origin,
-                        "input_policy": "untrusted",
-                        "output_policy": "untrusted",
-                        "call_count": 1,
-                        "created_by": created_by,
-                        "updated_by": created_by,
-                        "key_hash": key_hash,
-                        "team_id": team_id,
-                        "key_alias": key_alias,
-                        "user_agent": user_agent,
-                        "last_used_at": now,
-                    },
-                    "update": {
-                        "call_count": {"increment": 1},
-                        "updated_at": now,
-                        "last_used_at": now,
-                    },
-                },
+        await batch_upsert_tools(prisma_client=prisma_client, items=items)
+    except Exception as e:  # noqa: BLE001  # a lost tool row only delays its discovery, the scheduler job must not fail
+        verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e)
+
+
+async def commit_tool_discoveries(
+    prisma_client: "PrismaClient",
+    tool_discovery_queue: ToolDiscoveryQueue,
+    redis_update_buffer: RedisUpdateBuffer,
+    pod_lock_manager: PodLockManager,
+    use_redis_buffer: bool,
+) -> None:
+    """Scheduled alongside update_spend on its own lock, so tool table writes never ride the spend commit.
+
+    With the Redis buffer every worker hands its rows to Redis and the holder of the tool lock writes them all in
+    one statement. Without it each worker writes its own rows."""
+    if not use_redis_buffer:
+        await flush_tool_discovery_queue(prisma_client=prisma_client, tool_discovery_queue=tool_discovery_queue)
+        return
+    await redis_update_buffer.store_tool_discoveries_in_redis(tool_discovery_queue=tool_discovery_queue)
+    if not await pod_lock_manager.acquire_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME):
+        return
+    try:
+        await _commit_tool_discoveries_from_redis(prisma_client=prisma_client, redis_update_buffer=redis_update_buffer)
+    finally:
+        await pod_lock_manager.release_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME)
+
+
+async def _commit_tool_discoveries_from_redis(
+    prisma_client: "PrismaClient", redis_update_buffer: RedisUpdateBuffer
+) -> None:
+    """Tool lock holder only: drain every worker's tool rows from Redis and write them in one statement."""
+    items: Final = await redis_update_buffer.get_all_tool_discoveries_from_redis_buffer()
+    if not items:
+        return
+    try:
+        await batch_upsert_tools(prisma_client=prisma_client, items=items)
+    except Exception as e:  # noqa: BLE001  # the rows go back to redis or are dropped, the scheduler job must not fail
+        if db_write_failure_is_requeue_safe(e):
+            await redis_update_buffer.restore_tool_discoveries_to_redis(items)
+            verbose_proxy_logger.error(
+                "Tool discovery - %d tool rows went back to Redis for the next tick. Error: %s", len(items), str(e)
             )
-        verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(data))
-    except Exception as e:
-        verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e)
+            return
+        verbose_proxy_logger.error(
+            "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database "
+            "refused the data, so re-sending it is not safe. Error: %s",
+            len(items),
+            str(e),
+        )
 
 
 _NO_OWNERS: Final[Mapping[str, ToolDiscoveryUser]] = MappingProxyType({})

litellm/proxy/db/tool_table_bulk_upsert.py

@@ -0,0 +1,71 @@
+"""One multi-row ``INSERT ... ON CONFLICT ("tool_name") DO UPDATE`` per batch of discovered tools.
+
+LiteLLM_ToolTable holds one row per tool name the proxy has seen. A batch arrives already folded to one
+entry per tool name with a summed call_count, and this statement inserts the names the table has not seen
+with the untrusted policies the columns default to, and for the names it has adds the batch's call_count
+while leaving the policies and the discovering caller's fields alone. Same shape as daily_spend_bulk_upsert.
+"""
+
+import uuid
+from collections.abc import Mapping, Sequence
+from itertools import chain
+from types import MappingProxyType
+from typing import Final, TypeAlias
+
+from litellm.proxy._types import ToolDiscoveryQueueItem
+
+SqlValue: TypeAlias = str | int | None
+
+_TOOL_ROW_COLUMNS: Final[tuple[str, ...]] = (
+    "tool_id",
+    "tool_name",
+    "origin",
+    "call_count",
+    "created_by",
+    "updated_by",
+    "key_hash",
+    "team_id",
+    "key_alias",
+    "user_agent",
+)
+_CASTS: Final[Mapping[str, str]] = MappingProxyType({"call_count": "integer"})
+
+
+def _row_params(item: ToolDiscoveryQueueItem) -> tuple[SqlValue, ...]:
+    created_by: Final = item.get("created_by") or "system"
+    return (
+        str(uuid.uuid4()),
+        item.get("tool_name", ""),
+        item.get("origin") or "user_defined",
+        item.get("call_count", 1),
+        created_by,
+        created_by,
+        item.get("key_hash"),
+        item.get("team_id"),
+        item.get("key_alias"),
+        item.get("user_agent"),
+    )
+
+
+def build_tool_bulk_upsert(items: Sequence[ToolDiscoveryQueueItem]) -> tuple[str, tuple[SqlValue, ...]]:
+    """The single statement writing one batch of distinct tool_names, plus its positional arguments."""
+    width: Final = len(_TOOL_ROW_COLUMNS)
+    rows: Final = ", ".join(
+        "("
+        + ", ".join(
+            f"${row_index * width + offset + 1}::{_CASTS.get(column, 'text')}"
+            for offset, column in enumerate(_TOOL_ROW_COLUMNS)
+        )
+        + ", (NOW() AT TIME ZONE 'UTC'), (NOW() AT TIME ZONE 'UTC'))"
+        for row_index in range(len(items))
+    )
+    columns: Final = ", ".join(f'"{column}"' for column in _TOOL_ROW_COLUMNS)
+    sql: Final = (
+        f'INSERT INTO "LiteLLM_ToolTable" ({columns}, "last_used_at", "updated_at")\n'
+        f"VALUES {rows}\n"
+        'ON CONFLICT ("tool_name") DO UPDATE SET\n'
+        '  "call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count",\n'
+        '  "last_used_at" = EXCLUDED."last_used_at",\n'
+        '  "updated_at" = EXCLUDED."updated_at"'
+    )
+    return sql, tuple(chain.from_iterable(map(_row_params, items)))

litellm/proxy/management_endpoints/prompt_caching_requests.py

@@ -14,9 +14,7 @@ from litellm.proxy.spend_tracking.savings import (
     marks_gateway_injection,
     prompt_caching_savings_for_request,
 )
-from litellm.proxy.spend_tracking.spend_tracking_utils import (
-    _query_raw_rows,  # pyright: ignore[reportPrivateUsage]  # existing typed spend-query adapter; rows validated below
-)
+from litellm.proxy.utils import query_raw_within_statement_timeout
 from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.management_endpoints.prompt_caching_requests import (
@@ -158,7 +156,7 @@ async def get_prompt_caching_requests(
         else cursor_start_time
     )
     rows: Final = _REQUEST_ROWS.validate_python(
-        await _query_raw_rows(
+        await query_raw_within_statement_timeout(
             prisma_client,
             prompt_caching_requests_sql(filter),
             start.isoformat(),
@@ -166,8 +164,8 @@ async def get_prompt_caching_requests(
             page_size + 1,
             cursor_time.isoformat() if cursor_time is not None else None,
             cursor_request_id,
+            what="Prompt caching requests",
         )
-        or ()
     )
 
     def current_router() -> "Router | None":

litellm/proxy/management_endpoints/team_endpoints.py

@@ -205,7 +205,12 @@ from litellm.proxy.management_helpers.utils import (
     add_new_member,
     management_endpoint_wrapper,
 )
-from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
+from litellm.proxy.utils import (
+    PrismaClient,
+    ProxyLogging,
+    handle_exception_on_proxy,
+    query_raw_within_statement_timeout,
+)
 from litellm.repositories.budget_repository import BudgetRepository
 from litellm.repositories.chunked_in import find_many_in
 from litellm.repositories.organization_repository import OrganizationRepository
@@ -6969,6 +6974,9 @@ class _TeamUserSpendDbRow(TypedDict):
     failed_requests: ReadOnly[int]
 
 
+_TEAM_USER_SPEND_ROWS: Final = TypeAdapter(tuple[_TeamUserSpendDbRow, ...])
+
+
 @router.get(
     "/team/spend/by_user",
     response_model=TeamUserSpendResponse,
@@ -7019,12 +7027,16 @@ async def get_team_spend_by_user(
     own_user_only: Final = scope.api_key_filter is not None
     user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else ()
     async with db_span("team_user_spend", "LiteLLM_SpendLogs"):
-        rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw(
-            _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
-            start_date,
-            end_date,
-            *scoped_team_ids,
-            *user_param,
+        rows: Final = _TEAM_USER_SPEND_ROWS.validate_python(
+            await query_raw_within_statement_timeout(
+                prisma_client,
+                _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
+                start_date,
+                end_date,
+                *scoped_team_ids,
+                *user_param,
+                what="Spend per user",
+            )
         )
     results: Final = tuple(
         TeamUserSpendRow(

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -4259,10 +4259,16 @@ async def watsonx_proxy_route(
     if provider_config is None:
         raise HTTPException(status_code=404, detail="Watsonx passthrough config not found")
 
+    watsonx_credentials: Final = passthrough_endpoint_router.get_proxy_wide_deployment_credentials(
+        custom_llm_provider="watsonx"
+    )
+    api_key: Final = watsonx_credentials.api_key if watsonx_credentials is not None else None
+    api_base: Final = watsonx_credentials.api_base if watsonx_credentials is not None else None
+
     # Get complete URL with version parameter
     complete_url, _ = provider_config.get_complete_url(
-        api_base=None,
-        api_key=None,
+        api_base=api_base,
+        api_key=api_key,
         model="",
         endpoint=endpoint,
         request_query_params=None,
@@ -4276,8 +4282,8 @@ async def watsonx_proxy_route(
         messages=[],
         optional_params={},
         litellm_params={},
-        api_key=None,
-        api_base=None,
+        api_key=api_key,
+        api_base=api_base,
     )
 
     # Check for streaming

litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py

@@ -5,7 +5,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, cast
 
 import httpx
-from pydantic import ConfigDict, TypeAdapter
+from pydantic import ConfigDict, TypeAdapter, ValidationError
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -25,6 +25,7 @@ from litellm.llms.anthropic.pass_through.stream_assembly import (
     collapse_pure_text_chunks,
     split_sse_chunk_into_events,
 )
+from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD
 from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
 from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
@@ -55,6 +56,29 @@ else:
 _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
 
 
+def _event_payload(event: str) -> Mapping[str, object] | None:
+    data_lines: Final = tuple(
+        line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:")
+    )
+    if not data_lines or data_lines[0] in ("", "[DONE]"):
+        return None
+    try:
+        return _JSON_OBJECT.validate_json(data_lines[0])
+    except ValidationError:
+        return None
+
+
+def _bedrock_trace_payload(all_chunks: Sequence[str | bytes]) -> Mapping[str, object] | None:
+    payloads: Final = (
+        _event_payload(event)
+        for chunk in reversed(all_chunks)
+        for event in reversed(split_sse_chunk_into_events(chunk))
+    )  # comprehension-ok: chunk then event, read newest first
+    return next(
+        (payload for payload in payloads if payload is not None and BEDROCK_INVOKE_TRACE_FIELD in payload), None
+    )
+
+
 class AnthropicPassthroughLoggingHandler:
     @staticmethod
     def anthropic_passthrough_handler(
@@ -106,6 +130,7 @@ class AnthropicPassthroughLoggingHandler:
             json_mode=False,
             litellm_params={},
         )
+        logging_obj.record_native_bedrock_guardrail(response_body)  # pyright: ignore[reportUnknownArgumentType]  # response_body is an untyped dict param
 
         kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
             litellm_model_response=litellm_model_response,
@@ -522,6 +547,7 @@ class AnthropicPassthroughLoggingHandler:
                 "result": None,
                 "kwargs": {},
             }
+        litellm_logging_obj.record_native_bedrock_guardrail(_bedrock_trace_payload(all_chunks))
         kwargs: Final = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
             litellm_model_response=complete_streaming_response,
             model=model,

litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py

@@ -1,5 +1,6 @@
 import json
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Final
 
 from fastapi import HTTPException
@@ -57,6 +58,12 @@ def _resolve_oidc_reference(credential: str | None) -> str | None:
         ) from exc
 
 
+@dataclass(frozen=True, slots=True)
+class _PassThroughCredentials:
+    api_key: str
+    api_base: str | None
+
+
 class PassthroughEndpointRouter:
     """
     Use this class to Get credentials for pass-through endpoints
@@ -89,6 +96,35 @@ class PassthroughEndpointRouter:
         )
         return _resolve_oidc_reference(get_secret_str(_env_variable_name))
 
+    def get_proxy_wide_deployment_credentials(self, custom_llm_provider: str) -> _PassThroughCredentials | None:
+        llm_router: Final = self.llm_router_getter()
+        if llm_router is None:
+            return None
+        for deployment in llm_router.get_model_list() or ():
+            litellm_params: LiteLLMParamsTypedDict = deployment["litellm_params"]
+            api_key: str | None = self._resolve_matching_deployment_api_key(
+                litellm_params=litellm_params,
+                custom_llm_provider=custom_llm_provider,
+                region_name=None,
+            )
+            if api_key is None:
+                continue
+            resolved_api_key: str | None = _resolve_oidc_reference(api_key)
+            if not resolved_api_key:
+                return None
+            return _PassThroughCredentials(
+                api_key=resolved_api_key,
+                api_base=self._resolve_matching_deployment_api_base(litellm_params),
+            )
+        return None
+
+    def _resolve_matching_deployment_api_base(self, litellm_params: LiteLLMParamsTypedDict) -> str | None:
+        credential_name: Final = litellm_params.get("litellm_credential_name")
+        credential_values: Final = (
+            CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else None
+        )
+        return _get_str_value(credential_values, "api_base") or litellm_params.get("api_base")
+
     def _get_deployment_api_key(
         self,
         custom_llm_provider: str,

litellm/proxy/proxy_server.py

@@ -10888,6 +10888,19 @@ class ProxyStartupEvent:
             misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
         )
 
+        ### TOOL DISCOVERY (own lock, same interval, never rides the spend commit) ###
+        from litellm.proxy.utils import commit_tool_discoveries
+
+        scheduler.add_job(  # pyright: ignore[reportUnknownMemberType]  # apscheduler is untyped
+            commit_tool_discoveries,
+            "interval",
+            seconds=batch_writing_interval,
+            args=[prisma_client, proxy_logging_obj],
+            id="commit_tool_discoveries_job",
+            replace_existing=True,
+            misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
+        )
+
         ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ###
         ## Reduces QPS as there are more tags for a single request
         tag_spend_update_interval: Final = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER)

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -65,7 +65,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
     get_spend_by_team,
     get_spend_by_team_and_customer,
 )
-from litellm.proxy.utils import handle_exception_on_proxy
+from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout
 from litellm.repositories.prisma_protocols import TableActions
 from litellm.repositories.table_repositories import SpendLogsRepository
 from litellm.repositories.team_repository import TeamRepository
@@ -268,6 +268,20 @@ async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args:
     return await _query_raw(prisma_client, sql_query, *args)
 
 
+async def _query_spend_logs(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
+    """``_query_raw`` for raw ``LiteLLM_SpendLogs`` reads, which Postgres cancels at the statement timeout."""
+    return cast(  # cast-ok: callers declare the row shape their SELECT list produces, as with _query_raw
+        "Sequence[_RowT]", await query_raw_within_statement_timeout(prisma_client, sql_query, *args)
+    )
+
+
+async def _query_spend_logs_or_none(
+    prisma_client: PrismaClient, sql_query: str, *args: object
+) -> Sequence[_RowT] | None:
+    """``_query_spend_logs`` for the call sites that guard the result against ``None``."""
+    return await _query_spend_logs(prisma_client, sql_query, *args)
+
+
 class _TeamTable(Protocol):
     """The subset of the Prisma team table API this module uses."""
 
@@ -1567,7 +1581,9 @@ async def get_global_spend_report(
                 ORDER BY
                     total_cost DESC;
             """
-            db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj, api_key)
+            db_response = await _query_spend_logs_or_none(
+                prisma_client, sql_query, start_date_obj, end_date_obj, api_key
+            )
             if db_response is None:
                 return []
 
@@ -1610,7 +1626,7 @@ async def get_global_spend_report(
                 ORDER BY
                     total_cost DESC;
             """
-            db_response = await _query_raw_or_none(
+            db_response = await _query_spend_logs_or_none(
                 prisma_client, sql_query, start_date_obj, end_date_obj, internal_user_id
             )
             if db_response is None:
@@ -1677,7 +1693,7 @@ async def get_global_spend_report(
                 group_by_day;
                 """
 
-            db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj)
+            db_response = await _query_spend_logs_or_none(prisma_client, sql_query, start_date_obj, end_date_obj)
             if db_response is None:
                 return []
 
@@ -1718,12 +1734,14 @@ async def get_global_spend_report(
                 ORDER BY
                     total_cost DESC;
             """
-            db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj)
+            db_response = await _query_spend_logs_or_none(prisma_client, sql_query, start_date_obj, end_date_obj)
             if db_response is None:
                 return []
 
             return db_response
 
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -1967,7 +1985,7 @@ async def get_key_spend_report(
         caller_value=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
         scope_name="api_key",
     )
-    db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none(
+    db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none(
         prisma_client,
         _scoped_spend_report_sql(scope_column="api_key"),
         start_date_obj,
@@ -2009,7 +2027,7 @@ async def get_user_spend_report(
         caller_value=user_api_key_dict.user_id,
         scope_name="internal_user_id",
     )
-    db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none(
+    db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none(
         prisma_client,
         _scoped_spend_report_sql(scope_column="user"),
         start_date_obj,
@@ -2051,7 +2069,7 @@ async def get_team_spend_report(
         caller_value=user_api_key_dict.team_id,
         scope_name="team_id",
     )
-    db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none(
+    db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none(
         prisma_client,
         _scoped_spend_report_sql(scope_column="team_id"),
         start_date_obj,
@@ -2093,7 +2111,7 @@ async def get_organization_spend_report(
         organization_id=organization_id,
         prisma_client=prisma_client,
     )
-    db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none(
+    db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none(
         prisma_client,
         _ORG_SPEND_REPORT_SQL,
         start_date_obj,
@@ -3019,7 +3037,7 @@ async def ui_view_spend_logs(
                 LIMIT ${p}
             ) AS bounded_matches
         """
-        count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_raw_or_none(
+        count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_spend_logs_or_none(
             prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
         )
         raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0
@@ -3050,7 +3068,7 @@ async def ui_view_spend_logs(
         )
         sql_params.extend([page_size, skip])
 
-        data: Final = await prisma_client.db.query_raw(sql_query, *sql_params)
+        data: Final[Sequence[Mapping[str, object]]] = await _query_spend_logs(prisma_client, sql_query, *sql_params)
 
         if request_id is not None and not is_v2 and not is_admin_view:
             await _assert_user_owns_fetched_spend_rows(
@@ -3118,7 +3136,7 @@ async def _fetch_session_representatives(
             ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL}
         ) AS session_representatives
     """
-    rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw(  # mutable-ok: rows are enriched in place
+    rep_rows: Final[Sequence[dict[str, object]]] = await _query_spend_logs(  # mutable-ok: rows are enriched in place
         prisma_client,
         rep_query,
         *sql_params,
@@ -3148,7 +3166,7 @@ async def _count_grouped_sessions(
             LIMIT ${next_param_index}
         ) AS bounded_sessions
     """
-    count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw(
+    count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_spend_logs(
         prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
     )
     raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0
@@ -3222,7 +3240,9 @@ async def _ui_session_grouped_spend_logs(
     page_rows: Final[Sequence[_SessionPageRow]] = (
         ()
         if page_limit <= 0
-        else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params)
+        else await _query_spend_logs(
+            prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params
+        )
     )
 
     has_more: Final = len(page_rows) > page_limit
@@ -4707,7 +4727,7 @@ async def _build_ui_spend_logs_response(
                     if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None
                 }
             )
-            rows: Final[Sequence[_SessionSpendRow]] = await _query_raw(
+            rows: Final[Sequence[_SessionSpendRow]] = await _query_spend_logs(
                 prisma_client,
                 f"""
                 SELECT s.*, COALESCE(m.session_models, ARRAY[]::text[]) AS session_models
@@ -4778,7 +4798,7 @@ async def _build_ui_spend_logs_response(
                 if row.get("session_id") and row.get("api_key") is not None
                 for models in (TypeAdapter(list[str]).validate_python(row.get("session_models") or ()),)
             }
-        except PrismaError:
+        except (PrismaError, HTTPException):
             verbose_proxy_logger.debug(
                 "Failed to enrich session spend aggregates for spend logs UI",
                 exc_info=True,

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -63,7 +63,7 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.route_llm_request import ProxyModelNotFoundError
 from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
-from litellm.proxy.utils import PrismaClient, hash_token
+from litellm.proxy.utils import PrismaClient, hash_token, query_raw_within_statement_timeout
 from litellm.types.router import DeploymentTypedDict, LiteLLM_Params
 from litellm.types.utils import (
     BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
@@ -985,7 +985,8 @@ async def _query_raw_rows(
     sql_query: str,
     *args: object,
 ) -> Sequence[Mapping[str, object]] | None:
-    return await prisma_client.db.query_raw(sql_query, *args)
+    """A raw ``LiteLLM_SpendLogs`` read that Postgres cancels at the statement timeout."""
+    return await query_raw_within_statement_timeout(prisma_client, sql_query, *args)
 
 
 async def get_spend_by_team(

litellm/proxy/utils.py

@@ -64,6 +64,7 @@ from litellm.constants import (
     PROXY_REJECTED_BEFORE_ROUTING_KEY,
     REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT,
     SPEND_LOG_QUEUE_MAX_BYTES,
+    SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS,
     SPEND_LOG_WRITE_BATCH_MAX_BYTES,
     SPEND_LOG_WRITE_BATCH_MAX_ROWS,
 )
@@ -180,6 +181,7 @@ from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
 from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.db import tool_registry_writer
 from litellm.proxy.db.create_views import (
     create_missing_views,
     create_view_tolerating_race,
@@ -187,6 +189,7 @@ from litellm.proxy.db.create_views import (
 )
 from litellm.proxy.db.db_span import db_span
 from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
+from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
 from litellm.proxy.db.db_url_settings import (
     DatabaseURLSettings,
     add_missing_query_params,
@@ -5035,6 +5038,39 @@ def jsonify_object(data: dict) -> dict:
 _deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(max_size=1000)
 _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60
 _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5)
+_SPEND_LOG_READ_STATEMENT_TIMEOUT: Final = timedelta(seconds=SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS)
+_STATEMENT_TIMEOUT_TRANSACTION_GRACE: Final = timedelta(seconds=5)
+
+
+class _ReadTransactionClient(Protocol):
+    @property
+    def read_db(self) -> PrismaWrapper: ...
+
+
+async def query_raw_within_statement_timeout(
+    prisma_client: _ReadTransactionClient,
+    sql: str,
+    *params: object,
+    what: str = "Spend logs",
+    timeout: timedelta = _SPEND_LOG_READ_STATEMENT_TIMEOUT,
+) -> Sequence[Mapping[str, object]]:
+    """One raw read on the read target that Postgres cancels at ``timeout``.
+
+    Prisma's HTTP client gives up on a statement after 30s anyway, so without
+    this the Admin UI's error is the same but Postgres keeps running the scan.
+    A cancelled read surfaces as a 503 asking for a narrower date range.
+    """
+    try:
+        async with prisma_client.read_db.tx(timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE) as transaction:
+            await transaction.execute_raw(f"SET LOCAL statement_timeout = {int(timeout / timedelta(milliseconds=1))}")
+            return await transaction.query_raw(sql, *params)
+    except Exception as e:
+        if not PrismaDBExceptionHandler.is_statement_timeout_error(e):
+            raise
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail={"error": f"{what} took longer than {int(timeout.total_seconds())}s. Pick a shorter date range"},
+        ) from e
 
 
 async def _lookup_deprecated_key(
@@ -8289,6 +8325,18 @@ class ProxyUpdateSpend:
         return False
 
 
+async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None:
+    """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend."""
+    writer: Final = proxy_logging_obj.db_spend_update_writer
+    await tool_registry_writer.commit_tool_discoveries(
+        prisma_client=prisma_client,
+        tool_discovery_queue=writer.tool_discovery_queue,
+        redis_update_buffer=writer.redis_update_buffer,
+        pod_lock_manager=writer.pod_lock_manager,
+        use_redis_buffer=RedisUpdateBuffer.should_commit_spend_updates_to_redis(),
+    )
+
+
 async def update_spend(
     prisma_client: PrismaClient,
     db_writer_client: AsyncHTTPHandler | None,

litellm/responses/litellm_completion_transformation/handler.py

@@ -122,7 +122,6 @@ class LiteLLMCompletionTransformationHandler:
         litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion(
             **acompletion_args,
         )
-
         if isinstance(litellm_completion_response, ModelResponse):
             responses_api_response: Final[ResponsesAPIResponse] = (
                 LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(

litellm/router.py

@@ -1224,7 +1224,7 @@ class Router:
             },
         ]
 
-        router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": "openai-gpt-3.5-turbo"}])
+        router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}])
         ```
         """
 
@@ -1431,6 +1431,10 @@ class Router:
 
         if default_fallbacks is not None or litellm.default_fallbacks is not None:
             _fallbacks = default_fallbacks or litellm.default_fallbacks
+            if isinstance(_fallbacks, str):
+                raise TypeError(
+                    f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']"
+                )
             if self.fallbacks is not None:
                 self.fallbacks.append({"*": _fallbacks})
             else:
@@ -2564,7 +2568,7 @@ class Router:
         self._initialize_core_endpoints()
         self._initialize_specialized_endpoints()
 
-    def validate_fallbacks(self, fallback_param: list | None):
+    def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None):
         """
         Validate the fallbacks parameter.
         """
@@ -2577,6 +2581,13 @@ class Router:
                 raise ValueError(
                     f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys."
                 )
+            ((model_group, targets),) = fallback_dict.items()
+            if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]):
+                continue
+            raise TypeError(
+                f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. "
+                f"Use {{'{model_group}': ['{targets}']}}"
+            )
 
     def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None:
         def _move_before_deployment_affinity(

litellm/router_utils/cooldown_handlers.py

@@ -8,12 +8,12 @@ Router cooldown handlers
 
 import asyncio
 import math
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from datetime import datetime
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final
 
-from pydantic import TypeAdapter
+from pydantic import TypeAdapter, ValidationError
 
 import litellm
 from litellm._internal_context import service_target
@@ -176,7 +176,14 @@ def _should_cooldown_based_on_deployment_policy(
     override that safety net, matching `_has_explicit_allowed_fails_policy_for_exception`.
     """
     allowed_fails_from_policy: Final = _resolve_allowed_fails_from_policy(dep_policy, original_exception)
-    if allowed_fails_from_policy is None and dep_allowed_fails is not None and is_single_deployment_model_group:
+    if (
+        allowed_fails_from_policy is None
+        and is_single_deployment_model_group
+        and (
+            dep_allowed_fails is not None
+            or litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None
+        )
+    ):
         return False
 
     allowed_fails_override: Final[int | None] = (
@@ -390,6 +397,50 @@ def _should_run_cooldown_logic(
     return True
 
 
+_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...])
+_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str])
+
+
+def _fallback_target_name(target: str | Mapping[str, object]) -> str:
+    return target if isinstance(target, str) else str(target.get("model", ""))
+
+
+def _fallback_target_names(
+    targets: Sequence[str | Mapping[str, object]] | str,
+) -> list[str]:  # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
+    return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets]
+
+
+def _fallback_chain_names(
+    entry: object,
+) -> Mapping[str, list[str]] | None:  # mutable-ok: mirrors get_fallback_model_group's fallbacks shape
+    try:
+        validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry)
+    except ValidationError:
+        return None
+    return {key: _fallback_target_names(targets) for key, targets in validated.items()}
+
+
+def _fallback_entries(configured: object) -> tuple[object, ...]:
+    try:
+        return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured)
+    except ValidationError:
+        return ()
+
+
+def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool:
+    from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
+
+    configured: Final[object] = getattr(litellm_router_instance, "fallbacks", None)
+    if requested_model_group is None or not configured:
+        return False
+    chains: Final = tuple(
+        chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None
+    )
+    fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group)
+    return bool(fallback_group)
+
+
 def _should_cooldown_deployment(
     litellm_router_instance: LitellmRouter,
     deployment: str,
@@ -422,6 +473,9 @@ def _should_cooldown_deployment(
             litellm_router_instance.routing_group_has_alternatives(requested_model_group)
             or litellm_router_instance.team_model_has_alternatives(deployment)
         )
+    is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group(
+        litellm_router_instance, requested_model_group
+    )
 
     exception_status_int: Final = cast_exception_status_to_int(exception_status)
     dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment)
@@ -441,7 +495,7 @@ def _should_cooldown_deployment(
             original_exception,
             dep_policy,
             dep_allowed_fails,
-            is_single_deployment_model_group,
+            is_lone_deployment,
         )
 
     ## BASE CASE - single deployment
@@ -486,6 +540,8 @@ def _should_cooldown_deployment(
 
         return False
     else:
+        if is_lone_deployment and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None:
+            return False
         return should_cooldown_based_on_allowed_fails_policy(
             litellm_router_instance=litellm_router_instance,
             deployment=deployment,

litellm/types/llms/anthropic.py

@@ -209,6 +209,13 @@ class AnthropicMessagesTextParam(TypedDict, total=False):
     prompt_cache_breakpoint: ReadOnly[PromptCacheBreakpoint]
 
 
+class AnthropicSearchResultParam(TypedDict, total=False):
+    type: Required[ReadOnly[Literal["search_result"]]]
+    source: Required[ReadOnly[str]]
+    title: Required[ReadOnly[str]]
+    content: Required[ReadOnly[Iterable[AnthropicMessagesTextParam]]]
+
+
 class AnthropicMessagesToolUseParam(TypedDict, total=False):
     type: Required[Literal["tool_use"]]
     id: str
@@ -357,6 +364,7 @@ class AnthropicMessagesToolResultParam(TypedDict, total=False):
             AnthropicMessagesToolResultContent
             | AnthropicMessagesImageParam
             | AnthropicMessagesDocumentParam
+            | AnthropicSearchResultParam
             | ToolReference
         ]
     )
@@ -368,6 +376,7 @@ AnthropicMessagesUserMessageValues = (
     | AnthropicMessagesImageParam
     | AnthropicMessagesToolResultParam
     | AnthropicMessagesDocumentParam
+    | AnthropicSearchResultParam
     | AnthropicMessagesContainerUploadParam
 )
 

litellm/types/llms/custom_llm.py

@@ -1,11 +1,12 @@
 from typing import TYPE_CHECKING
 
-from typing_extensions import TypedDict
+from typing_extensions import ReadOnly, TypedDict
 
 if TYPE_CHECKING:
+    from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
     from litellm.llms.custom_llm import CustomLLM
 
 
 class CustomLLMItem(TypedDict):
     provider: str
-    custom_handler: "CustomLLM"
+    custom_handler: ReadOnly["CustomLLM | BaseVideoConfig"]

litellm/types/utils.py

@@ -321,6 +321,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     cache_creation_input_token_cost_ultrafast: ReadOnly[float | None]  # OpenAI ultrafast service tier pricing
     cache_read_input_token_cost: float | None
     cache_read_input_audio_token_cost: ReadOnly[float | None]
+    cache_read_input_audio_token_cost_flex: ReadOnly[float | None]
+    cache_read_input_audio_token_cost_priority: ReadOnly[float | None]
     cache_read_input_image_token_cost: ReadOnly[float | None]
     cache_read_input_token_cost_flex: float | None  # OpenAI flex service tier pricing
     cache_read_input_token_cost_priority: float | None  # OpenAI priority service tier pricing
@@ -335,6 +337,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     cache_read_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None]
     cache_read_input_token_cost_above_512k_tokens: float | None
     cache_read_input_token_cost_batches: ReadOnly[float | None]
+    cache_read_input_audio_token_cost_batches: ReadOnly[float | None]
     cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None]
     cache_read_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None]
     cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None]
@@ -348,6 +351,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     off_peak_pricing: ReadOnly[OffPeakPricing | None]  # time-windowed off-peak rates
     input_cost_per_character: float | None  # only for vertex ai models
     input_cost_per_audio_token: float | None
+    input_cost_per_audio_token_flex: ReadOnly[float | None]
+    input_cost_per_audio_token_above_200k_tokens: ReadOnly[float | None]
+    input_cost_per_audio_token_above_200k_tokens_flex: ReadOnly[float | None]
+    input_cost_per_audio_token_above_200k_tokens_priority: ReadOnly[float | None]
+    input_cost_per_audio_token_above_200k_tokens_batches: ReadOnly[float | None]
+    input_cost_per_audio_token_priority: ReadOnly[float | None]
     input_cost_per_token_above_128k_tokens: float | None  # only for vertex ai models
     input_cost_per_token_above_200k_tokens: float | None  # only for vertex ai gemini-2.5-pro models
     input_cost_per_token_above_100k_tokens: ReadOnly[float | None]
@@ -372,10 +381,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     input_cost_per_token_batches: float | None
     input_cost_per_video_token_batches: ReadOnly[float | None]
     input_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None]
+    input_cost_per_token_above_200k_tokens_flex: ReadOnly[float | None]
     input_cost_per_token_above_100k_tokens_batches: ReadOnly[float | None]
     input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
     output_cost_per_token_batches: float | None
     output_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None]
+    output_cost_per_token_above_200k_tokens_flex: ReadOnly[float | None]
     output_cost_per_token_above_100k_tokens_batches: ReadOnly[float | None]
     output_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
     output_cost_per_token: Required[float | None]
@@ -408,6 +419,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
     output_cost_per_pixel: ReadOnly[float | None]
     output_cost_per_image_token: float | None
     output_cost_per_image_token_batches: ReadOnly[float | None]
+    output_cost_per_image_token_priority: ReadOnly[float | None]
     output_cost_per_video_token: float | None  # for gemini omni models with video output
     output_vector_size: int | None
     output_cost_per_reasoning_token: float | None
@@ -3964,14 +3976,23 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
     cache_read_input_token_cost_above_100k_tokens_batches: float | None = None
     cache_read_input_token_cost_above_200k_tokens_batches: float | None = None
     cache_read_input_token_cost_above_272k_tokens_batches: float | None = None
+    cache_read_input_audio_token_cost_batches: float | None = None
     cache_creation_input_token_cost_batches: float | None = None
     cache_creation_input_token_cost_above_100k_tokens_batches: float | None = None
     cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None
     cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None
     cache_read_input_audio_token_cost: float | None = None
+    cache_read_input_audio_token_cost_flex: float | None = None
+    cache_read_input_audio_token_cost_priority: float | None = None
     cache_read_input_image_token_cost: float | None = None
     input_cost_per_character_above_128k_tokens: float | None = None
     input_cost_per_audio_token: float | None = None
+    input_cost_per_audio_token_flex: float | None = None
+    input_cost_per_audio_token_above_200k_tokens: float | None = None
+    input_cost_per_audio_token_above_200k_tokens_flex: float | None = None
+    input_cost_per_audio_token_above_200k_tokens_priority: float | None = None
+    input_cost_per_audio_token_above_200k_tokens_batches: float | None = None
+    input_cost_per_audio_token_priority: float | None = None
     input_cost_per_token_cache_hit: float | None = None
     input_cost_per_token_above_128k_tokens: float | None = None
     input_cost_per_token_above_100k_tokens: float | None = None
@@ -3981,6 +4002,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
     input_cost_per_token_above_272k_tokens_flex: float | None = None
     input_cost_per_token_above_272k_tokens_ultrafast: float | None = None
     input_cost_per_token_above_200k_tokens_batches: float | None = None
+    input_cost_per_token_above_200k_tokens_flex: float | None = None
     input_cost_per_token_above_100k_tokens_batches: float | None = None
     input_cost_per_token_above_272k_tokens_batches: float | None = None
     input_cost_per_query: float | None = None
@@ -4010,12 +4032,14 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
     output_cost_per_token_above_272k_tokens_flex: float | None = None
     output_cost_per_token_above_272k_tokens_ultrafast: float | None = None
     output_cost_per_token_above_200k_tokens_batches: float | None = None
+    output_cost_per_token_above_200k_tokens_flex: float | None = None
     output_cost_per_token_above_100k_tokens_batches: float | None = None
     output_cost_per_token_above_272k_tokens_batches: float | None = None
     output_cost_per_character_above_128k_tokens: float | None = None
     output_cost_per_image: float | None = None
     output_cost_per_image_token: float | None = None
     output_cost_per_image_token_batches: float | None = None
+    output_cost_per_image_token_priority: float | None = None
     output_cost_per_video_token: float | None = None
     output_cost_per_reasoning_token: float | None = None
     output_cost_per_reasoning_token_flex: float | None = None

litellm/utils.py

@@ -438,6 +438,7 @@ if TYPE_CHECKING:
         get_num_retries_from_retry_policy,
         reset_retry_policy,
     )
+    from litellm.llms.custom_llm import CustomLLM
     from litellm.types.llms.anthropic import (
         ANTHROPIC_API_ONLY_HEADERS,
         AnthropicThinkingParam,
@@ -6542,6 +6543,13 @@ def get_model_info_helper(
                 ),
                 cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None),
                 cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None),
+                cache_read_input_audio_token_cost_flex=_model_info.get("cache_read_input_audio_token_cost_flex", None),
+                cache_read_input_audio_token_cost_priority=_model_info.get(
+                    "cache_read_input_audio_token_cost_priority", None
+                ),
+                cache_read_input_audio_token_cost_batches=_model_info.get(
+                    "cache_read_input_audio_token_cost_batches", None
+                ),
                 prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None),
                 cache_read_input_token_cost_above_200k_tokens=_model_info.get(
                     "cache_read_input_token_cost_above_200k_tokens", None
@@ -6608,6 +6616,20 @@ def get_model_info_helper(
                 input_cost_per_second=_model_info.get("input_cost_per_second", None),
                 input_cost_per_text_message=_model_info.get("input_cost_per_text_message", None),
                 input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None),
+                input_cost_per_audio_token_flex=_model_info.get("input_cost_per_audio_token_flex", None),
+                input_cost_per_audio_token_above_200k_tokens=_model_info.get(
+                    "input_cost_per_audio_token_above_200k_tokens", None
+                ),
+                input_cost_per_audio_token_above_200k_tokens_flex=_model_info.get(
+                    "input_cost_per_audio_token_above_200k_tokens_flex", None
+                ),
+                input_cost_per_audio_token_above_200k_tokens_priority=_model_info.get(
+                    "input_cost_per_audio_token_above_200k_tokens_priority", None
+                ),
+                input_cost_per_audio_token_above_200k_tokens_batches=_model_info.get(
+                    "input_cost_per_audio_token_above_200k_tokens_batches", None
+                ),
+                input_cost_per_audio_token_priority=_model_info.get("input_cost_per_audio_token_priority", None),
                 input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None),
                 input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None),
                 input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None),
@@ -6620,6 +6642,9 @@ def get_model_info_helper(
                 input_cost_per_token_above_200k_tokens_batches=_model_info.get(
                     "input_cost_per_token_above_200k_tokens_batches"
                 ),
+                input_cost_per_token_above_200k_tokens_flex=_model_info.get(
+                    "input_cost_per_token_above_200k_tokens_flex", None
+                ),
                 input_cost_per_token_above_272k_tokens_batches=_model_info.get(
                     "input_cost_per_token_above_272k_tokens_batches"
                 ),
@@ -6627,6 +6652,9 @@ def get_model_info_helper(
                 output_cost_per_token_above_200k_tokens_batches=_model_info.get(
                     "output_cost_per_token_above_200k_tokens_batches"
                 ),
+                output_cost_per_token_above_200k_tokens_flex=_model_info.get(
+                    "output_cost_per_token_above_200k_tokens_flex", None
+                ),
                 output_cost_per_token_above_272k_tokens_batches=_model_info.get(
                     "output_cost_per_token_above_272k_tokens_batches"
                 ),
@@ -6688,6 +6716,7 @@ def get_model_info_helper(
                 output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None),
                 output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None),
                 output_cost_per_image_token_batches=_model_info.get("output_cost_per_image_token_batches", None),
+                output_cost_per_image_token_priority=_model_info.get("output_cost_per_image_token_priority", None),
                 output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None),
                 output_vector_size=_model_info.get("output_vector_size", None),
                 citation_cost_per_token=_model_info.get("citation_cost_per_token", None),
@@ -10140,8 +10169,13 @@ class ProviderConfigManager:
     @staticmethod
     def get_provider_video_config(
         model: str | None,
-        provider: LlmProviders,
+        provider: LlmProviders | str,
     ) -> BaseVideoConfig | None:
+        if not isinstance(provider, LlmProviders):
+            custom_video_config: Final = ProviderConfigManager.get_custom_provider_video_config(provider)
+            if custom_video_config is not None or provider not in LlmProvidersSet:
+                return custom_video_config
+            return ProviderConfigManager.get_provider_video_config(model=model, provider=LlmProviders(provider))
         if LlmProviders.OPENAI == provider:
             from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
 
@@ -10188,6 +10222,34 @@ class ProviderConfigManager:
             return litellm.EdenAIVideoConfig()
         return None
 
+    @staticmethod
+    def get_custom_llm_handler(provider: str | None) -> CustomLLM | None:
+        """The `CustomLLM` registered for `provider` in `litellm.custom_provider_map`, if its handler is one."""
+        from litellm.llms.custom_llm import CustomLLM
+
+        return next(
+            (
+                item["custom_handler"]
+                for item in litellm.custom_provider_map
+                if item["provider"] == provider and isinstance(item["custom_handler"], CustomLLM)
+            ),
+            None,
+        )
+
+    @staticmethod
+    def get_custom_provider_video_config(provider: str) -> BaseVideoConfig | None:
+        """The `BaseVideoConfig` registered for `provider` in `litellm.custom_provider_map`, if its handler is one."""
+        from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
+
+        return next(
+            (
+                item["custom_handler"]
+                for item in litellm.custom_provider_map
+                if item["provider"] == provider and isinstance(item["custom_handler"], BaseVideoConfig)
+            ),
+            None,
+        )
+
     @staticmethod
     def get_provider_container_config(
         provider: LlmProviders,

litellm/videos/main.py

@@ -26,6 +26,17 @@ from litellm.types.videos.utils import decode_video_id_with_provider
 from litellm.utils import ProviderConfigManager, client
 from litellm.videos.utils import VideoGenerationRequestUtils
 
+
+def _unsupported_video_call(
+    call_name: str, custom_llm_provider: str, model: str | None = None
+) -> litellm.BadRequestError:
+    return litellm.BadRequestError(
+        message=f"{call_name} is not supported for {custom_llm_provider}",
+        model=model or "",
+        llm_provider=custom_llm_provider,
+    )
+
+
 #################### Initialize provider clients ####################
 llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler()
 
@@ -40,7 +51,7 @@ async def avideo_generation(
     size: str | None = None,
     user: str | None = None,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -170,7 +181,7 @@ def video_generation(
     size: str | None = None,
     user: str | None = None,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -208,11 +219,11 @@ def video_generation(
         # get provider config
         video_generation_provider_config: BaseVideoConfig | None = ProviderConfigManager.get_provider_video_config(
             model=model,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if video_generation_provider_config is None:
-            raise ValueError(f"video generation is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video generation", custom_llm_provider, model)
 
         local_vars.update(kwargs)
         # Get VideoGenerationOptionalRequestParams with only valid parameters
@@ -329,11 +340,11 @@ def video_content(
         # get provider config
         video_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if video_provider_config is None:
-            raise ValueError(f"video support download is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video support download", custom_llm_provider)
 
         local_vars.update(kwargs)
         # For video content download, we don't need complex optional parameter handling
@@ -456,7 +467,7 @@ async def avideo_remix(
     video_id: str,
     prompt: str,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -559,7 +570,7 @@ def video_remix(
     video_id: str,
     prompt: str,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -598,11 +609,11 @@ def video_remix(
         # get provider config
         video_remix_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if video_remix_provider_config is None:
-            raise ValueError(f"video remix is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video remix", custom_llm_provider)
 
         local_vars.update(kwargs)
         # For video remix, we need the video_id and prompt
@@ -660,7 +671,7 @@ async def avideo_list(
     order: str | None = None,
     api_key: str | None = None,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -774,7 +785,7 @@ def video_list(
     limit: int | None = None,
     order: str | None = None,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -810,11 +821,11 @@ def video_list(
         # get provider config
         video_list_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if video_list_provider_config is None:
-            raise ValueError(f"video list is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video list", custom_llm_provider)
 
         local_vars.update(kwargs)
         # For video list, we need the query parameters
@@ -871,7 +882,7 @@ def video_list(
 async def avideo_status(
     video_id: str,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -970,7 +981,7 @@ def video_status(
 def video_status(
     video_id: str,
     timeout=600,  # default to 10 minutes
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: dict[str, object] | None = None,
@@ -1031,11 +1042,11 @@ def video_status(
         # get provider config
         video_status_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if video_status_provider_config is None:
-            raise ValueError(f"video status is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video status", custom_llm_provider)
 
         local_vars.update(kwargs)
         # For video status, we need the video_id
@@ -1088,7 +1099,7 @@ async def avideo_create_character(
     name: str,
     video: FileTypes,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1143,7 +1154,7 @@ def video_create_character(
     name: str,
     video: FileTypes,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1172,11 +1183,11 @@ def video_create_character(
 
         provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if provider_config is None:
-            raise ValueError(f"video create character is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video create character", custom_llm_provider)
 
         local_vars.update(kwargs)
         request_params: Final[dict] = {"name": name}
@@ -1218,7 +1229,7 @@ def video_create_character(
 async def avideo_get_character(
     character_id: str,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1268,7 +1279,7 @@ async def avideo_get_character(
 def video_get_character(
     character_id: str,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1297,11 +1308,11 @@ def video_get_character(
 
         provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if provider_config is None:
-            raise ValueError(f"video get character is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video get character", custom_llm_provider)
 
         local_vars.update(kwargs)
         request_params: Final[dict] = {"character_id": character_id}
@@ -1343,7 +1354,7 @@ async def avideo_edit(
     video_id: str,
     prompt: str,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1398,7 +1409,7 @@ def video_edit(
     video_id: str,
     prompt: str,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1430,11 +1441,11 @@ def video_edit(
 
         provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if provider_config is None:
-            raise ValueError(f"video edit is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video edit", custom_llm_provider)
 
         local_vars.update(kwargs)
         request_params: Final[dict] = {"video_id": video_id, "prompt": prompt}
@@ -1480,7 +1491,7 @@ async def avideo_extension(
     prompt: str,
     seconds: str,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1534,7 +1545,7 @@ def video_extension(
     prompt: str,
     seconds: str,
     timeout=600,
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     extra_headers: dict[str, object] | None = None,
     extra_query: dict[str, object] | None = None,
     extra_body: dict[str, object] | None = None,
@@ -1564,11 +1575,11 @@ def video_extension(
 
         provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config(
             model=None,
-            provider=litellm.LlmProviders(custom_llm_provider),
+            provider=custom_llm_provider,
         )
 
         if provider_config is None:
-            raise ValueError(f"video extension is not supported for {custom_llm_provider}")
+            raise _unsupported_video_call("video extension", custom_llm_provider)
 
         local_vars.update(kwargs)
         request_params: Final[dict] = {

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm"
-version = "1.108.0"
+version = "1.108.1"
 description = "Library to easily interface with LLM API providers"
 readme = "README.md"
 requires-python = ">=3.10, <3.15"
@@ -79,7 +79,7 @@ proxy = [
     "httpx2>=2.5.0,<3",
     "pydantic>=2.12.0,<3",
     "litellm-proxy-extras==0.4.110",
-    "litellm-enterprise==0.1.79",
+    "litellm-enterprise==0.1.79.post1",
     "RestrictedPython>=8.5,<9.0",
     "rich>=13.9.4,<14.0",
     "InquirerPy>=0.3.4,<1.0",
@@ -393,7 +393,7 @@ litellm-enterprise = { workspace = true }
 members = ["enterprise", "litellm-proxy-extras"]
 
 [tool.commitizen]
-version = "1.108.0"
+version = "1.108.1"
 version_files = [
     "pyproject.toml:^version",
 ]

tests/integration/_support/custom_video_provider.py

@@ -0,0 +1,181 @@
+"""A customer-owned video provider for the proxy's `custom_provider_map`: a job queue with its own wire shape."""
+
+from collections.abc import Mapping
+from types import MappingProxyType
+from typing import Final
+
+import httpx
+from httpx._types import RequestFiles
+
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
+from litellm.llms.custom_llm import CustomLLM
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject
+from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id
+
+_STATUS_BY_STATE: Final = MappingProxyType({"queued": "queued", "running": "in_progress", "done": "completed"})
+
+
+def _video_object(job: Mapping[str, object], model: str | None, custom_llm_provider: str | None) -> VideoObject:
+    job_id: Final = str(job["job_id"])
+    seconds: Final = job.get("seconds")
+    video: Final = VideoObject(
+        id=encode_video_id_with_provider(job_id, custom_llm_provider, model) if custom_llm_provider else job_id,
+        object="video",
+        status=_STATUS_BY_STATE[str(job["state"])],
+        created_at=0,
+        model=model,
+        seconds=str(seconds) if seconds is not None else None,
+    )
+    return video
+
+
+class QueueVideoConfig(BaseVideoConfig):
+    def get_supported_openai_params(self, model: str) -> list:  # mutable-ok: BaseVideoConfig contract
+        return ["seconds"]
+
+    def map_openai_params(
+        self,
+        video_create_optional_params: VideoCreateOptionalRequestParams,
+        model: str,
+        drop_params: bool,
+    ) -> dict:  # mutable-ok: BaseVideoConfig contract
+        return {key: value for key, value in video_create_optional_params.items() if key == "seconds"}
+
+    def validate_environment(
+        self,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+        model: str,
+        api_key: str | None = None,
+        litellm_params: GenericLiteLLMParams | None = None,
+    ) -> dict:  # mutable-ok: BaseVideoConfig contract
+        key: Final = (litellm_params.api_key if litellm_params is not None else None) or api_key or ""
+        return {**headers, "X-Api-Key": key}
+
+    def get_complete_url(
+        self,
+        model: str,
+        api_base: str | None,
+        litellm_params: dict,  # mutable-ok: BaseVideoConfig contract
+    ) -> str:
+        if api_base is None:
+            raise ValueError("api_base is required")
+        return f"{api_base.rstrip('/')}/jobs"
+
+    def transform_video_create_request(
+        self,
+        model: str,
+        prompt: str,
+        api_base: str,
+        video_create_optional_request_params: dict,  # mutable-ok: BaseVideoConfig contract
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+    ) -> tuple[dict, RequestFiles, str]:  # mutable-ok: BaseVideoConfig contract
+        seconds: Final = video_create_optional_request_params.get("seconds")
+        return ({"model": model, "prompt": prompt, "seconds": seconds}, [], api_base)
+
+    def transform_video_create_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+        custom_llm_provider: str | None = None,
+        request_data: dict | None = None,
+    ) -> VideoObject:
+        return _video_object(raw_response.json(), model, custom_llm_provider)
+
+    def transform_video_status_retrieve_request(
+        self,
+        video_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+    ) -> tuple[str, dict]:  # mutable-ok: BaseVideoConfig contract
+        return (f"{api_base}/{extract_original_video_id(video_id)}", {})
+
+    def transform_video_status_retrieve_response(
+        self,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+        custom_llm_provider: str | None = None,
+        client: object = None,
+    ) -> VideoObject:
+        return _video_object(raw_response.json(), None, custom_llm_provider)
+
+    def transform_video_content_request(
+        self,
+        video_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+        variant: str | None = None,
+    ) -> tuple[str, dict]:  # mutable-ok: BaseVideoConfig contract
+        return (f"{api_base}/{extract_original_video_id(video_id)}/file", {})
+
+    def transform_video_content_response(
+        self,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+        client: object = None,
+    ) -> bytes:
+        return raw_response.content
+
+    def transform_video_list_request(
+        self,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+        after: str | None = None,
+        limit: int | None = None,
+        order: str | None = None,
+        extra_query: dict[str, object] | None = None,
+    ) -> tuple[str, dict]:  # mutable-ok: BaseVideoConfig contract
+        return (api_base, {})
+
+    def transform_video_list_response(
+        self,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+        custom_llm_provider: str | None = None,
+    ) -> dict[str, str]:  # mutable-ok: BaseVideoConfig contract
+        return raw_response.json()
+
+    def transform_video_remix_request(
+        self,
+        video_id: str,
+        prompt: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+        extra_body: dict[str, object] | None = None,
+    ) -> tuple[str, dict]:  # mutable-ok: BaseVideoConfig contract
+        raise NotImplementedError("the queue backend has no remix")
+
+    def transform_video_remix_response(
+        self,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+        custom_llm_provider: str | None = None,
+    ) -> VideoObject:
+        raise NotImplementedError("the queue backend has no remix")
+
+    def transform_video_delete_request(
+        self,
+        video_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: BaseVideoConfig contract
+    ) -> tuple[str, dict]:  # mutable-ok: BaseVideoConfig contract
+        raise NotImplementedError("the queue backend has no delete")
+
+    def transform_video_delete_response(
+        self,
+        raw_response: httpx.Response,
+        logging_obj: LiteLLMLoggingObj,
+    ) -> VideoObject:
+        raise NotImplementedError("the queue backend has no delete")
+
+
+VIDEO_CONFIG: Final = QueueVideoConfig()
+IMAGE_ONLY_HANDLER: Final = CustomLLM()

tests/integration/observability/test_guardrail_effects.py

@@ -1,7 +1,9 @@
 import asyncio
 import json
 import re
+import socket
 import threading
+import time
 import uuid
 from collections.abc import Callable, Iterator, Mapping
 from concurrent.futures import ThreadPoolExecutor
@@ -11,6 +13,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from types import MappingProxyType
 from typing import Final
+from urllib.parse import urlsplit
 
 import anthropic
 import httpx
@@ -2928,3 +2931,60 @@ def test_disconnect_scans_keep_recording_after_a_worker_is_killed(gateway: Gatew
         assert all(rig.secret(index) in line for index, line in enumerate(received)), received
         assert _post_call_statuses(rig, rig.model, rows=8) == (("success",),) * 8
         assert rig.owned.process.poll() is None
+
+
+def _openai_upstream(request: Request) -> Reply:
+    if request.target.endswith("/vector_stores"):
+        return Reply(body=json.dumps({"id": "vs_" + uuid.uuid4().hex, "object": "vector_store", "name": None}).encode())
+    return Reply(body=_chat_completion_body("control-reply"))
+
+
+def _post_half_of_the_body_then_hang_up(gateway: Gateway, path: str, body: bytes) -> None:
+    address: Final = urlsplit(str(gateway.client.base_url))
+    assert address.hostname is not None and address.port is not None, address
+    head: Final = (
+        f"POST {path} HTTP/1.1\r\nHost: {address.netloc}\r\nAuthorization: Bearer {gateway.key}\r\n"
+        f"Content-Type: application/json\r\nContent-Length: {len(body)}\r\n\r\n"
+    ).encode()
+    with socket.create_connection((address.hostname, address.port)) as connection:
+        connection.sendall(head + body[: len(body) // 2])
+        time.sleep(0.5)
+
+
+@pytest.mark.parametrize(
+    ("path", "body"),
+    (
+        pytest.param(
+            "/v1/chat/completions",
+            b'{"model": "MODEL", "messages": [{"role": "user", "content": "Say OK"}], "max_tokens": 16}',
+            id="chat-completions-where-an-empty-body-is-invalid",
+        ),
+        pytest.param(
+            "/v1/vector_stores",
+            b'{"name": "integration-store", "metadata": {"owner": "integration", "purpose": "half upload"}}',
+            id="vector-stores-where-an-empty-body-is-valid",
+        ),
+    ),
+)
+def test_client_disconnect_during_the_body_upload_runs_no_endpoint_and_writes_no_spend_row(
+    gateway: Gateway, tmp_path: Path, path: str, body: bytes
+) -> None:
+    with (
+        wire_server(_openai_upstream) as upstream,
+        owned_proxy_process(
+            gateway, tmp_path, {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"}
+        ) as owned,
+        owned.gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(api_base=upstream.url + "/v1", api_key="synthetic-openai-key")
+        _post_half_of_the_body_then_hang_up(owned.gateway, path, body.replace(b"MODEL", model.encode()))
+        control: Final = owned.gateway.chat(model)
+        rows: Final = eventually(
+            lambda: owned.gateway.rows('SELECT request_id, call_type, status FROM "LiteLLM_SpendLogs"', ()),
+            lambda values: any(row["request_id"] == control["id"] for row in values),
+            seconds=70,
+        )
+        assert [(row["call_type"], row["status"]) for row in rows] == [("acompletion", "success")], rows
+        assert [request.target for request in upstream.drain()] == ["/v1/chat/completions"], (
+            "the half-uploaded request must never reach the provider"
+        )

tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py

@@ -0,0 +1,423 @@
+"""A Bedrock deployment whose ``guardrailConfig`` has ``trace: enabled`` bills the policy units Bedrock reports in the
+trace as guardrail cost, on every surface the deployment serves: ``/v1/chat/completions`` and ``/v1/responses`` over
+Converse and Invoke, streamed and not, and the four ``/bedrock/model/{model}`` pass-through actions. The cost rides
+``x-litellm-response-cost`` on non-streaming responses and the spend row's ``spend`` and ``guardrail_information``
+everywhere. ``trace: disabled`` leaves no usage, so nothing is billed beyond the tokens (LIT-10101)"""
+
+import base64
+import json
+import uuid
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from hashlib import sha256
+from pathlib import Path
+from typing import Final
+
+import boto3
+import pytest
+import yaml
+from botocore.config import Config
+from integration._support.budget import case_timeout_seconds
+from integration._support.client import Gateway, Scenario, eventually
+from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.upstream import _aws_event_frame
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+import litellm
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240)
+
+_MODEL_ID: Final = "anthropic.claude-sonnet-4-5-20250929-v1:0"
+_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
+_GUARDRAIL_ID: Final = "native-guardrail"
+_INPUT_COST_PER_TOKEN: Final = 0.001
+_OUTPUT_COST_PER_TOKEN: Final = 0.002
+_PROMPT_TOKENS: Final = 5
+_COMPLETION_TOKENS: Final = 1
+_INPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 1}
+_OUTPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 2, "sensitiveInformationPolicyUnits": 2}
+_SUMMED_USAGE: Final[Mapping[str, int]] = {
+    "topicPolicyUnits": 3,
+    "contentPolicyUnits": 1,
+    "wordPolicyUnits": 1,
+    "sensitiveInformationPolicyUnits": 2,
+}
+_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream")
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_SPEND_SQL: Final = 'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"'
+
+
+def _llm_cost() -> float:
+    return _PROMPT_TOKENS * _INPUT_COST_PER_TOKEN + _COMPLETION_TOKENS * _OUTPUT_COST_PER_TOKEN
+
+
+def _guardrail_cost() -> float:
+    """What the proxy's own cost map prices the traced units at, so the test owns the relation, not AWS's rates."""
+    rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
+    return sum(units * rates[counter] for counter, units in _SUMMED_USAGE.items())
+
+
+def _action(request: Request) -> str:
+    return request.target.partition("?")[0].rsplit("/", 1)[-1]
+
+
+def _traced(request: Request) -> bool:
+    if _action(request).startswith("converse"):
+        config: Final = json.loads(request.body).get("guardrailConfig") or {}
+        return config.get("trace") == "enabled"
+    return request.headers.get("x-amzn-bedrock-trace", "").upper() == "ENABLED"
+
+
+def _assessment(usage: Mapping[str, int]) -> dict[str, JsonValue]:
+    return {"invocationMetrics": {"usage": dict(usage), "guardrailProcessingLatency": 7}}
+
+
+def _converse_trace() -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "inputAssessment": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)},
+            "outputAssessments": {_GUARDRAIL_ID: [_assessment(_OUTPUT_USAGE)]},
+        }
+    }
+
+
+def _invoke_trace() -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "input": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)},
+            "outputs": [{_GUARDRAIL_ID: _assessment(_OUTPUT_USAGE)}],
+        }
+    }
+
+
+def _bedrock(request: Request) -> Reply:
+    """Bedrock's runtime as the caller sees it: a one-token answer, and the guardrail trace with its policy units only
+    when the request enabled tracing, in the shape of the action that served it."""
+    traced: Final = _traced(request)
+    action: Final = _action(request)
+    usage: Final = {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6}
+    converse_trace: Final = {"trace": _converse_trace()} if traced else {}
+    if action == "converse":
+        body: Final = {
+            "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
+            "stopReason": "end_turn",
+            "usage": usage,
+            "metrics": {"latencyMs": 3},
+            **converse_trace,
+        }
+        return Reply(body=json.dumps(body).encode())
+    if action == "converse-stream" and "toolConfig" in json.loads(request.body):
+        return Reply(
+            chunks=(
+                _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"),
+                _aws_event_frame(
+                    "contentBlockStart",
+                    {"start": {"toolUse": {"toolUseId": "tool-1", "name": "lookup"}}, "contentBlockIndex": 0},
+                    "sc",
+                    "u",
+                ),
+                _aws_event_frame(
+                    "contentBlockDelta",
+                    {"delta": {"toolUse": {"input": '{"q": "x"}'}}, "contentBlockIndex": 0},
+                    "sc",
+                    "u",
+                ),
+                _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"),
+                _aws_event_frame("messageStop", {"stopReason": "tool_use"}, "sc", "u"),
+                _aws_event_frame(
+                    "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u"
+                ),
+            ),
+            content_type=_EVENT_STREAM,
+        )
+    if action == "converse-stream":
+        return Reply(
+            chunks=(
+                _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"),
+                _aws_event_frame("contentBlockDelta", {"delta": {"text": "ok"}, "contentBlockIndex": 0}, "sc", "u"),
+                _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"),
+                _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"),
+                _aws_event_frame(
+                    "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u"
+                ),
+            ),
+            content_type=_EVENT_STREAM,
+        )
+    invoke_trace: Final = {"amazon-bedrock-trace": _invoke_trace()} if traced else {}
+    if action == "invoke":
+        message: Final = {
+            "id": f"msg_{uuid.uuid4().hex[:12]}",
+            "type": "message",
+            "role": "assistant",
+            "model": "claude",
+            "content": [{"type": "text", "text": "ok"}],
+            "stop_reason": "end_turn",
+            "stop_sequence": None,
+            "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": _COMPLETION_TOKENS},
+            **invoke_trace,
+        }
+        return Reply(body=json.dumps(message).encode())
+    events: Final = (
+        {
+            "type": "message_start",
+            "message": {
+                "id": f"msg_{uuid.uuid4().hex[:12]}",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude",
+                "content": [],
+                "stop_reason": None,
+                "stop_sequence": None,
+                "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": 0},
+            },
+        },
+        {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+        {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
+        {"type": "content_block_stop", "index": 0},
+        {
+            "type": "message_delta",
+            "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+            "usage": {"output_tokens": _COMPLETION_TOKENS},
+        },
+        {
+            "type": "message_stop",
+            "amazon-bedrock-invocationMetrics": {
+                "inputTokenCount": _PROMPT_TOKENS,
+                "outputTokenCount": _COMPLETION_TOKENS,
+                "invocationLatency": 3,
+                "firstByteLatency": 1,
+            },
+            **invoke_trace,
+        },
+    )
+    return Reply(
+        chunks=tuple(
+            _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u")
+            for event in events
+        ),
+        content_type=_EVENT_STREAM,
+    )
+
+
+def _deployment(wire: Wire, model: str, trace: str) -> Mapping[str, JsonValue]:
+    return {
+        "model": model,
+        "api_base": wire.url,
+        "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+        "aws_secret_access_key": "scripted-secret",
+        "aws_region_name": "us-east-1",
+        "use_in_pass_through": True,
+        "input_cost_per_token": _INPUT_COST_PER_TOKEN,
+        "output_cost_per_token": _OUTPUT_COST_PER_TOKEN,
+        "guardrailConfig": {"guardrailIdentifier": _GUARDRAIL_ID, "guardrailVersion": "1", "trace": trace},
+    }
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+    proxy: Gateway
+    scenario: Scenario
+    converse_model: str
+    invoke_model: str
+    untraced_model: str
+
+    def key(self) -> str:
+        return self.scenario.key(key_alias=f"native-trace-{uuid.uuid4().hex[:8]}")
+
+    def runtime(self, monkeypatch: pytest.MonkeyPatch, key: str):  # noqa: ANN201  # boto3 clients are untyped
+        monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key)
+        return boto3.client(
+            "bedrock-runtime",
+            endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock",
+            region_name="us-east-1",
+            config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}),
+        )
+
+
+@pytest.fixture(scope="module")
+def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+    suffix: Final = uuid.uuid4().hex[:12]
+    converse_model: Final = f"bedrock-traced-converse-{suffix}"
+    invoke_model: Final = f"bedrock-traced-invoke-{suffix}"
+    untraced_model: Final = f"bedrock-untraced-{suffix}"
+    directory: Final = tmp_path_factory.mktemp("bedrock-native-guardrail-trace")
+    base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    with wire_server(_bedrock) as wire:
+        config: Final = {
+            **base_config,
+            "model_list": [
+                *base_config.get("model_list", []),
+                {"model_name": converse_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "enabled")},
+                {
+                    "model_name": invoke_model,
+                    "litellm_params": _deployment(wire, f"bedrock/invoke/{_MODEL_ID}", "enabled"),
+                },
+                {"model_name": untraced_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "disabled")},
+            ],
+        }
+        config_path: Final = directory / "bedrock-native-guardrail-trace.yaml"
+        config_path.write_text(yaml.safe_dump(config))
+        with owned_proxy(session_gateway, directory, {}, config=config_path) as proxy, proxy.scenario() as scenario:
+            yield Rig(
+                proxy=proxy,
+                scenario=scenario,
+                converse_model=converse_model,
+                invoke_model=invoke_model,
+                untraced_model=untraced_model,
+            )
+
+
+@dataclass(frozen=True, slots=True)
+class SpendRow:
+    spend: float
+    guardrail_information: tuple[dict[str, JsonValue], ...]
+
+
+def _spend_row(proxy: Gateway, key: str) -> SpendRow:
+    digest: Final = sha256(key.encode()).hexdigest()
+    (row,) = eventually(lambda: proxy.rows(_SPEND_SQL, (digest,)), lambda rows: len(rows) == 1, seconds=70)
+    metadata: Final = _JSON_OBJECT.validate_python(
+        json.loads(row["metadata"]) if isinstance(row["metadata"], str) else row["metadata"]
+    )
+    raw_information: Final = metadata.get("guardrail_information")
+    entries: Final = raw_information if isinstance(raw_information, list) else [raw_information]
+    return SpendRow(
+        spend=float(str(row["spend"])),
+        guardrail_information=tuple(_JSON_OBJECT.validate_python(entry) for entry in entries if entry is not None),
+    )
+
+
+def _assert_billed_as_guardrail_cost(row: SpendRow) -> None:
+    assert row.spend == pytest.approx(_llm_cost() + _guardrail_cost())
+    (entry,) = row.guardrail_information
+    assert entry["guardrail_provider"] == "bedrock"
+    assert entry["guardrail_name"] == _GUARDRAIL_ID
+    assert entry["guardrail_usage"] == _SUMMED_USAGE
+    assert isinstance(entry["guardrail_cost"], float) and entry["guardrail_cost"] == pytest.approx(_guardrail_cost())
+
+
+def _chat_body(model: str, stream: bool) -> dict[str, JsonValue]:
+    return {"model": model, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 16, "stream": stream}
+
+
+def _responses_body(model: str, stream: bool) -> dict[str, JsonValue]:
+    return {"model": model, "input": "hello", "stream": stream}
+
+
+@pytest.mark.parametrize("api", ["converse", "invoke"])
+@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"])
+def test_non_streaming_response_prices_the_trace_into_the_cost_header_and_the_spend_row(
+    rig: Rig, endpoint: str, api: str
+) -> None:
+    key: Final = rig.key()
+    model: Final = rig.converse_model if api == "converse" else rig.invoke_model
+    body: Final = _chat_body(model, False) if endpoint == "/v1/chat/completions" else _responses_body(model, False)
+    response: Final = rig.proxy.request("POST", endpoint, body, key=key)
+    assert response.is_success, response.text[:400]
+    assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost() + _guardrail_cost())
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+@pytest.mark.parametrize("api", ["converse", "invoke"])
+@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"])
+def test_streamed_response_prices_the_trace_of_the_final_event_into_the_spend_row(
+    rig: Rig, endpoint: str, api: str
+) -> None:
+    key: Final = rig.key()
+    model: Final = rig.converse_model if api == "converse" else rig.invoke_model
+    body: Final = _chat_body(model, True) if endpoint == "/v1/chat/completions" else _responses_body(model, True)
+    with rig.proxy.client.stream(
+        "POST", endpoint, json=body, headers={"Authorization": f"Bearer {key}"}, timeout=case_timeout_seconds()
+    ) as response:
+        frames: Final = tuple(line for line in response.iter_lines() if line.startswith("data:"))
+    assert response.is_success and frames
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+@pytest.mark.parametrize("action", _ACTIONS)
+def test_passthrough_action_prices_the_trace_into_the_spend_row(
+    rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str
+) -> None:
+    key: Final = rig.key()
+    runtime: Final = rig.runtime(monkeypatch, key)
+    messages: Final = [{"role": "user", "content": [{"text": "hello"}]}]
+    invoke_body: Final = json.dumps(
+        {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 16,
+            "messages": [{"role": "user", "content": "hello"}],
+        }
+    )
+    if action == "converse":
+        getattr(runtime, "converse")(modelId=rig.converse_model, messages=messages)
+    elif action == "converse-stream":
+        list(getattr(runtime, "converse_stream")(modelId=rig.converse_model, messages=messages)["stream"])
+    elif action == "invoke":
+        getattr(runtime, "invoke_model")(modelId=rig.converse_model, body=invoke_body)["body"].read()
+    else:
+        list(
+            getattr(runtime, "invoke_model_with_response_stream")(modelId=rig.converse_model, body=invoke_body)["body"]
+        )
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+def test_disabled_trace_bills_only_the_tokens(rig: Rig) -> None:
+    key: Final = rig.key()
+    response: Final = rig.proxy.request("POST", "/v1/chat/completions", _chat_body(rig.untraced_model, False), key=key)
+    assert response.is_success, response.text[:400]
+    assert "trace" not in response.json()
+    assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost())
+    row: Final = _spend_row(rig.proxy, key)
+    assert row.spend == pytest.approx(_llm_cost())
+    assert row.guardrail_information == ()
+
+
+def _messages_body(model: str, stream: bool) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [{"role": "user", "content": f"hello from the messages api, stream={stream}"}],
+        "max_tokens": 16,
+        "stream": stream,
+    }
+
+
+def test_messages_non_streaming_response_applies_the_guardrail_and_prices_its_trace(rig: Rig) -> None:
+    key: Final = rig.key()
+    response: Final = rig.proxy.request("POST", "/v1/messages", _messages_body(rig.invoke_model, False), key=key)
+    assert response.is_success, response.text[:400]
+    assert "amazon-bedrock-trace" in response.json()
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+def test_messages_streamed_response_prices_the_trace_of_the_message_stop_event(rig: Rig) -> None:
+    key: Final = rig.key()
+    with rig.proxy.client.stream(
+        "POST",
+        "/v1/messages",
+        json=_messages_body(rig.invoke_model, True),
+        headers={"Authorization": f"Bearer {key}"},
+        timeout=case_timeout_seconds(),
+    ) as response:
+        body: Final = b"".join(response.iter_bytes())
+    assert response.is_success and b"message_stop" in body, body[:400]
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))
+
+
+def test_passthrough_streamed_tool_call_prices_the_trace_of_the_final_event(
+    rig: Rig, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    key: Final = rig.key()
+    runtime: Final = rig.runtime(monkeypatch, key)
+    tool_config: Final = {
+        "tools": [{"toolSpec": {"name": "lookup", "inputSchema": {"json": {"type": "object", "properties": {}}}}}]
+    }
+    events: Final = list(
+        getattr(runtime, "converse_stream")(
+            modelId=rig.converse_model,
+            messages=[{"role": "user", "content": [{"text": "hello"}]}],
+            toolConfig=tool_config,
+        )["stream"]
+    )
+    assert any("contentBlockStart" in event for event in events)
+    _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key))

tests/integration/providers/test_count_tokens_error_shapes_wire.py

@@ -68,6 +68,65 @@ def _never_called(request: Request) -> Reply:
     raise AssertionError(f"count reached the provider without a key: {request.target}")
 
 
+@pytest.mark.parametrize("nested", (False, True), ids=("user-content", "tool-result"))
+def test_local_messages_count_accepts_search_result_text(gateway: Gateway, nested: bool) -> None:
+    search: Final[JsonValue] = {
+        "type": "search_result",
+        "source": "https://example.invalid/notes",
+        "title": "Project notes",
+        "content": [{"type": "text", "text": "The code word is 7f3a."}],
+        "citations": {"enabled": True},
+    }
+    content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [search]}] if nested else [search]
+    text: Final[list[JsonValue]] = [
+        {"type": "text", "text": "https://example.invalid/notes"},
+        {"type": "text", "text": "Project notes"},
+        {"type": "text", "text": "The code word is 7f3a."},
+    ]
+    equivalent: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": text}] if nested else text
+    with gateway.scenario() as scenario:
+        model: Final = scenario.model(model="hosted_vllm/zai-org/GLM-5.3", api_base="http://127.0.0.1:9/v1")
+        counted: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": content}]}
+        )
+        control: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": equivalent}]}
+        )
+        assert counted.status_code == control.status_code == 200, (counted.text, control.text)
+        tokens: Final = _JSON_OBJECT.validate_json(counted.content)["input_tokens"]
+        assert tokens == _JSON_OBJECT.validate_json(control.content)["input_tokens"], (counted.text, control.text)
+
+
+@pytest.mark.parametrize("nested", (False, True), ids=("user-content", "tool-result"))
+@pytest.mark.parametrize(
+    ("block", "status", "error_type", "message"),
+    (
+        (
+            {"type": "future_block", "future_block": {"value": "ignore me"}},
+            400,
+            "invalid_request_error",
+            "future_block",
+        ),
+        ({"type": "image"}, 500, "api_error", "source"),
+    ),
+    ids=("unknown-block", "unrelated-counting-error"),
+)
+def test_local_messages_count_classifies_only_unsupported_content_as_an_invalid_request(
+    gateway: Gateway, nested: bool, block: JsonValue, status: int, error_type: str, message: str
+) -> None:
+    content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [block]}] if nested else [block]
+    with gateway.scenario() as scenario:
+        model: Final = scenario.model(model="hosted_vllm/zai-org/GLM-5.3", api_base="http://127.0.0.1:9/v1")
+        response: Final = gateway.request(
+            "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": content}]}
+        )
+        assert response.status_code == status, response.text
+        payload: Final = _JSON_OBJECT.validate_json(response.content)
+        error: Final = payload["error"]
+        assert payload["type"] == "error" and isinstance(error, dict), payload
+        assert error["type"] == error_type and message in str(error["message"]), payload
+
+
 @pytest.mark.parametrize(
     ("deployment", "missing"),
     (

tests/integration/providers/test_custom_video_provider_wire.py

@@ -0,0 +1,112 @@
+"""A provider registered in `custom_provider_map` with a `BaseVideoConfig` serves /v1/videos create, status and content."""
+
+import json
+import shutil
+from contextlib import ExitStack
+from pathlib import Path
+from typing import Final
+
+from integration._support.client import Gateway
+from integration._support.process import owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+_SUPPORT_MODULE: Final = Path(__file__).resolve().parents[1] / "_support" / "custom_video_provider.py"
+_PROMPT: Final = "a cat walking across a bridge"
+_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + b"queue-video-bytes"
+
+
+def _queue_backend(request: Request) -> Reply:
+    if (request.method, request.target) == ("POST", "/jobs"):
+        return Reply(body=json.dumps({"job_id": "job-7", "state": "queued"}).encode())
+    if (request.method, request.target) == ("GET", "/jobs/job-7"):
+        return Reply(body=json.dumps({"job_id": "job-7", "state": "done", "seconds": 8}).encode())
+    if (request.method, request.target) == ("GET", "/jobs/job-7/file"):
+        return Reply(body=_MP4, content_type="video/mp4")
+    return Reply(status=404, body=json.dumps({"error": f"no route {request.method} {request.target}"}).encode())
+
+
+def _write_config(path: Path, wire: Wire) -> Path:
+    shutil.copy(_SUPPORT_MODULE, path.parent / "custom_video_provider.py")
+    path.write_text(
+        json.dumps(
+            {
+                "model_list": [
+                    {
+                        "model_name": "queue-video",
+                        "litellm_params": {
+                            "model": "queue-video/seedance-1",
+                            "api_base": wire.url,
+                            "api_key": "queue-secret",
+                        },
+                    },
+                    {
+                        "model_name": "image-only-video",
+                        "litellm_params": {
+                            "model": "image-only/seedance-1",
+                            "api_base": wire.url,
+                            "api_key": "unused",
+                        },
+                    },
+                ],
+                "litellm_settings": {
+                    "custom_provider_map": [
+                        {"provider": "queue-video", "custom_handler": "custom_video_provider.VIDEO_CONFIG"},
+                        {"provider": "image-only", "custom_handler": "custom_video_provider.IMAGE_ONLY_HANDLER"},
+                    ]
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+    return path
+
+
+def _lower_headers(request: Request) -> dict[str, str]:
+    return {name.lower(): value for name, value in request.headers.items()}
+
+
+def test_custom_provider_map_base_video_config_serves_create_status_and_content(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with ExitStack() as stack:
+        wire: Final = stack.enter_context(wire_server(_queue_backend))
+        config: Final = _write_config(tmp_path / "custom-video.yaml", wire)
+        owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config))
+
+        created: Final = owned.gateway.request(
+            "POST", "/v1/videos", {"model": "queue-video", "prompt": _PROMPT, "seconds": "8"}
+        )
+        assert created.status_code == 200, created.text
+        created_body: Final = created.json()
+        assert created_body["status"] == "queued", created_body
+        video_id: Final = created_body["id"]
+        assert isinstance(video_id, str) and video_id.startswith("video_"), created_body
+
+        submits: Final = wire.drain()
+        assert [(r.method, r.target) for r in submits] == [("POST", "/jobs")], submits
+        assert json.loads(submits[0].body) == {"model": "seedance-1", "prompt": _PROMPT, "seconds": "8"}
+        assert _lower_headers(submits[0])["x-api-key"] == "queue-secret"
+
+        status: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}")
+        assert status.status_code == 200, status.text
+        assert status.json()["status"] == "completed", status.text
+
+        content: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}/content")
+        assert content.status_code == 200, content.text
+        assert content.content == _MP4
+
+        polls: Final = wire.drain()
+        assert [(r.method, r.target) for r in polls] == [("GET", "/jobs/job-7"), ("GET", "/jobs/job-7/file")], polls
+        assert all(_lower_headers(r)["x-api-key"] == "queue-secret" for r in polls)
+
+
+def test_custom_provider_without_a_video_config_is_refused_with_400(gateway: Gateway, tmp_path: Path) -> None:
+    with ExitStack() as stack:
+        wire: Final = stack.enter_context(wire_server(_queue_backend))
+        config: Final = _write_config(tmp_path / "custom-video.yaml", wire)
+        owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config))
+
+        refused: Final = owned.gateway.request("POST", "/v1/videos", {"model": "image-only-video", "prompt": _PROMPT})
+        assert refused.status_code == 400, refused.text
+        assert "video generation is not supported for image-only" in refused.text
+        assert wire.drain() == ()

tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py

@@ -0,0 +1,81 @@
+import json
+import uuid
+from pathlib import Path
+from typing import Final
+
+import yaml
+from integration._support.client import Gateway, JsonValue, object_value
+from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, wire_server
+
+_DEPLOYMENT_KEY: Final = "synthetic-watsonx-deployment-key"
+_IAM_TOKEN: Final = "synthetic-watsonx-iam-token"
+_WATSONX_ENV: Final = (
+    "WATSONX_API_KEY",
+    "WATSONX_APIKEY",
+    "WX_API_KEY",
+    "WATSONX_API_BASE",
+    "WATSONX_URL",
+    "WX_URL",
+    "WML_URL",
+    "WATSONX_TOKEN",
+    "WATSONX_ZENAPIKEY",
+)
+
+
+def _config(directory: Path, api_base: str) -> Path:
+    stock: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    config: Final[dict[str, JsonValue]] = {
+        **stock,
+        "model_list": [
+            {
+                "model_name": "granite",
+                "litellm_params": {
+                    "model": "watsonx/ibm/granite-3-8b-instruct",
+                    "api_key": _DEPLOYMENT_KEY,
+                    "api_base": api_base,
+                    "project_id": "synthetic-project",
+                    "use_in_pass_through": True,
+                },
+            }
+        ],
+    }
+    path: Final = directory / "watsonx-passthrough.yaml"
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+def test_watsonx_passthrough_uses_the_flagged_deployment_key_and_base(gateway: Gateway, tmp_path: Path) -> None:
+    marker: Final = "watsonx-" + uuid.uuid4().hex
+
+    def respond(request: Request) -> Reply:
+        if request.method == "POST" and request.target == "/identity/token":
+            assert f"apikey={_DEPLOYMENT_KEY}".encode() in request.body, request.body
+            return Reply(body=json.dumps({"access_token": _IAM_TOKEN, "expires_in": 3600}).encode())
+        if request.method == "POST" and request.target.startswith("/ml/v1/text/generation?"):
+            assert request.headers["authorization"] == f"Bearer {_IAM_TOKEN}", request.headers
+            assert marker.encode() in request.body, request.body[:400]
+            return Reply(
+                body=json.dumps(
+                    {"model_id": "ibm/granite-3-8b-instruct", "results": [{"generated_text": marker}]}
+                ).encode()
+            )
+        return Reply(status=404)
+
+    with wire_server(respond) as wire:
+        with owned_proxy(
+            gateway,
+            tmp_path,
+            {"WATSONX_IAM_URL": f"{wire.url}/identity/token"},
+            config=_config(tmp_path, wire.url),
+            remove_environment=_WATSONX_ENV,
+        ) as candidate:
+            response: Final = candidate.request(
+                "POST",
+                "/watsonx/ml/v1/text/generation",
+                {"model_id": "ibm/granite-3-8b-instruct", "input": marker, "project_id": "synthetic-project"},
+            )
+            assert response.status_code == 200, response.text
+            assert response.json()["results"][0]["generated_text"] == marker
+        forwarded: Final = tuple((request.method, request.target.split("?")[0]) for request in wire.drain())
+        assert forwarded == (("POST", "/identity/token"), ("POST", "/ml/v1/text/generation")), forwarded

tests/integration/routing/test_payment_required_cooldown.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 import json
 import re
 import uuid
@@ -390,6 +391,79 @@ def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_coun
         assert markers_received(refusing) == (first.marker, second.marker)
 
 
+def failing_reply(request: Request) -> Reply:
+    return model_list_reply() if is_model_info_probe(request) else anthropic_error(500, "upstream down")
+
+
+def assert_failed(outcome: Outcome) -> None:
+    assert outcome.status == 500, (outcome.status, outcome.detail)
+    assert "upstream down" in outcome.detail, outcome.detail
+
+
+@pytest.mark.timeout(900)
+def test_a_router_allowed_fails_count_leaves_a_lone_deployment_warm_but_benches_one_with_a_sibling_or_fallback(
+    tmp_path: Path,
+) -> None:
+    lone_group: Final = new_group()
+    pair_group: Final = new_group()
+    primary_group: Final = new_group()
+    backup_group: Final = new_group()
+    with ExitStack() as stack:
+        upstream: Final = stack.enter_context(gateway_from_environment())
+        cache: Final = stack.enter_context(owned_redis(tmp_path))
+        settings: Final = router_settings(cache, allowed_fails=1, fallbacks=[{primary_group: [backup_group]}])
+        owned: Final = stack.enter_context(
+            owned_proxy_process(
+                upstream,
+                tmp_path,
+                {},
+                config=cooldown_config(tmp_path, settings),
+                remove_environment=("PROMETHEUS_MULTIPROC_DIR",),
+            )
+        )
+        failing: Final = stack.enter_context(wire_server(failing_reply))
+        healthy: Final = stack.enter_context(wire_server(healthy_reply))
+        scenario: Final = stack.enter_context(owned.gateway.scenario())
+        lone_id: Final = deployment(scenario, lone_group, anthropic_params(failing))
+        failed: Final = tuple(raw_chat(owned.gateway, lone_group, new_marker()) for _ in range(4))
+        for outcome in failed:
+            assert_failed(outcome)
+        assert not cooled(cache, lone_id)
+        assert markers_received(failing) == markers(failed)
+
+        paired_id: Final = two_deployments(scenario, pair_group, failing, healthy)
+        history: Final = tuple(
+            itertools.takewhile(
+                lambda outcome: not cooled(cache, paired_id),
+                (raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(40)),
+            )
+        )
+        assert cooled(cache, paired_id), [(outcome.status, outcome.detail[:80]) for outcome in history]
+        assert sum(outcome.status == 500 for outcome in history) == 1
+        served: Final = tuple(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(10))
+        assert_served(served, chat_identity)
+        assert markers_received(healthy) == (*markers(o for o in history if o.status == 200), *markers(served))
+
+        primary_id: Final = deployment(scenario, primary_group, anthropic_params(failing))
+        deployment(scenario, backup_group, openai_params(healthy))
+        failing.drain()
+        healthy.drain()
+        fell_back: Final = tuple(
+            itertools.takewhile(
+                lambda outcome: not cooled(cache, primary_id),
+                (raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(40)),
+            )
+        )
+        assert cooled(cache, primary_id) and fell_back, [(o.status, o.detail[:80]) for o in fell_back]
+        assert_served(fell_back, chat_identity)
+        backed: Final = tuple(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(4))
+        assert_served(backed, chat_identity)
+        received: Final = markers_received(failing)
+        assert received[: len(fell_back)] == markers(fell_back)
+        assert len(received) <= len(fell_back) + 1, received
+        assert {*markers(fell_back), *markers(backed)} <= set(markers_received(healthy))
+
+
 def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None:
     group: Final = new_group()
     policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}

tests/integration/sdk/test_aiohttp_session_rebuild_wire.py

@@ -11,11 +11,30 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import Final
 
 import pytest
-from pydantic import JsonValue, TypeAdapter
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
 
 CONFIGURED_KEEPALIVE_SECONDS: Final = 1
 IDLE_SECONDS: Final = 2
+EVICTED_CLIENTS: Final = 30
 RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]])
+CHAT_COMPLETION: Final = {
+    "id": "chatcmpl-scripted",
+    "object": "chat.completion",
+    "created": 1700000000,
+    "model": "gpt-5.4-nano",
+    "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
+    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+}
+
+
+class EvictionReport(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    cache_size: int
+    cached_clients: int
+    established_connections: int
+    unclosed_messages: tuple[str, ...]
+
 
 REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent(
     """
@@ -38,6 +57,47 @@ REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent(
 )
 
 
+EVICTED_CLIENTS_EXCHANGE: Final = textwrap.dedent(
+    """
+    import asyncio, gc, json, sys
+    import psutil
+    import litellm
+    from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
+
+    unclosed = []
+
+    def established(port: int) -> int:
+        connections = psutil.Process().net_connections("tcp")
+        return sum(1 for c in connections if c.status == "ESTABLISHED" and c.raddr and c.raddr.port == port)
+
+    async def one_call(base_url: str, index: int) -> None:
+        await litellm.acompletion(
+            model="openai/gpt-5.4-nano",
+            api_key=f"sk-distinct-{index}",
+            api_base=base_url,
+            messages=[{"role": "user", "content": "hi"}],
+        )
+
+    async def main(base_url: str, port: int, evicted: int) -> None:
+        asyncio.get_running_loop().set_exception_handler(lambda loop, ctx: unclosed.append(ctx["message"]))
+        cache = litellm.in_memory_llm_clients_cache
+        for index in range(cache.max_size_in_memory + evicted):
+            await one_call(base_url, index)
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+        await asyncio.sleep(0)
+        print(json.dumps({
+            "cache_size": cache.max_size_in_memory,
+            "cached_clients": len(cache.cache_dict),
+            "established_connections": established(port),
+            "unclosed_messages": unclosed,
+        }))
+
+    asyncio.run(main(sys.argv[1], int(sys.argv[2]), int(sys.argv[3])))
+    """
+)
+
+
 class _ConnectionCountingPeer(ThreadingHTTPServer):
     daemon_threads = True
 
@@ -62,7 +122,10 @@ class _ConnectionHandler(BaseHTTPRequestHandler):
 
     def do_POST(self) -> None:
         self.rfile.read(int(self.headers["Content-Length"]))
-        body: Final = json.dumps({"connection": self.connection_number}).encode()
+        payload: Final = (
+            CHAT_COMPLETION if self.path.endswith("/chat/completions") else {"connection": self.connection_number}
+        )
+        body: Final = json.dumps(payload).encode()
         self.send_response(200)
         self.send_header("Content-Type", "application/json")
         self.send_header("Content-Length", str(len(body)))
@@ -107,3 +170,32 @@ def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive
 ) -> None:
     observed: Final = _rebuilt_session_exchange(connection_counting_peer)
     assert observed == [{"connection": 1}, {"connection": 2}], observed
+
+
+def _evicted_clients_exchange(base_url: str) -> EvictionReport:
+    completed: Final = subprocess.run(
+        [
+            sys.executable,
+            "-P",
+            "-c",
+            EVICTED_CLIENTS_EXCHANGE,
+            base_url,
+            base_url.rsplit(":", 1)[1],
+            str(EVICTED_CLIENTS),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=300,
+        check=False,
+    )
+    assert completed.returncode == 0, completed.stderr
+    return EvictionReport.model_validate_json(completed.stdout)
+
+
+def test_sdk_clients_evicted_from_a_full_client_cache_leave_no_open_session_or_connection(
+    connection_counting_peer: str,
+) -> None:
+    report: Final = _evicted_clients_exchange(connection_counting_peer)
+    assert report.cached_clients == report.cache_size, report
+    assert report.unclosed_messages == (), report
+    assert report.established_connections <= report.cache_size, report

tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py

@@ -0,0 +1,67 @@
+"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose
+``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the
+caller's dict is left as the caller wrote it (LIT-10239)"""
+
+import json
+from collections.abc import Mapping
+from typing import Final
+
+import pytest
+from integration._support.client import eventually
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
+
+import litellm
+
+_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0"
+_PROMPT_TOKENS: Final = 5
+_COMPLETION_TOKENS: Final = 1
+_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1}
+_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two")
+
+
+def _bedrock(request: Request) -> Reply:
+    """Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named."""
+    guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"]
+    body: Final[dict[str, JsonValue]] = {
+        "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
+        "stopReason": "end_turn",
+        "usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6},
+        "metrics": {"latencyMs": 3},
+        "trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}},
+    }
+    return Reply(body=json.dumps(body).encode())
+
+
+def _guardrail_cost() -> float:
+    rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
+    return sum(units * rates[counter] for counter, units in _USAGE.items())
+
+
+def _llm_cost() -> float:
+    prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")]
+    return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"]
+
+
+def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None:
+    billed: list[float] = []  # mutable-ok: the success callback appends what each call was billed
+    litellm.success_callback.append(
+        lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"])
+    )
+    shared_metadata: Final = {"tag": "x"}
+    with wire_server(_bedrock) as wire:
+        for guardrail_id in _GUARDRAIL_IDS:
+            litellm.completion(
+                model=_MODEL,
+                messages=[{"role": "user", "content": "hello"}],
+                api_base=wire.url,
+                aws_access_key_id="AKIASCRIPTED",
+                aws_secret_access_key="scripted",
+                aws_region_name="us-east-1",
+                guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"},
+                metadata=shared_metadata,
+            )
+    costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30)
+
+    assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS))
+    assert shared_metadata == {"tag": "x"}

tests/integration/sdk/test_router_cooldown_fallback_sdk.py

@@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter
 
 import litellm
 from litellm import Router
+from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments
 from litellm.types.router import DeploymentTypedDict
 
 _MODEL: Final = "gpt-5.4"
@@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No
         "/primary/chat/completions",
         "/backup/chat/completions",
     )
+
+
+@pytest.mark.parametrize(
+    "fallbacks",
+    [[{"primary": "backup"}], [{"*": "backup"}]],
+    ids=["string_fallback_target", "string_catch_all_target"],
+)
+@pytest.mark.asyncio
+async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(fallbacks: list[dict[str, object]]) -> None:
+    with wire_server(_peer) as wire:
+        router: Final = Router(
+            model_list=[
+                _deployment("primary", f"{wire.url}/primary", "primary-deployment"),
+                _deployment("backup", f"{wire.url}/backup", "backup-deployment"),
+            ],
+            allowed_fails=0,
+            num_retries=0,
+        )
+        router.fallbacks = fallbacks
+        with pytest.raises(litellm.AuthenticationError):
+            await router.acompletion(
+                model="primary",
+                messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}],
+            )
+        cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None)
+
+    assert tuple(cooled) == ("primary-deployment",)

tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py

@@ -0,0 +1,189 @@
+"""Two Vertex deployments share a base model but live in different projects with different
+buckets. An unmanaged gs:// batch created on the second one must be polled through that
+deployment (LIT-10326). Before the fix the cost poller matched the base model name in the
+gs:// path to the first deployment in the list and asked the wrong project for the job."""
+
+import json
+import re
+from collections.abc import Callable
+from contextlib import ExitStack
+from pathlib import Path
+from typing import Final
+from uuid import uuid4
+
+import pytest
+import yaml
+from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
+from integration._support.process import owned_proxy_process
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
+
+LOCATION: Final = "us-central1"
+BASE_MODEL: Final = "gemini-2.5-flash"
+TEAMS: Final = ("a", "b")
+_JOB_PATH: Final = re.compile(
+    r"^/v1/projects/(?P<project>[^/]+)/locations/(?P<location>[^/]+)/batchPredictionJobs(?:/(?P<job>\d+))?$"
+)
+
+
+def _project(team: str) -> str:
+    return f"scripted-project-{team}"
+
+
+def _deployment(team: str) -> str:
+    return f"{BASE_MODEL}-team-{team}"
+
+
+def _input_file(team: str) -> str:
+    return f"gs://scripted-bucket-{team}/litellm-vertex-files/publishers/google/models/{BASE_MODEL}/{uuid4().hex}/input.jsonl"
+
+
+def _job_number(team: str) -> str:
+    return f"{TEAMS.index(team) + 1}{uuid4().int % 10**17:017d}"
+
+
+def _vertex_peer(jobs: dict[str, tuple[str, str]]) -> Callable[[Request], Reply]:
+    """A Vertex double that remembers which project each job was created in. A retrieve
+    from any other project answers the 404 real Vertex gives (the job does not exist
+    there), a retrieve from the right one reports the job still running."""
+
+    def respond(request: Request) -> Reply:
+        path: Final = _JOB_PATH.match(request.target.split("?", 1)[0])
+        assert path is not None, request.target
+        assert request.headers.get("authorization") == "Bearer scripted-token", request.headers
+        project, location, job = path.group("project"), path.group("location"), path.group("job")
+        if request.method == "POST":
+            team: Final = project.rsplit("-", 1)[-1]
+            uri: Final = re.search(r'"(gs://[^"]+)"', request.body.decode())
+            assert uri is not None and uri.group(1).startswith(f"gs://scripted-bucket-{team}/"), (project, request.body)
+            created, created_uri = jobs.setdefault(project, (_job_number(team), uri.group(1)))
+            return Reply(body=json.dumps(_job(project, location, created, created_uri, "JOB_STATE_PENDING")).encode())
+        assert request.method == "GET" and job is not None, (request.method, request.target)
+        known: Final = jobs.get(project)
+        if known is None or known[0] != job:
+            return Reply(
+                status=404,
+                body=json.dumps(
+                    {
+                        "error": {
+                            "code": 404,
+                            "message": f"The BatchPredictionJob does not exist: {job}",
+                            "status": "NOT_FOUND",
+                        }
+                    }
+                ).encode(),
+            )
+        return Reply(body=json.dumps(_job(project, location, job, known[1], "JOB_STATE_RUNNING")).encode())
+
+    return respond
+
+
+def _job(project: str, location: str, job: str, input_uri: str, state: str) -> dict[str, JsonValue]:
+    return {
+        "name": f"projects/{project}/locations/{location}/batchPredictionJobs/{job}",
+        "displayName": "litellm-vertex-batch-scripted",
+        "model": f"publishers/google/models/{BASE_MODEL}",
+        "inputConfig": {"gcsSource": {"uris": [input_uri]}, "instancesFormat": "jsonl"},
+        "outputConfig": {
+            "predictionsFormat": "jsonl",
+            "gcsDestination": {"outputUriPrefix": input_uri.rsplit("/", 1)[0]},
+        },
+        "state": state,
+        "createTime": "2026-10-30T20:00:00.000000Z",
+        "updateTime": "2026-10-30T20:00:00.000000Z",
+    }
+
+
+def _owned_config(path: Path, gateway: Gateway, wire_url: str) -> Path:
+    config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    path.write_text(
+        yaml.safe_dump(
+            {
+                **config,
+                "model_list": [
+                    {
+                        "model_name": _deployment(team),
+                        "litellm_params": {
+                            "model": f"vertex_ai/{BASE_MODEL}",
+                            "api_base": wire_url,
+                            "vertex_project": _project(team),
+                            "vertex_location": LOCATION,
+                            "vertex_credentials": service_account_json(
+                                _project(team), gateway.upstream_url.rstrip("/")
+                            ),
+                            "gcs_bucket_name": f"scripted-bucket-{team}",
+                        },
+                    }
+                    for team in TEAMS
+                ],
+                "general_settings": {
+                    **object_value(config.get("general_settings") or {}),
+                    "track_unmanaged_batch_cost": True,
+                },
+                "litellm_settings": {
+                    **object_value(config.get("litellm_settings") or {}),
+                    "enable_loadbalancing_on_batch_endpoints": True,
+                    "user_url_allowed_hosts": ["127.0.0.1"],
+                },
+            }
+        )
+    )
+    return path
+
+
+def _create_batch(gateway: Gateway, key: str, team: str) -> str:
+    response: Final = gateway.request(
+        "POST",
+        "/v1/batches",
+        {
+            "input_file_id": _input_file(team),
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+            "model": _deployment(team),
+        },
+        key=key,
+    )
+    assert response.status_code == 200, response.text
+    return string_value(JSON_OBJECT.validate_json(response.content)["id"])
+
+
+def _retrieve_targets(requests: tuple[Request, ...], job: str) -> tuple[str, ...]:
+    return tuple(
+        request.target for request in requests if request.method == "GET" and request.target.endswith(f"/{job}")
+    )
+
+
+@pytest.mark.timeout(240)
+def test_unmanaged_batch_on_the_second_deployment_sharing_a_base_model_is_polled_through_its_own_project(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    jobs: Final[dict[str, tuple[str, str]]] = {}  # mutable-ok: the wire double records which project created each job
+    with ExitStack() as stack:
+        wire: Final = stack.enter_context(wire_server(_vertex_peer(jobs)))
+        config: Final = _owned_config(tmp_path / "two-vertex-teams.yaml", gateway, wire.url)
+        owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config))
+        scenario: Final = stack.enter_context(owned.gateway.scenario())
+        key: Final = scenario.key(models=[_deployment(team) for team in TEAMS])
+
+        batch_a: Final = _create_batch(owned.gateway, key, "a")
+        batch_b: Final = _create_batch(owned.gateway, key, "b")
+        assert (batch_a, batch_b) == (jobs[_project("a")][0], jobs[_project("b")][0]), (batch_a, batch_b, jobs)
+
+        seen: Final[
+            list[Request]
+        ] = []  # mutable-ok: drain empties the wire queue, so samples accumulate across retries
+
+        def polled_so_far() -> tuple[Request, ...]:
+            seen.extend(wire.drain())
+            return tuple(seen)
+
+        polled: Final = eventually(
+            polled_so_far,
+            lambda requests: bool(_retrieve_targets(requests, batch_b)) and bool(_retrieve_targets(requests, batch_a)),
+            seconds=90,
+        )
+        for team, batch in (("a", batch_a), ("b", batch_b)):
+            assert all(
+                target.startswith(f"/v1/projects/{_project(team)}/") for target in _retrieve_targets(polled, batch)
+            ), (team, _retrieve_targets(polled, batch))

tests/unit/caching/_event_loop_reuse.py

@@ -0,0 +1,20 @@
+from collections.abc import Callable
+from typing import Final, TypeVar
+
+import pytest
+
+LoopResult = TypeVar("LoopResult")
+
+
+def runs_until_a_loop_address_repeats(
+    run_on_new_loop: Callable[[], tuple[int, LoopResult]],
+    earlier: tuple[tuple[int, LoopResult], ...] = (),
+    limit: int = 50,
+) -> tuple[tuple[int, LoopResult], ...]:
+    if len(earlier) == limit:
+        pytest.skip(f"none of {limit} collected loops had its address handed to a later loop")
+    latest: Final = run_on_new_loop()
+    runs: Final = (*earlier, latest)
+    if latest[0] in {loop_id for loop_id, _ in earlier}:
+        return runs
+    return runs_until_a_loop_address_repeats(run_on_new_loop, runs, limit)

tests/unit/caching/test_evicted_client_closer.py

@@ -10,6 +10,7 @@ never closed, because litellm does not own its lifecycle.
 import asyncio
 import gc
 import weakref
+from typing import Final
 from unittest.mock import AsyncMock
 
 import httpx
@@ -18,6 +19,7 @@ from redis.asyncio import ConnectionPool, Redis
 
 from litellm.caching.evicted_client_closer import EvictedClientCloser
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
 
 
 class FakeClock:
@@ -41,6 +43,14 @@ class AsyncClient:
         self.closed = True
 
 
+class RecordingAsyncClient:
+    def __init__(self) -> None:
+        self.close_calls = 0
+
+    async def aclose(self) -> None:
+        self.close_calls += 1
+
+
 class SyncClient:
     def __init__(self) -> None:
         self.closed = False
@@ -266,6 +276,100 @@ def test_sync_client_evicted_outside_an_event_loop_is_still_closed():
     assert closer.pending_count == 0
 
 
+def test_async_client_from_collected_loop_is_not_closed_by_fresh_loops() -> None:
+    clock: Final = FakeClock()
+    closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
+
+    async def reap_then_schedule(client: RecordingAsyncClient) -> None:
+        closer.reap()
+        await asyncio.sleep(0)
+        closer.schedule(client)
+
+    def run_on_new_loop() -> tuple[int, RecordingAsyncClient]:
+        client: Final = RecordingAsyncClient()
+        closer.mark_owned(client)
+        loop: Final = asyncio.new_event_loop()
+        loop_id: Final = id(loop)
+        try:
+            loop.run_until_complete(reap_then_schedule(client))
+        finally:
+            loop.close()
+        del loop
+        gc.collect()
+        return loop_id, client
+
+    runs: Final = runs_until_a_loop_address_repeats(run_on_new_loop)
+
+    assert all(client.close_calls == 0 for _, client in runs)
+
+
+def test_async_client_queued_on_a_live_loop_is_closed_only_by_that_loop() -> None:
+    clock: Final = FakeClock()
+    closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock)
+    client: Final = RecordingAsyncClient()
+    closer.mark_owned(client)
+    owner_loop: Final = asyncio.new_event_loop()
+    other_loop: Final = asyncio.new_event_loop()
+
+    async def schedule() -> None:
+        closer.schedule(client)
+
+    async def reap() -> None:
+        closer.reap()
+        await asyncio.sleep(0)
+
+    try:
+        owner_loop.run_until_complete(schedule())
+        other_loop.run_until_complete(reap())
+        closes_after_other_loop_reaped: Final = client.close_calls
+        pending_after_other_loop_reaped: Final = closer.pending_count
+        owner_loop.run_until_complete(reap())
+    finally:
+        other_loop.close()
+        owner_loop.close()
+
+    assert closes_after_other_loop_reaped == 0
+    assert pending_after_other_loop_reaped == 1
+    assert client.close_calls == 1
+    assert closer.pending_count == 0
+
+
+def test_closes_queued_on_a_collected_loop_free_capacity() -> None:
+    clock: Final = FakeClock()
+    closer: Final = EvictedClientCloser(grace_seconds=0, max_pending=1, clock=clock)
+    first: Final = RecordingAsyncClient()
+    closer.mark_owned(first)
+    loop_a: Final = asyncio.new_event_loop()
+
+    async def schedule_first() -> None:
+        closer.schedule(first)
+
+    try:
+        loop_a.run_until_complete(schedule_first())
+    finally:
+        loop_a.close()
+    del loop_a
+    gc.collect()
+
+    second: Final = RecordingAsyncClient()
+    closer.mark_owned(second)
+    loop_b: Final = asyncio.new_event_loop()
+
+    async def schedule_and_reap_second() -> None:
+        closer.schedule(second)
+        closer.reap()
+        await asyncio.sleep(0)
+
+    try:
+        loop_b.run_until_complete(schedule_and_reap_second())
+    finally:
+        loop_b.close()
+
+    assert first.close_calls == 0
+    assert second.close_calls == 1
+    assert closer.pending_count == 0
+
+
 @pytest.mark.asyncio
 async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
     clock = FakeClock()
@@ -289,8 +393,7 @@ async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped():
 
 
 @pytest.mark.asyncio
-async def test_a_client_evicted_on_another_event_loop_is_left_alone():
-    """Closing a client bound to a different loop would schedule work on that loop."""
+async def test_a_client_evicted_on_a_collected_loop_is_dropped():
     clock = FakeClock()
     closer = make_closer(clock)
     client = AsyncClient()
@@ -310,7 +413,7 @@ async def test_a_client_evicted_on_another_event_loop_is_left_alone():
     await asyncio.sleep(0.05)
 
     assert client.closed is False
-    assert closer.pending_count == 1
+    assert closer.pending_count == 0
 
 
 @pytest.mark.asyncio

tests/unit/caching/test_llm_caching_handler.py

@@ -10,12 +10,16 @@ See: https://github.com/BerriAI/litellm/pull/22247
 
 import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref
 import warnings
+from typing import Final
 
+import aiohttp
 import pytest
+from openai import AsyncOpenAI
 
 
 from litellm.caching.evicted_client_closer import EvictedClientCloser
 from litellm.caching.llm_caching_handler import LLMClientCache
+from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.custom_httpx.http_handler import(
@@ -26,6 +30,41 @@ from litellm.llms.custom_httpx.http_handler import(
 from litellm.types.utils import LlmProviders
 from litellm.utils import _invalidate_model_cost_lowercase_map
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
+from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats
+
+
+async def _get_and_set_cache(cache: LLMClientCache) -> object | None:
+    cached: Final = await cache.async_get_cache("client")
+    await cache.async_set_cache("client", object())
+    return cached
+
+
+def _cache_result_on_new_loop(cache: LLMClientCache) -> tuple[int, object | None]:
+    loop: Final = asyncio.new_event_loop()
+    loop_id: Final = id(loop)
+    try:
+        result: Final = loop.run_until_complete(_get_and_set_cache(cache))
+    finally:
+        loop.close()
+    del loop
+    gc.collect()
+    return loop_id, result
+
+
+async def _get_cache_key(cache: LLMClientCache) -> str:
+    return cache.update_cache_key_with_event_loop("client")
+
+
+def _cache_key_on_new_loop(cache: LLMClientCache) -> tuple[int, str]:
+    loop: Final = asyncio.new_event_loop()
+    loop_id: Final = id(loop)
+    try:
+        key: Final = loop.run_until_complete(_get_cache_key(cache))
+    finally:
+        loop.close()
+    del loop
+    gc.collect()
+    return loop_id, key
 
 
 class MockAsyncClient:
@@ -48,6 +87,44 @@ class MockSyncClient:
         self.closed = True
 
 
+def test_async_client_cache_does_not_reuse_clients_between_collected_loops() -> None:
+    cache: Final = LLMClientCache()
+    runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_result_on_new_loop(cache))
+
+    assert all(result is None for _, result in runs)
+
+
+def test_async_client_cache_keys_differ_for_every_collected_loop() -> None:
+    cache: Final = LLMClientCache()
+    runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_key_on_new_loop(cache))
+    keys: Final = tuple(key for _, key in runs)
+
+    assert len(set(keys)) == len(keys)
+
+
+class _LoopWithoutWeakReferences:
+    __slots__ = ()
+
+
+def _cache_key_while_running(cache: LLMClientCache, loop: _LoopWithoutWeakReferences) -> str:
+    asyncio._set_running_loop(loop)
+    try:
+        return cache.update_cache_key_with_event_loop("client")
+    finally:
+        asyncio._set_running_loop(None)
+
+
+def test_cache_keys_stay_per_loop_on_loops_without_weak_references() -> None:
+    cache: Final = LLMClientCache()
+    loop_a: Final = _LoopWithoutWeakReferences()
+    loop_b: Final = _LoopWithoutWeakReferences()
+
+    first_key_on_a: Final = _cache_key_while_running(cache, loop_a)
+
+    assert _cache_key_while_running(cache, loop_a) == first_key_on_a
+    assert _cache_key_while_running(cache, loop_b) != first_key_on_a
+
+
 @pytest.mark.asyncio
 async def test_remove_key_does_not_close_async_client():
     """
@@ -204,6 +281,36 @@ async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elap
     assert client.closed is True
 
 
+@pytest.mark.asyncio
+async def test_an_evicted_owned_client_collected_inside_its_grace_window_still_closes_its_aiohttp_session():
+    """
+    The closer holds an evicted client weakly for the whole grace window. An SDK
+    process that exits, or a proxy that evicts and never calls the cache again,
+    collects the client first, so only its own finalizer can still close the
+    aiohttp session (LIT-9767).
+    """
+    clock = _FakeClock()
+    cache = LLMClientCache(
+        max_size_in_memory=2,
+        evicted_client_closer=EvictedClientCloser(grace_seconds=900.0, clock=clock),
+    )
+    session = aiohttp.ClientSession()
+    client = AsyncOpenAI(
+        api_key="sk-evicted",
+        http_client=httpx.AsyncClient(transport=LiteLLMAiohttpTransport(client=session)),
+    )
+    cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600)
+
+    cache.ttl_dict = {key: 0 for key in cache.ttl_dict}
+    cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap]
+    cache.evict_cache()
+    del client
+    gc.collect()
+    await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+
+    assert session.closed, "the evicted client was collected inside the grace window and left its session open"
+
+
 @pytest.mark.asyncio
 async def test_evicted_caller_supplied_client_is_never_closed():
     """litellm does not own a client the caller passed in, so it must stay open."""

tests/unit/enterprise/proxy/test_managed_files_hook.py

@@ -1420,6 +1420,8 @@ def _make_object_store_instance():
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock()
     mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock()
+    mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+    mock_prisma.db.writer = mock_prisma.db
 
     return (
         PROXY_LiteLLMManagedFiles(
@@ -2144,3 +2146,117 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets
     assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1")
     create_data = file_table.upsert.await_args.kwargs["data"]["create"]
     assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1")
+
+
+@pytest.mark.asyncio
+async def test_create_stores_the_creating_deployment_in_file_object():
+    """Regression (LIT-10326): an unmanaged gs:// batch stores the raw Vertex job id, which
+    says nothing about which of two deployments sharing a base model created it. The create
+    must keep the router's deployment id inside the stored file_object so the cost poller
+    reads the result through that deployment instead of the first base-model match."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    creator: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None)
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_object=_make_batch_response(status="validating"),
+        litellm_parent_otel_span=None,
+        model_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_purpose="batch",
+        user_api_key_dict=creator,
+        persist_attribution=True,
+        create_if_missing=True,
+        creating_model_id="deploy-team-b",
+    )
+
+    created: Final = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]
+    assert json.loads(created["file_object"])["model_id"] == "deploy-team-b"
+
+
+@pytest.mark.asyncio
+async def test_poll_refresh_keeps_the_stored_creating_deployment():
+    """A later observe of the row (a GET through whichever deployment the client named)
+    rewrites file_object; it must carry the create-time model_id forward, not drop it or
+    replace it with the deployment that served the retrieve."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    stored_row: Final = MagicMock()
+    stored_row.file_object = json.dumps(
+        {**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"}
+    )
+    mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row)
+    observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None)
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_object=_make_batch_response(status="completed"),
+        litellm_parent_otel_span=None,
+        model_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_purpose="batch",
+        user_api_key_dict=observer,
+        persist_attribution=False,
+        create_if_missing=False,
+    )
+
+    written: Final = json.loads(
+        mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]
+    )
+    assert written["model_id"] == "deploy-team-b"
+    assert written["status"] == "completed"
+
+
+@pytest.mark.asyncio
+async def test_poll_refresh_reads_the_stored_creating_deployment_from_the_writer():
+    """With a read replica configured, a retrieve that lands before the replica has the create
+    must still find the create-time model_id, so the lookup goes to the writer. Prisma may hand
+    the Json column back already decoded, so the row here is a dict rather than JSON text."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    stored_row: Final = MagicMock()
+    stored_row.file_object = {
+        **json.loads(_make_batch_response(status="validating").model_dump_json()),
+        "model_id": "deploy-team-b",
+    }
+    mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+    mock_prisma.db.writer = MagicMock()
+    mock_prisma.db.writer.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row)
+    observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None)
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_object=_make_batch_response(status="completed"),
+        litellm_parent_otel_span=None,
+        model_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_purpose="batch",
+        user_api_key_dict=observer,
+        persist_attribution=False,
+        create_if_missing=False,
+    )
+
+    written: Final = json.loads(
+        mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]
+    )
+    assert written["model_id"] == "deploy-team-b"
+    mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_id():
+    """The post-call hook on a batch create hands the router's model_id to the object store;
+    a retrieve of the same batch hands None so it cannot overwrite the creator."""
+    from litellm_enterprise.proxy.hooks.managed_files import BATCH_CREATE_HIDDEN_PARAM
+
+    managed_files: Final = _make_managed_files_instance()
+    created: Final = _make_batch_response(model_id="deploy-team-b", output_file_id=None, status="validating")
+    created._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True
+    retrieved: Final = _make_batch_response(model_id="deploy-team-a", output_file_id=None, status="in_progress")
+
+    with patch("litellm.proxy.proxy_server.llm_router", MagicMock()):
+        await managed_files.async_post_call_success_hook(
+            data={}, user_api_key_dict=_make_user_api_key_dict(), response=created
+        )
+        await managed_files.async_post_call_success_hook(
+            data={}, user_api_key_dict=_make_user_api_key_dict(), response=retrieved
+        )
+
+    first, second = managed_files.store_unified_object_id.await_args_list
+    assert first.kwargs["creating_model_id"] == "deploy-team-b"
+    assert second.kwargs["creating_model_id"] is None

tests/unit/integrations/test_agentops.py

@@ -1,10 +1,14 @@
 import os
-from unittest.mock import MagicMock, patch
+from typing import Final
+from unittest.mock import patch
 
 import pytest
+import respx
 
-
+import litellm
+from litellm.caching.llm_caching_handler import LLMClientCache
 from litellm.integrations.agentops.agentops import AgentOps, AgentOpsConfig
+from litellm.llms.custom_httpx.http_handler import get_httpx_client
 
 
 @pytest.fixture
@@ -111,3 +115,17 @@ def test_agentops_initialization_no_auth():
     assert agentops.resource_attributes["deployment.environment"] == "test_env"
     assert agentops.resource_attributes["telemetry.sdk.name"] == "agentops"
     assert "project.id" not in agentops.resource_attributes
+
+
+def test_auth_fetch_leaves_the_shared_sync_client_open_for_other_callers(
+    monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
+) -> None:
+    monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
+    auth_endpoint: Final = "https://agentops.example/v3/auth/token"
+    respx_mock.post(auth_endpoint).respond(json={"token": "jwt", "project_id": "project-1"})
+    client_another_caller_is_using: Final = get_httpx_client().client
+
+    agentops: Final = AgentOps(config=AgentOpsConfig(api_key="key", auth_endpoint=auth_endpoint))
+
+    assert agentops.resource_attributes["project.id"] == "project-1"
+    assert not client_another_caller_is_using.is_closed

tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py

@@ -1,7 +1,11 @@
+import json
+from typing import Final
 
 import pytest
+from pydantic import JsonValue
 
 import litellm
+from litellm.litellm_core_utils.llm_cost_calc import guardrail_cost
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
     bedrock_guardrail_cost,
     bedrock_guardrail_cost_by_unit,
@@ -230,3 +234,118 @@ def test_guardrail_information_cost_skips_malformed_entry_keeps_siblings():
     ]
     assert guardrail_information_cost(entries) == pytest.approx(0.0003)
     assert guardrail_information_cost({"guardrail_cost": 0.5, "guardrail_cost_in_spend": "maybe"}) == 0.0
+
+
+_TRACE_USAGE_IN: Final = {
+    "topicPolicyUnits": 1,
+    "contentPolicyUnits": 1,
+    "wordPolicyUnits": 1,
+    "sensitiveInformationPolicyUnits": 0,
+    "contextualGroundingPolicyUnits": 0,
+    "automatedReasoningPolicies": 0,
+}
+_TRACE_USAGE_OUT: Final = {
+    **_TRACE_USAGE_IN,
+    "topicPolicyUnits": 2,
+    "contentPolicyUnits": 2,
+    "sensitiveInformationPolicyUnits": 2,
+}
+_TRACE_PRICING: Final = {
+    "bedrock/guardrails": {
+        "litellm_provider": "bedrock",
+        "mode": "guardrails",
+        "guardrail_cost_per_unit": {
+            "topicPolicyUnits": 0.00015,
+            "contentPolicyUnits": 0.00015,
+            "wordPolicyUnits": 0.0,
+            "sensitiveInformationPolicyUnits": 0.0001,
+        },
+    }
+}
+
+
+def _converse_trace(action: str = "NONE") -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "inputAssessment": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}},
+            "outputAssessments": {
+                "gid": [
+                    {
+                        "topicPolicy": {"topics": [{"name": "t", "type": "DENY", "action": action}]},
+                        "invocationMetrics": {"usage": _TRACE_USAGE_OUT},
+                    }
+                ]
+            },
+        }
+    }
+
+
+def _invoke_trace() -> dict[str, JsonValue]:
+    return {
+        "guardrail": {
+            "input": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}},
+            "outputs": [{"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_OUT}}}],
+        }
+    }
+
+
+@pytest.mark.parametrize("trace", [_converse_trace(), _invoke_trace()], ids=["converse", "invoke"])
+def test_native_trace_sums_every_assessment_and_prices_it_like_the_proxy_guardrail(monkeypatch, trace) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    entry: Final = guardrail_cost.native_bedrock_guardrail_information(
+        trace, aws_region_name=None, guardrail_identifier=None
+    )
+    assert entry is not None
+    assert entry["guardrail_usage"] == {
+        "topicPolicyUnits": 3,
+        "contentPolicyUnits": 3,
+        "wordPolicyUnits": 2,
+        "sensitiveInformationPolicyUnits": 2,
+        "contextualGroundingPolicyUnits": 0,
+    }
+    assert entry["guardrail_cost"] == pytest.approx(3 * 0.00015 + 3 * 0.00015 + 2 * 0.0001)
+    assert entry["guardrail_cost_by_unit"]["contextualGroundingPolicyUnits"] is None
+    assert entry["guardrail_name"] == "gid"
+    assert entry["guardrail_provider"] == "bedrock"
+    assert entry["guardrail_status"] == "success"
+    assert guardrail_cost.guardrail_information_cost([entry]) == pytest.approx(entry["guardrail_cost"])
+
+
+def test_native_trace_marks_a_blocked_assessment_intervened_and_prefers_the_configured_identifier(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    entry: Final = guardrail_cost.native_bedrock_guardrail_information(
+        _converse_trace("BLOCKED"), aws_region_name="us-east-1", guardrail_identifier="configured-id"
+    )
+    assert entry is not None
+    assert entry["guardrail_status"] == "guardrail_intervened"
+    assert entry["guardrail_name"] == "configured-id"
+
+
+def test_native_trace_without_usage_is_not_billed(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    unpriced: Final = {"guardrail": {"inputAssessment": {"gid": {"topicPolicy": {"topics": []}}}}}
+    assert guardrail_cost.native_bedrock_guardrail_information(unpriced, None, None) is None
+    assert guardrail_cost.bedrock_guardrail_trace({"usage": {}}) is None
+
+
+def test_with_guardrail_entry_records_the_same_trace_once() -> None:
+    entry: Final = {"guardrail_name": "gid", "guardrail_cost": 0.001}
+    metadata: dict = {"standard_logging_guardrail_information": {"guardrail_name": "hook"}}
+    guardrail_cost.with_guardrail_entry(metadata, entry)  # pyright: ignore[reportArgumentType]  # partial entry
+    guardrail_cost.with_guardrail_entry(metadata, entry)  # pyright: ignore[reportArgumentType]  # partial entry
+    assert metadata["standard_logging_guardrail_information"] == [{"guardrail_name": "hook"}, entry]
+
+
+def test_native_trace_entry_redacts_the_pii_matches_before_they_reach_spend_logs(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING)
+    trace: Final = _converse_trace()
+    trace["guardrail"]["inputAssessment"]["gid"]["sensitiveInformationPolicy"] = {
+        "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}]
+    }
+    entry: Final = guardrail_cost.native_bedrock_guardrail_information(
+        trace, aws_region_name=None, guardrail_identifier=None
+    )
+    assert entry is not None
+    assert entry["guardrail_status"] == "guardrail_intervened"
+    assert "123-45-6789" not in json.dumps(entry["guardrail_response"])
+    assert entry["guardrail_usage"]["sensitiveInformationPolicyUnits"] == 2

tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py

@@ -11,7 +11,8 @@ import respx
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
+from litellm.types.llms.base import CachedTokensDetails
+from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, PromptTokensDetailsWrapper, Usage
 
 TIER_MODEL: Final = "tier-priced-test-model"
 TIER_ROW: Final[Mapping[str, float]] = MappingProxyType(
@@ -31,6 +32,204 @@ PROMPT_TOKENS: Final = 1000
 CACHED_TOKENS: Final = 200
 COMPLETION_TOKENS: Final = 500
 TIER_API_BASE: Final = "https://tier-pricing.invalid/v1"
+AUDIO_TIER_MODEL: Final = "tier-priced-gemini-audio-test-model"
+AUDIO_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType(
+    {
+        "input_cost_per_token": 4e-06,
+        "output_cost_per_token": 8e-06,
+        "cache_read_input_token_cost": 1e-06,
+        "input_cost_per_audio_token": 3e-06,
+        "cache_read_input_audio_token_cost": 2e-06,
+        "input_cost_per_audio_token_flex": 1e-06,
+        "input_cost_per_audio_token_priority": 2e-06,
+        "cache_read_input_audio_token_cost_flex": 4e-07,
+        "cache_read_input_audio_token_cost_priority": 8e-07,
+    }
+)
+
+
+@pytest.mark.parametrize("image_tokens", (60, 100))
+@pytest.mark.parametrize(
+    ("service_tier", "priority_rate", "image_rate"),
+    ((None, 11e-06, 7e-06), ("priority", 11e-06, 11e-06), ("flex", 11e-06, 7e-06), ("priority", 0.0, 0.0)),
+)
+def test_completion_cost_uses_priority_image_output_rate(
+    local_model_cost_map: None, service_tier: str | None, priority_rate: float, image_rate: float, image_tokens: int
+) -> None:
+    model: Final = "gemini-synthetic-image-priority"
+    litellm.register_model(
+        {
+            model: {
+                "litellm_provider": "vertex_ai",
+                "mode": "chat",
+                "input_cost_per_token": 3e-06,
+                "output_cost_per_token": 5e-06,
+                "output_cost_per_image_token": 7e-06,
+                "output_cost_per_image_token_priority": priority_rate,
+            }
+        }
+    )
+    response: Final = ModelResponse(
+        model=model,
+        usage=Usage(
+            prompt_tokens=20,
+            completion_tokens=100,
+            total_tokens=120,
+            completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=image_tokens),
+        ),
+    )
+    cost: Final = litellm.completion_cost(
+        completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier=service_tier
+    )
+    assert cost == pytest.approx(20 * 3e-06 + (100 - image_tokens) * 5e-06 + image_tokens * image_rate)
+
+
+@pytest.mark.parametrize("prompt_tokens", (199_999, 200_000, 200_001))
+@pytest.mark.parametrize(
+    ("service_tier", "flat_rate", "context_rate"),
+    ((None, 7e-06, 11e-06), ("priority", 13e-06, 17e-06), ("flex", 19e-06, 23e-06)),
+)
+def test_completion_cost_uses_total_prompt_for_audio_context_tiers(
+    local_model_cost_map: None, prompt_tokens: int, service_tier: str | None, flat_rate: float, context_rate: float
+) -> None:
+    model: Final = "gemini-synthetic-audio-context"
+    litellm.register_model(
+        {
+            model: {
+                "litellm_provider": "vertex_ai",
+                "mode": "chat",
+                "input_cost_per_token": 3e-06,
+                "output_cost_per_token": 5e-06,
+                "input_cost_per_audio_token": 7e-06,
+                "input_cost_per_audio_token_above_200k_tokens": 11e-06,
+                "input_cost_per_audio_token_priority": 13e-06,
+                "input_cost_per_audio_token_above_200k_tokens_priority": 17e-06,
+                "input_cost_per_audio_token_flex": 19e-06,
+                "input_cost_per_audio_token_above_200k_tokens_flex": 23e-06,
+            }
+        }
+    )
+    response: Final = ModelResponse(
+        model=model,
+        usage=Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=10,
+            total_tokens=prompt_tokens + 10,
+            prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=100),
+        ),
+    )
+    cost: Final = litellm.completion_cost(
+        completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier=service_tier
+    )
+    audio_rate: Final = context_rate if prompt_tokens > 200_000 else flat_rate
+    assert cost == pytest.approx((prompt_tokens - 100) * 3e-06 + 100 * audio_rate + 10 * 5e-06)
+
+
+@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001))
+def test_batch_cost_uses_audio_context_and_image_output_tiers(local_model_cost_map: None, prompt_tokens: int) -> None:
+    from litellm.cost_calculator import batch_cost_calculator
+
+    model: Final = "gemini-synthetic-batch-audio-context"
+    litellm.register_model(
+        {
+            model: {
+                "litellm_provider": "vertex_ai",
+                "mode": "chat",
+                "input_cost_per_token": 6e-06,
+                "output_cost_per_token": 10e-06,
+                "input_cost_per_token_batches": 3e-06,
+                "output_cost_per_token_batches": 5e-06,
+                "input_cost_per_token_above_200k_tokens_batches": 13e-06,
+                "output_cost_per_token_above_200k_tokens_batches": 17e-06,
+                "output_cost_per_image_token_batches": 19e-06,
+                "input_cost_per_audio_token_batches": 7e-06,
+                "input_cost_per_audio_token_above_200k_tokens_batches": 11e-06,
+            }
+        }
+    )
+    usage: Final = Usage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=20,
+        total_tokens=prompt_tokens + 20,
+        prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=100),
+        completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=10),
+    )
+    cost: Final = batch_cost_calculator(model=model, custom_llm_provider="vertex_ai", usage=usage)
+    audio_rate: Final = 11e-06 if prompt_tokens > 200_000 else 7e-06
+    input_rate: Final = 13e-06 if prompt_tokens > 200_000 else 3e-06
+    output_rate: Final = 17e-06 if prompt_tokens > 200_000 else 5e-06
+    assert sum(cost) == pytest.approx(
+        (prompt_tokens - 100) * input_rate + 100 * audio_rate + 10 * output_rate + 10 * 19e-06
+    )
+
+
+@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001))
+def test_completion_cost_uses_flex_context_input_and_output_rates(
+    local_model_cost_map: None, prompt_tokens: int
+) -> None:
+    model: Final = "gemini-synthetic-flex-context"
+    litellm.register_model(
+        {
+            model: {
+                "litellm_provider": "vertex_ai",
+                "mode": "chat",
+                "input_cost_per_token": 6e-06,
+                "output_cost_per_token": 10e-06,
+                "input_cost_per_token_flex": 3e-06,
+                "output_cost_per_token_flex": 5e-06,
+                "input_cost_per_token_above_200k_tokens_flex": 7e-06,
+                "output_cost_per_token_above_200k_tokens_flex": 11e-06,
+            }
+        }
+    )
+    response: Final = ModelResponse(
+        model=model, usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=10, total_tokens=prompt_tokens + 10)
+    )
+    cost: Final = litellm.completion_cost(
+        completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier="flex"
+    )
+    input_rate: Final = 7e-06 if prompt_tokens > 200_000 else 3e-06
+    output_rate: Final = 11e-06 if prompt_tokens > 200_000 else 5e-06
+    assert cost == pytest.approx(prompt_tokens * input_rate + 10 * output_rate)
+
+
+@pytest.mark.parametrize(
+    ("audio_fields", "audio_rate"),
+    (
+        ({}, 3e-06),
+        ({"input_cost_per_audio_token": 7e-06}, 7e-06),
+        ({"input_cost_per_audio_token_above_200k_tokens": 11e-06}, 11e-06),
+        ({"input_cost_per_audio_token_above_200k_tokens_priority": 0.0}, 0.0),
+    ),
+)
+def test_audio_context_rates_preserve_fallbacks_and_explicit_zero(
+    local_model_cost_map: None, audio_fields: Mapping[str, float], audio_rate: float
+) -> None:
+    model: Final = "gemini-synthetic-audio-fallback"
+    litellm.register_model(
+        {
+            model: {
+                "litellm_provider": "vertex_ai",
+                "mode": "chat",
+                "input_cost_per_token": 3e-06,
+                "output_cost_per_token": 5e-06,
+                **audio_fields,
+            }
+        }
+    )
+    response: Final = ModelResponse(
+        model=model,
+        usage=Usage(
+            prompt_tokens=200_001,
+            completion_tokens=10,
+            total_tokens=200_011,
+            prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=100),
+        ),
+    )
+    cost: Final = litellm.completion_cost(
+        completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier="priority"
+    )
+    assert cost == pytest.approx(199_901 * 3e-06 + 100 * audio_rate + 10 * 5e-06)
 
 
 def _cost_at(prices: Mapping[str, float], column_suffix: str) -> float:
@@ -78,6 +277,53 @@ def test_completion_cost_bills_the_price_columns_of_the_service_tier(
     assert cost == pytest.approx(_cost_at(TIER_ROW, column_suffix))
 
 
+@pytest.mark.parametrize(
+    ("service_tier", "audio_rate_key", "cache_audio_rate_key"),
+    (
+        pytest.param("flex", "input_cost_per_audio_token_flex", "cache_read_input_audio_token_cost_flex"),
+        pytest.param("priority", "input_cost_per_audio_token_priority", "cache_read_input_audio_token_cost_priority"),
+    ),
+)
+def test_gemini_audio_and_cached_audio_use_service_tier_rates(
+    local_model_cost_map: None, service_tier: str, audio_rate_key: str, cache_audio_rate_key: str
+) -> None:
+    litellm.register_model(
+        {
+            AUDIO_TIER_MODEL: {
+                "litellm_provider": "gemini",
+                "mode": "chat",
+                **AUDIO_TIER_ROW,
+            }
+        }
+    )
+    usage: Final = Usage(
+        prompt_tokens=40,
+        completion_tokens=0,
+        total_tokens=40,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            text_tokens=20,
+            audio_tokens=20,
+            cached_tokens=10,
+            cached_tokens_details=CachedTokensDetails(audio_tokens=10),
+        ),
+    )
+    response: Final = ModelResponse(model=AUDIO_TIER_MODEL, usage=usage)
+
+    cost: Final = litellm.completion_cost(
+        completion_response=response,
+        model=AUDIO_TIER_MODEL,
+        custom_llm_provider="gemini",
+        service_tier=service_tier,
+    )
+
+    expected: Final = (
+        20 * AUDIO_TIER_ROW["input_cost_per_token"]
+        + 10 * AUDIO_TIER_ROW[audio_rate_key]
+        + 10 * AUDIO_TIER_ROW[cache_audio_rate_key]
+    )
+    assert cost == pytest.approx(expected)
+
+
 LONG_CONTEXT_TIER_MODEL: Final = "long-context-tier-priced-test-model"
 LONG_CONTEXT_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType(
     {

tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py

@@ -9,6 +9,7 @@ from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import (
     used_pricing_keys,
     zero_cost_warning,
 )
+from litellm.types.llms.base import CachedTokensDetails
 from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage
 
 PER_SECOND_ENTRY: Final = {"cost_per_second": 0.00042}
@@ -135,6 +136,36 @@ def test_audio_tokens_need_the_audio_rates() -> None:
     assert diagnostic["missing_pricing_keys"] == ("input_cost_per_audio_token", "output_cost_per_audio_token")
 
 
+def test_cached_audio_uses_its_cache_rate_in_zero_cost_diagnostic() -> None:
+    usage: Final = Usage(
+        prompt_tokens=20,
+        completion_tokens=0,
+        total_tokens=20,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            text_tokens=0,
+            audio_tokens=20,
+            cached_tokens=20,
+            cached_tokens_details=CachedTokensDetails(audio_tokens=20),
+        ),
+    )
+    pricing_entry: Final = {
+        "input_cost_per_token": 1e-6,
+        "input_cost_per_audio_token": 2e-6,
+        "cache_read_input_audio_token_cost": 0,
+    }
+
+    assert used_pricing_keys(usage) == ("cache_read_input_audio_token_cost",)
+    assert (
+        diagnose_zero_cost(
+            usage=usage,
+            pricing_model="gemini-audio",
+            pricing_entry=pricing_entry,
+            calculation_failed=False,
+        )
+        is None
+    )
+
+
 def test_warning_names_the_request_the_entry_the_missing_keys_and_the_counter() -> None:
     diagnostic = diagnose_zero_cost(
         usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False

tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py

@@ -25,13 +25,13 @@ from typing import Any, Dict, Final, List, Optional, Tuple
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-
+from openai._legacy_response import LegacyAPIResponse
 
 import litellm
-from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.code_interpreter_interception.handler import (
     CodeInterpreterInterceptionLogger,
 )
+from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.chat_completion_agentic_loop import (
     maybe_run_chat_completion_agentic_loop,
 )
@@ -45,9 +45,9 @@ from litellm.types.integrations.custom_logger import (
     AgenticLoopRequestPatch,
 )
 from litellm.types.utils import (
+    ChatCompletionMessageToolCall,
     Choices,
     Function,
-    ChatCompletionMessageToolCall,
     Message,
     ModelResponse,
 )
@@ -142,7 +142,7 @@ def _raw_response_for(model_response: ModelResponse) -> MagicMock:
     (an object exposing `.headers` and `.parse()` -> something with model_dump)."""
     parsed = MagicMock()
     parsed.model_dump.return_value = model_response.model_dump()
-    raw = MagicMock()
+    raw = MagicMock(spec=LegacyAPIResponse)
     raw.headers = {}
     raw.parse.return_value = parsed
     return raw
@@ -171,7 +171,7 @@ async def test_internal_control_fields_never_leak_into_provider_body(restore_cal
         ]
     )
     mock_client = MagicMock()
-    mock_client.chat.completions.with_raw_response.create = create
+    mock_client.post = create
 
     response = await litellm.acompletion(
         model="openai/gpt-4o-mini",

tests/unit/litellm_core_utils/test_exception_mapping_utils.py

@@ -2462,7 +2462,6 @@ def test_openai_gateway_timeout_error():
     Test that the OpenAI gateway timeout error is raised
     """
     openai_client = OpenAI()
-    mapped_target = openai_client.chat.completions.with_raw_response  # type: ignore
 
     def _return_exception(*args, **kwargs):
 
@@ -2498,8 +2497,8 @@ def test_openai_gateway_timeout_error():
 
     with pytest.raises(litellm.Timeout) as exc_info:
         with patch.object(
-            mapped_target,
-            "create",
+            openai_client,
+            "post",
             side_effect=_return_exception,
         ):
             litellm.completion(

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -2444,6 +2444,44 @@ async def test_async_success_handler_runs_async_callbacks_in_the_post_response_p
     assert in_post_response_phase() is False, "the phase must not leak into the request task"
 
 
+@pytest.mark.parametrize("stream", [True, False], ids=["streaming", "non_streaming"])
+def test_reused_logging_obj_does_not_report_the_previous_attempts_first_token(logging_obj, stream: bool):
+    """A fallback reuses the request's logging object but runs in its own window.
+
+    When the first attempt already recorded a first-token time, the fallback must not
+    report it. Otherwise ``completion_start_time`` lands before the attempt's
+    ``start_time``, the callbacks see a negative TTFT, and latency-based routing caches
+    a negative sample that empties its buffer window.
+    """
+    logging_obj.stream = stream
+    logging_obj.model_call_details["litellm_params"] = {}
+    logging_obj.litellm_params = {}
+
+    first_attempt_start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0)
+    logging_obj.update_completion_start_time(first_attempt_start + datetime.timedelta(seconds=1))
+    fallback_start: Final = first_attempt_start + datetime.timedelta(seconds=600)
+    fallback_end: Final = fallback_start + datetime.timedelta(seconds=5)
+
+    observed: dict[str, object] = {}
+
+    class Recorder(CustomLogger):
+        def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
+            observed["completion_start_time"] = kwargs.get("completion_start_time")
+
+    with patch.object(logging_obj, "get_combined_callback_list", return_value=[Recorder()]):
+        logging_obj.success_handler(
+            result=_success_response(),
+            start_time=fallback_start,
+            end_time=fallback_end,
+        )
+
+    completion_start_time: Final = observed["completion_start_time"]
+    assert isinstance(completion_start_time, datetime.datetime)
+    assert completion_start_time >= fallback_start, (
+        "the fallback attempt reported a first-token time that predates its own start"
+    )
+
+
 def test_is_sync_litellm_request():
     assert LitellmLogging.is_sync_litellm_request({}) is True
     assert LitellmLogging.is_sync_litellm_request({"acompletion": True}) is False
@@ -12470,3 +12508,88 @@ def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_b
         )
     finally:
         litellm.model_cost.pop(deployment_id, None)
+
+
+def test_record_native_bedrock_guardrail_bills_when_the_metadata_is_none() -> None:
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    logging_obj = Logging(
+        model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
+        messages=[{"role": "user", "content": "hello"}],
+        stream=False,
+        call_type="completion",
+        start_time=datetime.datetime.now(),
+        litellm_call_id="call-none-metadata",
+        function_id="fn",
+    )
+    logging_obj.model_call_details["litellm_params"] = {"metadata": None}
+    trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}}
+
+    logging_obj.record_native_bedrock_guardrail({"trace": trace})
+
+    (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
+    assert entry["guardrail_usage"] == {"contentPolicyUnits": 2}
+    assert entry["guardrail_cost"] > 0
+
+
+def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_from_cache() -> None:
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    logging_obj = Logging(
+        model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
+        messages=[{"role": "user", "content": "hello"}],
+        stream=False,
+        call_type="completion",
+        start_time=datetime.datetime.now(),
+        litellm_call_id="call-cache-hit",
+        function_id="fn",
+    )
+    logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
+    trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}}
+    cached = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x")
+    setattr(cached, "trace", trace)
+    cached._hidden_params = {"cache_hit": True}
+
+    logging_obj.record_native_bedrock_guardrail(cached)
+    logging_obj.model_call_details["cache_hit"] = True
+    logging_obj.record_native_bedrock_guardrail({"trace": trace})
+
+    assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"]
+
+
+def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_calls_share_one_metadata_dict() -> None:
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
+
+    shared_metadata = {"tag": "x"}
+    traces = tuple(
+        {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": units}}}}}}
+        for units in (1, 2)
+    )
+    costs = []
+    for call_index, trace in enumerate(traces):
+        logging_obj = Logging(
+            model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
+            messages=[{"role": "user", "content": "hello"}],
+            stream=False,
+            call_type="completion",
+            start_time=datetime.datetime.now(),
+            litellm_call_id=f"call-shared-metadata-{call_index}",
+            function_id="fn",
+        )
+        logging_obj.model_call_details["litellm_params"] = {"metadata": shared_metadata}
+        logging_obj.record_native_bedrock_guardrail({"trace": trace})
+        recorded = logging_obj.model_call_details["litellm_params"]["metadata"][
+            "standard_logging_guardrail_information"
+        ]
+        costs.append(guardrail_information_cost(recorded))
+
+    (first_cost, second_cost) = costs
+    assert second_cost == pytest.approx(2 * first_cost)
+    assert shared_metadata == {"tag": "x"}

tests/unit/litellm_core_utils/test_redact_messages.py

@@ -1798,3 +1798,4 @@ async def test_redaction_with_metadata_completion_api():
     response = standard_logging_payload["response"]
     assert response["choices"][0]["message"]["content"] == "redacted-by-litellm"
     assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
+

tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -8,7 +8,6 @@ import httpx
 import pytest
 import respx
 
-
 import litellm
 from litellm import ChatCompletionUsageBlock, stream_chunk_builder
 from litellm.types.utils import GenericStreamingChunk
@@ -19,6 +18,7 @@ from litellm.types.utils import (
     ChatCompletionMessageToolCall,
     Delta,
     Function,
+    GenericStreamingChunk,
     ModelResponseStream,
     PromptTokensDetails,
     ServerToolUse,
@@ -194,11 +194,7 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
                 }
             ]
         ),
-        make_chunk(
-            thinking_blocks=[
-                {"type": "thinking", "thinking": None, "signature": "sig_block1"}
-            ]
-        ),
+        make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block1"}]),
         make_chunk(
             thinking_blocks=[
                 {
@@ -216,16 +212,10 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
                 }
             ]
         ),
-        make_chunk(
-            thinking_blocks=[
-                {"type": "thinking", "thinking": None, "signature": "sig_block2"}
-            ]
-        ),
+        make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block2"}]),
     ]
 
-    thinking_chunks = [
-        chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")
-    ]
+    thinking_chunks = [chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")]
     processor = ChunkProcessor(chunks=chunks)
     result = processor.get_combined_thinking_content(thinking_chunks)
 
@@ -295,9 +285,7 @@ def test_cache_read_input_tokens_retained():
             prompt_tokens=11779,
             total_tokens=11784,
             completion_tokens_details=None,
-            prompt_tokens_details=PromptTokensDetails(
-                audio_tokens=None, cached_tokens=11775
-            ),
+            prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=11775),
             cache_creation_input_tokens=4,
             cache_read_input_tokens=11775,
         ),
@@ -331,9 +319,7 @@ def test_cache_read_input_tokens_retained():
             prompt_tokens=0,
             total_tokens=214,
             completion_tokens_details=None,
-            prompt_tokens_details=PromptTokensDetails(
-                audio_tokens=None, cached_tokens=0
-            ),
+            prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=0),
             cache_creation_input_tokens=0,
             cache_read_input_tokens=0,
         ),
@@ -393,10 +379,7 @@ def test_streaming_preserves_anthropic_1hr_cache_creation_breakdown():
     )
     # Sanity: the delta event genuinely lacks the breakdown - this is the input
     # condition that used to defeat cost calc.
-    assert (
-        getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None)
-        is None
-    )
+    assert getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None) is None
 
     def _usage_chunk(usage, finish_reason):
         return ModelResponseStream(
@@ -515,9 +498,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk():
             prompt_tokens=1234,
             total_tokens=1239,
             completion_tokens_details=None,
-            prompt_tokens_details=PromptTokensDetails(
-                audio_tokens=None, cached_tokens=543
-            ).model_dump(),
+            prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=543).model_dump(),
         ),
         index=2,
     )
@@ -534,6 +515,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk():
 
     assert usage.prompt_tokens_details.cached_tokens == 543
 
+
 def test_stream_chunk_builder_litellm_usage_chunks():
     """
     Validate ChunkProcessor.calculate_usage uses provided usage fields from streaming chunks
@@ -607,9 +589,7 @@ def test_stream_chunk_builder_litellm_usage_chunks():
     chunks = [chunk1, chunk2]
     processor = ChunkProcessor(chunks=chunks)
 
-    usage = processor.calculate_usage(
-        chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output=""
-    )
+    usage = processor.calculate_usage(chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="")
 
     assert usage.prompt_tokens == 50
     assert usage.completion_tokens == 27
@@ -653,15 +633,11 @@ def test_calculate_usage_honors_openai_sdk_completion_usage_chunks():
         provider_specific_fields=None,
         stream_options={"include_usage": True},
     )
-    usage_chunk.usage = CompletionUsage(
-        prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704
-    )
+    usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704)
     assert type(usage_chunk.usage) is CompletionUsage
 
     chunks = [content_chunk, usage_chunk]
-    usage = ChunkProcessor(chunks=chunks).calculate_usage(
-        chunks=chunks, model="mantle-claude", completion_output=""
-    )
+    usage = ChunkProcessor(chunks=chunks).calculate_usage(chunks=chunks, model="mantle-claude", completion_output="")
 
     assert usage.prompt_tokens == 20
     assert usage.completion_tokens == 60
@@ -684,9 +660,7 @@ def test_get_model_from_chunks_azure_model_router():
         {"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []},
     ]
 
-    result = ChunkProcessor._get_model_from_chunks(
-        chunks=chunks, first_chunk_model="azure-model-router"
-    )
+    result = ChunkProcessor._get_model_from_chunks(chunks=chunks, first_chunk_model="azure-model-router")
 
     # Should return the actual model, not the request model
     assert result == "gpt-4.1-nano-2025-04-14"
@@ -697,9 +671,7 @@ def test_get_model_from_chunks_azure_model_router():
         {"model": "gpt-4", "id": "chatcmpl-456", "choices": []},
     ]
 
-    result_same = ChunkProcessor._get_model_from_chunks(
-        chunks=chunks_same_model, first_chunk_model="gpt-4"
-    )
+    result_same = ChunkProcessor._get_model_from_chunks(chunks=chunks_same_model, first_chunk_model="gpt-4")
 
     # Should return the first chunk's model when all are the same
     assert result_same == "gpt-4"
@@ -775,9 +747,7 @@ def test_stream_chunk_builder_anthropic_web_search():
     chunks = [chunk1, chunk2]
     processor = ChunkProcessor(chunks=chunks)
 
-    usage = processor.calculate_usage(
-        chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output=""
-    )
+    usage = processor.calculate_usage(chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output="")
 
     assert usage.prompt_tokens == 50
     assert usage.completion_tokens == 27
@@ -929,15 +899,11 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields():
         ],
     )
     chunk_dict = chunk.model_dump()
-    chunk_dict["_hidden_params"] = {
-        "provider_specific_fields": {"traffic_type": "default"}
-    }
+    chunk_dict["_hidden_params"] = {"provider_specific_fields": {"traffic_type": "default"}}
 
     response = stream_chunk_builder(chunks=[chunk_dict])
     assert response is not None
-    assert (
-        response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"
-    )
+    assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default"
 
 
 def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks():
@@ -982,10 +948,7 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks():
     assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata
     assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata
     assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata
-    assert (
-        response._hidden_params["vertex_ai_url_context_metadata"]
-        == url_context_metadata
-    )
+    assert response._hidden_params["vertex_ai_url_context_metadata"] == url_context_metadata
 
     dumped = response.model_dump()
     assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata
@@ -1032,9 +995,7 @@ def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata():
 
 def test_stream_chunk_builder_propagates_vertex_ai_safety_results():
     """Assembled response must expose safety data under the non-streaming field name."""
-    safety_ratings = [
-        [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}]
-    ]
+    safety_ratings = [[{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}]]
 
     chunk = ModelResponseStream(
         id="chatcmpl-vertex-safety",
@@ -1076,18 +1037,12 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks():
             )
         ],
     ).model_dump()
-    chunk_dict["_hidden_params"] = {
-        "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}]
-    }
+    chunk_dict["_hidden_params"] = {"vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}]}
 
     response = stream_chunk_builder(chunks=[chunk_dict])
     assert response is not None
-    assert getattr(response, "vertex_ai_grounding_metadata") == [
-        {"webSearchQueries": ["test query"]}
-    ]
-    assert response.model_dump()["vertex_ai_grounding_metadata"] == [
-        {"webSearchQueries": ["test query"]}
-    ]
+    assert getattr(response, "vertex_ai_grounding_metadata") == [{"webSearchQueries": ["test query"]}]
+    assert response.model_dump()["vertex_ai_grounding_metadata"] == [{"webSearchQueries": ["test query"]}]
 
 
 def test_cost_field_in_usage_chunks():
@@ -1096,29 +1051,21 @@ def test_cost_field_in_usage_chunks():
         id="chatcmpl-1",
         created=1745513206,
         model="openrouter/claude",
-        choices=[
-            StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
-        ],
+        choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))],
         usage=chunk1_usage,
     )
 
-    chunk2_usage = Usage(
-        completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
-    )
+    chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025)
     chunk2 = ModelResponseStream(
         id="chatcmpl-1",
         created=1745513207,
         model="openrouter/claude",
-        choices=[
-            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
-        ],
+        choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))],
         usage=chunk2_usage,
     )
 
     processor = ChunkProcessor(chunks=[chunk1, chunk2])
-    usage = processor.calculate_usage(
-        chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi"
-    )
+    usage = processor.calculate_usage(chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi")
 
     assert hasattr(usage, "cost")
     assert usage.cost == 0.00025
@@ -1201,25 +1148,19 @@ def test_prompt_tokens_details_survive_later_usage_chunk_without_details():
         id="chatcmpl-1",
         created=1745513206,
         model="openai/gpt-5.6-sol",
-        choices=[
-            StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
-        ],
+        choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))],
         usage=Usage(
             prompt_tokens=6017,
             completion_tokens=4,
             total_tokens=6021,
-            prompt_tokens_details=PromptTokensDetailsWrapper(
-                cached_tokens=6004, cache_write_tokens=10
-            ),
+            prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=6004, cache_write_tokens=10),
         ),
     )
     chunk_without_details = ModelResponseStream(
         id="chatcmpl-1",
         created=1745513207,
         model="openai/gpt-5.6-sol",
-        choices=[
-            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
-        ],
+        choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))],
         usage=Usage(prompt_tokens=6017, completion_tokens=4, total_tokens=6021),
     )
 
@@ -1557,9 +1498,7 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate(
     assert usage.completion_tokens_details.text_tokens == expected_text_tokens
 
 
-def _openai_chunk(
-    choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None
-) -> dict[str, object]:
+def _openai_chunk(choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None) -> dict[str, object]:
     base: Final = {
         "id": "chatcmpl-lit6552",
         "object": "chat.completion.chunk",
@@ -1994,7 +1933,11 @@ def test_stream_chunk_builder_litellm_tool_call(respx_mock: respx.MockRouter) ->
     )
 
     assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
-        {"function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, "id": "call_weather", "type": "function"}
+        {
+            "function": {"arguments": '{"city":"Boston"}', "name": "get_weather"},
+            "id": "call_weather",
+            "type": "function",
+        }
     ]
     assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (5, 3, 8)
 
@@ -2016,7 +1959,9 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message(respx_mock: resp
 
 
 def test_stream_chunk_builder_multiple_tool_calls() -> None:
-    def tool_call_chunk(*tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None) -> ModelResponseStream:
+    def tool_call_chunk(
+        *tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None
+    ) -> ModelResponseStream:
         return _legacy_stream_chunk("openai/gpt-4o", tool_calls=tool_calls, finish_reason=finish_reason)
 
     chunks: Final = (
@@ -2027,7 +1972,9 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None:
         ),
         tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='{"base": 3, '))),
         tool_call_chunk(
-            ChatCompletionDeltaToolCall(id="call_add", index=1, type="function", function=Function(name="add", arguments=""))
+            ChatCompletionDeltaToolCall(
+                id="call_add", index=1, type="function", function=Function(name="add", arguments="")
+            )
         ),
         tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"exponent": 5}'))),
         tool_call_chunk(
@@ -2041,7 +1988,11 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None:
     assert response is not None
     assert response.choices[0].finish_reason == "tool_calls"
     assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [
-        {"function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, "id": "call_exp", "type": "function"},
+        {
+            "function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"},
+            "id": "call_exp",
+            "type": "function",
+        },
         {"function": {"arguments": '{"x": 1, "y": 2}', "name": "add"}, "id": "call_add", "type": "function"},
     ]
 
@@ -2075,9 +2026,7 @@ def test_grok_bug() -> None:
     response: Final = stream_chunk_builder(chunks=chunks)
 
     assert response is not None
-    assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {
-        "location": "Boston"
-    }
+    assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {"location": "Boston"}
 
 
 def test_stream_chunk_builder_openai_audio_output_usage() -> None:
@@ -2180,3 +2129,50 @@ def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies():
     assert response is not None
     assert response.choices[0].message.content == "Hello world"
     assert response.choices[0].message.reasoning_items == reasoning_items
+
+
+def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it_on_the_logging_object(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A Bedrock stream reports its guardrail trace on the last chunk only; the built response carries it like the
+    non-streaming response and the call bills it as guardrail cost."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+    monkeypatch.setattr(
+        litellm,
+        "model_cost",
+        {
+            "bedrock/guardrails": {
+                "litellm_provider": "bedrock",
+                "mode": "guardrails",
+                "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015},
+            }
+        },
+    )
+    trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2}}}}}}
+    chunks = [
+        ModelResponseStream(id="s", model="m", choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]),
+        ModelResponseStream(
+            id="s",
+            model="m",
+            choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")],
+            provider_specific_fields={"trace": trace},
+        ),
+    ]
+    logging_obj = Logging(
+        model="bedrock/m",
+        messages=[{"role": "user", "content": "q"}],
+        stream=True,
+        call_type="acompletion",
+        start_time=None,
+        litellm_call_id="call",
+        function_id="fn",
+    )
+    logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
+    built = stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "q"}], logging_obj=logging_obj)
+    assert built is not None
+    assert getattr(built, "trace") == trace
+    (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
+    assert entry["guardrail_usage"] == {"topicPolicyUnits": 2}
+    assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015)

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -6621,3 +6621,102 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream():
     async with _wrapper_over(stream):
         pass
     assert stream.close_calls == 1
+
+
+def _final_chunk_with_dict_usage(choices) -> ModelResponseStream:
+    """A stream chunk whose ``usage`` is a plain dict.
+
+    ``ModelResponseStream.__init__`` coerces a dict ``usage`` kwarg into
+    ``litellm.Usage``, so the dict is attached after construction — the same
+    shape a custom ``streaming_decoder`` (or any third-party iterator that
+    bypasses litellm chunk construction) delivers to ``chunk_creator``.
+    """
+    chunk: Final = ModelResponseStream(
+        id="chatcmpl-dict-usage-test",
+        created=1754900000,
+        model="gpt-4o",
+        object="chat.completion.chunk",
+        choices=choices,
+    )
+    chunk.usage = {  # mutable-ok: wire-shaped usage payload under test
+        "prompt_tokens": 100,
+        "completion_tokens": 10,
+        "total_tokens": 110,
+        "prompt_tokens_details": {"cached_tokens": 80},  # mutable-ok: wire payload
+        "completion_tokens_details": {"reasoning_tokens": 5},  # mutable-ok: wire payload
+        "cache_creation_input_tokens": 42,
+    }
+    return chunk
+
+
+def test_dispatch_provider_chunk_dict_usage_preserves_token_details(
+    logging_obj: Logging,
+):
+    """The openai else-branch of ``_dispatch_provider_chunk`` handles three
+    usage shapes: ``Usage`` (passed through), ``BaseModel``
+    (``Usage(**model_dump())``), and plain ``dict``. The dict arm used to
+    rebuild ``litellm.Usage`` from only prompt/completion/total tokens,
+    silently dropping ``prompt_tokens_details`` (cached_tokens),
+    ``completion_tokens_details`` (reasoning_tokens), and cache token counts
+    that both sibling arms preserve.
+    """
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=None,
+        model="gpt-4o",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+        stream_options={"include_usage": True},  # mutable-ok: constructor payload
+    )
+    chunk: Final = _final_chunk_with_dict_usage(
+        choices=[
+            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=None))
+        ]  # mutable-ok: chunk payload
+    )
+    model_response: Final = wrapper.model_response_creator()
+    wrapper._dispatch_provider_chunk(
+        chunk=chunk,
+        model_response=model_response,
+        completion_obj={"content": ""},  # mutable-ok: dispatch scratch dict
+    )
+
+    usage: Final = model_response.usage
+    assert isinstance(usage, Usage)
+    assert usage.prompt_tokens == 100
+    assert usage.completion_tokens == 10
+    assert usage.total_tokens == 110
+    assert usage.prompt_tokens_details is not None, (
+        "dict-usage arm dropped prompt_tokens_details (its Usage/BaseModel sibling arms preserve it)"
+    )
+    assert usage.prompt_tokens_details.cached_tokens == 80
+    assert usage.completion_tokens_details is not None
+    assert usage.completion_tokens_details.reasoning_tokens == 5
+    assert getattr(usage, "cache_creation_input_tokens", None) == 42
+
+
+def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details(
+    logging_obj: Logging,
+):
+    """End-to-end through ``chunk_creator``: a usage-only final chunk (empty
+    ``choices``, ``stream_options.include_usage`` set) takes the early-return
+    path, so whatever the dict-usage arm wrote is exactly what the consumer
+    receives — there is no later aggregation to repair it.
+    """
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=None,
+        model="gpt-4o",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+        stream_options={"include_usage": True},  # mutable-ok: constructor payload
+    )
+    chunk: Final = _final_chunk_with_dict_usage(choices=[])  # mutable-ok: empty choices list
+    result: Final = wrapper.chunk_creator(chunk=chunk)
+
+    assert result is not None
+    usage: Final = result.usage
+    assert isinstance(usage, Usage)
+    assert usage.prompt_tokens_details is not None, (
+        "usage-only early return delivered a Usage stripped of prompt_tokens_details to the stream consumer"
+    )
+    assert usage.prompt_tokens_details.cached_tokens == 80
+    assert usage.completion_tokens_details is not None
+    assert usage.completion_tokens_details.reasoning_tokens == 5

tests/unit/litellm_core_utils/test_token_counter.py

@@ -1401,6 +1401,34 @@ def test_anthropic_image_block_with_empty_base64_data():
     assert tokens > 0
 
 
+def test_token_counter_counts_anthropic_search_result_text() -> None:
+    from litellm.litellm_core_utils.token_counter import _count_content_list
+
+    search_result: Final = {
+        "type": "search_result",
+        "source": "https://example.invalid/notes",
+        "title": "Project notes",
+        "content": [{"type": "text", "text": "The code word is 7f3a."}],
+        "citations": {"enabled": True},
+    }
+
+    assert _count_content_list(len, [search_result], False, None) == len(
+        "https://example.invalid/notesProject notesThe code word is 7f3a."
+    )
+
+
+@pytest.mark.parametrize("nested", (False, True))
+def test_unknown_content_type_remains_a_value_error_with_an_optional_fallback(nested: bool) -> None:
+    from litellm.litellm_core_utils.token_counter import _count_content_list
+
+    unknown: Final = {"type": "future_block"}
+    content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [unknown]}] if nested else [unknown]
+
+    with pytest.raises(ValueError, match="Invalid content item type: future_block"):
+        _count_content_list(len, content, False, None)
+    assert _count_content_list(len, content, False, 7) == 7
+
+
 def test_anthropic_image_block_without_source_raises():
     """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior."""
     from litellm.litellm_core_utils.token_counter import _count_content_list

tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py

@@ -1,9 +1,9 @@
 import json
-import traceback
-from typing import Callable, Optional
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import pytest
+from openai._legacy_response import LegacyAPIResponse
 
 import litellm
 from litellm import ModelResponse
@@ -130,7 +130,7 @@ def test_azure_o3_streaming():
         api_version="2024-02-15-preview",
     )
 
-    with patch.object(client.chat.completions.with_raw_response, "create") as mock_create:
+    with patch.object(client, "post") as mock_create:
         try:
             completion(
                 model="azure/o3-mini",
@@ -141,7 +141,7 @@ def test_azure_o3_streaming():
         except Exception as e:
             print(e)
         assert mock_create.call_count == 1
-        assert "stream" in mock_create.call_args.kwargs
+        assert "stream" in mock_create.call_args.kwargs["body"]
 
 
 def test_azure_o_series_routing():
@@ -157,7 +157,7 @@ def test_azure_o_series_routing():
         api_version="2024-02-15-preview",
     )
 
-    with patch.object(client.chat.completions.with_raw_response, "create") as mock_create:
+    with patch.object(client, "post") as mock_create:
         try:
             completion(
                 model="azure/o_series/my-random-deployment-name",
@@ -168,17 +168,19 @@ def test_azure_o_series_routing():
         except Exception as e:
             print(e)
         assert mock_create.call_count == 1
-        assert "stream" not in mock_create.call_args.kwargs
+        assert "stream" not in mock_create.call_args.kwargs["body"]
 
 
 @patch("litellm.main.azure_o1_chat_completions._get_openai_client")
 def test_openai_o_series_max_retries_0(mock_get_openai_client):
     import litellm
 
-    mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {}
-    mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = (
-        ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}])
-    )
+    raw_response: Final = MagicMock(spec=LegacyAPIResponse)
+    raw_response.headers = {}
+    raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}])
+    client_stub: Final = MagicMock()
+    client_stub.post.return_value = raw_response
+    mock_get_openai_client.return_value = client_stub
     litellm.set_verbose = True
     response = litellm.completion(
         model="azure/o1-preview",
@@ -233,7 +235,7 @@ async def test_azure_o1_series_response_format_extra_params():
     response_format = {"type": "json_object"}
     tool_choice = "auto"
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             await litellm.acompletion(
@@ -251,7 +253,7 @@ async def test_azure_o1_series_response_format_extra_params():
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         print("request_body: ", json.dumps(request_body, indent=4))
         assert request_body["tools"] == tools

tests/unit/llms/azure/test_azure.py

@@ -313,7 +313,7 @@ def test_azure_openai_gpt_4o_naming(monkeypatch):
         number: str = Field(description="total number of days in a week")
         days: list[str] = Field(description="name of days in a week")
 
-    with patch.object(client.chat.completions.with_raw_response, "create") as mock_post:
+    with patch.object(client, "post") as mock_post:
         try:
             completion(
                 model="azure/gpt4o",
@@ -326,9 +326,9 @@ def test_azure_openai_gpt_4o_naming(monkeypatch):
 
         mock_post.assert_called_once()
 
-        print(mock_post.call_args.kwargs)
+        print(mock_post.call_args.kwargs["body"])
 
-        assert "tool_calls" not in mock_post.call_args.kwargs
+        assert "tool_calls" not in mock_post.call_args.kwargs["body"]
 
 
 @pytest.mark.parametrize(
@@ -339,7 +339,6 @@ def test_azure_openai_gpt_4o_naming(monkeypatch):
     ],
 )
 def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
-    from litellm import completion
     from typing import Optional
     from pydantic import BaseModel
     import litellm
@@ -382,7 +381,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
         }
     ]
 
-    with patch.object(client.chat.completions.with_raw_response, "create") as mock_post:
+    with patch.object(client, "post") as mock_post:
         mock_post.return_value.headers = {}
         mock_post.return_value.parse.return_value = litellm.ModelResponse(
             choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}]
@@ -411,9 +410,9 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
         mock_post.assert_called_once()
 
         if api_version == "2024-10-21":
-            assert "response_format" in mock_post.call_args.kwargs
+            assert "response_format" in mock_post.call_args.kwargs["body"]
         else:
-            assert "response_format" not in mock_post.call_args.kwargs
+            assert "response_format" not in mock_post.call_args.kwargs["body"]
         assert response.choices[0].message.content == InvestigationOutput().model_dump_json()
 
 
@@ -609,7 +608,6 @@ def test_azure_with_content_safety_error():
     """
     Verify user can access innererror from the Azure OpenAI exception
     """
-    from litellm import completion
     from litellm.exceptions import ContentPolicyViolationError
     from litellm.litellm_core_utils.exception_mapping_utils import exception_type
     from unittest.mock import MagicMock

tests/unit/llms/bedrock/chat/test_invoke_handler.py

@@ -19,6 +19,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.llms.bedrock.chat.invoke_handler import (
+    AmazonAnthropicClaudeStreamDecoder,
     AmazonOpenAICompatibleStreamDecoder,
     AWSEventStreamDecoder,
     _stream_decoder,
@@ -1120,6 +1121,7 @@ async def test_converse_streaming_completion_yields_text_and_finish_reason(
             ),
         )
     )
+
     def handle_request(request: httpx.Request) -> httpx.Response:
         assert request.url.path.endswith("/converse-stream")
         return httpx.Response(200, content=event_stream_body, headers=_event_stream_headers())
@@ -1179,9 +1181,7 @@ def test_mistral_invoke_stream_yields_text_and_finish_reason(
     assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == (
         "A court case moves through the courts."
     )
-    assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == [
-        "stop"
-    ]
+    assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
 
 
 def test_converse_streaming_tool_call_arguments_are_valid_json(
@@ -1721,6 +1721,7 @@ def test_filter_headers_for_aws_signature():
 @pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
 def test_invoke_decoder_reports_only_missing_aws_dependency(missing):
     from unittest.mock import patch
+
     from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
 
     failure = ModuleNotFoundError("missing dependency", name=missing)
@@ -1835,9 +1836,7 @@ async def test_converse_stream_thinking_merges_into_content_with_think_tags() ->
             "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "c2lnbmF0dXJl"}}}
         ),
         _converse_event_frame("contentBlockStop", {"contentBlockIndex": 0}),
-        _converse_event_frame(
-            "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}}
-        ),
+        _converse_event_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}}),
         _converse_event_frame(
             "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "I am a helpful assistant"}}
         ),

tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py

@@ -3835,3 +3835,32 @@ def test_bedrock_removed_tool_change_does_not_add_beta(explicit_beta: bool) -> N
 
     assert result["messages"] == [{"role": "user", "content": "Reply with OK"}]
     assert result.get("anthropic_beta", []).count(beta) == int(explicit_beta)
+
+
+@pytest.mark.parametrize("trace", ["enabled", "disabled"])
+def test_bedrock_messages_applies_the_deployment_guardrail_config_as_invoke_headers(trace: str) -> None:
+    from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
+        AmazonAnthropicClaudeMessagesConfig,
+    )
+
+    headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
+        headers={},
+        model="anthropic.claude-sonnet-5-v1:0",
+        messages=[],
+        optional_params={},
+        litellm_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "2", "trace": trace}},
+    )
+    assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "gid"
+    assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "2"
+    assert headers["X-Amzn-Bedrock-Trace"] == trace.upper()
+
+
+def test_bedrock_messages_without_guardrail_config_sends_no_guardrail_headers() -> None:
+    from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
+        AmazonAnthropicClaudeMessagesConfig,
+    )
+
+    headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
+        headers={}, model="anthropic.claude-sonnet-5-v1:0", messages=[], optional_params={}, litellm_params={}
+    )
+    assert not any(name.startswith("X-Amzn-Bedrock-Guardrail") for name in headers)

tests/unit/llms/bedrock/test_cost_calculation.py

@@ -0,0 +1,41 @@
+import json
+from importlib.resources import files
+from typing import Final
+
+import pytest
+
+import litellm
+from litellm.types.utils import Usage
+
+CLAUDE_PLATFORM_PREFIX: Final = "bedrock/claude_platform/"
+CLAUDE_PLATFORM_KEYS: Final = tuple(
+    key
+    for key in json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text())
+    if key.startswith(CLAUDE_PLATFORM_PREFIX)
+)
+
+
+def test_claude_platform_rows_exist() -> None:
+    assert CLAUDE_PLATFORM_KEYS
+
+
+@pytest.mark.parametrize("inference_geo", ["global", "us"])
+@pytest.mark.parametrize("key", CLAUDE_PLATFORM_KEYS)
+def test_claude_platform_cost_matches_claude_api_cost(local_model_cost_map: None, key: str, inference_geo: str) -> None:
+    usage: Final = Usage(
+        prompt_tokens=1_000,
+        completion_tokens=200,
+        cache_read_input_tokens=300,
+        cache_creation_input_tokens=250,
+        inference_geo=inference_geo,
+    )
+
+    claude_platform_cost: Final = litellm.cost_per_token(
+        model=key.removeprefix("bedrock/"), custom_llm_provider="bedrock", usage_object=usage
+    )
+    claude_api_cost: Final = litellm.cost_per_token(
+        model=key.removeprefix(CLAUDE_PLATFORM_PREFIX), custom_llm_provider="anthropic", usage_object=usage
+    )
+
+    assert claude_platform_cost == pytest.approx(claude_api_cost)
+    assert sum(claude_platform_cost) > 0

tests/unit/llms/custom_httpx/test_aiohttp_transport.py

@@ -1,8 +1,11 @@
 import aiohttp as aiohttp_aiohttp_handler, asyncio, importlib, litellm
 import concurrent.futures
+import gc
 import socket
 import sys
 import struct
+import threading
+import weakref
 from typing import Final
 
 import aiohttp
@@ -1310,3 +1313,128 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat
             assert len(received_requests) == 1
     finally:
         await handler.close()
+
+
+async def _new_session() -> aiohttp.ClientSession:
+    return aiohttp.ClientSession()
+
+
+async def _until_closed(session: aiohttp.ClientSession) -> None:
+    while not session.closed:
+        await asyncio.sleep(0)
+
+
+@pytest.mark.asyncio
+async def test_a_collected_transport_closes_the_session_it_owns():
+    """An SDK client evicted from the cache is collected long before the closer's grace window ends."""
+    session = aiohttp.ClientSession()
+    transport = LiteLLMAiohttpTransport(client=session)
+
+    del transport
+    gc.collect()
+    await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+
+    assert session.closed, "the owned session outlived the transport that was collected before closing it"
+
+
+@pytest.mark.asyncio
+async def test_a_collected_transport_leaves_a_shared_session_open():
+    session = aiohttp.ClientSession()
+    try:
+        transport = LiteLLMAiohttpTransport(client=session, owns_session=False)
+
+        del transport
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+
+        assert not session.closed, "a transport closed a session it does not own"
+    finally:
+        await session.close()
+
+
+def test_a_transport_collected_after_its_loop_closed_still_closes_its_session(caplog: pytest.LogCaptureFixture):
+    """asyncio.run() closes the loop the session lives on before the evicted client is collected."""
+    loop = asyncio.new_event_loop()
+    session = loop.run_until_complete(_new_session())
+    loop.close()
+    transport = LiteLLMAiohttpTransport(client=session)
+
+    del transport
+    gc.collect()
+
+    assert session.closed, "the owned session stayed open after its transport was collected on a closed loop"
+    del session
+    gc.collect()
+    assert "Unclosed client session" not in caplog.text, caplog.text
+
+
+def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_on_that_loop():
+    loop = asyncio.new_event_loop()
+    thread = threading.Thread(target=loop.run_forever, daemon=True)
+    thread.start()
+    session = asyncio.run_coroutine_threadsafe(_new_session(), loop).result(timeout=10)
+    transport = LiteLLMAiohttpTransport(client=session)
+    try:
+        del transport
+        gc.collect()
+        asyncio.run_coroutine_threadsafe(_until_closed(session), loop).result(timeout=10)
+
+        assert session.closed, "the owned session stayed open after its transport was collected from another thread"
+    finally:
+        loop.call_soon_threadsafe(loop.stop)
+        thread.join(timeout=10)
+        loop.close()
+
+
+@pytest.mark.asyncio
+async def test_a_stream_keeps_reading_after_the_client_that_opened_it_is_collected():
+    from aiohttp import web
+
+    release: Final = asyncio.Event()
+    body: Final = tuple(f"chunk{i}\n".encode() for i in range(4))
+
+    async def handler(request: web.Request) -> web.StreamResponse:
+        response: Final = web.StreamResponse()
+        await response.prepare(request)
+        await response.write(body[0])
+        await release.wait()
+        for chunk in body[1:]:
+            await response.write(chunk)
+        await response.write_eof()
+        return response
+
+    app: Final = web.Application()
+    app.router.add_get("/stream", handler)
+    runner: Final = web.AppRunner(app)
+    await runner.setup()
+    sock: Final = socket.socket()
+    sock.bind(("127.0.0.1", 0))
+    port: Final[int] = sock.getsockname()[1]
+    await web.SockSite(runner, sock).start()
+
+    transport: Final = LiteLLMAiohttpTransport(client=aiohttp.ClientSession)
+    client: Final = httpx.AsyncClient(transport=transport)
+    transport_ref: Final = weakref.ref(transport)
+    try:
+        response: Final = await client.send(client.build_request("GET", f"http://127.0.0.1:{port}/stream"), stream=True)
+        session: Final = transport.client
+        assert isinstance(session, aiohttp.ClientSession)
+        chunks: Final = response.aiter_bytes()
+        first: Final = await chunks.__anext__()
+
+        del client, transport
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+        release.set()
+
+        rest: Final = [chunk async for chunk in chunks]
+        assert first + b"".join(rest) == b"".join(body), (first, rest)
+        assert not session.closed, "the session was closed under a stream that was still being read"
+
+        del response, chunks
+        gc.collect()
+        await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks)
+        assert transport_ref() is None, "the finished stream kept the transport alive"
+        assert session.closed, "the owned session stayed open after its last stream finished"
+    finally:
+        await runner.cleanup()

tests/unit/llms/custom_httpx/test_http_handler.py

@@ -1,6 +1,5 @@
 import asyncio, importlib, json, time
 import gc
-import io
 import os
 import pathlib
 import ssl
@@ -1388,6 +1387,55 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
     assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled)
 
 
+async def _drop_unused_handler() -> None:
+    AsyncHTTPHandler(timeout=61.0)
+    gc.collect()
+
+
+def test_finalizer_close_does_not_keep_a_loop_closed_right_after_it_alive():
+    loop: Final = asyncio.new_event_loop()
+    loop.run_until_complete(_drop_unused_handler())
+    loop.close()
+    closed_loop: Final = weakref.ref(loop)
+    del loop
+    gc.collect()
+
+    assert closed_loop() is None
+
+
+class _TransportFailingToClose(httpx.AsyncBaseTransport):
+    async def aclose(self) -> None:
+        raise RuntimeError("close failed")
+
+
+async def _drop_handler_whose_close_fails() -> None:
+    AsyncHTTPHandler(timeout=61.0, transport=_TransportFailingToClose())
+    gc.collect()
+
+
+def test_failed_finalizer_close_in_a_loops_last_step_reports_no_unretrieved_exception():
+    loop: Final = asyncio.new_event_loop()
+    exception_handler: Final = MagicMock()
+    loop.set_exception_handler(exception_handler)
+    loop.run_until_complete(_drop_handler_whose_close_fails())
+    loop.close()
+    del loop
+    gc.collect()
+
+    exception_handler.assert_not_called()
+
+
+async def _sync_client_with_timeout_62() -> HTTPHandler:
+    return get_httpx_client(params={"timeout": 62.0})
+
+
+def test_sync_client_is_shared_by_short_lived_loops_and_callers_outside_a_loop():
+    built_in_first_loop: Final = asyncio.run(_sync_client_with_timeout_62())
+
+    assert asyncio.run(_sync_client_with_timeout_62()) is built_in_first_loop
+    assert get_httpx_client(params={"timeout": 62.0}) is built_in_first_loop
+
+
 @pytest.mark.asyncio
 async def test_sync_close_helper_respects_session_ownership():
     """The loop-less fallback closes only sessions the transport owns; a

tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py

@@ -1,15 +1,16 @@
 import json
+import re
 from io import BytesIO
 from pathlib import Path
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
+import pytest
 from openai import AsyncOpenAI, OpenAI
+from openai._legacy_response import LegacyAPIResponse
 from openai.types import CreateEmbeddingResponse, Embedding
 from openai.types.create_embedding_response import Usage
-import pytest
-import re
 
 import litellm
 from litellm import completion
@@ -18,7 +19,6 @@ from litellm.llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfi
 from tests._vcr_conftest_common import rewound_new_episodes_cassette
 from tests.capturing_transport import CapturingTransport
 
-
 _GATEWAY_EMBEDDING_RESPONSE: Final = CreateEmbeddingResponse(
     object="list",
     data=(Embedding(object="embedding", index=0, embedding=(0.1, 0.2, 0.3)),),
@@ -98,7 +98,7 @@ async def test_litellm_gateway_from_sdk():
     openai_client = OpenAI(api_key="fake-key")
 
     with patch.object(
-        openai_client.chat.completions.with_raw_response, "create", new=MagicMock()
+        openai_client, "post", new=MagicMock()
     ) as mock_call:
         try:
             completion(
@@ -114,9 +114,9 @@ async def test_litellm_gateway_from_sdk():
 
         mock_call.assert_called_once()
 
-        print("Call KWARGS - {}".format(mock_call.call_args.kwargs))
+        print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"]))
 
-        assert "hello" in mock_call.call_args.kwargs["extra_body"]
+        assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"]
 
 
 @pytest.mark.parametrize("is_async", (False, True))
@@ -264,7 +264,7 @@ async def test_litellm_gateway_from_sdk_rerank(is_async):
         patch_target = client.post
 
     with patch.object(client, "post", new=mock_method):
-        mock_response = MagicMock()
+        mock_response = MagicMock(spec=LegacyAPIResponse)
 
         # Create a mock response similar to OpenAI's rerank response
         mock_response.text = json.dumps(
@@ -408,7 +408,7 @@ async def test_litellm_gateway_from_sdk_structured_output():
     openai_client = OpenAI(api_key="fake-key")
 
     with patch.object(
-        openai_client.chat.completions, "create", new=MagicMock()
+        openai_client, "post", new=MagicMock()
     ) as mock_call:
         try:
             litellm.completion(
@@ -427,8 +427,8 @@ async def test_litellm_gateway_from_sdk_structured_output():
 
         mock_call.assert_called_once()
 
-        print("Call KWARGS - {}".format(mock_call.call_args.kwargs))
-        json_schema = mock_call.call_args.kwargs["response_format"]
+        print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"]))
+        json_schema = mock_call.call_args.kwargs["body"]["response_format"]
         assert "json_schema" in json_schema
 
 
@@ -486,7 +486,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers():
     openai_client = OpenAI(api_key="fake-key")
 
     # Create mock response object
-    mock_response = MagicMock()
+    mock_response = MagicMock(spec=LegacyAPIResponse)
     mock_response.headers = {"x-litellm-response-cost": "120"}
     mock_response.parse.return_value = litellm.ModelResponse(
         **{
@@ -528,8 +528,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers():
     )
 
     with patch.object(
-        openai_client.chat.completions.with_raw_response,
-        "create",
+        openai_client, "post",
         return_value=mock_response,
     ) as mock_call:
         response = litellm.completion(
@@ -564,7 +563,7 @@ async def test_litellm_gateway_image_generation_direct(is_async):
         "created": 1,
         "data": [{"url": "https://example.com/image.png"}],
     }
-    mock_raw_response = MagicMock()
+    mock_raw_response = MagicMock(spec=LegacyAPIResponse)
     mock_raw_response.parse.return_value = mock_openai_response
     mock_raw_response.headers = {}
 

tests/unit/llms/nvidia_nim/test_nvidia_nim.py

@@ -1,6 +1,6 @@
 import json
 from typing import Final
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, patch
 
 import httpx
 import pytest
@@ -8,7 +8,7 @@ from openai.types import CreateEmbeddingResponse, Embedding
 from openai.types.create_embedding_response import Usage as EmbeddingUsage
 
 import litellm
-from litellm import ModelResponse, completion
+from litellm import completion
 from tests.capturing_transport import CapturingTransport
 
 
@@ -94,7 +94,7 @@ def test_chat_completion_nvidia_nim_with_tools():
     ]
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             completion(
@@ -116,7 +116,7 @@ def test_chat_completion_nvidia_nim_with_tools():
 
         # Add assertions to check the request
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         print("request_body: ", request_body)
 
@@ -208,7 +208,7 @@ def test_completion_nvidia_nim():
         api_key="fake-api-key",
     )
 
-    with patch.object(client.chat.completions.with_raw_response, "create") as mock_client:
+    with patch.object(client, "post") as mock_client:
         try:
             completion(
                 model=model_name,
@@ -226,7 +226,7 @@ def test_completion_nvidia_nim():
             print(e)
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         assert request_body["messages"] == [
             {

tests/unit/llms/openai/test_o_series_transformation.py

@@ -1,5 +1,4 @@
-from typing import Final
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 
@@ -70,7 +69,7 @@ async def test_o1_handle_system_role(model):
     client = AsyncOpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             await litellm.acompletion(
@@ -83,7 +82,7 @@ async def test_o1_handle_system_role(model):
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         print("request_body: ", request_body)
 
@@ -141,7 +140,7 @@ async def test_o1_max_completion_tokens(model: str):
     client = AsyncOpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             await litellm.acompletion(
@@ -154,7 +153,7 @@ async def test_o1_max_completion_tokens(model: str):
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         print("request_body: ", request_body)
 

tests/unit/llms/openai/test_openai.py

@@ -5,12 +5,13 @@ import os
 from collections.abc import Mapping
 from types import ModuleType
 from typing import Final, Literal, Protocol
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, patch
 
 import httpx
 import pytest
 import respx
 from openai import AsyncOpenAI, OpenAI
+from openai._legacy_response import LegacyAPIResponse
 from openai._types import Response as SDKResponse
 from pydantic import BaseModel, ConfigDict, JsonValue
 
@@ -1596,7 +1597,7 @@ async def test_openai_prediction_param_mock():
     client = AsyncOpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             await litellm.acompletion(
@@ -1615,7 +1616,7 @@ async def test_openai_prediction_param_mock():
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         # Verify the request contains the prediction parameter
         assert "prediction" in request_body
@@ -1627,10 +1628,12 @@ async def test_openai_prediction_param_mock():
 def test_openai_max_retries_0(mock_get_openai_client):
     import litellm
 
-    mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {}
-    mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = (
-        ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}])
-    )
+    raw_response: Final = MagicMock(spec=LegacyAPIResponse)
+    raw_response.headers = {}
+    raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}])
+    client_stub: Final = MagicMock()
+    client_stub.post.return_value = raw_response
+    mock_get_openai_client.return_value = client_stub
     litellm.set_verbose = True
     response = litellm.completion(
         model="gpt-4o-mini",
@@ -1739,7 +1742,7 @@ async def test_openai_safety_identifier_parameter():
     client = AsyncOpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             await litellm.acompletion(
@@ -1752,7 +1755,7 @@ async def test_openai_safety_identifier_parameter():
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         # Verify the request contains the safety_identifier parameter
         assert "safety_identifier" in request_body
@@ -1768,7 +1771,7 @@ def test_openai_safety_identifier_parameter_sync():
     client = OpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             litellm.completion(
@@ -1781,7 +1784,7 @@ def test_openai_safety_identifier_parameter_sync():
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         # Verify the request contains the safety_identifier parameter
         assert "safety_identifier" in request_body
@@ -1798,7 +1801,7 @@ async def test_openai_service_tier_parameter():
     client = AsyncOpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             await litellm.acompletion(
@@ -1811,7 +1814,7 @@ async def test_openai_service_tier_parameter():
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         # Verify the request contains the service_tier parameter
         assert "service_tier" in request_body, "service_tier should be in request body"
@@ -1829,7 +1832,7 @@ def test_openai_service_tier_parameter_sync():
     client = OpenAI(api_key="fake-api-key")
 
     with patch.object(
-        client.chat.completions.with_raw_response, "create"
+        client, "post"
     ) as mock_client:
         try:
             litellm.completion(
@@ -1842,7 +1845,7 @@ def test_openai_service_tier_parameter_sync():
             print(f"Error: {e}")
 
         mock_client.assert_called_once()
-        request_body = mock_client.call_args.kwargs
+        request_body = mock_client.call_args.kwargs["body"]
 
         # Verify the request contains the service_tier parameter
         assert "service_tier" in request_body, "service_tier should be in request body"
@@ -1874,3 +1877,63 @@ def test_responses_gpt54_with_xhigh_reasoning():
         assert request_body["model"] == "openai/gpt-5.4"
 
         assert request_body["reasoning"] == {"effort": "xhigh"}
+
+
+_CHAT_RESPONSE_JSON: Final = {
+    "id": "chatcmpl-wire",
+    "object": "chat.completion",
+    "created": 1,
+    "model": "gpt-5.4",
+    "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+}
+_WIRE_REQUEST: Final = {
+    "model": "gpt-5.4",
+    "messages": [
+        {"role": "system", "content": "sys"},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "what is in the image"},
+                {"type": "image_url", "image_url": {"url": "https://example.com/a.png", "detail": "low"}},
+            ],
+        },
+    ],
+    "tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}],
+    "tool_choice": "auto",
+    "response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}},
+    "max_completion_tokens": 7,
+    "reasoning_effort": "none",
+    "n": 2,
+    "seed": 3,
+}
+
+
+@pytest.mark.parametrize("is_async", (False, True))
+@respx.mock
+async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_headers_unchanged(
+    is_async: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=_CHAT_RESPONSE_JSON)
+    )
+    kwargs: Final = {
+        **_WIRE_REQUEST,
+        "model": "openai/gpt-5.4",
+        "api_base": "https://api.openai.com/v1",
+        "api_key": "sk-test",
+        "num_retries": 0,
+        "extra_body": {"vendor_flag": True},
+        "extra_headers": {"x-caller": "unit"},
+    }
+
+    if is_async:
+        await litellm.acompletion(**kwargs)
+    else:
+        litellm.completion(**kwargs)
+
+    sent: Final = route.calls.last.request
+    assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True}
+    assert sent.headers["x-caller"] == "unit"
+    assert sent.headers["authorization"] == "Bearer sk-test"