Conversation
fix(cli): Write shell command output to a file and limit memory buffered in UI Fixes. Checkpoint. fix(core, cli): await outputStream.end() to prevent race conditions This commit fixes a critical race condition where was called synchronously without being awaited. This led to potential file truncation or EBUSY errors on Windows when attempting to manipulate the file immediately after the call. Additionally, this change removes fixed wait times (`setTimeout`) that were previously used in test files as a band-aid. fix(core): stream processed xterm output to file to remove spurious escape codes test(core): update shell regression tests to use file_data events
Summary of ChangesHello @jacob314, 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 significantly overhauls how shell command outputs are handled, particularly for large data streams. The primary goal is to enhance stability and performance by offloading extensive output from memory to temporary files, thereby preventing UI freezes and memory exhaustion. It also refines the output processing to ensure clean, readable logs and introduces robust error handling and cleanup for these temporary files, making the system more resilient and user-friendly for long-running or verbose commands. Highlights
Changelog
Activity
Using Gemini Code AssistThe 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 by creating a comment using either
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 Limitations & Feedback Gemini Code Assist 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. 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
|
|
Size Change: +12.6 kB (+0.05%) Total Size: 25.7 MB
ℹ️ View Unchanged
|
There was a problem hiding this comment.
Code Review
This pull request introduces a more efficient way to handle large shell command outputs by streaming them to a file instead of holding them in memory, which is a great improvement for performance and stability. However, a critical security vulnerability was identified in the ToolExecutor class, which trusts file paths returned by tools without validation. This could allow a malicious tool or MCP server to read or delete arbitrary files on the user's system. Additionally, there is a critical data loss regression in the binary output handling and an inconsistency in how log files are created for different types of shell commands that need to be addressed to ensure robustness and a consistent user experience.
| if (toolResult.fullOutputFilePath) { | ||
| const threshold = this.config.getTruncateToolOutputThreshold(); | ||
| if ( | ||
| threshold > 0 && | ||
| typeof content === 'string' && | ||
| content.length > threshold | ||
| ) { | ||
| const { outputFile: savedPath } = await moveToolOutputToFile( | ||
| toolResult.fullOutputFilePath, | ||
| toolName, | ||
| callId, | ||
| this.config.storage.getProjectTempDir(), | ||
| this.config.getSessionId(), | ||
| ); | ||
| outputFile = savedPath; | ||
| content = formatTruncatedToolOutput(content, outputFile, threshold); | ||
| } else { | ||
| // If the content is not truncated, we don't need the temporary file. | ||
| try { | ||
| await fsPromises.unlink(toolResult.fullOutputFilePath); | ||
| } catch (error) { | ||
| debugLogger.warn( | ||
| `Failed to delete temporary tool output file: ${toolResult.fullOutputFilePath}`, | ||
| error, | ||
| ); | ||
| } | ||
| } |
There was a problem hiding this comment.
The ToolExecutor class trusts the fullOutputFilePath field returned in a ToolResult by any tool (including third-party MCP tools). This path is used in several dangerous file operations without validation:
- Arbitrary File Read/Exfiltration: On line 243,
moveToolOutputToFileis called withtoolResult.fullOutputFilePath. This function renames or copies the file to the project's temporary directory. The new path is then returned to the LLM in the function response (line 290). A malicious tool could return a sensitive path (e.g.,/etc/passwdor~/.ssh/id_rsa), causing the CLI to move it to an allowed directory and expose its contents to the LLM. - Arbitrary File Deletion: On lines 127, 144, and 255, the code calls
fsPromises.unlink(toolResult.fullOutputFilePath). A malicious tool could return a path to a critical system file, which would then be deleted by the CLI if the execution is aborted or fails.
Remediation:
Before performing any file operations on toolResult.fullOutputFilePath, verify that the path is located within a safe, designated temporary directory (e.g., os.tmpdir()).
| case 'raw_data': | ||
| if (!isBinaryStream) { | ||
| outputStream.write(event.chunk); | ||
| totalBytesWritten += Buffer.byteLength(event.chunk); | ||
| } | ||
| break; |
There was a problem hiding this comment.
This case handles raw_data events for streaming output to a file. However, raw_data contains unprocessed PTY escape codes, which will result in a log file that is not clean and may be difficult for users to read.
For consistency and better user experience, this should be changed to handle the file_data event instead. The file_data event provides processed, clean text suitable for logging, which is already being correctly used by the run_shell_command tool implementation.
| case 'raw_data': | |
| if (!isBinaryStream) { | |
| outputStream.write(event.chunk); | |
| totalBytesWritten += Buffer.byteLength(event.chunk); | |
| } | |
| break; | |
| case 'file_data': | |
| if (!isBinaryStream) { | |
| outputStream.write(event.chunk); | |
| totalBytesWritten += Buffer.byteLength(event.chunk); | |
| } | |
| break; |
|
|
||
| const onAbort = () => combinedController.abort(); | ||
|
|
||
| const outputFileName = `gemini_shell_output_${crypto.randomBytes(6).toString('hex')}.log`; |
There was a problem hiding this comment.
I've gotten a bit lost in the code. seems like shellExecutionService.ts has the core clever logic for this but maybe I'm misunderstanding how this all works.
Summary
Optimize shell output support to avoid memory usage crashes and improve performance by 30X
Before demo on the same file:
https://screencast.googleplex.com/cast/NTE1NTk3NTcyMjk1ODg0OHxkMjE4ODdjNy0xYg
Performance is so slow you may not run out of memory because the heat death of the universe may occur first.
After demo:
You can now read a 5GB file streamed to stdout and not crash. Performance is now generally aligned with VSCode for the same task instead of being 30X slower.
https://screencast.googleplex.com/cast/NjUyMTA1MjIyNTA3NzI0OHw5YTZiMGIwNS1jNg
Fixes #18953
How to test
Test to make sure this hasn't regressed any crucial use cases related to streaming shell output. In particular make sure content passed to the model is still behaving as you would expect. Ensure the saved file still has the full content for very large outputs.