feat: Add --yolo mode that automatically accepts all tools executions (#695)

Co-authored-by: N. Taylor Mullen <[email protected]>

4 files changed, 236 insertions, 77 deletions

diff --git a/packages/cli/src/ui/hooks/useAutoAcceptIndicator.test.ts b/packages/cli/src/ui/hooks/useAutoAcceptIndicator.test.ts
index 8c611ccc..520262f5 100644
--- a/packages/cli/src/ui/hooks/useAutoAcceptIndicator.test.ts
+++ b/packages/cli/src/ui/hooks/useAutoAcceptIndicator.test.ts
@@ -16,7 +16,11 @@ import {
 import { renderHook, act } from '@testing-library/react';
 import { useAutoAcceptIndicator } from './useAutoAcceptIndicator.js';
 
-import type { Config as ActualConfigType } from '@gemini-code/core';
+import {
+  Config,
+  Config as ActualConfigType,
+  ApprovalMode,
+} from '@gemini-code/core';
 import { useInput, type Key as InkKey } from 'ink';
 
 vi.mock('ink');
@@ -31,11 +35,9 @@ vi.mock('@gemini-code/core', async () => {
   };
 });
 
-import { Config } from '@gemini-code/core';
-
 interface MockConfigInstanceShape {
-  getAlwaysSkipModificationConfirmation: Mock<() => boolean>;
-  setAlwaysSkipModificationConfirmation: Mock<(value: boolean) => void>;
+  getApprovalMode: Mock<() => ApprovalMode>;
+  setApprovalMode: Mock<(value: ApprovalMode) => void>;
   getCoreTools: Mock<() => string[]>;
   getToolDiscoveryCommand: Mock<() => string | undefined>;
   getTargetDir: Mock<() => string>;
@@ -65,14 +67,16 @@ describe('useAutoAcceptIndicator', () => {
     (
       Config as unknown as MockedFunction<() => MockConfigInstanceShape>
     ).mockImplementation(() => {
-      const instanceGetAlwaysSkipMock = vi.fn();
-      const instanceSetAlwaysSkipMock = vi.fn();
+      const instanceGetApprovalModeMock = vi.fn();
+      const instanceSetApprovalModeMock = vi.fn();
 
       const instance: MockConfigInstanceShape = {
-        getAlwaysSkipModificationConfirmation:
-          instanceGetAlwaysSkipMock as Mock<() => boolean>,
-        setAlwaysSkipModificationConfirmation:
-          instanceSetAlwaysSkipMock as Mock<(value: boolean) => void>,
+        getApprovalMode: instanceGetApprovalModeMock as Mock<
+          () => ApprovalMode
+        >,
+        setApprovalMode: instanceSetApprovalModeMock as Mock<
+          (value: ApprovalMode) => void
+        >,
         getCoreTools: vi.fn().mockReturnValue([]) as Mock<() => string[]>,
         getToolDiscoveryCommand: vi.fn().mockReturnValue(undefined) as Mock<
           () => string | undefined
@@ -101,8 +105,8 @@ describe('useAutoAcceptIndicator', () => {
           () => { discoverTools: Mock<() => void> }
         >,
       };
-      instanceSetAlwaysSkipMock.mockImplementation((value: boolean) => {
-        instanceGetAlwaysSkipMock.mockReturnValue(value);
+      instanceSetApprovalModeMock.mockImplementation((value: ApprovalMode) => {
+        instanceGetApprovalModeMock.mockReturnValue(value);
       });
       return instance;
     });
@@ -116,68 +120,99 @@ describe('useAutoAcceptIndicator', () => {
     mockConfigInstance = new (Config as any)() as MockConfigInstanceShape;
   });
 
-  it('should initialize with true if config.getAlwaysSkipModificationConfirmation returns true', () => {
-    mockConfigInstance.getAlwaysSkipModificationConfirmation.mockReturnValue(
-      true,
-    );
+  it('should initialize with ApprovalMode.AUTO_EDIT if config.getApprovalMode returns ApprovalMode.AUTO_EDIT', () => {
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.AUTO_EDIT);
     const { result } = renderHook(() =>
       useAutoAcceptIndicator({
         config: mockConfigInstance as unknown as ActualConfigType,
       }),
     );
-    expect(result.current).toBe(true);
-    expect(
-      mockConfigInstance.getAlwaysSkipModificationConfirmation,
-    ).toHaveBeenCalledTimes(1);
+    expect(result.current).toBe(ApprovalMode.AUTO_EDIT);
+    expect(mockConfigInstance.getApprovalMode).toHaveBeenCalledTimes(1);
   });
 
-  it('should initialize with false if config.getAlwaysSkipModificationConfirmation returns false', () => {
-    mockConfigInstance.getAlwaysSkipModificationConfirmation.mockReturnValue(
-      false,
-    );
+  it('should initialize with ApprovalMode.DEFAULT if config.getApprovalMode returns ApprovalMode.DEFAULT', () => {
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT);
     const { result } = renderHook(() =>
       useAutoAcceptIndicator({
         config: mockConfigInstance as unknown as ActualConfigType,
       }),
     );
-    expect(result.current).toBe(false);
-    expect(
-      mockConfigInstance.getAlwaysSkipModificationConfirmation,
-    ).toHaveBeenCalledTimes(1);
+    expect(result.current).toBe(ApprovalMode.DEFAULT);
+    expect(mockConfigInstance.getApprovalMode).toHaveBeenCalledTimes(1);
   });
 
-  it('should toggle the indicator and update config when Shift+Tab is pressed', () => {
-    mockConfigInstance.getAlwaysSkipModificationConfirmation.mockReturnValue(
-      false,
+  it('should initialize with ApprovalMode.YOLO if config.getApprovalMode returns ApprovalMode.YOLO', () => {
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.YOLO);
+    const { result } = renderHook(() =>
+      useAutoAcceptIndicator({
+        config: mockConfigInstance as unknown as ActualConfigType,
+      }),
     );
+    expect(result.current).toBe(ApprovalMode.YOLO);
+    expect(mockConfigInstance.getApprovalMode).toHaveBeenCalledTimes(1);
+  });
+
+  it('should toggle the indicator and update config when Shift+Tab or Ctrl+Y is pressed', () => {
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT);
     const { result } = renderHook(() =>
       useAutoAcceptIndicator({
         config: mockConfigInstance as unknown as ActualConfigType,
       }),
     );
-    expect(result.current).toBe(false);
+    expect(result.current).toBe(ApprovalMode.DEFAULT);
+
+    act(() => {
+      capturedUseInputHandler('', { tab: true, shift: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith(
+      ApprovalMode.AUTO_EDIT,
+    );
+    expect(result.current).toBe(ApprovalMode.AUTO_EDIT);
+
+    act(() => {
+      capturedUseInputHandler('y', { ctrl: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith(
+      ApprovalMode.YOLO,
+    );
+    expect(result.current).toBe(ApprovalMode.YOLO);
+
+    act(() => {
+      capturedUseInputHandler('y', { ctrl: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith(
+      ApprovalMode.DEFAULT,
+    );
+    expect(result.current).toBe(ApprovalMode.DEFAULT);
+
+    act(() => {
+      capturedUseInputHandler('y', { ctrl: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith(
+      ApprovalMode.YOLO,
+    );
+    expect(result.current).toBe(ApprovalMode.YOLO);
 
     act(() => {
       capturedUseInputHandler('', { tab: true, shift: true } as InkKey);
     });
-    expect(
-      mockConfigInstance.setAlwaysSkipModificationConfirmation,
-    ).toHaveBeenCalledWith(true);
-    expect(result.current).toBe(true);
+    expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith(
+      ApprovalMode.AUTO_EDIT,
+    );
+    expect(result.current).toBe(ApprovalMode.AUTO_EDIT);
 
     act(() => {
       capturedUseInputHandler('', { tab: true, shift: true } as InkKey);
     });
-    expect(
-      mockConfigInstance.setAlwaysSkipModificationConfirmation,
-    ).toHaveBeenCalledWith(false);
-    expect(result.current).toBe(false);
+    expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith(
+      ApprovalMode.DEFAULT,
+    );
+    expect(result.current).toBe(ApprovalMode.DEFAULT);
   });
 
-  it('should not toggle if only Tab, only Shift, or other keys are pressed', () => {
-    mockConfigInstance.getAlwaysSkipModificationConfirmation.mockReturnValue(
-      false,
-    );
+  it('should not toggle if only one key or other keys combinations are pressed', () => {
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT);
     renderHook(() =>
       useAutoAcceptIndicator({
         config: mockConfigInstance as unknown as ActualConfigType,
@@ -187,29 +222,41 @@ describe('useAutoAcceptIndicator', () => {
     act(() => {
       capturedUseInputHandler('', { tab: true, shift: false } as InkKey);
     });
-    expect(
-      mockConfigInstance.setAlwaysSkipModificationConfirmation,
-    ).not.toHaveBeenCalled();
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
 
     act(() => {
       capturedUseInputHandler('', { tab: false, shift: true } as InkKey);
     });
-    expect(
-      mockConfigInstance.setAlwaysSkipModificationConfirmation,
-    ).not.toHaveBeenCalled();
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
 
     act(() => {
       capturedUseInputHandler('a', { tab: false, shift: false } as InkKey);
     });
-    expect(
-      mockConfigInstance.setAlwaysSkipModificationConfirmation,
-    ).not.toHaveBeenCalled();
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
+
+    act(() => {
+      capturedUseInputHandler('y', { tab: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
+
+    act(() => {
+      capturedUseInputHandler('a', { ctrl: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
+
+    act(() => {
+      capturedUseInputHandler('y', { shift: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
+
+    act(() => {
+      capturedUseInputHandler('a', { ctrl: true, shift: true } as InkKey);
+    });
+    expect(mockConfigInstance.setApprovalMode).not.toHaveBeenCalled();
   });
 
   it('should update indicator when config value changes externally (useEffect dependency)', () => {
-    mockConfigInstance.getAlwaysSkipModificationConfirmation.mockReturnValue(
-      false,
-    );
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT);
     const { result, rerender } = renderHook(
       (props: { config: ActualConfigType }) => useAutoAcceptIndicator(props),
       {
@@ -218,16 +265,12 @@ describe('useAutoAcceptIndicator', () => {
         },
       },
     );
-    expect(result.current).toBe(false);
+    expect(result.current).toBe(ApprovalMode.DEFAULT);
 
-    mockConfigInstance.getAlwaysSkipModificationConfirmation.mockReturnValue(
-      true,
-    );
+    mockConfigInstance.getApprovalMode.mockReturnValue(ApprovalMode.AUTO_EDIT);
 
     rerender({ config: mockConfigInstance as unknown as ActualConfigType });
-    expect(result.current).toBe(true);
-    expect(
-      mockConfigInstance.getAlwaysSkipModificationConfirmation,
-    ).toHaveBeenCalledTimes(3);
+    expect(result.current).toBe(ApprovalMode.AUTO_EDIT);
+    expect(mockConfigInstance.getApprovalMode).toHaveBeenCalledTimes(3);
   });
 });
diff --git a/packages/cli/src/ui/hooks/useAutoAcceptIndicator.ts b/packages/cli/src/ui/hooks/useAutoAcceptIndicator.ts
index 5af1783b..aaa1dc68 100644
--- a/packages/cli/src/ui/hooks/useAutoAcceptIndicator.ts
+++ b/packages/cli/src/ui/hooks/useAutoAcceptIndicator.ts
@@ -6,7 +6,7 @@
 
 import { useState, useEffect } from 'react';
 import { useInput } from 'ink';
-import type { Config } from '@gemini-code/core';
+import { ApprovalMode, type Config } from '@gemini-code/core';
 
 export interface UseAutoAcceptIndicatorArgs {
   config: Config;
@@ -14,8 +14,8 @@ export interface UseAutoAcceptIndicatorArgs {
 
 export function useAutoAcceptIndicator({
   config,
-}: UseAutoAcceptIndicatorArgs): boolean {
-  const currentConfigValue = config.getAlwaysSkipModificationConfirmation();
+}: UseAutoAcceptIndicatorArgs): ApprovalMode {
+  const currentConfigValue = config.getApprovalMode();
   const [showAutoAcceptIndicator, setShowAutoAcceptIndicator] =
     useState(currentConfigValue);
 
@@ -23,15 +23,25 @@ export function useAutoAcceptIndicator({
     setShowAutoAcceptIndicator(currentConfigValue);
   }, [currentConfigValue]);
 
-  useInput((_input, key) => {
-    if (key.tab && key.shift) {
-      const alwaysAcceptModificationConfirmations =
-        !config.getAlwaysSkipModificationConfirmation();
-      config.setAlwaysSkipModificationConfirmation(
-        alwaysAcceptModificationConfirmations,
-      );
+  useInput((input, key) => {
+    let nextApprovalMode: ApprovalMode | undefined;
+
+    if (key.ctrl && input === 'y') {
+      nextApprovalMode =
+        config.getApprovalMode() === ApprovalMode.YOLO
+          ? ApprovalMode.DEFAULT
+          : ApprovalMode.YOLO;
+    } else if (key.tab && key.shift) {
+      nextApprovalMode =
+        config.getApprovalMode() === ApprovalMode.AUTO_EDIT
+          ? ApprovalMode.DEFAULT
+          : ApprovalMode.AUTO_EDIT;
+    }
+
+    if (nextApprovalMode) {
+      config.setApprovalMode(nextApprovalMode);
       // Update local state immediately for responsiveness
-      setShowAutoAcceptIndicator(alwaysAcceptModificationConfirmations);
+      setShowAutoAcceptIndicator(nextApprovalMode);
     }
   });
 
diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
index 12333d92..e681e972 100644
--- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts
+++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
@@ -134,6 +134,7 @@ export function useReactToolScheduler(
       outputUpdateHandler,
       onAllToolCallsComplete: allToolCallsCompleteHandler,
       onToolCallsUpdate: toolCallsUpdateHandler,
+      approvalMode: config.getApprovalMode(),
     });
   }, [config, onComplete, setPendingHistoryItem]);
 
diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts
index 92bff2bc..30880ba6 100644
--- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts
+++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts
@@ -28,7 +28,8 @@ import {
   ToolCallResponseInfo,
   formatLlmContentForFunctionResponse, // Import from core
   ToolCall, // Import from core
-  Status as ToolCallStatusType, // Import from core
+  Status as ToolCallStatusType,
+  ApprovalMode, // Import from core
 } from '@gemini-code/core';
 import {
   HistoryItemWithoutId,
@@ -52,6 +53,7 @@ const mockToolRegistry = {
 
 const mockConfig = {
   getToolRegistry: vi.fn(() => mockToolRegistry as unknown as ToolRegistry),
+  getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT),
 };
 
 const mockTool: Tool = {
@@ -205,6 +207,109 @@ describe('formatLlmContentForFunctionResponse', () => {
   });
 });
 
+describe('useReactToolScheduler in YOLO Mode', () => {
+  let onComplete: Mock;
+  let setPendingHistoryItem: Mock;
+
+  beforeEach(() => {
+    onComplete = vi.fn();
+    setPendingHistoryItem = vi.fn();
+    mockToolRegistry.getTool.mockClear();
+    (mockToolRequiresConfirmation.execute as Mock).mockClear();
+    (mockToolRequiresConfirmation.shouldConfirmExecute as Mock).mockClear();
+
+    // IMPORTANT: Enable YOLO mode for this test suite
+    (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO);
+
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.clearAllTimers();
+    vi.useRealTimers();
+    // IMPORTANT: Disable YOLO mode after this test suite
+    (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT);
+  });
+
+  const renderSchedulerInYoloMode = () =>
+    renderHook(() =>
+      useReactToolScheduler(
+        onComplete,
+        mockConfig as unknown as Config,
+        setPendingHistoryItem,
+      ),
+    );
+
+  it('should skip confirmation and execute tool directly when yoloMode is true', async () => {
+    mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation);
+    const expectedOutput = 'YOLO Confirmed output';
+    (mockToolRequiresConfirmation.execute as Mock).mockResolvedValue({
+      llmContent: expectedOutput,
+      returnDisplay: 'YOLO Formatted tool output',
+    } as ToolResult);
+
+    const { result } = renderSchedulerInYoloMode();
+    const schedule = result.current[1];
+    const request: ToolCallRequestInfo = {
+      callId: 'yoloCall',
+      name: 'mockToolRequiresConfirmation',
+      args: { data: 'any data' },
+    };
+
+    act(() => {
+      schedule(request);
+    });
+
+    await act(async () => {
+      await vi.runAllTimersAsync(); // Process validation
+    });
+    await act(async () => {
+      await vi.runAllTimersAsync(); // Process scheduling
+    });
+    await act(async () => {
+      await vi.runAllTimersAsync(); // Process execution
+    });
+
+    // Check that shouldConfirmExecute was NOT called
+    expect(
+      mockToolRequiresConfirmation.shouldConfirmExecute,
+    ).not.toHaveBeenCalled();
+
+    // Check that execute WAS called
+    expect(mockToolRequiresConfirmation.execute).toHaveBeenCalledWith(
+      request.args,
+      expect.any(AbortSignal),
+      undefined,
+    );
+
+    // Check that onComplete was called with success
+    expect(onComplete).toHaveBeenCalledWith([
+      expect.objectContaining({
+        status: 'success',
+        request,
+        response: expect.objectContaining({
+          resultDisplay: 'YOLO Formatted tool output',
+          responseParts: expect.arrayContaining([
+            expect.objectContaining({
+              functionResponse: expect.objectContaining({
+                response: { output: expectedOutput },
+              }),
+            }),
+          ]),
+        }),
+      }),
+    ]);
+
+    // Ensure no confirmation UI was triggered (setPendingHistoryItem should not have been called with confirmation details)
+    const setPendingHistoryItemCalls = setPendingHistoryItem.mock.calls;
+    const confirmationCall = setPendingHistoryItemCalls.find((call) => {
+      const item = typeof call[0] === 'function' ? call[0]({}) : call[0];
+      return item?.tools?.[0]?.confirmationDetails;
+    });
+    expect(confirmationCall).toBeUndefined();
+  });
+});
+
 describe('useReactToolScheduler', () => {
   // TODO(ntaylormullen): The following tests are skipped due to difficulties in
   // reliably testing the asynchronous state updates and interactions with timers.