feat(cli): add fork-session resume flag (#4159)

* feat(cli): add fork-session resume flag

* fix(cli): address fork-session review feedback

* fix(cli): handle fork session copy failures

* fix(cli): guard sandbox session handoff flag
This commit is contained in:
qqqys 2026-05-17 00:27:52 +08:00 committed by GitHub
parent b9590283c0
commit daaa85e98e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 204 additions and 6 deletions

View file

@ -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 () => {

View file

@ -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<CliArgs> {
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<CliArgs> {
if (argv['continue'] && argv['resume']) {
return 'Cannot use both --continue and --resume together. Use --continue to resume the latest session, or --resume <sessionId> 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