diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts index 1948583fe16..548f4920ef3 100644 --- a/packages/core/src/services/shellExecutionService.test.ts +++ b/packages/core/src/services/shellExecutionService.test.ts @@ -15,6 +15,7 @@ import { } from 'vitest'; import os from 'node:os'; +import path from 'node:path'; import EventEmitter from 'node:events'; import type { Readable } from 'node:stream'; import { type ChildProcess } from 'node:child_process'; @@ -999,6 +1000,77 @@ describe('ShellExecutionService', () => { ShellExecutionService.listBackgroundProcesses(undefined as any), ).toThrow('Session ID is required'); }); + + it('should accept tempDir in background() and delete it when the background process exits', async () => { + const actualFs = + await vi.importActual('node:fs'); + const actualOs = + await vi.importActual('node:os'); + const tempDir = actualFs.mkdtempSync( + path.join(actualOs.tmpdir(), 'gemini-shell-bg-unit-'), + ); + actualFs.writeFileSync(path.join(tempDir, 'bgpids.tmp'), '123\n'); + + let triggerExit: + | ((args: { exitCode: number; signal?: number }) => void) + | undefined; + + await simulateExecution('sleep 1', (pty) => { + triggerExit = pty.onExit.mock.calls[0][0]; + + ShellExecutionService.background( + pty.pid, + 'default', + 'sleep 1', + tempDir, + ); + }); + + const history = ( + ShellExecutionService as unknown as { + backgroundProcessHistory: Map< + string, + Map + >; + } + ).backgroundProcessHistory.get('default'); + expect(history?.get(12345)?.tempDir).toBe(tempDir); + expect(actualFs.existsSync(tempDir)).toBe(true); + + triggerExit?.({ exitCode: 0 }); + + await vi.waitFor(() => { + expect(actualFs.existsSync(tempDir)).toBe(false); + }); + }); + + it('should return early and clean up tempDir if background() is called for an untracked or already-exited process', async () => { + const actualFs = + await vi.importActual('node:fs'); + const actualOs = + await vi.importActual('node:os'); + const tempDir = actualFs.mkdtempSync( + path.join(actualOs.tmpdir(), 'gemini-shell-bg-untracked-'), + ); + + ShellExecutionService.background( + 99999, + 'default', + 'echo exited', + tempDir, + ); + + const backgroundTempDirs = ( + ShellExecutionService as unknown as { + backgroundTempDirs: Map; + } + ).backgroundTempDirs; + expect(backgroundTempDirs.has(99999)).toBe(false); + + await vi.waitFor(() => { + expect(actualFs.existsSync(tempDir)).toBe(false); + }); + }); }); describe('Binary Output', () => { diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index ca80036a6c1..a07eca843e1 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/packages/core/src/services/shellExecutionService.ts @@ -11,6 +11,7 @@ import { TextDecoder } from 'node:util'; import type { Writable } from 'node:stream'; import os from 'node:os'; import fs, { mkdirSync } from 'node:fs'; +import fsPromises from 'node:fs/promises'; import path from 'node:path'; import type { IPty } from '@lydell/node-pty'; import { @@ -142,6 +143,7 @@ export interface ShellExecutionConfig { originalCommand?: string; sessionId?: string; env?: Record; + tempDir?: string; } /** @@ -187,6 +189,7 @@ interface ActivePty { command: string; sessionId?: string; cancelRender?: () => void; + tempDir?: string; } interface ActiveChildProcess { @@ -199,6 +202,7 @@ interface ActiveChildProcess { }; command: string; sessionId?: string; + tempDir?: string; } const isAnsiOutputEqual = ( @@ -369,6 +373,7 @@ export type BackgroundProcess = { export type BackgroundProcessRecord = Omit & { startTime: number; endTime?: number; + tempDir?: string; }; export class ShellExecutionService { @@ -376,6 +381,7 @@ export class ShellExecutionService { private static activeChildProcesses = new Map(); private static backgroundLogPids = new Set(); private static backgroundLogStreams = new Map(); + private static backgroundTempDirs = new Map(); private static backgroundProcessHistory = new Map< string, // sessionId Map @@ -417,6 +423,12 @@ export class ShellExecutionService { } private static async cleanupLogStream(pid: number): Promise { + const tempDir = this.backgroundTempDirs.get(pid); + this.backgroundTempDirs.delete(pid); + const rmPromise = tempDir + ? fsPromises.rm(tempDir, { recursive: true, force: true }).catch(() => {}) + : Promise.resolve(); + const stream = this.backgroundLogStreams.get(pid); if (stream) { await new Promise((resolve) => { @@ -426,6 +438,7 @@ export class ShellExecutionService { } this.backgroundLogPids.delete(pid); + await rmPromise; } /** @@ -704,6 +717,7 @@ export class ShellExecutionService { state, command: shellExecutionConfig.originalCommand ?? commandToExecute, sessionId: shellExecutionConfig.sessionId, + tempDir: shellExecutionConfig.tempDir, }); } @@ -1246,6 +1260,7 @@ export class ShellExecutionService { command: shellExecutionConfig.originalCommand ?? commandToExecute, sessionId: shellExecutionConfig.sessionId, cancelRender, + tempDir: shellExecutionConfig.tempDir, }); const result = ExecutionLifecycleService.attachExecution(assignedPid, { @@ -1811,15 +1826,22 @@ export class ShellExecutionService { * This resolves the execution promise but keeps the PTY active. * * @param pid The process ID of the target PTY. + * @param sessionId Optional session ID for process history tracking. + * @param command Optional command string for display. + * @param tempDir Optional temporary directory path owned by this execution to clean up on exit. */ - static background(pid: number, sessionId?: string, command?: string): void { - if (this.backgroundLogPids.has(pid)) { - return; - } - + static background( + pid: number, + sessionId?: string, + command?: string, + tempDir?: string, + ): void { const activePty = this.activePtys.get(pid); const activeChild = this.activeChildProcesses.get(pid); + const resolvedTempDir = + tempDir ?? activePty?.tempDir ?? activeChild?.tempDir; + const resolvedSessionId = sessionId ?? activePty?.sessionId ?? activeChild?.sessionId; const resolvedCommand = @@ -1832,20 +1854,27 @@ export class ShellExecutionService { throw new Error('Session ID is required for background operations'); } + if (!activePty && !activeChild) { + if (resolvedTempDir) { + fsPromises + .rm(resolvedTempDir, { recursive: true, force: true }) + .catch(() => {}); + } + return; + } + + if (resolvedTempDir) { + this.backgroundTempDirs.set(pid, resolvedTempDir); + } + + if (this.backgroundLogPids.has(pid)) { + return; + } + const MAX_BACKGROUND_PROCESS_HISTORY_SIZE = 100; const history = this.backgroundProcessHistory.get(resolvedSessionId) ?? - new Map< - number, - { - command: string; - status: 'running' | 'exited'; - exitCode?: number | null; - signal?: number | null; - startTime: number; - endTime?: number; - } - >(); + new Map(); if (history.size >= MAX_BACKGROUND_PROCESS_HISTORY_SIZE) { const oldestPid = history.keys().next().value; @@ -1858,6 +1887,7 @@ export class ShellExecutionService { command: resolvedCommand, status: 'running', startTime: Date.now(), + ...(resolvedTempDir ? { tempDir: resolvedTempDir } : {}), }); this.backgroundProcessHistory.set(resolvedSessionId, history); @@ -2045,8 +2075,16 @@ export class ShellExecutionService { // ignored } } + for (const tempDir of this.backgroundTempDirs.values()) { + try { + fs.rmSync(tempDir, { recursive: true, force: true }); + } catch { + // ignored + } + } this.backgroundLogPids.clear(); this.backgroundLogStreams.clear(); + this.backgroundTempDirs.clear(); this.backgroundProcessHistory.clear(); } } diff --git a/packages/core/src/tools/shell.test.ts b/packages/core/src/tools/shell.test.ts index 378f25adc79..bb9adb3581e 100644 --- a/packages/core/src/tools/shell.test.ts +++ b/packages/core/src/tools/shell.test.ts @@ -247,6 +247,12 @@ describe('ShellTool', () => { }); afterEach(() => { + if (extractedTmpFile) { + const extractedDir = path.dirname(extractedTmpFile); + if (fs.existsSync(extractedDir)) { + fs.rmSync(extractedDir, { recursive: true, force: true }); + } + } if (fs.existsSync(tempRootDir)) { fs.rmSync(tempRootDir, { recursive: true, force: true }); } @@ -482,16 +488,20 @@ describe('ShellTool', () => { await vi.advanceTimersByTimeAsync(250); + const expectedTempDir = path.dirname(extractedTmpFile); expect(mockShellBackground).toHaveBeenCalledWith( 12345, 'default', 'sleep 10', + expectedTempDir, ); await promise; + // Ownership was transferred to ShellExecutionService, so shell.ts should not delete it prematurely + expect(fs.existsSync(expectedTempDir)).toBe(true); }); - it('should cancel the promotion timer when the command completes before the delay elapses', async () => { + it('should cancel the promotion timer and clean up tempDir when the command completes before the delay elapses', async () => { vi.useFakeTimers(); const invocation = shellTool.build({ command: 'echo done', @@ -506,6 +516,27 @@ describe('ShellTool', () => { expect(mockShellBackground).not.toHaveBeenCalled(); await promise; + const expectedTempDir = path.dirname(extractedTmpFile); + expect(fs.existsSync(expectedTempDir)).toBe(false); + }); + + it('should clean up tempDir in finally if ShellExecutionService.background throws an error', async () => { + vi.useFakeTimers(); + mockShellBackground.mockImplementationOnce(() => { + throw new Error('Background failed'); + }); + + const invocation = shellTool.build({ + command: 'sleep 10', + is_background: true, + }); + const promise = invocation.execute({ abortSignal: mockAbortSignal }); + + await vi.advanceTimersByTimeAsync(250); + await promise; + + const expectedTempDir = path.dirname(extractedTmpFile); + expect(fs.existsSync(expectedTempDir)).toBe(false); }); itWindowsOnly( @@ -963,6 +994,7 @@ EOF`; 12345, 'default', 'sleep 10', + path.dirname(extractedTmpFile), ); await promise; diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 3b9c0ed5ba0..0b4ab4ad421 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -551,6 +551,7 @@ export class ShellToolInvocation extends BaseToolInvocation< const isWindows = os.platform() === 'win32'; let tempFilePath = ''; let tempDir = ''; + let isBackgrounded = false; const timeoutMs = this.context.config.getShellToolInactivityTimeout(); const timeoutController = new AbortController(); @@ -749,6 +750,7 @@ export class ShellToolInvocation extends BaseToolInvocation< backgroundCompletionBehavior: this.context.config.getShellBackgroundCompletionBehavior(), originalCommand: strippedCommand, + tempDir, }, ); @@ -765,11 +767,20 @@ export class ShellToolInvocation extends BaseToolInvocation< () => { promotionTimer = null; if (!completed) { - ShellExecutionService.background( - pid, - sessionId, - strippedCommand, - ); + try { + ShellExecutionService.background( + pid, + sessionId, + strippedCommand, + tempDir, + ); + isBackgrounded = true; + } catch (err) { + debugLogger.error( + 'Failed to background shell execution:', + err, + ); + } } }, delay, @@ -813,7 +824,9 @@ export class ShellToolInvocation extends BaseToolInvocation< } const result = await resultPromise; - if (!result.backgrounded) { + if (result.backgrounded) { + isBackgrounded = true; + } else { flushOutput(); } @@ -1154,9 +1167,9 @@ export class ShellToolInvocation extends BaseToolInvocation< timeoutController.signal.removeEventListener('abort', onAbort); // Only clean up if NOT running in background. - // Background processes need the temp directory and PID file to remain - // available until they exit. - if (!this.params.is_background) { + // Background processes transfer ownership of the temp directory to + // ShellExecutionService, which removes it once the process exits. + if (!isBackgrounded) { if (tempFilePath) { try { await fsPromises.unlink(tempFilePath); diff --git a/packages/core/src/tools/shellBackgroundTools.integration.test.ts b/packages/core/src/tools/shellBackgroundTools.integration.test.ts index ab96df73831..d46a620b51a 100644 --- a/packages/core/src/tools/shellBackgroundTools.integration.test.ts +++ b/packages/core/src/tools/shellBackgroundTools.integration.test.ts @@ -10,6 +10,7 @@ import { ListBackgroundProcessesTool, ReadBackgroundOutputTool, } from './shellBackgroundTools.js'; +import { ShellTool } from './shell.js'; import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; import { NoopSandboxManager } from '../services/sandboxManager.js'; import type { AgentLoopContext } from '../config/agent-loop-context.js'; @@ -120,4 +121,69 @@ describe('Background Tools Integration', () => { await ShellExecutionService.kill(pid); controller.abort(); }); + + it('should delete the temporary directory when a short-lived background shell command exits', async () => { + const scriptPath = path.join(tempRootDir, 'short-bg.js'); + fs.writeFileSync(scriptPath, 'setTimeout(() => process.exit(0), 400);'); + + const mockContext = { + config: { + getSessionId: () => 'default', + getTargetDir: () => tempRootDir, + validatePathAccess: () => null, + getShellToolInactivityTimeout: () => 5000, + isInteractiveShellEnabled: () => false, + getEnableShellOutputEfficiency: () => true, + getSandboxEnabled: () => false, + getShellBackgroundCompletionBehavior: () => 'silent', + getSummarizeToolOutputConfig: () => undefined, + getDebugMode: () => false, + sanitizationConfig: { + allowedEnvironmentVariables: [], + blockedEnvironmentVariables: [], + enableEnvironmentVariableRedaction: false, + }, + sandboxManager: new NoopSandboxManager(), + }, + } as unknown as AgentLoopContext; + + const shellTool = new ShellTool(mockContext, bus); + const invocation = shellTool.build({ + command: `node "${scriptPath}"`, + is_background: true, + delay_ms: 100, + }); + + let assignedPid: number | undefined; + const result = await invocation.execute({ + abortSignal: new AbortController().signal, + setExecutionIdCallback: (pid) => { + assignedPid = pid; + }, + }); + + expect(result.llmContent).toContain('Command moved to background'); + expect(assignedPid).toBeDefined(); + + const history = ( + ShellExecutionService as unknown as { + backgroundProcessHistory: Map< + string, + Map + >; + } + ).backgroundProcessHistory.get('default'); + const record = history?.get(assignedPid!); + expect(record?.tempDir).toBeDefined(); + expect(record?.tempDir).toContain('gemini-shell-'); + expect(fs.existsSync(record!.tempDir!)).toBe(true); + + await vi.waitFor( + () => { + expect(record?.status).toBe('exited'); + expect(fs.existsSync(record!.tempDir!)).toBe(false); + }, + { timeout: 5000 }, + ); + }); });