diff --git a/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py new file mode 100644 index 00000000000..09d56be30ba --- /dev/null +++ b/tests/test_litellm/proxy/test_tools_allowlist_enforcement.py @@ -0,0 +1,548 @@ +""" +Tests for tool allowlist enforcement by team/key (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. +""" + +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 + + +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, + ) + 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.""" + + @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?", + "tools": [ + { + "type": "function", + "name": "get_current_weather", + "description": "Get current weather", + "parameters": {"type": "object"}, + } + ], + } + token = get_mock_user_token(metadata={"allowed_tools": ["get_current_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_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?", + "tools": [ + { + "type": "function", + "name": "get_current_weather", + "description": "Get current weather", + "parameters": {"type": "object"}, + } + ], + } + 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)) 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", + "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": ["dmcp"]}) + 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_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"}]} + ], + "tools": [ + { + "functionDeclarations": [ + { + "name": "schedule_meeting", + "description": "Schedules a meeting", + "parameters": {"type": "object", "properties": {}}, + } + ] + } + ], + } + token = get_mock_user_token(metadata={"allowed_tools": ["schedule_meeting"]}) + 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_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, + ) + assert result is True + + @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"]}) + + 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 diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index 6e62370b0c6..aa005495575 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -1,14 +1,41 @@ "use client"; -import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; +import React, { useCallback, useDeferredValue, useEffect, useMemo, useState } from "react"; import { 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"; +// --- Date helpers (UTC) for "new tools" counts --- +function getUTCDateKey(date: Date): string { + return `${date.getUTCFullYear()}-${String(date.getUTCMonth() + 1).padStart(2, "0")}-${String(date.getUTCDate()).padStart(2, "0")}`; +} + +function isCreatedInUTCDay(createdAt: string | undefined, utcDateKey: string): boolean { + if (!createdAt) return false; + try { + const d = new Date(createdAt); + return getUTCDateKey(d) === utcDateKey; + } catch { + return false; + } +} + +function countToolsInUTCDay(tools: ToolRow[], utcDateKey: string): number { + return tools.filter((t) => isCreatedInUTCDay(t.created_at, utcDateKey)).length; +} + +function getTrendSubtitle(newToday: number, newYesterday: number): string | undefined { + const diff = newToday - newYesterday; + if (diff === 0) return undefined; + if (diff > 0) return `+${diff} since yesterday`; + return `${diff} since yesterday`; +} + const POLICY_OPTIONS = [ { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, @@ -197,6 +224,41 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { }, ]; + // Derived counts for summary cards and "Needs Review" (UTC today/yesterday) + const { newToday, newYesterday, trendSubtitle, totalTools, blockedCount, activeTeamsCount, needsReviewTools } = + useMemo(() => { + const now = new Date(); + const todayKey = getUTCDateKey(now); + const yesterday = new Date(now); + yesterday.setUTCDate(yesterday.getUTCDate() - 1); + const yesterdayKey = getUTCDateKey(yesterday); + + const newToday = countToolsInUTCDay(tools, todayKey); + const newYesterday = countToolsInUTCDay(tools, yesterdayKey); + const trendSubtitle = getTrendSubtitle(newToday, newYesterday); + + const totalTools = tools.length; + const blockedCount = tools.filter((t) => t.call_policy === "blocked").length; + const activeTeamsCount = new Set(tools.map((t) => t.team_id).filter(Boolean)).size; + + // New in period (today) and not yet decided — untrusted or dual_llm + const needsReviewTools = tools.filter( + (t) => + isCreatedInUTCDay(t.created_at, todayKey) && + (t.call_policy === "untrusted" || t.call_policy === "dual_llm") + ); + + return { + newToday, + newYesterday, + trendSubtitle, + totalTools, + blockedCount, + activeTeamsCount, + needsReviewTools, + }; + }, [tools]); + const SortHeader = ({ label, field }: { label: string; field: SortField }) => (
{label} @@ -235,9 +297,76 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { const totalPages = Math.max(1, Math.ceil(sorted.length / pageSize)); const paginated = sorted.slice((currentPage - 1) * pageSize, currentPage * pageSize); + const scrollToToolRow = (toolId: string) => { + const idx = sorted.findIndex((t) => t.tool_id === toolId); + if (idx >= 0) { + const page = Math.floor(idx / pageSize) + 1; + if (page !== currentPage) setCurrentPage(page); + // Scroll after a short delay so the table has re-rendered with the new page + requestAnimationFrame(() => { + setTimeout(() => { + document.getElementById(`tool-row-${toolId}`)?.scrollIntoView({ behavior: "smooth", block: "center" }); + }, 100); + }); + } + }; + return (

Tool Policies

+ + {/* Summary cards */} +
+ + + + } + /> + + 0 ? "text-red-600" : undefined} + /> + 0 ? activeTeamsCount : "—"} /> +
+ + {/* Needs Review */} + {needsReviewTools.length > 0 && ( +
+

Needs Review

+

+ {needsReviewTools.length} new tool{needsReviewTools.length !== 1 ? "s" : ""} discovered that require + policy decisions. +

+
+ {needsReviewTools.map((t) => ( + + + {t.tool_name} + + + + ))} +
+
+ )} +
{/* Toolbar */}
@@ -389,7 +518,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { ) : ( paginated.map((tool) => ( - +