fix(cli): use newline in shell command wrapping to avoid breaking heredocs (#25537)

2026-04-23 11:34:44 -07:00 · 2026-04-21 15:12:50 -04:00
parent cdc5cccc13
commit 93a8d9001c
4 changed files with 138 additions and 73 deletions
@@ -96,6 +96,7 @@ describe('ShellTool', () => {
  let mockShellOutputCallback: (event: ShellOutputEvent) => void;
  let resolveExecutionPromise: (result: ShellExecutionResult) => void;
  let tempRootDir: string;
+  let extractedTmpFile: string;

  beforeEach(() => {
    vi.clearAllMocks();
@@ -197,16 +198,28 @@ describe('ShellTool', () => {
    process.env['ComSpec'] =
      'C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe';

+    extractedTmpFile = '';
+
    // Capture the output callback to simulate streaming events from the service
-    mockShellExecutionService.mockImplementation((_cmd, _cwd, callback) => {
-      mockShellOutputCallback = callback;
-      return {
-        pid: 12345,
-        result: new Promise((resolve) => {
-          resolveExecutionPromise = resolve;
-        }),
-      };
-    });
+    mockShellExecutionService.mockImplementation(
+      (
+        cmd: string,
+        _cwd: string,
+        callback: (event: ShellOutputEvent) => void,
+      ) => {
+        mockShellOutputCallback = callback;
+        const match = cmd.match(/pgrep -g 0 >([^ ]+)/);
+        if (match) {
+          extractedTmpFile = match[1].replace(/['"]/g, ''); // remove any quotes if present
+        }
+        return {
+          pid: 12345,
+          result: new Promise((resolve) => {
+            resolveExecutionPromise = resolve;
+          }),
+        };
+      },
+    );

    mockShellBackground.mockImplementation(() => {
      resolveExecutionPromise({
@@ -293,17 +306,16 @@ describe('ShellTool', () => {
    it('should wrap command on linux and parse pgrep output', async () => {
      const invocation = shellTool.build({ command: 'my-command &' });
      const promise = invocation.execute({ abortSignal: mockAbortSignal });
-      resolveShellExecution({ pid: 54321 });

      // Simulate pgrep output file creation by the shell command
-      const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-      fs.writeFileSync(tmpFile, `54321${os.EOL}54322${os.EOL}`);
+      fs.writeFileSync(extractedTmpFile, `54321${os.EOL}54322${os.EOL}`);
+
+      resolveShellExecution({ pid: 54321 });

      const result = await promise;

-      const wrappedCommand = `(\n${'my-command &'}\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
      expect(mockShellExecutionService).toHaveBeenCalledWith(
-        wrappedCommand,
+        expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
        tempRootDir,
        expect.any(Function),
        expect.any(AbortSignal),
@@ -316,7 +328,7 @@ describe('ShellTool', () => {
      );
      expect(result.llmContent).toContain('Background PIDs: 54322');
      // The file should be deleted by the tool
-      expect(fs.existsSync(tmpFile)).toBe(false);
+      expect(fs.existsSync(extractedTmpFile)).toBe(false);
    });

    it('should add a space when command ends with a backslash to prevent escaping newline', async () => {
@@ -325,10 +337,8 @@ describe('ShellTool', () => {
      resolveShellExecution();
      await promise;

-      const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-      const wrappedCommand = `(\nls\\ \n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
      expect(mockShellExecutionService).toHaveBeenCalledWith(
-        wrappedCommand,
+        expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
        tempRootDir,
        expect.any(Function),
        expect.any(AbortSignal),
@@ -343,10 +353,8 @@ describe('ShellTool', () => {
      resolveShellExecution();
      await promise;

-      const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-      const wrappedCommand = `(\nls # comment\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
      expect(mockShellExecutionService).toHaveBeenCalledWith(
-        wrappedCommand,
+        expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
        tempRootDir,
        expect.any(Function),
        expect.any(AbortSignal),
@@ -365,10 +373,8 @@ describe('ShellTool', () => {
      resolveShellExecution();
      await promise;

-      const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-      const wrappedCommand = `(\n${'ls'}\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
      expect(mockShellExecutionService).toHaveBeenCalledWith(
-        wrappedCommand,
+        expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
        subdir,
        expect.any(Function),
        expect.any(AbortSignal),
@@ -390,10 +396,8 @@ describe('ShellTool', () => {
      resolveShellExecution();
      await promise;

-      const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-      const wrappedCommand = `(\n${'ls'}\n); __code=$?; pgrep -g 0 >${tmpFile} 2>&1; exit $__code;`;
      expect(mockShellExecutionService).toHaveBeenCalledWith(
-        wrappedCommand,
+        expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
        path.join(tempRootDir, 'subdir'),
        expect.any(Function),
        expect.any(AbortSignal),
@@ -462,6 +466,26 @@ describe('ShellTool', () => {
      20000,
    );

+    it('should correctly wrap heredoc commands', async () => {
+      const command = `cat << 'EOF'
+hello world
+EOF`;
+      const invocation = shellTool.build({ command });
+      const promise = invocation.execute({ abortSignal: mockAbortSignal });
+      resolveShellExecution();
+      await promise;
+
+      expect(mockShellExecutionService).toHaveBeenCalledWith(
+        expect.stringMatching(/pgrep -g 0 >.*gemini-shell-.*[/\\]pgrep\.tmp/),
+        tempRootDir,
+        expect.any(Function),
+        expect.any(AbortSignal),
+        false,
+        expect.any(Object),
+      );
+      expect(mockShellExecutionService.mock.calls[0][0]).toMatch(/\nEOF\n\)\n/);
+    });
+
    it('should format error messages correctly', async () => {
      const error = new Error('wrapped command failed');
      const invocation = shellTool.build({ command: 'user-command' });
@@ -562,10 +586,13 @@ describe('ShellTool', () => {

    it('should clean up the temp file on synchronous execution error', async () => {
      const error = new Error('sync spawn error');
-      mockShellExecutionService.mockImplementation(() => {
-        // Create the temp file before throwing to simulate it being left behind
-        const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-        fs.writeFileSync(tmpFile, '');
+      mockShellExecutionService.mockImplementation((cmd: string) => {
+        const match = cmd.match(/pgrep -g 0 >([^ ]+)/);
+        if (match) {
+          extractedTmpFile = match[1].replace(/['"]/g, ''); // remove any quotes if present
+          // Create the temp file before throwing to simulate it being left behind
+          fs.writeFileSync(extractedTmpFile, '');
+        }
        throw error;
      });

@@ -574,8 +601,7 @@ describe('ShellTool', () => {
        invocation.execute({ abortSignal: mockAbortSignal }),
      ).rejects.toThrow(error);

-      const tmpFile = path.join(os.tmpdir(), 'shell_pgrep_abcdef.tmp');
-      expect(fs.existsSync(tmpFile)).toBe(false);
+      expect(fs.existsSync(extractedTmpFile)).toBe(false);
    });

    it('should not log "missing pgrep output" when process is backgrounded', async () => {
@@ -8,7 +8,6 @@ import fsPromises from 'node:fs/promises';
 import fs from 'node:fs';
 import path from 'node:path';
 import os from 'node:os';
-import crypto from 'node:crypto';
 import { debugLogger } from '../index.js';
 import { type SandboxPermissions } from '../services/sandboxManager.js';
 import { ToolErrorType } from './tool-error.js';
@@ -42,6 +41,7 @@ import {
  parseCommandDetails,
  hasRedirection,
  normalizeCommand,
+  escapeShellArg,
 } from '../utils/shell-utils.js';
 import { SHELL_TOOL_NAME } from './tool-names.js';
 import { PARAM_ADDITIONAL_PERMISSIONS } from './definitions/base-declarations.js';
@@ -111,7 +111,8 @@ export class ShellToolInvocation extends BaseToolInvocation<
    if (trimmed.endsWith('\\')) {
      trimmed += ' ';
    }
-    return `(\n${trimmed}\n); __code=$?; pgrep -g 0 >${tempFilePath} 2>&1; exit $__code;`;
+    const escapedTempFilePath = escapeShellArg(tempFilePath, 'bash');
+    return `(\n${trimmed}\n)\n__code=$?; pgrep -g 0 >${escapedTempFilePath} 2>&1; exit $__code;`;
  }

  private getContextualDetails(): string {
@@ -450,10 +451,8 @@ export class ShellToolInvocation extends BaseToolInvocation<
    }

    const isWindows = os.platform() === 'win32';
-    const tempFileName = `shell_pgrep_${crypto
-      .randomBytes(6)
-      .toString('hex')}.tmp`;
-    const tempFilePath = path.join(os.tmpdir(), tempFileName);
+    let tempFilePath = '';
+    let tempDir = '';

    const timeoutMs = this.context.config.getShellToolInactivityTimeout();
    const timeoutController = new AbortController();
@@ -463,8 +462,10 @@ export class ShellToolInvocation extends BaseToolInvocation<
    const combinedController = new AbortController();

    const onAbort = () => combinedController.abort();
-
    try {
+      tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'gemini-shell-'));
+      tempFilePath = path.join(tempDir, 'pgrep.tmp');
+
      // pgrep is not available on Windows, so we can't get background PIDs
      const commandToExecute = this.wrapCommandForPgrep(
        strippedCommand,
@@ -638,7 +639,10 @@ export class ShellToolInvocation extends BaseToolInvocation<

        if (tempFileExists) {
          const pgrepContent = await fsPromises.readFile(tempFilePath, 'utf8');
-          const pgrepLines = pgrepContent.split(os.EOL).filter(Boolean);
+          const pgrepLines = pgrepContent
+            .split('\n')
+            .map((line) => line.trim())
+            .filter(Boolean);
          for (const line of pgrepLines) {
            if (!/^\d+$/.test(line)) {
              if (
@@ -935,10 +939,19 @@ export class ShellToolInvocation extends BaseToolInvocation<
      if (timeoutTimer) clearTimeout(timeoutTimer);
      signal.removeEventListener('abort', onAbort);
      timeoutController.signal.removeEventListener('abort', onAbort);
-      try {
-        await fsPromises.unlink(tempFilePath);
-      } catch {
-        // Ignore errors during unlink
+      if (tempFilePath) {
+        try {
+          await fsPromises.unlink(tempFilePath);
+        } catch {
+          // Ignore errors during unlink
+        }
+      }
+      if (tempDir) {
+        try {
+          await fsPromises.rm(tempDir, { recursive: true, force: true });
+        } catch {
+          // Ignore errors during rm
+        }
      }
    }
  }