fix(core): Sanitize tool parameters to fix 400 API errors#3300
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @BigUncle, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request provides a critical fix for API errors encountered when registering tools with the Gemini API. It introduces a robust sanitization process for tool parameter schemas, ensuring that only compatible formats are sent, and significantly improves the overall reliability of the tool discovery process. The changes centralize tool-related utility functions and are backed by comprehensive new tests.
Highlights
- API Compatibility Fix: Implemented a
sanitizeParametersfunction to remove unsupportedformatproperties fromSTRINGtype parameters within tool function declarations. This directly addresses persistent 400 API errors when interacting with the Gemini API, which strictly requiresenumordate-timeformats for strings. - Tool Discovery Robustness: Enhanced the tool discovery mechanism in
ToolRegistryby addingArray.isArraychecks forfunction_declarationsandfunctionDeclarationsto ensure correct structure, preventing 'missing field' errors. Error handling for failed discovery commands and warnings for tools without names have also been added. - Code Reorganization: The
sanitizeParametersutility function has been relocated frommcp-client.tstotool-registry.tsto centralize tool-related logic and improve modularity. This function is now called directly withinToolRegistryduring tool discovery. - Comprehensive Testing: Added extensive unit tests for the new
sanitizeParametersfunction intool-registry.test.ts, covering various schema structures including simple properties, nested objects, arrays, and circular references, ensuring its correctness and resilience. ExistingToolRegistrytests were also updated to reflect the new sanitization behavior.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request effectively addresses 400 API errors by sanitizing tool parameter schemas and improves the tool discovery process. The sanitizeParameters function is enhanced to handle cyclical schemas and remove unsupported format values. The tool discovery is more resilient with improved error handling and validation. I've identified a high-severity issue related to robustness in the tool discovery logic that could cause it to fail if the external discovery command returns malformed data, and have included a suggestion to address this.
ba5d6fb to
b00583b
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request is a great improvement, fixing a persistent 400 API error by sanitizing tool schemas before they are sent to the Gemini API. The introduction of sanitizeParameters with cycle detection and the added robustness in discoverTools are excellent changes. I've identified one high-severity issue where insufficient validation of the discovered tool's parameters property could re-introduce the very API errors this PR aims to solve. My suggestion addresses this by ensuring parameters is always an object.
Sanitizes tool function declaration schemas before they are sent to the Gemini API. This resolves a persistent '400 Bad Request' error that occurred when a tool parameter of type STRING had a `format` property other than 'enum' or 'date-time'. This commit inspects all discovered tools and sets unsupported `format` properties to `undefined`, ensuring API compatibility. It also adds array validation for tool declarations to prevent `missing field` errors. This single commit also includes the necessary test coverage and fixes for related flaky tests. Fixes google-gemini#2141 Fixes google-gemini#2706 Fixes google-gemini#3178 Fixes google-gemini#2237 Related to google-gemini#2522
b00583b to
b21ce8f
Compare
|
/gemini review |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request enhances the security and stability of the tool discovery process. The replacement of execSync with spawn mitigates command injection risks, and the sanitizeParameters function addresses API errors and prevents infinite recursion. A potential resource exhaustion issue due to unbounded output buffering has been identified.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request provides a robust fix for API errors by sanitizing tool parameters and significantly improves the tool discovery process by replacing execSync with a more secure and resilient asynchronous spawn implementation. The introduction of cycle detection in the schema sanitizer is a critical stability improvement. My review includes one high-severity recommendation to further enhance the new spawn implementation by adding a size limit for stderr and ensuring correct handling of multi-byte character streams to prevent potential memory issues and data corruption.
|
/gemini review |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request provides a robust solution to fix 400 API errors from invalid tool parameter schemas by introducing a sanitizeParameters function. The changes are well-implemented, significantly improving security and robustness by replacing execSync with spawn, adding shell command parsing, and implementing I/O size limits. The new sanitization logic is well-tested, including for complex cases like cyclic schemas. I have one high-severity suggestion to improve maintainability by using the public API for sanitization instead of its internal implementation detail.
This commit introduces a comprehensive set of improvements to the tool discovery mechanism in the `ToolRegistry`. The key changes include: - Replacing `execSync` with `spawn` to mitigate command injection risks. - Implementing robust error handling for the `spawn` process, including an `error` event listener and re-throwing errors to prevent silent failures. - Adding comprehensive input validation for the tool discovery command and its outputs to ensure stability. - Fixing a memory leak in `sanitizeParameters` by creating a new `Set` for each top-level call, preventing the `visited` set from growing indefinitely. - Adding a size limit to the `stdout` buffer in `discoverTools` to prevent potential memory exhaustion. - Fixing a broken test in `getFolderStructure.test.ts` by adding the missing `path` and `parentPath` properties to the `createDirent` mock. - Consolidating multiple related commits into a single, cohesive update. feat(tool-registry): improve robustness of tool discovery This commit enhances the robustness of the tool discovery process in several ways: - Replaces `execSync` with `spawn` to prevent command injection vulnerabilities. - Adds comprehensive error handling for the `spawn` process, including listening for the 'error' event. - Validates the discovery command to prevent the use of shell operators and to handle empty commands gracefully. - Implements size limits for both `stdout` and `stderr` to prevent excessive memory usage. - Uses `StringDecoder` to correctly handle multi-byte UTF-8 characters in the output of the discovery command. - Ensures that `func.parameters` is a valid object before use. - Refactors `sanitizeParameters` to use a non-exported helper function, avoiding issues with mutable default values. - Re-throws errors caught during tool discovery to prevent silent failures. These changes make the tool discovery process more secure, resilient, and reliable. feat(tool-discovery): Enhance tool discovery robustness This commit addresses multiple feedback points from `gemini-code-assist` in GitHub PR google-gemini#3300, focusing on improving the robustness, security, and error handling of the tool discovery mechanism. Key changes include: - **Memory Usage Optimization**: Enforced size limits for both `stdout` and `stderr` to prevent potential memory exhaustion from excessive output by external commands. - **Secure Execution**: Replaced the `execSync` call with the more secure `spawn` to mitigate command injection risks. - **Enhanced Robustness**: - Added null checks for the `tool` object returned by `discoveryCmd` to prevent potential `TypeError`. - Implemented strict type validation for `func.parameters`, defaulting to an empty object `{}` if it's not a valid object, thus avoiding API errors from invalid schemas. - Corrected the `visited` parameter handling in `sanitizeParameters` by creating a new `Set` for each call, resolving a potential memory leak from a mutable default argument. - **Improved Error Handling**: - Strengthened error handling by listening for the `error` event on the `spawn` process, ensuring process launch failures are caught and properly handled. - Validated the output of `shell-quote.parse` to ensure all command parts are strings, preventing runtime crashes from invalid arguments. - Elevated the log level from `console.warn` to `console.error` for tool discovery failures and re-threw the exception to make issues more visible. - **Encoding Correctness**: Introduced `StringDecoder` to process data streams from `stdout` and `stderr`, ensuring correct parsing of multi-byte UTF-8 characters and preventing `JSON.parse` errors. - **Byte Length Correction**: Fixed the size-limiting logic for `stdout` and `stderr`, which incorrectly compared character length with byte length, to ensure accurate enforcement of the limits. Together, these changes significantly improve the reliability and security of the tool discovery feature, providing a more stable user experience. feat(tool-discovery): Enhance tool discovery robustness This commit addresses multiple feedback points from `gemini-code-assist` in GitHub PR google-gemini#3300, focusing on improving the robustness, security, and error handling of the tool discovery mechanism. Key changes include: - **Memory Usage Optimization**: Enforced size limits for both `stdout` and `stderr` to prevent potential memory exhaustion from excessive output by external commands. - **Secure Execution**: Replaced the `execSync` call with the more secure `spawn` to mitigate command injection risks. - **Enhanced Robustness**: - Added null checks for the `tool` object returned by `discoveryCmd` to prevent potential `TypeError`. - Implemented strict type validation for `func.parameters`, defaulting to an empty object `{}` if it's not a valid object, thus avoiding API errors from invalid schemas. - Corrected the `visited` parameter handling in `sanitizeParameters` by creating a new `Set` for each call, resolving a potential memory leak from a mutable default argument. - **Improved Error Handling**: - Strengthened error handling by listening for the `error` event on the `spawn` process, ensuring process launch failures are caught and properly handled. - Validated the output of `shell-quote.parse` to ensure all command parts are strings, preventing runtime crashes from invalid arguments. - Elevated the log level from `console.warn` to `console.error` for tool discovery failures and re-threw the exception to make issues more visible. - **Encoding Correctness**: Introduced `StringDecoder` to process data streams from `stdout` and `stderr`, ensuring correct parsing of multi-byte UTF-8 characters and preventing `JSON.parse` errors. - **Byte Length Correction**: Fixed the size-limiting logic for `stdout` and `stderr`, which incorrectly compared character length with byte length, to ensure accurate enforcement of the limits. - **Encapsulation**: Switched from calling the internal `_sanitizeParameters` function to the public `sanitizeParameters` function to respect encapsulation and improve maintainability. Together, these changes significantly improve the reliability and security of the tool discovery feature, providing a more stable user experience.
|
/gemini review |
There was a problem hiding this comment.
Code Review
The pull request introduces a fix for 400 API errors by sanitizing tool parameters and improving tool registry robustness. The changes include sanitizing the format property for STRING type parameters and adding Array.isArray checks. The use of spawn instead of execSync is a good improvement. I've added comments regarding potential character encoding issues and command injection vulnerabilities.
This commit implements a series of improvements to the tool discovery mechanism in `packages/core/src/tools/tool-registry.ts`, addressing feedback from the `gemini-code-assist` bot to enhance robustness, security, and error handling.
The following changes have been made:
- **Command Injection Prevention**: Replaced the use of `execSync` with `spawn` to prevent potential command injection vulnerabilities when executing the tool discovery command. This ensures that the command and its arguments are treated as separate entities, mitigating the risk of malicious code execution.
- **Graceful Error Handling**: Implemented more robust error handling for the `spawn` process. This includes listening for the 'error' event to catch process launch failures and rejecting the promise accordingly. This prevents unhandled exceptions from crashing the application.
- **Input Validation**: Added validation to ensure that the parsed command from `shell-quote` is an array of strings, preventing runtime errors when the command contains shell operators. An informative error message is now thrown to guide users on how to correctly format complex commands.
- **Output Size Limiting**: Introduced a size limit for both `stdout` and `stderr` to prevent excessive memory usage from commands that produce large amounts of data. If the output exceeds the defined limit, the child process is terminated, and an appropriate error is thrown.
- **Correct Character Encoding**: Implemented the use of `StringDecoder` for both `stdout` and `stderr` streams to correctly handle multi-byte UTF-8 characters. This prevents potential data corruption and ensures that JSON output is parsed correctly.
- **Safe Parameter Handling**: Added validation to ensure that `func.parameters` is a valid object before it is used. If it is not a valid object, it now defaults to an empty object `{}` to prevent runtime errors and invalid schemas from being sent to the Gemini API.
- **Improved Error Logging**: Changed the error logging for failed discovery commands from `console.warn` to `console.error` and re-throwing the error to make failures more prominent and prevent silent failures.
- **Decoder Flushing**: Ensured that the `StringDecoder` is properly flushed by calling `.end()` in the `close` event handler. This guarantees that any remaining buffered characters are correctly decoded before parsing, preventing data corruption.
These changes collectively make the tool discovery process more secure, resilient, and user-friendly by providing clearer error messages and preventing common sources of runtime failures.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request provides a robust fix for the 400 API errors by sanitizing tool parameters, and it significantly hardens the security of tool discovery by replacing execSync with spawn and shell-quote. The new sanitizeParameters function is well-implemented and includes comprehensive tests, including for edge cases like circular schemas.
My main concern is the removal of tests for the MCP tool discovery functionality within the discoverTools method, which has led to a regression in test coverage. I've left a comment with details on this.
Overall, this is a high-quality contribution that addresses the reported issues and improves security and robustness.
| describe('discoverTools', () => { | ||
| let mockConfigGetToolDiscoveryCommand: ReturnType<typeof vi.spyOn>; | ||
| let mockConfigGetMcpServers: ReturnType<typeof vi.spyOn>; | ||
| let mockConfigGetMcpServerCommand: ReturnType<typeof vi.spyOn>; | ||
| let mockExecSync: ReturnType<typeof vi.mocked<typeof execSync>>; | ||
|
|
||
| beforeEach(() => { | ||
| mockConfigGetToolDiscoveryCommand = vi.spyOn( | ||
| config, | ||
| 'getToolDiscoveryCommand', | ||
| ); | ||
| mockConfigGetMcpServers = vi.spyOn(config, 'getMcpServers'); | ||
| mockConfigGetMcpServerCommand = vi.spyOn(config, 'getMcpServerCommand'); | ||
| mockExecSync = vi.mocked(execSync); | ||
| toolRegistry = new ToolRegistry(config); // Reset registry | ||
| // Reset the mock for discoverMcpTools before each test in this suite | ||
| mockDiscoverMcpTools.mockReset().mockResolvedValue(undefined); | ||
| }); | ||
|
|
||
| it('should discover tools using discovery command', async () => { | ||
| // ... this test remains largely the same | ||
| it('should sanitize tool parameters during discovery from command', async () => { | ||
| const discoveryCommand = 'my-discovery-command'; | ||
| mockConfigGetToolDiscoveryCommand.mockReturnValue(discoveryCommand); | ||
| const mockToolDeclarations: FunctionDeclaration[] = [ | ||
| { | ||
| name: 'discovered-tool-1', | ||
| description: 'A discovered tool', | ||
| parameters: { type: Type.OBJECT, properties: {} }, | ||
|
|
||
| const unsanitizedToolDeclaration: FunctionDeclaration = { | ||
| name: 'tool-with-bad-format', | ||
| description: 'A tool with an invalid format property', | ||
| parameters: { | ||
| type: Type.OBJECT, | ||
| properties: { | ||
| some_string: { | ||
| type: Type.STRING, | ||
| format: 'uuid', // This is an unsupported format | ||
| }, | ||
| }, | ||
| }, | ||
| ]; | ||
| mockExecSync.mockReturnValue( | ||
| Buffer.from( | ||
| JSON.stringify([{ function_declarations: mockToolDeclarations }]), | ||
| ), | ||
| ); | ||
| await toolRegistry.discoverTools(); | ||
| expect(execSync).toHaveBeenCalledWith(discoveryCommand); | ||
| const discoveredTool = toolRegistry.getTool('discovered-tool-1'); | ||
| expect(discoveredTool).toBeInstanceOf(DiscoveredTool); | ||
| }); | ||
| }; | ||
|
|
||
| it('should discover tools using MCP servers defined in getMcpServers', async () => { | ||
| mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined); | ||
| mockConfigGetMcpServerCommand.mockReturnValue(undefined); | ||
| const mcpServerConfigVal = { | ||
| 'my-mcp-server': { | ||
| command: 'mcp-server-cmd', | ||
| args: ['--port', '1234'], | ||
| trust: true, | ||
| } as MCPServerConfig, | ||
| const mockSpawn = vi.mocked(spawn); | ||
| const mockChildProcess = { | ||
| stdout: { on: vi.fn() }, | ||
| stderr: { on: vi.fn() }, | ||
| on: vi.fn(), | ||
| }; | ||
| mockConfigGetMcpServers.mockReturnValue(mcpServerConfigVal); | ||
| mockSpawn.mockReturnValue(mockChildProcess as any); | ||
|
|
||
| // Simulate stdout data | ||
| mockChildProcess.stdout.on.mockImplementation((event, callback) => { | ||
| if (event === 'data') { | ||
| callback( | ||
| Buffer.from( | ||
| JSON.stringify([ | ||
| { function_declarations: [unsanitizedToolDeclaration] }, | ||
| ]), | ||
| ), | ||
| ); | ||
| } | ||
| return mockChildProcess as any; | ||
| }); | ||
|
|
||
| // Simulate process close | ||
| mockChildProcess.on.mockImplementation((event, callback) => { | ||
| if (event === 'close') { | ||
| callback(0); | ||
| } | ||
| return mockChildProcess as any; | ||
| }); | ||
|
|
||
| await toolRegistry.discoverTools(); | ||
|
|
||
| expect(mockDiscoverMcpTools).toHaveBeenCalledWith( | ||
| mcpServerConfigVal, | ||
| const discoveredTool = toolRegistry.getTool('tool-with-bad-format'); | ||
| expect(discoveredTool).toBeDefined(); | ||
|
|
||
| const registeredParams = (discoveredTool as DiscoveredTool).schema | ||
| .parameters as Schema; | ||
| expect(registeredParams.properties?.['some_string']).toBeDefined(); | ||
| expect(registeredParams.properties?.['some_string']).toHaveProperty( | ||
| 'format', | ||
| undefined, | ||
| toolRegistry, | ||
| ); | ||
| // We no longer check these as discoverMcpTools is mocked | ||
| // expect(vi.mocked(mcpToTool)).toHaveBeenCalledTimes(1); | ||
| // expect(Client).toHaveBeenCalledTimes(1); | ||
| // expect(StdioClientTransport).toHaveBeenCalledWith({ | ||
| // command: 'mcp-server-cmd', | ||
| // args: ['--port', '1234'], | ||
| // env: expect.any(Object), | ||
| // stderr: 'pipe', | ||
| // }); | ||
| // expect(mockMcpClientConnect).toHaveBeenCalled(); | ||
|
|
||
| // To verify that tools *would* have been registered, we'd need mockDiscoverMcpTools | ||
| // to call toolRegistry.registerTool, or we test that separately. | ||
| // For now, we just check that the delegation happened. | ||
| }); | ||
| }); |
There was a problem hiding this comment.
This refactoring has removed several tests for discoverTools, specifically those covering tool discovery from MCP servers. The discoverTools method in tool-registry.ts still contains the logic for MCP discovery, but this functionality is now untested.
The previous test suite included checks for:
- Discovery via
getMcpServersconfig. - Discovery via
getMcpServerCommandconfig. - Graceful error handling during MCP client connection.
While the new test for parameter sanitization is a great addition, removing the tests for other discovery paths is a significant regression in test coverage for a core feature. This could allow future bugs in MCP tool discovery to go unnoticed.
Please reintroduce tests for the MCP discovery paths to ensure this functionality remains covered.
This commit reintroduces the previously removed tests for MCP (Model Context Protocol) tool discovery functionality in `packages/core/src/tools/tool-registry.test.ts`. The restored tests ensure that the `discoverTools` method correctly handles tool discovery from MCP servers configured via `getMcpServers`. This improves the test coverage for a core feature of the tool registry and helps prevent future regressions.
NTaylorMullen
left a comment
There was a problem hiding this comment.
@BigUncle we're super close! Thank you for the updates. Left a few more comments
| expect(WITTY_LOADING_PHRASES).toContain( | ||
| result.current.currentLoadingPhrase, | ||
| ); | ||
| const initialPhrase = result.current.currentLoadingPhrase; |
There was a problem hiding this comment.
Intentional removals?
There was a problem hiding this comment.
I made a mistake in the operation. I have restored it.
| functions.push(...tool['functionDeclarations']); | ||
| } else if (tool['name']) { | ||
| functions.push(tool); | ||
| try { |
There was a problem hiding this comment.
Could you extract this block into a specific function to resolve tools for readability?
There was a problem hiding this comment.
I have extracted it into a new function called discoverAndRegisterToolsFromCommand, thank you!
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a fix for 400 API errors by sanitizing tool parameters. However, there are critical security vulnerabilities related to command injection in the tool discovery mechanism that must be addressed. The use of spawn without proper sanitization and the insufficient check for shell operators pose significant risks.
This commit implements a series of improvements to the tool registry based on feedback from @NTaylorMullen and gemini-code-assist. Key changes: - Refactored the `discoverTools` method by extracting the command-based discovery logic into a new private method, `discoverAndRegisterToolsFromCommand`, to improve readability and maintainability. - Moved the `shell-quote` dependency from the root `package.json` to `packages/core/package.json` where it is actually used. - Restored the unintentionally removed test code in `useLoadingIndicator.test.ts`. - Added a `null` check for `discoveredItems` in `tool-registry.ts` to prevent potential crashes. - Re-instated the previously removed MCP tool discovery tests in `tool-registry.test.ts` to ensure proper test coverage.
Hi @NTaylorMullen, Thank you so much for your detailed feedback! I've addressed all your points in the latest commit:
I really appreciate you taking the time to review my code. Your suggestions have made it much better! |
|
@BigUncle it looks like you may need to rebase your branch on latest main. It wont let me do it on your behalf because it looks like there are some merge conflicts. Almost there! |
|
Thank you for all the back and forth and the contribution @BigUncle ! |
Thank you so much, @NTaylorMullen! I really appreciate you taking the time for such a detailed review and for all the guidance. A special thanks for helping get this over the finish line by resolving the final merge conflict. As this is my first PR to be merged in a major open-source project, I'm absolutely thrilled. I learned a ton from the process. Looking forward to contributing more in the future! |
TLDR
This pull request provides a robust fix for a persistent 400 API error caused by improper tool parameter schemas being sent to the Gemini API. It directly resolves issue #2141 and several duplicates by sanitizing tool function declarations before they are sent, ensuring compatibility with the API's requirements.
Key Fixes:
formatproperty: ForSTRINGtype parameters, it removes anyformatvalue that is notenumordate-time, which was the primary cause of the 400 errors.Array.isArraychecks to ensurefunction_declarationsare correctly structured, which likely prevents themissing fielderrors also reported.This PR closes: #2141, #2706, #3178, #2237
This PR likely also resolves: #2522
Dive Deeper
The root cause of the
Persistent 400 API errorwas that the tool discovery mechanism would generate function declarations withformatproperties (e.g.,uri,email) that are not supported by the Gemini API forSTRINGparameters. The API strictly requires that theformatfor a string be eitherenumordate-time.This PR introduces a
sanitizeParametersfunction that recursively traverses the tool schema. It inspects everySTRINGparameter and sets itsformattoundefinedif it's not one of the allowed values. This sanitization happens right before the tool is registered, ensuring that any request sent to the Gemini API contains a valid tool definition.Additionally, issue #2141 and the related #2522 also reported errors like
...items: missing field.. This often happens when an expected array structure is malformed. The secondary fix in this PR, which validates thatfunction_declarationsandfunctionDeclarationsare indeed arrays before processing them, adds a layer of type safety that should prevent these malformed definitions from causing API errors.Reviewer Test Plan
npm run preflight. All tests should pass consistently."format": "uri").Testing Matrix
upstream/mainwas successful.Linked issues / bugs