From e43def100bd7e8ce22abdecb313219c8753b7ed9 Mon Sep 17 00:00:00 2001 From: qqqys Date: Wed, 22 Apr 2026 19:13:53 +0800 Subject: [PATCH 1/6] feat(session): add /branch to fork the current conversation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces `/branch` (alias `/fork`), mirroring Claude Code's fork-session command. Writes a new JSONL under a fresh sessionId with every record stamped `forkedFrom: { sessionId, messageUuid }`, rebuilds `parentUuid` in write order so the fork is a clean linear descendant, and swaps the CLI into the new session with a Claude-style two-line announcement plus a `/resume ` hint. Core: - `SessionService.forkSession(src, new)` performs the copy. Uses `fs.openSync(path, 'wx', 0o600)` for exclusive create — atomic existence + open in one syscall, no TOCTOU window. Rejects invalid sessionId patterns, missing/empty sources, cross-project sources, and pre-existing targets. - `ChatRecord.forkedFrom` optional field records per-message lineage. - `SessionStartSource.Branch` lets hook consumers distinguish fork from resume. CLI: - `branchCommand` guards on `isIdleRef` so mid-stream forks can't tear the parent chain, and on `sessionExists` so empty sessions can't be forked. - `useBranchCommand` orchestrates finalize → fork → load → core swap → init → UI swap, in that order: anything that can still fail runs while the UI is still on the parent, so a throw leaves the user safely on the parent session instead of stranded with a cleared history. - Branch title is ` (Branch)` with `(Branch N)` collision bump (cap 99, then timestamp fallback). When no name is given it's derived from the first real user `ChatRecord` (skipping cron/notification subtypes), falling back to `Branched conversation`. - `/branch` is added to `SLASH_COMMANDS_SKIP_RECORDING` so the command itself doesn't bleed into the fork's tail. Tests cover: command guards; hook ordering; title collision bump; synthetic-record skip; empty-transcript fallback; core-throws-after-fork UI-preservation invariant; forkSession disk I/O including invalid ids, cross-project rejection, already-exists rejection. 🤖 Generated with [Qwen Code](https://github.com/QwenLM/qwen-code) Co-authored-by: Qwen-Coder --- .../cli/src/services/BuiltinCommandLoader.ts | 2 + packages/cli/src/ui/AppContainer.tsx | 15 + .../cli/src/ui/commands/branchCommand.test.ts | 76 +++++ packages/cli/src/ui/commands/branchCommand.ts | 60 ++++ packages/cli/src/ui/commands/types.ts | 4 + .../cli/src/ui/contexts/UIActionsContext.tsx | 2 + .../ui/hooks/slashCommandProcessor.test.ts | 52 ++++ .../cli/src/ui/hooks/slashCommandProcessor.ts | 5 + .../cli/src/ui/hooks/useBranchCommand.test.ts | 283 ++++++++++++++++++ packages/cli/src/ui/hooks/useBranchCommand.ts | 230 ++++++++++++++ packages/core/src/hooks/types.ts | 1 + .../core/src/services/chatRecordingService.ts | 18 ++ .../core/src/services/sessionService.test.ts | 207 +++++++++++++ packages/core/src/services/sessionService.ts | 86 ++++++ 14 files changed, 1041 insertions(+) create mode 100644 packages/cli/src/ui/commands/branchCommand.test.ts create mode 100644 packages/cli/src/ui/commands/branchCommand.ts create mode 100644 packages/cli/src/ui/hooks/useBranchCommand.test.ts create mode 100644 packages/cli/src/ui/hooks/useBranchCommand.ts diff --git a/packages/cli/src/services/BuiltinCommandLoader.ts b/packages/cli/src/services/BuiltinCommandLoader.ts index cdde266b85e..d78ea21fbff 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.ts @@ -12,6 +12,7 @@ import { agentsCommand } from '../ui/commands/agentsCommand.js'; import { arenaCommand } from '../ui/commands/arenaCommand.js'; import { approvalModeCommand } from '../ui/commands/approvalModeCommand.js'; import { authCommand } from '../ui/commands/authCommand.js'; +import { branchCommand } from '../ui/commands/branchCommand.js'; import { btwCommand } from '../ui/commands/btwCommand.js'; import { bugCommand } from '../ui/commands/bugCommand.js'; import { clearCommand } from '../ui/commands/clearCommand.js'; @@ -93,6 +94,7 @@ export class BuiltinCommandLoader implements ICommandLoader { arenaCommand, approvalModeCommand, authCommand, + branchCommand, btwCommand, bugCommand, clearCommand, diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 89bb17befaf..fdec91bf967 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -71,6 +71,7 @@ import { useSettingsCommand } from './hooks/useSettingsCommand.js'; import { useModelCommand } from './hooks/useModelCommand.js'; import { useArenaCommand } from './hooks/useArenaCommand.js'; import { useApprovalModeCommand } from './hooks/useApprovalModeCommand.js'; +import { useBranchCommand } from './hooks/useBranchCommand.js'; import { useResumeCommand } from './hooks/useResumeCommand.js'; import { useDeleteCommand } from './hooks/useDeleteCommand.js'; import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js'; @@ -592,6 +593,14 @@ export const AppContainer = (props: AppContainerProps) => { remount: refreshStatic, }); + const { handleBranch } = useBranchCommand({ + config, + historyManager, + startNewSession, + setSessionName, + remount: refreshStatic, + }); + const { isDeleteDialogOpen, openDeleteDialog, @@ -652,6 +661,7 @@ export const AppContainer = (props: AppContainerProps) => { openHooksDialog, openResumeDialog, handleResume, + handleBranch, openDeleteDialog, }), [ @@ -675,6 +685,7 @@ export const AppContainer = (props: AppContainerProps) => { openHooksDialog, openResumeDialog, handleResume, + handleBranch, openDeleteDialog, ], ); @@ -2336,6 +2347,8 @@ export const AppContainer = (props: AppContainerProps) => { openResumeDialog, closeResumeDialog, handleResume, + // Branch (fork) session + handleBranch, // Delete session dialog openDeleteDialog, closeDeleteDialog, @@ -2400,6 +2413,8 @@ export const AppContainer = (props: AppContainerProps) => { openResumeDialog, closeResumeDialog, handleResume, + // Branch (fork) session + handleBranch, // Delete session dialog openDeleteDialog, closeDeleteDialog, diff --git a/packages/cli/src/ui/commands/branchCommand.test.ts b/packages/cli/src/ui/commands/branchCommand.test.ts new file mode 100644 index 00000000000..df81527389f --- /dev/null +++ b/packages/cli/src/ui/commands/branchCommand.test.ts @@ -0,0 +1,76 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi } from 'vitest'; +import { branchCommand } from './branchCommand.js'; +import type { CommandContext } from './types.js'; + +function makeCtx( + overrides: { + isIdle?: boolean; + sessionExists?: boolean; + noConfig?: boolean; + } = {}, +): CommandContext { + const sessionService = { + sessionExists: vi.fn().mockResolvedValue(overrides.sessionExists ?? true), + }; + const config = overrides.noConfig + ? null + : ({ + getSessionId: () => '11111111-1111-1111-1111-111111111111', + getSessionService: () => sessionService, + } as unknown as NonNullable); + return { + services: { config, settings: {} as never, git: undefined, logger: null }, + ui: { + isIdleRef: { current: overrides.isIdle ?? true }, + } as unknown as CommandContext['ui'], + session: { stats: {} as never, sessionShellAllowlist: new Set() }, + } as unknown as CommandContext; +} + +describe('branchCommand', () => { + it('rejects when config is unavailable', async () => { + const result = await branchCommand.action!(makeCtx({ noConfig: true }), ''); + expect(result).toMatchObject({ type: 'message', messageType: 'error' }); + }); + + it('rejects when no conversation exists to branch from', async () => { + const result = await branchCommand.action!( + makeCtx({ sessionExists: false }), + '', + ); + expect(result).toMatchObject({ type: 'message', messageType: 'error' }); + expect((result as { content: string }).content).toMatch( + /No conversation to branch/, + ); + }); + + it('rejects while streaming or awaiting a tool confirmation', async () => { + const result = await branchCommand.action!(makeCtx({ isIdle: false }), ''); + expect(result).toMatchObject({ type: 'message', messageType: 'error' }); + expect((result as { content: string }).content).toMatch(/in progress/); + }); + + it('returns dialog action with no name when args are empty', async () => { + const result = await branchCommand.action!(makeCtx(), ' '); + expect(result).toEqual({ type: 'dialog', dialog: 'branch' }); + }); + + it('returns dialog action with trimmed name when args are provided', async () => { + const result = await branchCommand.action!(makeCtx(), ' my-branch '); + expect(result).toEqual({ + type: 'dialog', + dialog: 'branch', + name: 'my-branch', + }); + }); + + it('exposes /fork as an alias', () => { + expect(branchCommand.altNames).toContain('fork'); + }); +}); diff --git a/packages/cli/src/ui/commands/branchCommand.ts b/packages/cli/src/ui/commands/branchCommand.ts new file mode 100644 index 00000000000..1d298aa1b73 --- /dev/null +++ b/packages/cli/src/ui/commands/branchCommand.ts @@ -0,0 +1,60 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { SlashCommand, SlashCommandActionReturn } from './types.js'; +import { CommandKind } from './types.js'; +import { t } from '../../i18n/index.js'; + +export const branchCommand: SlashCommand = { + name: 'branch', + altNames: ['fork'], + kind: CommandKind.BUILT_IN, + commandType: 'local', + get description() { + return t('Fork the current conversation into a new session'); + }, + action: async (context, args): Promise => { + const { config } = context.services; + if (!config) { + return { + type: 'message', + messageType: 'error', + content: t('Config is not available.'), + }; + } + + // Guard: streaming or awaiting tool confirmation — forking mid-flight + // would tear the new session's parent chain. + if (context.ui.isIdleRef?.current === false) { + return { + type: 'message', + messageType: 'error', + content: t( + 'Cannot branch while a response or tool call is in progress. Wait for it to finish or resolve the pending tool call.', + ), + }; + } + + // Guard: nothing to fork from. + const sessionService = config.getSessionService(); + const currentId = config.getSessionId(); + const hasRecords = await sessionService.sessionExists(currentId); + if (!hasRecords) { + return { + type: 'message', + messageType: 'error', + content: t('No conversation to branch.'), + }; + } + + const name = args.trim().replace(/[\r\n]+/g, ' '); + return ( + name + ? { type: 'dialog', dialog: 'branch', name } + : { type: 'dialog', dialog: 'branch' } + ) as SlashCommandActionReturn; + }, +}; diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index ee7a1c096fa..7ea453e4b69 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -160,6 +160,9 @@ export interface OpenDialogActionReturn { /** Pre-filtered sessions for the picker (e.g., multiple title matches from /resume ). */ matchedSessions?: SessionListItem[]; + /** Optional session name for /branch — passed through to handleBranch. */ + name?: string; + dialog: | 'help' | 'arena_start' @@ -180,6 +183,7 @@ export interface OpenDialogActionReturn { | 'approval-mode' | 'resume' | 'delete' + | 'branch' | 'extensions_manage' | 'hooks' | 'mcp'; diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index 5aac4e66a2c..0a7157c6bf5 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -101,6 +101,8 @@ export interface UIActions { openResumeDialog: () => void; closeResumeDialog: () => void; handleResume: (sessionId: string) => void; + // Branch (fork) session + handleBranch: (name?: string) => void; // Delete session dialog openDeleteDialog: () => void; closeDeleteDialog: () => void; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 0893c8e2869..e9dcc5e3e57 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -165,6 +165,9 @@ describe('useSlashCommandProcessor', () => { openPermissionsDialog: vi.fn(), openApprovalModeDialog: vi.fn(), openResumeDialog: vi.fn(), + handleResume: vi.fn(), + handleBranch: vi.fn(), + openDeleteDialog: vi.fn(), quit: mockSetQuittingMessages, setDebugMessage: vi.fn(), dispatchExtensionStateUpdate: vi.fn(), @@ -1176,4 +1179,53 @@ describe('useSlashCommandProcessor', () => { ).toBeNull(); }); }); + + describe('SLASH_COMMANDS_SKIP_RECORDING', () => { + // Why these live in the skip set: the fork itself is the side effect + // (new JSONL file with full parent history), so also writing a + // `/branch <name>` slash-command record into the parent session would + // bleed into the fork's tail as a trailing user input — user-visible + // noise with no semantic value. Same rationale for /new, /resume, + // /delete, /clear: session-level commands whose outcome is the new + // session state, not a conversation turn. + it('does not record /branch via the chat recorder', async () => { + const branchCmd = createTestCommand({ + name: 'branch', + action: vi.fn().mockResolvedValue({ type: 'dialog', dialog: 'branch' }), + }); + const result = setupProcessorHook([branchCmd]); + await waitFor(() => expect(result.current.slashCommands).toHaveLength(1)); + + const recorder = mockConfig.getChatRecordingService() as { + recordSlashCommand: ReturnType<typeof vi.fn>; + }; + recorder.recordSlashCommand.mockClear(); + + await act(async () => { + await result.current.handleSlashCommand('/branch my-branch'); + }); + + expect(recorder.recordSlashCommand).not.toHaveBeenCalled(); + }); + + it('still records unrelated commands via the chat recorder (control)', async () => { + const testCmd = createTestCommand({ + name: 'regular', + action: vi.fn().mockResolvedValue(undefined), + }); + const result = setupProcessorHook([testCmd]); + await waitFor(() => expect(result.current.slashCommands).toHaveLength(1)); + + const recorder = mockConfig.getChatRecordingService() as { + recordSlashCommand: ReturnType<typeof vi.fn>; + }; + recorder.recordSlashCommand.mockClear(); + + await act(async () => { + await result.current.handleSlashCommand('/regular'); + }); + + expect(recorder.recordSlashCommand).toHaveBeenCalled(); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 5d9cb0ba0ac..6f446d51e4f 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -74,6 +74,7 @@ const SLASH_COMMANDS_SKIP_RECORDING = new Set([ 'new', 'resume', 'delete', + 'branch', 'btw', ]); @@ -90,6 +91,7 @@ interface SlashCommandProcessorActions { openApprovalModeDialog: () => void; openResumeDialog: (matchedSessions?: SessionListItem[]) => void; handleResume: (sessionId: string) => void; + handleBranch: (name?: string) => void; openDeleteDialog: () => void; quit: (messages: HistoryItem[]) => void; setDebugMessage: (message: string) => void; @@ -573,6 +575,9 @@ export const useSlashCommandProcessor = ( actions.openResumeDialog(result.matchedSessions); } return { type: 'handled' }; + case 'branch': + actions.handleBranch(result.name); + return { type: 'handled' }; case 'delete': actions.openDeleteDialog(); return { type: 'handled' }; diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts new file mode 100644 index 00000000000..9440c3ab939 --- /dev/null +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -0,0 +1,283 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { renderHook, act } from '@testing-library/react'; +import { SessionStartSource } from '@qwen-code/qwen-code-core'; +import { useBranchCommand } from './useBranchCommand.js'; + +describe('useBranchCommand', () => { + let forkSession: ReturnType<typeof vi.fn>; + let loadSession: ReturnType<typeof vi.fn>; + let finalize: ReturnType<typeof vi.fn>; + let startNewSessionConfig: ReturnType<typeof vi.fn>; + let startNewSessionUI: ReturnType<typeof vi.fn>; + let recordCustomTitle: ReturnType<typeof vi.fn>; + let findSessionsByTitle: ReturnType<typeof vi.fn>; + let fireSessionStartEvent: ReturnType<typeof vi.fn>; + let clearItems: ReturnType<typeof vi.fn>; + let loadHistory: ReturnType<typeof vi.fn>; + let setSessionName: ReturnType<typeof vi.fn>; + let remount: ReturnType<typeof vi.fn>; + let addItem: ReturnType<typeof vi.fn>; + // Mock Config shape covers only what useBranchCommand touches. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let config: any; + + const makeOptions = () => ({ + config, + historyManager: { clearItems, loadHistory, addItem }, + startNewSession: startNewSessionUI, + setSessionName, + remount, + }); + + // Helper to build a ChatRecord-shaped user message for loadSession mocks. + // Keeps intent explicit at each call site (genuine user msg vs. synthetic + // subtype vs. non-text) without pulling in the full ChatRecord type here. + const userRecord = (text: string, subtype?: string) => ({ + uuid: 'u' + text.slice(0, 3), + parentUuid: null, + sessionId: 'sid', + type: 'user' as const, + ...(subtype ? { subtype } : {}), + timestamp: 't', + cwd: '/', + version: 'v', + message: { role: 'user', parts: [{ text }] }, + }); + + beforeEach(() => { + forkSession = vi + .fn() + .mockResolvedValue({ filePath: '/tmp/new.jsonl', copiedCount: 2 }); + loadSession = vi.fn().mockResolvedValue({ + conversation: { + messages: [userRecord('help me fix the login bug')], + }, + filePath: '/tmp/new.jsonl', + lastCompletedUuid: 'u2', + }); + finalize = vi.fn(); + recordCustomTitle = vi.fn().mockReturnValue(true); + findSessionsByTitle = vi.fn().mockResolvedValue([]); + fireSessionStartEvent = vi.fn(); + startNewSessionConfig = vi.fn(); + startNewSessionUI = vi.fn(); + clearItems = vi.fn(); + loadHistory = vi.fn(); + setSessionName = vi.fn(); + remount = vi.fn(); + addItem = vi.fn(); + config = { + getSessionId: () => '12345678-aaaa-bbbb-cccc-dddddddddddd', + getSessionService: () => ({ + forkSession, + loadSession, + findSessionsByTitle, + }), + getChatRecordingService: () => ({ finalize, recordCustomTitle }), + getGeminiClient: () => ({ initialize: vi.fn() }), + getHookSystem: () => ({ fireSessionStartEvent }), + startNewSession: startNewSessionConfig, + getModel: () => 'test-model', + getApprovalMode: () => 'default', + getDebugLogger: () => ({ warn: vi.fn() }), + }; + }); + + it('runs finalize → forkSession → loadSession → config.startNewSession in order', async () => { + const order: string[] = []; + finalize.mockImplementation(() => order.push('finalize')); + forkSession.mockImplementation(async () => { + order.push('fork'); + return { filePath: '/tmp/new.jsonl', copiedCount: 2 }; + }); + loadSession.mockImplementation(async () => { + order.push('load'); + return { + conversation: { messages: [] }, + filePath: '/tmp/new.jsonl', + lastCompletedUuid: 'u', + }; + }); + startNewSessionConfig.mockImplementation(() => order.push('config.start')); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(order).toEqual(['finalize', 'fork', 'load', 'config.start']); + }); + + it('records the user-provided name with a (Branch) suffix', async () => { + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + expect(recordCustomTitle).toHaveBeenCalledWith('my-branch (Branch)'); + expect(setSessionName).toHaveBeenCalledWith('my-branch (Branch)'); + }); + + it('bumps to (Branch N) when the default suffix is already taken', async () => { + findSessionsByTitle.mockImplementation(async (title: string) => { + if (title === 'my-branch (Branch)') { + return [{ sessionId: 'other', customTitle: title } as unknown]; + } + return []; + }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + expect(recordCustomTitle).toHaveBeenCalledWith('my-branch (Branch 2)'); + expect(setSessionName).toHaveBeenCalledWith('my-branch (Branch 2)'); + }); + + it('derives the base title from the first user ChatRecord when no name is given', async () => { + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch(); + }); + // deriveFirstPrompt collapses whitespace and truncates to 100 chars; + // "help me fix the login bug" fits, then + " (Branch)" + expect(recordCustomTitle).toHaveBeenCalledWith( + 'help me fix the login bug (Branch)', + ); + }); + + it('falls back to "Branched conversation (Branch)" when the transcript has no user records', async () => { + loadSession.mockResolvedValue({ + conversation: { messages: [] }, + filePath: '/tmp/new.jsonl', + lastCompletedUuid: null, + }); + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch(); + }); + expect(recordCustomTitle).toHaveBeenCalledWith( + 'Branched conversation (Branch)', + ); + }); + + it('skips synthetic user-role records (cron, notification, etc.) and picks the first real prompt', async () => { + loadSession.mockResolvedValue({ + conversation: { + messages: [ + userRecord('scheduled task ran', 'cron'), + userRecord('agent finished X', 'notification'), + userRecord('what does this codebase do'), + ], + }, + filePath: '/tmp/new.jsonl', + lastCompletedUuid: null, + }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch(); + }); + expect(recordCustomTitle).toHaveBeenCalledWith( + 'what does this codebase do (Branch)', + ); + }); + + it('emits the Claude-style success pair naming the branch and the resume hint with the old sessionId', async () => { + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'info', + text: 'Branched conversation "my-branch". You are now in the branch.', + }), + expect.any(Number), + ); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'info', + text: 'To resume the original: /resume 12345678-aaaa-bbbb-cccc-dddddddddddd', + }), + expect.any(Number), + ); + }); + + it('fires SessionStart with SessionStartSource.Branch (not Resume)', async () => { + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + expect(fireSessionStartEvent).toHaveBeenCalledTimes(1); + expect(fireSessionStartEvent).toHaveBeenCalledWith( + SessionStartSource.Branch, + expect.any(String), + expect.any(String), + ); + }); + + it('omits the quoted-title fragment when no name is provided', async () => { + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch(); + }); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'info', + text: 'Branched conversation. You are now in the branch.', + }), + expect.any(Number), + ); + }); + + it('surfaces an error item and does not switch sessions when forkSession throws', async () => { + forkSession.mockRejectedValue(new Error('disk full')); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('x'); + }); + + expect(startNewSessionConfig).not.toHaveBeenCalled(); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringMatching(/Failed to branch conversation.*disk full/), + }), + expect.any(Number), + ); + }); + + it('does not clear or swap the UI when core startNewSession throws post-fork', async () => { + // Guards the "swap core first" invariant: if core swap fails after the + // disk fork succeeds, the UI must stay on the parent — no cleared + // history, no new UI sessionId — so the user is not stranded. + startNewSessionConfig.mockImplementation(() => { + throw new Error('core boom'); + }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('x'); + }); + + expect(forkSession).toHaveBeenCalledTimes(1); + expect(clearItems).not.toHaveBeenCalled(); + expect(loadHistory).not.toHaveBeenCalled(); + expect(startNewSessionUI).not.toHaveBeenCalled(); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringMatching(/Failed to branch conversation.*core boom/), + }), + expect.any(Number), + ); + }); +}); diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts new file mode 100644 index 00000000000..c0a6abbdfa0 --- /dev/null +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -0,0 +1,230 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useCallback } from 'react'; +import { randomUUID } from 'node:crypto'; +import { + type Config, + type SessionService, + type ChatRecord, + SessionStartSource, + type PermissionMode, +} from '@qwen-code/qwen-code-core'; +import { buildResumedHistoryItems } from '../utils/resumeHistoryUtils.js'; +import type { UseHistoryManagerReturn } from './useHistoryManager.js'; +import { t } from '../../i18n/index.js'; + +/** + * Cap for the `(Branch N)` collision scan. Each probe is a project-wide + * scan via `findSessionsByTitle`; 99 is generous for realistic use and + * bounds the worst case. + */ +const MAX_BRANCH_COLLISION_SCAN = 99; + +/** + * Derives a short one-line title from the first *real* user message in the + * transcript. Mirrors Claude Code's `deriveFirstPrompt` (see + * claude-code/src/commands/branch/branch.ts): collapse whitespace, truncate + * to 100 chars, fall back to "Branched conversation" when the transcript + * has no user text. + * + * Reads ChatRecord[] — the JSONL-level transcript — NOT the Gemini API + * `Content[]` history. The latter is prepended with environment / CLAUDE.md / + * context injections by the runtime; its first role=user entry is a + * synthetic bootstrap message, not anything the user typed. + * + * Records with a `subtype` are skipped — those are cron-fired prompts, + * notifications, slash-command echoes, etc., not genuine user input. + */ +function deriveFirstPrompt(messages: ChatRecord[]): string { + for (const record of messages) { + if (record.type !== 'user') continue; + if (record.subtype) continue; + const parts = record.message?.parts; + if (!parts) continue; + for (const part of parts) { + if ('text' in part && typeof part.text === 'string' && part.text) { + const collapsed = part.text.replace(/\s+/g, ' ').trim().slice(0, 100); + if (collapsed) return collapsed; + } + } + } + return 'Branched conversation'; +} + +/** + * Appends ` (Branch)` to `baseName`, bumping to ` (Branch 2)`, ` (Branch 3)`, + * ... when the exact name is already taken by another session's customTitle + * in the current project. Mirrors Claude's `getUniqueForkName`. + */ +async function computeUniqueBranchTitle( + baseName: string, + sessionService: SessionService, +): Promise<string> { + const trimmed = baseName.trim(); + const first = `${trimmed} (Branch)`; + if ((await sessionService.findSessionsByTitle(first)).length === 0) { + return first; + } + for (let n = 2; n <= MAX_BRANCH_COLLISION_SCAN; n++) { + const candidate = `${trimmed} (Branch ${n})`; + if ((await sessionService.findSessionsByTitle(candidate)).length === 0) { + return candidate; + } + } + // Pathological density — timestamp fallback keeps the fork unique. + return `${trimmed} (Branch ${Date.now()})`; +} + +export interface UseBranchCommandOptions { + config: Config | null; + historyManager: Pick< + UseHistoryManagerReturn, + 'clearItems' | 'loadHistory' | 'addItem' + >; + startNewSession: (sessionId: string) => void; + setSessionName?: (name: string | null) => void; + remount?: () => void; +} + +export interface UseBranchCommandResult { + handleBranch: (name?: string) => Promise<void>; +} + +/** + * Orchestrates `/branch`: + * 1. Capture the current (soon-to-be-parent) sessionId for the resume hint. + * 2. Finalize the outgoing ChatRecordingService so the last metadata is on disk. + * 3. Call `SessionService.forkSession` to write a new JSONL under a new id. + * 4. Load the fork back via `loadSession` and switch the UI + core config. + * 5. Compute the customTitle — user-provided name OR `deriveFirstPrompt` — + * always suffixed with ` (Branch)` (bumping to `(Branch N)` on collision). + * 6. Fire the SessionStart hook. + * 7. Announce the fork with Claude-style two-line info item: + * `Branched conversation "foo". You are now in the branch.` + * `To resume the original: /resume <oldSessionId>` + * + * Mirrors claude-code/src/commands/branch/branch.ts. + */ +export function useBranchCommand( + options: UseBranchCommandOptions, +): UseBranchCommandResult { + const { config, historyManager, startNewSession, setSessionName, remount } = + options; + + const handleBranch = useCallback( + async (name?: string) => { + if (!config) return; + + const oldSessionId = config.getSessionId(); + const newSessionId = randomUUID(); + const sessionService = config.getSessionService(); + + try { + // 1. Flush outgoing recorder. + try { + config.getChatRecordingService()?.finalize(); + } catch { + // best-effort + } + + // 2. Fork the JSONL on disk. + await sessionService.forkSession(oldSessionId, newSessionId); + + // 3. Load the new file. + const resumed = await sessionService.loadSession(newSessionId); + if (!resumed) { + throw new Error('Failed to load newly forked session'); + } + + // 4. Swap core first. Anything that can still fail (startNewSession, + // client init) runs while the UI is still showing the parent + // session, so a throw leaves the user safely on the parent + // instead of stranded with a cleared history and a half-live + // client. + config.startNewSession(newSessionId, resumed); + await config.getGeminiClient()?.initialize?.(); + + // 5. Swap UI. + const uiHistoryItems = buildResumedHistoryItems(resumed, config); + startNewSession(newSessionId); + historyManager.clearItems(); + historyManager.loadHistory(uiHistoryItems); + + // 6. Compute and apply the branch customTitle. + // The forked transcript is identical to the parent's, so reading + // the first real user message from `resumed.conversation.messages` + // mirrors Claude's "use the first parent message" behavior. + const baseName = + name ?? deriveFirstPrompt(resumed.conversation.messages); + const effectiveTitle = await computeUniqueBranchTitle( + baseName, + sessionService, + ); + config.getChatRecordingService()?.recordCustomTitle(effectiveTitle); + setSessionName?.(effectiveTitle); + + // 7. Fire SessionStart for the new session. A fork is semantically + // distinct from a resume — the sessionId is new and the transcript + // is a derivative — so we use the dedicated `Branch` source value + // to let hook consumers distinguish the two. + try { + await config + .getHookSystem() + ?.fireSessionStartEvent( + SessionStartSource.Branch, + config.getModel() ?? '', + String(config.getApprovalMode()) as PermissionMode, + ); + } catch (err) { + config.getDebugLogger().warn(`SessionStart hook failed: ${err}`); + } + + // 8. Refresh terminal UI. + remount?.(); + + // 9. Announce. Two history items mirror Claude's success message + // (branched line + resume hint). The quoted name is the raw + // user-provided `name`; no `(Branch)` suffix — that decoration + // belongs in the picker/prompt bar, not in the user-facing + // announcement. + const titleInfo = name ? ` "${name}"` : ''; + historyManager.addItem( + { + type: 'info', + text: t( + 'Branched conversation{{titleInfo}}. You are now in the branch.', + { titleInfo }, + ), + }, + Date.now(), + ); + historyManager.addItem( + { + type: 'info', + text: t('To resume the original: /resume {{sessionId}}', { + sessionId: oldSessionId, + }), + }, + Date.now(), + ); + } catch (err) { + historyManager.addItem( + { + type: 'error', + text: t('Failed to branch conversation: {{message}}', { + message: err instanceof Error ? err.message : String(err), + }), + }, + Date.now(), + ); + } + }, + [config, historyManager, startNewSession, setSessionName, remount], + ); + + return { handleBranch }; +} diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index 7b74b6cee89..f3715f1b750 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -688,6 +688,7 @@ export enum SessionStartSource { Resume = 'resume', Clear = 'clear', Compact = 'compact', + Branch = 'branch', } export enum PermissionMode { diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 0565c7ad2e8..d7595e57682 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -103,6 +103,24 @@ export interface ChatRecord { | AtCommandRecordPayload | CustomTitleRecordPayload | NotificationRecordPayload; + + /** + * Set on every record of a forked session to record its lineage. + * `sessionId` is the parent (source) session id; `messageUuid` is the + * uuid of the equivalent message in the parent — the same value as + * this record's `uuid`, since /branch copies each message verbatim + * except for rewriting `sessionId` and rebuilding `parentUuid` by + * write order. + * + * Written by /branch on every copied record; never consumed by any + * feature at read time — it exists purely as per-message audit trail + * so that when a record is inspected in isolation its origin is + * self-contained (mirrors Claude Code's /branch behavior). + */ + forkedFrom?: { + sessionId: string; + messageUuid: string; + }; } export interface NotificationRecordPayload { diff --git a/packages/core/src/services/sessionService.test.ts b/packages/core/src/services/sessionService.test.ts index 77f04012e0c..6f32b50b713 100644 --- a/packages/core/src/services/sessionService.test.ts +++ b/packages/core/src/services/sessionService.test.ts @@ -784,4 +784,211 @@ describe('SessionService', () => { ]); }); }); + + describe('forkSession', () => { + // forkSession uses real disk I/O through `jsonl.read` and `fs.*`. + // The outer describe hoist-mocks `node:path`, `../utils/paths.js`, and + // `../utils/jsonl-utils.js`; restore the real implementations inside this + // describe's setup so the fork actually reads/writes tmp files. + let realTmpDir: string; + let realOs: typeof import('node:os'); + let realPath: typeof import('node:path'); + let service: SessionService; + let cwd: string; + + beforeEach(async () => { + realOs = await import('node:os'); + realPath = await vi.importActual<typeof import('node:path')>('node:path'); + const actualPaths = + await vi.importActual<typeof import('../utils/paths.js')>( + '../utils/paths.js', + ); + const actualJsonl = await vi.importActual< + typeof import('../utils/jsonl-utils.js') + >('../utils/jsonl-utils.js'); + + vi.mocked(path.join).mockImplementation( + realPath.join as unknown as typeof path.join, + ); + vi.mocked(path.dirname).mockImplementation( + realPath.dirname as unknown as typeof path.dirname, + ); + // Storage.resolveRuntimeBaseDir uses isAbsolute and resolve; both are + // auto-mocked to return undefined, which silently falls back to + // `~/.qwen` and makes the fork write outside the tmp sandbox. + vi.mocked(path.isAbsolute).mockImplementation( + realPath.isAbsolute as unknown as typeof path.isAbsolute, + ); + vi.mocked(path.resolve).mockImplementation( + realPath.resolve as unknown as typeof path.resolve, + ); + vi.mocked(getProjectHash).mockImplementation(actualPaths.getProjectHash); + // Storage.getProjectDir calls sanitizeCwd via a non-spied namespace import; + // restore it module-globally so getChatsDir() returns a real path. + const mockedPaths = (await import('../utils/paths.js')) as unknown as { + sanitizeCwd: (cwd: string) => string; + }; + mockedPaths.sanitizeCwd = actualPaths.sanitizeCwd; + vi.mocked(jsonl.read).mockImplementation(actualJsonl.read); + vi.mocked(jsonl.readLines).mockImplementation(actualJsonl.readLines); + + // Restore any fs spies installed by the outer beforeEach. + vi.mocked(readdirSyncSpy).mockRestore?.(); + vi.mocked(statSyncSpy).mockRestore?.(); + vi.mocked(unlinkSyncSpy).mockRestore?.(); + + realTmpDir = fs.mkdtempSync( + realPath.join(realOs.tmpdir(), 'fork-session-'), + ); + process.env['QWEN_RUNTIME_DIR'] = realTmpDir; + cwd = process.cwd(); + service = new SessionService(cwd); + }); + + afterEach(() => { + delete process.env['QWEN_RUNTIME_DIR']; + try { + fs.rmSync(realTmpDir, { recursive: true, force: true }); + } catch { + // best-effort + } + }); + + const seedSession = (sessionId: string) => { + const chatsDir = realPath.join( + service['storage'].getProjectDir(), + 'chats', + ); + fs.mkdirSync(chatsDir, { recursive: true }); + const file = realPath.join(chatsDir, `${sessionId}.jsonl`); + const lines = [ + { + uuid: 'u1', + parentUuid: null, + sessionId, + type: 'user', + timestamp: '2026-04-22T00:00:00.000Z', + cwd, + version: 'test', + message: { role: 'user', parts: [{ text: 'hello' }] }, + }, + { + uuid: 'u2', + parentUuid: 'u1', + sessionId, + type: 'assistant', + timestamp: '2026-04-22T00:00:01.000Z', + cwd, + version: 'test', + message: { role: 'model', parts: [{ text: 'hi' }] }, + }, + ]; + fs.writeFileSync( + file, + lines.map((l) => JSON.stringify(l)).join('\n') + '\n', + ); + return { file, lines }; + }; + + it('rewrites sessionId, rebuilds parentUuid, and stamps forkedFrom on every record', async () => { + const oldId = '11111111-1111-1111-1111-111111111111'; + const newId = '22222222-2222-2222-2222-222222222222'; + const { file: srcPath } = seedSession(oldId); + + const result = await service.forkSession(oldId, newId); + expect(result.copiedCount).toBe(2); + expect(result.filePath).toContain(`${newId}.jsonl`); + + const written = fs + .readFileSync(result.filePath, 'utf8') + .trim() + .split('\n') + .map((l) => JSON.parse(l)); + + expect(written).toHaveLength(2); + expect(written[0]).toMatchObject({ + uuid: 'u1', + parentUuid: null, + sessionId: newId, + forkedFrom: { sessionId: oldId, messageUuid: 'u1' }, + }); + expect(written[1]).toMatchObject({ + uuid: 'u2', + parentUuid: 'u1', // rebuilt in write order + sessionId: newId, + forkedFrom: { sessionId: oldId, messageUuid: 'u2' }, + }); + // Source file is untouched. + expect(fs.existsSync(srcPath)).toBe(true); + const srcLines = fs + .readFileSync(srcPath, 'utf8') + .trim() + .split('\n') + .map((l) => JSON.parse(l)); + expect(srcLines.every((r) => r.sessionId === oldId)).toBe(true); + expect(srcLines.every((r) => !r.forkedFrom)).toBe(true); + }); + + it('throws when the source session does not exist', async () => { + const oldId = '33333333-3333-3333-3333-333333333333'; + const newId = '44444444-4444-4444-4444-444444444444'; + await expect(service.forkSession(oldId, newId)).rejects.toThrow(); + }); + + it('throws when the target session file already exists', async () => { + const oldId = '55555555-5555-5555-5555-555555555555'; + const newId = '66666666-6666-6666-6666-666666666666'; + seedSession(oldId); + const chatsDir = realPath.join( + service['storage'].getProjectDir(), + 'chats', + ); + fs.writeFileSync(realPath.join(chatsDir, `${newId}.jsonl`), 'x'); + + await expect(service.forkSession(oldId, newId)).rejects.toThrow( + /already exists/, + ); + }); + + it('throws when the source session belongs to a different project', async () => { + // Defensive guard: a file can physically sit in this project's chats + // dir but carry a record whose cwd hashes to a different project + // (manual file move, corrupted state). Fork must refuse rather than + // silently cross project boundaries. + const oldId = '77777777-7777-7777-7777-777777777777'; + const newId = '88888888-8888-8888-8888-888888888888'; + const chatsDir = realPath.join( + service['storage'].getProjectDir(), + 'chats', + ); + fs.mkdirSync(chatsDir, { recursive: true }); + fs.writeFileSync( + realPath.join(chatsDir, `${oldId}.jsonl`), + JSON.stringify({ + uuid: 'u1', + parentUuid: null, + sessionId: oldId, + type: 'user', + timestamp: '2026-04-22T00:00:00.000Z', + cwd: '/some/other/project', + version: 'test', + message: { role: 'user', parts: [{ text: 'hi' }] }, + }) + '\n', + ); + + await expect(service.forkSession(oldId, newId)).rejects.toThrow( + /does not belong to current project/, + ); + }); + + it('rejects invalid sessionId patterns before touching disk', async () => { + const valid = '99999999-9999-9999-9999-999999999999'; + await expect(service.forkSession('bogus', valid)).rejects.toThrow( + /Invalid source sessionId/, + ); + await expect(service.forkSession(valid, 'bogus')).rejects.toThrow( + /Invalid new sessionId/, + ); + }); + }); }); diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index 19b7af1f5a4..4640e53e9ca 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -638,6 +638,92 @@ export class SessionService { } } + /** + * Forks a session to a new sessionId. + * + * Reads the source JSONL into memory, rewrites every record's `sessionId` + * to `newSessionId`, rebuilds the `parentUuid` chain in write order so the + * fork is a linear continuation, stamps `forkedFrom: { sessionId, messageUuid }` + * on every copied record for audit, and writes the result to `<newId>.jsonl`. + * + * Mirrors Claude Code's `/branch` storage model: full in-memory copy + per- + * message forkedFrom (see claude-code/src/commands/branch/branch.ts). + * + * The source file is not modified. + * + * @throws If source does not exist, source is empty, source belongs to a + * different project, or the target file already exists. + */ + async forkSession( + sourceSessionId: string, + newSessionId: string, + ): Promise<{ filePath: string; copiedCount: number }> { + if (!SESSION_FILE_PATTERN.test(`${sourceSessionId}.jsonl`)) { + throw new Error(`Invalid source sessionId: ${sourceSessionId}`); + } + if (!SESSION_FILE_PATTERN.test(`${newSessionId}.jsonl`)) { + throw new Error(`Invalid new sessionId: ${newSessionId}`); + } + + const chatsDir = this.getChatsDir(); + const sourcePath = path.join(chatsDir, `${sourceSessionId}.jsonl`); + const targetPath = path.join(chatsDir, `${newSessionId}.jsonl`); + + // Read + parse the full source transcript. + const records = await jsonl.read<ChatRecord>(sourcePath); + if (records.length === 0) { + throw new Error(`Source session not found or empty: ${sourceSessionId}`); + } + + // Verify project ownership via the first record's cwd. + if (getProjectHash(records[0].cwd) !== this.projectHash) { + throw new Error( + `Source session does not belong to current project: ${sourceSessionId}`, + ); + } + + // Rebuild the parentUuid chain in write order so the fork is a clean + // linear descendant. `forkedFrom` captures the origin of each message. + let prevUuid: string | null = null; + const forked: ChatRecord[] = records.map((record) => { + const next: ChatRecord = { + ...record, + sessionId: newSessionId, + parentUuid: prevUuid, + forkedFrom: { + sessionId: sourceSessionId, + messageUuid: record.uuid, + }, + }; + prevUuid = record.uuid; + return next; + }); + + fs.mkdirSync(chatsDir, { recursive: true }); + const body = forked.map((r) => JSON.stringify(r)).join('\n') + '\n'; + + // Exclusive create: one syscall that both asserts "file doesn't exist" + // and opens for writing, eliminating the TOCTOU window between a + // separate existsSync check and writeFileSync. Also guarantees we + // never silently overwrite an existing session file. + let fd: number; + try { + fd = fs.openSync(targetPath, 'wx', 0o600); + } catch (err) { + if ((err as NodeJS.ErrnoException).code === 'EEXIST') { + throw new Error(`Target session file already exists: ${newSessionId}`); + } + throw err; + } + try { + fs.writeFileSync(fd, body, { encoding: 'utf8' }); + } finally { + fs.closeSync(fd); + } + + return { filePath: targetPath, copiedCount: forked.length }; + } + /** * Gets the custom title for a session by reading from its JSONL file. * From 5467afac8d6e8d7b780b4a69b596e5f8edaaaa8b Mon Sep 17 00:00:00 2001 From: qqqys <qys177@gmail.com> Date: Thu, 23 Apr 2026 09:59:46 +0800 Subject: [PATCH 2/6] fix(session): drop stale `commandType` field from branchCommand MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `commandType: 'local'` field was added referencing the Phase 1 slash-command redesign draft, but the field never made it onto `SlashCommand` — Phase 1 landed with `supportedModes` / `userInvocable` instead. After merging main, strict tsc rejects the unknown property with TS2353 and the CLI package fails to build. 🤖 Generated with [Qwen Code](https://github.com/QwenLM/qwen-code) Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com> --- packages/cli/src/ui/commands/branchCommand.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/cli/src/ui/commands/branchCommand.ts b/packages/cli/src/ui/commands/branchCommand.ts index 1d298aa1b73..e5ca498814f 100644 --- a/packages/cli/src/ui/commands/branchCommand.ts +++ b/packages/cli/src/ui/commands/branchCommand.ts @@ -12,7 +12,6 @@ export const branchCommand: SlashCommand = { name: 'branch', altNames: ['fork'], kind: CommandKind.BUILT_IN, - commandType: 'local', get description() { return t('Fork the current conversation into a new session'); }, From e526c5dc794173ab91c8e9f08a2184de13631582 Mon Sep 17 00:00:00 2001 From: qqqys <qys177@gmail.com> Date: Fri, 24 Apr 2026 17:01:01 +0800 Subject: [PATCH 3/6] fix(session): roll core back to parent when /branch post-fork init throws MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `useBranchCommand` swapped core onto the fork via `config.startNewSession` before `getGeminiClient().initialize()` resolved. If init rejected, the catch only surfaced an error item: UI was still on the parent, but `sessionId` + `ChatRecordingService` were already pointing at the orphan fork JSONL, so the next user message would silently record into the fork while appearing to belong to the parent conversation. Snapshot the parent session's `ResumedSessionData` up front, gate the rollback on a `coreSwapped` flag, and in the catch run `startNewSession(oldSessionId, prevSessionData)` + re-`initialize()` so sessionId, recorder (with the correct parentUuid chain tail), and chat history all return to the parent. Rollback re-init is best-effort — if it throws again we log and still surface the original failure, since sessionId + recorder are the load-bearing invariant. Regression tests: (1) initialize rejects after swap → two `startNewSessionConfig` calls (fork then rollback-with-parent-data), two `initialize` calls, no UI swap, original error surfaced; (2) rollback's own init also rejects → sessionId still lands on parent, debug logger warns, original error still surfaced. Reported by gpt-5.5 via Qwen Code `/review` on #3539. Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com> --- .../cli/src/ui/hooks/useBranchCommand.test.ts | 120 +++++++++++++++++- packages/cli/src/ui/hooks/useBranchCommand.ts | 41 +++++- 2 files changed, 158 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index 9440c3ab939..0037d5e6539 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -89,7 +89,10 @@ describe('useBranchCommand', () => { }; }); - it('runs finalize → forkSession → loadSession → config.startNewSession in order', async () => { + it('runs snapshot → finalize → forkSession → loadSession → config.startNewSession in order', async () => { + // The leading snapshot load (of the parent session) exists so the catch + // block can roll core back to the parent with the correct parentUuid + // chain tail if getGeminiClient().initialize() rejects after the swap. const order: string[] = []; finalize.mockImplementation(() => order.push('finalize')); forkSession.mockImplementation(async () => { @@ -111,7 +114,13 @@ describe('useBranchCommand', () => { await result.current.handleBranch('my-branch'); }); - expect(order).toEqual(['finalize', 'fork', 'load', 'config.start']); + expect(order).toEqual([ + 'load', // parent snapshot for rollback + 'finalize', + 'fork', + 'load', // forked session + 'config.start', + ]); }); it('records the user-provided name with a (Branch) suffix', async () => { @@ -255,6 +264,113 @@ describe('useBranchCommand', () => { ); }); + it('rolls core back to the parent session when getGeminiClient().initialize() rejects after swap', async () => { + // The reviewer's scenario: config.startNewSession succeeds (core is now + // on the fork), but then getGeminiClient().initialize() rejects. Without + // rollback, core stays on the fork while UI is still on the parent, so + // the recorder silently writes subsequent user input into an orphan + // JSONL. This test pins the rollback invariant — after the failure core + // must be back on the parent sessionId with the parent's ResumedSessionData. + const oldSessionId = '12345678-aaaa-bbbb-cccc-dddddddddddd'; + const parentResumed = { + conversation: { messages: [userRecord('parent msg')] }, + filePath: `/tmp/${oldSessionId}.jsonl`, + lastCompletedUuid: 'uparent', + }; + const forkResumed = { + conversation: { messages: [userRecord('parent msg')] }, + filePath: '/tmp/new.jsonl', + lastCompletedUuid: 'uparent', + }; + // Called twice: once up front to snapshot the parent for rollback, + // once after forkSession to load the fork. + loadSession.mockImplementation(async (sid: string) => + sid === oldSessionId ? parentResumed : forkResumed, + ); + + const initialize = vi + .fn() + .mockRejectedValueOnce(new Error('init boom')) // fork init fails + .mockResolvedValueOnce(undefined); // rollback re-init succeeds + config.getGeminiClient = () => ({ initialize }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('x'); + }); + + // Core was swapped to the fork, then rolled back to the parent. + expect(startNewSessionConfig).toHaveBeenNthCalledWith( + 1, + expect.not.stringMatching(oldSessionId), + forkResumed, + ); + expect(startNewSessionConfig).toHaveBeenNthCalledWith( + 2, + oldSessionId, + parentResumed, + ); + // Client was re-initialized after rollback so chat history re-hydrates + // against the parent session. + expect(initialize).toHaveBeenCalledTimes(2); + // UI never switched — no cleared history, no UI sessionId swap. + expect(clearItems).not.toHaveBeenCalled(); + expect(loadHistory).not.toHaveBeenCalled(); + expect(startNewSessionUI).not.toHaveBeenCalled(); + expect(setSessionName).not.toHaveBeenCalled(); + // User sees the failure. + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringMatching(/Failed to branch conversation.*init boom/), + }), + expect.any(Number), + ); + }); + + it('still surfaces the error and leaves core on the parent when rollback re-init also throws', async () => { + // If the rollback initialize() itself rejects, the swap of sessionId + + // recorder has still happened — that is the load-bearing invariant — + // so we just log and surface the original failure without crashing. + const oldSessionId = '12345678-aaaa-bbbb-cccc-dddddddddddd'; + loadSession.mockResolvedValue({ + conversation: { messages: [userRecord('parent msg')] }, + filePath: '/tmp/new.jsonl', + lastCompletedUuid: 'u2', + }); + const debugWarn = vi.fn(); + config.getDebugLogger = () => ({ warn: debugWarn }); + + const initialize = vi + .fn() + .mockRejectedValueOnce(new Error('init boom')) + .mockRejectedValueOnce(new Error('rollback boom')); + config.getGeminiClient = () => ({ initialize }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('x'); + }); + + // Core was still rolled back to the parent sessionId. + expect(startNewSessionConfig).toHaveBeenNthCalledWith( + 2, + oldSessionId, + expect.any(Object), + ); + expect(debugWarn).toHaveBeenCalledWith( + expect.stringContaining('Rollback after failed /branch init failed'), + ); + // Original failure is what the user sees, not the rollback failure. + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringMatching(/Failed to branch conversation.*init boom/), + }), + expect.any(Number), + ); + }); + it('does not clear or swap the UI when core startNewSession throws post-fork', async () => { // Guards the "swap core first" invariant: if core swap fails after the // disk fork succeeds, the UI must stay on the parent — no cleared diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index c0a6abbdfa0..1e93fb6d0b1 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -10,6 +10,7 @@ import { type Config, type SessionService, type ChatRecord, + type ResumedSessionData, SessionStartSource, type PermissionMode, } from '@qwen-code/qwen-code-core'; @@ -123,6 +124,21 @@ export function useBranchCommand( const newSessionId = randomUUID(); const sessionService = config.getSessionService(); + // Snapshot the parent JSONL state up front so a post-swap failure + // (see the catch block) can faithfully restore sessionId + recorder + // with the correct parentUuid chain tail. `/branch` is guarded on + // `isIdleRef`, so the file isn't being mutated concurrently. + let prevSessionData: ResumedSessionData | undefined; + try { + prevSessionData = await sessionService.loadSession(oldSessionId); + } catch { + // Best-effort snapshot. Falling back to undefined still rolls + // back sessionId + recorder, which is the load-bearing invariant; + // we just lose the parentUuid chain on the restored recorder. + } + + let coreSwapped = false; + try { // 1. Flush outgoing recorder. try { @@ -144,8 +160,12 @@ export function useBranchCommand( // client init) runs while the UI is still showing the parent // session, so a throw leaves the user safely on the parent // instead of stranded with a cleared history and a half-live - // client. + // client. `coreSwapped` gates the rollback path in the catch + // block below — without it, a failure between swap and UI + // update would leave core on the fork while UI still shows + // the parent, silently recording user input into an orphan. config.startNewSession(newSessionId, resumed); + coreSwapped = true; await config.getGeminiClient()?.initialize?.(); // 5. Swap UI. @@ -212,6 +232,25 @@ export function useBranchCommand( Date.now(), ); } catch (err) { + if (coreSwapped) { + // Core already switched to the fork before the failure — put it + // back on the parent, otherwise the recorder would keep writing + // new user messages into the orphan fork JSONL while UI still + // shows the parent. + try { + config.startNewSession(oldSessionId, prevSessionData); + // Re-hydrate chat history against the restored session. Best- + // effort: if this throws too, sessionId + recorder are still + // back on the parent, which is the load-bearing invariant. + await config.getGeminiClient()?.initialize?.(); + } catch (rollbackErr) { + config + .getDebugLogger() + .warn( + `Rollback after failed /branch init failed: ${rollbackErr}`, + ); + } + } historyManager.addItem( { type: 'error', From 8ac4af285daa763f811d9746516089196b409e6b Mon Sep 17 00:00:00 2001 From: qqqys <qys177@gmail.com> Date: Mon, 27 Apr 2026 19:11:21 +0800 Subject: [PATCH 4/6] fix(session): close /branch transactional swap holes flagged in review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three related correctness issues in the /branch core+UI swap, all reported by gpt-5.5 via Qwen Code /review on PR #3539: 1. Snapshot-before-finalize. ChatRecordingService.finalize() appends a trailing `system/custom_title` record that advances `lastRecordUuid`. Loading the parent ResumedSessionData snapshot before that ran captured a stale `lastCompletedUuid`; on rollback the restored recorder would chain its next record's parentUuid to a record that's no longer the JSONL tail, orphaning the custom_title from the parent chain. Move the snapshot to AFTER finalize(). 2. Reverse split-brain after UI swap. The catch block was gated solely on `coreSwapped`, so any failure AFTER the UI commits to the branch (recordCustomTitle, hook fire, remount, announcement render) would roll core back to the parent — leaving UI on the branch while the recorder writes new prompts into the parent JSONL. Track `uiSwapped` separately and skip the rollback once UI is committed; surface the failure as an error item without unwinding the swap. Pinned by a new regression test. 3. Slash dispatcher dropped the handleBranch promise. The `branch` case in slashCommandProcessor returned `{type: 'handled'}` while handleBranch was still in flight, so a fast follow-up prompt could interleave with the swap and be recorded against the wrong session. Await it and tighten the action type from `=> void` to `=> Promise<void>` (both in SlashCommandProcessorActions and UIActionsContext) so this cannot silently regress. Tests: vitest packages/cli/src/ui/hooks/useBranchCommand.test.ts 15 ✓ vitest packages/cli/src/ui/hooks/slashCommandProcessor.test.ts 41 ✓ vitest packages/cli/src/ui/commands/branchCommand.test.ts 6 ✓ vitest packages/core/src/services/sessionService.test.ts 32 ✓ tsc --noEmit clean eslint clean Co-Authored-By: Qwen-Coder <noreply@alibabacloud.com> --- .../cli/src/ui/contexts/UIActionsContext.tsx | 2 +- .../cli/src/ui/hooks/slashCommandProcessor.ts | 10 ++- .../cli/src/ui/hooks/useBranchCommand.test.ts | 57 +++++++++++++++-- packages/cli/src/ui/hooks/useBranchCommand.ts | 64 ++++++++++++------- 4 files changed, 101 insertions(+), 32 deletions(-) diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index 0a7157c6bf5..c89f77c9b85 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -102,7 +102,7 @@ export interface UIActions { closeResumeDialog: () => void; handleResume: (sessionId: string) => void; // Branch (fork) session - handleBranch: (name?: string) => void; + handleBranch: (name?: string) => Promise<void>; // Delete session dialog openDeleteDialog: () => void; closeDeleteDialog: () => void; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 413a407819b..642530af16d 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -92,7 +92,7 @@ interface SlashCommandProcessorActions { openApprovalModeDialog: () => void; openResumeDialog: (matchedSessions?: SessionListItem[]) => void; handleResume: (sessionId: string) => void; - handleBranch: (name?: string) => void; + handleBranch: (name?: string) => Promise<void>; openDeleteDialog: () => void; quit: (messages: HistoryItem[]) => void; setDebugMessage: (message: string) => void; @@ -625,7 +625,13 @@ export const useSlashCommandProcessor = ( } return { type: 'handled' }; case 'branch': - actions.handleBranch(result.name); + // Must be awaited: `/branch` swaps core + UI session + // state asynchronously, and a non-awaited call lets + // this dispatcher return `handled` while the swap is + // still in flight. A fast follow-up prompt could then + // interleave with the swap and be recorded against + // the wrong session. + await actions.handleBranch(result.name); return { type: 'handled' }; case 'delete': actions.openDeleteDialog(); diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index 0037d5e6539..fcc1ad58ab5 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -89,10 +89,13 @@ describe('useBranchCommand', () => { }; }); - it('runs snapshot → finalize → forkSession → loadSession → config.startNewSession in order', async () => { - // The leading snapshot load (of the parent session) exists so the catch - // block can roll core back to the parent with the correct parentUuid - // chain tail if getGeminiClient().initialize() rejects after the swap. + it('runs finalize → snapshot → forkSession → loadSession → config.startNewSession in order', async () => { + // The parent snapshot must come AFTER finalize(): finalize() appends a + // trailing custom_title record to the parent JSONL, advancing the + // recorder's lastCompletedUuid. A snapshot taken before that captures + // a stale tail; on rollback the restored recorder would chain its next + // record's parentUuid to a record that's no longer the JSONL tail, + // orphaning the custom_title record from the parent chain. const order: string[] = []; finalize.mockImplementation(() => order.push('finalize')); forkSession.mockImplementation(async () => { @@ -115,8 +118,8 @@ describe('useBranchCommand', () => { }); expect(order).toEqual([ - 'load', // parent snapshot for rollback 'finalize', + 'load', // parent snapshot for rollback (after finalize so it captures the custom_title append) 'fork', 'load', // forked session 'config.start', @@ -371,6 +374,50 @@ describe('useBranchCommand', () => { ); }); + it('does not roll core back to parent when a post-UI-swap step throws', async () => { + // The reviewer's reverse split-brain: once the UI commits to the branch, + // any subsequent failure (recordCustomTitle, hook fire, remount, + // announcement render) must NOT trigger the catch block's core rollback. + // If it did, the user would see the branch UI but every new prompt + // would be recorded into the parent's JSONL. + // + // Pin the invariant by making remount() — which runs after the UI swap — + // throw, then assert: only ONE config.startNewSession call (to the + // branch), no second call resetting it back to the parent. + const oldSessionId = '12345678-aaaa-bbbb-cccc-dddddddddddd'; + remount.mockImplementation(() => { + throw new Error('remount boom'); + }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('x'); + }); + + // UI did swap. + expect(startNewSessionUI).toHaveBeenCalledTimes(1); + expect(clearItems).toHaveBeenCalled(); + expect(loadHistory).toHaveBeenCalled(); + // Core did NOT roll back to the parent — only the initial swap to + // the branch. A second call with `oldSessionId` would mean the catch + // block reverted core while UI stayed on the branch. + expect(startNewSessionConfig).toHaveBeenCalledTimes(1); + expect(startNewSessionConfig).not.toHaveBeenCalledWith( + oldSessionId, + expect.anything(), + ); + // The user still sees the failure surfaced as an error item. + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringMatching( + /Failed to branch conversation.*remount boom/, + ), + }), + expect.any(Number), + ); + }); + it('does not clear or swap the UI when core startNewSession throws post-fork', async () => { // Guards the "swap core first" invariant: if core swap fails after the // disk fork succeeds, the UI must stay on the parent — no cleared diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index 1e93fb6d0b1..0468e03b652 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -124,39 +124,43 @@ export function useBranchCommand( const newSessionId = randomUUID(); const sessionService = config.getSessionService(); - // Snapshot the parent JSONL state up front so a post-swap failure - // (see the catch block) can faithfully restore sessionId + recorder - // with the correct parentUuid chain tail. `/branch` is guarded on - // `isIdleRef`, so the file isn't being mutated concurrently. - let prevSessionData: ResumedSessionData | undefined; - try { - prevSessionData = await sessionService.loadSession(oldSessionId); - } catch { - // Best-effort snapshot. Falling back to undefined still rolls - // back sessionId + recorder, which is the load-bearing invariant; - // we just lose the parentUuid chain on the restored recorder. - } - let coreSwapped = false; + let uiSwapped = false; + let prevSessionData: ResumedSessionData | undefined; try { - // 1. Flush outgoing recorder. + // 1. Flush outgoing recorder. Must happen BEFORE the parent snapshot + // so the snapshot captures `finalize()`'s trailing custom_title + // record — without that, a rollback restores the recorder with + // a stale `lastCompletedUuid` and the next user message attaches + // its parentUuid to a record that's no longer the JSONL tail. try { config.getChatRecordingService()?.finalize(); } catch { // best-effort } - // 2. Fork the JSONL on disk. + // 2. Snapshot the parent JSONL state for rollback. `/branch` is + // guarded on `isIdleRef`, so the file isn't being mutated + // concurrently between this load and the swap below. + try { + prevSessionData = await sessionService.loadSession(oldSessionId); + } catch { + // Best-effort snapshot. Falling back to undefined still rolls + // back sessionId + recorder, which is the load-bearing invariant; + // we just lose the parentUuid chain on the restored recorder. + } + + // 3. Fork the JSONL on disk. await sessionService.forkSession(oldSessionId, newSessionId); - // 3. Load the new file. + // 4. Load the new file. const resumed = await sessionService.loadSession(newSessionId); if (!resumed) { throw new Error('Failed to load newly forked session'); } - // 4. Swap core first. Anything that can still fail (startNewSession, + // 5. Swap core first. Anything that can still fail (startNewSession, // client init) runs while the UI is still showing the parent // session, so a throw leaves the user safely on the parent // instead of stranded with a cleared history and a half-live @@ -168,13 +172,19 @@ export function useBranchCommand( coreSwapped = true; await config.getGeminiClient()?.initialize?.(); - // 5. Swap UI. + // 6. Swap UI. Once this commits, rolling core back is unsafe — + // it would leave UI on the branch but recorder writing into + // the parent JSONL (the inverse split-brain). `uiSwapped` is + // set immediately after the UI commits so any subsequent + // failure (title, hook, remount, announce) skips the catch + // block's core rollback. const uiHistoryItems = buildResumedHistoryItems(resumed, config); startNewSession(newSessionId); historyManager.clearItems(); historyManager.loadHistory(uiHistoryItems); + uiSwapped = true; - // 6. Compute and apply the branch customTitle. + // 7. Compute and apply the branch customTitle. // The forked transcript is identical to the parent's, so reading // the first real user message from `resumed.conversation.messages` // mirrors Claude's "use the first parent message" behavior. @@ -187,7 +197,7 @@ export function useBranchCommand( config.getChatRecordingService()?.recordCustomTitle(effectiveTitle); setSessionName?.(effectiveTitle); - // 7. Fire SessionStart for the new session. A fork is semantically + // 8. Fire SessionStart for the new session. A fork is semantically // distinct from a resume — the sessionId is new and the transcript // is a derivative — so we use the dedicated `Branch` source value // to let hook consumers distinguish the two. @@ -203,10 +213,10 @@ export function useBranchCommand( config.getDebugLogger().warn(`SessionStart hook failed: ${err}`); } - // 8. Refresh terminal UI. + // 9. Refresh terminal UI. remount?.(); - // 9. Announce. Two history items mirror Claude's success message + // 10. Announce. Two history items mirror Claude's success message // (branched line + resume hint). The quoted name is the raw // user-provided `name`; no `(Branch)` suffix — that decoration // belongs in the picker/prompt bar, not in the user-facing @@ -232,11 +242,17 @@ export function useBranchCommand( Date.now(), ); } catch (err) { - if (coreSwapped) { - // Core already switched to the fork before the failure — put it + if (coreSwapped && !uiSwapped) { + // Core switched to the fork but UI hasn't swapped yet — put core // back on the parent, otherwise the recorder would keep writing // new user messages into the orphan fork JSONL while UI still // shows the parent. + // + // Skipped once `uiSwapped` is true: at that point UI is already + // on the branch, so reverting core would create the inverse + // split-brain (UI on branch, recorder on parent). Post-UI-swap + // failures (title, hook, remount, announce) are non-fatal and + // surfaced as an error item without unwinding the swap. try { config.startNewSession(oldSessionId, prevSessionData); // Re-hydrate chat history against the restored session. Best- From e03596dfcec95443df7705354768b42248791348 Mon Sep 17 00:00:00 2001 From: qqqys <qys177@gmail.com> Date: Wed, 6 May 2026 11:29:15 +0800 Subject: [PATCH 5/6] perf(session): fold /branch (Branch N) collision lookup into one scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `computeUniqueBranchTitle` was probing each `(Branch N)` candidate via `SessionService.findSessionsByTitle`, and that helper rescans the project's chats directory on every call. In dense title spaces /branch could end up doing the scan up to 99 times in a row before settling on a free suffix, which was visibly stalling the command. Add `SessionService.findSessionTitlesByPrefix(prefix)` — one project- wide scan that uses the cheap tail-read to extract each session's custom_title, filters to titles starting with the prefix, and applies the same project-scope filter as `findSessionsByTitle`. Heavy hydration steps (message count, prompt extraction) are skipped because collision lookup only needs the title. `computeUniqueBranchTitle` now does ONE call with prefix `${trimmed} (Branch`, builds an in-memory Set of taken titles, and picks the first free `(Branch)` / `(Branch N)` slot. Worst-case disk work drops from O(N) scans to one. Tests: new `findSessionTitlesByPrefix` describe in sessionService.test covers prefix match (case-insensitive), missing chats dir, project isolation, and files without a custom_title. useBranchCommand.test gains a perf invariant — even when 4 slots are taken, only ONE prefix-scan is issued. Reported by gpt-5.5 via Qwen Code \`/review\` on #3539. Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com> --- .../cli/src/ui/hooks/useBranchCommand.test.ts | 38 +++- packages/cli/src/ui/hooks/useBranchCommand.ts | 26 ++- .../core/src/services/sessionService.test.ts | 175 ++++++++++++++++++ packages/core/src/services/sessionService.ts | 61 ++++++ 4 files changed, 282 insertions(+), 18 deletions(-) diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index fcc1ad58ab5..434ded4db57 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -16,7 +16,7 @@ describe('useBranchCommand', () => { let startNewSessionConfig: ReturnType<typeof vi.fn>; let startNewSessionUI: ReturnType<typeof vi.fn>; let recordCustomTitle: ReturnType<typeof vi.fn>; - let findSessionsByTitle: ReturnType<typeof vi.fn>; + let findSessionTitlesByPrefix: ReturnType<typeof vi.fn>; let fireSessionStartEvent: ReturnType<typeof vi.fn>; let clearItems: ReturnType<typeof vi.fn>; let loadHistory: ReturnType<typeof vi.fn>; @@ -63,7 +63,7 @@ describe('useBranchCommand', () => { }); finalize = vi.fn(); recordCustomTitle = vi.fn().mockReturnValue(true); - findSessionsByTitle = vi.fn().mockResolvedValue([]); + findSessionTitlesByPrefix = vi.fn().mockResolvedValue([]); fireSessionStartEvent = vi.fn(); startNewSessionConfig = vi.fn(); startNewSessionUI = vi.fn(); @@ -77,7 +77,7 @@ describe('useBranchCommand', () => { getSessionService: () => ({ forkSession, loadSession, - findSessionsByTitle, + findSessionTitlesByPrefix, }), getChatRecordingService: () => ({ finalize, recordCustomTitle }), getGeminiClient: () => ({ initialize: vi.fn() }), @@ -136,12 +136,10 @@ describe('useBranchCommand', () => { }); it('bumps to (Branch N) when the default suffix is already taken', async () => { - findSessionsByTitle.mockImplementation(async (title: string) => { - if (title === 'my-branch (Branch)') { - return [{ sessionId: 'other', customTitle: title } as unknown]; - } - return []; - }); + // `findSessionTitlesByPrefix` returns every existing title under the + // `${name} (Branch` prefix in one shot, so the bump logic picks the + // first free slot in memory — no per-candidate disk probe. + findSessionTitlesByPrefix.mockResolvedValue(['my-branch (Branch)']); const { result } = renderHook(() => useBranchCommand(makeOptions())); await act(async () => { @@ -151,6 +149,28 @@ describe('useBranchCommand', () => { expect(setSessionName).toHaveBeenCalledWith('my-branch (Branch 2)'); }); + it('does ONE prefix scan even when many (Branch N) slots are taken', async () => { + // Pin the perf invariant: regardless of collision density, the + // collision lookup must be a single project-wide scan, not N probes. + // Reviewer's concern was that 99 sequential probes can stall /branch + // on dense title spaces. + findSessionTitlesByPrefix.mockResolvedValue([ + 'my-branch (Branch)', + 'my-branch (Branch 2)', + 'my-branch (Branch 3)', + 'my-branch (Branch 4)', + ]); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(findSessionTitlesByPrefix).toHaveBeenCalledTimes(1); + expect(findSessionTitlesByPrefix).toHaveBeenCalledWith('my-branch (Branch'); + expect(recordCustomTitle).toHaveBeenCalledWith('my-branch (Branch 5)'); + }); + it('derives the base title from the first user ChatRecord when no name is given', async () => { const { result } = renderHook(() => useBranchCommand(makeOptions())); await act(async () => { diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index 0468e03b652..4b20e218653 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -19,9 +19,10 @@ import type { UseHistoryManagerReturn } from './useHistoryManager.js'; import { t } from '../../i18n/index.js'; /** - * Cap for the `(Branch N)` collision scan. Each probe is a project-wide - * scan via `findSessionsByTitle`; 99 is generous for realistic use and - * bounds the worst case. + * Cap for the `(Branch N)` collision suffix. We scan all matching titles + * once via `findSessionTitlesByPrefix` and then pick the first free slot + * in memory; 99 is generous for realistic use and bounds the timestamp- + * fallback path on pathologically dense title spaces. */ const MAX_BRANCH_COLLISION_SCAN = 99; @@ -60,21 +61,28 @@ function deriveFirstPrompt(messages: ChatRecord[]): string { * Appends ` (Branch)` to `baseName`, bumping to ` (Branch 2)`, ` (Branch 3)`, * ... when the exact name is already taken by another session's customTitle * in the current project. Mirrors Claude's `getUniqueForkName`. + * + * Does ONE prefix scan instead of probing each candidate via + * `findSessionsByTitle`: in dense title spaces the per-probe scanner could + * walk the project's chat directory up to {@link MAX_BRANCH_COLLISION_SCAN} + * times, and `/branch` would visibly stall. We collect every existing + * `${trimmed} (Branch...` title once, then pick the first free slot in memory. */ async function computeUniqueBranchTitle( baseName: string, sessionService: SessionService, ): Promise<string> { const trimmed = baseName.trim(); + const taken = new Set( + (await sessionService.findSessionTitlesByPrefix(`${trimmed} (Branch`)).map( + (t) => t.toLowerCase().trim(), + ), + ); const first = `${trimmed} (Branch)`; - if ((await sessionService.findSessionsByTitle(first)).length === 0) { - return first; - } + if (!taken.has(first.toLowerCase())) return first; for (let n = 2; n <= MAX_BRANCH_COLLISION_SCAN; n++) { const candidate = `${trimmed} (Branch ${n})`; - if ((await sessionService.findSessionsByTitle(candidate)).length === 0) { - return candidate; - } + if (!taken.has(candidate.toLowerCase())) return candidate; } // Pathological density — timestamp fallback keeps the fork unique. return `${trimmed} (Branch ${Date.now()})`; diff --git a/packages/core/src/services/sessionService.test.ts b/packages/core/src/services/sessionService.test.ts index 9671c3944b4..feab2be20c4 100644 --- a/packages/core/src/services/sessionService.test.ts +++ b/packages/core/src/services/sessionService.test.ts @@ -1112,4 +1112,179 @@ describe('SessionService', () => { ); }); }); + + describe('findSessionTitlesByPrefix', () => { + // Uses real disk like forkSession — readSessionTitleInfoFromFile reads + // the file tail for the custom_title record, so mocks would defeat the + // method. Mirrors the forkSession describe's setup verbatim so the tmp + // sandbox + un-mocked path/jsonl utilities are in place. + let realTmpDir: string; + let realPath: typeof import('node:path'); + let service: SessionService; + let cwd: string; + + beforeEach(async () => { + const realOs = await import('node:os'); + realPath = await vi.importActual<typeof import('node:path')>('node:path'); + const actualPaths = + await vi.importActual<typeof import('../utils/paths.js')>( + '../utils/paths.js', + ); + const actualJsonl = await vi.importActual< + typeof import('../utils/jsonl-utils.js') + >('../utils/jsonl-utils.js'); + + vi.mocked(path.join).mockImplementation( + realPath.join as unknown as typeof path.join, + ); + vi.mocked(path.dirname).mockImplementation( + realPath.dirname as unknown as typeof path.dirname, + ); + vi.mocked(path.isAbsolute).mockImplementation( + realPath.isAbsolute as unknown as typeof path.isAbsolute, + ); + vi.mocked(path.resolve).mockImplementation( + realPath.resolve as unknown as typeof path.resolve, + ); + vi.mocked(getProjectHash).mockImplementation(actualPaths.getProjectHash); + const mockedPaths = (await import('../utils/paths.js')) as unknown as { + sanitizeCwd: (cwd: string) => string; + }; + mockedPaths.sanitizeCwd = actualPaths.sanitizeCwd; + vi.mocked(jsonl.read).mockImplementation(actualJsonl.read); + vi.mocked(jsonl.readLines).mockImplementation(actualJsonl.readLines); + + vi.mocked(readdirSyncSpy).mockRestore?.(); + vi.mocked(statSyncSpy).mockRestore?.(); + vi.mocked(unlinkSyncSpy).mockRestore?.(); + + realTmpDir = fs.mkdtempSync( + realPath.join(realOs.tmpdir(), 'find-titles-prefix-'), + ); + process.env['QWEN_RUNTIME_DIR'] = realTmpDir; + cwd = process.cwd(); + service = new SessionService(cwd); + }); + + afterEach(() => { + delete process.env['QWEN_RUNTIME_DIR']; + try { + fs.rmSync(realTmpDir, { recursive: true, force: true }); + } catch { + // best-effort + } + }); + + const seedSessionWithTitle = ( + sessionId: string, + title: string, + sessionCwd: string = cwd, + ) => { + const chatsDir = realPath.join( + service['storage'].getProjectDir(), + 'chats', + ); + fs.mkdirSync(chatsDir, { recursive: true }); + const file = realPath.join(chatsDir, `${sessionId}.jsonl`); + const lines = [ + { + uuid: 'u1', + parentUuid: null, + sessionId, + type: 'user', + timestamp: '2026-04-22T00:00:00.000Z', + cwd: sessionCwd, + version: 'test', + message: { role: 'user', parts: [{ text: 'hello' }] }, + }, + { + uuid: 'u2', + parentUuid: 'u1', + sessionId, + type: 'system', + subtype: 'custom_title', + timestamp: '2026-04-22T00:00:01.000Z', + cwd: sessionCwd, + version: 'test', + systemPayload: { customTitle: title, titleSource: 'manual' }, + }, + ]; + fs.writeFileSync( + file, + lines.map((l) => JSON.stringify(l)).join('\n') + '\n', + ); + return file; + }; + + it('returns titles whose custom_title starts with the prefix (case-insensitive)', async () => { + seedSessionWithTitle( + '11111111-1111-1111-1111-111111111111', + 'my-branch (Branch)', + ); + seedSessionWithTitle( + '22222222-2222-2222-2222-222222222222', + 'My-Branch (Branch 2)', + ); + seedSessionWithTitle( + '33333333-3333-3333-3333-333333333333', + 'unrelated session', + ); + + const titles = + await service.findSessionTitlesByPrefix('my-branch (Branch'); + + expect(new Set(titles)).toEqual( + new Set(['my-branch (Branch)', 'My-Branch (Branch 2)']), + ); + }); + + it('returns empty when chats directory does not exist', async () => { + const titles = await service.findSessionTitlesByPrefix('anything'); + expect(titles).toEqual([]); + }); + + it('skips sessions from other projects (collisions are project-scoped)', async () => { + seedSessionWithTitle( + '11111111-1111-1111-1111-111111111111', + 'shared (Branch)', + cwd, + ); + // Same chats dir (sessions are stored under projectHash anyway), but + // the record's cwd belongs to another project → must be skipped. + seedSessionWithTitle( + '22222222-2222-2222-2222-222222222222', + 'shared (Branch 2)', + '/some/other/project', + ); + + const titles = await service.findSessionTitlesByPrefix('shared (Branch'); + expect(titles).toEqual(['shared (Branch)']); + }); + + it('skips files without a custom_title record', async () => { + const sessionId = '11111111-1111-1111-1111-111111111111'; + const chatsDir = realPath.join( + service['storage'].getProjectDir(), + 'chats', + ); + fs.mkdirSync(chatsDir, { recursive: true }); + const file = realPath.join(chatsDir, `${sessionId}.jsonl`); + fs.writeFileSync( + file, + JSON.stringify({ + uuid: 'u1', + parentUuid: null, + sessionId, + type: 'user', + timestamp: '2026-04-22T00:00:00.000Z', + cwd, + version: 'test', + message: { role: 'user', parts: [{ text: 'hi' }] }, + }) + '\n', + ); + + const titles = await service.findSessionTitlesByPrefix('anything'); + expect(titles).toEqual([]); + }); + }); }); diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index 04f360257ed..f456e067d1a 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -904,6 +904,67 @@ export class SessionService { return matches; } + /** + * Returns the customTitles in this project that start with `prefix` + * (case-insensitive). Single project-wide scan — meant to replace + * repeated `findSessionsByTitle()` probes when the caller needs to + * pick the first free `(Branch N)` slot in memory. + * + * Skips the heavy hydration steps (message count, prompt extraction) + * that `findSessionsByTitle` does — collision lookup only needs the + * title and a project filter, so we read the first record only when + * the title actually matches the prefix. + * + * @param prefix Case-insensitive title prefix to match. + */ + async findSessionTitlesByPrefix(prefix: string): Promise<string[]> { + const normalizedPrefix = prefix.toLowerCase().trim(); + const titles: string[] = []; + const chatsDir = this.getChatsDir(); + + let fileNames: string[]; + try { + fileNames = fs.readdirSync(chatsDir); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') { + return titles; + } + throw error; + } + + let filesProcessed = 0; + for (const name of fileNames) { + if (!SESSION_FILE_PATTERN.test(name)) continue; + if (filesProcessed >= MAX_FILES_TO_PROCESS) break; + filesProcessed++; + + const filePath = path.join(chatsDir, name); + + // Cheap tail-read to extract the title before doing any project- + // filter work. Saves a per-file jsonl.readLines on the common + // case where most sessions don't share this prefix. + const titleInfo = this.readSessionTitleInfoFromFile(filePath); + if (!titleInfo.title) continue; + const normalizedTitle = titleInfo.title.toLowerCase().trim(); + if (!normalizedTitle.startsWith(normalizedPrefix)) continue; + + // Project filter — same semantics as findSessionsByTitle: scope + // collisions to the current project so a fork in another project + // can't make this one bump unnecessarily. + try { + const records = await jsonl.readLines<ChatRecord>(filePath, 1); + if (records.length === 0) continue; + if (getProjectHash(records[0].cwd) !== this.projectHash) continue; + } catch { + continue; + } + + titles.push(titleInfo.title); + } + + return titles; + } + /** * Loads the most recent session for the current project. * Combines listSessions and loadSession for convenience. From 36693a32ee50df38ae05450cfe95108eb62457f8 Mon Sep 17 00:00:00 2001 From: qqqys <qys177@gmail.com> Date: Fri, 8 May 2026 19:13:58 +0800 Subject: [PATCH 6/6] test(cli): tighten mocks and drop dead assertion in slashCommandProcessor tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses today's review feedback on #3539 plus two tsc gaps the IDE flagged in the same file. 1. ChatRecordingService cast (TS2352) — route through `unknown` at the two `recorder = mockConfig.getChatRecordingService() as { recordSlashCommand }` sites in SLASH_COMMANDS_SKIP_RECORDING. Insufficient overlap between `ChatRecordingService | undefined` and the inline mock shape; the existing single-step cast doesn't compile under strict. 2. SlashCommandProcessorActions mock missing `handleBranch` — this PR added `handleBranch: (name?: string) => Promise<void>` to the actions surface (commit 8ac4af285), but `createMockActions()` was never updated, so the mock failed to satisfy the type. Added `handleBranch: vi.fn().mockResolvedValue(undefined)`. 3. `stripThoughtsFromHistory` cleanup in load_history tests — `GeminiClient` has no `stripThoughtsFromHistory` method (the helper lives inside `sessionService.ts` and is never called from the slash processor), so the mocked field was a zombie and the assertion `expect(mockClient.stripThoughtsFromHistory).not.toHaveBeenCalled()` was vacuously true — it could never fail and provided zero regression guard. Replaced with `expect(mockClient.setHistory).toHaveBeenCalledWith(historyWithThoughts)`, which is what "preserve thoughts" actually means: the `thoughtSignature` inside `clientHistory` reaches `setHistory` untouched. This will fail the day someone reintroduces strip-on-load. Tests: vitest packages/cli/src/ui/hooks/slashCommandProcessor.test.ts 42 ✓ tsc -p packages/cli/tsconfig.json --noEmit clean Co-Authored-By: Qwen-Coder <noreply@alibabacloud.com> --- packages/cli/src/ui/hooks/slashCommandProcessor.test.ts | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 1f3be04ac90..7279ddd4e83 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -138,6 +138,7 @@ describe('useSlashCommandProcessor', () => { openApprovalModeDialog: vi.fn(), openResumeDialog: vi.fn(), handleResume: vi.fn(), + handleBranch: vi.fn().mockResolvedValue(undefined), openDeleteDialog: vi.fn(), quit: mockSetQuittingMessages, setDebugMessage: vi.fn(), @@ -503,7 +504,6 @@ describe('useSlashCommandProcessor', () => { it('should handle "load_history" action', async () => { const mockClient = { setHistory: vi.fn(), - stripThoughtsFromHistory: vi.fn(), } as unknown as GeminiClient; vi.spyOn(mockConfig, 'getGeminiClient').mockReturnValue(mockClient); @@ -532,7 +532,6 @@ describe('useSlashCommandProcessor', () => { it('should preserve 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); @@ -559,7 +558,7 @@ describe('useSlashCommandProcessor', () => { }); expect(mockClient.setHistory).toHaveBeenCalledTimes(1); - expect(mockClient.stripThoughtsFromHistory).not.toHaveBeenCalled(); + expect(mockClient.setHistory).toHaveBeenCalledWith(historyWithThoughts); }); it('should handle a "quit" action', async () => { @@ -1204,7 +1203,7 @@ describe('useSlashCommandProcessor', () => { const result = setupProcessorHook([branchCmd]); await waitFor(() => expect(result.current.slashCommands).toHaveLength(1)); - const recorder = mockConfig.getChatRecordingService() as { + const recorder = mockConfig.getChatRecordingService() as unknown as { recordSlashCommand: ReturnType<typeof vi.fn>; }; recorder.recordSlashCommand.mockClear(); @@ -1224,7 +1223,7 @@ describe('useSlashCommandProcessor', () => { const result = setupProcessorHook([testCmd]); await waitFor(() => expect(result.current.slashCommands).toHaveLength(1)); - const recorder = mockConfig.getChatRecordingService() as { + const recorder = mockConfig.getChatRecordingService() as unknown as { recordSlashCommand: ReturnType<typeof vi.fn>; }; recorder.recordSlashCommand.mockClear();