From d25ce0e143b712d2c509c0a1b8a0019d81e8d3ad Mon Sep 17 00:00:00 2001 From: Spencer Date: Thu, 26 Mar 2026 17:16:07 -0400 Subject: [PATCH] fix(core): remove shell outputChunks buffer caching to prevent memory bloat and sanitize prompt input (#23751) --- .../cli/src/ui/hooks/shellCommandProcessor.ts | 18 ++++---- .../src/services/executionLifecycleService.ts | 2 +- .../services/shellExecutionService.test.ts | 10 +--- .../src/services/shellExecutionService.ts | 46 +++++++++++-------- 4 files changed, 37 insertions(+), 39 deletions(-) diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts index 7e33d37d1f..3e67ad84b7 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts @@ -45,20 +45,18 @@ function addShellCommandToGeminiHistory( ? resultText.substring(0, MAX_OUTPUT_LENGTH) + '\n... (truncated)' : resultText; + // Escape backticks to prevent prompt injection breakouts + const safeQuery = rawQuery.replace(/\\/g, '\\\\').replace(/\x60/g, '\\\x60'); + const safeModelContent = modelContent + .replace(/\\/g, '\\\\') + .replace(/\x60/g, '\\\x60'); + // eslint-disable-next-line @typescript-eslint/no-floating-promises geminiClient.addHistory({ role: 'user', parts: [ { - text: `I ran the following shell command: -\`\`\`sh -${rawQuery} -\`\`\` - -This produced the following result: -\`\`\` -${modelContent} -\`\`\``, + text: `I ran the following shell command:\n\`\`\`sh\n${safeQuery}\n\`\`\`\n\nThis produced the following result:\n\`\`\`\n${safeModelContent}\n\`\`\``, }, ], }); @@ -444,7 +442,7 @@ export const useShellCommandProcessor = ( } let mainContent: string; - if (isBinary(result.rawOutput)) { + if (isBinaryStream || isBinary(result.rawOutput)) { mainContent = '[Command produced binary output, which is not shown.]'; } else { diff --git a/packages/core/src/services/executionLifecycleService.ts b/packages/core/src/services/executionLifecycleService.ts index 6df693fccb..5efe26c375 100644 --- a/packages/core/src/services/executionLifecycleService.ts +++ b/packages/core/src/services/executionLifecycleService.ts @@ -16,7 +16,7 @@ export type ExecutionMethod = | 'none'; export interface ExecutionResult { - rawOutput: Buffer; + rawOutput?: Buffer; output: string; exitCode: number | null; signal: number | null; diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts index 6a0371b68d..adb519d087 100644 --- a/packages/core/src/services/shellExecutionService.test.ts +++ b/packages/core/src/services/shellExecutionService.test.ts @@ -880,15 +880,12 @@ describe('ShellExecutionService', () => { const binaryChunk1 = Buffer.from([0x89, 0x50, 0x4e, 0x47]); const binaryChunk2 = Buffer.from([0x0d, 0x0a, 0x1a, 0x0a]); - const { result } = await simulateExecution('cat image.png', (pty) => { + await simulateExecution('cat image.png', (pty) => { pty.onData.mock.calls[0][0](binaryChunk1); pty.onData.mock.calls[0][0](binaryChunk2); pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); }); - expect(result.rawOutput).toEqual( - Buffer.concat([binaryChunk1, binaryChunk2]), - ); expect(onOutputEventMock).toHaveBeenCalledTimes(4); expect(onOutputEventMock.mock.calls[0][0]).toEqual({ type: 'binary_detected', @@ -1464,15 +1461,12 @@ describe('ShellExecutionService child_process fallback', () => { const binaryChunk1 = Buffer.from([0x89, 0x50, 0x4e, 0x47]); const binaryChunk2 = Buffer.from([0x0d, 0x0a, 0x1a, 0x0a]); - const { result } = await simulateExecution('cat image.png', (cp) => { + await simulateExecution('cat image.png', (cp) => { cp.stdout?.emit('data', binaryChunk1); cp.stdout?.emit('data', binaryChunk2); cp.emit('exit', 0, null); }); - expect(result.rawOutput).toEqual( - Buffer.concat([binaryChunk1, binaryChunk2]), - ); expect(onOutputEventMock).toHaveBeenCalledTimes(4); expect(onOutputEventMock.mock.calls[0][0]).toEqual({ type: 'binary_detected', diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index a5697104ec..6184354a2a 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/packages/core/src/services/shellExecutionService.ts @@ -120,7 +120,8 @@ interface ActiveChildProcess { state: { output: string; truncated: boolean; - outputChunks: Buffer[]; + sniffChunks: Buffer[]; + binaryBytesReceived: number; }; } @@ -493,7 +494,8 @@ export class ShellExecutionService { const state = { output: '', truncated: false, - outputChunks: [] as Buffer[], + sniffChunks: [] as Buffer[], + binaryBytesReceived: 0, }; if (child.pid) { @@ -563,14 +565,19 @@ export class ShellExecutionService { } } - state.outputChunks.push(data); + if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) { + state.sniffChunks.push(data); + } else if (!isStreamingRawContent) { + state.binaryBytesReceived += data.length; + } if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) { - const sniffBuffer = Buffer.concat(state.outputChunks.slice(0, 20)); + const sniffBuffer = Buffer.concat(state.sniffChunks.slice(0, 20)); sniffedBytes = sniffBuffer.length; if (isBinary(sniffBuffer)) { isStreamingRawContent = false; + state.binaryBytesReceived = sniffBuffer.length; const event: ShellOutputEvent = { type: 'binary_detected' }; onOutputEvent(event); if (child.pid) { @@ -610,10 +617,7 @@ export class ShellExecutionService { } } } else { - const totalBytes = state.outputChunks.reduce( - (sum, chunk) => sum + chunk.length, - 0, - ); + const totalBytes = state.binaryBytesReceived; const event: ShellOutputEvent = { type: 'binary_progress', bytesReceived: totalBytes, @@ -629,7 +633,7 @@ export class ShellExecutionService { code: number | null, signal: NodeJS.Signals | null, ) => { - const { finalBuffer } = cleanup(); + cleanup(); let combinedOutput = state.output; if (state.truncated) { @@ -644,7 +648,7 @@ export class ShellExecutionService { const exitSignal = signal ? os.constants.signals[signal] : null; const resultPayload: ShellExecutionResult = { - rawOutput: finalBuffer, + rawOutput: Buffer.from(''), output: finalStrippedOutput, exitCode, signal: exitSignal, @@ -733,8 +737,7 @@ export class ShellExecutionService { } } - const finalBuffer = Buffer.concat(state.outputChunks); - return { finalBuffer }; + return; } return { pid: child.pid, result }; @@ -864,7 +867,8 @@ export class ShellExecutionService { let processingChain = Promise.resolve(); let decoder: TextDecoder | null = null; let output: string | AnsiOutput | null = null; - const outputChunks: Buffer[] = []; + const sniffChunks: Buffer[] = []; + let binaryBytesReceived = 0; const error: Error | null = null; let exited = false; @@ -995,14 +999,19 @@ export class ShellExecutionService { } } - outputChunks.push(data); + if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) { + sniffChunks.push(data); + } else if (!isStreamingRawContent) { + binaryBytesReceived += data.length; + } if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) { - const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20)); + const sniffBuffer = Buffer.concat(sniffChunks.slice(0, 20)); sniffedBytes = sniffBuffer.length; if (isBinary(sniffBuffer)) { isStreamingRawContent = false; + binaryBytesReceived = sniffBuffer.length; const event: ShellOutputEvent = { type: 'binary_detected' }; onOutputEvent(event); ExecutionLifecycleService.emitEvent(ptyPid, event); @@ -1027,10 +1036,7 @@ export class ShellExecutionService { resolveChunk(); }); } else { - const totalBytes = outputChunks.reduce( - (sum, chunk) => sum + chunk.length, - 0, - ); + const totalBytes = binaryBytesReceived; const event: ShellOutputEvent = { type: 'binary_progress', bytesReceived: totalBytes, @@ -1076,7 +1082,7 @@ export class ShellExecutionService { }); ExecutionLifecycleService.completeWithResult(ptyPid, { - rawOutput: Buffer.concat(outputChunks), + rawOutput: Buffer.from(''), output: getFullBufferText(headlessTerminal), exitCode, signal: signal ?? null,