confirm mcp tool executions from untrusted servers (per "trust" setting) (#631)
This commit is contained in:
parent
a60e51f44d
commit
2e57989aec
|
@ -12,6 +12,7 @@ import {
|
||||||
ToolCallConfirmationDetails,
|
ToolCallConfirmationDetails,
|
||||||
ToolConfirmationOutcome,
|
ToolConfirmationOutcome,
|
||||||
ToolExecuteConfirmationDetails,
|
ToolExecuteConfirmationDetails,
|
||||||
|
ToolMcpConfirmationDetails,
|
||||||
} from '@gemini-code/server';
|
} from '@gemini-code/server';
|
||||||
import {
|
import {
|
||||||
RadioButtonSelect,
|
RadioButtonSelect,
|
||||||
|
@ -64,7 +65,7 @@ export const ToolConfirmationMessage: React.FC<
|
||||||
},
|
},
|
||||||
{ label: 'No (esc)', value: ToolConfirmationOutcome.Cancel },
|
{ label: 'No (esc)', value: ToolConfirmationOutcome.Cancel },
|
||||||
);
|
);
|
||||||
} else {
|
} else if (confirmationDetails.type === 'exec') {
|
||||||
const executionProps =
|
const executionProps =
|
||||||
confirmationDetails as ToolExecuteConfirmationDetails;
|
confirmationDetails as ToolExecuteConfirmationDetails;
|
||||||
|
|
||||||
|
@ -88,6 +89,33 @@ export const ToolConfirmationMessage: React.FC<
|
||||||
},
|
},
|
||||||
{ label: 'No (esc)', value: ToolConfirmationOutcome.Cancel },
|
{ label: 'No (esc)', value: ToolConfirmationOutcome.Cancel },
|
||||||
);
|
);
|
||||||
|
} else {
|
||||||
|
// mcp tool confirmation
|
||||||
|
const mcpProps = confirmationDetails as ToolMcpConfirmationDetails;
|
||||||
|
|
||||||
|
bodyContent = (
|
||||||
|
<Box flexDirection="column" paddingX={1} marginLeft={1}>
|
||||||
|
<Text color={Colors.AccentCyan}>MCP Server: {mcpProps.serverName}</Text>
|
||||||
|
<Text color={Colors.AccentCyan}>Tool: {mcpProps.toolName}</Text>
|
||||||
|
</Box>
|
||||||
|
);
|
||||||
|
|
||||||
|
question = `Allow execution of MCP tool "${mcpProps.toolName}" from server "${mcpProps.serverName}"?`;
|
||||||
|
options.push(
|
||||||
|
{
|
||||||
|
label: 'Yes, allow once',
|
||||||
|
value: ToolConfirmationOutcome.ProceedOnce,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: `Yes, always allow tool "${mcpProps.toolName}" from server "${mcpProps.serverName}"`,
|
||||||
|
value: ToolConfirmationOutcome.ProceedAlwaysTool, // Cast until types are updated
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: `Yes, always allow all tools from server "${mcpProps.serverName}"`,
|
||||||
|
value: ToolConfirmationOutcome.ProceedAlwaysServer,
|
||||||
|
},
|
||||||
|
{ label: 'No (esc)', value: ToolConfirmationOutcome.Cancel },
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
return (
|
return (
|
||||||
|
|
|
@ -33,6 +33,7 @@ export class MCPServerConfig {
|
||||||
readonly url?: string,
|
readonly url?: string,
|
||||||
// Common
|
// Common
|
||||||
readonly timeout?: number,
|
readonly timeout?: number,
|
||||||
|
readonly trust?: boolean,
|
||||||
) {}
|
) {}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -134,11 +134,13 @@ async function connectAndDiscover(
|
||||||
toolRegistry.registerTool(
|
toolRegistry.registerTool(
|
||||||
new DiscoveredMCPTool(
|
new DiscoveredMCPTool(
|
||||||
mcpClient,
|
mcpClient,
|
||||||
|
mcpServerName,
|
||||||
toolNameForModel,
|
toolNameForModel,
|
||||||
tool.description ?? '',
|
tool.description ?? '',
|
||||||
tool.inputSchema,
|
tool.inputSchema,
|
||||||
tool.name,
|
tool.name,
|
||||||
mcpServerConfig.timeout,
|
mcpServerConfig.timeout,
|
||||||
|
mcpServerConfig.trust,
|
||||||
),
|
),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
|
@ -55,6 +55,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
it('should set properties correctly and augment description', () => {
|
it('should set properties correctly and augment description', () => {
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
baseDescription,
|
baseDescription,
|
||||||
inputSchema,
|
inputSchema,
|
||||||
|
@ -78,6 +79,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
const customTimeout = 5000;
|
const customTimeout = 5000;
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
baseDescription,
|
baseDescription,
|
||||||
inputSchema,
|
inputSchema,
|
||||||
|
@ -92,6 +94,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
it('should call mcpClient.callTool with correct parameters and default timeout', async () => {
|
it('should call mcpClient.callTool with correct parameters and default timeout', async () => {
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
baseDescription,
|
baseDescription,
|
||||||
inputSchema,
|
inputSchema,
|
||||||
|
@ -122,6 +125,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
const customTimeout = 15000;
|
const customTimeout = 15000;
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
baseDescription,
|
baseDescription,
|
||||||
inputSchema,
|
inputSchema,
|
||||||
|
@ -146,6 +150,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
it('should propagate rejection if mcpClient.callTool rejects', async () => {
|
it('should propagate rejection if mcpClient.callTool rejects', async () => {
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
baseDescription,
|
baseDescription,
|
||||||
inputSchema,
|
inputSchema,
|
||||||
|
|
|
@ -5,20 +5,30 @@
|
||||||
*/
|
*/
|
||||||
|
|
||||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
||||||
import { BaseTool, ToolResult } from './tools.js';
|
import {
|
||||||
|
BaseTool,
|
||||||
|
ToolResult,
|
||||||
|
ToolCallConfirmationDetails,
|
||||||
|
ToolConfirmationOutcome,
|
||||||
|
ToolMcpConfirmationDetails,
|
||||||
|
} from './tools.js';
|
||||||
|
|
||||||
type ToolParams = Record<string, unknown>;
|
type ToolParams = Record<string, unknown>;
|
||||||
|
|
||||||
export const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
|
export const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
|
||||||
|
|
||||||
export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
|
export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
|
||||||
|
private static readonly whitelist: Set<string> = new Set();
|
||||||
|
|
||||||
constructor(
|
constructor(
|
||||||
private readonly mcpClient: Client,
|
private readonly mcpClient: Client,
|
||||||
|
private readonly serverName: string, // Added for server identification
|
||||||
readonly name: string,
|
readonly name: string,
|
||||||
readonly description: string,
|
readonly description: string,
|
||||||
readonly parameterSchema: Record<string, unknown>,
|
readonly parameterSchema: Record<string, unknown>,
|
||||||
readonly serverToolName: string,
|
readonly serverToolName: string,
|
||||||
readonly timeout?: number,
|
readonly timeout?: number,
|
||||||
|
readonly trust?: boolean,
|
||||||
) {
|
) {
|
||||||
description += `
|
description += `
|
||||||
|
|
||||||
|
@ -37,6 +47,41 @@ Returns the MCP server response as a json string.
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async shouldConfirmExecute(
|
||||||
|
_params: ToolParams,
|
||||||
|
_abortSignal: AbortSignal,
|
||||||
|
): Promise<ToolCallConfirmationDetails | false> {
|
||||||
|
const serverWhitelistKey = this.serverName;
|
||||||
|
const toolWhitelistKey = `${this.serverName}.${this.serverToolName}`;
|
||||||
|
|
||||||
|
if (this.trust) {
|
||||||
|
return false; // server is trusted, no confirmation needed
|
||||||
|
}
|
||||||
|
|
||||||
|
if (
|
||||||
|
DiscoveredMCPTool.whitelist.has(serverWhitelistKey) ||
|
||||||
|
DiscoveredMCPTool.whitelist.has(toolWhitelistKey)
|
||||||
|
) {
|
||||||
|
return false; // server and/or tool already whitelisted
|
||||||
|
}
|
||||||
|
|
||||||
|
const confirmationDetails: ToolMcpConfirmationDetails = {
|
||||||
|
type: 'mcp',
|
||||||
|
title: 'Confirm MCP Tool Execution',
|
||||||
|
serverName: this.serverName,
|
||||||
|
toolName: this.serverToolName,
|
||||||
|
toolDisplayName: this.name,
|
||||||
|
onConfirm: async (outcome: ToolConfirmationOutcome) => {
|
||||||
|
if (outcome === ToolConfirmationOutcome.ProceedAlwaysServer) {
|
||||||
|
DiscoveredMCPTool.whitelist.add(serverWhitelistKey);
|
||||||
|
} else if (outcome === ToolConfirmationOutcome.ProceedAlwaysTool) {
|
||||||
|
DiscoveredMCPTool.whitelist.add(toolWhitelistKey);
|
||||||
|
}
|
||||||
|
},
|
||||||
|
};
|
||||||
|
return confirmationDetails;
|
||||||
|
}
|
||||||
|
|
||||||
async execute(params: ToolParams): Promise<ToolResult> {
|
async execute(params: ToolParams): Promise<ToolResult> {
|
||||||
const result = await this.mcpClient.callTool(
|
const result = await this.mcpClient.callTool(
|
||||||
{
|
{
|
||||||
|
|
|
@ -729,6 +729,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
it('constructor should set up properties correctly and enhance description', () => {
|
it('constructor should set up properties correctly and enhance description', () => {
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
toolDescription,
|
toolDescription,
|
||||||
toolInputSchema,
|
toolInputSchema,
|
||||||
|
@ -744,6 +745,7 @@ describe('DiscoveredMCPTool', () => {
|
||||||
it('execute should call mcpClient.callTool with correct params and return serialized result', async () => {
|
it('execute should call mcpClient.callTool with correct params and return serialized result', async () => {
|
||||||
const tool = new DiscoveredMCPTool(
|
const tool = new DiscoveredMCPTool(
|
||||||
mockMcpClient,
|
mockMcpClient,
|
||||||
|
'mock-mcp-server',
|
||||||
toolName,
|
toolName,
|
||||||
toolDescription,
|
toolDescription,
|
||||||
toolInputSchema,
|
toolInputSchema,
|
||||||
|
|
|
@ -212,12 +212,24 @@ export interface ToolExecuteConfirmationDetails {
|
||||||
rootCommand: string;
|
rootCommand: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface ToolMcpConfirmationDetails {
|
||||||
|
type: 'mcp';
|
||||||
|
title: string;
|
||||||
|
serverName: string;
|
||||||
|
toolName: string;
|
||||||
|
toolDisplayName: string;
|
||||||
|
onConfirm: (outcome: ToolConfirmationOutcome) => Promise<void> | void;
|
||||||
|
}
|
||||||
|
|
||||||
export type ToolCallConfirmationDetails =
|
export type ToolCallConfirmationDetails =
|
||||||
| ToolEditConfirmationDetails
|
| ToolEditConfirmationDetails
|
||||||
| ToolExecuteConfirmationDetails;
|
| ToolExecuteConfirmationDetails
|
||||||
|
| ToolMcpConfirmationDetails;
|
||||||
|
|
||||||
export enum ToolConfirmationOutcome {
|
export enum ToolConfirmationOutcome {
|
||||||
ProceedOnce,
|
ProceedOnce = 'proceed_once',
|
||||||
ProceedAlways,
|
ProceedAlways = 'proceed_always',
|
||||||
Cancel,
|
ProceedAlwaysServer = 'proceed_always_server',
|
||||||
|
ProceedAlwaysTool = 'proceed_always_tool',
|
||||||
|
Cancel = 'cancel',
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue