feat(core): improve shell execution service reliability

packages/cli/src/ui/hooks/shellCommandProcessor.test.ts

@@ -219,7 +219,7 @@ describe('useShellCommandProcessor', () => {
       vi.useRealTimers();
     });
 
-    it('should throttle pending UI updates for text streams (non-interactive)', async () => {
+    it('should update UI for text streams (non-interactive)', async () => {
       const { result } = renderProcessorHook();
       act(() => {
         result.current.handleShellCommand(
@@ -243,61 +243,43 @@ describe('useShellCommandProcessor', () => {
       );
 
       // Wait for the async PID update to happen.
+      // Call 1: Initial, Call 2: PID update
       await vi.waitFor(() => {
-        // It's called once for initial, and once for the PID update.
         expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(2);
       });
 
-      // Simulate rapid output
+      // Get the state after the PID update to feed into the stream updaters
+      const pidUpdateFn = setPendingHistoryItemMock.mock.calls[1][0];
+      const initialState = setPendingHistoryItemMock.mock.calls[0][0];
+      const stateAfterPid = pidUpdateFn(initialState);
+
+      // Simulate first output chunk
       act(() => {
         mockShellOutputCallback({
           type: 'data',
           chunk: 'hello',
         });
       });
-      // The count should still be 2, as throttling is in effect.
-      expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(2);
+      // A UI update should have occurred.
+      expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(3);
 
-      // Simulate more rapid output
-      act(() => {
-        mockShellOutputCallback({
-          type: 'data',
-          chunk: ' world',
-        });
-      });
-      expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(2);
+      const streamUpdateFn1 = setPendingHistoryItemMock.mock.calls[2][0];
+      const stateAfterStream1 = streamUpdateFn1(stateAfterPid);
+      expect(stateAfterStream1.tools[0].resultDisplay).toBe('hello');
 
-      // Advance time, but the update won't happen until the next event
-      await act(async () => {
-        await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1);
-      });
-
-      // Trigger one more event to cause the throttled update to fire.
+      // Simulate second output chunk
       act(() => {
         mockShellOutputCallback({
           type: 'data',
-          chunk: '',
+          chunk: ' world',
         });
       });
+      // Another UI update should have occurred.
+      expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(4);
 
-      // Now the cumulative update should have occurred.
-      // Call 1: Initial, Call 2: PID update, Call 3: Throttled stream update
-      expect(setPendingHistoryItemMock).toHaveBeenCalledTimes(3);
-
-      const streamUpdateFn = setPendingHistoryItemMock.mock.calls[2][0];
-      if (!streamUpdateFn || typeof streamUpdateFn !== 'function') {
-        throw new Error(
-          'setPendingHistoryItem was not called with a stream updater function',
-        );
-      }
-
-      // Get the state after the PID update to feed into the stream updater
-      const pidUpdateFn = setPendingHistoryItemMock.mock.calls[1][0];
-      const initialState = setPendingHistoryItemMock.mock.calls[0][0];
-      const stateAfterPid = pidUpdateFn(initialState);
-
-      const stateAfterStream = streamUpdateFn(stateAfterPid);
-      expect(stateAfterStream.tools[0].resultDisplay).toBe('hello world');
+      const streamUpdateFn2 = setPendingHistoryItemMock.mock.calls[3][0];
+      const stateAfterStream2 = streamUpdateFn2(stateAfterStream1);
+      expect(stateAfterStream2.tools[0].resultDisplay).toBe('hello world');
     });
 
     it('should show binary progress messages correctly', async () => {

packages/cli/src/ui/hooks/shellCommandProcessor.ts

@@ -109,7 +109,6 @@ export const useShellCommandProcessor = (
       const executeCommand = async (
         resolve: (value: void | PromiseLike<void>) => void,
       ) => {
-        let lastUpdateTime = Date.now();
         let cumulativeStdout: string | AnsiOutput = '';
         let isBinaryStream = false;
         let binaryBytesReceived = 0;
@@ -168,6 +167,7 @@ export const useShellCommandProcessor = (
                     typeof cumulativeStdout === 'string'
                   ) {
                     cumulativeStdout += event.chunk;
+                    shouldUpdate = true;
                   }
                   break;
                 case 'binary_detected':
@@ -178,6 +178,7 @@ export const useShellCommandProcessor = (
                 case 'binary_progress':
                   isBinaryStream = true;
                   binaryBytesReceived = event.bytesReceived;
+                  shouldUpdate = true;
                   break;
                 default: {
                   throw new Error('An unhandled ShellOutputEvent was found.');
@@ -200,10 +201,7 @@ export const useShellCommandProcessor = (
               }
 
               // Throttle pending UI updates, but allow forced updates.
-              if (
-                shouldUpdate ||
-                Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS
-              ) {
+              if (shouldUpdate) {
                 setPendingHistoryItem((prevItem) => {
                   if (prevItem?.type === 'tool_group') {
                     return {
@@ -217,7 +215,6 @@ export const useShellCommandProcessor = (
                   }
                   return prevItem;
                 });
-                lastUpdateTime = Date.now();
               }
             },
             abortSignal,

packages/core/src/services/shellExecutionService.test.ts

@@ -153,7 +153,7 @@ describe('ShellExecutionService', () => {
     simulation: (
       ptyProcess: typeof mockPtyProcess,
       ac: AbortController,
-    ) => void,
+    ) => void | Promise<void>,
     config = shellExecutionConfig,
   ) => {
     const abortController = new AbortController();
@@ -167,7 +167,7 @@ describe('ShellExecutionService', () => {
     );
 
     await new Promise((resolve) => process.nextTick(resolve));
-    simulation(mockPtyProcess, abortController);
+    await simulation(mockPtyProcess, abortController);
     const result = await handle.result;
     return { result, handle, abortController };
   };
@@ -356,6 +356,63 @@ describe('ShellExecutionService', () => {
       expect(result.aborted).toBe(true);
       // The process kill is mocked, so we just check that the flag is set.
     });
+
+    it('should send SIGTERM and then SIGKILL on abort', async () => {
+      const sigkillPromise = new Promise<void>((resolve) => {
+        mockProcessKill.mockImplementation((pid, signal) => {
+          if (signal === 'SIGKILL' && pid === -mockPtyProcess.pid) {
+            resolve();
+          }
+          return true;
+        });
+      });
+
+      const { result } = await simulateExecution(
+        'long-running-process',
+        async (pty, abortController) => {
+          abortController.abort();
+          await sigkillPromise; // Wait for SIGKILL to be sent before exiting.
+          pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: 9 });
+        },
+      );
+
+      expect(result.aborted).toBe(true);
+
+      // Verify the calls were made in the correct order.
+      const killCalls = mockProcessKill.mock.calls;
+      const sigtermCallIndex = killCalls.findIndex(
+        (call) => call[0] === -mockPtyProcess.pid && call[1] === 'SIGTERM',
+      );
+      const sigkillCallIndex = killCalls.findIndex(
+        (call) => call[0] === -mockPtyProcess.pid && call[1] === 'SIGKILL',
+      );
+
+      expect(sigtermCallIndex).toBe(0);
+      expect(sigkillCallIndex).toBe(1);
+      expect(sigtermCallIndex).toBeLessThan(sigkillCallIndex);
+
+      expect(result.signal).toBe(9);
+    });
+
+    it('should resolve without waiting for the processing chain on abort', async () => {
+      const { result } = await simulateExecution(
+        'long-output',
+        (pty, abortController) => {
+          // Simulate a lot of data being in the queue to be processed
+          for (let i = 0; i < 1000; i++) {
+            pty.onData.mock.calls[0][0]('some data');
+          }
+          abortController.abort();
+          pty.onExit.mock.calls[0][0]({ exitCode: 1, signal: null });
+        },
+      );
+
+      // The main assertion here is implicit: the `await` for the result above
+      // should complete without timing out. This proves that the resolution
+      // was not blocked by the long chain of data processing promises,
+      // which is the desired behavior on abort.
+      expect(result.aborted).toBe(true);
+    });
   });
 
   describe('Binary Output', () => {
@@ -633,6 +690,36 @@ describe('ShellExecutionService child_process fallback', () => {
       expect(result.output.trim()).toBe('');
       expect(onOutputEventMock).not.toHaveBeenCalled();
     });
+
+    it('should truncate stdout using a sliding window and show a warning', async () => {
+      const MAX_SIZE = 16 * 1024 * 1024;
+      const chunk1 = 'a'.repeat(MAX_SIZE / 2 - 5);
+      const chunk2 = 'b'.repeat(MAX_SIZE / 2 - 5);
+      const chunk3 = 'c'.repeat(20);
+
+      const { result } = await simulateExecution('large-output', (cp) => {
+        cp.stdout?.emit('data', Buffer.from(chunk1));
+        cp.stdout?.emit('data', Buffer.from(chunk2));
+        cp.stdout?.emit('data', Buffer.from(chunk3));
+        cp.emit('exit', 0, null);
+      });
+
+      const truncationMessage =
+        '[GEMINI_CLI_WARNING: Output truncated. The buffer is limited to 16MB.]';
+      expect(result.output).toContain(truncationMessage);
+
+      const outputWithoutMessage = result.output
+        .substring(0, result.output.indexOf(truncationMessage))
+        .trimEnd();
+
+      expect(outputWithoutMessage.length).toBe(MAX_SIZE);
+
+      const expectedStart = (chunk1 + chunk2 + chunk3).slice(-MAX_SIZE);
+      expect(
+        outputWithoutMessage.startsWith(expectedStart.substring(0, 10)),
+      ).toBe(true);
+      expect(outputWithoutMessage.endsWith('c'.repeat(20))).toBe(true);
+    }, 20000);
   });
 
   describe('Failed Execution', () => {