fix(core): remove shell outputChunks buffer caching to prevent memory bloat and sanitize prompt input (#23751)

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:
+        text: `I ran the following shell command:\n\`\`\`sh\n${safeQuery}\n\`\`\`\n\nThis produced the following result:\n\`\`\`\n${safeModelContent}\n\`\`\``,
-\`\`\`sh
-${rawQuery}
-\`\`\`
-
-This produced the following result:
-\`\`\`
-${modelContent}
-\`\`\``,
       },
     ],
   });
@@ -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 {

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;

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',

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(
+          const totalBytes = state.binaryBytesReceived;
-            (sum, chunk) => sum + chunk.length,
-            0,
-          );
           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;
-        return { finalBuffer };
       }
 
       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(
+                const totalBytes = binaryBytesReceived;
-                  (sum, chunk) => sum + chunk.length,
-                  0,
-                );
                 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,