diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index b8c2c6e29..42ea0632e 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -21,6 +21,15 @@ import { isWorkspaceTrusted } from './trustedFolders.js'; const mockWriteStderrLine = vi.hoisted(() => vi.fn()); const mockWriteStdoutLine = vi.hoisted(() => vi.fn()); +const mockSessionServiceInstance = vi.hoisted(() => ({ + loadLastSession: vi.fn(), + loadSession: vi.fn(), + forkSession: vi.fn(), + sessionExists: vi.fn(), +})); +const mockSessionServiceCtor = vi.hoisted(() => + vi.fn(() => mockSessionServiceInstance), +); vi.mock('../utils/stdioHelpers.js', () => ({ writeStderrLine: mockWriteStderrLine, @@ -139,6 +148,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { NativeLspService: vi .fn() .mockImplementation(() => createNativeLspServiceInstance()), + SessionService: mockSessionServiceCtor, SkillManager: SkillManagerMock, IdeClient: { getInstance: vi.fn().mockResolvedValue({ @@ -412,6 +422,46 @@ describe('parseArguments', () => { expect(argv.continue).toBe(true); }); + it('should parse --fork-session with --resume', async () => { + process.argv = [ + 'node', + 'script.js', + '--resume', + '123e4567-e89b-12d3-a456-426614174000', + '--fork-session', + ]; + const argv = await parseArguments(); + expect(argv.resume).toBe('123e4567-e89b-12d3-a456-426614174000'); + expect(argv.forkSession).toBe(true); + }); + + it('should parse --fork-session with the --resume picker form', async () => { + process.argv = ['node', 'script.js', '--resume', '--fork-session']; + const argv = await parseArguments(); + // Empty string is the existing yargs shape for picker form: --resume + // without an explicit session ID. + expect(argv.resume).toBe(''); + expect(argv.forkSession).toBe(true); + }); + + it('should reject --fork-session without --resume or --continue', async () => { + process.argv = ['node', 'script.js', '--fork-session']; + const mockExit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit called'); + }); + mockWriteStderrLine.mockClear(); + + await expect(parseArguments()).rejects.toThrow('process.exit called'); + + expect(mockWriteStderrLine).toHaveBeenCalledWith( + expect.stringContaining( + '--fork-session must be used with --resume or --continue', + ), + ); + + mockExit.mockRestore(); + }); + it('should convert positional query argument to prompt by default', async () => { process.argv = ['node', 'script.js', 'Hi Gemini']; const argv = await parseArguments(); @@ -787,6 +837,14 @@ describe('loadCliConfig', () => { nativeLspServiceMock.mockImplementation( () => createNativeLspServiceInstance() as unknown as NativeLspService, ); + mockSessionServiceCtor.mockImplementation(() => mockSessionServiceInstance); + mockSessionServiceInstance.loadLastSession.mockResolvedValue(undefined); + mockSessionServiceInstance.loadSession.mockResolvedValue(undefined); + mockSessionServiceInstance.forkSession.mockResolvedValue({ + filePath: '/mock/fork.jsonl', + copiedCount: 1, + }); + mockSessionServiceInstance.sessionExists.mockResolvedValue(false); vi.mocked(os.homedir).mockReturnValue('/mock/home/user'); vi.stubEnv('GEMINI_API_KEY', 'test-api-key'); }); @@ -853,19 +911,115 @@ describe('loadCliConfig', () => { expect(config.getIncludePartialMessages()).toBe(true); }); + it('should fork and load a new session when --resume is combined with --fork-session', async () => { + const sourceSessionId = '123e4567-e89b-42d3-a456-426614174000'; + const sourceData = { + conversation: { sessionId: sourceSessionId, messages: [] }, + uiHistory: [], + }; + const forkedData = { + conversation: { sessionId: 'forked-session-id', messages: [] }, + uiHistory: [], + }; + mockSessionServiceInstance.loadSession.mockImplementation( + async (sessionId: string) => { + if (sessionId === sourceSessionId) return sourceData; + return forkedData; + }, + ); + + const config = await loadCliConfig({}, { + resume: sourceSessionId, + forkSession: true, + } as CliArgs); + + expect(mockSessionServiceInstance.forkSession).toHaveBeenCalledWith( + sourceSessionId, + config.getSessionId(), + ); + expect(config.getSessionId()).toBe( + mockSessionServiceInstance.forkSession.mock.calls[0]?.[1], + ); + expect(mockSessionServiceInstance.loadSession).toHaveBeenCalledWith( + config.getSessionId(), + ); + }); + + it('should explain when --fork-session fails to copy the source session', async () => { + const sourceSessionId = '123e4567-e89b-42d3-a456-426614174000'; + const sourceData = { + conversation: { sessionId: sourceSessionId, messages: [] }, + uiHistory: [], + }; + mockSessionServiceInstance.loadSession.mockResolvedValue(sourceData); + mockSessionServiceInstance.forkSession.mockRejectedValue( + new Error('source session belongs to another project'), + ); + const mockExit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit called'); + }); + + await expect( + loadCliConfig({}, { + resume: sourceSessionId, + forkSession: true, + } as CliArgs), + ).rejects.toThrow('process.exit called'); + + expect(mockWriteStderrLine).toHaveBeenCalledWith( + `Failed to fork session ${sourceSessionId}: source session belongs to another project`, + ); + expect(mockExit).toHaveBeenCalledWith(1); + }); + + it('should explain when --continue --fork-session has no saved session to fork', async () => { + const mockExit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit called'); + }); + + await expect( + loadCliConfig({}, { + continue: true, + forkSession: true, + } as CliArgs), + ).rejects.toThrow('process.exit called'); + + expect(mockWriteStderrLine).toHaveBeenCalledWith( + 'Cannot use --fork-session with --continue: no saved session found to fork.', + ); + expect(mockExit).toHaveBeenCalledWith(1); + }); + it('should use internal sandbox session ID without treating it as a new session', async () => { const sessionId = '123e4567-e89b-12d3-a456-426614174000'; + vi.stubEnv('SANDBOX', 'sandbox-exec'); process.argv = ['node', 'script.js', '--sandbox-session-id', sessionId]; - const sessionExistsSpy = vi.spyOn( - ServerConfig.SessionService.prototype, - 'sessionExists', - ); const argv = await parseArguments(); const settings: Settings = {}; const config = await loadCliConfig(settings, argv); expect(config.getSessionId()).toBe(sessionId); - expect(sessionExistsSpy).not.toHaveBeenCalled(); + expect(mockSessionServiceInstance.sessionExists).not.toHaveBeenCalled(); + }); + + it('should reject direct use of the internal sandbox session ID flag', async () => { + const sessionId = '123e4567-e89b-12d3-a456-426614174000'; + const mockExit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit called'); + }); + + process.argv = ['node', 'script.js', '--sandbox-session-id', sessionId]; + const argv = await parseArguments(); + + await expect(loadCliConfig({}, argv)).rejects.toThrow( + 'process.exit called', + ); + + expect(mockWriteStderrLine).toHaveBeenCalledWith( + '--sandbox-session-id is for internal sandbox use only.', + ); + expect(mockExit).toHaveBeenCalledWith(1); + expect(mockSessionServiceInstance.sessionExists).not.toHaveBeenCalled(); }); it('should reset context filenames to defaults when context.fileName is not configured', async () => { diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 1a56c3edc..17e7a8bee 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -47,6 +47,7 @@ import { hideBin } from 'yargs/helpers'; import * as fs from 'node:fs'; import * as path from 'node:path'; import { homedir } from 'node:os'; +import { randomUUID } from 'node:crypto'; import stripJsonComments from 'strip-json-comments'; import { resolvePath } from '../utils/resolvePath.js'; @@ -159,6 +160,11 @@ export interface CliArgs { resume: string | undefined; /** Specify a session ID without session resumption */ sessionId: string | undefined; + /** + * Create a new forked session from the resumed session. Must be used with + * --resume or --continue. + */ + forkSession?: boolean | undefined; /** Internal: preserve the outer session ID when relaunching in a sandbox */ sandboxSessionId?: string | undefined; maxSessionTurns: number | undefined; @@ -805,6 +811,12 @@ export async function parseArguments(): Promise { type: 'string', description: 'Specify a session ID for this run.', }) + .option('fork-session', { + type: 'boolean', + description: + 'Create a new forked session from the resumed session. Must be used with --resume or --continue.', + default: false, + }) .option('sandbox-session-id', { type: 'string', hidden: true, @@ -906,9 +918,13 @@ export async function parseArguments(): Promise { if (argv['continue'] && argv['resume']) { return 'Cannot use both --continue and --resume together. Use --continue to resume the latest session, or --resume to resume a specific session.'; } - if (argv['sessionId'] && (argv['continue'] || argv['resume'])) { + const hasResume = argv['resume'] !== undefined; + if (argv['sessionId'] && (argv['continue'] || hasResume)) { return 'Cannot use --session-id with --continue or --resume. Use --session-id to start a new session with a specific ID, or use --continue/--resume to resume an existing session.'; } + if (argv['forkSession'] && !(argv['continue'] || hasResume)) { + return '--fork-session must be used with --resume or --continue.'; + } if ( argv['sandboxSessionId'] && (argv['sessionId'] || argv['continue'] || argv['resume']) @@ -1535,6 +1551,11 @@ export async function loadCliConfig( sessionData = await sessionService.loadLastSession(); if (sessionData) { sessionId = sessionData.conversation.sessionId; + } else if (argv.forkSession) { + writeStderrLine( + 'Cannot use --fork-session with --continue: no saved session found to fork.', + ); + process.exit(1); } } @@ -1550,7 +1571,30 @@ export async function loadCliConfig( process.exit(1); } } + + if (argv.forkSession && sessionId) { + const sourceSessionId = sessionId; + const forkedSessionId = randomUUID(); + try { + await sessionService.forkSession(sourceSessionId, forkedSessionId); + } catch (err) { + writeStderrLine( + `Failed to fork session ${sourceSessionId}: ${err instanceof Error ? err.message : String(err)}`, + ); + process.exit(1); + } + sessionId = forkedSessionId; + sessionData = await sessionService.loadSession(forkedSessionId); + if (!sessionData) { + writeStderrLine(`Failed to load forked session ${forkedSessionId}.`); + process.exit(1); + } + } } else if (argv.sandboxSessionId) { + if (!process.env['SANDBOX']) { + writeStderrLine('--sandbox-session-id is for internal sandbox use only.'); + process.exit(1); + } sessionId = argv.sandboxSessionId; } else if (argv['sessionId']) { // Use provided session ID without session resumption