scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

feat: ensure tool allowlist runs correctly for tool names + mcp's

Browse source
This commit is contained in:
Krrish Dholakia 2026-02-25 22:12:23 -08:00
parent 2487943846
commit fffe8253ea
19 changed files with 1477 additions and 710 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm-proxy-extras/litellm_proxy_extras/migrations/20260225220000_add_tool_policy_override_table/migration.sql Normal file
View file

@ -0,0 +1,27 @@
-- CreateTable
CREATE TABLE "LiteLLM_ToolPolicyOverrideTable" (
"override_id" TEXT NOT NULL,
"tool_name" TEXT NOT NULL,
"team_id" TEXT NOT NULL DEFAULT '',
"key_hash" TEXT NOT NULL DEFAULT '',
"call_policy" TEXT NOT NULL DEFAULT 'blocked',
"key_alias" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"created_by" TEXT,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_by" TEXT,
CONSTRAINT "LiteLLM_ToolPolicyOverrideTable_pkey" PRIMARY KEY ("override_id")
);
-- CreateIndex
CREATE INDEX "LiteLLM_ToolPolicyOverrideTable_tool_name_idx" ON "LiteLLM_ToolPolicyOverrideTable"("tool_name");
-- CreateIndex
CREATE INDEX "LiteLLM_ToolPolicyOverrideTable_team_id_idx" ON "LiteLLM_ToolPolicyOverrideTable"("team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_ToolPolicyOverrideTable_key_hash_idx" ON "LiteLLM_ToolPolicyOverrideTable"("key_hash");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_ToolPolicyOverrideTable_tool_name_team_id_key_hash_key" ON "LiteLLM_ToolPolicyOverrideTable"("tool_name", "team_id", "key_hash");

19
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1076,6 +1076,25 @@ model LiteLLM_ToolTable {
@@index([team_id])
}
// Per-(tool, team/key) policy overrides. When present, override replaces global tool policy for that scope.
model LiteLLM_ToolPolicyOverrideTable {
override_id String @id @default(uuid())
tool_name String
team_id String @default("") // "" = not scoped to team; non-empty = override for this team only
key_hash String @default("") // "" = not scoped to key; non-empty = override for this key only
call_policy String @default("blocked")
key_alias String? // human-readable key alias for UI
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
@@unique([tool_name, team_id, key_hash])
@@index([tool_name])
@@index([team_id])
@@index([key_hash])
}
//Unified Access Groups table for storing unified access groups
model LiteLLM_AccessGroupTable {
access_group_id String @id @default(uuid())

10
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -75,7 +75,7 @@ class AnthropicMessagesHandler(BaseTranslation):
if messages is None:
return data
chat_completion_compatible_request, tool_name_mapping = (
chat_completion_compatible_request, _tool_name_mapping = (
LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
# Use a shallow copy to avoid mutating request data (pop on litellm_metadata).
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
@ -141,6 +141,14 @@ class AnthropicMessagesHandler(BaseTranslation):
return data
def extract_request_tool_names(self, data: dict) -> List[str]:
"""Extract tool names from Anthropic messages request (tools[].name)."""
names: List[str] = []
for tool in data.get("tools") or []:
if isinstance(tool, dict) and tool.get("name"):
names.append(str(tool["name"]))
return names
def _extract_input_text_and_images(
self,
message: Dict[str, Any],

7
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -98,3 +98,10 @@ class BaseTranslation(ABC):
Optional to override in subclasses.
"""
return responses_so_far
def extract_request_tool_names(self, data: dict) -> List[str]:
"""
Extract tool names from the request body for allowlist/policy checks.
Override in tool-capable handlers; default returns [].
"""
return []

13
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -135,6 +135,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return data
def extract_request_tool_names(self, data: dict) -> List[str]:
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
names: List[str] = []
for tool in data.get("tools") or []:
if isinstance(tool, dict) and tool.get("type") == "function":
fn = tool.get("function")
if isinstance(fn, dict) and fn.get("name"):
names.append(str(fn["name"]))
for fn in data.get("functions") or []:
if isinstance(fn, dict) and fn.get("name"):
names.append(str(fn["name"]))
return names
def _extract_inputs(
self,
message: Dict[str, Any],

39
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -30,27 +30,22 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from openai.types.responses.response_function_tool_call import \
ResponseFunctionToolCall
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
OpenAiResponsesToChatCompletionStreamIterator,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
)
from litellm.types.responses.main import (
GenericResponseOutputItem,
OutputFunctionToolCall,
OutputText,
)
OpenAiResponsesToChatCompletionStreamIterator)
from litellm.llms.base_llm.guardrail_translation.base_translation import \
BaseTranslation
from litellm.responses.litellm_completion_transformation.transformation import \
LiteLLMCompletionResponsesConfig
from litellm.types.llms.openai import (ChatCompletionToolCallChunk,
ChatCompletionToolParam)
from litellm.types.responses.main import (GenericResponseOutputItem,
OutputFunctionToolCall, OutputText)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -188,6 +183,18 @@ class OpenAIResponsesHandler(BaseTranslation):
return data
def extract_request_tool_names(self, data: dict) -> List[str]:
"""Extract tool names from Responses API request (tools[].name for function, tools[].server_label for mcp)."""
names: List[str] = []
for tool in data.get("tools") or []:
if not isinstance(tool, dict):
continue
if tool.get("type") == "function" and tool.get("name"):
names.append(str(tool["name"]))
elif tool.get("type") == "mcp" and tool.get("server_label"):
names.append(str(tool["server_label"]))
return names
def _extract_and_transform_tools(
self,
tools: List[Dict[str, Any]],

80
litellm/proxy/_types.py
View file

@ -1,40 +1,60 @@
import enum
import json
from datetime import datetime
from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
Optional, Union)
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
import httpx
from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator,
model_validator)
from pydantic import (
BaseModel,
ConfigDict,
Field,
Json,
field_validator,
model_validator,
)
from typing_extensions import Required, TypedDict
from litellm._uuid import uuid
from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject,
ResponsesAPIResponse)
from litellm.types.mcp import (MCPAuth, MCPAuthType, MCPCredentials,
MCPTransport, MCPTransportType)
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIFileObject,
ResponsesAPIResponse,
)
from litellm.types.mcp import (
MCPAuth,
MCPAuthType,
MCPCredentials,
MCPTransport,
MCPTransportType,
)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo
from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse,
GenericBudgetConfigType, ImageResponse,
LiteLLMBatch, LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase, ModelResponse,
ProviderField, StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse)
from litellm.types.utils import (
CallTypes,
CostBreakdown,
EmbeddingResponse,
GenericBudgetConfigType,
ImageResponse,
LiteLLMBatch,
LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase,
ModelResponse,
ProviderField,
StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse,
)
from litellm.types.videos.main import VideoObject
from .types_utils.utils import (get_instance_fn,
validate_custom_validate_return_type)
from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -2349,8 +2369,7 @@ class UserAPIKeyAuth(
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import \
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@ -2382,8 +2401,7 @@ class UserAPIKeyAuth(
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import \
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
@ -2774,8 +2792,7 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
@model_validator(mode="after")
def mask_api_keys(self):
from litellm.litellm_core_utils.sensitive_data_masker import \
SensitiveDataMasker
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
masker = SensitiveDataMasker(sensitive_patterns={"key"})
@ -3324,6 +3341,11 @@ class ProxyErrorTypes(str, enum.Enum):
Team member is already in team
"""
tool_access_denied = "tool_access_denied"
"""
Tool is not in the allowed tools list for this key/team
"""
@classmethod
def get_model_access_error_type_for_object(
cls, object_type: Literal["key", "user", "team", "org", "project"]

139
litellm/proxy/auth/auth_checks.py
View file

@ -11,7 +11,8 @@ Run checks for:
import asyncio
import re
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
cast)
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
@ -20,44 +21,33 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.dual_cache import LimitedSizeOrderedDict
from litellm.constants import (
CLI_JWT_EXPIRATION_HOURS,
CLI_JWT_TOKEN_NAME,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE)
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
RBAC_ROLES,
CallInfo,
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_JWTAuth,
LiteLLM_ObjectPermissionTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TagTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLMRoutes,
LitellmUserRoles,
NewTeamRequest,
ProxyErrorTypes,
ProxyException,
RoleBasedPermissions,
SpecialModelNames,
UserAPIKeyAuth,
)
from litellm.proxy._types import (RBAC_ROLES, CallInfo,
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable, LiteLLM_EndUserTable,
Litellm_EntityType, LiteLLM_JWTAuth,
LiteLLM_ObjectPermissionTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TagTable, LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable, LiteLLMRoutes,
LitellmUserRoles, NewTeamRequest,
ProxyErrorTypes, ProxyException,
RoleBasedPermissions, SpecialModelNames,
UserAPIKeyAuth)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names)
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
from litellm.router import Router
@ -220,7 +210,47 @@ async def _run_project_checks(
)
async def common_checks(
async def check_tools_allowlist(
request_body: dict,
valid_token: Optional[UserAPIKeyAuth],
team_object: Optional[LiteLLM_TeamTable],
route: str,
) -> None:
"""
Enforce key/team tool allowlist (metadata.allowed_tools). No DB in hot path —
effective allowlist is read from valid_token.metadata and valid_token.team_metadata.
Raises ProxyException with tool_access_denied if a tool is not allowed.
"""
from litellm.litellm_core_utils.api_route_to_call_types import \
get_call_types_for_route
if valid_token is None:
return
call_types = get_call_types_for_route(route)
if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types):
return
tool_names = extract_request_tool_names(route, request_body)
if not tool_names:
return
key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {}
team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {}
key_allowed = key_meta.get("allowed_tools")
team_allowed = team_meta.get("allowed_tools")
effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed
if not isinstance(effective, list) or len(effective) == 0:
return
allowed_set = {str(t) for t in effective}
disallowed = [n for n in tool_names if n not in allowed_set]
if disallowed:
raise ProxyException(
message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.",
type=ProxyErrorTypes.tool_access_denied,
param="tools",
code=status.HTTP_403_FORBIDDEN,
)
async def common_checks( # noqa: PLR0915
request_body: dict,
team_object: Optional[LiteLLM_TeamTable],
user_object: Optional[LiteLLM_UserTable],
@ -435,7 +465,8 @@ async def common_checks(
_request_metadata: dict = request_body.get("metadata", {}) or {}
if _request_metadata.get("guardrails"):
# check if team allowed to modify guardrails
from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
from litellm.proxy.guardrails.guardrail_helpers import \
can_modify_guardrails
can_modify: bool = can_modify_guardrails(team_object)
if can_modify is False:
@ -473,6 +504,14 @@ async def common_checks(
valid_token=valid_token,
)
# 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
await check_tools_allowlist(
request_body=request_body,
valid_token=valid_token,
team_object=team_object,
route=route,
)
return True
@ -1877,9 +1916,8 @@ class ExperimentalUIJWTToken:
def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
encrypt_value_helper
if user_info.user_role is None:
raise Exception("User role is required for experimental UI login")
@ -1925,9 +1963,8 @@ class ExperimentalUIJWTToken:
"""
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
encrypt_value_helper
if user_info.user_role is None:
raise Exception("User role is required for CLI JWT login")
@ -1966,9 +2003,8 @@ class ExperimentalUIJWTToken:
import json
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
decrypt_value_helper
decrypted_token = decrypt_value_helper(
hashed_token, key="ui_hash_key", exception_type="debug"
@ -2263,8 +2299,10 @@ async def _get_resources_from_access_groups(
# Lazy import to avoid circular imports
if prisma_client is None or user_api_key_cache is None:
from litellm.proxy.proxy_server import prisma_client as _prisma_client
from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
from litellm.proxy.proxy_server import \
proxy_logging_obj as _proxy_logging_obj
from litellm.proxy.proxy_server import \
user_api_key_cache as _user_api_key_cache
prisma_client = prisma_client or _prisma_client
user_api_key_cache = user_api_key_cache or _user_api_key_cache
@ -3220,7 +3258,8 @@ async def _tag_max_budget_check(
BudgetExceededError if any tag is over its max budget.
Triggers a budget alert if any tag is over its max budget.
"""
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
if prisma_client is None:
return

255
litellm/proxy/db/tool_registry_writer.py
View file

@ -10,12 +10,23 @@ from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
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.proxy._types import ToolDiscoveryQueueItem
from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy
from litellm.types.tool_management import (
LiteLLM_ToolTableRow,
ToolCallPolicy,
ToolPolicyOverrideRow,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
# Sentinel for "not scoped" in override table (DB unique constraint needs non-null)
_TOOL_OVERRIDE_ANY = ""
TOOL_POLICY_CACHE_KEY_PREFIX = "tool_policy:"
def _row_to_model(row: Union[dict, Any]) -> LiteLLM_ToolTableRow:
"""Convert a Prisma model instance or dict to LiteLLM_ToolTableRow."""
@ -204,3 +215,245 @@ async def get_tools_by_names(
"tool_registry_writer get_tools_by_names error: %s", e
)
return {}
def _override_row_to_model(row: Any) -> ToolPolicyOverrideRow:
"""Convert a Prisma override row to ToolPolicyOverrideRow."""
model_dump = getattr(row, "model_dump", None)
if callable(model_dump):
row = model_dump()
elif not isinstance(row, dict):
row = {
k: getattr(row, k, None)
for k in (
"override_id",
"tool_name",
"team_id",
"key_hash",
"call_policy",
"key_alias",
"created_at",
"updated_at",
)
}
def _norm(s: Optional[str]) -> Optional[str]:
if s is None or s == _TOOL_OVERRIDE_ANY:
return None
return s or None
return ToolPolicyOverrideRow(
override_id=row.get("override_id", ""),
tool_name=row.get("tool_name", ""),
team_id=_norm(row.get("team_id")),
key_hash=_norm(row.get("key_hash")),
call_policy=row.get("call_policy", "blocked"),
key_alias=row.get("key_alias"),
created_at=row.get("created_at"),
updated_at=row.get("updated_at"),
)
async def list_overrides_for_tool(
prisma_client: "PrismaClient",
tool_name: str,
) -> List[ToolPolicyOverrideRow]:
"""Return all policy overrides for a tool."""
try:
rows = await prisma_client.db.litellm_toolpolicyoverridetable.find_many(
where={"tool_name": tool_name},
order={"created_at": "desc"},
)
return [_override_row_to_model(row) for row in rows]
except Exception as e:
verbose_proxy_logger.error(
"tool_registry_writer list_overrides_for_tool error: %s", e
)
return []
async def upsert_tool_policy_override(
prisma_client: "PrismaClient",
tool_name: str,
call_policy: ToolCallPolicy,
team_id: Optional[str] = None,
key_hash: Optional[str] = None,
key_alias: Optional[str] = None,
updated_by: Optional[str] = None,
) -> Optional[ToolPolicyOverrideRow]:
"""Create or update a per-(tool, team, key) policy override."""
try:
_team = (team_id or "").strip() or _TOOL_OVERRIDE_ANY
_key = (key_hash or "").strip() or _TOOL_OVERRIDE_ANY
_updated_by = updated_by or "system"
now = datetime.now(timezone.utc)
table = prisma_client.db.litellm_toolpolicyoverridetable
await table.upsert(
where={
"tool_name_team_id_key_hash": {
"tool_name": tool_name,
"team_id": _team,
"key_hash": _key,
}
},
data={
"create": {
"override_id": str(uuid.uuid4()),
"tool_name": tool_name,
"team_id": _team,
"key_hash": _key,
"call_policy": call_policy,
"key_alias": key_alias,
"created_by": _updated_by,
"updated_by": _updated_by,
"created_at": now,
"updated_at": now,
},
"update": {
"call_policy": call_policy,
"key_alias": key_alias,
"updated_by": _updated_by,
"updated_at": now,
},
},
)
row = await table.find_unique(
where={
"tool_name_team_id_key_hash": {
"tool_name": tool_name,
"team_id": _team,
"key_hash": _key,
}
}
)
if row is None:
return None
return _override_row_to_model(row)
except Exception as e:
verbose_proxy_logger.error(
"tool_registry_writer upsert_tool_policy_override error: %s", e
)
return None
async def delete_tool_policy_override(
prisma_client: "PrismaClient",
tool_name: str,
team_id: Optional[str] = None,
key_hash: Optional[str] = None,
) -> bool:
"""Remove a policy override. Exactly one of team_id or key_hash should be set for a specific override."""
try:
_team = (team_id or "").strip() or _TOOL_OVERRIDE_ANY
_key = (key_hash or "").strip() or _TOOL_OVERRIDE_ANY
result = await prisma_client.db.litellm_toolpolicyoverridetable.delete_many(
where={
"tool_name": tool_name,
"team_id": _team,
"key_hash": _key,
}
)
return (result or 0) > 0
except Exception as e:
verbose_proxy_logger.error(
"tool_registry_writer delete_tool_policy_override error: %s", e
)
return False
async def get_effective_policies(
prisma_client: "PrismaClient",
tool_names: List[str],
team_id: Optional[str],
key_hash: Optional[str],
) -> Dict[str, str]:
"""
Return effective call_policy per tool: override for (tool, team_id, key_hash) if present,
otherwise global policy from LiteLLM_ToolTable.
"""
if not tool_names:
return {}
_team = (team_id or "").strip() or _TOOL_OVERRIDE_ANY
_key = (key_hash or "").strip() or _TOOL_OVERRIDE_ANY
try:
# Fetch overrides for (tool_name, team_id, key_hash) - exact match
overrides = await prisma_client.db.litellm_toolpolicyoverridetable.find_many(
where={
"tool_name": {"in": tool_names},
"team_id": _team,
"key_hash": _key,
}
)
override_map = {row.tool_name: row.call_policy for row in overrides}
# Global policies for tools that have no override
missing = [n for n in tool_names if n not in override_map]
if not missing:
return override_map
global_map = await get_tools_by_names(
prisma_client=prisma_client, tool_names=missing
)
override_map.update(global_map)
return override_map
except Exception as e:
verbose_proxy_logger.error(
"tool_registry_writer get_effective_policies error: %s", e
)
return {}
def _effective_cache_suffix(team_id: Optional[str], key_hash: Optional[str]) -> str:
"""Cache key suffix so different request contexts (team/key) get correct policies."""
return f":{team_id or ''}:{key_hash or ''}"
async def get_tool_policies_cached(
tool_names: List[str],
cache: DualCache,
prisma_client: Optional["PrismaClient"],
team_id: Optional[str] = None,
key_hash: Optional[str] = None,
) -> Dict[str, str]:
"""
Return effective call_policy per tool (override for team/key if present, else global).
Cache-first; cache key includes team_id and key_hash when provided.
"""
if not tool_names:
return {}
suffix = _effective_cache_suffix(team_id, key_hash)
result: Dict[str, str] = {}
cache_misses: List[str] = []
for name in tool_names:
key = f"{TOOL_POLICY_CACHE_KEY_PREFIX}{name}{suffix}"
cached = await cache.async_get_cache(key=key)
if cached is not None and isinstance(cached, str):
result[name] = cached
else:
cache_misses.append(name)
if cache_misses and prisma_client is not None:
try:
if team_id is not None or key_hash is not None:
fetched = await get_effective_policies(
prisma_client=prisma_client,
tool_names=cache_misses,
team_id=team_id,
key_hash=key_hash,
)
else:
fetched = await get_tools_by_names(
prisma_client=prisma_client, tool_names=cache_misses
)
for name, policy in fetched.items():
result[name] = policy
await cache.async_set_cache(
key=f"{TOOL_POLICY_CACHE_KEY_PREFIX}{name}{suffix}",
value=policy,
ttl=TOOL_POLICY_CACHE_TTL_SECONDS,
)
verbose_proxy_logger.debug(
"get_tool_policies_cached: fetched %d from DB (hits: %d)",
len(cache_misses),
len(tool_names) - len(cache_misses),
)
except Exception as e:
verbose_proxy_logger.error(
"tool_registry_writer get_tool_policies_cached error: %s", e
)
return result

151
litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py
View file

@ -23,32 +23,70 @@ or both pre and post call:
mode: during_call # runs before LLM and on response
"""
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
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.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
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
GUARDRAIL_NAME = "tool_policy"
def _get_request_team_and_key(request_data: dict) -> Tuple[Optional[str], Optional[str]]:
"""Extract team_id and key hash from request_data (litellm_metadata or metadata)."""
if not request_data:
return None, None
for key in ("litellm_metadata", "metadata"):
meta = request_data.get(key)
if not isinstance(meta, dict):
continue
team_id = meta.get("user_api_key_team_id")
key_hash = meta.get("user_api_key_api_key") or meta.get("user_api_key")
if team_id is not None or key_hash is not None:
return (
str(team_id).strip() if team_id else None,
str(key_hash).strip() if key_hash else None,
)
return None, None
def _get_request_route_from_data(request_data: dict) -> Optional[str]:
"""Get request route from request_data (metadata or top-level)."""
route = request_data.get("user_api_key_request_route")
if route:
return route
meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {}
return meta.get("user_api_key_request_route")
def _get_effective_allowed_tools_from_request(request_data: dict) -> Optional[List[str]]:
"""Key allowed_tools overrides team; empty/missing means no restriction."""
meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {}
key_meta = meta.get("user_api_key_metadata") or {}
team_meta = meta.get("user_api_key_team_metadata") or {}
key_allowed = key_meta.get("allowed_tools") if isinstance(key_meta, dict) else None
team_allowed = team_meta.get("allowed_tools") if isinstance(team_meta, dict) else None
if isinstance(key_allowed, list) and len(key_allowed) > 0:
return key_allowed
if isinstance(team_allowed, list) and len(team_allowed) > 0:
return team_allowed
return None
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.
Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable
and key/team allowed_tools (allowlist). No DB in hot path for policy lookup —
uses shared cache (user_api_key_cache).
"""
def __init__(self, **kwargs: Any) -> None:
@ -70,12 +108,7 @@ class ToolPolicyGuardrail(CustomGuardrail):
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".
Enforce key/team allowlist then DB call_policy on request tools / response tool_calls.
"""
if input_type == "request":
tools = inputs.get("tools") or []
@ -86,6 +119,12 @@ class ToolPolicyGuardrail(CustomGuardrail):
and isinstance(t.get("function"), dict)
and t["function"].get("name")
]
if not tool_names:
route = _get_request_route_from_data(request_data)
if route:
from litellm.proxy.guardrails.tool_name_extraction import \
extract_request_tool_names
tool_names = extract_request_tool_names(route, request_data)
else: # response
tool_calls = inputs.get("tool_calls") or []
tool_names = []
@ -101,8 +140,26 @@ class ToolPolicyGuardrail(CustomGuardrail):
if not tool_names:
return inputs
policy_map = await self._get_policies_cached(tool_names)
allowed_tools = _get_effective_allowed_tools_from_request(request_data)
if isinstance(allowed_tools, list) and len(allowed_tools) > 0:
allowed_set = {str(t) for t in allowed_tools}
disallowed = [n for n in tool_names if n not in allowed_set]
if disallowed:
verbose_proxy_logger.warning(
"ToolPolicyGuardrail: tool(s) %s not in key/team allowed_tools",
disallowed,
)
raise HTTPException(
status_code=400,
detail={
"error": "Violated tool allowlist",
"disallowed_tools": disallowed,
"message": f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.",
},
)
team_id, key_hash = _get_request_team_and_key(request_data)
policy_map = await self._get_policies_cached(tool_names, team_id, key_hash)
blocked = [name for name in tool_names if policy_map.get(name) == "blocked"]
if blocked:
verbose_proxy_logger.warning(
@ -119,45 +176,27 @@ class ToolPolicyGuardrail(CustomGuardrail):
return inputs
async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]:
async def _get_policies_cached(
self,
tool_names: List[str],
team_id: Optional[str] = None,
key_hash: Optional[str] = None,
) -> Dict[str, str]:
"""
Batch-fetch call_policy for the given tool names.
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.
Fetch effective call_policy (override for team/key if present, else global)
via shared cache to avoid DB in hot path.
"""
from litellm.proxy.db.tool_registry_writer import get_tools_by_names
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.db.tool_registry_writer import \
get_tool_policies_cached
from litellm.proxy.proxy_server import (prisma_client,
user_api_key_cache)
if not tool_names or prisma_client is None:
if not tool_names:
return {}
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: fetched %d policies from DB (cache hits: %d)",
len(cache_misses),
len(tool_names) - len(cache_misses),
)
return result
return await get_tool_policies_cached(
tool_names=tool_names,
cache=user_api_key_cache,
prisma_client=prisma_client,
team_id=team_id,
key_hash=key_hash,
)

85
litellm/proxy/guardrails/tool_name_extraction.py Normal file
View file

@ -0,0 +1,85 @@
"""
Extract tool names from request body by route/call type.
Used by auth (check_tools_allowlist) and ToolPolicyGuardrail so tool-format
knowledge lives in one place. Uses guardrail translation handlers where available,
with standalone extractors for generate_content and MCP.
"""
from typing import Any, Dict, List
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.llms import load_guardrail_translation_mappings
from litellm.types.utils import CallTypes
# Call types that have no guardrail translation handler; we use standalone extractors
STANDALONE_EXTRACTORS: Dict[str, Any] = {}
def _extract_generate_content_tool_names(data: dict) -> List[str]:
"""Google generateContent: tools[].functionDeclarations[].name"""
names: List[str] = []
for tool in data.get("tools") or []:
if not isinstance(tool, dict):
continue
for decl in tool.get("functionDeclarations") or []:
if isinstance(decl, dict) and decl.get("name"):
names.append(str(decl["name"]))
return names
def _extract_mcp_tool_names(data: dict) -> List[str]:
"""MCP call_tool: name or mcp_tool_name in body"""
names: List[str] = []
name = data.get("name") or data.get("mcp_tool_name")
if name:
names.append(str(name))
return names
def _register_standalone_extractors() -> None:
if STANDALONE_EXTRACTORS:
return
STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = _extract_generate_content_tool_names
STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = _extract_generate_content_tool_names
STANDALONE_EXTRACTORS[CallTypes.call_mcp_tool.value] = _extract_mcp_tool_names
# Tool-capable call types (routes that can send tools in the request)
TOOL_CAPABLE_CALL_TYPES = frozenset({
CallTypes.completion.value,
CallTypes.acompletion.value,
CallTypes.responses.value,
CallTypes.aresponses.value,
CallTypes.anthropic_messages.value,
CallTypes.generate_content.value,
CallTypes.agenerate_content.value,
CallTypes.call_mcp_tool.value,
})
def extract_request_tool_names(route: str, data: dict) -> List[str]:
"""
Extract tool names from the request body for the given route.
Uses guardrail translation handlers when available, else standalone extractors
for generate_content and MCP. Returns [] for non-tool-capable routes or when
no tools are present.
"""
call_types = get_call_types_for_route(route)
if not call_types:
return []
_register_standalone_extractors()
mappings = load_guardrail_translation_mappings()
for call_type in call_types:
if not isinstance(call_type, CallTypes):
continue
if call_type.value not in TOOL_CAPABLE_CALL_TYPES:
continue
if call_type.value in STANDALONE_EXTRACTORS:
return STANDALONE_EXTRACTORS[call_type.value](data)
handler_cls = mappings.get(call_type)
if handler_cls is not None:
names = handler_cls().extract_request_tool_names(data)
if names:
return names
return []

50
litellm/proxy/litellm_pre_call_utils.py
View file

@ -10,16 +10,12 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles,
SpecialHeaders,
TeamCallbackMetadata,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles, SpecialHeaders,
TeamCallbackMetadata, UserAPIKeyAuth)
from litellm.proxy.common_utils.http_parsing_utils import \
_safe_get_request_headers
# Cache special headers as a frozenset for O(1) lookup performance
_SPECIAL_HEADERS_CACHE = frozenset(
@ -28,12 +24,9 @@ _SPECIAL_HEADERS_CACHE = frozenset(
from litellm.router import Router
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
LlmProviders,
ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls,
)
from litellm.types.utils import (LlmProviders, ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls)
service_logger_obj = ServiceLogging() # used for tracking latency on OTEL
@ -667,8 +660,7 @@ class LiteLLMProxyRequestSetup:
return data
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
)
LiteLLM_ManagementEndpoint_MetadataFields_Premium)
# ignore any special fields
added_metadata = {}
@ -1058,6 +1050,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
] = user_api_key_dict.user_max_budget
data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata
data[_metadata_variable_name]["user_api_key_team_metadata"] = (
user_api_key_dict.team_metadata
)
data[_metadata_variable_name]["headers"] = _headers
data[_metadata_variable_name]["endpoint"] = str(request.url)
@ -1501,7 +1496,8 @@ async def move_guardrails_to_metadata(
# Only check policy engine if no local config (avoid import + registry lookup)
if not (has_key_config or has_team_config or has_request_config):
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_registry import \
get_policy_registry
if not get_policy_registry().is_initialized():
# Nothing configured anywhere - clean up request body fields and return
@ -1565,14 +1561,16 @@ async def move_guardrails_to_metadata(
def _is_policy_version_id(s: str) -> bool:
"""Return True if string is a policy version ID (starts with policy_<uuid> prefix)."""
from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
from litellm.proxy.policy_engine.policy_registry import \
POLICY_VERSION_ID_PREFIX
return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX)
def _extract_policy_id(s: str) -> Optional[str]:
"""Extract raw UUID from policy_<uuid> string, or None if not a valid version ID."""
from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
from litellm.proxy.policy_engine.policy_registry import \
POLICY_VERSION_ID_PREFIX
if not _is_policy_version_id(s):
return None
@ -1593,10 +1591,9 @@ def _match_and_track_policies(
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.callback_utils import (
add_policy_sources_to_metadata,
add_policy_to_applied_policies_header,
)
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
add_policy_sources_to_metadata, add_policy_to_applied_policies_header)
from litellm.proxy.policy_engine.attachment_registry import \
get_attachment_registry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
# Get matching policies via attachments (with match reasons for attribution)
@ -1741,7 +1738,8 @@ async def add_guardrails_from_policy_engine(
user_api_key_dict: The user's API key authentication info
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import PolicyMatchContext

128
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

@ -10,7 +10,7 @@ POST /v1/tool/policy - Update the call_policy for a tool
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException
from fastapi import APIRouter, Depends, HTTPException, Query
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
@ -18,6 +18,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.tool_management import (
LiteLLM_ToolTableRow,
ToolCallPolicy,
ToolDetailResponse,
ToolListResponse,
ToolPolicyUpdateRequest,
ToolPolicyUpdateResponse,
@ -58,6 +59,48 @@ async def list_tools(
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/v1/tool/{tool_name:path}/detail",
tags=["tool management"],
dependencies=[Depends(user_api_key_auth)],
response_model=ToolDetailResponse,
)
async def get_tool_detail(
tool_name: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get a single tool with its policy overrides (for UI detail view).
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.db.tool_registry_writer import list_overrides_for_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"
)
overrides = await list_overrides_for_tool(
prisma_client=prisma_client, tool_name=tool_name
)
return ToolDetailResponse(tool=tool, overrides=overrides)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error getting tool detail: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/v1/tool/{tool_name:path}",
tags=["tool management"],
@ -107,18 +150,23 @@ async def update_tool_policy(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Set the call policy for a tool.
Set the call policy for a tool (global) or for a specific team/key (override).
Parameters:
- tool_name: str - The tool to update
- call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked"
- team_id: optional - if set, create/update override for this team only
- key_hash: optional - if set, create/update override for this key only
- key_alias: optional - human-readable key alias for UI
Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove
that tool_call from LLM responses before returning them to the client.
If both team_id and key_hash are omitted, updates the global tool policy.
Setting a tool to "blocked" will cause the ToolPolicyGuardrail to reject
that tool_call for the relevant scope.
"""
from litellm.proxy.db.tool_registry_writer import (
update_tool_policy as db_update_tool_policy,
)
from litellm.proxy.db.tool_registry_writer import upsert_tool_policy_override
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
@ -127,6 +175,28 @@ async def update_tool_policy(
)
try:
if data.team_id is not None or data.key_hash is not None:
override = await upsert_tool_policy_override(
prisma_client=prisma_client,
tool_name=data.tool_name,
call_policy=data.call_policy,
team_id=data.team_id,
key_hash=data.key_hash,
key_alias=data.key_alias,
updated_by=user_api_key_dict.user_id,
)
if override is None:
raise HTTPException(
status_code=500,
detail=f"Failed to update policy override for tool '{data.tool_name}'",
)
return ToolPolicyUpdateResponse(
tool_name=override.tool_name,
call_policy=override.call_policy,
updated=True,
team_id=override.team_id,
key_hash=override.key_hash,
)
updated = await db_update_tool_policy(
prisma_client=prisma_client,
tool_name=data.tool_name,
@ -135,7 +205,8 @@ async def update_tool_policy(
)
if updated is None:
raise HTTPException(
status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'"
status_code=500,
detail=f"Failed to update policy for tool '{data.tool_name}'",
)
return ToolPolicyUpdateResponse(
tool_name=updated.tool_name,
@ -147,3 +218,50 @@ async def update_tool_policy(
except Exception as e:
verbose_proxy_logger.exception("Error updating tool policy: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.delete(
"/v1/tool/{tool_name:path}/overrides",
tags=["tool management"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_tool_policy_override(
tool_name: str,
team_id: Optional[str] = Query(None, description="Team ID of the override to remove"),
key_hash: Optional[str] = Query(None, description="Key hash of the override to remove"),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Remove a policy override for a tool. Specify the override by team_id and/or key_hash
(must match the override that was created; use empty string for unscoped dimension).
"""
from litellm.proxy.db.tool_registry_writer import delete_tool_policy_override
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
)
if team_id is None and key_hash is None:
raise HTTPException(
status_code=400,
detail="At least one of team_id or key_hash is required to identify the override",
)
try:
deleted = await delete_tool_policy_override(
prisma_client=prisma_client,
tool_name=tool_name,
team_id=team_id,
key_hash=key_hash,
)
if not deleted:
raise HTTPException(
status_code=404,
detail=f"No override found for tool '{tool_name}' with the given scope",
)
return {"deleted": True, "tool_name": tool_name}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error deleting tool policy override: %s", e)
raise HTTPException(status_code=500, detail=str(e))

19
litellm/proxy/schema.prisma
View file

@ -1076,6 +1076,25 @@ model LiteLLM_ToolTable {
@@index([team_id])
}
// Per-(tool, team/key) policy overrides. When present, override replaces global tool policy for that scope.
model LiteLLM_ToolPolicyOverrideTable {
override_id String @id @default(uuid())
tool_name String
team_id String @default("") // "" = not scoped to team; non-empty = override for this team only
key_hash String @default("") // "" = not scoped to key; non-empty = override for this key only
call_policy String @default("blocked")
key_alias String? // human-readable key alias for UI
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
@@unique([tool_name, team_id, key_hash])
@@index([tool_name])
@@index([team_id])
@@index([key_hash])
}
//Unified Access Groups table for storing unified access groups
model LiteLLM_AccessGroupTable {
access_group_id String @id @default(uuid())

23
litellm/types/tool_management.py
View file

@ -5,7 +5,7 @@ Pydantic models for Tool Policy management endpoints.
from datetime import datetime
from typing import Dict, List, Literal, Optional
from pydantic import BaseModel
from pydantic import BaseModel, Field
ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"]
@ -34,9 +34,30 @@ class ToolListResponse(BaseModel):
class ToolPolicyUpdateRequest(BaseModel):
tool_name: str
call_policy: ToolCallPolicy
team_id: Optional[str] = None # if set, create/update override for this team
key_hash: Optional[str] = None # if set, create/update override for this key
key_alias: Optional[str] = None # human-readable key alias for UI
class ToolPolicyUpdateResponse(BaseModel):
tool_name: str
call_policy: ToolCallPolicy
updated: bool
team_id: Optional[str] = None
key_hash: Optional[str] = None
class ToolPolicyOverrideRow(BaseModel):
override_id: str
tool_name: str
team_id: Optional[str] = None
key_hash: Optional[str] = None
call_policy: ToolCallPolicy
key_alias: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
class ToolDetailResponse(BaseModel):
tool: LiteLLM_ToolTableRow
overrides: List[ToolPolicyOverrideRow] = Field(default_factory=list)

116
scripts/test_tool_allowlist_script.py Normal file
View file

@ -0,0 +1,116 @@
#!/usr/bin/env python3
"""
Standalone script to test tool allowlist enforcement and tool name extraction.
Run from repo root:
poetry run python scripts/test_tool_allowlist_script.py
Or run the unit tests:
poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v
"""
import asyncio
import sys
from pathlib import Path
# Ensure repo root is on path
repo_root = Path(__file__).resolve().parent.parent
if str(repo_root) not in sys.path:
sys.path.insert(0, str(repo_root))
def test_extraction():
"""Test extract_request_tool_names for each API shape."""
from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names
cases = [
("OpenAI chat tools", "/v1/chat/completions", {"tools": [{"type": "function", "function": {"name": "get_weather"}}]}),
("OpenAI chat functions", "/v1/chat/completions", {"functions": [{"name": "run_sql"}]}),
("OpenAI responses function", "/v1/responses", {"tools": [{"type": "function", "name": "get_current_weather"}]}),
("OpenAI responses MCP", "/v1/responses", {"tools": [{"type": "mcp", "server_label": "dmcp"}]}),
("Anthropic", "/v1/messages", {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}),
("Google generateContent", "/generate_content", {"tools": [{"functionDeclarations": [{"name": "schedule_meeting"}]}]}),
("MCP call_tool", "/mcp/call_tool", {"name": "my_tool", "arguments": {}}),
("Non-tool route", "/v1/embeddings", {"tools": [{"type": "function", "function": {"name": "x"}}]}),
]
print("=== extract_request_tool_names(route, data) ===\n")
for label, route, data in cases:
names = extract_request_tool_names(route, data)
print(f" {label}: {names}")
print()
async def test_check_tools_allowlist():
"""Test check_tools_allowlist with mock tokens."""
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import check_tools_allowlist
def token(metadata=None, team_metadata=None):
return UserAPIKeyAuth(
api_key="test-key",
user_id="user",
team_id="team",
org_id=None,
models=["*"],
metadata=metadata or {},
team_metadata=team_metadata or {},
)
print("=== check_tools_allowlist (auth) ===\n")
# No allowlist -> pass
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
valid_token=token(),
team_object=None,
route="/v1/chat/completions",
)
print(" No allowlist, body has tools: PASS")
# Allowed tool -> pass
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
valid_token=token(metadata={"allowed_tools": ["get_weather"]}),
team_object=None,
route="/v1/chat/completions",
)
print(" allowed_tools=['get_weather'], body has get_weather: PASS")
# Disallowed tool -> raise
try:
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
valid_token=token(metadata={"allowed_tools": ["other_tool"]}),
team_object=None,
route="/v1/chat/completions",
)
print(" DISALLOWED: expected ProxyException")
except ProxyException as e:
if e.type == ProxyErrorTypes.tool_access_denied:
print(" allowed_tools=['other_tool'], body has get_weather: PASS (raised tool_access_denied)")
else:
print(f" Unexpected ProxyException type: {e.type}")
except Exception as e:
print(f" Unexpected: {e}")
# Team allowlist when key empty
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "get_weather"}}]},
valid_token=token(team_metadata={"allowed_tools": ["get_weather"]}),
team_object=None,
route="/v1/chat/completions",
)
print(" team_metadata.allowed_tools=['get_weather']: PASS")
print()
def main():
print("Tool allowlist / tool name extraction – script checks\n")
test_extraction()
asyncio.run(test_check_tools_allowlist())
print("Done. For full unit tests run:")
print(" poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v")
if __name__ == "__main__":
main()

674
tests/test_litellm/proxy/test_tools_allowlist_enforcement.py
View file

@ -1,548 +1,200 @@
"""
Tests for tool allowlist enforcement by team/key (metadata.allowed_tools).
Tests for tool allowlist enforcement (key/team metadata.allowed_tools).
No implementation yet; these tests define expected behavior. When check_tools_allowlist
is implemented in common_checks, disallowed-tool tests should raise; allowed and
no-allowlist tests should pass.
Covers:
- check_tools_allowlist: allowed, disallowed, no allowlist, non-tool routes
- extract_request_tool_names: OpenAI chat, responses, Anthropic, generate_content, MCP
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import common_checks
from litellm.proxy._types import (ProxyErrorTypes, ProxyException,
UserAPIKeyAuth)
from litellm.proxy.auth.auth_checks import check_tools_allowlist
from litellm.proxy.guardrails.tool_name_extraction import (
TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names)
class MockRequest:
"""Mock request with method attribute."""
def __init__(self, method: str = "POST"):
self.method = method
def get_mock_user_token(metadata=None, team_metadata=None) -> UserAPIKeyAuth:
"""Build UserAPIKeyAuth with optional metadata and team_metadata for allowlist."""
kwargs = {
"api_key": "test-key",
"user_id": "test-user",
"team_id": "test-team",
"org_id": "test-org",
"models": ["*"],
"metadata": metadata or {},
}
if team_metadata is not None:
kwargs["team_metadata"] = team_metadata
return UserAPIKeyAuth(**kwargs)
def _tools_allowlist_patches():
"""Patches so only tool-allowlist behavior is under test; heavy/DB parts no-op."""
p1 = patch(
"litellm.proxy.auth.auth_checks._is_api_route_allowed",
new_callable=AsyncMock,
return_value=True,
def _token(metadata=None, team_metadata=None):
return UserAPIKeyAuth(
api_key="test-key",
user_id="user",
team_id="team",
org_id=None,
models=["*"],
metadata=metadata or {},
team_metadata=team_metadata or {},
)
p2 = patch(
"litellm.proxy.auth.auth_checks.vector_store_access_check",
new_callable=AsyncMock,
return_value=None,
)
p3 = patch(
"litellm.proxy.auth.auth_checks._run_project_checks",
new_callable=AsyncMock,
return_value=None,
)
return p1, p2, p3
class TestOpenAIChatCompletionsToolsAllowlist:
"""Tool allowlist enforcement for /v1/chat/completions."""
class TestExtractRequestToolNames:
"""Test tool name extraction per API format."""
@pytest.mark.asyncio
async def test_chat_completions_allowed_tool_passes(self):
"""Request with tools in allowed_tools passes."""
route = "/v1/chat/completions"
request_body = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hi"}],
"tools": [{"type": "function", "function": {"name": "get_weather"}}],
}
token = get_mock_user_token(metadata={"allowed_tools": ["get_weather"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
assert result is True
@pytest.mark.asyncio
async def test_chat_completions_disallowed_tool_raises(self):
"""Request with tool not in allowed_tools raises."""
route = "/v1/chat/completions"
request_body = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hi"}],
"tools": [{"type": "function", "function": {"name": "get_weather"}}],
}
token = get_mock_user_token(metadata={"allowed_tools": ["other_tool"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
with pytest.raises((Exception, ProxyException)) as exc_info:
await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
msg = str(exc_info.value).lower()
assert "tool" in msg or "allowed" in msg
@pytest.mark.asyncio
async def test_chat_completions_legacy_functions_allowed(self):
"""Legacy 'functions' (no tools) with allowed name passes."""
route = "/v1/chat/completions"
request_body = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hi"}],
"functions": [{"name": "get_weather"}],
}
token = get_mock_user_token(metadata={"allowed_tools": ["get_weather"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
assert result is True
@pytest.mark.asyncio
async def test_chat_completions_no_allowlist_passes(self):
"""Request with tools but no metadata.allowed_tools / team_metadata passes."""
route = "/v1/chat/completions"
request_body = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hi"}],
"tools": [{"type": "function", "function": {"name": "get_weather"}}],
}
token = get_mock_user_token(metadata={}, team_metadata={})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
assert result is True
class TestOpenAIResponsesAPIToolsAllowlist:
"""Tool allowlist enforcement for /v1/responses."""
@pytest.mark.asyncio
async def test_responses_function_tool_allowed_passes(self):
"""Responses request with function tool in allowed_tools passes."""
route = "/v1/responses"
request_body = {
"model": "gpt-4",
"input": "What is the weather?",
def test_openai_chat_tools(self):
data = {
"tools": [
{
"type": "function",
"name": "get_current_weather",
"description": "Get current weather",
"parameters": {"type": "object"},
}
],
{"type": "function", "function": {"name": "get_weather"}},
{"type": "function", "function": {"name": "run_sql"}},
]
}
token = get_mock_user_token(metadata={"allowed_tools": ["get_current_weather"]})
request = MockRequest("POST")
assert extract_request_tool_names("/v1/chat/completions", data) == [
"get_weather",
"run_sql",
]
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
assert result is True
def test_openai_chat_functions_legacy(self):
data = {"functions": [{"name": "get_weather"}, {"name": "run_sql"}]}
assert extract_request_tool_names("/v1/chat/completions", data) == [
"get_weather",
"run_sql",
]
@pytest.mark.asyncio
async def test_responses_function_tool_disallowed_raises(self):
"""Responses request with function tool not in allowed_tools raises."""
route = "/v1/responses"
request_body = {
"model": "gpt-4",
"input": "What is the weather?",
def test_openai_responses_function_tools(self):
data = {
"tools": [
{
"type": "function",
"name": "get_current_weather",
"description": "Get current weather",
"parameters": {"type": "object"},
}
],
{"type": "function", "name": "get_current_weather", "description": "x"},
]
}
token = get_mock_user_token(metadata={"allowed_tools": ["other"]})
request = MockRequest("POST")
assert extract_request_tool_names("/v1/responses", data) == [
"get_current_weather"
]
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
with pytest.raises((Exception, ProxyException)) as exc_info:
await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
msg = str(exc_info.value).lower()
assert "tool" in msg or "allowed" in msg
@pytest.mark.asyncio
async def test_responses_mcp_server_allowed_passes(self):
"""Responses request with MCP server in allowed_tools passes."""
route = "/v1/responses"
request_body = {
"model": "gpt-4",
"input": "Hi",
def test_openai_responses_mcp_tools(self):
data = {
"tools": [
{
"type": "mcp",
"server_label": "dmcp",
"server_description": "Example MCP server",
"server_url": "https://example.com",
"require_approval": "never",
}
],
{"type": "mcp", "server_label": "dmcp", "server_url": "http://x"},
]
}
token = get_mock_user_token(metadata={"allowed_tools": ["dmcp"]})
request = MockRequest("POST")
assert extract_request_tool_names("/v1/responses", data) == ["dmcp"]
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
assert result is True
def test_anthropic_tools(self):
data = {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}
assert extract_request_tool_names("/v1/messages", data) == [
"get_weather",
"run_sql",
]
@pytest.mark.asyncio
async def test_responses_mcp_server_disallowed_raises(self):
"""Responses request with MCP server not in allowed_tools raises."""
route = "/v1/responses"
request_body = {
"model": "gpt-4",
"input": "Hi",
"tools": [
{
"type": "mcp",
"server_label": "dmcp",
"server_description": "Example MCP server",
"server_url": "https://example.com",
"require_approval": "never",
}
],
}
token = get_mock_user_token(metadata={"allowed_tools": ["other"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
with pytest.raises((Exception, ProxyException)):
await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
class TestAnthropicMessagesToolsAllowlist:
"""Tool allowlist enforcement for Anthropic /v1/messages."""
@pytest.mark.asyncio
async def test_anthropic_allowed_tool_passes(self):
"""Request with Anthropic-style tools in allowed_tools passes."""
route = "/v1/messages"
request_body = {
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 1024,
"messages": [{"role": "user", "content": "Hi"}],
"tools": [{"name": "get_weather", "description": "Get weather"}],
}
token = get_mock_user_token(metadata={"allowed_tools": ["get_weather"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
assert result is True
@pytest.mark.asyncio
async def test_anthropic_disallowed_tool_raises(self):
"""Request with Anthropic-style tool not in allowed_tools raises."""
route = "/v1/messages"
request_body = {
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 1024,
"messages": [{"role": "user", "content": "Hi"}],
"tools": [{"name": "get_weather", "description": "Get weather"}],
}
token = get_mock_user_token(metadata={"allowed_tools": ["other_tool"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
with pytest.raises((Exception, ProxyException)) as exc_info:
await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
msg = str(exc_info.value).lower()
assert "tool" in msg or "allowed" in msg
class TestGoogleGenerateContentToolsAllowlist:
"""Tool allowlist enforcement for Google generateContent."""
@pytest.mark.asyncio
async def test_google_allowed_tool_passes(self):
"""Request with tools[].functionDeclarations[].name in allowed_tools passes."""
route = "/v1beta/models/gemini-3-flash-preview:generateContent"
request_body = {
"contents": [
{"role": "user", "parts": [{"text": "Schedule a meeting"}]}
],
def test_generate_content_tools(self):
data = {
"tools": [
{
"functionDeclarations": [
{
"name": "schedule_meeting",
"description": "Schedules a meeting",
"parameters": {"type": "object", "properties": {}},
}
{"name": "schedule_meeting", "description": "x"},
]
}
],
},
]
}
token = get_mock_user_token(metadata={"allowed_tools": ["schedule_meeting"]})
request = MockRequest("POST")
assert extract_request_tool_names("/generate_content", data) == [
"schedule_meeting"
]
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
def test_mcp_call_tool_name(self):
data = {"name": "my_tool", "arguments": {}}
assert extract_request_tool_names("/mcp/call_tool", data) == ["my_tool"]
def test_mcp_call_tool_mcp_tool_name(self):
data = {"mcp_tool_name": "other_tool"}
assert extract_request_tool_names("/mcp/call_tool", data) == ["other_tool"]
def test_non_tool_route_returns_empty(self):
data = {"tools": [{"type": "function", "function": {"name": "x"}}]}
assert extract_request_tool_names("/v1/embeddings", data) == []
class TestCheckToolsAllowlist:
"""Test allowlist enforcement in auth (no DB in hot path)."""
@pytest.mark.asyncio
async def test_no_allowlist_passes(self):
token = _token(metadata={}, team_metadata={})
body = {
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
}
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/chat/completions",
)
@pytest.mark.asyncio
async def test_allowed_tool_passes(self):
token = _token(metadata={"allowed_tools": ["get_weather"]})
body = {
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
}
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/chat/completions",
)
@pytest.mark.asyncio
async def test_disallowed_tool_raises(self):
token = _token(metadata={"allowed_tools": ["other_tool"]})
body = {
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
}
with pytest.raises(ProxyException) as exc_info:
await check_tools_allowlist(
request_body=body,
valid_token=token,
request=request,
)
assert result is True
@pytest.mark.asyncio
async def test_google_disallowed_tool_raises(self):
"""Request with tools[].functionDeclarations[].name not in allowed_tools raises."""
route = "/v1beta/models/gemini-3-flash-preview:generateContent"
request_body = {
"contents": [
{"role": "user", "parts": [{"text": "Schedule a meeting"}]}
],
"tools": [
{
"functionDeclarations": [
{
"name": "schedule_meeting",
"description": "Schedules a meeting",
"parameters": {"type": "object", "properties": {}},
}
]
}
],
}
token = get_mock_user_token(metadata={"allowed_tools": ["other_tool"]})
request = MockRequest("POST")
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
with pytest.raises((Exception, ProxyException)) as exc_info:
await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=request,
)
msg = str(exc_info.value).lower()
assert "tool" in msg or "allowed" in msg
# MCP REST tools/call body shape: server_id, name (tool name), arguments.
# See litellm/proxy/_experimental/mcp_server/rest_endpoints.py call_tool_rest_api.
# The exact field for tool name in the request body should match the implementation.
MCP_TOOL_CALL_BODY_ALLOWED = {
"server_id": "srv",
"name": "roll_dice",
"arguments": {},
}
class TestMCPToolCallToolsAllowlist:
"""Test that MCP tool call routes (/mcp/tools/call, /mcp-rest/tools/call) enforce token allowed_tools via common_checks."""
@pytest.mark.asyncio
async def test_mcp_tool_call_allowed_passes(self):
"""Route /mcp-rest/tools/call with tool in token allowed_tools passes common_checks."""
request = MockRequest("POST")
request_body = dict(MCP_TOOL_CALL_BODY_ALLOWED)
valid_token = get_mock_user_token(metadata={"allowed_tools": ["roll_dice"]})
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
result = await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/mcp-rest/tools/call",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=valid_token,
request=request,
route="/v1/chat/completions",
)
assert result is True
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
assert "get_weather" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_mcp_tool_call_disallowed_raises(self):
"""Route /mcp-rest/tools/call with tool not in token allowed_tools raises."""
request = MockRequest("POST")
request_body = dict(MCP_TOOL_CALL_BODY_ALLOWED)
valid_token = get_mock_user_token(metadata={"allowed_tools": ["other"]})
async def test_team_allowlist_used_when_key_empty(self):
token = _token(
metadata={},
team_metadata={"allowed_tools": ["get_weather"]},
)
body = {
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
}
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/chat/completions",
)
p1, p2, p3 = _tools_allowlist_patches()
with p1, p2, p3:
with pytest.raises((Exception, ProxyException)) as exc_info:
await common_checks(
request_body=request_body,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/mcp-rest/tools/call",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=valid_token,
request=request,
)
exc_str = (
getattr(exc_info.value, "message", None) or str(exc_info.value) or ""
).lower()
assert "tool" in exc_str or "allowed" in exc_str
@pytest.mark.asyncio
async def test_key_allowlist_overrides_team(self):
token = _token(
metadata={"allowed_tools": ["get_weather"]},
team_metadata={"allowed_tools": ["other_tool"]},
)
body = {
"tools": [{"type": "function", "function": {"name": "get_weather"}}]
}
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/chat/completions",
)
@pytest.mark.asyncio
async def test_valid_token_none_skips(self):
await check_tools_allowlist(
request_body={"tools": [{"type": "function", "function": {"name": "x"}}]},
valid_token=None,
team_object=None,
route="/v1/chat/completions",
)
@pytest.mark.asyncio
async def test_no_tools_in_body_passes(self):
token = _token(metadata={"allowed_tools": ["get_weather"]})
await check_tools_allowlist(
request_body={"messages": []},
valid_token=token,
team_object=None,
route="/v1/chat/completions",
)

266
ui/litellm-dashboard/src/components/ToolPolicies.tsx
View file

@ -1,14 +1,24 @@
"use client";
import React, { useCallback, useDeferredValue, useEffect, useMemo, useState } from "react";
import { Select, Switch, Tooltip } from "antd";
import { Button, Modal, Select, Switch, 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 { MetricCard } from "./GuardrailsMonitor/MetricCard";
import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking";
import TeamDropdown from "./common_components/team_dropdown";
import {
fetchToolDetail,
fetchToolsList,
updateToolPolicy,
deleteToolPolicyOverride,
ToolRow,
ToolDetailResponse,
ToolPolicyOverrideRow,
} from "./networking";
import { teamListCall, keyListCall } from "./networking";
// --- Date helpers (UTC) for "new tools" counts ---
function getUTCDateKey(date: Date): string {
@ -119,6 +129,16 @@ const PolicySelect: React.FC<{
);
};
interface TeamOption {
team_id: string;
team_alias?: string;
}
interface KeyOption {
token: string;
key_alias?: string;
}
export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
const [tools, setTools] = useState<ToolRow[]>([]);
const [loading, setLoading] = useState(true);
@ -126,6 +146,17 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
const [error, setError] = useState<string | null>(null);
const [saving, setSaving] = useState<string | null>(null);
const [detailModalOpen, setDetailModalOpen] = useState(false);
const [detailToolName, setDetailToolName] = useState<string | null>(null);
const [detail, setDetail] = useState<ToolDetailResponse | null>(null);
const [detailLoading, setDetailLoading] = useState(false);
const [teams, setTeams] = useState<TeamOption[]>([]);
const [keys, setKeys] = useState<KeyOption[]>([]);
const [overrideSaving, setOverrideSaving] = useState(false);
const [blockScope, setBlockScope] = useState<"team" | "key">("team");
const [blockTeamId, setBlockTeamId] = useState<string | null>(null);
const [blockKey, setBlockKey] = useState<{ token: string; key_alias?: string } | null>(null);
const [searchTerm, setSearchTerm] = useState("");
const [sortField, setSortField] = useState<SortField>("created_at");
const [sortOrder, setSortOrder] = useState<"asc" | "desc">("desc");
@ -168,6 +199,9 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
try {
await updateToolPolicy(accessToken, toolName, newPolicy);
setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t)));
if (detailToolName === toolName && detail) {
setDetail((d) => (d ? { ...d, tool: { ...d.tool, call_policy: newPolicy } } : null));
}
} catch (e: any) {
alert(`Failed to update policy: ${e.message}`);
} finally {
@ -175,6 +209,91 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
}
};
const openDetailModal = useCallback(
async (toolName: string) => {
if (!accessToken) return;
setDetailToolName(toolName);
setDetailModalOpen(true);
setDetail(null);
setDetailLoading(true);
setBlockTeamId(null);
setBlockKey(null);
try {
const [detailRes, teamsRes, keysRes] = await Promise.all([
fetchToolDetail(accessToken, toolName),
teamListCall(accessToken, null, null),
keyListCall(accessToken, null, null, null, null, null, 1, 100),
]);
setDetail(detailRes);
const teamsArray = Array.isArray(teamsRes) ? teamsRes : teamsRes?.data ?? [];
setTeams(
teamsArray.map((t: any) => ({ team_id: t.team_id ?? t.id, team_alias: t.team_alias ?? t.team_id }))
);
const keysArray = keysRes?.keys ?? keysRes?.data ?? [];
setKeys(
keysArray.map((k: any) => ({
token: k.token ?? k.api_key ?? k.key_hash ?? "",
key_alias: k.key_alias ?? k.token?.substring?.(0, 8),
}))
);
} catch (e: any) {
setError(e.message ?? "Failed to load tool detail");
} finally {
setDetailLoading(false);
}
},
[accessToken]
);
const closeDetailModal = useCallback(() => {
setDetailModalOpen(false);
setDetailToolName(null);
setDetail(null);
}, []);
const handleAddOverride = useCallback(async () => {
if (!accessToken || !detailToolName) return;
const isTeam = blockScope === "team";
if (isTeam && !blockTeamId) return;
if (!isTeam && !blockKey?.token) return;
setOverrideSaving(true);
try {
await updateToolPolicy(accessToken, detailToolName, "blocked", {
team_id: isTeam ? blockTeamId! : undefined,
key_hash: !isTeam ? blockKey!.token : undefined,
key_alias: !isTeam ? blockKey!.key_alias : undefined,
});
const refreshed = await fetchToolDetail(accessToken, detailToolName);
setDetail(refreshed);
setBlockTeamId(null);
setBlockKey(null);
} catch (e: any) {
alert(`Failed to add override: ${e.message}`);
} finally {
setOverrideSaving(false);
}
}, [accessToken, detailToolName, blockScope, blockTeamId, blockKey]);
const handleRemoveOverride = useCallback(
async (override: ToolPolicyOverrideRow) => {
if (!accessToken || !detailToolName) return;
setOverrideSaving(true);
try {
await deleteToolPolicyOverride(accessToken, detailToolName, {
team_id: override.team_id ?? undefined,
key_hash: override.key_hash ?? undefined,
});
const refreshed = await fetchToolDetail(accessToken, detailToolName);
setDetail(refreshed);
} catch (e: any) {
alert(`Failed to remove override: ${e.message}`);
} finally {
setOverrideSaving(false);
}
},
[accessToken, detailToolName]
);
const handleSortChange = (field: SortField, newState: SortState) => {
if (newState === false) {
setSortField("created_at");
@ -523,11 +642,15 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
<TimeCell utcTime={tool.created_at ?? ""} />
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden">
<Tooltip title={tool.tool_name}>
<span className="font-mono text-xs max-w-[20ch] truncate block font-medium">
{tool.tool_name}
</span>
</Tooltip>
<button
type="button"
onClick={() => openDetailModal(tool.tool_name)}
className="text-left w-full font-mono text-xs max-w-[20ch] truncate block font-medium text-blue-600 hover:text-blue-800 hover:underline focus:outline-none focus:ring-0"
>
<Tooltip title="Click to view details and block for team/key">
<span>{tool.tool_name}</span>
</Tooltip>
</button>
</TableCell>
<TableCell className="py-0.5 max-h-8">
<PolicySelect
@ -594,6 +717,135 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
</div>
)}
</div>
{/* Tool detail modal: view tool, global policy, overrides, block for team/key */}
<Modal
title={detailToolName ? `Tool: ${detailToolName}` : "Tool details"}
open={detailModalOpen}
onCancel={closeDetailModal}
footer={null}
width={640}
destroyOnClose
>
{detailLoading ? (
<p className="text-gray-500 py-4">Loading…</p>
) : detail ? (
<div className="space-y-4">
<div className="flex flex-wrap gap-4 text-sm">
<span>
<strong>Origin:</strong> {detail.tool.origin ?? "—"}
</span>
<span>
<strong># Calls:</strong> {(detail.tool.call_count ?? 0).toLocaleString()}
</span>
</div>
<div>
<strong className="block text-sm text-gray-700 mb-1">Global policy</strong>
<PolicySelect
value={detail.tool.call_policy}
toolName={detail.tool.tool_name}
saving={saving === detail.tool.tool_name}
onChange={handlePolicyChange}
/>
</div>
{detail.overrides.length > 0 && (
<div>
<strong className="block text-sm text-gray-700 mb-2">Blocked for team/key</strong>
<ul className="border rounded-md divide-y divide-gray-100">
{detail.overrides.map((ov) => (
<li
key={ov.override_id}
className="flex items-center justify-between px-3 py-2 text-sm bg-red-50/50"
>
<span>
{ov.team_id ? `Team: ${ov.team_id}` : ""}
{ov.team_id && ov.key_hash ? " · " : ""}
{ov.key_hash ? `Key: ${ov.key_alias || ov.key_hash.substring(0, 8)}` : ""}
{!ov.team_id && !ov.key_hash ? "—" : ""}
</span>
<Button
type="link"
danger
size="small"
disabled={overrideSaving}
onClick={() => handleRemoveOverride(ov)}
>
Remove
</Button>
</li>
))}
</ul>
</div>
)}
<div className="border-t pt-4">
<strong className="block text-sm text-gray-700 mb-2">Block for team or key</strong>
<div className="flex flex-wrap items-end gap-3">
<div className="flex items-center gap-2">
<label className="text-sm">
<input
type="radio"
checked={blockScope === "team"}
onChange={() => setBlockScope("team")}
className="mr-1"
/>
Team
</label>
<label className="text-sm">
<input
type="radio"
checked={blockScope === "key"}
onChange={() => setBlockScope("key")}
className="mr-1"
/>
Key
</label>
</div>
{blockScope === "team" ? (
<div className="min-w-[200px]">
<TeamDropdown
teams={teams}
value={blockTeamId}
onChange={(id) => setBlockTeamId(id ?? null)}
/>
</div>
) : (
<Select
placeholder="Select key"
allowClear
showSearch
optionFilterProp="label"
value={blockKey ? blockKey.token : undefined}
onChange={(token) => {
const k = keys.find((x) => x.token === token);
setBlockKey(k ?? null);
}}
options={keys.map((k) => ({
value: k.token,
label: k.key_alias || k.token?.substring?.(0, 12) || k.token,
}))}
style={{ minWidth: 200 }}
/>
)}
<Button
type="primary"
danger
disabled={
overrideSaving || (blockScope === "team" ? !blockTeamId : !blockKey?.token)
}
loading={overrideSaving}
onClick={handleAddOverride}
>
Block for {blockScope}
</Button>
</div>
</div>
</div>
) : (
<p className="text-gray-500 py-4">No data</p>
)}
</Modal>
</div>
);
};

86
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -9995,19 +9995,91 @@ export const fetchToolsList = async (accessToken: string): Promise<ToolRow[]> =>
return data.tools ?? [];
};
export const updateToolPolicy = async (
export interface ToolPolicyOverrideRow {
override_id: string;
tool_name: string;
team_id?: string | null;
key_hash?: string | null;
call_policy: string;
key_alias?: string | null;
created_at?: string;
updated_at?: string;
}
export interface ToolDetailResponse {
tool: ToolRow;
overrides: ToolPolicyOverrideRow[];
}
export const fetchToolDetail = async (
accessToken: string,
toolName: string,
callPolicy: string
): Promise<ToolRow> => {
const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/policy` : `/v1/tool/policy`;
toolName: string
): Promise<ToolDetailResponse> => {
const encoded = encodeURIComponent(toolName);
const url = proxyBaseUrl
? `${proxyBaseUrl}/v1/tool/${encoded}/detail`
: `/v1/tool/${encoded}/detail`;
const response = await fetch(url, {
method: "POST",
method: "GET",
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();
};
export const updateToolPolicy = async (
accessToken: string,
toolName: string,
callPolicy: string,
options?: { team_id?: string | null; key_hash?: string | null; key_alias?: string | null }
): Promise<ToolRow & { team_id?: string; key_hash?: string }> => {
const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/policy` : `/v1/tool/policy`;
const body: Record<string, string | undefined | null> = {
tool_name: toolName,
call_policy: callPolicy,
};
if (options?.team_id != null) body.team_id = options.team_id || undefined;
if (options?.key_hash != null) body.key_hash = options.key_hash || undefined;
if (options?.key_alias != null) body.key_alias = options.key_alias || undefined;
const response = await fetch(url, {
method: "POST",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify(body),
});
if (!response.ok) {
const errorData = await response.text();
throw new Error(errorData);
}
return response.json();
};
export const deleteToolPolicyOverride = async (
accessToken: string,
toolName: string,
params: { team_id?: string | null; key_hash?: string | null }
): Promise<{ deleted: boolean; tool_name: string }> => {
const encoded = encodeURIComponent(toolName);
const q = new URLSearchParams();
if (params.team_id != null && params.team_id !== "") q.set("team_id", params.team_id);
if (params.key_hash != null && params.key_hash !== "") q.set("key_hash", params.key_hash);
const query = q.toString();
const url = proxyBaseUrl
? `${proxyBaseUrl}/v1/tool/${encoded}/overrides${query ? `?${query}` : ""}`
: `/v1/tool/${encoded}/overrides${query ? `?${query}` : ""}`;
const response = await fetch(url, {
method: "DELETE",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
},
});
if (!response.ok) {
const errorData = await response.text();