feat(proxy): tool policies - auto-discover tools + policy enforcement guardrail

litellm/constants.py

@@ -244,6 +244,7 @@
 MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
 # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
 LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
+TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60))
 # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
 # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
 MAX_SIZE_IN_MEMORY_QUEUE = int(

litellm/proxy/_types.py

@@ -183,6 +183,7 @@ class LitellmTableNames(str, enum.Enum):
     KEY_TABLE_NAME = "LiteLLM_VerificationToken"
     PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable"
     MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable"
+    TOOL_TABLE_NAME = "LiteLLM_ToolTable"
 
 
 class Litellm_EntityType(enum.Enum):
@@ -4123,6 +4124,15 @@ class SpendUpdateQueueItem(TypedDict, total=False):
     response_cost: Optional[float]
 
 
+class ToolDiscoveryQueueItem(TypedDict, total=False):
+    tool_name: str
+    origin: Optional[str]   # MCP server name or "user_defined"
+    created_by: Optional[str]
+    key_hash: Optional[str]   # hash of virtual key that triggered discovery
+    team_id: Optional[str]    # team that triggered discovery
+    key_alias: Optional[str]  # human-readable key alias
+
+
 class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase):
     unified_file_id: str
     file_object: Optional[OpenAIFileObject] = None

litellm/proxy/db/db_spend_update_writer.py

@@ -13,7 +13,17 @@
 import time
 import traceback
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Union,
+    cast,
+    overload,
+)
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -23,25 +33,29 @@
 from litellm.proxy._types import (
     DB_CONNECTION_ERROR_TYPES,
     BaseDailySpendTransaction,
-    DailyTagSpendTransaction,
+    DailyAgentSpendTransaction,
+    DailyEndUserSpendTransaction,
     DailyOrganizationSpendTransaction,
+    DailyTagSpendTransaction,
     DailyTeamSpendTransaction,
-    DailyEndUserSpendTransaction,
     DailyUserSpendTransaction,
-    DailyAgentSpendTransaction,
     DBSpendUpdateTransactions,
     Litellm_EntityType,
     LiteLLM_UserTable,
     SpendLogsMetadata,
     SpendLogsPayload,
     SpendUpdateQueueItem,
+    ToolDiscoveryQueueItem,
 )
 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.tool_discovery_queue import (
+    ToolDiscoveryQueue,
+)
 from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
 
 if TYPE_CHECKING:
@@ -67,6 +81,7 @@ def __init__(
         self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache)
         self.pod_lock_manager = PodLockManager()
         self.spend_update_queue = SpendUpdateQueue()
+        self.tool_discovery_queue = ToolDiscoveryQueue()
         self.daily_spend_update_queue = DailySpendUpdateQueue()
         self.daily_team_spend_update_queue = DailySpendUpdateQueue()
         self.daily_end_user_spend_update_queue = DailySpendUpdateQueue()
@@ -222,6 +237,13 @@ async def update_database(
                 )
             )
 
+            self._enqueue_tool_registry_upsert(
+                kwargs=kwargs,
+                completion_response=completion_response,
+                hashed_token=hashed_token,
+                team_id=team_id,
+            )
+
             verbose_proxy_logger.debug("Runs spend update on all tables")
         except Exception:
             verbose_proxy_logger.error(
@@ -237,6 +259,104 @@ async def update_database(
                 traceback.format_exc(),
             )
 
+    def _enqueue_tool_registry_upsert(
+        self,
+        kwargs: Optional[dict],
+        completion_response: Optional[Any],
+        hashed_token: Optional[str] = None,
+        team_id: Optional[str] = None,
+    ) -> None:
+        """
+        Extract tool names from the LLM request and response and enqueue them
+        for upsert into LiteLLM_ToolTable via ToolDiscoveryQueue.
+
+        Handles four sources:
+        - MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name
+        - Response tool_calls (OpenAI / Anthropic pass-through converted to OpenAI format):
+            completion_response.choices[].message.tool_calls[].function.name
+        - Request tools array (OpenAI format): kwargs["tools"][].function.name
+        - Request tools array (Anthropic /messages format): kwargs["passthrough_logging_payload"]
+            ["request_body"]["tools"][].name
+        """
+        try:
+            if kwargs is None:
+                return
+
+            # Extract key_alias from kwargs metadata if available
+            key_alias: Optional[str] = None
+            _litellm_params = kwargs.get("litellm_params") or {}
+            _metadata = _litellm_params.get("metadata") or {}
+            key_alias = _metadata.get("user_api_key_alias") or None
+
+            def _enqueue(tool_name: str, origin: str = "user_defined") -> None:
+                self.tool_discovery_queue.add_update(
+                    ToolDiscoveryQueueItem(
+                        tool_name=tool_name,
+                        origin=origin,
+                        key_hash=hashed_token,
+                        team_id=team_id,
+                        key_alias=key_alias,
+                    )
+                )
+
+            # --- MCP tool calls ---
+            sl_object = kwargs.get("standard_logging_object")
+            if sl_object is not None:
+                mcp_metadata = (
+                    sl_object.get("metadata", {}) or {}
+                ).get("mcp_tool_call_metadata")
+                if mcp_metadata and isinstance(mcp_metadata, dict):
+                    tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name")
+                    mcp_server_name = mcp_metadata.get("mcp_server_name")
+                    if tool_name:
+                        _enqueue(tool_name, origin=mcp_server_name or "user_defined")
+
+            # --- Tools from request body (OpenAI format: tools[].function.name) ---
+            request_tools = kwargs.get("tools") or []
+            for tool_def in request_tools:
+                if not isinstance(tool_def, dict):
+                    continue
+                fn = tool_def.get("function") or {}
+                name = fn.get("name") if isinstance(fn, dict) else None
+                if name:
+                    _enqueue(name)
+
+            # --- Tools from Anthropic /messages pass-through request body
+            #     (Anthropic format: tools[].name, no "function" wrapper) ---
+            passthrough_payload = kwargs.get("passthrough_logging_payload") or {}
+            request_body = (
+                passthrough_payload.get("request_body")
+                if isinstance(passthrough_payload, dict)
+                else None
+            ) or {}
+            for tool_def in request_body.get("tools") or []:
+                if not isinstance(tool_def, dict):
+                    continue
+                name = tool_def.get("name")
+                if name:
+                    _enqueue(name)
+
+            # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) ---
+            if completion_response is not None and hasattr(completion_response, "choices"):
+                for choice in completion_response.choices or []:
+                    message = getattr(choice, "message", None)
+                    if message is None:
+                        continue
+                    tool_calls = getattr(message, "tool_calls", None)
+                    if not tool_calls:
+                        continue
+                    for tc in tool_calls:
+                        fn = getattr(tc, "function", None)
+                        if fn is None:
+                            continue
+                        tool_name = getattr(fn, "name", None)
+                        if tool_name:
+                            _enqueue(tool_name)
+        except Exception as e:
+            verbose_proxy_logger.debug(
+                "_enqueue_tool_registry_upsert error (non-blocking): %s", e
+            )
+
     async def _update_key_db(
         self,
         response_cost: Optional[float],
@@ -752,6 +872,25 @@ async def _commit_spend_updates_to_db_without_redis_buffer(
             daily_spend_transactions=daily_agent_spend_update_transactions,
         )
 
+        ################## Tool Registry Upserts ##################
+        await self._flush_tool_discovery_queue(prisma_client=prisma_client)
+
+    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 = 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
+            )
+
     async def _commit_spend_updates_to_db(  # noqa: PLR0915
         self,
         prisma_client: PrismaClient,

litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py

@@ -0,0 +1,54 @@
+"""
+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.
+"""
+
+from typing import List, Set
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import ToolDiscoveryQueueItem
+
+
+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.
+    """
+
+    def __init__(self) -> None:
+        self._seen_tool_names: Set[str] = set()
+        self._pending: List[ToolDiscoveryQueueItem] = []
+
+    def add_update(self, item: ToolDiscoveryQueueItem) -> None:
+        """Enqueue a tool discovery item if tool_name has not been seen before."""
+        tool_name = 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)
+        verbose_proxy_logger.debug(
+            "ToolDiscoveryQueue: queued new tool %s (origin=%s)",
+            tool_name,
+            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 items