From d68cc3a88f222f6445b5388ab20e70e06a3ca7bd Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Thu, 12 Mar 2026 14:59:32 -0400 Subject: [PATCH] feat(cli): make background task UI agnostic to execution type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add onBackground event to ExecutionLifecycleService that fires when any execution is moved to the background. The CLI subscribes to this event and automatically registers background tasks in the UI — no per-tool changes needed. Any tool that calls ExecutionLifecycleService.createExecution() or attachExecution() now automatically gets Ctrl+B support. Shell-specific concerns (PTY log files) stay in ShellExecutionService. Forward setExecutionIdCallback through SubAgentInvocation so agents can expose their execution ID to the scheduler for backgrounding. Route registerBackgroundTask and dismissBackgroundTask through ExecutionLifecycleService instead of ShellExecutionService for agnostic subscribe/onExit/kill support. --- .../ui/hooks/shellCommandProcessor.test.tsx | 71 ++++++----- .../cli/src/ui/hooks/shellCommandProcessor.ts | 112 +++++++++++++----- .../core/src/agents/subagent-tool.test.ts | 1 + packages/core/src/agents/subagent-tool.ts | 4 +- .../executionLifecycleService.test.ts | 75 ++++++++++++ .../src/services/executionLifecycleService.ts | 48 ++++++++ 6 files changed, 252 insertions(+), 59 deletions(-) diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.test.tsx b/packages/cli/src/ui/hooks/shellCommandProcessor.test.tsx index b8486bc378..9208d59153 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.test.tsx +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.test.tsx @@ -34,6 +34,23 @@ const mockShellOnExit = vi.hoisted(() => ) => () => void >(() => vi.fn()), ); +const mockLifecycleSubscribe = vi.hoisted(() => + vi.fn< + (pid: number, listener: (event: ShellOutputEvent) => void) => () => void + >(() => vi.fn()), +); +const mockLifecycleOnExit = vi.hoisted(() => + vi.fn< + ( + pid: number, + callback: (exitCode: number, signal?: number) => void, + ) => () => void + >(() => vi.fn()), +); +const mockLifecycleKill = vi.hoisted(() => vi.fn()); +const mockLifecycleBackground = vi.hoisted(() => vi.fn()); +const mockLifecycleOnBackground = vi.hoisted(() => vi.fn()); +const mockLifecycleOffBackground = vi.hoisted(() => vi.fn()); vi.mock('@google/gemini-cli-core', async (importOriginal) => { const actual = @@ -47,6 +64,14 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { subscribe: mockShellSubscribe, onExit: mockShellOnExit, }, + ExecutionLifecycleService: { + subscribe: mockLifecycleSubscribe, + onExit: mockLifecycleOnExit, + kill: mockLifecycleKill, + background: mockLifecycleBackground, + onBackground: mockLifecycleOnBackground, + offBackground: mockLifecycleOffBackground, + }, isBinary: mockIsBinary, }; }); @@ -777,8 +802,11 @@ describe('useShellCommandProcessor', () => { output: 'initial', }), ); - expect(mockShellOnExit).toHaveBeenCalledWith(1001, expect.any(Function)); - expect(mockShellSubscribe).toHaveBeenCalledWith( + expect(mockLifecycleOnExit).toHaveBeenCalledWith( + 1001, + expect.any(Function), + ); + expect(mockLifecycleSubscribe).toHaveBeenCalledWith( 1001, expect.any(Function), ); @@ -816,7 +844,7 @@ describe('useShellCommandProcessor', () => { expect(addItemToHistoryMock).toHaveBeenCalledWith( expect.objectContaining({ type: 'info', - text: 'No background shells are currently active.', + text: 'No background tasks are currently active.', }), expect.any(Number), ); @@ -834,7 +862,7 @@ describe('useShellCommandProcessor', () => { await result.current.dismissBackgroundShell(1001); }); - expect(mockShellKill).toHaveBeenCalledWith(1001); + expect(mockLifecycleKill).toHaveBeenCalledWith(1001); expect(result.current.backgroundShellCount).toBe(0); expect(result.current.backgroundShells.has(1001)).toBe(false); }); @@ -884,7 +912,7 @@ describe('useShellCommandProcessor', () => { expect(result.current.activeShellPtyId).toBeNull(); }); - it('should persist background shell on successful exit and mark as exited', async () => { + it('should auto-dismiss background task on successful exit', async () => { const { result } = renderProcessorHook(); act(() => { @@ -892,7 +920,7 @@ describe('useShellCommandProcessor', () => { }); // Find the exit callback registered - const exitCallback = mockShellOnExit.mock.calls.find( + const exitCallback = mockLifecycleOnExit.mock.calls.find( (call) => call[0] === 888, )?.[1]; expect(exitCallback).toBeDefined(); @@ -903,22 +931,19 @@ describe('useShellCommandProcessor', () => { }); } - // Should NOT be removed, but updated - expect(result.current.backgroundShellCount).toBe(0); // Badge count is 0 - expect(result.current.backgroundShells.has(888)).toBe(true); // Map has it - const shell = result.current.backgroundShells.get(888); - expect(shell?.status).toBe('exited'); - expect(shell?.exitCode).toBe(0); + // Should be auto-dismissed from the panel + expect(result.current.backgroundShellCount).toBe(0); + expect(result.current.backgroundShells.has(888)).toBe(false); }); - it('should persist background shell on failed exit', async () => { + it('should auto-dismiss background task on failed exit', async () => { const { result } = renderProcessorHook(); act(() => { result.current.registerBackgroundShell(999, 'fail-exit', ''); }); - const exitCallback = mockShellOnExit.mock.calls.find( + const exitCallback = mockLifecycleOnExit.mock.calls.find( (call) => call[0] === 999, )?.[1]; expect(exitCallback).toBeDefined(); @@ -929,17 +954,9 @@ describe('useShellCommandProcessor', () => { }); } - // Should NOT be removed, but updated - expect(result.current.backgroundShellCount).toBe(0); // Badge count is 0 - const shell = result.current.backgroundShells.get(999); - expect(shell?.status).toBe('exited'); - expect(shell?.exitCode).toBe(1); - - // Now dismiss it - await act(async () => { - await result.current.dismissBackgroundShell(999); - }); + // Should be auto-dismissed from the panel expect(result.current.backgroundShellCount).toBe(0); + expect(result.current.backgroundShells.has(999)).toBe(false); }); it('should NOT trigger re-render on background shell output when visible', async () => { @@ -956,7 +973,7 @@ describe('useShellCommandProcessor', () => { const initialRenderCount = getRenderCount(); - const subscribeCallback = mockShellSubscribe.mock.calls.find( + const subscribeCallback = mockLifecycleSubscribe.mock.calls.find( (call) => call[0] === 1001, )?.[1]; expect(subscribeCallback).toBeDefined(); @@ -982,7 +999,7 @@ describe('useShellCommandProcessor', () => { // Ensure background shells are hidden (default) const initialRenderCount = getRenderCount(); - const subscribeCallback = mockShellSubscribe.mock.calls.find( + const subscribeCallback = mockLifecycleSubscribe.mock.calls.find( (call) => call[0] === 1001, )?.[1]; expect(subscribeCallback).toBeDefined(); @@ -1012,7 +1029,7 @@ describe('useShellCommandProcessor', () => { const initialRenderCount = getRenderCount(); - const subscribeCallback = mockShellSubscribe.mock.calls.find( + const subscribeCallback = mockLifecycleSubscribe.mock.calls.find( (call) => call[0] === 1001, )?.[1]; expect(subscribeCallback).toBeDefined(); diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts index 7e33d37d1f..2fb6910b0a 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts @@ -13,6 +13,7 @@ import type { AnsiOutput, Config, GeminiClient } from '@google/gemini-cli-core'; import { isBinary, ShellExecutionService, + ExecutionLifecycleService, CoreToolCallStatus, } from '@google/gemini-cli-core'; import { type PartListUnion } from '@google/genai'; @@ -144,7 +145,7 @@ export const useShellCommandProcessor = ( useEffect( () => () => { - // Unsubscribe from all background shell events on unmount + // Unsubscribe from all background task events on unmount for (const unsubscribe of m.subscriptions.values()) { unsubscribe(); } @@ -176,7 +177,7 @@ export const useShellCommandProcessor = ( addItemToHistory( { type: 'info', - text: 'No background shells are currently active.', + text: 'No background tasks are currently active.', }, Date.now(), ); @@ -191,12 +192,19 @@ export const useShellCommandProcessor = ( dispatch, ]); - const backgroundCurrentShell = useCallback(() => { + const backgroundCurrentExecution = useCallback(() => { const pidToBackground = state.activeShellPtyId ?? activeBackgroundExecutionId; if (pidToBackground) { - ShellExecutionService.background(pidToBackground); + // Use ShellExecutionService for shell PTYs (handles log files, etc.), + // fall back to ExecutionLifecycleService for non-shell executions + // (e.g. remote agents, MCP tools, local agents). m.backgroundedPids.add(pidToBackground); + if (state.activeShellPtyId) { + ShellExecutionService.background(pidToBackground); + } else { + ExecutionLifecycleService.background(pidToBackground); + } // Ensure backgrounding is silent and doesn't trigger restoration m.wasVisibleBeforeForeground = false; if (m.restoreTimeout) { @@ -206,12 +214,14 @@ export const useShellCommandProcessor = ( } }, [state.activeShellPtyId, activeBackgroundExecutionId, m]); - const dismissBackgroundShell = useCallback( + const dismissBackgroundTask = useCallback( async (pid: number) => { const shell = state.backgroundShells.get(pid); if (shell) { if (shell.status === 'running') { - await ShellExecutionService.kill(pid); + // ExecutionLifecycleService.kill handles both shell and non-shell + // executions. For shells, ShellExecutionService.kill delegates to it. + ExecutionLifecycleService.kill(pid); } dispatch({ type: 'DISMISS_SHELL', pid }); m.backgroundedPids.delete(pid); @@ -227,37 +237,55 @@ export const useShellCommandProcessor = ( [state.backgroundShells, dispatch, m], ); - const registerBackgroundShell = useCallback( + const registerBackgroundTask = useCallback( (pid: number, command: string, initialOutput: string | AnsiOutput) => { dispatch({ type: 'REGISTER_SHELL', pid, command, initialOutput }); - // Subscribe to process exit directly - const exitUnsubscribe = ShellExecutionService.onExit(pid, (code) => { + // Subscribe to exit via ExecutionLifecycleService (works for all execution types) + const exitUnsubscribe = ExecutionLifecycleService.onExit(pid, (code) => { dispatch({ type: 'UPDATE_SHELL', pid, update: { status: 'exited', exitCode: code }, }); + // Auto-dismiss completed tasks from the background panel. + dispatch({ type: 'DISMISS_SHELL', pid }); + const unsub = m.subscriptions.get(pid); + if (unsub) { + unsub(); + m.subscriptions.delete(pid); + } m.backgroundedPids.delete(pid); }); - // Subscribe to future updates (data only) - const dataUnsubscribe = ShellExecutionService.subscribe(pid, (event) => { - if (event.type === 'data') { - dispatch({ type: 'APPEND_SHELL_OUTPUT', pid, chunk: event.chunk }); - } else if (event.type === 'binary_detected') { - dispatch({ type: 'UPDATE_SHELL', pid, update: { isBinary: true } }); - } else if (event.type === 'binary_progress') { - dispatch({ - type: 'UPDATE_SHELL', - pid, - update: { - isBinary: true, - binaryBytesReceived: event.bytesReceived, - }, - }); - } - }); + // Subscribe to output via ExecutionLifecycleService (works for all execution types) + const dataUnsubscribe = ExecutionLifecycleService.subscribe( + pid, + (event) => { + if (event.type === 'data') { + dispatch({ + type: 'APPEND_SHELL_OUTPUT', + pid, + chunk: event.chunk, + }); + } else if (event.type === 'binary_detected') { + dispatch({ + type: 'UPDATE_SHELL', + pid, + update: { isBinary: true }, + }); + } else if (event.type === 'binary_progress') { + dispatch({ + type: 'UPDATE_SHELL', + pid, + update: { + isBinary: true, + binaryBytesReceived: event.bytesReceived, + }, + }); + } + }, + ); m.subscriptions.set(pid, () => { exitUnsubscribe(); @@ -267,6 +295,28 @@ export const useShellCommandProcessor = ( [dispatch, m], ); + // Auto-register any execution that gets backgrounded, regardless of type. + // This is the agnostic hook: any tool that calls + // ExecutionLifecycleService.createExecution() or attachExecution() + // automatically gets Ctrl+B support — no UI changes needed per tool. + useEffect(() => { + const listener = (info: { + executionId: number; + label: string; + output: string; + }) => { + // Skip if already registered (e.g. shells register via their own flow) + if (m.backgroundedPids.has(info.executionId)) { + return; + } + registerBackgroundTask(info.executionId, info.label, info.output); + }; + ExecutionLifecycleService.onBackground(listener); + return () => { + ExecutionLifecycleService.offBackground(listener); + }; + }, [registerBackgroundTask, m]); + const handleShellCommand = useCallback( (rawQuery: PartListUnion, abortSignal: AbortSignal): boolean => { if (typeof rawQuery !== 'string' || rawQuery.trim() === '') { @@ -439,7 +489,7 @@ export const useShellCommandProcessor = ( setPendingHistoryItem(null); if (result.backgrounded && result.pid) { - registerBackgroundShell(result.pid, rawQuery, cumulativeStdout); + registerBackgroundTask(result.pid, rawQuery, cumulativeStdout); dispatch({ type: 'SET_ACTIVE_PTY', pid: null }); } @@ -531,7 +581,7 @@ export const useShellCommandProcessor = ( setShellInputFocused, terminalHeight, terminalWidth, - registerBackgroundShell, + registerBackgroundTask, m, dispatch, ], @@ -548,9 +598,9 @@ export const useShellCommandProcessor = ( backgroundShellCount, isBackgroundShellVisible: state.isBackgroundShellVisible, toggleBackgroundShell, - backgroundCurrentShell, - registerBackgroundShell, - dismissBackgroundShell, + backgroundCurrentShell: backgroundCurrentExecution, + registerBackgroundShell: registerBackgroundTask, + dismissBackgroundShell: dismissBackgroundTask, backgroundShells: state.backgroundShells, }; }; diff --git a/packages/core/src/agents/subagent-tool.test.ts b/packages/core/src/agents/subagent-tool.test.ts index 438df59cd3..6f6bc490fa 100644 --- a/packages/core/src/agents/subagent-tool.test.ts +++ b/packages/core/src/agents/subagent-tool.test.ts @@ -189,6 +189,7 @@ describe('SubAgentInvocation', () => { expect(mockInnerInvocation.execute).toHaveBeenCalledWith( abortSignal, updateOutput, + undefined, ); expect(runInDevTraceSpan).toHaveBeenCalledWith( diff --git a/packages/core/src/agents/subagent-tool.ts b/packages/core/src/agents/subagent-tool.ts index 0c4f19ee8b..34daf3b141 100644 --- a/packages/core/src/agents/subagent-tool.ts +++ b/packages/core/src/agents/subagent-tool.ts @@ -13,6 +13,7 @@ import { type ToolCallConfirmationDetails, isTool, type ToolLiveOutput, + type ExecuteOptions, } from '../tools/tools.js'; import type { Config } from '../config/config.js'; import { type AgentLoopContext } from '../config/agent-loop-context.js'; @@ -161,6 +162,7 @@ class SubAgentInvocation extends BaseToolInvocation { async execute( signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, + options?: ExecuteOptions, ): Promise { const validationError = SchemaValidator.validate( this.definition.inputConfig.inputSchema, @@ -188,7 +190,7 @@ class SubAgentInvocation extends BaseToolInvocation { }, async ({ metadata }) => { metadata.input = this.params; - const result = await invocation.execute(signal, updateOutput); + const result = await invocation.execute(signal, updateOutput, options); metadata.output = result; return result; }, diff --git a/packages/core/src/services/executionLifecycleService.test.ts b/packages/core/src/services/executionLifecycleService.test.ts index 0d800c6e55..b5b78a8e86 100644 --- a/packages/core/src/services/executionLifecycleService.test.ts +++ b/packages/core/src/services/executionLifecycleService.test.ts @@ -296,6 +296,81 @@ describe('ExecutionLifecycleService', () => { }).toThrow('Execution 4324 is already attached.'); }); + describe('Background Start Listeners', () => { + it('fires onBackground when an execution is backgrounded', async () => { + const listener = vi.fn(); + ExecutionLifecycleService.onBackground(listener); + + const handle = ExecutionLifecycleService.createExecution( + '', + undefined, + 'remote_agent', + undefined, + 'My Remote Agent', + ); + const executionId = handle.pid!; + + ExecutionLifecycleService.appendOutput(executionId, 'some output'); + ExecutionLifecycleService.background(executionId); + await handle.result; + + expect(listener).toHaveBeenCalledTimes(1); + const info = listener.mock.calls[0][0]; + expect(info.executionId).toBe(executionId); + expect(info.executionMethod).toBe('remote_agent'); + expect(info.label).toBe('My Remote Agent'); + expect(info.output).toBe('some output'); + + ExecutionLifecycleService.offBackground(listener); + }); + + it('uses fallback label when none is provided', async () => { + const listener = vi.fn(); + ExecutionLifecycleService.onBackground(listener); + + const handle = ExecutionLifecycleService.createExecution( + '', + undefined, + 'none', + ); + const executionId = handle.pid!; + + ExecutionLifecycleService.background(executionId); + await handle.result; + + const info = listener.mock.calls[0][0]; + expect(info.label).toContain('none'); + expect(info.label).toContain(String(executionId)); + + ExecutionLifecycleService.offBackground(listener); + }); + + it('does not fire onBackground for non-backgrounded completions', async () => { + const listener = vi.fn(); + ExecutionLifecycleService.onBackground(listener); + + const handle = ExecutionLifecycleService.createExecution(); + ExecutionLifecycleService.completeExecution(handle.pid!); + await handle.result; + + expect(listener).not.toHaveBeenCalled(); + + ExecutionLifecycleService.offBackground(listener); + }); + + it('offBackground removes the listener', async () => { + const listener = vi.fn(); + ExecutionLifecycleService.onBackground(listener); + ExecutionLifecycleService.offBackground(listener); + + const handle = ExecutionLifecycleService.createExecution(); + ExecutionLifecycleService.background(handle.pid!); + await handle.result; + + expect(listener).not.toHaveBeenCalled(); + }); + }); + describe('Background Completion Listeners', () => { it('fires onBackgroundComplete with formatInjection text when backgrounded execution settles', async () => { const listener = vi.fn(); diff --git a/packages/core/src/services/executionLifecycleService.ts b/packages/core/src/services/executionLifecycleService.ts index 6df693fccb..a0212bbed5 100644 --- a/packages/core/src/services/executionLifecycleService.ts +++ b/packages/core/src/services/executionLifecycleService.ts @@ -59,6 +59,8 @@ export interface ExecutionCompletionOptions { export interface ExternalExecutionRegistration { executionMethod: ExecutionMethod; + /** Human-readable label for the background task UI (e.g. the command string). */ + label?: string; initialOutput?: string; getBackgroundOutput?: () => string; getSubscriptionSnapshot?: () => string | AnsiOutput | undefined; @@ -79,6 +81,7 @@ export type FormatInjectionFn = ( interface ManagedExecutionBase { executionMethod: ExecutionMethod; + label?: string; output: string; backgrounded?: boolean; formatInjection?: FormatInjectionFn; @@ -86,6 +89,18 @@ interface ManagedExecutionBase { getSubscriptionSnapshot?: () => string | AnsiOutput | undefined; } +/** + * Payload emitted when an execution is moved to the background. + */ +export interface BackgroundStartInfo { + executionId: number; + executionMethod: ExecutionMethod; + label: string; + output: string; +} + +export type BackgroundStartListener = (info: BackgroundStartInfo) => void; + /** * Payload emitted when a previously-backgrounded execution settles. */ @@ -150,6 +165,23 @@ export class ExecutionLifecycleService { this.injectionService = service; } + private static backgroundStartListeners = new Set(); + + /** + * Registers a listener that fires when any execution is moved to the background. + * This is the hook for the UI to automatically discover backgrounded executions. + */ + static onBackground(listener: BackgroundStartListener): void { + this.backgroundStartListeners.add(listener); + } + + /** + * Unregisters a background start listener. + */ + static offBackground(listener: BackgroundStartListener): void { + this.backgroundStartListeners.delete(listener); + } + /** * Registers a listener that fires when a previously-backgrounded * execution settles (completes or errors). @@ -222,6 +254,7 @@ export class ExecutionLifecycleService { this.exitedExecutionInfo.clear(); this.backgroundCompletionListeners.clear(); this.injectionService = null; + this.backgroundStartListeners.clear(); this.nextExecutionId = NON_PROCESS_EXECUTION_ID_START; } @@ -239,6 +272,7 @@ export class ExecutionLifecycleService { this.activeExecutions.set(executionId, { executionMethod: registration.executionMethod, + label: registration.label, output: registration.initialOutput ?? '', kind: 'external', getBackgroundOutput: registration.getBackgroundOutput, @@ -259,11 +293,13 @@ export class ExecutionLifecycleService { onKill?: () => void, executionMethod: ExecutionMethod = 'none', formatInjection?: FormatInjectionFn, + label?: string, ): ExecutionHandle { const executionId = this.allocateExecutionId(); this.activeExecutions.set(executionId, { executionMethod, + label, output: initialOutput, kind: 'virtual', onKill, @@ -434,6 +470,18 @@ export class ExecutionLifecycleService { this.activeResolvers.delete(executionId); execution.backgrounded = true; + + // Notify listeners that an execution was moved to the background. + const info: BackgroundStartInfo = { + executionId, + executionMethod: execution.executionMethod, + label: + execution.label ?? `${execution.executionMethod} (ID: ${executionId})`, + output, + }; + for (const listener of this.backgroundStartListeners) { + listener(info); + } } static subscribe(