quantum-cli-core
Version:
Quantum CLI Core - Multi-LLM Collaboration System
173 lines • 9.72 kB
JavaScript
/**
* @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, } from 'vitest';
import { DiscoveredMCPTool } from './mcp-tool.js'; // Added getStringifiedResultForDisplay
import { ToolConfirmationOutcome } from './tools.js'; // Added ToolConfirmationOutcome
// 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 = {
tool: mockToolMethod, // Not directly used by DiscoveredMCPTool instance methods
callTool: mockCallTool,
// Add other methods if DiscoveredMCPTool starts using them
};
describe('DiscoveredMCPTool', () => {
const serverName = 'mock-mcp-server';
const toolNameForModel = 'test-mcp-tool-for-model';
const serverToolName = 'actual-server-tool-name';
const baseDescription = 'A test MCP tool.';
const inputSchema = {
type: 'object',
properties: { param: { type: 'string' } },
required: ['param'],
};
beforeEach(() => {
mockCallTool.mockClear();
mockToolMethod.mockClear();
// Clear allowlist before each relevant test, especially for shouldConfirmExecute
DiscoveredMCPTool.allowlist.clear();
});
afterEach(() => {
vi.restoreAllMocks();
});
describe('constructor', () => {
it('should set properties correctly (non-generic server)', () => {
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, // serverName is 'mock-mcp-server', not 'mcp'
toolNameForModel, baseDescription, inputSchema, serverToolName);
expect(tool.name).toBe(toolNameForModel);
expect(tool.schema.name).toBe(toolNameForModel);
expect(tool.schema.description).toBe(baseDescription);
expect(tool.schema.parameters).toEqual(inputSchema);
expect(tool.serverToolName).toBe(serverToolName);
expect(tool.timeout).toBeUndefined();
});
it('should set properties correctly (generic "mcp" server)', () => {
const genericServerName = 'mcp';
const tool = new DiscoveredMCPTool(mockCallableToolInstance, genericServerName, // serverName is 'mcp'
toolNameForModel, baseDescription, inputSchema, serverToolName);
expect(tool.schema.description).toBe(baseDescription);
});
it('should accept and store a custom timeout', () => {
const customTimeout = 5000;
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName, customTimeout);
expect(tool.timeout).toBe(customTimeout);
});
});
describe('execute', () => {
it('should call mcpTool.callTool with correct parameters and format display output', async () => {
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
const params = { param: 'testValue' };
const mockToolSuccessResultObject = {
success: true,
details: 'executed',
};
const mockFunctionResponseContent = [
{ text: JSON.stringify(mockToolSuccessResultObject) },
];
const mockMcpToolResponseParts = [
{
functionResponse: {
name: serverToolName,
response: { content: mockFunctionResponseContent },
},
},
];
mockCallTool.mockResolvedValue(mockMcpToolResponseParts);
const toolResult = await tool.execute(params);
expect(mockCallTool).toHaveBeenCalledWith([
{ name: serverToolName, args: params },
]);
expect(toolResult.llmContent).toEqual(mockMcpToolResponseParts);
const stringifiedResponseContent = JSON.stringify(mockToolSuccessResultObject);
expect(toolResult.returnDisplay).toBe(stringifiedResponseContent);
});
it('should handle empty result from getStringifiedResultForDisplay', async () => {
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
const params = { param: 'testValue' };
const mockMcpToolResponsePartsEmpty = [];
mockCallTool.mockResolvedValue(mockMcpToolResponsePartsEmpty);
const toolResult = await tool.execute(params);
expect(toolResult.returnDisplay).toBe('```json\n[]\n```');
});
it('should propagate rejection if mcpTool.callTool rejects', async () => {
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
const params = { param: 'failCase' };
const expectedError = new Error('MCP call failed');
mockCallTool.mockRejectedValue(expectedError);
await expect(tool.execute(params)).rejects.toThrow(expectedError);
});
});
describe('shouldConfirmExecute', () => {
// beforeEach is already clearing allowlist
it('should return false if trust is true', async () => {
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName, undefined, true);
expect(await tool.shouldConfirmExecute({}, new AbortController().signal)).toBe(false);
});
it('should return false if server is allowlisted', async () => {
DiscoveredMCPTool.allowlist.add(serverName);
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
expect(await tool.shouldConfirmExecute({}, new AbortController().signal)).toBe(false);
});
it('should return false if tool is allowlisted', async () => {
const toolAllowlistKey = `${serverName}.${serverToolName}`;
DiscoveredMCPTool.allowlist.add(toolAllowlistKey);
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
expect(await tool.shouldConfirmExecute({}, new AbortController().signal)).toBe(false);
});
it('should return confirmation details if not trusted and not allowlisted', async () => {
const tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
const confirmation = await tool.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 tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
const confirmation = await tool.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(DiscoveredMCPTool.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 tool = new DiscoveredMCPTool(mockCallableToolInstance, serverName, toolNameForModel, baseDescription, inputSchema, serverToolName);
const toolAllowlistKey = `${serverName}.${serverToolName}`;
const confirmation = await tool.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(DiscoveredMCPTool.allowlist.has(toolAllowlistKey)).toBe(true);
}
else {
throw new Error('Confirmation details or onConfirm not in expected format');
}
});
});
});
//# sourceMappingURL=mcp-tool.test.js.map