From 0795e55f0e7d2f2822bcd83eaf066eb99c67f858 Mon Sep 17 00:00:00 2001 From: Tolik Malibroda <54813212+tolikmalibroda@users.noreply.github.com> Date: Mon, 2 Jun 2025 22:05:45 +0200 Subject: [PATCH] feat: Add --yolo mode that automatically accepts all tools executions (#695) Co-authored-by: N. Taylor Mullen --- packages/cli/src/config/config.ts | 12 +- packages/cli/src/ui/App.test.tsx | 17 +- packages/cli/src/ui/App.tsx | 10 +- .../src/ui/components/AutoAcceptIndicator.tsx | 45 ++++- .../ui/hooks/useAutoAcceptIndicator.test.ts | 175 +++++++++++------- .../src/ui/hooks/useAutoAcceptIndicator.ts | 32 ++-- .../cli/src/ui/hooks/useReactToolScheduler.ts | 1 + .../cli/src/ui/hooks/useToolScheduler.test.ts | 107 ++++++++++- packages/core/src/config/config.ts | 21 ++- packages/core/src/core/coreToolScheduler.ts | 52 +++--- packages/core/src/tools/edit.test.ts | 20 +- packages/core/src/tools/edit.ts | 6 +- packages/core/src/tools/tool-registry.test.ts | 4 +- packages/core/src/tools/write-file.test.ts | 12 +- packages/core/src/tools/write-file.ts | 6 +- 15 files changed, 364 insertions(+), 156 deletions(-) diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 44057fad..ee1c9d36 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -15,6 +15,7 @@ import { ConfigParameters, setGeminiMdFilename as setServerGeminiMdFilename, getCurrentGeminiMdFilename, + ApprovalMode, } from '@gemini-code/core'; import { Settings } from './settings.js'; import { readPackageUp } from 'read-package-up'; @@ -38,6 +39,7 @@ interface CliArgs { prompt: string | undefined; all_files: boolean | undefined; show_memory_usage: boolean | undefined; + yolo: boolean | undefined; } async function parseArguments(): Promise { @@ -75,6 +77,13 @@ async function parseArguments(): Promise { description: 'Show memory usage in status bar', default: false, }) + .option('yolo', { + alias: 'y', + type: 'boolean', + description: + 'Automatically accept all actions (aka YOLO mode, see https://www.youtube.com/watch?v=xvFZjo5PgG0 for more details)?', + default: false, + }) .version() // This will enable the --version flag based on package.json .help() .alias('h', 'help') @@ -158,7 +167,7 @@ export async function loadCliConfig(settings: Settings): Promise { const configParams: ConfigParameters = { apiKey: apiKeyForServer, model: argv.model || DEFAULT_GEMINI_MODEL, - sandbox: argv.sandbox ?? settings.sandbox ?? false, + sandbox: argv.sandbox ?? settings.sandbox ?? argv.yolo ?? false, targetDir: process.cwd(), debugMode, question: argv.prompt || '', @@ -171,6 +180,7 @@ export async function loadCliConfig(settings: Settings): Promise { userAgent, userMemory: memoryContent, geminiMdFileCount: fileCount, + approvalMode: argv.yolo || false ? ApprovalMode.YOLO : ApprovalMode.DEFAULT, vertexai: useVertexAI, showMemoryUsage: argv.show_memory_usage || settings.showMemoryUsage || false, diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 17d9b459..82c28934 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -8,7 +8,7 @@ import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest'; import { render } from 'ink-testing-library'; import { App } from './App.js'; import { Config as ServerConfig, MCPServerConfig } from '@gemini-code/core'; -import type { ToolRegistry } from '@gemini-code/core'; +import { ApprovalMode, ToolRegistry } from '@gemini-code/core'; import { LoadedSettings, SettingsFile, Settings } from '../config/settings.js'; // Define a more complete mock server config based on actual Config @@ -28,7 +28,7 @@ interface MockServerConfig { userAgent: string; userMemory: string; geminiMdFileCount: number; - alwaysSkipModificationConfirmation: boolean; + approvalMode: ApprovalMode; vertexai?: boolean; showMemoryUsage?: boolean; @@ -50,8 +50,8 @@ interface MockServerConfig { setUserMemory: Mock<(newUserMemory: string) => void>; getGeminiMdFileCount: Mock<() => number>; setGeminiMdFileCount: Mock<(count: number) => void>; - getAlwaysSkipModificationConfirmation: Mock<() => boolean>; - setAlwaysSkipModificationConfirmation: Mock<(skip: boolean) => void>; + getApprovalMode: Mock<() => ApprovalMode>; + setApprovalMode: Mock<(skip: ApprovalMode) => void>; getVertexAI: Mock<() => boolean | undefined>; getShowMemoryUsage: Mock<() => boolean>; } @@ -80,8 +80,7 @@ vi.mock('@gemini-code/core', async (importOriginal) => { userAgent: opts.userAgent || 'test-agent', userMemory: opts.userMemory || '', geminiMdFileCount: opts.geminiMdFileCount || 0, - alwaysSkipModificationConfirmation: - opts.alwaysSkipModificationConfirmation ?? false, + approvalMode: opts.approvalMode ?? ApprovalMode.DEFAULT, vertexai: opts.vertexai, showMemoryUsage: opts.showMemoryUsage ?? false, @@ -105,10 +104,8 @@ vi.mock('@gemini-code/core', async (importOriginal) => { setUserMemory: vi.fn(), getGeminiMdFileCount: vi.fn(() => opts.geminiMdFileCount || 0), setGeminiMdFileCount: vi.fn(), - getAlwaysSkipModificationConfirmation: vi.fn( - () => opts.alwaysSkipModificationConfirmation ?? false, - ), - setAlwaysSkipModificationConfirmation: vi.fn(), + getApprovalMode: vi.fn(() => opts.approvalMode ?? ApprovalMode.DEFAULT), + setApprovalMode: vi.fn(), getVertexAI: vi.fn(() => opts.vertexai), getShowMemoryUsage: vi.fn(() => opts.showMemoryUsage ?? false), }; diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 73643bd5..7e0e19dd 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -44,6 +44,7 @@ import { getErrorMessage, type Config, getCurrentGeminiMdFilename, + ApprovalMode, } from '@gemini-code/core'; import { useLogger } from './hooks/useLogger.js'; import { StreamingContext } from './contexts/StreamingContext.js'; @@ -412,9 +413,12 @@ export const App = ({ )} - {showAutoAcceptIndicator && !shellModeActive && ( - - )} + {showAutoAcceptIndicator !== ApprovalMode.DEFAULT && + !shellModeActive && ( + + )} {shellModeActive && } diff --git a/packages/cli/src/ui/components/AutoAcceptIndicator.tsx b/packages/cli/src/ui/components/AutoAcceptIndicator.tsx index dc23d5e2..dbdb3e9a 100644 --- a/packages/cli/src/ui/components/AutoAcceptIndicator.tsx +++ b/packages/cli/src/ui/components/AutoAcceptIndicator.tsx @@ -7,12 +7,41 @@ import React from 'react'; import { Box, Text } from 'ink'; import { Colors } from '../colors.js'; +import { ApprovalMode } from '@gemini-code/core'; -export const AutoAcceptIndicator: React.FC = () => ( - - - accepting edits - (shift + tab to disable) - - -); +interface AutoAcceptIndicatorProps { + approvalMode: ApprovalMode; +} + +export const AutoAcceptIndicator: React.FC = ({ + approvalMode, +}) => { + let textColor = ''; + let textContent = ''; + let subText = ''; + + switch (approvalMode) { + case ApprovalMode.AUTO_EDIT: + textColor = Colors.AccentGreen; + textContent = 'accepting edits'; + subText = ' (shift + tab to toggle)'; + break; + case ApprovalMode.YOLO: + textColor = Colors.AccentRed; + textContent = 'YOLO mode'; + subText = ' (ctrl + y to toggle)'; + break; + case ApprovalMode.DEFAULT: + default: + break; + } + + return ( + + + {textContent} + {subText && {subText}} + + + ); +}; 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(false); + 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(ApprovalMode.DEFAULT); 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('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(false); - expect(result.current).toBe(false); + expect(mockConfigInstance.setApprovalMode).toHaveBeenCalledWith( + ApprovalMode.AUTO_EDIT, + ); + expect(result.current).toBe(ApprovalMode.AUTO_EDIT); + + act(() => { + capturedUseInputHandler('', { tab: true, shift: true } as InkKey); + }); + 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. diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index a6279e2e..71c4a7d2 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -22,6 +22,12 @@ import { ReadManyFilesTool } from '../tools/read-many-files.js'; import { MemoryTool, setGeminiMdFilename } from '../tools/memoryTool.js'; import { WebSearchTool } from '../tools/web-search.js'; +export enum ApprovalMode { + DEFAULT = 'default', + AUTO_EDIT = 'autoEdit', + YOLO = 'yolo', +} + export class MCPServerConfig { constructor( // For stdio transport @@ -53,7 +59,7 @@ export interface ConfigParameters { userAgent: string; userMemory?: string; geminiMdFileCount?: number; - alwaysSkipModificationConfirmation?: boolean; + approvalMode?: ApprovalMode; vertexai?: boolean; showMemoryUsage?: boolean; contextFileName?: string; @@ -76,7 +82,7 @@ export class Config { private readonly userAgent: string; private userMemory: string; private geminiMdFileCount: number; - private alwaysSkipModificationConfirmation: boolean; + private approvalMode: ApprovalMode; private readonly vertexai: boolean | undefined; private readonly showMemoryUsage: boolean; @@ -96,8 +102,7 @@ export class Config { this.userAgent = params.userAgent; this.userMemory = params.userMemory ?? ''; this.geminiMdFileCount = params.geminiMdFileCount ?? 0; - this.alwaysSkipModificationConfirmation = - params.alwaysSkipModificationConfirmation ?? false; + this.approvalMode = params.approvalMode ?? ApprovalMode.DEFAULT; this.vertexai = params.vertexai; this.showMemoryUsage = params.showMemoryUsage ?? false; @@ -179,12 +184,12 @@ export class Config { this.geminiMdFileCount = count; } - getAlwaysSkipModificationConfirmation(): boolean { - return this.alwaysSkipModificationConfirmation; + getApprovalMode(): ApprovalMode { + return this.approvalMode; } - setAlwaysSkipModificationConfirmation(skip: boolean): void { - this.alwaysSkipModificationConfirmation = skip; + setApprovalMode(mode: ApprovalMode): void { + this.approvalMode = mode; } getVertexAI(): boolean | undefined { diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 58f821c5..3fe6562e 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -12,6 +12,7 @@ import { ToolCallConfirmationDetails, ToolResult, ToolRegistry, + ApprovalMode, } from '../index.js'; import { Part, PartUnion, PartListUnion } from '@google/genai'; @@ -159,6 +160,7 @@ interface CoreToolSchedulerOptions { outputUpdateHandler?: OutputUpdateHandler; onAllToolCallsComplete?: AllToolCallsCompleteHandler; onToolCallsUpdate?: ToolCallsUpdateHandler; + approvalMode?: ApprovalMode; } export class CoreToolScheduler { @@ -168,12 +170,14 @@ export class CoreToolScheduler { private outputUpdateHandler?: OutputUpdateHandler; private onAllToolCallsComplete?: AllToolCallsCompleteHandler; private onToolCallsUpdate?: ToolCallsUpdateHandler; + private approvalMode: ApprovalMode; constructor(options: CoreToolSchedulerOptions) { this.toolRegistry = options.toolRegistry; this.outputUpdateHandler = options.outputUpdateHandler; this.onAllToolCallsComplete = options.onAllToolCallsComplete; this.onToolCallsUpdate = options.onToolCallsUpdate; + this.approvalMode = options.approvalMode ?? ApprovalMode.DEFAULT; this.abortController = new AbortController(); } @@ -324,29 +328,33 @@ export class CoreToolScheduler { const { request: reqInfo, tool: toolInstance } = toolCall; try { - const confirmationDetails = await toolInstance.shouldConfirmExecute( - reqInfo.args, - this.abortController.signal, - ); - - if (confirmationDetails) { - const originalOnConfirm = confirmationDetails.onConfirm; - const wrappedConfirmationDetails: ToolCallConfirmationDetails = { - ...confirmationDetails, - onConfirm: (outcome: ToolConfirmationOutcome) => - this.handleConfirmationResponse( - reqInfo.callId, - originalOnConfirm, - outcome, - ), - }; - this.setStatusInternal( - reqInfo.callId, - 'awaiting_approval', - wrappedConfirmationDetails, - ); - } else { + if (this.approvalMode === ApprovalMode.YOLO) { this.setStatusInternal(reqInfo.callId, 'scheduled'); + } else { + const confirmationDetails = await toolInstance.shouldConfirmExecute( + reqInfo.args, + this.abortController.signal, + ); + + if (confirmationDetails) { + const originalOnConfirm = confirmationDetails.onConfirm; + const wrappedConfirmationDetails: ToolCallConfirmationDetails = { + ...confirmationDetails, + onConfirm: (outcome: ToolConfirmationOutcome) => + this.handleConfirmationResponse( + reqInfo.callId, + originalOnConfirm, + outcome, + ), + }; + this.setStatusInternal( + reqInfo.callId, + 'awaiting_approval', + wrappedConfirmationDetails, + ); + } else { + this.setStatusInternal(reqInfo.callId, 'scheduled'); + } } } catch (error) { this.setStatusInternal( diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index c6c2ba63..3b93708a 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -25,7 +25,7 @@ import { FileDiff } from './tools.js'; import path from 'path'; import fs from 'fs'; import os from 'os'; -import { Config } from '../config/config.js'; +import { ApprovalMode, Config } from '../config/config.js'; import { Content, Part, SchemaUnion } from '@google/genai'; describe('EditTool', () => { @@ -41,8 +41,8 @@ describe('EditTool', () => { mockConfig = { getTargetDir: () => rootDir, - getAlwaysSkipModificationConfirmation: vi.fn(() => false), - setAlwaysSkipModificationConfirmation: vi.fn(), + getApprovalMode: vi.fn(() => false), + setApprovalMode: vi.fn(), // getGeminiConfig: () => ({ apiKey: 'test-api-key' }), // This was not a real Config method // Add other properties/methods of Config if EditTool uses them // Minimal other methods to satisfy Config type if needed by EditTool constructor or other direct uses: @@ -65,12 +65,10 @@ describe('EditTool', () => { } as unknown as Config; // Reset mocks before each test - (mockConfig.getAlwaysSkipModificationConfirmation as Mock).mockClear(); - (mockConfig.setAlwaysSkipModificationConfirmation as Mock).mockClear(); + (mockConfig.getApprovalMode as Mock).mockClear(); + (mockConfig.getApprovalMode as Mock).mockClear(); // Default to not skipping confirmation - (mockConfig.getAlwaysSkipModificationConfirmation as Mock).mockReturnValue( - false, - ); + (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT); // Reset mocks and set default implementation for ensureCorrectEdit mockEnsureCorrectEdit.mockReset(); @@ -439,9 +437,9 @@ describe('EditTool', () => { new_string: fileContent, }; - ( - mockConfig.getAlwaysSkipModificationConfirmation as Mock - ).mockReturnValueOnce(true); + (mockConfig.getApprovalMode as Mock).mockReturnValueOnce( + ApprovalMode.AUTO_EDIT, + ); const result = await tool.execute(params, new AbortController().signal); expect(result.llmContent).toMatch(/Created new file/); diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 53b12480..b2f648f8 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -20,7 +20,7 @@ import { makeRelative, shortenPath } from '../utils/paths.js'; import { isNodeError } from '../utils/errors.js'; import { ReadFileTool } from './read-file.js'; import { GeminiClient } from '../core/client.js'; -import { Config } from '../config/config.js'; +import { Config, ApprovalMode } from '../config/config.js'; import { ensureCorrectEdit } from '../utils/editCorrector.js'; import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js'; @@ -281,7 +281,7 @@ Expectation for required parameters: params: EditToolParams, abortSignal: AbortSignal, ): Promise { - if (this.config.getAlwaysSkipModificationConfirmation()) { + if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { return false; } const validationError = this.validateToolParams(params); @@ -356,7 +356,7 @@ Expectation for required parameters: fileDiff, onConfirm: async (outcome: ToolConfirmationOutcome) => { if (outcome === ToolConfirmationOutcome.ProceedAlways) { - this.config.setAlwaysSkipModificationConfirmation(true); + this.config.setApprovalMode(ApprovalMode.AUTO_EDIT); } }, }; diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index 121e91c8..9aaa7e5a 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -16,7 +16,7 @@ import { } from 'vitest'; import { ToolRegistry, DiscoveredTool } from './tool-registry.js'; import { DiscoveredMCPTool } from './mcp-tool.js'; -import { Config, ConfigParameters } from '../config/config.js'; +import { ApprovalMode, Config, ConfigParameters } from '../config/config.js'; import { BaseTool, ToolResult } from './tools.js'; import { FunctionDeclaration } from '@google/genai'; import { execSync, spawn } from 'node:child_process'; // Import spawn here @@ -85,7 +85,7 @@ const baseConfigParams: ConfigParameters = { userAgent: 'TestAgent/1.0', userMemory: '', geminiMdFileCount: 0, - alwaysSkipModificationConfirmation: false, + approvalMode: ApprovalMode.DEFAULT, vertexai: false, }; diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index 3fd97c9e..c94edfd1 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -20,7 +20,7 @@ import { ToolEditConfirmationDetails, } from './tools.js'; import { type EditToolParams } from './edit.js'; -import { Config } from '../config/config.js'; +import { ApprovalMode, Config } from '../config/config.js'; import { ToolRegistry } from './tool-registry.js'; import path from 'path'; import fs from 'fs'; @@ -51,8 +51,8 @@ vi.mocked(ensureCorrectFileContent).mockImplementation( // Mock Config const mockConfigInternal = { getTargetDir: () => rootDir, - getAlwaysSkipModificationConfirmation: vi.fn(() => false), - setAlwaysSkipModificationConfirmation: vi.fn(), + getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT), + setApprovalMode: vi.fn(), getApiKey: () => 'test-key', getModel: () => 'test-model', getSandbox: () => false, @@ -100,10 +100,8 @@ describe('WriteFileTool', () => { tool = new WriteFileTool(mockConfig); // Reset mocks before each test - mockConfigInternal.getAlwaysSkipModificationConfirmation.mockReturnValue( - false, - ); - mockConfigInternal.setAlwaysSkipModificationConfirmation.mockClear(); + mockConfigInternal.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT); + mockConfigInternal.setApprovalMode.mockClear(); mockEnsureCorrectEdit.mockReset(); mockEnsureCorrectFileContent.mockReset(); diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index 2285c819..2e04a10a 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -7,7 +7,7 @@ import fs from 'fs'; import path from 'path'; import * as Diff from 'diff'; -import { Config } from '../config/config.js'; +import { Config, ApprovalMode } from '../config/config.js'; import { BaseTool, ToolResult, @@ -143,7 +143,7 @@ export class WriteFileTool extends BaseTool { params: WriteFileToolParams, abortSignal: AbortSignal, ): Promise { - if (this.config.getAlwaysSkipModificationConfirmation()) { + if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { return false; } @@ -186,7 +186,7 @@ export class WriteFileTool extends BaseTool { fileDiff, onConfirm: async (outcome: ToolConfirmationOutcome) => { if (outcome === ToolConfirmationOutcome.ProceedAlways) { - this.config.setAlwaysSkipModificationConfirmation(true); + this.config.setApprovalMode(ApprovalMode.AUTO_EDIT); } }, };