fix(proxy): buffer tool discoveries in Redis and write LiteLLM_ToolTable in one statement per drain, on its own lock (#48089)

* fix(proxy): buffer tool discoveries in redis and let the leader write LiteLLM_ToolTable in one statement

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): declare the redis family on the tool discovery buffer and drop the blind except and cast

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): satisfy LIT010 and LIT014 in the tool discovery buffer

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): mark the aggregation fold Final

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): requeue drained tool rows to redis when the write provably did not apply

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): drain tool discoveries on their own lock and scheduler job, not the spend commit

* chore(proxy): suppress the untyped apscheduler add_job on the tool discovery job

* refactor(proxy): move the LiteLLM_ToolTable bulk upsert into its own module next to the other bulk upserts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): name the LiteLLM_ToolTable bulk write span so the prisma span check resolves it

* refactor(proxy): move the tool discovery drain out of DBSpendUpdateWriter into tool_registry_writer

* refactor(proxy): share db_write_failure_is_requeue_safe from exception_handler so the spend writer stays untouched

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/constants.py

@@ -480,6 +480,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
@@ -1844,6 +1845,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"

litellm/proxy/_types.py

@@ -5262,6 +5262,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/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
@@ -251,16 +250,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",
@@ -298,7 +287,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(
@@ -1635,8 +1624,6 @@ class DBSpendUpdateWriter:
             window_spend_update_queue=self.window_spend_update_queue,
         )
 
-        await self._flush_tool_discovery_queue(prisma_client=prisma_client)
-
         # Only commit from redis to db if this pod is the leader
         if await self.pod_lock_manager.acquire_lock(
             cronjob_id=DB_SPEND_UPDATE_JOB_NAME,
@@ -1781,7 +1768,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 "
@@ -1947,7 +1934,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",
@@ -1965,9 +1952,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,
@@ -2099,20 +2083,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,
@@ -2124,7 +2094,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 = (
@@ -2633,7 +2602,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,6 +9,7 @@ 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
@@ -27,6 +28,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,
 )
@@ -42,12 +44,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,
@@ -135,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
@@ -559,6 +569,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.
+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.
-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.
 """
 
+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, one pending row per tool_name with a running call_count."""
-    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.
-    """
 
     def __init__(self) -> None:
-        self._seen_tool_names: set[str] = set()
+        self._pending: dict[str, ToolDiscoveryQueueItem] = {}  # mutable-ok: folded per add_update, cleared on flush
-        self._pending: list[ToolDiscoveryQueueItem] = []
 
     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:
+        self._pending[tool_name] = _merged(self._pending.get(tool_name), item)
-            verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name)
-            return
-        self._seen_tool_names.add(tool_name)
-        self._pending.append(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
+        """Return and clear all pending rows."""
-        flush cycle can re-count the same tools."""
+        items, self._pending = list(self._pending.values()), {}
-        items, self._pending = self._pending, []
-        self._seen_tool_names.clear()
         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,
 )
@@ -637,3 +638,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:
-    """
+    """Write a batch of discovered tools to LiteLLM_ToolTable in a single statement. Raises on failure so the
-    Batch-upsert tool registry rows via Prisma.
+    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")]
+        await batch_upsert_tools(prisma_client=prisma_client, items=items)
-        if not data:
+    except Exception as e:  # noqa: BLE001  # a lost tool row only delays its discovery, the scheduler job must not fail
-            return
+        verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e)
-        now: Final = datetime.now(timezone.utc)
+
-        table: Final = _tool_table_actions(prisma_client)
+
-        for item in data:
+async def commit_tool_discoveries(
-            tool_name = item.get("tool_name", "")
+    prisma_client: "PrismaClient",
-            origin = item.get("origin") or "user_defined"
+    tool_discovery_queue: ToolDiscoveryQueue,
-            created_by = item.get("created_by") or "system"
+    redis_update_buffer: RedisUpdateBuffer,
-            key_hash = item.get("key_hash")
+    pod_lock_manager: PodLockManager,
-            team_id = item.get("team_id")
+    use_redis_buffer: bool,
-            key_alias = item.get("key_alias")
+) -> None:
-            user_agent = item.get("user_agent")
+    """Scheduled alongside update_spend on its own lock, so tool table writes never ride the spend commit.
-            await table.upsert(
+
-                where={"tool_name": tool_name},
+    With the Redis buffer every worker hands its rows to Redis and the holder of the tool lock writes them all in
-                data={
+    one statement. Without it each worker writes its own rows."""
-                    "create": {
+    if not use_redis_buffer:
-                        "tool_id": str(uuid.uuid4()),
+        await flush_tool_discovery_queue(prisma_client=prisma_client, tool_discovery_queue=tool_discovery_queue)
-                        "tool_name": tool_name,
+        return
-                        "origin": origin,
+    await redis_update_buffer.store_tool_discoveries_in_redis(tool_discovery_queue=tool_discovery_queue)
-                        "input_policy": "untrusted",
+    if not await pod_lock_manager.acquire_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME):
-                        "output_policy": "untrusted",
+        return
-                        "call_count": 1,
+    try:
-                        "created_by": created_by,
+        await _commit_tool_discoveries_from_redis(prisma_client=prisma_client, redis_update_buffer=redis_update_buffer)
-                        "updated_by": created_by,
+    finally:
-                        "key_hash": key_hash,
+        await pod_lock_manager.release_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME)
-                        "team_id": team_id,
+
-                        "key_alias": key_alias,
+
-                        "user_agent": user_agent,
+async def _commit_tool_discoveries_from_redis(
-                        "last_used_at": now,
+    prisma_client: "PrismaClient", redis_update_buffer: RedisUpdateBuffer
-                    },
+) -> None:
-                    "update": {
+    """Tool lock holder only: drain every worker's tool rows from Redis and write them in one statement."""
-                        "call_count": {"increment": 1},
+    items: Final = await redis_update_buffer.get_all_tool_discoveries_from_redis_buffer()
-                        "updated_at": now,
+    if not items:
-                        "last_used_at": now,
+        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))
+            return
-    except Exception as e:
+        verbose_proxy_logger.error(
-        verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e)
+            "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/proxy_server.py

@@ -11296,6 +11296,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/utils.py

@@ -184,6 +184,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,
@@ -191,6 +192,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,
@@ -8182,6 +8184,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,

tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py

@@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-
 from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import DailyTagSpendTransaction
@@ -769,3 +768,18 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_
 
     assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False
     assert redis.rows == []
+
+
+@pytest.mark.asyncio
+async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails():
+    from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue
+
+    redis_cache = MagicMock()
+    redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down"))
+    queue = ToolDiscoveryQueue()
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+
+    await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue)
+
+    assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}]

tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py

@@ -5,9 +5,9 @@ Unit tests for ToolDiscoveryQueue.
 
 import pytest
 
-
 from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
     ToolDiscoveryQueue,
+    aggregate_tool_discoveries,
 )
 
 
@@ -16,57 +16,50 @@ def queue():
     return ToolDiscoveryQueue()
 
 
-def test_add_single_tool(queue):
+def test_a_tool_seen_once_flushes_with_call_count_one(queue):
     queue.add_update({"tool_name": "my_tool", "origin": "user_defined"})
-    items = queue.flush()
+    assert queue.flush() == [{"tool_name": "my_tool", "origin": "user_defined", "call_count": 1}]
-    assert len(items) == 1
-    assert items[0]["tool_name"] == "my_tool"
-    assert items[0]["origin"] == "user_defined"
 
 
-def test_deduplication_same_name(queue):
+def test_repeated_sightings_fold_into_one_row_counting_every_call(queue):
-    """Adding the same tool_name twice should only keep the first."""
+    queue.add_update({"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice"})
-    queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"})
+    queue.add_update({"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob"})
-    queue.add_update({"tool_name": "tool_a", "origin": "user_defined"})
+    queue.add_update({"tool_name": "tool_a"})
-    items = queue.flush()
+    assert queue.flush() == [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice", "call_count": 3}]
-    assert len(items) == 1
-    assert items[0]["origin"] == "mcp_server"  # first wins
 
 
-def test_deduplication_different_names(queue):
+def test_distinct_tools_each_get_their_own_row(queue):
     queue.add_update({"tool_name": "tool_a"})
     queue.add_update({"tool_name": "tool_b"})
-    items = queue.flush()
-    assert len(items) == 2
-    names = {i["tool_name"] for i in items}
-    assert names == {"tool_a", "tool_b"}
-
-
-def test_flush_clears_pending(queue):
-    queue.add_update({"tool_name": "tool_x"})
-    items1 = queue.flush()
-    assert len(items1) == 1
-    items2 = queue.flush()
-    assert len(items2) == 0
-
-
-def test_seen_names_reset_after_flush(queue):
-    """Seen-set is cleared on flush so the same tool can re-enter the next cycle."""
     queue.add_update({"tool_name": "tool_a"})
-    queue.flush()
+    assert queue.flush() == [
-    queue.add_update({"tool_name": "tool_a"})  # same tool, new cycle
+        {"tool_name": "tool_a", "call_count": 2},
-    items = queue.flush()
+        {"tool_name": "tool_b", "call_count": 1},
-    assert len(items) == 1
+    ]
-    assert items[0]["tool_name"] == "tool_a"
+
+
+def test_flush_clears_pending_and_the_next_cycle_counts_from_zero(queue):
+    queue.add_update({"tool_name": "tool_x"})
+    queue.add_update({"tool_name": "tool_x"})
+    assert queue.flush() == [{"tool_name": "tool_x", "call_count": 2}]
+    assert queue.flush() == []
+    queue.add_update({"tool_name": "tool_x"})
+    assert queue.flush() == [{"tool_name": "tool_x", "call_count": 1}]
 
 
 def test_empty_tool_name_ignored(queue):
     queue.add_update({"tool_name": ""})
     queue.add_update({"tool_name": None})  # type: ignore[arg-type]
-    items = queue.flush()
+    assert queue.flush() == []
-    assert len(items) == 0
 
 
-def test_flush_returns_list(queue):
+def test_aggregate_sums_call_counts_across_workers_flushes():
-    result = queue.flush()
+    worker_a = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 3}]
-    assert isinstance(result, list)
+    worker_b = [
+        {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 2},
+        {"tool_name": "search", "call_count": 1},
+    ]
+    assert aggregate_tool_discoveries([*worker_a, *worker_b]) == (
+        {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5},
+        {"tool_name": "search", "call_count": 1},
+    )

tests/unit/proxy/db/test_db_spend_update_writer.py

@@ -19,7 +19,7 @@ from redis.exceptions import DataError
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm._service_logger import ServiceTypes
-from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_UPDATE_BUFFER_KEY, SPEND_ROLLUP_LOCK_TIMEOUT_MS
+from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS
 from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem
 from litellm.proxy.db.db_spend_update_writer import (
     _TEAM_ADVISORY_LOCK_SQL,
@@ -29,7 +29,6 @@ from litellm.proxy.db.db_spend_update_writer import (
     _SpendTableName,
 )
 from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue
-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.spend_update_queue import SpendUpdateQueue
 from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
@@ -2514,31 +2513,6 @@ class _InMemoryRedis:
         return [self.lists.pop(op["key"], None) for op in lpop_list]
 
 
-@pytest.mark.asyncio
-@pytest.mark.parametrize("lock_held_by_another_pod", [False, True])
-async def test_commit_with_redis_writes_discovered_tools_on_leader_and_non_leader(lock_held_by_another_pod: bool):
-    redis = _InMemoryRedis()
-    if lock_held_by_another_pod:
-        assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=DB_SPEND_UPDATE_JOB_NAME)
-    db_writer = DBSpendUpdateWriter(redis_cache=redis)
-    db_writer.pod_lock_manager.redis_cache = redis
-    db_writer.tool_discovery_queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_tooltable.upsert = AsyncMock()
-
-    await db_writer._commit_spend_updates_to_db_with_redis(
-        prisma_client=prisma_client,
-        n_retry_times=0,
-        proxy_logging_obj=MagicMock(),
-    )
-
-    upserted = [c.kwargs["where"] for c in prisma_client.db.litellm_tooltable.upsert.await_args_list]
-    assert upserted == [{"tool_name": "get_weather"}]
-    assert (REDIS_UPDATE_BUFFER_KEY in redis.lists) is lock_held_by_another_pod, (
-        f"spend buffer left in redis={list(redis.lists)}, only the lock holder may drain it"
-    )
-
-
 @pytest.mark.asyncio
 async def test_commit_with_redis_requeues_all_on_db_failure():
     """
@@ -3384,7 +3358,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_
     await db_writer.window_spend_update_queue.add_update(transaction)
     db = _WindowSpendFakeDB()
     db.query_raw = AsyncMock(side_effect=Exception("connection reset"))
-    db_writer._flush_tool_discovery_queue = AsyncMock()
 
     await db_writer._commit_spend_updates_to_db_without_redis_buffer(
         prisma_client=_WindowSpendFakePrisma(db),
@@ -3392,7 +3365,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_
         proxy_logging_obj=MagicMock(),
     )
 
-    db_writer._flush_tool_discovery_queue.assert_called_once()
     requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
     assert requeued == (transaction,)
 
@@ -3468,7 +3440,6 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov
     db_writer = DBSpendUpdateWriter()
     await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")})
     db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure)
-    db_writer._flush_tool_discovery_queue = AsyncMock()
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.failure_handler = AsyncMock()
 
@@ -3494,7 +3465,6 @@ async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent():
         {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)}
     )
     db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply"))
-    db_writer._flush_tool_discovery_queue = AsyncMock()
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.failure_handler = AsyncMock()
 
@@ -3522,7 +3492,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other
     team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"}
     await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn})
     db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend")
-    db_writer._flush_tool_discovery_queue = AsyncMock()
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.failure_handler = AsyncMock()
 
@@ -3533,7 +3502,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other
     assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == []
     (team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend")
     assert _row_values(team_upsert, "team_id") == ["team-1"]
-    db_writer._flush_tool_discovery_queue.assert_called_once()
 
     db.failing_table = None
     await db_writer._commit_spend_updates_to_db_without_redis_buffer(
@@ -3760,7 +3728,6 @@ async def test_failed_window_spend_commit_is_requeued_only_when_the_rows_are_pro
     )
     await db_writer.window_spend_update_queue.add_update(transaction)
     db = _WindowSpendFailureDB(failure)
-    db_writer._flush_tool_discovery_queue = AsyncMock()
 
     await db_writer._commit_spend_updates_to_db_without_redis_buffer(
         prisma_client=_WindowSpendFakePrisma(db),

tests/unit/proxy/db/test_tool_registry_writer.py

@@ -3,22 +3,28 @@ Unit tests for tool_registry_writer.py — uses a mock prisma client
 that exposes litellm_tooltable.upsert / find_many / find_unique.
 """
 
+import json
 from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-from prisma.errors import PrismaError
+from prisma.errors import PrismaError, RawQueryError
-
 
+from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_TOOL_DISCOVERY_BUFFER_KEY, TOOL_DISCOVERY_COMMIT_JOB_NAME
+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
 from litellm.proxy.db.tool_registry_writer import (
     ToolPolicyRegistry,
     batch_upsert_tools,
+    commit_tool_discoveries,
     get_tool,
     get_tool_policy_registry,
     get_tools_by_names,
     list_tools,
     update_tool_policy,
 )
+from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
 
 
 def _mock_row(**kwargs):
@@ -62,6 +68,7 @@ def _make_prisma(
     prisma = MagicMock()
     prisma.db.litellm_tooltable = MagicMock()
     prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return)
+    prisma.db.execute_raw = AsyncMock(return_value=1)
     prisma.db.litellm_tooltable.find_many = AsyncMock(
         return_value=find_many_rows if find_many_rows is not None else []
     )
@@ -73,28 +80,32 @@ def _make_prisma(
     return prisma
 
 
+def _statement(prisma) -> tuple[str, tuple[object, ...]]:
+    prisma.db.execute_raw.assert_awaited_once()
+    sql, *params = prisma.db.execute_raw.call_args.args
+    return sql, tuple(params)
+
+
 @pytest.mark.asyncio
-async def test_batch_upsert_tools_calls_upsert():
+async def test_batch_upsert_tools_inserts_a_new_tool_with_the_batch_call_count_and_adds_it_on_conflict():
     prisma = _make_prisma()
-    items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}]
+    await batch_upsert_tools(
-    await batch_upsert_tools(prisma, items)
+        prisma, [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}]
-    prisma.db.litellm_tooltable.upsert.assert_awaited_once()
+    )
-    call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs
+    sql, params = _statement(prisma)
-    assert call_kw["where"] == {"tool_name": "tool_a"}
+    assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"')
-    assert call_kw["data"]["create"]["tool_name"] == "tool_a"
+    assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql
-    assert call_kw["data"]["create"]["origin"] == "mcp_server"
+    assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql
-    assert call_kw["data"]["create"]["input_policy"] == "untrusted"
+    assert '"input_policy"' not in sql, "policies come from the table defaults and are never overwritten"
-    assert call_kw["data"]["create"]["output_policy"] == "untrusted"
+    assert params[1:] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None)
-    assert call_kw["data"]["create"]["call_count"] == 1
+    prisma.db.litellm_tooltable.upsert.assert_not_awaited()
-    assert call_kw["data"]["update"]["call_count"] == {"increment": 1}
-    assert "updated_at" in call_kw["data"]["update"]
 
 
 @pytest.mark.asyncio
 async def test_batch_upsert_tools_empty_list():
     prisma = _make_prisma()
     await batch_upsert_tools(prisma, [])
-    prisma.db.litellm_tooltable.upsert.assert_not_awaited()
+    prisma.db.execute_raw.assert_not_awaited()
 
 
 @pytest.mark.asyncio
@@ -102,21 +113,22 @@ async def test_batch_upsert_tools_skips_empty_names():
     prisma = _make_prisma()
     items = [{"tool_name": "", "origin": None}, {"tool_name": None}]  # type: ignore[list-item]
     await batch_upsert_tools(prisma, items)
-    prisma.db.litellm_tooltable.upsert.assert_not_awaited()
+    prisma.db.execute_raw.assert_not_awaited()
 
 
 @pytest.mark.asyncio
-async def test_batch_upsert_multiple_tools_calls_upsert_per_tool():
+async def test_batch_upsert_many_tools_is_one_statement_with_one_row_per_distinct_tool():
     prisma = _make_prisma()
     items = [
-        {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None},
+        {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 2},
-        {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"},
+        {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice", "call_count": 1},
+        {"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob", "call_count": 5},
     ]
     await batch_upsert_tools(prisma, items)
-    assert prisma.db.litellm_tooltable.upsert.await_count == 2
+    sql, params = _statement(prisma)
-    calls = prisma.db.litellm_tooltable.upsert.call_args_list
+    assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql
-    assert calls[0].kwargs["where"]["tool_name"] == "tool_a"
+    assert params[1:10] == ("tool_a", "mcp_server", 7, "system", "system", None, None, None, None)
-    assert calls[1].kwargs["where"]["tool_name"] == "tool_b"
+    assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", None, None, None, None)
 
 
 @pytest.mark.asyncio
@@ -264,20 +276,14 @@ async def test_update_tool_policy_calls_upsert_then_get_tool():
     assert call_kw["where"] == {"tool_name": "my_tool"}
     assert call_kw["data"]["update"]["input_policy"] == "blocked"
     assert call_kw["data"]["update"]["updated_by"] == "admin"
-    prisma.db.litellm_tooltable.find_unique.assert_awaited_with(
+    prisma.db.litellm_tooltable.find_unique.assert_awaited_with(where={"tool_name": "my_tool"})
-        where={"tool_name": "my_tool"}
-    )
 
 
 @pytest.mark.asyncio
 async def test_get_tools_by_names_returns_policy_map():
     rows = [
-        _mock_row(
+        _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"),
-            tool_name="tool_a", input_policy="trusted", output_policy="untrusted"
+        _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"),
-        ),
-        _mock_row(
-            tool_name="tool_b", input_policy="blocked", output_policy="untrusted"
-        ),
     ]
     prisma = _make_prisma(find_many_rows=rows)
     result = await get_tools_by_names(prisma, ["tool_a", "tool_b"])
@@ -285,9 +291,7 @@ async def test_get_tools_by_names_returns_policy_map():
         "tool_a": ("trusted", "untrusted"),
         "tool_b": ("blocked", "untrusted"),
     }
-    prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(
+    prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(where={"tool_name": {"in": ["tool_a", "tool_b"]}})
-        where={"tool_name": {"in": ["tool_a", "tool_b"]}}
-    )
 
 
 @pytest.mark.asyncio
@@ -382,12 +386,8 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read():
         return []
 
     mock_prisma_client = MagicMock()
-    mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(
+    mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(side_effect=_flaky_find_many)
-        side_effect=_flaky_find_many
+    mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(return_value=[])
-    )
-    mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(
-        return_value=[]
-    )
     mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
     mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0
     mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1
@@ -397,10 +397,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read():
     assert len(invocations) == 2
     mock_prisma_client.attempt_db_reconnect.assert_awaited_once()
     reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs
-    assert (
+    assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_tools_lookup_failure"
-        reconnect_kwargs["reason"]
-        == "sync_tool_policy_from_db_tools_lookup_failure"
-    )
     assert registry.is_initialized()
 
 
@@ -421,9 +418,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read()
 
     mock_prisma_client = MagicMock()
     mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[])
-    mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(
+    mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(side_effect=_flaky_perms_find_many)
-        side_effect=_flaky_perms_find_many
-    )
     mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
     mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0
     mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1
@@ -433,7 +428,142 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read()
     assert len(perms_invocations) == 2
     mock_prisma_client.attempt_db_reconnect.assert_awaited_once()
     reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs
-    assert (
+    assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_perms_lookup_failure"
-        reconnect_kwargs["reason"]
+
-        == "sync_tool_policy_from_db_perms_lookup_failure"
+
+class _InMemoryRedis:
+    def __init__(self) -> None:
+        self.store: dict[str, str] = {}
+        self.lists: dict[str, list[object]] = {}
+
+    async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None:
+        if nx and key in self.store:
+            return None
+        self.store[key] = json.dumps(value)
+        return True
+
+    async def async_get_cache(self, key: str, **kwargs: object) -> str | None:
+        raw = self.store.get(key)
+        return json.loads(raw) if raw is not None else None
+
+    async def async_delete_cache(self, key: str) -> int:
+        return 1 if self.store.pop(key, None) is not None else 0
+
+    async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]:
+        for op in rpush_list:
+            self.lists.setdefault(op["key"], []).extend(op["values"])
+        return [len(self.lists[op["key"]]) for op in rpush_list]
+
+    async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]:
+        return [self.lists.pop(op["key"], None) for op in lpop_list]
+
+
+def _postgres_rejection(sqlstate: str) -> RawQueryError:
+    return RawQueryError(
+        data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}}
     )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "lock_held_by_another_pod",
+    [None, DB_SPEND_UPDATE_JOB_NAME, TOOL_DISCOVERY_COMMIT_JOB_NAME],
+    ids=["free", "spend_lock_held", "tool_lock_held"],
+)
+async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_from_the_tool_lock_holder_only(
+    lock_held_by_another_pod: str | None,
+):
+    redis = _InMemoryRedis()
+    if lock_held_by_another_pod is not None:
+        assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=lock_held_by_another_pod)
+    other_worker = ToolDiscoveryQueue()
+    other_worker.add_update({"tool_name": "get_weather", "origin": "mcp_server"})
+    other_worker.add_update({"tool_name": "search", "origin": "user_defined"})
+    await RedisUpdateBuffer(redis_cache=redis).store_tool_discoveries_in_redis(other_worker)
+
+    queue = ToolDiscoveryQueue()
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+    prisma_client = MagicMock()
+    prisma_client.db.execute_raw = AsyncMock(return_value=2)
+    prisma_client.db.litellm_tooltable.upsert = AsyncMock()
+
+    await commit_tool_discoveries(
+        prisma_client=prisma_client,
+        tool_discovery_queue=queue,
+        redis_update_buffer=RedisUpdateBuffer(redis_cache=redis),
+        pod_lock_manager=PodLockManager(redis_cache=redis),
+        use_redis_buffer=True,
+    )
+
+    assert queue.flush() == [], "every worker hands its tools to redis each tick"
+    prisma_client.db.litellm_tooltable.upsert.assert_not_awaited()
+    if lock_held_by_another_pod == TOOL_DISCOVERY_COMMIT_JOB_NAME:
+        prisma_client.db.execute_raw.assert_not_awaited()
+        assert len(redis.lists[REDIS_TOOL_DISCOVERY_BUFFER_KEY]) == 2, "non-holders leave their tools for the holder"
+        return
+    assert REDIS_TOOL_DISCOVERY_BUFFER_KEY not in redis.lists, "the spend lock must not gate the tool drain"
+    prisma_client.db.execute_raw.assert_awaited_once()
+    sql, *params = prisma_client.db.execute_raw.call_args.args
+    assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') and 'ON CONFLICT ("tool_name")' in sql
+    rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)}
+    assert rows == {"get_weather": 3, "search": 1}, f"one row per tool summed across workers, got {params}"
+
+
+@pytest.mark.asyncio
+async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statement_and_never_touches_redis():
+    redis = _InMemoryRedis()
+    queue = ToolDiscoveryQueue()
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+    queue.add_update({"tool_name": "search", "origin": "user_defined"})
+    prisma_client = MagicMock()
+    prisma_client.db.execute_raw = AsyncMock(return_value=2)
+
+    await commit_tool_discoveries(
+        prisma_client=prisma_client,
+        tool_discovery_queue=queue,
+        redis_update_buffer=RedisUpdateBuffer(redis_cache=redis),
+        pod_lock_manager=PodLockManager(redis_cache=redis),
+        use_redis_buffer=False,
+    )
+
+    assert redis.lists == {}, "without the buffer flag nothing goes through redis"
+    assert queue.flush() == []
+    prisma_client.db.execute_raw.assert_awaited_once()
+    sql, *params = prisma_client.db.execute_raw.call_args.args
+    assert 'ON CONFLICT ("tool_name")' in sql
+    rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)}
+    assert rows == {"get_weather": 2, "search": 1}, f"got {params}"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("failure", "requeued"),
+    [
+        (Exception("connection reset"), True),
+        (_postgres_rejection("23502"), False),
+    ],
+    ids=["connection_error_goes_back_to_redis", "rejected_row_is_dropped"],
+)
+async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably_did_not_apply(
+    failure: Exception, requeued: bool
+):
+    redis = _InMemoryRedis()
+    queue = ToolDiscoveryQueue()
+    queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
+    prisma_client = MagicMock()
+    prisma_client.db.execute_raw = AsyncMock(side_effect=failure)
+
+    await commit_tool_discoveries(
+        prisma_client=prisma_client,
+        tool_discovery_queue=queue,
+        redis_update_buffer=RedisUpdateBuffer(redis_cache=redis),
+        pod_lock_manager=PodLockManager(redis_cache=redis),
+        use_redis_buffer=True,
+    )
+
+    prisma_client.db.execute_raw.assert_awaited_once()
+    left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])]
+    expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else []
+    assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}"

tests/unit/proxy/db/test_tool_table_bulk_upsert.py

@@ -0,0 +1,38 @@
+"""Tests for the single-statement LiteLLM_ToolTable upsert (LIT-9997)."""
+
+from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert
+
+
+def test_one_row_per_item_and_the_conflict_update_only_moves_the_counters():
+    sql, params = build_tool_bulk_upsert(
+        [
+            {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4},
+            {
+                "tool_name": "tool_b",
+                "origin": "user_defined",
+                "created_by": "alice",
+                "call_count": 1,
+                "key_hash": "hash",
+                "team_id": "team",
+                "key_alias": "alias",
+                "user_agent": "agent",
+            },
+        ]
+    )
+    assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"')
+    assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql
+    assert "$4::integer" in sql and "$14::integer" in sql
+    assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql
+    assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql
+    assert '"last_used_at" = EXCLUDED."last_used_at"' in sql
+    assert '"input_policy"' not in sql and '"output_policy"' not in sql, "policies stay on the table defaults"
+    assert '"origin" =' not in sql and '"key_hash" =' not in sql, "an existing row keeps who discovered it"
+    assert len(params) == 20
+    assert params[1:10] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None)
+    assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", "hash", "team", "alias", "agent")
+    assert params[0] != params[10], "each row gets its own tool_id"
+
+
+def test_defaults_fill_origin_created_by_and_call_count():
+    _, params = build_tool_bulk_upsert([{"tool_name": "tool_a"}])
+    assert params[1:6] == ("tool_a", "user_defined", 1, "system", "system")