fix(add-new-tiles-to-tool-policies): allow developer to see what's available

This commit is contained in:
Krrish Dholakia 2026-02-25 21:53:42 -08:00
parent 4b4018b0b2
commit 2487943846
2 changed files with 679 additions and 2 deletions

View file

@ -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

View file

@ -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<ToolPoliciesProps> = ({ 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 }) => (
<div className="flex items-center gap-1">
<span>{label}</span>
@ -235,9 +297,76 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ 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 (
<div className="p-6 w-full">
<h1 className="text-2xl font-semibold text-gray-900 mb-6">Tool Policies</h1>
{/* Summary cards */}
<div className="grid grid-cols-2 lg:grid-cols-4 gap-4 mb-6">
<MetricCard
label="New Today"
value={newToday}
valueColor="text-green-600"
subtitle={trendSubtitle}
icon={
<svg className="w-4 h-4 text-green-500" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M13 7h8m0 0v8m0-8l-8 8-4-4-6 6" />
</svg>
}
/>
<MetricCard label="Total Tools Discovered" value={totalTools} />
<MetricCard
label="Blocked Tools"
value={blockedCount}
valueColor={blockedCount > 0 ? "text-red-600" : undefined}
/>
<MetricCard label="Active Teams" value={activeTeamsCount > 0 ? activeTeamsCount : "—"} />
</div>
{/* Needs Review */}
{needsReviewTools.length > 0 && (
<div className="bg-amber-50 border border-amber-200 rounded-lg p-4 mb-6">
<h2 className="text-sm font-semibold text-amber-900 mb-1">Needs Review</h2>
<p className="text-sm text-amber-800 mb-3">
{needsReviewTools.length} new tool{needsReviewTools.length !== 1 ? "s" : ""} discovered that require
policy decisions.
</p>
<div className="flex flex-wrap gap-2">
{needsReviewTools.map((t) => (
<span
key={t.tool_id}
className="inline-flex items-center gap-2 px-3 py-1.5 bg-white border border-amber-200 rounded-md text-sm"
>
<span className="font-mono text-amber-900 truncate max-w-[200px]" title={t.tool_name}>
{t.tool_name}
</span>
<button
type="button"
onClick={() => scrollToToolRow(t.tool_id)}
className="text-amber-700 hover:text-amber-900 font-medium text-xs whitespace-nowrap"
>
Review
</button>
</span>
))}
</div>
</div>
)}
<div className="bg-white rounded-lg shadow w-full max-w-full box-border">
{/* Toolbar */}
<div className="border-b px-6 py-4 w-full max-w-full box-border">
@ -389,7 +518,7 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
</TableRow>
) : (
paginated.map((tool) => (
<TableRow key={tool.tool_id} className="h-8 hover:bg-gray-50">
<TableRow key={tool.tool_id} id={`tool-row-${tool.tool_id}`} className="h-8 hover:bg-gray-50">
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<TimeCell utcTime={tool.created_at ?? ""} />
</TableCell>