From a3b557222a95965de39920922df21439d69def57 Mon Sep 17 00:00:00 2001 From: Olcan Date: Fri, 30 May 2025 01:35:03 -0700 Subject: [PATCH] tweaks to shell abort logic based on feedback (#618) --- .../ui/hooks/shellCommandProcessor.test.ts | 22 ++++++---- .../cli/src/ui/hooks/shellCommandProcessor.ts | 44 ++++++------------- packages/cli/src/ui/hooks/useGeminiStream.ts | 12 ++--- packages/server/src/tools/shell.ts | 17 +++---- 4 files changed, 38 insertions(+), 57 deletions(-) diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts index 92bc98ea..2fd32087 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.test.ts @@ -105,9 +105,10 @@ describe('useShellCommandProcessor', () => { it('should return false for non-string input', () => { const { result } = setupHook(); - const handled = result.current.handleShellCommand([ - { text: 'not a string' }, - ] as PartListUnion); + const handled = result.current.handleShellCommand( + [{ text: 'not a string' }] as PartListUnion, + new AbortController().signal, + ); expect(handled).toBe(false); expect(mockAddItemToHistory).not.toHaveBeenCalled(); }); @@ -115,7 +116,10 @@ describe('useShellCommandProcessor', () => { it('should handle empty shell command', () => { const { result } = setupHook(); act(() => { - const handled = result.current.handleShellCommand(''); + const handled = result.current.handleShellCommand( + '', + new AbortController().signal, + ); expect(handled).toBe(true); }); expect(mockAddItemToHistory).toHaveBeenNthCalledWith( @@ -144,7 +148,7 @@ describe('useShellCommandProcessor', () => { existsSyncSpy.mockReturnValue(false); await act(async () => { - result.current.handleShellCommand(command); + result.current.handleShellCommand(command, new AbortController().signal); await new Promise(process.nextTick); }); @@ -182,7 +186,7 @@ describe('useShellCommandProcessor', () => { existsSyncSpy.mockReturnValue(false); await act(async () => { - result.current.handleShellCommand(command); + result.current.handleShellCommand(command, new AbortController().signal); await new Promise(process.nextTick); }); expect(mockAddItemToHistory).toHaveBeenNthCalledWith( @@ -210,7 +214,7 @@ describe('useShellCommandProcessor', () => { existsSyncSpy.mockReturnValue(false); await act(async () => { - result.current.handleShellCommand(command); + result.current.handleShellCommand(command, new AbortController().signal); await new Promise(process.nextTick); }); expect(mockAddItemToHistory).toHaveBeenNthCalledWith( @@ -237,7 +241,7 @@ describe('useShellCommandProcessor', () => { existsSyncSpy.mockReturnValue(false); await act(async () => { - result.current.handleShellCommand(command); + result.current.handleShellCommand(command, new AbortController().signal); await new Promise(process.nextTick); }); expect(mockAddItemToHistory).toHaveBeenNthCalledWith( @@ -264,7 +268,7 @@ describe('useShellCommandProcessor', () => { existsSyncSpy.mockReturnValue(false); await act(async () => { - result.current.handleShellCommand(command); + result.current.handleShellCommand(command, new AbortController().signal); await new Promise(process.nextTick); }); diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts index d6c837d1..1325fd1d 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts @@ -37,7 +37,7 @@ export const useShellCommandProcessor = ( * @returns True if the query was handled as a shell command, false otherwise. */ const handleShellCommand = useCallback( - (rawQuery: PartListUnion, signal?: AbortSignal): boolean => { + (rawQuery: PartListUnion, abortSignal: AbortSignal): boolean => { if (typeof rawQuery !== 'string') { return false; } @@ -149,27 +149,19 @@ export const useShellCommandProcessor = ( error = err; }); - const abortHandler = () => { + const abortHandler = async () => { if (child.pid && !exited) { onDebugMessage( `Aborting shell command (PID: ${child.pid}) due to signal.`, ); try { // attempt to SIGTERM process group (negative PID) - // if SIGTERM fails after 200ms, attempt SIGKILL + // fall back to SIGKILL (to group) after 200ms process.kill(-child.pid, 'SIGTERM'); - setTimeout(() => { - // if SIGTERM fails, attempt SIGKILL - try { - if (child.pid && !exited) { - process.kill(-child.pid, 'SIGKILL'); - } - } catch (_e) { - console.error( - `failed to kill shell process ${child.pid}: ${_e}`, - ); - } - }, 200); + await new Promise((resolve) => setTimeout(resolve, 200)); + if (child.pid && !exited) { + process.kill(-child.pid, 'SIGKILL'); + } } catch (_e) { // if group kill fails, fall back to killing just the main process try { @@ -185,20 +177,11 @@ export const useShellCommandProcessor = ( } }; - if (signal) { - if (signal.aborted) { - abortHandler(); - // No need to add listener if already aborted - } else { - signal.addEventListener('abort', abortHandler, { once: true }); - } - } + abortSignal.addEventListener('abort', abortHandler, { once: true }); - child.on('exit', (code, processSignal) => { + child.on('exit', (code, signal) => { exited = true; - if (signal) { - signal.removeEventListener('abort', abortHandler); - } + abortSignal.removeEventListener('abort', abortHandler); setPendingHistoryItem(null); output = output.trim() || '(Command produced no output)'; if (error) { @@ -207,7 +190,7 @@ export const useShellCommandProcessor = ( } else if (code !== null && code !== 0) { const text = `Command exited with code ${code}\n${output}`; addItemToHistory({ type: 'error', text }, userMessageTimestamp); - } else if (signal?.aborted) { + } else if (abortSignal.aborted) { addItemToHistory( { type: 'info', @@ -215,9 +198,8 @@ export const useShellCommandProcessor = ( }, userMessageTimestamp, ); - } else if (processSignal) { - const text = `Command terminated with signal ${processSignal} -${output}`; + } else if (signal) { + const text = `Command terminated with signal ${signal}.\n${output}`; addItemToHistory({ type: 'error', text }, userMessageTimestamp); } else { addItemToHistory( diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 66fb317a..9dd6f728 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -161,7 +161,7 @@ export const useGeminiStream = ( async ( query: PartListUnion, userMessageTimestamp: number, - signal: AbortSignal, + abortSignal: AbortSignal, ): Promise<{ queryToSend: PartListUnion | null; shouldProceed: boolean; @@ -199,7 +199,7 @@ export const useGeminiStream = ( return { queryToSend: null, shouldProceed: false }; // Handled by scheduling the tool } - if (shellModeActive && handleShellCommand(trimmedQuery, signal)) { + if (shellModeActive && handleShellCommand(trimmedQuery, abortSignal)) { return { queryToSend: null, shouldProceed: false }; } @@ -211,7 +211,7 @@ export const useGeminiStream = ( addItem, onDebugMessage, messageId: userMessageTimestamp, - signal, + signal: abortSignal, }); if (!atCommandResult.shouldProceed) { return { queryToSend: null, shouldProceed: false }; @@ -493,12 +493,12 @@ export const useGeminiStream = ( setShowHelp(false); abortControllerRef.current = new AbortController(); - const signal = abortControllerRef.current.signal; + const abortSignal = abortControllerRef.current.signal; const { queryToSend, shouldProceed } = await prepareQueryForGemini( query, userMessageTimestamp, - signal, + abortSignal, ); if (!shouldProceed || queryToSend === null) { @@ -515,7 +515,7 @@ export const useGeminiStream = ( setInitError(null); try { - const stream = client.sendMessageStream(chat, queryToSend, signal); + const stream = client.sendMessageStream(chat, queryToSend, abortSignal); const processingStatus = await processGeminiStreamEvents( stream, userMessageTimestamp, diff --git a/packages/server/src/tools/shell.ts b/packages/server/src/tools/shell.ts index a4634a3e..b43b2d0a 100644 --- a/packages/server/src/tools/shell.ts +++ b/packages/server/src/tools/shell.ts @@ -200,21 +200,16 @@ export class ShellTool extends BaseTool { }; shell.on('exit', exitHandler); - const abortHandler = () => { + const abortHandler = async () => { if (shell.pid && !exited) { try { // attempt to SIGTERM process group (negative PID) - // if SIGTERM fails after 200ms, attempt SIGKILL + // fall back to SIGKILL (to group) after 200ms process.kill(-shell.pid, 'SIGTERM'); - setTimeout(() => { - try { - if (shell.pid && !exited) { - process.kill(-shell.pid, 'SIGKILL'); - } - } catch (_e) { - console.error(`failed to kill shell process ${shell.pid}: ${_e}`); - } - }, 200); + await new Promise((resolve) => setTimeout(resolve, 200)); + if (shell.pid && !exited) { + process.kill(-shell.pid, 'SIGKILL'); + } } catch (_e) { // if group kill fails, fall back to killing just the main process try {