mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-05-19 16:28:28 +00:00
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:
parent
b9590283c0
commit
daaa85e98e
2 changed files with 204 additions and 6 deletions
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue