From 83c7677d8e775e9ac201ab7bd6c578420499dfa0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 15:03:54 -0800 Subject: [PATCH 1/7] feat(proxy): tool policies - auto-discover tools, manage policies, guardrail enforcement - New LiteLLM_ToolTable in schema.prisma to store discovered tools - Auto-discovery: tools seen in LLM responses get upserted via ToolDiscoveryQueue (hooks into DBSpendUpdateWriter, same pipeline as spend tracking) - Management endpoints: GET /v1/tool/list, GET /v1/tool/{name}, POST /v1/tool/policy - ToolPolicyGuardrail: blocks tool_calls in responses based on policy setting - UI: Tool Policies page under Guardrails section with policy selector, filters by policy/team/key, live tail, sortable table - Unit tests for queue, writer, endpoints, guardrail --- litellm/constants.py | 1 + litellm/proxy/_types.py | 10 + litellm/proxy/db/db_spend_update_writer.py | 123 +++++- .../tool_discovery_queue.py | 49 +++ litellm/proxy/db/tool_registry_writer.py | 173 ++++++++ .../guardrail_hooks/tool_policy/__init__.py | 16 + .../tool_policy/tool_policy_guardrail.py | 153 +++++++ .../tool_management_endpoints.py | 149 +++++++ litellm/proxy/proxy_server.py | 4 + litellm/proxy/schema.prisma | 20 + litellm/types/tool_management.py | 41 ++ .../test_tool_discovery_queue.py | 74 ++++ .../proxy/db/test_tool_registry_writer.py | 159 +++++++ .../test_tool_policy_guardrail.py | 181 ++++++++ .../test_tool_management_endpoints.py | 151 +++++++ ui/litellm-dashboard/src/app/page.tsx | 3 + .../src/components/ToolPolicies.tsx | 410 ++++++++++++++++++ .../src/components/leftnav.tsx | 6 + .../src/components/networking.tsx | 50 +++ 19 files changed, 1769 insertions(+), 4 deletions(-) create mode 100644 litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py create mode 100644 litellm/proxy/db/tool_registry_writer.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py create mode 100644 litellm/proxy/management_endpoints/tool_management_endpoints.py create mode 100644 litellm/types/tool_management.py create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py create mode 100644 tests/test_litellm/proxy/db/test_tool_registry_writer.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py create mode 100644 ui/litellm-dashboard/src/components/ToolPolicies.tsx diff --git a/litellm/constants.py b/litellm/constants.py index ee79f2fa56fd..b1a0021bcc67 100644 --- a/litellm/constants.py +++ b/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( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f354e28acd76..75b9f91acd9c 100644 --- a/litellm/proxy/_types.py +++ b/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 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d7d6b2b1eb09..d3c9b14c2ec7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/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,18 +33,19 @@ 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, @@ -42,6 +53,9 @@ 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,80 @@ 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 response and enqueue them for upsert + into LiteLLM_ToolTable via ToolDiscoveryQueue. + + Handles two sources: + - MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name + - Regular function calls: completion_response.choices[].message.tool_calls[].function.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 + + # --- 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: + self.tool_discovery_queue.add_update( + ToolDiscoveryQueueItem( + tool_name=tool_name, + origin=mcp_server_name or "user_defined", + key_hash=hashed_token, + team_id=team_id, + key_alias=key_alias, + ) + ) + + # --- Regular function/tool calls in LLM response --- + 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: + self.tool_discovery_queue.add_update( + ToolDiscoveryQueueItem( + tool_name=tool_name, + origin="user_defined", + key_hash=hashed_token, + team_id=team_id, + key_alias=key_alias, + ) + ) + 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 +848,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, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py new file mode 100644 index 000000000000..6c04a10c0a46 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -0,0 +1,49 @@ +""" +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 pod +lifetime, so DB upserts stop entirely after warmup. +""" + +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 — once a tool has been seen in this process, + it is never enqueued again (the DB row already exists or will be created + during the current flush cycle). + """ + + 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.""" + items, self._pending = self._pending, [] + return items diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py new file mode 100644 index 000000000000..8aedf65a6dd4 --- /dev/null +++ b/litellm/proxy/db/tool_registry_writer.py @@ -0,0 +1,173 @@ +""" +DB helpers for LiteLLM_ToolTable — the global tool registry. + +Tools are auto-discovered from LLM responses and upserted here. +Admins use the management endpoints to read and update call_policy. + +NOTE: Uses raw SQL (query_raw / execute_raw) instead of Prisma model methods +because the generated Prisma Python client may not have LiteLLM_ToolTable +when running against an older generated schema. +""" + +from typing import TYPE_CHECKING, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +def _row_to_model(row: dict) -> LiteLLM_ToolTableRow: + return LiteLLM_ToolTableRow( + tool_id=row.get("tool_id", ""), + tool_name=row.get("tool_name", ""), + origin=row.get("origin"), + call_policy=row.get("call_policy", "untrusted"), + assignments=row.get("assignments"), + key_hash=row.get("key_hash"), + team_id=row.get("team_id"), + key_alias=row.get("key_alias"), + created_at=row.get("created_at"), + updated_at=row.get("updated_at"), + created_by=row.get("created_by"), + updated_by=row.get("updated_by"), + ) + + +async def batch_upsert_tools( + prisma_client: "PrismaClient", + items: List[ToolDiscoveryQueueItem], +) -> None: + """ + Batch-upsert tool registry rows via raw SQL. + + Uses INSERT ON CONFLICT DO NOTHING so that: + - First insert sets call_policy to "untrusted" (the schema default). + - Subsequent upserts are no-ops (existing policy is preserved). + """ + if not items: + return + try: + data = [ + item + for item in items + if item.get("tool_name") + ] + if not data: + return + 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") + await prisma_client.db.execute_raw( + 'INSERT INTO "LiteLLM_ToolTable" ' + "(tool_id, tool_name, origin, call_policy, created_by, updated_by, key_hash, team_id, key_alias) " + "VALUES (gen_random_uuid()::text, $1, $2, 'untrusted', $3, $3, $4, $5, $6) " + "ON CONFLICT (tool_name) DO NOTHING", + tool_name, + origin, + created_by, + key_hash, + team_id, + key_alias, + ) + 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) + + +async def list_tools( + prisma_client: "PrismaClient", + call_policy: Optional[ToolCallPolicy] = None, +) -> List[LiteLLM_ToolTableRow]: + """Return all tools, optionally filtered by call_policy.""" + try: + if call_policy is not None: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC', + call_policy, + ) + else: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC', + ) + return [_row_to_model(row) for row in rows] + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer list_tools error: %s", e) + return [] + + +async def get_tool( + prisma_client: "PrismaClient", + tool_name: str, +) -> Optional[LiteLLM_ToolTableRow]: + """Return a single tool row by tool_name.""" + try: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" WHERE tool_name = $1', + tool_name, + ) + if not rows: + return None + return _row_to_model(rows[0]) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e) + return None + + +async def update_tool_policy( + prisma_client: "PrismaClient", + tool_name: str, + call_policy: ToolCallPolicy, + updated_by: Optional[str], +) -> Optional[LiteLLM_ToolTableRow]: + """Update the call_policy for a tool. Upserts the row if it does not exist yet.""" + try: + _updated_by = updated_by or "system" + await prisma_client.db.execute_raw( + 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by) ' + "VALUES (gen_random_uuid()::text, $1, $2, $3, $3) " + "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = NOW()", + tool_name, + call_policy, + _updated_by, + ) + return await get_tool(prisma_client, tool_name) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e) + return None + + +async def get_tools_by_names( + prisma_client: "PrismaClient", + tool_names: List[str], +) -> Dict[str, str]: + """ + Return a {tool_name: call_policy} map for the given tool names. + Used by the policy enforcement guardrail — single batch query, never N+1. + """ + if not tool_names: + return {} + try: + placeholders = ", ".join(f"${i+1}" for i in range(len(tool_names))) + rows = await prisma_client.db.query_raw( + f'SELECT tool_name, call_policy FROM "LiteLLM_ToolTable" WHERE tool_name IN ({placeholders})', + *tool_names, + ) + return {row["tool_name"]: row["call_policy"] for row in rows} + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e) + return {} diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py new file mode 100644 index 000000000000..5a43006e23ce --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py @@ -0,0 +1,16 @@ +import litellm +from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail): + from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, + ) + + _callback = ToolPolicyGuardrail( + guardrail_name=guardrail.get("guardrail_name", "tool_policy"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py new file mode 100644 index 000000000000..ec3e6ede1ae9 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -0,0 +1,153 @@ +""" +Tool Policy Guardrail + +Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses. + +Policy values: + "trusted" - allow through (no action) + "untrusted" - allow through (no action; default for newly discovered tools) + "blocked" - raise HTTPException, preventing the tool call + "dual_llm" - (Phase 3) send to second LLM for verification; currently treated as allowed + +Configuration in proxy config YAML: + guardrails: + - guardrail_name: "tool_policy" + litellm_params: + guardrail: tool_policy + mode: post_call + +or both pre and post call: + - guardrail_name: "tool_policy" + litellm_params: + guardrail: tool_policy + mode: during_call # runs before LLM and on response +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.constants import TOOL_POLICY_CACHE_TTL_SECONDS +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +GUARDRAIL_NAME = "tool_policy" + + +class ToolPolicyGuardrail(CustomGuardrail): + """ + Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable. + + Tools with call_policy="blocked" are rejected before/after the LLM call. + Tools with call_policy="trusted" or "untrusted" pass through unchanged. + """ + + def __init__(self, **kwargs: Any) -> None: + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ] + super().__init__(**kwargs) + self._policy_cache: DualCache = DualCache() + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Enforce tool policies on both request tools and response tool_calls. + + - input_type="request": check inputs["tools"] (tool definitions in the LLM request) + - input_type="response": check inputs["tool_calls"] (tool_calls in the LLM response) + + Raises HTTPException (400) if any tool is "blocked". + """ + if input_type == "request": + tools = inputs.get("tools") or [] + tool_names = [ + t["function"]["name"] + for t in tools + if isinstance(t, dict) + and t.get("type") == "function" + and isinstance(t.get("function"), dict) + and t["function"].get("name") + ] + else: # response + tool_calls = inputs.get("tool_calls") or [] + tool_names = [] + for tc in tool_calls: + fn = None + if isinstance(tc, dict): + fn = (tc.get("function") or {}).get("name") + elif hasattr(tc, "function"): + fn = getattr(tc.function, "name", None) + if fn: + tool_names.append(fn) + + if not tool_names: + return inputs + + policy_map = await self._get_policies_cached(tool_names) + + blocked = [name for name in tool_names if policy_map.get(name) == "blocked"] + if blocked: + verbose_proxy_logger.warning( + "ToolPolicyGuardrail: blocking tool(s) %s (policy=blocked)", blocked + ) + raise HTTPException( + status_code=400, + detail={ + "error": "Violated tool policy", + "blocked_tools": blocked, + "message": f"Tool(s) {blocked} are blocked by policy.", + }, + ) + + return inputs + + async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: + """ + Batch-fetch call_policy for the given tool names. + Uses DualCache (in-mem + optional Redis) with TTL to avoid N+1 DB reads. + """ + from litellm.proxy.db.tool_registry_writer import get_tools_by_names + from litellm.proxy.proxy_server import prisma_client + + if not tool_names or prisma_client is None: + return {} + + cache_key = f"tool_policies:{':'.join(sorted(tool_names))}" + cached = await self._policy_cache.async_get_cache(cache_key) + if cached is not None and isinstance(cached, dict): + verbose_proxy_logger.debug( + "ToolPolicyGuardrail: cache hit for tools %s", tool_names + ) + return cached + + policy_map = await get_tools_by_names( + prisma_client=prisma_client, tool_names=tool_names + ) + await self._policy_cache.async_set_cache( + key=cache_key, + value=policy_map, + ttl=TOOL_POLICY_CACHE_TTL_SECONDS, + ) + verbose_proxy_logger.debug( + "ToolPolicyGuardrail: fetched policies %s", policy_map + ) + return policy_map diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py new file mode 100644 index 000000000000..89880c9a4ecc --- /dev/null +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +TOOL POLICY MANAGEMENT + +All /tool management endpoints + +GET /v1/tool/list - List all discovered tools and their policies +GET /v1/tool/{tool_name} - Get a single tool's details +POST /v1/tool/policy - Update the call_policy for a tool +""" + +from typing import Optional + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.tool_management import ( + LiteLLM_ToolTableRow, + ToolCallPolicy, + ToolListResponse, + ToolPolicyUpdateRequest, + ToolPolicyUpdateResponse, +) + +router = APIRouter() + + +@router.get( + "/v1/tool/list", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolListResponse, +) +async def list_tools( + call_policy: Optional[ToolCallPolicy] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List all auto-discovered tools and their call policies. + + Parameters: + - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked" + """ + from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy) + return ToolListResponse(tools=tools, total=len(tools)) + except Exception as e: + verbose_proxy_logger.exception("Error listing tools: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/tool/{tool_name:path}", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ToolTableRow, +) +async def get_tool( + tool_name: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get details for a single tool. + + Parameters: + - tool_name: The tool name (supports namespaced names with slashes) + """ + from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) + if tool is None: + raise HTTPException( + status_code=404, detail=f"Tool '{tool_name}' not found" + ) + return tool + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting tool: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/tool/policy", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolPolicyUpdateResponse, +) +async def update_tool_policy( + data: ToolPolicyUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Set the call policy for a tool. + + Parameters: + - tool_name: str - The tool to update + - call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked" + + Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove + that tool_call from LLM responses before returning them to the client. + """ + from litellm.proxy.db.tool_registry_writer import ( + update_tool_policy as db_update_tool_policy, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + updated = await db_update_tool_policy( + prisma_client=prisma_client, + tool_name=data.tool_name, + call_policy=data.call_policy, + updated_by=user_api_key_dict.user_id, + ) + if updated is None: + raise HTTPException( + status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'" + ) + return ToolPolicyUpdateResponse( + tool_name=updated.tool_name, + call_policy=updated.call_policy, + updated=True, + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error updating tool policy: %s", e) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1983a601e13b..d8bfaaf22b80 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -409,6 +409,9 @@ def generate_feedback_box(): update_team, validate_membership, ) +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + router as tool_management_router, +) from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -12881,6 +12884,7 @@ async def get_routes(): app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) +app.include_router(tool_management_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 50c0a55a8751..2bfd10b01581 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1051,6 +1051,26 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } +// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +model LiteLLM_ToolTable { + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? + + @@index([tool_name]) + @@index([call_policy]) + @@index([team_id]) +} + //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py new file mode 100644 index 000000000000..b0bf0fc1ebe3 --- /dev/null +++ b/litellm/types/tool_management.py @@ -0,0 +1,41 @@ +""" +Pydantic models for Tool Policy management endpoints. +""" + +from datetime import datetime +from typing import Dict, List, Literal, Optional + +from pydantic import BaseModel + +ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"] + + +class LiteLLM_ToolTableRow(BaseModel): + tool_id: str + tool_name: str + origin: Optional[str] = None + call_policy: ToolCallPolicy = "untrusted" + assignments: Optional[Dict] = None + key_hash: Optional[str] = None + team_id: Optional[str] = None + key_alias: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + + +class ToolListResponse(BaseModel): + tools: List[LiteLLM_ToolTableRow] + total: int + + +class ToolPolicyUpdateRequest(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + + +class ToolPolicyUpdateResponse(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + updated: bool diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py new file mode 100644 index 000000000000..ae19845d3e7e --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -0,0 +1,74 @@ +""" +Unit tests for ToolDiscoveryQueue. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) + + +@pytest.fixture +def queue(): + return ToolDiscoveryQueue() + + +def test_add_single_tool(queue): + queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) + items = queue.flush() + assert len(items) == 1 + assert items[0]["tool_name"] == "my_tool" + assert items[0]["origin"] == "user_defined" + + +def test_deduplication_same_name(queue): + """Adding the same tool_name twice should only keep the first.""" + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) + items = queue.flush() + assert len(items) == 1 + assert items[0]["origin"] == "mcp_server" # first wins + + +def test_deduplication_different_names(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_persist_across_flushes(queue): + """Process-local dedup should prevent re-queuing even after a flush.""" + queue.add_update({"tool_name": "tool_a"}) + queue.flush() + queue.add_update({"tool_name": "tool_a"}) # already seen + items = queue.flush() + assert len(items) == 0 + + +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 len(items) == 0 + + +def test_flush_returns_list(queue): + result = queue.flush() + assert isinstance(result, list) diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py new file mode 100644 index 000000000000..b975aa8e852e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -0,0 +1,159 @@ +""" +Unit tests for tool_registry_writer.py — tests use a mock prisma client +to avoid requiring a real DB connection. +""" + +import os +import sys +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.tool_registry_writer import ( + batch_upsert_tools, + get_tool, + get_tools_by_names, + list_tools, + update_tool_policy, +) + + +def _make_prisma(rows=None): + """Return a minimal mock prisma_client.""" + now = datetime.now(timezone.utc) + default_row = MagicMock( + tool_id="uuid-1", + tool_name="my_tool", + origin="user_defined", + call_policy="untrusted", + assignments={}, + created_at=now, + updated_at=now, + created_by=None, + updated_by=None, + ) + table = MagicMock() + table.create_many = AsyncMock(return_value=None) + table.find_many = AsyncMock(return_value=rows if rows is not None else [default_row]) + table.find_unique = AsyncMock(return_value=default_row) + table.upsert = AsyncMock(return_value=default_row) + + prisma = MagicMock() + prisma.db.litellm_tooltable = table + return prisma + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_calls_create_many(): + prisma = _make_prisma() + items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] + await batch_upsert_tools(prisma, items) + prisma.db.litellm_tooltable.create_many.assert_awaited_once() + call_kwargs = prisma.db.litellm_tooltable.create_many.call_args + assert call_kwargs.kwargs["skip_duplicates"] is True + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_empty_list(): + prisma = _make_prisma() + await batch_upsert_tools(prisma, []) + prisma.db.litellm_tooltable.create_many.assert_not_awaited() + + +@pytest.mark.asyncio +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.create_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_list_tools_no_filter(): + now = datetime.now(timezone.utc) + row = MagicMock( + tool_id="id1", + tool_name="tool_a", + origin="mcp", + call_policy="untrusted", + assignments={}, + created_at=now, + updated_at=now, + created_by=None, + updated_by=None, + ) + prisma = _make_prisma(rows=[row]) + result = await list_tools(prisma) + assert len(result) == 1 + assert result[0].tool_name == "tool_a" + + +@pytest.mark.asyncio +async def test_list_tools_with_policy_filter(): + now = datetime.now(timezone.utc) + row = MagicMock( + tool_id="id1", + tool_name="blocked_tool", + origin=None, + call_policy="blocked", + assignments=None, + created_at=now, + updated_at=now, + created_by=None, + updated_by=None, + ) + prisma = _make_prisma(rows=[row]) + result = await list_tools(prisma, call_policy="blocked") + assert result[0].call_policy == "blocked" + call_kwargs = prisma.db.litellm_tooltable.find_many.call_args + assert call_kwargs.kwargs["where"] == {"call_policy": "blocked"} + + +@pytest.mark.asyncio +async def test_get_tool_found(): + prisma = _make_prisma() + result = await get_tool(prisma, "my_tool") + assert result is not None + assert result.tool_name == "my_tool" + + +@pytest.mark.asyncio +async def test_get_tool_not_found(): + prisma = _make_prisma() + prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=None) + result = await get_tool(prisma, "nonexistent") + assert result is None + + +@pytest.mark.asyncio +async def test_update_tool_policy_upsert(): + prisma = _make_prisma() + result = await update_tool_policy(prisma, "my_tool", "blocked", "admin") + assert result is not None + prisma.db.litellm_tooltable.upsert.assert_awaited_once() + upsert_call = prisma.db.litellm_tooltable.upsert.call_args + assert upsert_call.kwargs["data"]["update"]["call_policy"] == "blocked" + assert upsert_call.kwargs["data"]["update"]["updated_by"] == "admin" + + +@pytest.mark.asyncio +async def test_get_tools_by_names_returns_policy_map(): + now = datetime.now(timezone.utc) + rows = [ + MagicMock(tool_name="tool_a", call_policy="trusted"), + MagicMock(tool_name="tool_b", call_policy="blocked"), + ] + prisma = _make_prisma(rows=rows) + result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) + assert result == {"tool_a": "trusted", "tool_b": "blocked"} + + +@pytest.mark.asyncio +async def test_get_tools_by_names_empty_list(): + prisma = _make_prisma() + result = await get_tools_by_names(prisma, []) + assert result == {} + prisma.db.litellm_tooltable.find_many.assert_not_awaited() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py new file mode 100644 index 000000000000..c6a81efbf0b5 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -0,0 +1,181 @@ +""" +Unit tests for ToolPolicyGuardrail. +""" + +import os +import sys +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, +) +from litellm.types.guardrails import GuardrailEventHooks + + +@pytest.fixture +def guardrail(): + return ToolPolicyGuardrail() + + +# --- helpers --- + +def _tool_request_inputs(tool_names: list) -> dict: + return { + "tools": [ + {"type": "function", "function": {"name": name, "description": ""}} + for name in tool_names + ] + } + + +def _tool_response_inputs(tool_names: list) -> dict: + return { + "tool_calls": [ + {"type": "function", "function": {"name": name}} + for name in tool_names + ] + } + + +# --- tests --- + + +def test_guardrail_supports_pre_and_post_call(guardrail): + hooks = guardrail.supported_event_hooks + assert GuardrailEventHooks.pre_call in hooks + assert GuardrailEventHooks.post_call in hooks + + +@pytest.mark.asyncio +async def test_no_tools_in_request_passes_through(guardrail): + inputs: Any = {"tools": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_no_tool_calls_in_response_passes_through(guardrail): + inputs: Any = {"tool_calls": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_untrusted_tools_pass_through(guardrail): + policy_map = {"search": "untrusted", "read_file": "trusted"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["search", "read_file"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_blocked_tool_in_request_raises_http_exception(guardrail): + policy_map = {"dangerous_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["dangerous_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert exc_info.value.status_code == 400 + assert "dangerous_tool" in exc_info.value.detail["blocked_tools"] + + +@pytest.mark.asyncio +async def test_blocked_tool_in_response_raises_http_exception(guardrail): + policy_map = {"exfil_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_response_inputs(["exfil_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert exc_info.value.status_code == 400 + assert "exfil_tool" in exc_info.value.detail["blocked_tools"] + + +@pytest.mark.asyncio +async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail): + policy_map = {"safe_tool": "trusted", "bad_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + blocked = exc_info.value.detail["blocked_tools"] + assert "bad_tool" in blocked + assert "safe_tool" not in blocked + + +@pytest.mark.asyncio +async def test_tool_not_in_db_passes_through(guardrail): + """Tools not found in the DB (no entry) should not be blocked.""" + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value={})): + inputs: Any = _tool_request_inputs(["unknown_tool"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_get_policies_cached_uses_cache(guardrail): + """Second call with same tool names should return the cached result.""" + policy_map = {"tool_a": "trusted"} + with patch( + "litellm.proxy.db.tool_registry_writer.get_tools_by_names", + new=AsyncMock(return_value=policy_map), + ) as mock_db, patch( + "litellm.proxy.proxy_server.prisma_client", + new=MagicMock(), + ): + # first call — should hit DB + result1 = await guardrail._get_policies_cached(["tool_a"]) + assert result1 == policy_map + + # second call — should hit cache, not DB again + result2 = await guardrail._get_policies_cached(["tool_a"]) + assert result2 == policy_map + + assert mock_db.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_policies_cached_no_prisma(guardrail): + """Without a prisma client, returns empty dict.""" + with patch( + "litellm.proxy.proxy_server.prisma_client", + None, + ): + result = await guardrail._get_policies_cached(["tool_a"]) + assert result == {} + + +@pytest.mark.asyncio +async def test_response_tool_calls_as_objects(guardrail): + """tool_calls that are objects (not dicts) with .function.name should work.""" + policy_map = {"obj_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + fn = MagicMock() + fn.name = "obj_tool" + tc = MagicMock() + tc.function = fn + inputs: Any = {"tool_calls": [tc]} + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py new file mode 100644 index 000000000000..66297434ff1a --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -0,0 +1,151 @@ +""" +Unit tests for tool management endpoints (/v1/tool/*). +Uses FastAPI TestClient with mocked DB functions. +""" + +import os +import sys +from datetime import datetime, timezone +from typing import Optional +from unittest.mock import MagicMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.management_endpoints.tool_management_endpoints import router +from litellm.types.tool_management import LiteLLM_ToolTableRow + +# --- helpers --- + + +def _make_tool_row( + tool_name: str = "my_tool", + call_policy: str = "untrusted", + origin: Optional[str] = None, +) -> LiteLLM_ToolTableRow: + now = datetime.now(timezone.utc) + return LiteLLM_ToolTableRow( + tool_id="uuid-1", + tool_name=tool_name, + origin=origin, + call_policy=call_policy, # type: ignore[arg-type] + assignments={}, + created_at=now, + updated_at=now, + ) + + +def _make_app() -> FastAPI: + """Build a minimal FastAPI app with the tool management router.""" + app = FastAPI() + app.include_router(router) + return app + + +# Stub the auth dependency so we don't need a real proxy running. +def _override_auth(): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(api_key="sk-test", user_id="admin") + + +# --- test class --- + + +class TestToolManagementEndpoints: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_list_tools") + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", + new_callable=MagicMock, + ) + def test_list_tools_returns_200(self, mock_prisma, mock_db_list): + mock_db_list.return_value = [_make_tool_row()] + mock_prisma.__bool__ = MagicMock(return_value=True) + + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 200 + body = resp.json() + assert body["total"] == 1 + assert body["tools"][0]["tool_name"] == "my_tool" + + @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_list_tools") + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", + new_callable=MagicMock, + ) + def test_list_tools_with_policy_filter(self, mock_prisma, mock_db_list): + mock_db_list.return_value = [_make_tool_row(call_policy="blocked")] + mock_prisma.__bool__ = MagicMock(return_value=True) + + resp = self.client.get("/v1/tool/list?call_policy=blocked") + assert resp.status_code == 200 + assert resp.json()["tools"][0]["call_policy"] == "blocked" + + @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_get_tool") + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", + new_callable=MagicMock, + ) + def test_get_tool_found(self, mock_prisma, mock_db_get): + mock_db_get.return_value = _make_tool_row(tool_name="tool_a") + mock_prisma.__bool__ = MagicMock(return_value=True) + + resp = self.client.get("/v1/tool/tool_a") + assert resp.status_code == 200 + assert resp.json()["tool_name"] == "tool_a" + + @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_get_tool") + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", + new_callable=MagicMock, + ) + def test_get_tool_not_found_returns_404(self, mock_prisma, mock_db_get): + mock_db_get.return_value = None + mock_prisma.__bool__ = MagicMock(return_value=True) + + resp = self.client.get("/v1/tool/nonexistent", follow_redirects=True) + assert resp.status_code == 404 + + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.db_update_tool_policy" + ) + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", + new_callable=MagicMock, + ) + def test_update_tool_policy_blocked(self, mock_prisma, mock_db_update): + mock_db_update.return_value = _make_tool_row(call_policy="blocked") + mock_prisma.__bool__ = MagicMock(return_value=True) + + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "blocked"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["call_policy"] == "blocked" + assert body["updated"] is True + + @patch( + "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", + None, + ) + def test_list_tools_no_db_returns_500(self): + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 500 + + def test_update_tool_policy_invalid_policy_returns_422(self): + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "invalid_value"}, + ) + assert resp.status_code == 422 diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index fb749d7afb07..258c2ccb0e03 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -38,6 +38,7 @@ import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage"; import VectorStoreManagement from "@/components/vector_store_management"; +import ToolPolicies from "@/components/ToolPolicies"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -548,6 +549,8 @@ function CreateKeyPageContent() { ) : page == "vector-stores" ? ( + ) : page == "tool-policies" ? ( + ) : page == "guardrails-monitor" ? ( ) : page == "new_usage" ? ( diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx new file mode 100644 index 000000000000..b12c324849f6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -0,0 +1,410 @@ +"use client"; + +import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; +import { Switch } from "@tremor/react"; +import { Select, Tooltip } from "antd"; +import { + Table, + TableHead, + TableHeaderCell, + TableBody, + TableRow, + TableCell, +} from "@tremor/react"; +import { TimeCell } from "./view_logs/time_cell"; +import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import FilterComponent, { FilterOption } from "./molecules/filter"; +import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking"; + +const POLICY_OPTIONS = [ + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, +] as const; + +type PolicyValue = "trusted" | "blocked"; + +const policyStyle = (p: string) => + POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; + +type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at"; + +interface FilterValues { + [key: string]: string; +} + +interface ToolPoliciesProps { + accessToken: string | null; + userRole?: string; +} + +const PolicySelect: React.FC<{ + value: string; + toolName: string; + saving: boolean; + onChange: (toolName: string, policy: string) => void; +}> = ({ value, toolName, saving, onChange }) => { + const style = policyStyle(value); + return ( + { setSearchTerm(e.target.value); setCurrentPage(1); }} + /> + + + + + +
+ Live Tail + +
+ + + + +
+ + Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + + Page {currentPage} of {totalPages} +
+ + +
+
+ + + {/* Filter row */} +
+ +
+ + + {/* Auto-refresh banner */} + {isLiveTail && ( +
+ Auto-refreshing every 15 seconds + +
+ )} + + {error && ( +
{error}
+ )} + + {/* Table */} + + + + + + + + Key Hash + + Origin + + + + {loading ? ( + + Loading tools… + + ) : paginated.length === 0 ? ( + + + No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery. + + + ) : ( + paginated.map((tool) => ( + + + + + + + + {tool.tool_name} + + + + + + + + + {tool.team_id ?? "-"} + + + + + + {tool.key_hash ?? "-"} + + + + + + {tool.key_alias ?? "-"} + + + + + {tool.origin ?? "-"} + + + + )) + )} + +
+ + {/* Bottom pagination (only when > 1 page) */} + {totalPages > 1 && ( +
+ Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length} +
+ + +
+
+ )} + + + ); +}; + +export default ToolPolicies; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index da3ca2a8baef..2cbeb22ec811 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -134,6 +134,12 @@ const menuGroups: MenuGroup[] = [ label: "Vector Stores", icon: , }, + { + key: "tool-policies", + page: "tool-policies", + label: "Tool Policies", + icon: , + }, ], }, ], diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6ffd744cce99..b18fe3ccb25f 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9854,3 +9854,53 @@ export const checkGdprCompliance = async ( } return response.json(); }; + +export interface ToolRow { + tool_id: string; + tool_name: string; + origin?: string; + call_policy: string; + assignments?: Record; + created_at?: string; + updated_at?: string; + created_by?: string; + updated_by?: string; +} + +export const fetchToolsList = async (accessToken: string): Promise => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/list` : `/v1/tool/list`; + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.text(); + throw new Error(errorData); + } + const data = await response.json(); + return data.tools ?? []; +}; + +export const updateToolPolicy = async ( + accessToken: string, + toolName: string, + callPolicy: string +): Promise => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/policy` : `/v1/tool/policy`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ tool_name: toolName, call_policy: callPolicy }), + }); + if (!response.ok) { + const errorData = await response.text(); + throw new Error(errorData); + } + return response.json(); +}; From 2d28f2199afe28b2e8c53af96fe7e0810a969539 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 15:09:23 -0800 Subject: [PATCH 2/7] feat(tool-policies): track call_count + discover tools from request body and /messages API - Add call_count column to LiteLLM_ToolTable; incremented on every flush - Extract tools from request body too (not just response tool_calls): - OpenAI /chat/completions: tools[].function.name - Anthropic /messages pass-through: request_body.tools[].name - Show call_count column in UI table (sortable) - UI: drop dual_llm option, keep only trusted/blocked --- litellm/proxy/db/db_spend_update_writer.py | 72 ++++++++++++------- litellm/proxy/db/tool_registry_writer.py | 26 ++++--- litellm/proxy/schema.prisma | 1 + litellm/types/tool_management.py | 1 + .../src/components/ToolPolicies.tsx | 10 ++- .../src/components/networking.tsx | 4 ++ 6 files changed, 73 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d3c9b14c2ec7..edf0cf0d3977 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -267,12 +267,16 @@ def _enqueue_tool_registry_upsert( team_id: Optional[str] = None, ) -> None: """ - Extract tool names from the LLM response and enqueue them for upsert - into LiteLLM_ToolTable via ToolDiscoveryQueue. + Extract tool names from the LLM request and response and enqueue them + for upsert into LiteLLM_ToolTable via ToolDiscoveryQueue. - Handles two sources: + Handles four sources: - MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name - - Regular function calls: completion_response.choices[].message.tool_calls[].function.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: @@ -284,6 +288,17 @@ def _enqueue_tool_registry_upsert( _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: @@ -294,17 +309,34 @@ def _enqueue_tool_registry_upsert( 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: - self.tool_discovery_queue.add_update( - ToolDiscoveryQueueItem( - tool_name=tool_name, - origin=mcp_server_name or "user_defined", - key_hash=hashed_token, - team_id=team_id, - key_alias=key_alias, - ) - ) - - # --- Regular function/tool calls in LLM response --- + _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) @@ -319,15 +351,7 @@ def _enqueue_tool_registry_upsert( continue tool_name = getattr(fn, "name", None) if tool_name: - self.tool_discovery_queue.add_update( - ToolDiscoveryQueueItem( - tool_name=tool_name, - origin="user_defined", - key_hash=hashed_token, - team_id=team_id, - key_alias=key_alias, - ) - ) + _enqueue(tool_name) except Exception as e: verbose_proxy_logger.debug( "_enqueue_tool_registry_upsert error (non-blocking): %s", e diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 8aedf65a6dd4..efb411fea77b 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -25,6 +25,7 @@ def _row_to_model(row: dict) -> LiteLLM_ToolTableRow: tool_name=row.get("tool_name", ""), origin=row.get("origin"), call_policy=row.get("call_policy", "untrusted"), + call_count=int(row.get("call_count") or 0), assignments=row.get("assignments"), key_hash=row.get("key_hash"), team_id=row.get("team_id"), @@ -43,18 +44,13 @@ async def batch_upsert_tools( """ Batch-upsert tool registry rows via raw SQL. - Uses INSERT ON CONFLICT DO NOTHING so that: - - First insert sets call_policy to "untrusted" (the schema default). - - Subsequent upserts are no-ops (existing policy is preserved). + On first insert: sets call_policy = "untrusted" (schema default), call_count = 1. + On conflict: increments call_count; preserves existing call_policy. """ if not items: return try: - data = [ - item - for item in items - if item.get("tool_name") - ] + data = [item for item in items if item.get("tool_name")] if not data: return for item in data: @@ -66,9 +62,11 @@ async def batch_upsert_tools( key_alias = item.get("key_alias") await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" ' - "(tool_id, tool_name, origin, call_policy, created_by, updated_by, key_hash, team_id, key_alias) " - "VALUES (gen_random_uuid()::text, $1, $2, 'untrusted', $3, $3, $4, $5, $6) " - "ON CONFLICT (tool_name) DO NOTHING", + "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) " + "VALUES (gen_random_uuid()::text, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " + "ON CONFLICT (tool_name) DO UPDATE SET " + "call_count = \"LiteLLM_ToolTable\".call_count + 1, " + "updated_at = NOW()", tool_name, origin, created_by, @@ -91,14 +89,14 @@ async def list_tools( try: if call_policy is not None: rows = await prisma_client.db.query_raw( - 'SELECT tool_id, tool_name, origin, call_policy, assignments, ' + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' 'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC', call_policy, ) else: rows = await prisma_client.db.query_raw( - 'SELECT tool_id, tool_name, origin, call_policy, assignments, ' + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' 'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC', ) @@ -115,7 +113,7 @@ async def get_tool( """Return a single tool row by tool_name.""" try: rows = await prisma_client.db.query_raw( - 'SELECT tool_id, tool_name, origin, call_policy, assignments, ' + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' 'FROM "LiteLLM_ToolTable" WHERE tool_name = $1', tool_name, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2bfd10b01581..0effe92cdf55 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1057,6 +1057,7 @@ model LiteLLM_ToolTable { tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" origin String? // MCP server name or "user_defined" call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" + call_count Int @default(0) // cumulative number of times this tool was seen assignments Json? @default("{}") key_hash String? // hash of the virtual key that first called this tool team_id String? // team that first called this tool diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index b0bf0fc1ebe3..8704ff27759a 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -15,6 +15,7 @@ class LiteLLM_ToolTableRow(BaseModel): tool_name: str origin: Optional[str] = None call_policy: ToolCallPolicy = "untrusted" + call_count: int = 0 assignments: Optional[Dict] = None key_hash: Optional[str] = None team_id: Optional[str] = None diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index b12c324849f6..ee29e81b03de 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -27,7 +27,7 @@ type PolicyValue = "trusted" | "blocked"; const policyStyle = (p: string) => POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; -type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at"; +type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; interface FilterValues { [key: string]: string; @@ -324,6 +324,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { + Key Hash @@ -333,11 +334,11 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { {loading ? ( - Loading tools… + Loading tools… ) : paginated.length === 0 ? ( - + No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery. @@ -362,6 +363,9 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { onChange={handlePolicyChange} /> + + {(tool.call_count ?? 0).toLocaleString()} + {tool.team_id ?? "-"} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index b18fe3ccb25f..8536a584ee12 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9860,7 +9860,11 @@ export interface ToolRow { tool_name: string; origin?: string; call_policy: string; + call_count?: number; assignments?: Record; + key_hash?: string; + team_id?: string; + key_alias?: string; created_at?: string; updated_at?: string; created_by?: string; From ab46c1bd5b7e1a29f1e63c78e93516a0657361cc Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 15:32:30 -0800 Subject: [PATCH 3/7] fix: address greptile review feedback - Remove redundant @@index([tool_name]) from schema.prisma (tool_name has @unique which already creates an index) - Replace gen_random_uuid()::text with str(uuid.uuid4()) for portability - Rewrite test_tool_registry_writer.py to mock execute_raw/query_raw (actual implementation) instead of Prisma model methods - Fix test patches in test_tool_management_endpoints.py to target source modules since imports are inside function bodies - Add "Tool Policies" page title to ToolPolicies.tsx --- litellm/proxy/db/tool_registry_writer.py | 7 +- litellm/proxy/schema.prisma | 1 - .../proxy/db/test_tool_registry_writer.py | 176 +++++++++++------- .../test_tool_management_endpoints.py | 62 +++--- .../src/components/ToolPolicies.tsx | 1 + 5 files changed, 143 insertions(+), 104 deletions(-) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index efb411fea77b..d566a4308fc9 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -9,6 +9,7 @@ when running against an older generated schema. """ +import uuid from typing import TYPE_CHECKING, Dict, List, Optional from litellm._logging import verbose_proxy_logger @@ -63,7 +64,7 @@ async def batch_upsert_tools( await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" ' "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) " - "VALUES (gen_random_uuid()::text, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " + "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " "ON CONFLICT (tool_name) DO UPDATE SET " "call_count = \"LiteLLM_ToolTable\".call_count + 1, " "updated_at = NOW()", @@ -73,6 +74,7 @@ async def batch_upsert_tools( key_hash, team_id, key_alias, + str(uuid.uuid4()), ) verbose_proxy_logger.debug( "tool_registry_writer: upserted %d tool(s)", len(data) @@ -137,11 +139,12 @@ async def update_tool_policy( _updated_by = updated_by or "system" await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by) ' - "VALUES (gen_random_uuid()::text, $1, $2, $3, $3) " + "VALUES ($4, $1, $2, $3, $3) " "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = NOW()", tool_name, call_policy, _updated_by, + str(uuid.uuid4()), ) return await get_tool(prisma_client, tool_name) except Exception as e: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 0effe92cdf55..23917cf7c7fd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1067,7 +1067,6 @@ model LiteLLM_ToolTable { updated_at DateTime @default(now()) @updatedAt updated_by String? - @@index([tool_name]) @@index([call_policy]) @@index([team_id]) } diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py index b975aa8e852e..44f9e32058a5 100644 --- a/tests/test_litellm/proxy/db/test_tool_registry_writer.py +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -1,6 +1,6 @@ """ -Unit tests for tool_registry_writer.py — tests use a mock prisma client -to avoid requiring a real DB connection. +Unit tests for tool_registry_writer.py — uses a mock prisma client +that exposes execute_raw / query_raw (matching the actual raw-SQL implementation). """ import os @@ -21,46 +21,48 @@ ) -def _make_prisma(rows=None): - """Return a minimal mock prisma_client.""" - now = datetime.now(timezone.utc) - default_row = MagicMock( - tool_id="uuid-1", - tool_name="my_tool", - origin="user_defined", - call_policy="untrusted", - assignments={}, - created_at=now, - updated_at=now, - created_by=None, - updated_by=None, - ) - table = MagicMock() - table.create_many = AsyncMock(return_value=None) - table.find_many = AsyncMock(return_value=rows if rows is not None else [default_row]) - table.find_unique = AsyncMock(return_value=default_row) - table.upsert = AsyncMock(return_value=default_row) +def _make_prisma(query_rows=None): + """Return a minimal mock prisma_client with execute_raw / query_raw.""" + default_row = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "untrusted", + "call_count": 1, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + rows = query_rows if query_rows is not None else [default_row] prisma = MagicMock() - prisma.db.litellm_tooltable = table + prisma.db.execute_raw = AsyncMock(return_value=None) + prisma.db.query_raw = AsyncMock(return_value=rows) return prisma @pytest.mark.asyncio -async def test_batch_upsert_tools_calls_create_many(): +async def test_batch_upsert_tools_calls_execute_raw(): prisma = _make_prisma() items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.create_many.assert_awaited_once() - call_kwargs = prisma.db.litellm_tooltable.create_many.call_args - assert call_kwargs.kwargs["skip_duplicates"] is True + prisma.db.execute_raw.assert_awaited_once() + call_args = prisma.db.execute_raw.call_args + sql = call_args.args[0] + assert "LiteLLM_ToolTable" in sql + assert "ON CONFLICT" in sql @pytest.mark.asyncio async def test_batch_upsert_tools_empty_list(): prisma = _make_prisma() await batch_upsert_tools(prisma, []) - prisma.db.litellm_tooltable.create_many.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio @@ -68,48 +70,68 @@ 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.create_many.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_upsert_multiple_tools_calls_execute_raw_per_tool(): + prisma = _make_prisma() + items = [ + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, + {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, + ] + await batch_upsert_tools(prisma, items) + assert prisma.db.execute_raw.await_count == 2 @pytest.mark.asyncio async def test_list_tools_no_filter(): - now = datetime.now(timezone.utc) - row = MagicMock( - tool_id="id1", - tool_name="tool_a", - origin="mcp", - call_policy="untrusted", - assignments={}, - created_at=now, - updated_at=now, - created_by=None, - updated_by=None, - ) - prisma = _make_prisma(rows=[row]) + row = { + "tool_id": "id1", + "tool_name": "tool_a", + "origin": "mcp", + "call_policy": "untrusted", + "call_count": 5, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + prisma = _make_prisma(query_rows=[row]) result = await list_tools(prisma) assert len(result) == 1 assert result[0].tool_name == "tool_a" + assert result[0].call_count == 5 + prisma.db.query_raw.assert_awaited_once() @pytest.mark.asyncio async def test_list_tools_with_policy_filter(): - now = datetime.now(timezone.utc) - row = MagicMock( - tool_id="id1", - tool_name="blocked_tool", - origin=None, - call_policy="blocked", - assignments=None, - created_at=now, - updated_at=now, - created_by=None, - updated_by=None, - ) - prisma = _make_prisma(rows=[row]) + row = { + "tool_id": "id1", + "tool_name": "blocked_tool", + "origin": None, + "call_policy": "blocked", + "call_count": 2, + "assignments": None, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + prisma = _make_prisma(query_rows=[row]) result = await list_tools(prisma, call_policy="blocked") assert result[0].call_policy == "blocked" - call_kwargs = prisma.db.litellm_tooltable.find_many.call_args - assert call_kwargs.kwargs["where"] == {"call_policy": "blocked"} + call_args = prisma.db.query_raw.call_args + sql = call_args.args[0] + assert "WHERE call_policy" in sql @pytest.mark.asyncio @@ -118,35 +140,51 @@ async def test_get_tool_found(): result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" + prisma.db.query_raw.assert_awaited_once() @pytest.mark.asyncio async def test_get_tool_not_found(): - prisma = _make_prisma() - prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=None) + prisma = _make_prisma(query_rows=[]) result = await get_tool(prisma, "nonexistent") assert result is None @pytest.mark.asyncio -async def test_update_tool_policy_upsert(): - prisma = _make_prisma() +async def test_update_tool_policy_calls_execute_raw(): + row = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "blocked", + "call_count": 1, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": "admin", + } + prisma = _make_prisma(query_rows=[row]) result = await update_tool_policy(prisma, "my_tool", "blocked", "admin") assert result is not None - prisma.db.litellm_tooltable.upsert.assert_awaited_once() - upsert_call = prisma.db.litellm_tooltable.upsert.call_args - assert upsert_call.kwargs["data"]["update"]["call_policy"] == "blocked" - assert upsert_call.kwargs["data"]["update"]["updated_by"] == "admin" + assert result.call_policy == "blocked" + prisma.db.execute_raw.assert_awaited_once() + call_args = prisma.db.execute_raw.call_args + sql = call_args.args[0] + assert "ON CONFLICT" in sql + assert "call_policy" in sql @pytest.mark.asyncio async def test_get_tools_by_names_returns_policy_map(): - now = datetime.now(timezone.utc) rows = [ - MagicMock(tool_name="tool_a", call_policy="trusted"), - MagicMock(tool_name="tool_b", call_policy="blocked"), + {"tool_name": "tool_a", "call_policy": "trusted"}, + {"tool_name": "tool_b", "call_policy": "blocked"}, ] - prisma = _make_prisma(rows=rows) + prisma = _make_prisma(query_rows=rows) result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) assert result == {"tool_a": "trusted", "tool_b": "blocked"} @@ -156,4 +194,4 @@ async def test_get_tools_by_names_empty_list(): prisma = _make_prisma() result = await get_tools_by_names(prisma, []) assert result == {} - prisma.db.litellm_tooltable.find_many.assert_not_awaited() + prisma.db.query_raw.assert_not_awaited() diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index 66297434ff1a..6f1d373fdeef 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -1,13 +1,17 @@ """ Unit tests for tool management endpoints (/v1/tool/*). Uses FastAPI TestClient with mocked DB functions. + +Patches target the source modules (litellm.proxy.db.tool_registry_writer.* +and litellm.proxy.proxy_server.prisma_client) because the endpoint code +imports these inside function bodies to avoid circular imports. """ import os import sys from datetime import datetime, timezone from typing import Optional -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from fastapi import FastAPI from fastapi.testclient import TestClient @@ -51,6 +55,10 @@ def _override_auth(): return UserAPIKeyAuth(api_key="sk-test", user_id="admin") +# A real (non-None) prisma stub for truthiness checks. +_MOCK_PRISMA = MagicMock() + + # --- test class --- @@ -62,14 +70,13 @@ def setup_method(self): app.dependency_overrides[user_api_key_auth] = _override_auth self.client = TestClient(app, raise_server_exceptions=True) - @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_list_tools") @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", - new_callable=MagicMock, + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, ) - def test_list_tools_returns_200(self, mock_prisma, mock_db_list): + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_returns_200(self, mock_db_list): mock_db_list.return_value = [_make_tool_row()] - mock_prisma.__bool__ = MagicMock(return_value=True) resp = self.client.get("/v1/tool/list") assert resp.status_code == 200 @@ -77,54 +84,48 @@ def test_list_tools_returns_200(self, mock_prisma, mock_db_list): assert body["total"] == 1 assert body["tools"][0]["tool_name"] == "my_tool" - @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_list_tools") @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", - new_callable=MagicMock, + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, ) - def test_list_tools_with_policy_filter(self, mock_prisma, mock_db_list): + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_with_policy_filter(self, mock_db_list): mock_db_list.return_value = [_make_tool_row(call_policy="blocked")] - mock_prisma.__bool__ = MagicMock(return_value=True) resp = self.client.get("/v1/tool/list?call_policy=blocked") assert resp.status_code == 200 assert resp.json()["tools"][0]["call_policy"] == "blocked" - @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_get_tool") @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", - new_callable=MagicMock, + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, ) - def test_get_tool_found(self, mock_prisma, mock_db_get): + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_found(self, mock_db_get): mock_db_get.return_value = _make_tool_row(tool_name="tool_a") - mock_prisma.__bool__ = MagicMock(return_value=True) resp = self.client.get("/v1/tool/tool_a") assert resp.status_code == 200 assert resp.json()["tool_name"] == "tool_a" - @patch("litellm.proxy.management_endpoints.tool_management_endpoints.db_get_tool") @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", - new_callable=MagicMock, + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, ) - def test_get_tool_not_found_returns_404(self, mock_prisma, mock_db_get): + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_not_found_returns_404(self, mock_db_get): mock_db_get.return_value = None - mock_prisma.__bool__ = MagicMock(return_value=True) resp = self.client.get("/v1/tool/nonexistent", follow_redirects=True) assert resp.status_code == 404 @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.db_update_tool_policy" + "litellm.proxy.db.tool_registry_writer.update_tool_policy", + new_callable=AsyncMock, ) - @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", - new_callable=MagicMock, - ) - def test_update_tool_policy_blocked(self, mock_prisma, mock_db_update): + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_update_tool_policy_blocked(self, mock_db_update): mock_db_update.return_value = _make_tool_row(call_policy="blocked") - mock_prisma.__bool__ = MagicMock(return_value=True) resp = self.client.post( "/v1/tool/policy", @@ -135,10 +136,7 @@ def test_update_tool_policy_blocked(self, mock_prisma, mock_db_update): assert body["call_policy"] == "blocked" assert body["updated"] is True - @patch( - "litellm.proxy.management_endpoints.tool_management_endpoints.prisma_client", - None, - ) + @patch("litellm.proxy.proxy_server.prisma_client", None) def test_list_tools_no_db_returns_500(self): resp = self.client.get("/v1/tool/list") assert resp.status_code == 500 diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index ee29e81b03de..56139c33fbaa 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -244,6 +244,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { return (
+

Tool Policies

{/* Toolbar */} From c5d4977a505f5f7b7ea87f1574958ccf6a518c08 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 15:48:28 -0800 Subject: [PATCH 4/7] fix: address greptile review round 2 - Replace NOW() with Python datetime parameter in tool_registry_writer (SQLite portability) - Fix cache key collision in tool_policy_guardrail: use null-byte separator instead of colon - Remove type==function filter from request-side tool extraction to match response-side behavior - Clear seen_tool_names on flush so call_count increments per batch cycle not per pod lifetime --- .../db/db_transaction_queue/tool_discovery_queue.py | 12 ++++++++---- litellm/proxy/db/tool_registry_writer.py | 9 +++++++-- .../tool_policy/tool_policy_guardrail.py | 3 +-- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index 6c04a10c0a46..5a27aae4cf8e 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -16,9 +16,11 @@ class ToolDiscoveryQueue: """ In-memory buffer for tool registry upserts. - Deduplicates by tool_name — once a tool has been seen in this process, - it is never enqueued again (the DB row already exists or will be created - during the current flush cycle). + 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: @@ -44,6 +46,8 @@ def add_update(self, item: ToolDiscoveryQueueItem) -> None: ) def flush(self) -> List[ToolDiscoveryQueueItem]: - """Return and clear all pending items.""" + """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 diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index d566a4308fc9..82b27249a00a 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -10,6 +10,7 @@ """ import uuid +from datetime import datetime, timezone from typing import TYPE_CHECKING, Dict, List, Optional from litellm._logging import verbose_proxy_logger @@ -61,13 +62,14 @@ async def batch_upsert_tools( key_hash = item.get("key_hash") team_id = item.get("team_id") key_alias = item.get("key_alias") + now = datetime.now(timezone.utc).isoformat() await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" ' "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) " "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " "ON CONFLICT (tool_name) DO UPDATE SET " "call_count = \"LiteLLM_ToolTable\".call_count + 1, " - "updated_at = NOW()", + "updated_at = $8", tool_name, origin, created_by, @@ -75,6 +77,7 @@ async def batch_upsert_tools( team_id, key_alias, str(uuid.uuid4()), + now, ) verbose_proxy_logger.debug( "tool_registry_writer: upserted %d tool(s)", len(data) @@ -137,14 +140,16 @@ async def update_tool_policy( """Update the call_policy for a tool. Upserts the row if it does not exist yet.""" try: _updated_by = updated_by or "system" + now = datetime.now(timezone.utc).isoformat() await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by) ' "VALUES ($4, $1, $2, $3, $3) " - "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = NOW()", + "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5", tool_name, call_policy, _updated_by, str(uuid.uuid4()), + now, ) return await get_tool(prisma_client, tool_name) except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index ec3e6ede1ae9..fbd72747e614 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -83,7 +83,6 @@ async def apply_guardrail( t["function"]["name"] for t in tools if isinstance(t, dict) - and t.get("type") == "function" and isinstance(t.get("function"), dict) and t["function"].get("name") ] @@ -131,7 +130,7 @@ async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: if not tool_names or prisma_client is None: return {} - cache_key = f"tool_policies:{':'.join(sorted(tool_names))}" + cache_key = f"tool_policies:{chr(0).join(sorted(tool_names))}" cached = await self._policy_cache.async_get_cache(cache_key) if cached is not None and isinstance(cached, dict): verbose_proxy_logger.debug( From c197fe91ed515f78840b2d87d9ee902a9fd5e0cc Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 16:11:35 -0800 Subject: [PATCH 5/7] fix: address greptile review round 3 - Fix test_seen_names_persist_across_flushes to match actual per-flush-cycle behavior - Update module docstring in tool_discovery_queue.py to accurately describe flush behavior - Add created_at/updated_at to raw SQL INSERT in batch_upsert_tools and update_tool_policy --- .../db/db_transaction_queue/tool_discovery_queue.py | 5 +++-- litellm/proxy/db/tool_registry_writer.py | 8 ++++---- .../db/db_transaction_queue/test_tool_discovery_queue.py | 9 +++++---- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index 5a27aae4cf8e..16a3ada40f20 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -2,8 +2,9 @@ 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 pod -lifetime, so DB upserts stop entirely after warmup. +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 diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 82b27249a00a..4e0a8095a08e 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -65,8 +65,8 @@ async def batch_upsert_tools( now = datetime.now(timezone.utc).isoformat() await prisma_client.db.execute_raw( 'INSERT INTO "LiteLLM_ToolTable" ' - "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias) " - "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6) " + "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) " + "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) " "ON CONFLICT (tool_name) DO UPDATE SET " "call_count = \"LiteLLM_ToolTable\".call_count + 1, " "updated_at = $8", @@ -142,8 +142,8 @@ async def update_tool_policy( _updated_by = updated_by or "system" now = datetime.now(timezone.utc).isoformat() await prisma_client.db.execute_raw( - 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by) ' - "VALUES ($4, $1, $2, $3, $3) " + 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) ' + "VALUES ($4, $1, $2, $3, $3, $5, $5) " "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5", tool_name, call_policy, diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index ae19845d3e7e..defdb3834d87 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -53,13 +53,14 @@ def test_flush_clears_pending(queue): assert len(items2) == 0 -def test_seen_names_persist_across_flushes(queue): - """Process-local dedup should prevent re-queuing even after a flush.""" +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() - queue.add_update({"tool_name": "tool_a"}) # already seen + queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle items = queue.flush() - assert len(items) == 0 + assert len(items) == 1 + assert items[0]["tool_name"] == "tool_a" def test_empty_tool_name_ignored(queue): From d4fac55ad5b431ddf39c203cac91c8187eb62327 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 16:17:37 -0800 Subject: [PATCH 6/7] fix: cache tool policies per tool name not per combination Previously the cache key was built from the full set of tool names in a request, so each unique combination of tools got its own cold cache entry and triggered a separate DB query. With N distinct tools across requests this was effectively a DB hit on every request. Now each tool name is cached individually. Cache hits are checked per tool, only missing tools are fetched from DB in a single batch query, and each result is cached separately. Once a tool's policy is warm, any subsequent request using that tool benefits from the cache regardless of what other tools are in the request. --- .../tool_policy/tool_policy_guardrail.py | 49 ++++++++++++------- 1 file changed, 30 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index fbd72747e614..87558566c42d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -122,7 +122,10 @@ async def apply_guardrail( async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: """ Batch-fetch call_policy for the given tool names. - Uses DualCache (in-mem + optional Redis) with TTL to avoid N+1 DB reads. + + Caches per individual tool name (not per combination) so that adding + a new tool to a request doesn't invalidate the cached policies for all + the other tools already in the cache. """ from litellm.proxy.db.tool_registry_writer import get_tools_by_names from litellm.proxy.proxy_server import prisma_client @@ -130,23 +133,31 @@ async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: if not tool_names or prisma_client is None: return {} - cache_key = f"tool_policies:{chr(0).join(sorted(tool_names))}" - cached = await self._policy_cache.async_get_cache(cache_key) - if cached is not None and isinstance(cached, dict): + result: Dict[str, str] = {} + cache_misses: List[str] = [] + + for name in tool_names: + cached = await self._policy_cache.async_get_cache(f"tool_policy:{name}") + if cached is not None and isinstance(cached, str): + result[name] = cached + else: + cache_misses.append(name) + + if cache_misses: + fetched = await get_tools_by_names( + prisma_client=prisma_client, tool_names=cache_misses + ) + for name, policy in fetched.items(): + result[name] = policy + await self._policy_cache.async_set_cache( + key=f"tool_policy:{name}", + value=policy, + ttl=TOOL_POLICY_CACHE_TTL_SECONDS, + ) verbose_proxy_logger.debug( - "ToolPolicyGuardrail: cache hit for tools %s", tool_names + "ToolPolicyGuardrail: fetched %d policies from DB (cache hits: %d)", + len(cache_misses), + len(tool_names) - len(cache_misses), ) - return cached - - policy_map = await get_tools_by_names( - prisma_client=prisma_client, tool_names=tool_names - ) - await self._policy_cache.async_set_cache( - key=cache_key, - value=policy_map, - ttl=TOOL_POLICY_CACHE_TTL_SECONDS, - ) - verbose_proxy_logger.debug( - "ToolPolicyGuardrail: fetched policies %s", policy_map - ) - return policy_map + + return result From 43b9a2af8b15ee6cfa80139e6a2f130492ab3dc3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 24 Feb 2026 16:26:57 -0800 Subject: [PATCH 7/7] Update ui/litellm-dashboard/src/components/ToolPolicies.tsx Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/components/ToolPolicies.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index 56139c33fbaa..0e3f5434e7fd 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -1,7 +1,7 @@ "use client"; import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; -import { Switch } from "@tremor/react"; +import { Select, Switch, Tooltip } from "antd"; import { Select, Tooltip } from "antd"; import { Table,