| |
| |
| |
| |
| |
|
|
| import { |
| describe, |
| it, |
| expect, |
| vi, |
| beforeEach, |
| type Mocked, |
| type Mock, |
| } from 'vitest'; |
| import { GeminiAgent } from './acpRpcDispatcher.js'; |
| import * as acp from '@agentclientprotocol/sdk'; |
| import { |
| ApprovalMode, |
| AuthType, |
| type Config, |
| CoreToolCallStatus, |
| } from '@google/gemini-cli-core'; |
| import { loadCliConfig, type CliArgs } from '../config/config.js'; |
| import { |
| SessionSelector, |
| convertSessionToHistoryFormats, |
| } from '../utils/sessionUtils.js'; |
| import { convertSessionToClientHistory } from '@google/gemini-cli-core'; |
| import type { LoadedSettings } from '../config/settings.js'; |
| import { waitFor } from '../test-utils/async.js'; |
|
|
| vi.mock('../config/config.js', () => ({ |
| loadCliConfig: vi.fn(), |
| })); |
|
|
| vi.mock('../utils/sessionUtils.js', async (importOriginal) => { |
| const actual = |
| await importOriginal<typeof import('../utils/sessionUtils.js')>(); |
| return { |
| ...actual, |
| SessionSelector: vi.fn(), |
| convertSessionToHistoryFormats: vi.fn(), |
| }; |
| }); |
|
|
| vi.mock('@google/gemini-cli-core', async (importOriginal) => { |
| const actual = |
| await importOriginal<typeof import('@google/gemini-cli-core')>(); |
| return { |
| ...actual, |
| CoreToolCallStatus: { |
| Validating: 'validating', |
| Scheduled: 'scheduled', |
| Error: 'error', |
| Success: 'success', |
| Executing: 'executing', |
| Cancelled: 'cancelled', |
| AwaitingApproval: 'awaiting_approval', |
| }, |
| LlmRole: { |
| MAIN: 'main', |
| SUBAGENT: 'subagent', |
| UTILITY_TOOL: 'utility_tool', |
| USER: 'user', |
| MODEL: 'model', |
| SYSTEM: 'system', |
| TOOL: 'tool', |
| }, |
| convertSessionToClientHistory: vi.fn(), |
| }; |
| }); |
|
|
| describe('GeminiAgent Session Resume', () => { |
| let mockConfig: Mocked<Config>; |
| let mockSettings: Mocked<LoadedSettings>; |
| let mockArgv: CliArgs; |
| let mockConnection: Mocked<acp.AgentSideConnection>; |
| let agent: GeminiAgent; |
|
|
| beforeEach(() => { |
| mockConfig = { |
| refreshAuth: vi.fn().mockResolvedValue(undefined), |
| initialize: vi.fn().mockResolvedValue(undefined), |
| getFileSystemService: vi.fn(), |
| setFileSystemService: vi.fn(), |
| getGeminiClient: vi.fn().mockReturnValue({ |
| initialize: vi.fn().mockResolvedValue(undefined), |
| resumeChat: vi.fn().mockResolvedValue(undefined), |
| getChat: vi.fn().mockReturnValue({}), |
| }), |
| storage: { |
| getProjectTempDir: vi.fn().mockReturnValue('/tmp/project'), |
| }, |
| getPolicyEngine: vi.fn().mockReturnValue({ |
| addRule: vi.fn(), |
| }), |
| messageBus: { |
| publish: vi.fn(), |
| subscribe: vi.fn(), |
| unsubscribe: vi.fn(), |
| }, |
| getMessageBus: vi.fn().mockReturnValue({ |
| publish: vi.fn(), |
| subscribe: vi.fn(), |
| unsubscribe: vi.fn(), |
| }), |
| getApprovalMode: vi.fn().mockReturnValue('default'), |
| isAutoMemoryEnabled: vi.fn().mockReturnValue(false), |
| isPlanEnabled: vi.fn().mockReturnValue(true), |
| getModel: vi.fn().mockReturnValue('gemini-pro'), |
| getHasAccessToPreviewModel: vi.fn().mockReturnValue(false), |
| getGemini31LaunchedSync: vi.fn().mockReturnValue(false), |
| getCheckpointingEnabled: vi.fn().mockReturnValue(false), |
| toolRegistry: { |
| getTool: vi.fn().mockReturnValue({ kind: 'read' }), |
| }, |
| get config() { |
| return this; |
| }, |
| } as unknown as Mocked<Config>; |
| mockSettings = { |
| merged: { |
| security: { auth: { selectedType: AuthType.LOGIN_WITH_GOOGLE } }, |
| mcpServers: {}, |
| }, |
| setValue: vi.fn(), |
| } as unknown as Mocked<LoadedSettings>; |
| mockArgv = {} as unknown as CliArgs; |
| mockConnection = { |
| sessionUpdate: vi.fn().mockResolvedValue(undefined), |
| } as unknown as Mocked<acp.AgentSideConnection>; |
|
|
| (loadCliConfig as Mock).mockResolvedValue(mockConfig); |
|
|
| agent = new GeminiAgent(mockConfig, mockSettings, mockArgv, mockConnection); |
| }); |
|
|
| it('should advertise loadSession capability', async () => { |
| const response = await agent.initialize({ |
| protocolVersion: acp.PROTOCOL_VERSION, |
| }); |
| expect(response.agentCapabilities?.loadSession).toBe(true); |
| }); |
|
|
| it('should load a session, resume chat, and stream all message types', async () => { |
| const sessionId = 'existing-session-id'; |
| const sessionData = { |
| sessionId, |
| messages: [ |
| { type: 'user', content: [{ text: 'Hello' }] }, |
| { |
| type: 'gemini', |
| content: [{ text: 'Hi there' }], |
| thoughts: [{ subject: 'Thinking', description: 'about greeting' }], |
| toolCalls: [ |
| { |
| id: 'call-1', |
| name: 'test_tool', |
| displayName: 'Test Tool', |
| status: CoreToolCallStatus.Success, |
| resultDisplay: 'Tool output', |
| }, |
| ], |
| }, |
| { |
| type: 'gemini', |
| content: [{ text: 'Trying a write' }], |
| toolCalls: [ |
| { |
| id: 'call-2', |
| name: 'write_file', |
| displayName: 'Write File', |
| status: CoreToolCallStatus.Error, |
| resultDisplay: 'Permission denied', |
| }, |
| ], |
| }, |
| ], |
| }; |
|
|
| (SessionSelector as unknown as Mock).mockImplementation(() => ({ |
| resolveSession: vi.fn().mockResolvedValue({ |
| sessionData, |
| sessionPath: '/path/to/session.json', |
| }), |
| })); |
|
|
| const mockClientHistory = [ |
| { role: 'user', parts: [{ text: 'Hello' }] }, |
| { role: 'model', parts: [{ text: 'Hi there' }] }, |
| ]; |
| (convertSessionToHistoryFormats as unknown as Mock).mockReturnValue({ |
| uiHistory: [], |
| }); |
| (convertSessionToClientHistory as unknown as Mock).mockReturnValue( |
| mockClientHistory, |
| ); |
|
|
| const response = await agent.loadSession({ |
| sessionId, |
| cwd: '/tmp', |
| mcpServers: [], |
| }); |
|
|
| expect(response).toEqual({ |
| modes: { |
| availableModes: [ |
| { |
| id: ApprovalMode.DEFAULT, |
| name: 'Default', |
| description: 'Prompts for approval', |
| }, |
| { |
| id: ApprovalMode.AUTO_EDIT, |
| name: 'Auto Edit', |
| description: 'Auto-approves edit tools', |
| }, |
| { |
| id: ApprovalMode.YOLO, |
| name: 'YOLO', |
| description: 'Auto-approves all tools', |
| }, |
| { |
| id: ApprovalMode.PLAN, |
| name: 'Plan', |
| description: 'Read-only mode', |
| }, |
| ], |
| currentModeId: ApprovalMode.DEFAULT, |
| }, |
| models: { |
| availableModels: expect.any(Array) as unknown, |
| currentModelId: 'gemini-pro', |
| }, |
| }); |
|
|
| |
| expect(mockConfig.getGeminiClient().resumeChat).toHaveBeenCalledWith( |
| mockClientHistory, |
| expect.objectContaining({ |
| conversation: sessionData, |
| filePath: '/path/to/session.json', |
| }), |
| ); |
|
|
| await waitFor(() => { |
| |
| expect(mockConnection.sessionUpdate).toHaveBeenCalledWith( |
| expect.objectContaining({ |
| update: expect.objectContaining({ |
| sessionUpdate: 'user_message_chunk', |
| content: expect.objectContaining({ text: 'Hello' }), |
| }), |
| }), |
| ); |
|
|
| |
| expect(mockConnection.sessionUpdate).toHaveBeenCalledWith( |
| expect.objectContaining({ |
| update: expect.objectContaining({ |
| sessionUpdate: 'agent_thought_chunk', |
| content: expect.objectContaining({ |
| text: '**Thinking**\nabout greeting', |
| }), |
| }), |
| }), |
| ); |
|
|
| |
| expect(mockConnection.sessionUpdate).toHaveBeenCalledWith( |
| expect.objectContaining({ |
| update: expect.objectContaining({ |
| sessionUpdate: 'agent_message_chunk', |
| content: expect.objectContaining({ text: 'Hi there' }), |
| }), |
| }), |
| ); |
|
|
| |
| expect(mockConnection.sessionUpdate).toHaveBeenCalledWith( |
| expect.objectContaining({ |
| update: expect.objectContaining({ |
| sessionUpdate: 'tool_call', |
| toolCallId: 'call-1', |
| status: 'completed', |
| title: 'Test Tool', |
| kind: 'read', |
| content: [ |
| { |
| type: 'content', |
| content: { type: 'text', text: 'Tool output' }, |
| }, |
| ], |
| }), |
| }), |
| ); |
|
|
| |
| expect(mockConnection.sessionUpdate).toHaveBeenCalledWith( |
| expect.objectContaining({ |
| update: expect.objectContaining({ |
| sessionUpdate: 'tool_call', |
| toolCallId: 'call-2', |
| status: 'failed', |
| title: 'Write File', |
| kind: 'read', |
| }), |
| }), |
| ); |
| }); |
| }); |
| }); |
|
|
