/** * @license * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ import { describe, it, expect, vi, beforeEach, afterEach, type Mock, type Mocked, } from 'vitest'; import { AcpSessionManager } from './acpSessionManager.js'; import type * as acp from '@agentclientprotocol/sdk'; import { AuthType, type Config, GEMINI_MODEL_ALIAS_AUTO, type MessageBus, type Storage, } from '@google/gemini-cli-core'; import type { LoadedSettings } from '../config/settings.js'; import { loadCliConfig, type CliArgs } from '../config/config.js'; import { loadSettings } from '../config/settings.js'; vi.mock('../config/config.js', () => ({ loadCliConfig: vi.fn(), })); vi.mock('../config/settings.js', async (importOriginal) => { const actual = await importOriginal(); return { ...actual, loadSettings: vi.fn(), }; }); const startAutoMemoryIfEnabledMock = vi.fn(); vi.mock('../utils/autoMemory.js', () => ({ startAutoMemoryIfEnabled: (config: Config) => startAutoMemoryIfEnabledMock(config), })); describe('AcpSessionManager', () => { let mockConfig: Mocked; let mockSettings: Mocked; let mockArgv: CliArgs; let mockConnection: Mocked; let manager: AcpSessionManager; beforeEach(() => { mockConfig = { refreshAuth: vi.fn(), initialize: vi.fn(), waitForMcpInit: vi.fn(), getFileSystemService: vi.fn(), setFileSystemService: vi.fn(), getContentGeneratorConfig: vi.fn(), getActiveModel: vi.fn().mockReturnValue('gemini-pro'), getModel: vi.fn().mockReturnValue('gemini-pro'), getGeminiClient: vi.fn().mockReturnValue({ startChat: vi.fn().mockResolvedValue({}), }), getMessageBus: vi.fn().mockReturnValue({ publish: vi.fn(), subscribe: vi.fn(), unsubscribe: vi.fn(), }), getApprovalMode: vi.fn().mockReturnValue('default'), isPlanEnabled: vi.fn().mockReturnValue(true), getGemini31LaunchedSync: vi.fn().mockReturnValue(false), getHasAccessToPreviewModel: vi.fn().mockReturnValue(false), getCheckpointingEnabled: vi.fn().mockReturnValue(false), getDisableAlwaysAllow: vi.fn().mockReturnValue(false), validatePathAccess: vi.fn().mockReturnValue(null), getWorkspaceContext: vi.fn().mockReturnValue({ addReadOnlyPath: vi.fn(), }), getPolicyEngine: vi.fn().mockReturnValue({ addRule: vi.fn(), }), messageBus: { publish: vi.fn(), subscribe: vi.fn(), unsubscribe: vi.fn(), } as unknown as MessageBus, storage: { getWorkspaceAutoSavedPolicyPath: vi.fn(), getAutoSavedPolicyPath: vi.fn(), } as unknown as Storage, get config() { return this; }, } as unknown as Mocked; mockSettings = { merged: { security: { auth: { selectedType: 'login_with_google' } }, mcpServers: {}, }, setValue: vi.fn(), } as unknown as Mocked; mockArgv = {} as unknown as CliArgs; mockConnection = { sessionUpdate: vi.fn(), requestPermission: vi.fn(), } as unknown as Mocked; (loadCliConfig as unknown as Mock).mockResolvedValue(mockConfig); (loadSettings as unknown as Mock).mockImplementation(() => ({ merged: { security: { auth: { selectedType: AuthType.LOGIN_WITH_GOOGLE }, enablePermanentToolApproval: true, }, mcpServers: {}, }, setValue: vi.fn(), })); manager = new AcpSessionManager(mockSettings, mockArgv, mockConnection); vi.mock('node:crypto', () => ({ randomUUID: () => 'test-session-id', })); }); afterEach(() => { vi.restoreAllMocks(); }); it('should create a new session', async () => { vi.useFakeTimers(); mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); const response = await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); expect(response.sessionId).toBe('test-session-id'); expect(loadCliConfig).toHaveBeenCalled(); expect(mockConfig.initialize).toHaveBeenCalled(); expect(mockConfig.getGeminiClient).toHaveBeenCalled(); // Verify deferred call (sendAvailableCommands) await vi.runAllTimersAsync(); expect(mockConnection.sessionUpdate).toHaveBeenCalledWith( expect.objectContaining({ update: expect.objectContaining({ sessionUpdate: 'available_commands_update', }), }), ); vi.useRealTimers(); }); it('should return modes without plan mode when plan is disabled', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); mockConfig.isPlanEnabled = vi.fn().mockReturnValue(false); mockConfig.getApprovalMode = vi.fn().mockReturnValue('default'); const response = await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); expect(response.modes).toEqual({ availableModes: [ { id: 'default', name: 'Default', description: 'Prompts for approval' }, { id: 'autoEdit', name: 'Auto Edit', description: 'Auto-approves edit tools', }, { id: 'yolo', name: 'YOLO', description: 'Auto-approves all tools' }, ], currentModeId: 'default', }); }); it('should include preview models when user has access', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); mockConfig.getHasAccessToPreviewModel = vi.fn().mockReturnValue(true); mockConfig.getGemini31LaunchedSync = vi.fn().mockReturnValue(true); const response = await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); expect(response.models?.availableModels).toEqual( expect.arrayContaining([ expect.objectContaining({ modelId: GEMINI_MODEL_ALIAS_AUTO, name: expect.stringContaining('Auto'), }), ]), ); }); it('should NOT include retired preview models (none) in available models', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); mockConfig.getHasAccessToPreviewModel = vi.fn().mockReturnValue(true); mockConfig.getGemini31LaunchedSync = vi.fn().mockReturnValue(true); const response = await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); const modelIds = response.models?.availableModels?.map((m) => m.modelId) ?? []; expect(modelIds).not.toContain('none'); }); it('should return modes with plan mode when plan is enabled', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); mockConfig.isPlanEnabled = vi.fn().mockReturnValue(true); mockConfig.getApprovalMode = vi.fn().mockReturnValue('plan'); const response = await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); expect(response.modes).toEqual({ availableModes: [ { id: 'default', name: 'Default', description: 'Prompts for approval' }, { id: 'autoEdit', name: 'Auto Edit', description: 'Auto-approves edit tools', }, { id: 'yolo', name: 'YOLO', description: 'Auto-approves all tools' }, { id: 'plan', name: 'Plan', description: 'Read-only mode' }, ], currentModeId: 'plan', }); }); it('should fail session creation if Gemini API key is missing', async () => { (loadSettings as unknown as Mock).mockImplementation(() => ({ merged: { security: { auth: { selectedType: AuthType.USE_GEMINI } }, mcpServers: {}, }, setValue: vi.fn(), })); mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: undefined, }); await expect( manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ), ).rejects.toMatchObject({ message: 'Gemini API key is missing or not configured.', }); }); it('should create a new session with mcp servers', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); const mcpServers = [ { name: 'test-server', command: 'node', args: ['server.js'], env: [{ name: 'KEY', value: 'VALUE' }], }, ]; await manager.newSession( { cwd: '/tmp', mcpServers, }, {}, ); expect(loadCliConfig).toHaveBeenCalledWith( expect.objectContaining({ mcpServers: expect.objectContaining({ 'test-server': expect.objectContaining({ command: 'node', args: ['server.js'], env: { KEY: 'VALUE' }, }), }), }), 'test-session-id', mockArgv, { cwd: '/tmp' }, ); }); it('should handle authentication failure gracefully', async () => { mockConfig.refreshAuth.mockRejectedValue(new Error('Auth failed')); await expect( manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ), ).rejects.toMatchObject({ message: 'Auth failed', }); }); it('should initialize file system service if client supports it', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); manager.setClientCapabilities({ fs: { readTextFile: true, writeTextFile: true }, }); await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); expect(mockConfig.setFileSystemService).toHaveBeenCalled(); }); it('should start auto memory for new ACP sessions', async () => { mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({ apiKey: 'test-key', }); await manager.newSession( { cwd: '/tmp', mcpServers: [], }, {}, ); expect(startAutoMemoryIfEnabledMock).toHaveBeenCalledWith(mockConfig); }); });