From 48e3932f65f23ead61f26a52d808b526f24f216d Mon Sep 17 00:00:00 2001 From: joshualitt Date: Thu, 13 Nov 2025 15:13:39 -0800 Subject: [PATCH] feat(core, cli): Add auth type to history checkpoint. (#13023) --- .../cli/src/ui/commands/chatCommand.test.ts | 59 +++++++++++++++++-- packages/cli/src/ui/commands/chatCommand.ts | 25 ++++++-- .../ui/hooks/slashCommandProcessor.test.tsx | 33 ----------- .../cli/src/ui/hooks/slashCommandProcessor.ts | 1 - packages/core/src/core/logger.test.ts | 57 ++++++++++++------ packages/core/src/core/logger.ts | 41 +++++++++---- 6 files changed, 143 insertions(+), 73 deletions(-) diff --git a/packages/cli/src/ui/commands/chatCommand.test.ts b/packages/cli/src/ui/commands/chatCommand.test.ts index 0eceb42e42..866d763995 100644 --- a/packages/cli/src/ui/commands/chatCommand.test.ts +++ b/packages/cli/src/ui/commands/chatCommand.test.ts @@ -10,7 +10,7 @@ import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest'; import type { SlashCommand, CommandContext } from './types.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; import type { Content } from '@google/genai'; -import type { GeminiClient } from '@google/gemini-cli-core'; +import { AuthType, type GeminiClient } from '@google/gemini-cli-core'; import * as fsPromises from 'node:fs/promises'; import { chatCommand, serializeHistoryToMarkdown } from './chatCommand.js'; @@ -52,7 +52,7 @@ describe('chatCommand', () => { getHistory: mockGetHistory, }); mockSaveCheckpoint = vi.fn().mockResolvedValue(undefined); - mockLoadCheckpoint = vi.fn().mockResolvedValue([]); + mockLoadCheckpoint = vi.fn().mockResolvedValue({ history: [] }); mockDeleteCheckpoint = vi.fn().mockResolvedValue(true); mockContext = createMockCommandContext({ @@ -66,6 +66,9 @@ describe('chatCommand', () => { storage: { getProjectTempDir: () => '/project/root/.gemini/tmp/mockhash', }, + getContentGeneratorConfig: () => ({ + authType: AuthType.LOGIN_WITH_GOOGLE, + }), }, logger: { saveCheckpoint: mockSaveCheckpoint, @@ -215,7 +218,10 @@ describe('chatCommand', () => { const result = await saveCommand?.action?.(mockContext, tag); expect(mockCheckpointExists).not.toHaveBeenCalled(); // Should skip existence check - expect(mockSaveCheckpoint).toHaveBeenCalledWith(history, tag); + expect(mockSaveCheckpoint).toHaveBeenCalledWith( + { history, authType: AuthType.LOGIN_WITH_GOOGLE }, + tag, + ); expect(result).toEqual({ type: 'message', messageType: 'info', @@ -244,7 +250,7 @@ describe('chatCommand', () => { }); it('should inform if checkpoint is not found', async () => { - mockLoadCheckpoint.mockResolvedValue([]); + mockLoadCheckpoint.mockResolvedValue({ history: [] }); const result = await resumeCommand?.action?.(mockContext, badTag); @@ -255,12 +261,53 @@ describe('chatCommand', () => { }); }); - it('should resume a conversation', async () => { + it('should resume a conversation with matching authType', async () => { const conversation: Content[] = [ { role: 'user', parts: [{ text: 'hello gemini' }] }, { role: 'model', parts: [{ text: 'hello world' }] }, ]; - mockLoadCheckpoint.mockResolvedValue(conversation); + mockLoadCheckpoint.mockResolvedValue({ + history: conversation, + authType: AuthType.LOGIN_WITH_GOOGLE, + }); + + const result = await resumeCommand?.action?.(mockContext, goodTag); + + expect(result).toEqual({ + type: 'load_history', + history: [ + { type: 'user', text: 'hello gemini' }, + { type: 'gemini', text: 'hello world' }, + ] as HistoryItemWithoutId[], + clientHistory: conversation, + }); + }); + + it('should block resuming a conversation with mismatched authType', async () => { + const conversation: Content[] = [ + { role: 'user', parts: [{ text: 'hello gemini' }] }, + { role: 'model', parts: [{ text: 'hello world' }] }, + ]; + mockLoadCheckpoint.mockResolvedValue({ + history: conversation, + authType: AuthType.USE_GEMINI, + }); + + const result = await resumeCommand?.action?.(mockContext, goodTag); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: `Cannot resume chat. It was saved with a different authentication method (${AuthType.USE_GEMINI}) than the current one (${AuthType.LOGIN_WITH_GOOGLE}).`, + }); + }); + + it('should resume a legacy conversation without authType', async () => { + const conversation: Content[] = [ + { role: 'user', parts: [{ text: 'hello gemini' }] }, + { role: 'model', parts: [{ text: 'hello world' }] }, + ]; + mockLoadCheckpoint.mockResolvedValue({ history: conversation }); const result = await resumeCommand?.action?.(mockContext, goodTag); diff --git a/packages/cli/src/ui/commands/chatCommand.ts b/packages/cli/src/ui/commands/chatCommand.ts index d26199b2b7..149351976a 100644 --- a/packages/cli/src/ui/commands/chatCommand.ts +++ b/packages/cli/src/ui/commands/chatCommand.ts @@ -128,11 +128,14 @@ const saveCommand: SlashCommand = { const history = chat.getHistory(); if (history.length > 2) { - await logger.saveCheckpoint(history, tag); + const authType = config?.getContentGeneratorConfig()?.authType; + await logger.saveCheckpoint({ history, authType }, tag); return { type: 'message', messageType: 'info', - content: `Conversation checkpoint saved with tag: ${decodeTagName(tag)}.`, + content: `Conversation checkpoint saved with tag: ${decodeTagName( + tag, + )}.`, }; } else { return { @@ -160,9 +163,10 @@ const resumeCommand: SlashCommand = { }; } - const { logger } = context.services; + const { logger, config } = context.services; await logger.initialize(); - const conversation = await logger.loadCheckpoint(tag); + const checkpoint = await logger.loadCheckpoint(tag); + const conversation = checkpoint.history; if (conversation.length === 0) { return { @@ -172,6 +176,19 @@ const resumeCommand: SlashCommand = { }; } + const currentAuthType = config?.getContentGeneratorConfig()?.authType; + if ( + checkpoint.authType && + currentAuthType && + checkpoint.authType !== currentAuthType + ) { + return { + type: 'message', + messageType: 'error', + content: `Cannot resume chat. It was saved with a different authentication method (${checkpoint.authType}) than the current one (${currentAuthType}).`, + }; + } + const rolemap: { [key: string]: MessageType } = { user: MessageType.USER, model: MessageType.GEMINI, diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx b/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx index 4b247c9361..99afe37b10 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx @@ -495,39 +495,6 @@ describe('useSlashCommandProcessor', () => { ); }); - it('should strip thoughts when handling "load_history" action', async () => { - const mockClient = { - setHistory: vi.fn(), - stripThoughtsFromHistory: vi.fn(), - } as unknown as GeminiClient; - vi.spyOn(mockConfig, 'getGeminiClient').mockReturnValue(mockClient); - - const historyWithThoughts = [ - { - role: 'model', - parts: [{ text: 'response', thoughtSignature: 'CikB...' }], - }, - ]; - const command = createTestCommand({ - name: 'loadwiththoughts', - action: vi.fn().mockResolvedValue({ - type: 'load_history', - history: [{ type: MessageType.GEMINI, text: 'response' }], - clientHistory: historyWithThoughts, - }), - }); - - const result = await setupProcessorHook([command]); - await waitFor(() => expect(result.current.slashCommands).toHaveLength(1)); - - await act(async () => { - await result.current.handleSlashCommand('/loadwiththoughts'); - }); - - expect(mockClient.setHistory).toHaveBeenCalledTimes(1); - expect(mockClient.stripThoughtsFromHistory).toHaveBeenCalledWith(); - }); - it('should handle a "quit" action', async () => { const quitAction = vi .fn() diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 0fc0570188..692869f167 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -413,7 +413,6 @@ export const useSlashCommandProcessor = ( } case 'load_history': { config?.getGeminiClient()?.setHistory(result.clientHistory); - config?.getGeminiClient()?.stripThoughtsFromHistory(); fullCommandContext.ui.clear(); result.history.forEach((item, index) => { fullCommandContext.ui.addItem(item, index); diff --git a/packages/core/src/core/logger.test.ts b/packages/core/src/core/logger.test.ts index 70697b6f98..b389dcbe6b 100644 --- a/packages/core/src/core/logger.test.ts +++ b/packages/core/src/core/logger.test.ts @@ -20,6 +20,7 @@ import { encodeTagName, decodeTagName, } from './logger.js'; +import { AuthType } from './contentGenerator.js'; import { Storage } from '../config/storage.js'; import { promises as fs, existsSync } from 'node:fs'; import path from 'node:path'; @@ -432,13 +433,19 @@ describe('Logger', () => { encodedTag: '..%2F..%2Fsecret', }, ])('should save a checkpoint', async ({ tag, encodedTag }) => { - await logger.saveCheckpoint(conversation, tag); + await logger.saveCheckpoint( + { history: conversation, authType: AuthType.LOGIN_WITH_GOOGLE }, + tag, + ); const taggedFilePath = path.join( TEST_GEMINI_DIR, `checkpoint-${encodedTag}.json`, ); const fileContent = await fs.readFile(taggedFilePath, 'utf-8'); - expect(JSON.parse(fileContent)).toEqual(conversation); + expect(JSON.parse(fileContent)).toEqual({ + history: conversation, + authType: AuthType.LOGIN_WITH_GOOGLE, + }); }); it('should not throw if logger is not initialized', async () => { @@ -452,7 +459,7 @@ describe('Logger', () => { .mockImplementation(() => {}); await expect( - uninitializedLogger.saveCheckpoint(conversation, 'tag'), + uninitializedLogger.saveCheckpoint({ history: conversation }, 'tag'), ).resolves.not.toThrow(); expect(consoleErrorSpy).toHaveBeenCalledWith( 'Logger not initialized or checkpoint file path not set. Cannot save a checkpoint.', @@ -492,10 +499,13 @@ describe('Logger', () => { encodedTag: '..%2F..%2Fsecret', }, ])('should load from a checkpoint', async ({ tag, encodedTag }) => { - const taggedConversation = [ - ...conversation, - { role: 'user', parts: [{ text: 'hello' }] }, - ]; + const taggedConversation = { + history: [ + ...conversation, + { role: 'user', parts: [{ text: 'hello' }] }, + ], + authType: AuthType.USE_GEMINI, + }; const taggedFilePath = path.join( TEST_GEMINI_DIR, `checkpoint-${encodedTag}.json`, @@ -511,18 +521,31 @@ describe('Logger', () => { expect(decodeTagName(encodedTag)).toBe(tag); }); - it('should return an empty array if a tagged checkpoint file does not exist', async () => { - const loaded = await logger.loadCheckpoint('nonexistent-tag'); - expect(loaded).toEqual([]); + it('should load a legacy checkpoint without authType', async () => { + const tag = 'legacy-tag'; + const encodedTag = 'legacy-tag'; + const taggedFilePath = path.join( + TEST_GEMINI_DIR, + `checkpoint-${encodedTag}.json`, + ); + await fs.writeFile(taggedFilePath, JSON.stringify(conversation, null, 2)); + + const loaded = await logger.loadCheckpoint(tag); + expect(loaded).toEqual({ history: conversation }); }); - it('should return an empty array if the checkpoint file does not exist', async () => { + it('should return an empty history if a tagged checkpoint file does not exist', async () => { + const loaded = await logger.loadCheckpoint('nonexistent-tag'); + expect(loaded).toEqual({ history: [] }); + }); + + it('should return an empty history if the checkpoint file does not exist', async () => { await fs.unlink(TEST_CHECKPOINT_FILE_PATH); // Ensure it's gone const loaded = await logger.loadCheckpoint('missing'); - expect(loaded).toEqual([]); + expect(loaded).toEqual({ history: [] }); }); - it('should return an empty array if the file contains invalid JSON', async () => { + it('should return an empty history if the file contains invalid JSON', async () => { const tag = 'invalid-json-tag'; const encodedTag = 'invalid-json-tag'; const taggedFilePath = path.join( @@ -534,14 +557,14 @@ describe('Logger', () => { .spyOn(console, 'error') .mockImplementation(() => {}); const loadedCheckpoint = await logger.loadCheckpoint(tag); - expect(loadedCheckpoint).toEqual([]); + expect(loadedCheckpoint).toEqual({ history: [] }); expect(consoleErrorSpy).toHaveBeenCalledWith( expect.stringContaining('Failed to read or parse checkpoint file'), expect.any(Error), ); }); - it('should return an empty array if logger is not initialized', async () => { + it('should return an empty history if logger is not initialized', async () => { const uninitializedLogger = new Logger( testSessionId, new Storage(process.cwd()), @@ -551,7 +574,7 @@ describe('Logger', () => { .spyOn(console, 'error') .mockImplementation(() => {}); const loadedCheckpoint = await uninitializedLogger.loadCheckpoint('tag'); - expect(loadedCheckpoint).toEqual([]); + expect(loadedCheckpoint).toEqual({ history: [] }); expect(consoleErrorSpy).toHaveBeenCalledWith( 'Logger not initialized or checkpoint file path not set. Cannot load checkpoint.', ); @@ -726,7 +749,7 @@ describe('Logger', () => { ); const loaded = await logger.loadCheckpoint(tag); - expect(loaded).toEqual(taggedConversation); + expect(loaded.history).toEqual(taggedConversation); }); }); diff --git a/packages/core/src/core/logger.ts b/packages/core/src/core/logger.ts index ce7812b6d8..003891010c 100644 --- a/packages/core/src/core/logger.ts +++ b/packages/core/src/core/logger.ts @@ -7,6 +7,7 @@ import path from 'node:path'; import { promises as fs } from 'node:fs'; import type { Content } from '@google/genai'; +import type { AuthType } from './contentGenerator.js'; import type { Storage } from '../config/storage.js'; import { debugLogger } from '../utils/debugLogger.js'; import { coreEvents } from '../utils/events.js'; @@ -25,6 +26,11 @@ export interface LogEntry { message: string; } +export interface Checkpoint { + history: Content[]; + authType?: AuthType; +} + // This regex matches any character that is NOT a letter (a-z, A-Z), // a number (0-9), a hyphen (-), an underscore (_), or a dot (.). @@ -314,7 +320,7 @@ export class Logger { return newPath; } - async saveCheckpoint(conversation: Content[], tag: string): Promise { + async saveCheckpoint(checkpoint: Checkpoint, tag: string): Promise { if (!this.initialized) { debugLogger.error( 'Logger not initialized or checkpoint file path not set. Cannot save a checkpoint.', @@ -324,42 +330,53 @@ export class Logger { // Always save with the new encoded path. const path = this._checkpointPath(tag); try { - await fs.writeFile(path, JSON.stringify(conversation, null, 2), 'utf-8'); + await fs.writeFile(path, JSON.stringify(checkpoint, null, 2), 'utf-8'); } catch (error) { debugLogger.error('Error writing to checkpoint file:', error); } } - async loadCheckpoint(tag: string): Promise { + async loadCheckpoint(tag: string): Promise { if (!this.initialized) { debugLogger.error( 'Logger not initialized or checkpoint file path not set. Cannot load checkpoint.', ); - return []; + return { history: [] }; } const path = await this._getCheckpointPath(tag); try { const fileContent = await fs.readFile(path, 'utf-8'); const parsedContent = JSON.parse(fileContent); - if (!Array.isArray(parsedContent)) { - debugLogger.warn( - `Checkpoint file at ${path} is not a valid JSON array. Returning empty checkpoint.`, - ); - return []; + + // Handle legacy format (just an array of Content) + if (Array.isArray(parsedContent)) { + return { history: parsedContent as Content[] }; } - return parsedContent as Content[]; + + if ( + typeof parsedContent === 'object' && + parsedContent !== null && + 'history' in parsedContent + ) { + return parsedContent as Checkpoint; + } + + debugLogger.warn( + `Checkpoint file at ${path} has an unknown format. Returning empty checkpoint.`, + ); + return { history: [] }; } catch (error) { const nodeError = error as NodeJS.ErrnoException; if (nodeError.code === 'ENOENT') { // This is okay, it just means the checkpoint doesn't exist in either format. - return []; + return { history: [] }; } debugLogger.error( `Failed to read or parse checkpoint file ${path}:`, error, ); - return []; + return { history: [] }; } }