diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225220000_add_tool_policy_override_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225220000_add_tool_policy_override_table/migration.sql new file mode 100644 index 00000000000..642508b7b20 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260225220000_add_tool_policy_override_table/migration.sql @@ -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"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 440c9c1d829..e48c0fe3027 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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()) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 98650a238e9..a6df346e8a8 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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], diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 7106c207bd6..a7982cb606e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -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 [] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 683e165c315..0658a953318 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -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], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 6b092911d3c..7c3354cf88e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -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]], diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 88cfa5f6c11..65ea9cb42d8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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"] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 500a39d9455..7b44699825c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 332607f0305..a0dffefdc59 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -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 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 87558566c42..213b248929e 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 @@ -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, + ) diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py new file mode 100644 index 00000000000..db24fa2277c --- /dev/null +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -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 [] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d2312a00c3b..c3d58df83ea 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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_ 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_ 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 diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 89880c9a4ec..bc8fc9165bb 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -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)) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 440c9c1d829..e48c0fe3027 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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()) diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 8704ff27759..90fc4a6ec7d 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -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) diff --git a/scripts/test_tool_allowlist_script.py b/scripts/test_tool_allowlist_script.py new file mode 100644 index 00000000000..75a50d09b84 --- /dev/null +++ b/scripts/test_tool_allowlist_script.py @@ -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() diff --git a/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py index 09d56be30ba..4adc5acde8b 100644 --- a/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py +++ b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py @@ -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", + ) diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index aa005495575..f39b21aed76 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -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 = ({ accessToken }) => { const [tools, setTools] = useState([]); const [loading, setLoading] = useState(true); @@ -126,6 +146,17 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { const [error, setError] = useState(null); const [saving, setSaving] = useState(null); + const [detailModalOpen, setDetailModalOpen] = useState(false); + const [detailToolName, setDetailToolName] = useState(null); + const [detail, setDetail] = useState(null); + const [detailLoading, setDetailLoading] = useState(false); + const [teams, setTeams] = useState([]); + const [keys, setKeys] = useState([]); + const [overrideSaving, setOverrideSaving] = useState(false); + const [blockScope, setBlockScope] = useState<"team" | "key">("team"); + const [blockTeamId, setBlockTeamId] = useState(null); + const [blockKey, setBlockKey] = useState<{ token: string; key_alias?: string } | null>(null); + const [searchTerm, setSearchTerm] = useState(""); const [sortField, setSortField] = useState("created_at"); const [sortOrder, setSortOrder] = useState<"asc" | "desc">("desc"); @@ -168,6 +199,9 @@ export const ToolPolicies: React.FC = ({ 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 = ({ 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 = ({ accessToken }) => { - - - {tool.tool_name} - - + = ({ accessToken }) => { )} + + {/* Tool detail modal: view tool, global policy, overrides, block for team/key */} + + {detailLoading ? ( +

Loading…

+ ) : detail ? ( +
+
+ + Origin: {detail.tool.origin ?? "—"} + + + # Calls: {(detail.tool.call_count ?? 0).toLocaleString()} + +
+
+ Global policy + +
+ + {detail.overrides.length > 0 && ( +
+ Blocked for team/key +
    + {detail.overrides.map((ov) => ( +
  • + + {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 ? "—" : ""} + + +
    • + ))} +
+
+ )} + +
+ Block for team or key +
+
+ + +
+ {blockScope === "team" ? ( +
+ setBlockTeamId(id ?? null)} + /> +
+ ) : ( +