feat: Strip schema props from MCP tool definitions
- This change modifies the tool discovery process for MCP (Model Context Protocol) tools. - When tools are fetched from an MCP server, the `additionalProperties` and `$schema` fields are now recursively removed from their input schemas. This ensures cleaner and more concise tool definitions within the CLI, aligning with the expected schema structure and preventing potential conflicts or verbose outputs. - The corresponding tests in `tool-registry.test.ts` have been updated to reflect this new behavior and verify the correct stripping of these properties. Workaround for https://github.com/google-gemini/gemini-cli/issues/398
This commit is contained in:
parent
0e25fdd56e
commit
e0b88dc8da
|
@ -263,7 +263,7 @@ describe('ToolRegistry', () => {
|
||||||
);
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should discover tools using MCP servers defined in getMcpServers', async () => {
|
it('should discover tools using MCP servers defined in getMcpServers and strip schema properties', async () => {
|
||||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined); // No regular discovery
|
mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined); // No regular discovery
|
||||||
mockConfigGetMcpServerCommand.mockReturnValue(undefined); // No command-based MCP
|
mockConfigGetMcpServerCommand.mockReturnValue(undefined); // No command-based MCP
|
||||||
mockConfigGetMcpServers.mockReturnValue({
|
mockConfigGetMcpServers.mockReturnValue({
|
||||||
|
@ -279,16 +279,26 @@ describe('ToolRegistry', () => {
|
||||||
{
|
{
|
||||||
name: 'mcp-tool-1',
|
name: 'mcp-tool-1',
|
||||||
description: 'An MCP tool',
|
description: 'An MCP tool',
|
||||||
inputSchema: { type: 'object' },
|
inputSchema: {
|
||||||
}, // Corrected: Add type: 'object'
|
type: 'object',
|
||||||
|
properties: {
|
||||||
|
param1: { type: 'string', $schema: 'remove-me' },
|
||||||
|
param2: {
|
||||||
|
type: 'object',
|
||||||
|
additionalProperties: false,
|
||||||
|
properties: {
|
||||||
|
nested: { type: 'number' },
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
additionalProperties: true,
|
||||||
|
$schema: 'http://json-schema.org/draft-07/schema#',
|
||||||
|
},
|
||||||
|
},
|
||||||
],
|
],
|
||||||
});
|
});
|
||||||
mockMcpClientInstance.connect.mockResolvedValue(undefined);
|
mockMcpClientInstance.connect.mockResolvedValue(undefined);
|
||||||
|
|
||||||
// discoverTools has an async IIFE, so we need to wait for it.
|
|
||||||
// A simple way is to use a short timeout, but a more robust way would be to
|
|
||||||
// have discoverTools return a promise that resolves when all async operations are done.
|
|
||||||
// For now, using a timeout.
|
|
||||||
toolRegistry.discoverTools();
|
toolRegistry.discoverTools();
|
||||||
await new Promise((resolve) => setTimeout(resolve, 100)); // Wait for async operations
|
await new Promise((resolve) => setTimeout(resolve, 100)); // Wait for async operations
|
||||||
|
|
||||||
|
@ -302,11 +312,26 @@ describe('ToolRegistry', () => {
|
||||||
expect(mockMcpClientInstance.connect).toHaveBeenCalled();
|
expect(mockMcpClientInstance.connect).toHaveBeenCalled();
|
||||||
expect(mockMcpClientInstance.listTools).toHaveBeenCalled();
|
expect(mockMcpClientInstance.listTools).toHaveBeenCalled();
|
||||||
|
|
||||||
const discoveredTool = toolRegistry.getTool('mcp-tool-1'); // Name is not prefixed if only one MCP server from getMcpServers
|
const discoveredTool = toolRegistry.getTool('mcp-tool-1');
|
||||||
expect(discoveredTool).toBeInstanceOf(DiscoveredMCPTool);
|
expect(discoveredTool).toBeInstanceOf(DiscoveredMCPTool);
|
||||||
expect(discoveredTool?.name).toBe('mcp-tool-1');
|
expect(discoveredTool?.name).toBe('mcp-tool-1');
|
||||||
expect(discoveredTool?.description).toContain('An MCP tool');
|
expect(discoveredTool?.description).toContain('An MCP tool');
|
||||||
expect(discoveredTool?.description).toContain('mcp-tool-1');
|
expect(discoveredTool?.description).toContain('mcp-tool-1');
|
||||||
|
|
||||||
|
// Verify that $schema and additionalProperties are removed
|
||||||
|
const cleanedSchema = discoveredTool?.schema.parameters;
|
||||||
|
expect(cleanedSchema).not.toHaveProperty('$schema');
|
||||||
|
expect(cleanedSchema).not.toHaveProperty('additionalProperties');
|
||||||
|
expect(cleanedSchema?.properties?.param1).not.toHaveProperty('$schema');
|
||||||
|
expect(cleanedSchema?.properties?.param2).not.toHaveProperty(
|
||||||
|
'additionalProperties',
|
||||||
|
);
|
||||||
|
expect(
|
||||||
|
cleanedSchema?.properties?.param2?.properties?.nested,
|
||||||
|
).not.toHaveProperty('$schema');
|
||||||
|
expect(
|
||||||
|
cleanedSchema?.properties?.param2?.properties?.nested,
|
||||||
|
).not.toHaveProperty('additionalProperties');
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should discover tools using MCP server command from getMcpServerCommand', async () => {
|
it('should discover tools using MCP server command from getMcpServerCommand', async () => {
|
||||||
|
|
|
@ -226,6 +226,22 @@ export class ToolRegistry {
|
||||||
});
|
});
|
||||||
const result = await mcpClient.listTools();
|
const result = await mcpClient.listTools();
|
||||||
for (const tool of result.tools) {
|
for (const tool of result.tools) {
|
||||||
|
// Recursively remove additionalProperties and $schema from the inputSchema
|
||||||
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any -- This function recursively navigates a deeply nested and potentially heterogeneous JSON schema object. Using 'any' is a pragmatic choice here to avoid overly complex type definitions for all possible schema variations.
|
||||||
|
const removeSchemaProps = (obj: any) => {
|
||||||
|
if (typeof obj !== 'object' || obj === null) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if (Array.isArray(obj)) {
|
||||||
|
obj.forEach(removeSchemaProps);
|
||||||
|
} else {
|
||||||
|
delete obj.additionalProperties;
|
||||||
|
delete obj.$schema;
|
||||||
|
Object.values(obj).forEach(removeSchemaProps);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
removeSchemaProps(tool.inputSchema);
|
||||||
|
|
||||||
this.registerTool(
|
this.registerTool(
|
||||||
new DiscoveredMCPTool(
|
new DiscoveredMCPTool(
|
||||||
mcpClient,
|
mcpClient,
|
||||||
|
|
Loading…
Reference in New Issue