655 lines
21 KiB
TypeScript
655 lines
21 KiB
TypeScript
/**
|
|
* @license
|
|
* Copyright 2025 Google LLC
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*/
|
|
|
|
/* eslint-disable @typescript-eslint/no-explicit-any */
|
|
import {
|
|
describe,
|
|
it,
|
|
expect,
|
|
vi,
|
|
beforeEach,
|
|
afterEach,
|
|
Mocked,
|
|
} from 'vitest';
|
|
import { DiscoveredMCPTool, generateValidName } from './mcp-tool.js'; // Added getStringifiedResultForDisplay
|
|
import { ToolResult, ToolConfirmationOutcome } from './tools.js'; // Added ToolConfirmationOutcome
|
|
import { CallableTool, Part } from '@google/genai';
|
|
|
|
// Mock @google/genai mcpToTool and CallableTool
|
|
// We only need to mock the parts of CallableTool that DiscoveredMCPTool uses.
|
|
const mockCallTool = vi.fn();
|
|
const mockToolMethod = vi.fn();
|
|
|
|
const mockCallableToolInstance: Mocked<CallableTool> = {
|
|
tool: mockToolMethod as any, // Not directly used by DiscoveredMCPTool instance methods
|
|
callTool: mockCallTool as any,
|
|
// Add other methods if DiscoveredMCPTool starts using them
|
|
};
|
|
|
|
describe('generateValidName', () => {
|
|
it('should return a valid name for a simple function', () => {
|
|
expect(generateValidName('myFunction')).toBe('myFunction');
|
|
});
|
|
|
|
it('should replace invalid characters with underscores', () => {
|
|
expect(generateValidName('invalid-name with spaces')).toBe(
|
|
'invalid-name_with_spaces',
|
|
);
|
|
});
|
|
|
|
it('should truncate long names', () => {
|
|
expect(generateValidName('x'.repeat(80))).toBe(
|
|
'xxxxxxxxxxxxxxxxxxxxxxxxxxxx___xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
|
|
);
|
|
});
|
|
|
|
it('should handle names with only invalid characters', () => {
|
|
expect(generateValidName('!@#$%^&*()')).toBe('__________');
|
|
});
|
|
|
|
it('should handle names that are exactly 63 characters long', () => {
|
|
expect(generateValidName('a'.repeat(63)).length).toBe(63);
|
|
});
|
|
|
|
it('should handle names that are exactly 64 characters long', () => {
|
|
expect(generateValidName('a'.repeat(64)).length).toBe(63);
|
|
});
|
|
|
|
it('should handle names that are longer than 64 characters', () => {
|
|
expect(generateValidName('a'.repeat(80)).length).toBe(63);
|
|
});
|
|
});
|
|
|
|
describe('DiscoveredMCPTool', () => {
|
|
const serverName = 'mock-mcp-server';
|
|
const serverToolName = 'actual-server-tool-name';
|
|
const baseDescription = 'A test MCP tool.';
|
|
const inputSchema: Record<string, unknown> = {
|
|
type: 'object' as const,
|
|
properties: { param: { type: 'string' } },
|
|
required: ['param'],
|
|
};
|
|
|
|
let tool: DiscoveredMCPTool;
|
|
|
|
beforeEach(() => {
|
|
mockCallTool.mockClear();
|
|
mockToolMethod.mockClear();
|
|
tool = new DiscoveredMCPTool(
|
|
mockCallableToolInstance,
|
|
serverName,
|
|
serverToolName,
|
|
baseDescription,
|
|
inputSchema,
|
|
);
|
|
// Clear allowlist before each relevant test, especially for shouldConfirmExecute
|
|
const invocation = tool.build({}) as any;
|
|
invocation.constructor.allowlist.clear();
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
describe('constructor', () => {
|
|
it('should set properties correctly', () => {
|
|
expect(tool.name).toBe(serverToolName);
|
|
expect(tool.schema.name).toBe(serverToolName);
|
|
expect(tool.schema.description).toBe(baseDescription);
|
|
expect(tool.schema.parameters).toBeUndefined();
|
|
expect(tool.schema.parametersJsonSchema).toEqual(inputSchema);
|
|
expect(tool.serverToolName).toBe(serverToolName);
|
|
expect(tool.timeout).toBeUndefined();
|
|
});
|
|
|
|
it('should accept and store a custom timeout', () => {
|
|
const customTimeout = 5000;
|
|
const toolWithTimeout = new DiscoveredMCPTool(
|
|
mockCallableToolInstance,
|
|
serverName,
|
|
serverToolName,
|
|
baseDescription,
|
|
inputSchema,
|
|
customTimeout,
|
|
);
|
|
expect(toolWithTimeout.timeout).toBe(customTimeout);
|
|
});
|
|
});
|
|
|
|
describe('execute', () => {
|
|
it('should call mcpTool.callTool with correct parameters and format display output', async () => {
|
|
const params = { param: 'testValue' };
|
|
const mockToolSuccessResultObject = {
|
|
success: true,
|
|
details: 'executed',
|
|
};
|
|
const mockFunctionResponseContent = [
|
|
{
|
|
type: 'text',
|
|
text: JSON.stringify(mockToolSuccessResultObject),
|
|
},
|
|
];
|
|
const mockMcpToolResponseParts: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: { content: mockFunctionResponseContent },
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(mockMcpToolResponseParts);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult: ToolResult = await invocation.execute(
|
|
new AbortController().signal,
|
|
);
|
|
|
|
expect(mockCallTool).toHaveBeenCalledWith([
|
|
{ name: serverToolName, args: params },
|
|
]);
|
|
|
|
const stringifiedResponseContent = JSON.stringify(
|
|
mockToolSuccessResultObject,
|
|
);
|
|
expect(toolResult.llmContent).toEqual([
|
|
{ text: stringifiedResponseContent },
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe(stringifiedResponseContent);
|
|
});
|
|
|
|
it('should handle empty result from getStringifiedResultForDisplay', async () => {
|
|
const params = { param: 'testValue' };
|
|
const mockMcpToolResponsePartsEmpty: Part[] = [];
|
|
mockCallTool.mockResolvedValue(mockMcpToolResponsePartsEmpty);
|
|
const invocation = tool.build(params);
|
|
const toolResult: ToolResult = await invocation.execute(
|
|
new AbortController().signal,
|
|
);
|
|
expect(toolResult.returnDisplay).toBe('```json\n[]\n```');
|
|
expect(toolResult.llmContent).toEqual([
|
|
{ text: '[Error: Could not parse tool response]' },
|
|
]);
|
|
});
|
|
|
|
it('should propagate rejection if mcpTool.callTool rejects', async () => {
|
|
const params = { param: 'failCase' };
|
|
const expectedError = new Error('MCP call failed');
|
|
mockCallTool.mockRejectedValue(expectedError);
|
|
|
|
const invocation = tool.build(params);
|
|
await expect(
|
|
invocation.execute(new AbortController().signal),
|
|
).rejects.toThrow(expectedError);
|
|
});
|
|
|
|
it('should handle a simple text response correctly', async () => {
|
|
const params = { query: 'test' };
|
|
const successMessage = 'This is a success message.';
|
|
|
|
// Simulate the response from the GenAI SDK, which wraps the MCP
|
|
// response in a functionResponse Part.
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
// The `content` array contains MCP ContentBlocks.
|
|
content: [{ type: 'text', text: successMessage }],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
// 1. Assert that the llmContent sent to the scheduler is a clean Part array.
|
|
expect(toolResult.llmContent).toEqual([{ text: successMessage }]);
|
|
|
|
// 2. Assert that the display output is the simple text message.
|
|
expect(toolResult.returnDisplay).toBe(successMessage);
|
|
|
|
// 3. Verify that the underlying callTool was made correctly.
|
|
expect(mockCallTool).toHaveBeenCalledWith([
|
|
{ name: serverToolName, args: params },
|
|
]);
|
|
});
|
|
|
|
it('should handle an AudioBlock response', async () => {
|
|
const params = { action: 'play' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{
|
|
type: 'audio',
|
|
data: 'BASE64_AUDIO_DATA',
|
|
mimeType: 'audio/mp3',
|
|
},
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([
|
|
{
|
|
text: `[Tool '${serverToolName}' provided the following audio data with mime-type: audio/mp3]`,
|
|
},
|
|
{
|
|
inlineData: {
|
|
mimeType: 'audio/mp3',
|
|
data: 'BASE64_AUDIO_DATA',
|
|
},
|
|
},
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe('[Audio: audio/mp3]');
|
|
});
|
|
|
|
it('should handle a ResourceLinkBlock response', async () => {
|
|
const params = { resource: 'get' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{
|
|
type: 'resource_link',
|
|
uri: 'file:///path/to/thing',
|
|
name: 'resource-name',
|
|
title: 'My Resource',
|
|
},
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([
|
|
{
|
|
text: 'Resource Link: My Resource at file:///path/to/thing',
|
|
},
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe(
|
|
'[Link to My Resource: file:///path/to/thing]',
|
|
);
|
|
});
|
|
|
|
it('should handle an embedded text ResourceBlock response', async () => {
|
|
const params = { resource: 'get' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{
|
|
type: 'resource',
|
|
resource: {
|
|
uri: 'file:///path/to/text.txt',
|
|
text: 'This is the text content.',
|
|
mimeType: 'text/plain',
|
|
},
|
|
},
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([
|
|
{ text: 'This is the text content.' },
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe('This is the text content.');
|
|
});
|
|
|
|
it('should handle an embedded binary ResourceBlock response', async () => {
|
|
const params = { resource: 'get' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{
|
|
type: 'resource',
|
|
resource: {
|
|
uri: 'file:///path/to/data.bin',
|
|
blob: 'BASE64_BINARY_DATA',
|
|
mimeType: 'application/octet-stream',
|
|
},
|
|
},
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([
|
|
{
|
|
text: `[Tool '${serverToolName}' provided the following embedded resource with mime-type: application/octet-stream]`,
|
|
},
|
|
{
|
|
inlineData: {
|
|
mimeType: 'application/octet-stream',
|
|
data: 'BASE64_BINARY_DATA',
|
|
},
|
|
},
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe(
|
|
'[Embedded Resource: application/octet-stream]',
|
|
);
|
|
});
|
|
|
|
it('should handle a mix of content block types', async () => {
|
|
const params = { action: 'complex' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{ type: 'text', text: 'First part.' },
|
|
{
|
|
type: 'image',
|
|
data: 'BASE64_IMAGE_DATA',
|
|
mimeType: 'image/jpeg',
|
|
},
|
|
{ type: 'text', text: 'Second part.' },
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([
|
|
{ text: 'First part.' },
|
|
{
|
|
text: `[Tool '${serverToolName}' provided the following image data with mime-type: image/jpeg]`,
|
|
},
|
|
{
|
|
inlineData: {
|
|
mimeType: 'image/jpeg',
|
|
data: 'BASE64_IMAGE_DATA',
|
|
},
|
|
},
|
|
{ text: 'Second part.' },
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe(
|
|
'First part.\n[Image: image/jpeg]\nSecond part.',
|
|
);
|
|
});
|
|
|
|
it('should ignore unknown content block types', async () => {
|
|
const params = { action: 'test' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{ type: 'text', text: 'Valid part.' },
|
|
{ type: 'future_block', data: 'some-data' },
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([{ text: 'Valid part.' }]);
|
|
expect(toolResult.returnDisplay).toBe(
|
|
'Valid part.\n[Unknown content type: future_block]',
|
|
);
|
|
});
|
|
|
|
it('should handle a complex mix of content block types', async () => {
|
|
const params = { action: 'super-complex' };
|
|
const sdkResponse: Part[] = [
|
|
{
|
|
functionResponse: {
|
|
name: serverToolName,
|
|
response: {
|
|
content: [
|
|
{ type: 'text', text: 'Here is a resource.' },
|
|
{
|
|
type: 'resource_link',
|
|
uri: 'file:///path/to/resource',
|
|
name: 'resource-name',
|
|
title: 'My Resource',
|
|
},
|
|
{
|
|
type: 'resource',
|
|
resource: {
|
|
uri: 'file:///path/to/text.txt',
|
|
text: 'Embedded text content.',
|
|
mimeType: 'text/plain',
|
|
},
|
|
},
|
|
{
|
|
type: 'image',
|
|
data: 'BASE64_IMAGE_DATA',
|
|
mimeType: 'image/jpeg',
|
|
},
|
|
],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
mockCallTool.mockResolvedValue(sdkResponse);
|
|
|
|
const invocation = tool.build(params);
|
|
const toolResult = await invocation.execute(new AbortController().signal);
|
|
|
|
expect(toolResult.llmContent).toEqual([
|
|
{ text: 'Here is a resource.' },
|
|
{
|
|
text: 'Resource Link: My Resource at file:///path/to/resource',
|
|
},
|
|
{ text: 'Embedded text content.' },
|
|
{
|
|
text: `[Tool '${serverToolName}' provided the following image data with mime-type: image/jpeg]`,
|
|
},
|
|
{
|
|
inlineData: {
|
|
mimeType: 'image/jpeg',
|
|
data: 'BASE64_IMAGE_DATA',
|
|
},
|
|
},
|
|
]);
|
|
expect(toolResult.returnDisplay).toBe(
|
|
'Here is a resource.\n[Link to My Resource: file:///path/to/resource]\nEmbedded text content.\n[Image: image/jpeg]',
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('shouldConfirmExecute', () => {
|
|
it('should return false if trust is true', async () => {
|
|
const trustedTool = new DiscoveredMCPTool(
|
|
mockCallableToolInstance,
|
|
serverName,
|
|
serverToolName,
|
|
baseDescription,
|
|
inputSchema,
|
|
undefined,
|
|
true,
|
|
);
|
|
const invocation = trustedTool.build({});
|
|
expect(
|
|
await invocation.shouldConfirmExecute(new AbortController().signal),
|
|
).toBe(false);
|
|
});
|
|
|
|
it('should return false if server is allowlisted', async () => {
|
|
const invocation = tool.build({}) as any;
|
|
invocation.constructor.allowlist.add(serverName);
|
|
expect(
|
|
await invocation.shouldConfirmExecute(new AbortController().signal),
|
|
).toBe(false);
|
|
});
|
|
|
|
it('should return false if tool is allowlisted', async () => {
|
|
const toolAllowlistKey = `${serverName}.${serverToolName}`;
|
|
const invocation = tool.build({}) as any;
|
|
invocation.constructor.allowlist.add(toolAllowlistKey);
|
|
expect(
|
|
await invocation.shouldConfirmExecute(new AbortController().signal),
|
|
).toBe(false);
|
|
});
|
|
|
|
it('should return confirmation details if not trusted and not allowlisted', async () => {
|
|
const invocation = tool.build({});
|
|
const confirmation = await invocation.shouldConfirmExecute(
|
|
new AbortController().signal,
|
|
);
|
|
expect(confirmation).not.toBe(false);
|
|
if (confirmation && confirmation.type === 'mcp') {
|
|
// Type guard for ToolMcpConfirmationDetails
|
|
expect(confirmation.type).toBe('mcp');
|
|
expect(confirmation.serverName).toBe(serverName);
|
|
expect(confirmation.toolName).toBe(serverToolName);
|
|
} else if (confirmation) {
|
|
// Handle other possible confirmation types if necessary, or strengthen test if only MCP is expected
|
|
throw new Error(
|
|
'Confirmation was not of expected type MCP or was false',
|
|
);
|
|
} else {
|
|
throw new Error(
|
|
'Confirmation details not in expected format or was false',
|
|
);
|
|
}
|
|
});
|
|
|
|
it('should add server to allowlist on ProceedAlwaysServer', async () => {
|
|
const invocation = tool.build({}) as any;
|
|
const confirmation = await invocation.shouldConfirmExecute(
|
|
new AbortController().signal,
|
|
);
|
|
expect(confirmation).not.toBe(false);
|
|
if (
|
|
confirmation &&
|
|
typeof confirmation === 'object' &&
|
|
'onConfirm' in confirmation &&
|
|
typeof confirmation.onConfirm === 'function'
|
|
) {
|
|
await confirmation.onConfirm(
|
|
ToolConfirmationOutcome.ProceedAlwaysServer,
|
|
);
|
|
expect(invocation.constructor.allowlist.has(serverName)).toBe(true);
|
|
} else {
|
|
throw new Error(
|
|
'Confirmation details or onConfirm not in expected format',
|
|
);
|
|
}
|
|
});
|
|
|
|
it('should add tool to allowlist on ProceedAlwaysTool', async () => {
|
|
const toolAllowlistKey = `${serverName}.${serverToolName}`;
|
|
const invocation = tool.build({}) as any;
|
|
const confirmation = await invocation.shouldConfirmExecute(
|
|
new AbortController().signal,
|
|
);
|
|
expect(confirmation).not.toBe(false);
|
|
if (
|
|
confirmation &&
|
|
typeof confirmation === 'object' &&
|
|
'onConfirm' in confirmation &&
|
|
typeof confirmation.onConfirm === 'function'
|
|
) {
|
|
await confirmation.onConfirm(ToolConfirmationOutcome.ProceedAlwaysTool);
|
|
expect(invocation.constructor.allowlist.has(toolAllowlistKey)).toBe(
|
|
true,
|
|
);
|
|
} else {
|
|
throw new Error(
|
|
'Confirmation details or onConfirm not in expected format',
|
|
);
|
|
}
|
|
});
|
|
|
|
it('should handle Cancel confirmation outcome', async () => {
|
|
const invocation = tool.build({}) as any;
|
|
const confirmation = await invocation.shouldConfirmExecute(
|
|
new AbortController().signal,
|
|
);
|
|
expect(confirmation).not.toBe(false);
|
|
if (
|
|
confirmation &&
|
|
typeof confirmation === 'object' &&
|
|
'onConfirm' in confirmation &&
|
|
typeof confirmation.onConfirm === 'function'
|
|
) {
|
|
// Cancel should not add anything to allowlist
|
|
await confirmation.onConfirm(ToolConfirmationOutcome.Cancel);
|
|
expect(invocation.constructor.allowlist.has(serverName)).toBe(false);
|
|
expect(
|
|
invocation.constructor.allowlist.has(
|
|
`${serverName}.${serverToolName}`,
|
|
),
|
|
).toBe(false);
|
|
} else {
|
|
throw new Error(
|
|
'Confirmation details or onConfirm not in expected format',
|
|
);
|
|
}
|
|
});
|
|
|
|
it('should handle ProceedOnce confirmation outcome', async () => {
|
|
const invocation = tool.build({}) as any;
|
|
const confirmation = await invocation.shouldConfirmExecute(
|
|
new AbortController().signal,
|
|
);
|
|
expect(confirmation).not.toBe(false);
|
|
if (
|
|
confirmation &&
|
|
typeof confirmation === 'object' &&
|
|
'onConfirm' in confirmation &&
|
|
typeof confirmation.onConfirm === 'function'
|
|
) {
|
|
// ProceedOnce should not add anything to allowlist
|
|
await confirmation.onConfirm(ToolConfirmationOutcome.ProceedOnce);
|
|
expect(invocation.constructor.allowlist.has(serverName)).toBe(false);
|
|
expect(
|
|
invocation.constructor.allowlist.has(
|
|
`${serverName}.${serverToolName}`,
|
|
),
|
|
).toBe(false);
|
|
} else {
|
|
throw new Error(
|
|
'Confirmation details or onConfirm not in expected format',
|
|
);
|
|
}
|
|
});
|
|
});
|
|
});
|