diff --git a/desktop/src/api/sessions.test.ts b/desktop/src/api/sessions.test.ts new file mode 100644 index 00000000..fb0bde89 --- /dev/null +++ b/desktop/src/api/sessions.test.ts @@ -0,0 +1,42 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { setBaseUrl } from './client' +import { sessionsApi } from './sessions' + +describe('sessionsApi', () => { + afterEach(() => { + setBaseUrl('http://127.0.0.1:3456') + vi.restoreAllMocks() + }) + + it('posts branch requests to the session branch endpoint', async () => { + const fetchMock = vi.spyOn(globalThis, 'fetch') + fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({ + sessionId: 'branch-session', + title: 'Branch', + workDir: '/workspace/repo', + sourceSessionId: 'source-session', + targetMessageId: 'message-1', + }), { + status: 201, + headers: { 'Content-Type': 'application/json' }, + })) + + setBaseUrl('http://127.0.0.1:49237') + const result = await sessionsApi.branch('source-session', { + targetMessageId: 'message-1', + title: 'Branch', + }) + + expect(result.sessionId).toBe('branch-session') + expect(fetchMock).toHaveBeenCalledOnce() + const [url, init] = fetchMock.mock.calls[0]! + expect(url).toBe('http://127.0.0.1:49237/api/sessions/source-session/branch') + expect(init).toMatchObject({ + method: 'POST', + body: JSON.stringify({ + targetMessageId: 'message-1', + title: 'Branch', + }), + }) + }) +}) diff --git a/desktop/src/api/sessions.ts b/desktop/src/api/sessions.ts index 6687b149..d1ac7ffb 100644 --- a/desktop/src/api/sessions.ts +++ b/desktop/src/api/sessions.ts @@ -40,6 +40,17 @@ export type CreateSessionRequest = { workDir?: string repository?: CreateSessionRepositoryOptions } +export type BranchSessionRequest = { + targetMessageId: string + title?: string +} +export type BranchSessionResponse = { + sessionId: string + title: string + workDir: string | null + sourceSessionId: string + targetMessageId: string +} export type RepositoryBranchInfo = { name: string current: boolean @@ -314,6 +325,10 @@ export const sessionsApi = { return api.post('/api/sessions', body) }, + branch(sessionId: string, body: BranchSessionRequest) { + return api.post(`/api/sessions/${sessionId}/branch`, body) + }, + delete(sessionId: string) { return api.delete<{ ok: true }>(`/api/sessions/${sessionId}`) }, diff --git a/desktop/src/components/chat/MessageList.test.tsx b/desktop/src/components/chat/MessageList.test.tsx index e4ae602b..da2c1577 100644 --- a/desktop/src/components/chat/MessageList.test.tsx +++ b/desktop/src/components/chat/MessageList.test.tsx @@ -6,6 +6,7 @@ import { sessionsApi } from '../../api/sessions' import { useChatStore } from '../../stores/chatStore' import { useWorkspaceChatContextStore } from '../../stores/workspaceChatContextStore' import { useSettingsStore } from '../../stores/settingsStore' +import { useSessionStore } from '../../stores/sessionStore' import { useTabStore } from '../../stores/tabStore' import { useUIStore } from '../../stores/uiStore' import type { UIMessage } from '../../types/chat' @@ -105,6 +106,7 @@ describe('MessageList nested tool calls', () => { useSettingsStore.setState({ locale: 'en' }) useUIStore.setState({ pendingSettingsTab: null }) useTabStore.setState({ activeTabId: ACTIVE_TAB, tabs: [{ sessionId: ACTIVE_TAB, title: 'Test', type: 'session' as const, status: 'idle' }] }) + useSessionStore.setState({ sessions: [], activeSessionId: null, isLoading: false, error: null }) useChatStore.setState({ sessions: { [ACTIVE_TAB]: makeSessionState() } }) useWorkspaceChatContextStore.setState(useWorkspaceChatContextStore.getInitialState(), true) vi.spyOn(sessionsApi, 'getTurnCheckpoints').mockImplementation( @@ -2290,6 +2292,107 @@ describe('MessageList nested tool calls', () => { expect(screen.queryByRole('button', { name: 'Rewind to here' })).toBeNull() }) + it('branches from completed transcript-backed chat messages using the original transcript id', async () => { + const branchSession = vi.fn().mockResolvedValue({ + sessionId: 'branched-session-1', + title: 'Branched session', + workDir: '/tmp/branched-session-1', + }) + const connectToSession = vi.fn() + useSessionStore.setState({ + sessions: [{ + id: ACTIVE_TAB, + title: 'Source session', + createdAt: '2026-05-19T00:00:00.000Z', + modifiedAt: '2026-05-19T00:00:00.000Z', + messageCount: 2, + projectPath: '/tmp/source-project', + projectRoot: '/tmp/source-project', + workDir: '/tmp/source-project', + workDirExists: true, + }], + branchSession: branchSession as never, + }) + useChatStore.setState({ + connectToSession: connectToSession as never, + sessions: { + [ACTIVE_TAB]: makeSessionState({ + messages: [ + { + id: 'local-user-1', + transcriptMessageId: 'transcript-user-1', + type: 'user_text', + content: '从这里开始', + timestamp: 1, + }, + { + id: 'local-assistant-1', + transcriptMessageId: 'transcript-assistant-1', + type: 'assistant_text', + content: '这是完成的答复。', + timestamp: 2, + }, + ], + }), + }, + }) + + render() + + const branchButtons = screen.getAllByRole('button', { name: 'Branch from here' }) + expect(branchButtons).toHaveLength(2) + + fireEvent.click(branchButtons[1]!) + + await waitFor(() => { + expect(branchSession).toHaveBeenCalledWith(ACTIVE_TAB, 'transcript-assistant-1') + }) + expect(connectToSession).toHaveBeenCalledWith('branched-session-1') + expect(useTabStore.getState().activeTabId).toBe('branched-session-1') + const tabs = useTabStore.getState().tabs + expect(tabs[tabs.length - 1]).toMatchObject({ + sessionId: 'branched-session-1', + title: 'Branched session', + type: 'session', + }) + const toasts = useUIStore.getState().toasts + expect(toasts[toasts.length - 1]).toMatchObject({ + type: 'success', + message: 'Created branched session "Branched session".', + }) + }) + + it('hides branch actions while the current session is still running', () => { + useChatStore.setState({ + sessions: { + [ACTIVE_TAB]: makeSessionState({ + chatState: 'streaming', + streamingText: 'partial', + messages: [ + { + id: 'local-user-1', + transcriptMessageId: 'transcript-user-1', + type: 'user_text', + content: '从这里开始', + timestamp: 1, + }, + { + id: 'local-assistant-1', + transcriptMessageId: 'transcript-assistant-1', + type: 'assistant_text', + content: '这是完成的答复。', + timestamp: 2, + }, + ], + }), + }, + }) + + render() + + expect(screen.queryByRole('button', { name: 'Branch from here' })).toBeNull() + }) + it('keeps historical sessions readable when turn checkpoint payloads are missing', async () => { vi.spyOn(sessionsApi, 'getTurnCheckpoints').mockResolvedValue({} as never) diff --git a/desktop/src/components/chat/MessageList.tsx b/desktop/src/components/chat/MessageList.tsx index e88fdf57..35fda7ed 100644 --- a/desktop/src/components/chat/MessageList.tsx +++ b/desktop/src/components/chat/MessageList.tsx @@ -1,8 +1,9 @@ import { useRef, useEffect, useMemo, memo, useState, useCallback, useLayoutEffect, type ReactNode } from 'react' -import { ArrowDown, BookMarked, Bot, CheckCircle2, ChevronDown, ChevronRight, CircleStop, LoaderCircle, MessageCircle, Settings, Target, XCircle } from 'lucide-react' +import { ArrowDown, BookMarked, Bot, CheckCircle2, ChevronDown, ChevronRight, CircleStop, GitBranch, LoaderCircle, MessageCircle, Settings, Target, XCircle } from 'lucide-react' import { ApiError } from '../../api/client' import { sessionsApi, type SessionTurnCheckpoint } from '../../api/sessions' import { useChatStore } from '../../stores/chatStore' +import { useSessionStore } from '../../stores/sessionStore' import { useWorkspaceChatContextStore } from '../../stores/workspaceChatContextStore' import { SETTINGS_TAB_ID, useTabStore } from '../../stores/tabStore' import { useTeamStore } from '../../stores/teamStore' @@ -48,6 +49,11 @@ type RewindTurnTarget = { attachments?: Extract['attachments'] } +type BranchableMessageTarget = { + uiMessageId: string + transcriptMessageId: string +} + type TurnChangeCardModel = { target: RewindTurnTarget checkpoint: SessionTurnCheckpoint @@ -353,6 +359,37 @@ function SelectableChatMessage({ ) } +function MessageBranchAction({ + align, + label, + loading, + onClick, +}: { + align: ChatMessageRole + label: string + loading?: boolean + onClick: () => void +}) { + return ( +
+ +
+ ) +} + function appendChildToolCall( childToolCallsByParent: Map, parentToolUseId: string, @@ -450,6 +487,48 @@ function isTurnResponseMessage(message: UIMessage) { ) } +function getBranchableMessageTargets(messages: UIMessage[]): Map { + const branchableTargets = new Map() + let currentTurnCandidates: Array> = [] + let hasResponseForCurrentTurn = false + + const markCurrentTurnBranchable = () => { + if (!hasResponseForCurrentTurn) return + for (const candidate of currentTurnCandidates) { + if (!candidate.transcriptMessageId) continue + branchableTargets.set(candidate.id, { + uiMessageId: candidate.id, + transcriptMessageId: candidate.transcriptMessageId, + }) + } + } + + for (const message of messages) { + if (message.type === 'user_text') { + markCurrentTurnBranchable() + currentTurnCandidates = [] + hasResponseForCurrentTurn = false + if (!message.pending && message.transcriptMessageId) { + currentTurnCandidates = [message] + } + continue + } + + if (currentTurnCandidates.length === 0) continue + + if (isTurnResponseMessage(message)) { + hasResponseForCurrentTurn = true + } + + if (message.type === 'assistant_text' && message.transcriptMessageId) { + currentTurnCandidates.push(message) + } + } + + markCurrentTurnBranchable() + return branchableTargets +} + export function getCompletedTurnTargets(messages: UIMessage[]): RewindTurnTarget[] { let userMessageIndex = -1 const completedTurns: RewindTurnTarget[] = [] @@ -628,6 +707,7 @@ type MessageListProps = { const AUTO_SCROLL_BOTTOM_THRESHOLD_PX = 48 const MAX_SCROLL_SNAPSHOTS = 100 +const EMPTY_MESSAGES: UIMessage[] = [] const CHAT_SCROLL_AREA_CLASS = [ 'chat-scroll-area', '[scrollbar-width:auto]', @@ -685,6 +765,7 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { const sessionState = useChatStore((s) => resolvedSessionId ? s.sessions[resolvedSessionId] : undefined, ) + const branchSession = useSessionStore((s) => s.branchSession) const stopGeneration = useChatStore((s) => s.stopGeneration) const reloadHistory = useChatStore((s) => s.reloadHistory) const queueComposerPrefill = useChatStore((s) => s.queueComposerPrefill) @@ -692,7 +773,7 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { resolvedSessionId ? Boolean(s.getMemberBySessionId(resolvedSessionId)) : false, ) const addToast = useUIStore((s) => s.addToast) - const messages = sessionState?.messages ?? [] + const messages = sessionState?.messages ?? EMPTY_MESSAGES const chatState = sessionState?.chatState ?? 'idle' const streamingText = sessionState?.streamingText ?? '' const activeThinkingId = sessionState?.activeThinkingId ?? null @@ -713,9 +794,17 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { const [turnChangeLoadError, setTurnChangeLoadError] = useState(null) const [turnActionErrors, setTurnActionErrors] = useState>({}) const [isLoadingTurnChangeCards, setIsLoadingTurnChangeCards] = useState(false) + const [branchingMessageId, setBranchingMessageId] = useState(null) const [rewindingTurnId, setRewindingTurnId] = useState(null) const [turnUndoConfirmTargetId, setTurnUndoConfirmTargetId] = useState(null) const [showJumpToLatest, setShowJumpToLatest] = useState(false) + const branchActionsDisabled = + isMemberSession || + chatState !== 'idle' || + streamingText.trim().length > 0 || + Boolean(activeThinkingId) || + Boolean(sessionState?.activeToolUseId) || + Boolean(sessionState?.activeToolName) const scrollToBottom = useCallback((behavior: ScrollBehavior) => { shouldAutoScrollRef.current = true @@ -842,6 +931,10 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { () => buildRenderModel(messages), [messages], ) + const branchableMessageTargets = useMemo( + () => branchActionsDisabled ? new Map() : getBranchableMessageTargets(messages), + [branchActionsDisabled, messages], + ) const completedTurnTargets = useMemo(() => getCompletedTurnTargets(messages), [messages]) const latestCompletedTurnId = completedTurnTargets.length > 0 @@ -982,6 +1075,29 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { t, ]) + const handleBranchMessage = useCallback(async (target: BranchableMessageTarget) => { + if (!resolvedSessionId || branchingMessageId) return + + setBranchingMessageId(target.uiMessageId) + try { + const result = await branchSession(resolvedSessionId, target.transcriptMessageId) + const title = result.title.trim() || t('sidebar.newSession') + useTabStore.getState().openTab(result.sessionId, title) + useChatStore.getState().connectToSession(result.sessionId) + addToast({ + type: 'success', + message: t('chat.branchSuccess', { title }), + }) + } catch (error) { + addToast({ + type: 'error', + message: t('chat.branchError', { detail: getApiErrorMessage(error) }), + }) + } finally { + setBranchingMessageId(null) + } + }, [addToast, branchSession, branchingMessageId, resolvedSessionId, t]) + return (
{ + const branchTarget = branchableMessageTargets.get(item.message.id) + if (!branchTarget) return undefined + return { + label: t('chat.branchFromHere'), + loading: branchingMessageId === branchTarget.uiMessageId, + onBranch: () => { void handleBranchMessage(branchTarget) }, + } + })() + } /> )} @@ -1112,12 +1239,18 @@ export const MessageBlock = memo(function MessageBlock({ activeThinkingId, agentTaskNotifications, toolResult, + branchAction, }: { sessionId?: string | null message: UIMessage activeThinkingId: string | null agentTaskNotifications: Record toolResult?: { content: unknown; isError: boolean } | null + branchAction?: { + label: string + loading?: boolean + onBranch: () => void + } }) { const t = useTranslation() @@ -1130,10 +1263,20 @@ export const MessageBlock = memo(function MessageBlock({ role="user" content={message.content} > - +
+ + {branchAction ? ( + + ) : null} +
) case 'assistant_text': @@ -1144,7 +1287,17 @@ export const MessageBlock = memo(function MessageBlock({ role="assistant" content={message.content} > - +
+ + {branchAction ? ( + + ) : null} +
) case 'thinking': diff --git a/desktop/src/i18n/locales/en.ts b/desktop/src/i18n/locales/en.ts index 58f25da2..f31d0b6d 100644 --- a/desktop/src/i18n/locales/en.ts +++ b/desktop/src/i18n/locales/en.ts @@ -951,6 +951,9 @@ export const en = { 'chat.workspaceReferencesOnly': 'Added {count} workspace references', 'chat.contextReferencesOnly': 'Added {count} references', 'chat.addSelectionToChat': 'Add to chat', + 'chat.branchFromHere': 'Branch from here', + 'chat.branchSuccess': 'Created branched session "{title}".', + 'chat.branchError': 'Failed to branch from this message. Detail: {detail}', 'chat.userMessageReference': 'User message', 'chat.assistantMessageReference': 'Assistant message', 'chat.slashCommands': 'Slash commands', diff --git a/desktop/src/i18n/locales/zh.ts b/desktop/src/i18n/locales/zh.ts index 9fb6e331..350509b0 100644 --- a/desktop/src/i18n/locales/zh.ts +++ b/desktop/src/i18n/locales/zh.ts @@ -953,6 +953,9 @@ export const zh: Record = { 'chat.workspaceReferencesOnly': '已添加 {count} 个工作区引用', 'chat.contextReferencesOnly': '已添加 {count} 个引用', 'chat.addSelectionToChat': '添加到对话', + 'chat.branchFromHere': '从这里派生', + 'chat.branchSuccess': '已创建派生会话“{title}”。', + 'chat.branchError': '从该消息派生失败。详情：{detail}', 'chat.userMessageReference': '用户消息', 'chat.assistantMessageReference': 'AI 回复', 'chat.slashCommands': '斜杠命令', diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index 2434e28f..f287fe77 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -172,6 +172,8 @@ describe('chatStore history mapping', () => { updateTabTitleMock.mockReset() updateTabStatusMock.mockReset() updateSessionTitleMock.mockReset() + vi.mocked(sessionsApi.getMessages).mockReset() + vi.mocked(sessionsApi.getMessages).mockResolvedValue({ messages: [] }) sessionStoreSnapshot.sessions = [] cliTaskStoreSnapshot.tasks = [] cliTaskStoreSnapshot.sessionId = null @@ -251,6 +253,44 @@ describe('chatStore history mapping', () => { ]) }) + it('preserves transcript message ids on natural-language history messages', () => { + const messages: MessageEntry[] = [ + { + id: 'transcript-user-1', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: '请从这里继续', + }, + { + id: 'transcript-assistant-1', + type: 'assistant', + timestamp: '2026-04-06T00:00:01.000Z', + model: 'opus', + content: [ + { type: 'text', text: '这里是答复。' }, + { type: 'tool_use', name: 'Read', id: 'tool-1', input: { file_path: 'src/App.tsx' } }, + ], + }, + ] + + const mapped = mapHistoryMessagesToUiMessages(messages) + + expect(mapped).toMatchObject([ + { + id: 'transcript-user-1', + type: 'user_text', + transcriptMessageId: 'transcript-user-1', + }, + { + type: 'assistant_text', + transcriptMessageId: 'transcript-assistant-1', + }, + { + type: 'tool_use', + }, + ]) + }) + it('restores /goal local command output from transcript history', () => { const messages: MessageEntry[] = [ { @@ -458,6 +498,143 @@ describe('chatStore history mapping', () => { }) }) + it('hydrates transcript ids for a just-completed live turn', async () => { + vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({ + messages: [ + { + id: 'transcript-user-1', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: 'live prompt', + }, + { + id: 'transcript-assistant-1', + type: 'assistant', + timestamp: '2026-04-06T00:00:01.000Z', + content: 'live answer', + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + messages: [ + { + id: 'live-user', + type: 'user_text', + content: 'live prompt', + timestamp: 1, + }, + ], + streamingText: 'live answer', + chatState: 'streaming', + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 2 }, + }) + + await vi.waitFor(() => { + expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([ + { + type: 'user_text', + transcriptMessageId: 'transcript-user-1', + }, + { + type: 'assistant_text', + transcriptMessageId: 'transcript-assistant-1', + }, + ]) + }) + }) + + it('retries transcript id hydration after the assistant message is persisted', async () => { + vi.useFakeTimers() + vi.mocked(sessionsApi.getMessages) + .mockResolvedValueOnce({ + messages: [ + { + id: 'transcript-user-1', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: 'live prompt', + }, + ], + }) + .mockResolvedValueOnce({ + messages: [ + { + id: 'transcript-user-1', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: 'live prompt', + }, + { + id: 'transcript-assistant-1', + type: 'assistant', + timestamp: '2026-04-06T00:00:01.000Z', + content: 'live answer', + }, + ], + }) + + useChatStore.setState({ + sessions: { + [TEST_SESSION_ID]: makeSession({ + messages: [ + { + id: 'live-user', + type: 'user_text', + content: 'live prompt', + timestamp: 1, + }, + ], + streamingText: 'live answer', + chatState: 'streaming', + }), + }, + }) + + useChatStore.getState().handleServerMessage(TEST_SESSION_ID, { + type: 'message_complete', + usage: { input_tokens: 1, output_tokens: 2 }, + }) + + await Promise.resolve() + await Promise.resolve() + + expect(sessionsApi.getMessages).toHaveBeenCalledTimes(1) + const firstHydrationMessages = useChatStore.getState().sessions[TEST_SESSION_ID]?.messages ?? [] + expect(firstHydrationMessages[0]).toMatchObject({ + type: 'user_text', + transcriptMessageId: 'transcript-user-1', + }) + expect(firstHydrationMessages[1]).toMatchObject({ + type: 'assistant_text', + }) + expect(firstHydrationMessages[1]).not.toHaveProperty('transcriptMessageId') + + await vi.advanceTimersByTimeAsync(750) + + expect(sessionsApi.getMessages).toHaveBeenCalledTimes(2) + const secondHydrationMessages = useChatStore.getState().sessions[TEST_SESSION_ID]?.messages ?? [] + expect(secondHydrationMessages[0]).toMatchObject({ + type: 'user_text', + transcriptMessageId: 'transcript-user-1', + }) + expect(secondHydrationMessages[1]).toMatchObject({ + type: 'assistant_text', + transcriptMessageId: 'transcript-assistant-1', + }) + + vi.runOnlyPendingTimers() + vi.useRealTimers() + }) + it('merges consecutive assistant text blocks when restoring transcript history', () => { const messages: MessageEntry[] = [ { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 921985d3..f64d2350 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -233,15 +233,26 @@ function appendAssistantTextMessage( content: string, timestamp: number, model?: string, + transcriptMessageId?: string, ): UIMessage[] { if (!content.trim()) return messages const last = messages[messages.length - 1] - if (last?.type === 'assistant_text') { + const canMergeIntoLast = + last?.type === 'assistant_text' && + ( + transcriptMessageId + ? last.transcriptMessageId === transcriptMessageId + : !last.transcriptMessageId + ) + if (canMergeIntoLast) { const merged: UIMessage = { ...last, content: last.content + content, ...(model ?? last.model ? { model: model ?? last.model } : {}), + ...(transcriptMessageId ?? last.transcriptMessageId + ? { transcriptMessageId: transcriptMessageId ?? last.transcriptMessageId } + : {}), } return [...messages.slice(0, -1), merged] } @@ -253,6 +264,7 @@ function appendAssistantTextMessage( type: 'assistant_text', content, timestamp, + ...(transcriptMessageId ? { transcriptMessageId } : {}), ...(model ? { model } : {}), }, ] @@ -330,6 +342,91 @@ function mergeRestoredTerminalGoalEvents( : messages } +function mergeRestoredTranscriptMessageIds( + messages: UIMessage[], + restoredMessages: UIMessage[], +): UIMessage[] { + const restoredCandidates = restoredMessages.filter(( + message, + ): message is Extract => + (message.type === 'user_text' || message.type === 'assistant_text') && + typeof message.transcriptMessageId === 'string' && + message.transcriptMessageId.length > 0) + + if (restoredCandidates.length === 0) return messages + + let restoredCursor = 0 + let changed = false + const merged = messages.map((message) => { + if ( + (message.type !== 'user_text' && message.type !== 'assistant_text') || + message.transcriptMessageId + ) { + return message + } + + const matchIndex = restoredCandidates.findIndex((candidate, index) => + index >= restoredCursor && + candidate.type === message.type && + candidate.content.trim() === message.content.trim()) + + if (matchIndex === -1) return message + + restoredCursor = matchIndex + 1 + changed = true + return { + ...message, + transcriptMessageId: restoredCandidates[matchIndex]!.transcriptMessageId, + } + }) + + return changed ? merged : messages +} + +function mergeRestoredHistoryIntoLiveMessages( + messages: UIMessage[], + restoredMessages: UIMessage[], +): UIMessage[] { + return mergeRestoredTerminalGoalEvents( + mergeRestoredTranscriptMessageIds(messages, restoredMessages), + restoredMessages, + ) +} + +function needsTranscriptIdHydrationRetry(session: PerSessionState | undefined): boolean { + if (!session || session.chatState !== 'idle') return false + + let currentTurnHasHydratedUser = false + for (const message of session.messages) { + if (message.type === 'user_text') { + currentTurnHasHydratedUser = Boolean(message.transcriptMessageId) + continue + } + if ( + currentTurnHasHydratedUser && + message.type === 'assistant_text' && + !message.transcriptMessageId + ) { + return true + } + } + + return false +} + +function refreshCompletedTranscriptHistory( + get: () => ChatStore, + sessionId: string, +): void { + void get().loadHistory(sessionId).then(() => { + if (!needsTranscriptIdHydrationRetry(get().sessions[sessionId])) return + setTimeout(() => { + if (!needsTranscriptIdHydrationRetry(get().sessions[sessionId])) return + void get().loadHistory(sessionId) + }, 750) + }) +} + function normalizeMemoryEventFiles(data: unknown): MemoryEventFile[] { if (!data || typeof data !== 'object') return [] const writtenPaths = (data as { writtenPaths?: unknown }).writtenPaths @@ -705,7 +802,7 @@ export const useChatStore = create((set, get) => ({ s.backgroundAgentTasks ?? {}, restoredBackgroundTasks, ), - messages: mergeRestoredTerminalGoalEvents( + messages: mergeRestoredHistoryIntoLiveMessages( mergeBackgroundTaskMessages(s.messages, restoredBackgroundTasks), uiMessages, ), @@ -1072,6 +1169,7 @@ export const useChatStore = create((set, get) => ({ target: { type: 'session', sessionId }, }) } + refreshCompletedTranscriptHistory(get, sessionId) break } @@ -1666,13 +1764,24 @@ function pushAssistantHistoryText( content: string, timestamp: number, model?: string, + transcriptMessageId?: string, ): void { if (!content.trim()) return const last = messages[messages.length - 1] - if (last?.type === 'assistant_text') { + const canMergeIntoLast = + last?.type === 'assistant_text' && + ( + transcriptMessageId + ? last.transcriptMessageId === transcriptMessageId + : !last.transcriptMessageId + ) + if (canMergeIntoLast) { last.content += content if (model && !last.model) last.model = model + if (transcriptMessageId && !last.transcriptMessageId) { + last.transcriptMessageId = transcriptMessageId + } return } @@ -1681,6 +1790,7 @@ function pushAssistantHistoryText( type: 'assistant_text', content, timestamp, + ...(transcriptMessageId ? { transcriptMessageId } : {}), ...(model ? { model } : {}), }) } @@ -1881,6 +1991,7 @@ export function mapHistoryMessagesToUiMessages( id: msg.id || nextId(), type: 'user_text', content: teammateContents.join('\n\n'), + ...(msg.id ? { transcriptMessageId: msg.id } : {}), timestamp, }) continue @@ -1890,6 +2001,7 @@ export function mapHistoryMessagesToUiMessages( id: msg.id || nextId(), type: 'user_text', content: parsed.content, + ...(msg.id ? { transcriptMessageId: msg.id } : {}), ...(parsed.modelContent ? { modelContent: parsed.modelContent } : {}), ...(parsed.attachments ? { attachments: parsed.attachments } : {}), timestamp, @@ -1898,13 +2010,22 @@ export function mapHistoryMessagesToUiMessages( } if (msg.type === 'assistant' && typeof msg.content === 'string') { if (!msg.content.trim()) continue - uiMessages.push({ id: msg.id || nextId(), type: 'assistant_text', content: msg.content, timestamp, model: msg.model }) + uiMessages.push({ + id: msg.id || nextId(), + type: 'assistant_text', + content: msg.content, + ...(msg.id ? { transcriptMessageId: msg.id } : {}), + timestamp, + model: msg.model, + }) continue } if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) { for (const block of msg.content as AssistantHistoryBlock[]) { if (block.type === 'thinking' && block.thinking) uiMessages.push({ id: nextId(), type: 'thinking', content: block.thinking, timestamp }) - else if (block.type === 'text' && block.text) pushAssistantHistoryText(uiMessages, block.text, timestamp, msg.model) + else if (block.type === 'text' && block.text) { + pushAssistantHistoryText(uiMessages, block.text, timestamp, msg.model, msg.id || undefined) + } else if (block.type === 'tool_use') uiMessages.push({ id: nextId(), type: 'tool_use', toolName: block.name ?? 'unknown', toolUseId: block.id ?? '', input: block.input, timestamp, parentToolUseId: msg.parentToolUseId }) } continue @@ -1930,6 +2051,7 @@ export function mapHistoryMessagesToUiMessages( id: msg.id || nextId(), type: 'user_text', content: parsed.content, + ...(msg.id ? { transcriptMessageId: msg.id } : {}), ...(parsed.modelContent ? { modelContent: parsed.modelContent } : {}), attachments: allAttachments.length > 0 ? allAttachments : undefined, timestamp, diff --git a/desktop/src/stores/sessionStore.test.ts b/desktop/src/stores/sessionStore.test.ts index d74dc9b9..1eb6fc8d 100644 --- a/desktop/src/stores/sessionStore.test.ts +++ b/desktop/src/stores/sessionStore.test.ts @@ -1,12 +1,14 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -const { createMock, listMock } = vi.hoisted(() => ({ +const { branchMock, createMock, listMock } = vi.hoisted(() => ({ + branchMock: vi.fn(), createMock: vi.fn(), listMock: vi.fn(), })) vi.mock('../api/sessions', () => ({ sessionsApi: { + branch: branchMock, create: createMock, list: listMock, delete: vi.fn(), @@ -33,6 +35,7 @@ function createDeferred() { describe('sessionStore', () => { beforeEach(() => { + branchMock.mockReset() createMock.mockReset() listMock.mockReset() useSessionStore.setState({ @@ -135,7 +138,7 @@ describe('sessionStore', () => { it('forwards direct branch switch repository options when creating a session', async () => { createMock.mockResolvedValue({ sessionId: 'session-branch-switch', workDir: '/workspace/repo' }) - listMock.mockResolvedValue({ sessions: [], total: 0 }) + listMock.mockImplementation(() => new Promise(() => {})) await useSessionStore.getState().createSession('/workspace/repo', { repository: { branch: 'feature/rail', worktree: false }, @@ -165,4 +168,100 @@ describe('sessionStore', () => { expect(useSessionStore.getState().sessions[0]?.workDir) .toBe('/workspace/repo/.claude/worktrees/desktop-feature-rail-12345678') }) + + it('returns the branched session before the background refresh completes', async () => { + branchMock.mockResolvedValue({ + sessionId: 'session-branch-1', + title: 'Branch from here', + workDir: '/workspace/repo/branches/session-branch-1', + sourceSessionId: 'session-source-1', + targetMessageId: 'transcript-message-1', + }) + listMock.mockImplementation(() => new Promise(() => {})) + useSessionStore.setState({ + sessions: [{ + id: 'session-source-1', + title: 'Source session', + createdAt: '2026-05-19T00:00:00.000Z', + modifiedAt: '2026-05-19T00:00:00.000Z', + messageCount: 4, + projectPath: '/workspace/repo', + projectRoot: '/workspace/repo', + workDir: '/workspace/repo', + workDirExists: true, + }], + }) + + const result = await Promise.race([ + useSessionStore.getState().branchSession('session-source-1', 'transcript-message-1'), + delay(100).then(() => 'timed-out'), + ]) + + expect(result).toMatchObject({ + sessionId: 'session-branch-1', + title: 'Branch from here', + workDir: '/workspace/repo/branches/session-branch-1', + }) + expect(branchMock).toHaveBeenCalledWith('session-source-1', { + targetMessageId: 'transcript-message-1', + }) + expect(useSessionStore.getState().activeSessionId).toBe('session-branch-1') + expect(useSessionStore.getState().sessions[0]).toMatchObject({ + id: 'session-branch-1', + title: 'Branch from here', + projectPath: '/workspace/repo', + workDir: '/workspace/repo/branches/session-branch-1', + projectRoot: '/workspace/repo', + workDirExists: true, + }) + expect(listMock).toHaveBeenCalledOnce() + }) + + it('updates an existing optimistic branch row when the branch session id is already present', async () => { + branchMock.mockResolvedValue({ + sessionId: 'session-branch-existing', + title: 'Updated branch', + workDir: '/workspace/repo/branches/session-branch-existing', + sourceSessionId: 'session-source-1', + targetMessageId: 'transcript-message-2', + }) + listMock.mockImplementation(() => new Promise(() => {})) + useSessionStore.setState({ + sessions: [ + { + id: 'session-branch-existing', + title: 'Old branch title', + createdAt: '2026-05-18T00:00:00.000Z', + modifiedAt: '2026-05-18T00:00:00.000Z', + messageCount: 3, + projectPath: '/workspace/old', + projectRoot: '/workspace/old', + workDir: '/workspace/old', + workDirExists: true, + }, + { + id: 'session-source-1', + title: 'Source session', + createdAt: '2026-05-19T00:00:00.000Z', + modifiedAt: '2026-05-19T00:00:00.000Z', + messageCount: 4, + projectPath: '/workspace/repo', + projectRoot: '/workspace/repo', + workDir: '/workspace/repo', + workDirExists: true, + }, + ], + }) + + await useSessionStore.getState().branchSession('session-source-1', 'transcript-message-2') + + expect(useSessionStore.getState().sessions).toHaveLength(2) + expect(useSessionStore.getState().sessions[0]).toMatchObject({ + id: 'session-branch-existing', + title: 'Updated branch', + projectPath: '/workspace/repo', + projectRoot: '/workspace/repo', + workDir: '/workspace/repo/branches/session-branch-existing', + }) + }) }) diff --git a/desktop/src/stores/sessionStore.ts b/desktop/src/stores/sessionStore.ts index 246d7ad7..e987bb04 100644 --- a/desktop/src/stores/sessionStore.ts +++ b/desktop/src/stores/sessionStore.ts @@ -1,5 +1,10 @@ import { create } from 'zustand' -import { sessionsApi, type BatchDeleteSessionsResponse, type CreateSessionRepositoryOptions } from '../api/sessions' +import { + sessionsApi, + type BatchDeleteSessionsResponse, + type BranchSessionResponse, + type CreateSessionRepositoryOptions, +} from '../api/sessions' import { useSessionRuntimeStore } from './sessionRuntimeStore' import { useTabStore } from './tabStore' import type { SessionListItem } from '../types/session' @@ -9,6 +14,8 @@ type CreateSessionOptions = { repository?: CreateSessionRepositoryOptions } +type BranchSessionResult = Pick + type SessionStore = { sessions: SessionListItem[] activeSessionId: string | null @@ -19,6 +26,11 @@ type SessionStore = { fetchSessions: (project?: string) => Promise createSession: (workDir?: string, options?: CreateSessionOptions) => Promise + branchSession: ( + sourceSessionId: string, + targetMessageId: string, + options?: { title?: string }, + ) => Promise deleteSession: (id: string) => Promise deleteSessions: (ids: string[]) => Promise enterBatchMode: () => void @@ -96,6 +108,43 @@ export const useSessionStore = create((set, get) => ({ return id }, + branchSession: async (sourceSessionId, targetMessageId, options) => { + const result = await sessionsApi.branch(sourceSessionId, { + targetMessageId, + ...(options?.title ? { title: options.title } : {}), + }) + const sourceSession = get().sessions.find((session) => session.id === sourceSessionId) + const now = new Date().toISOString() + const optimisticSession: SessionListItem = { + id: result.sessionId, + title: result.title || 'New Session', + createdAt: now, + modifiedAt: now, + messageCount: 0, + projectPath: sourceSession?.projectPath ?? '', + projectRoot: sourceSession?.projectRoot ?? sourceSession?.workDir ?? result.workDir ?? null, + workDir: result.workDir ?? sourceSession?.workDir ?? null, + workDirExists: true, + } + + set((state) => ({ + sessions: state.sessions.some((session) => session.id === result.sessionId) + ? state.sessions.map((session) => + session.id === result.sessionId + ? { ...session, ...optimisticSession } + : session) + : [optimisticSession, ...state.sessions], + activeSessionId: result.sessionId, + })) + + void get().fetchSessions() + return { + sessionId: result.sessionId, + title: result.title, + workDir: result.workDir, + } + }, + deleteSession: async (id: string) => { await sessionsApi.delete(id) useSessionRuntimeStore.getState().clearSelection(id) diff --git a/desktop/src/types/chat.ts b/desktop/src/types/chat.ts index 8503dd65..dd23e714 100644 --- a/desktop/src/types/chat.ts +++ b/desktop/src/types/chat.ts @@ -212,8 +212,8 @@ export type TaskSummaryItem = { } export type UIMessage = - | { id: string; type: 'user_text'; content: string; modelContent?: string; timestamp: number; attachments?: UIAttachment[]; pending?: boolean } - | { id: string; type: 'assistant_text'; content: string; timestamp: number; model?: string } + | { id: string; type: 'user_text'; content: string; modelContent?: string; transcriptMessageId?: string; timestamp: number; attachments?: UIAttachment[]; pending?: boolean } + | { id: string; type: 'assistant_text'; content: string; transcriptMessageId?: string; timestamp: number; model?: string } | { id: string; type: 'thinking'; content: string; timestamp: number } | { id: string; type: 'tool_use'; toolName: string; toolUseId: string; input: unknown; timestamp: number; parentToolUseId?: string } | { id: string; type: 'tool_result'; toolUseId: string; content: unknown; isError: boolean; timestamp: number; parentToolUseId?: string } diff --git a/src/commands/branch/branch.ts b/src/commands/branch/branch.ts index 4a7c2779..5a930714 100644 --- a/src/commands/branch/branch.ts +++ b/src/commands/branch/branch.ts @@ -1,223 +1,13 @@ -import { randomUUID, type UUID } from 'crypto' -import { mkdir, readFile, writeFile } from 'fs/promises' import { getOriginalCwd, getSessionId } from '../../bootstrap/state.js' import type { LocalJSXCommandContext } from '../../commands.js' import { logEvent } from '../../services/analytics/index.js' import type { LocalJSXCommandOnDone } from '../../types/command.js' -import type { - ContentReplacementEntry, - Entry, - LogOption, - SerializedMessage, - TranscriptMessage, -} from '../../types/logs.js' -import { parseJSONL } from '../../utils/json.js' +import type { LogOption } from '../../types/logs.js' +import { getTranscriptPath } from '../../utils/sessionStorage.js' import { - getProjectDir, - getTranscriptPath, - getTranscriptPathForSession, - isTranscriptMessage, - saveCustomTitle, - searchSessionsByCustomTitle, -} from '../../utils/sessionStorage.js' -import { jsonStringify } from '../../utils/slowOperations.js' -import { escapeRegExp } from '../../utils/stringUtils.js' - -type TranscriptEntry = TranscriptMessage & { - forkedFrom?: { - sessionId: string - messageUuid: UUID - } -} - -/** - * Derive a single-line title base from the first user message. - * Collapses whitespace — multiline first messages (pasted stacks, code) - * otherwise flow into the saved title and break the resume hint. - */ -export function deriveFirstPrompt( - firstUserMessage: Extract | undefined, -): string { - const content = firstUserMessage?.message?.content - if (!content) return 'Branched conversation' - const raw = - typeof content === 'string' - ? content - : content.find( - (block): block is { type: 'text'; text: string } => - block.type === 'text', - )?.text - if (!raw) return 'Branched conversation' - return ( - raw.replace(/\s+/g, ' ').trim().slice(0, 100) || 'Branched conversation' - ) -} - -/** - * Creates a fork of the current conversation by copying from the transcript file. - * Preserves all original metadata (timestamps, gitBranch, etc.) while updating - * sessionId and adding forkedFrom traceability. - */ -async function createFork(customTitle?: string): Promise<{ - sessionId: UUID - title: string | undefined - forkPath: string - serializedMessages: SerializedMessage[] - contentReplacementRecords: ContentReplacementEntry['replacements'] -}> { - const forkSessionId = randomUUID() as UUID - const originalSessionId = getSessionId() - const projectDir = getProjectDir(getOriginalCwd()) - const forkSessionPath = getTranscriptPathForSession(forkSessionId) - const currentTranscriptPath = getTranscriptPath() - - // Ensure project directory exists - await mkdir(projectDir, { recursive: true, mode: 0o700 }) - - // Read current transcript file - let transcriptContent: Buffer - try { - transcriptContent = await readFile(currentTranscriptPath) - } catch { - throw new Error('No conversation to branch') - } - - if (transcriptContent.length === 0) { - throw new Error('No conversation to branch') - } - - // Parse all transcript entries (messages + metadata entries like content-replacement) - const entries = parseJSONL(transcriptContent) - - // Filter to only main conversation messages (exclude sidechains and non-message entries) - const mainConversationEntries = entries.filter( - (entry): entry is TranscriptMessage => - isTranscriptMessage(entry) && !entry.isSidechain, - ) - - // Content-replacement entries for the original session. These record which - // tool_result blocks were replaced with previews by the per-message budget. - // Without them in the fork JSONL, `claude -r {forkId}` reconstructs state - // with an empty replacements Map → previously-replaced results are classified - // as FROZEN and sent as full content (prompt cache miss + permanent overage). - // sessionId must be rewritten since loadTranscriptFile keys lookup by the - // session's messages' sessionId. - const contentReplacementRecords = entries - .filter( - (entry): entry is ContentReplacementEntry => - entry.type === 'content-replacement' && - entry.sessionId === originalSessionId, - ) - .flatMap(entry => entry.replacements) - - if (mainConversationEntries.length === 0) { - throw new Error('No messages to branch') - } - - // Build forked entries with new sessionId and preserved metadata - let parentUuid: UUID | null = null - const lines: string[] = [] - const serializedMessages: SerializedMessage[] = [] - - for (const entry of mainConversationEntries) { - // Create forked transcript entry preserving all original metadata - const forkedEntry: TranscriptEntry = { - ...entry, - sessionId: forkSessionId, - parentUuid, - isSidechain: false, - forkedFrom: { - sessionId: originalSessionId, - messageUuid: entry.uuid, - }, - } - - // Build serialized message for LogOption - const serialized: SerializedMessage = { - ...entry, - sessionId: forkSessionId, - } - - serializedMessages.push(serialized) - lines.push(jsonStringify(forkedEntry)) - if (entry.type !== 'progress') { - parentUuid = entry.uuid - } - } - - // Append content-replacement entry (if any) with the fork's sessionId. - // Written as a SINGLE entry (same shape as insertContentReplacement) so - // loadTranscriptFile's content-replacement branch picks it up. - if (contentReplacementRecords.length > 0) { - const forkedReplacementEntry: ContentReplacementEntry = { - type: 'content-replacement', - sessionId: forkSessionId, - replacements: contentReplacementRecords, - } - lines.push(jsonStringify(forkedReplacementEntry)) - } - - // Write the fork session file - await writeFile(forkSessionPath, lines.join('\n') + '\n', { - encoding: 'utf8', - mode: 0o600, - }) - - return { - sessionId: forkSessionId, - title: customTitle, - forkPath: forkSessionPath, - serializedMessages, - contentReplacementRecords, - } -} - -/** - * Generates a unique fork name by checking for collisions with existing session names. - * If "baseName (Branch)" already exists, tries "baseName (Branch 2)", "baseName (Branch 3)", etc. - */ -async function getUniqueForkName(baseName: string): Promise { - const candidateName = `${baseName} (Branch)` - - // Check if this exact name already exists - const existingWithExactName = await searchSessionsByCustomTitle( - candidateName, - { exact: true }, - ) - - if (existingWithExactName.length === 0) { - return candidateName - } - - // Name collision - find a unique numbered suffix - // Search for all sessions that start with the base pattern - const existingForks = await searchSessionsByCustomTitle(`${baseName} (Branch`) - - // Extract existing fork numbers to find the next available - const usedNumbers = new Set([1]) // Consider " (Branch)" as number 1 - const forkNumberPattern = new RegExp( - `^${escapeRegExp(baseName)} \\(Branch(?: (\\d+))?\\)$`, - ) - - for (const session of existingForks) { - const match = session.customTitle?.match(forkNumberPattern) - if (match) { - if (match[1]) { - usedNumbers.add(parseInt(match[1], 10)) - } else { - usedNumbers.add(1) // " (Branch)" without number is treated as 1 - } - } - } - - // Find the next available number - let nextNumber = 2 - while (usedNumbers.has(nextNumber)) { - nextNumber++ - } - - return `${baseName} (Branch ${nextNumber})` -} + createSessionBranch, + deriveFirstPrompt, +} from '../../utils/sessionBranching.js' export async function call( onDone: LocalJSXCommandOnDone, @@ -225,7 +15,6 @@ export async function call( args: string, ): Promise { const customTitle = args?.trim() || undefined - const originalSessionId = getSessionId() try { @@ -233,29 +22,23 @@ export async function call( sessionId, title, forkPath, + workDir, serializedMessages, contentReplacementRecords, - } = await createFork(customTitle) - - // Build LogOption for resume - const now = new Date() - const firstPrompt = deriveFirstPrompt( - serializedMessages.find(m => m.type === 'user'), - ) - - // Save custom title - use provided title or firstPrompt as default - // This ensures /status and /resume show the same session name - // Always add " (Branch)" suffix to make it clear this is a branched session - // Handle collisions by adding a number suffix (e.g., " (Branch 2)", " (Branch 3)") - const baseName = title ?? firstPrompt - const effectiveTitle = await getUniqueForkName(baseName) - await saveCustomTitle(sessionId, effectiveTitle, forkPath) - - logEvent('tengu_conversation_forked', { - message_count: serializedMessages.length, - has_custom_title: !!title, + } = await createSessionBranch({ + sourceSessionId: originalSessionId, + sourceTranscriptPath: getTranscriptPath(), + title: customTitle, + sourceWorkDir: getOriginalCwd(), }) + const now = new Date() + const firstPrompt = deriveFirstPrompt( + serializedMessages.find( + (message): message is Extract => + message.type === 'user', + ), + ) const forkLog: LogOption = { date: now.toISOString().split('T')[0]!, messages: serializedMessages, @@ -267,12 +50,16 @@ export async function call( messageCount: serializedMessages.length, isSidechain: false, sessionId, - customTitle: effectiveTitle, + customTitle: title, contentReplacements: contentReplacementRecords, } - // Resume into the fork - const titleInfo = title ? ` "${title}"` : '' + logEvent('tengu_conversation_forked', { + message_count: serializedMessages.length, + has_custom_title: !!customTitle, + }) + + const titleInfo = customTitle ? ` "${customTitle}"` : '' const resumeHint = `\nTo resume the original: claude -r ${originalSessionId}` const successMessage = `Branched conversation${titleInfo}. You are now in the branch.${resumeHint}` @@ -280,7 +67,6 @@ export async function call( await context.resume(sessionId, forkLog, 'fork') onDone(successMessage, { display: 'system' }) } else { - // Fallback if resume not available onDone( `Branched conversation${titleInfo}. Resume with: /resume ${sessionId}`, ) diff --git a/src/server/__tests__/sessions.test.ts b/src/server/__tests__/sessions.test.ts index 27ab6deb..8b341590 100644 --- a/src/server/__tests__/sessions.test.ts +++ b/src/server/__tests__/sessions.test.ts @@ -14,6 +14,8 @@ import { } from '../services/repositoryLaunchService.js' import { conversationService } from '../services/conversationService.js' import { clearCommandsCache } from '../../commands.js' +import { parseJSONL } from '../../utils/json.js' +import { createSessionBranch } from '../../utils/sessionBranching.js' import { sanitizePath } from '../../utils/sessionStoragePortable.js' import { clearInstalledPluginsCache } from '../../utils/plugins/installedPluginsManager.js' import { clearPluginCache } from '../../utils/plugins/pluginLoader.js' @@ -265,6 +267,33 @@ function makeAssistantToolUseEntry( } } +function makeToolResultUserEntry( + toolUseId: string, + content: string, + uuid?: string, + parentUuid?: string, + sessionId = 'test-session', +): Record { + return { + parentUuid: parentUuid || null, + isSidechain: false, + type: 'user', + message: { + role: 'user', + content: [{ + type: 'tool_result', + tool_use_id: toolUseId, + content, + }], + }, + uuid: uuid || crypto.randomUUID(), + timestamp: '2026-01-01T00:02:30.000Z', + userType: 'external', + cwd: '/tmp/test', + sessionId, + } +} + function makeMetaUserEntry(): Record { return { parentUuid: null, @@ -306,6 +335,17 @@ function makeWorktreeStateEntry( } } +function makeContentReplacementEntry( + sessionId: string, + replacements: Array<{ kind: 'tool-result'; toolUseId: string; replacement: string }>, +): Record { + return { + type: 'content-replacement', + sessionId, + replacements, + } +} + async function writeFileHistoryBackup( sessionId: string, backupFileName: string, @@ -1564,6 +1604,154 @@ describe('SessionService', () => { expect(launchInfo!.workDir).toBe(expectedWorkDir) expect(launchInfo!.transcriptMessageCount).toBe(2) }) + + it('createSessionBranch should preserve branch metadata, copied snapshots, and filtered replacements', async () => { + const sessionId = 'branch-source-session' + const workDir = path.join(tmpDir, 'branch-source') + const worktreePath = path.join(workDir, '.claude', 'worktrees', 'desktop-main-12345678') + const firstUserId = crypto.randomUUID() + const firstAssistantId = crypto.randomUUID() + const firstToolResultId = crypto.randomUUID() + const laterUserId = crypto.randomUUID() + const laterAssistantId = crypto.randomUUID() + const repository = { + branch: 'feature/rail', + worktree: true, + baseRef: 'feature/rail', + repoRoot: workDir, + } + const sourceProjectDir = sanitizePath(workDir) + const sourcePath = await writeSessionFile(sourceProjectDir, sessionId, [ + makeSessionMetaEntry(workDir), + { + type: 'session-meta', + isMeta: true, + workDir, + repository, + timestamp: '2026-01-01T00:00:00.000Z', + }, + makeWorktreeStateEntry(sessionId, worktreePath, { + originalCwd: workDir, + }), + makeFileHistorySnapshotEntry(firstUserId, { + 'src/step.js': { + backupFileName: 'branch-source-step@v1', + version: 1, + backupTime: '2026-01-01T00:00:00.000Z', + }, + }), + { + ...makeUserEntry('branch this conversation', firstUserId), + cwd: workDir, + sessionId, + }, + { + ...makeAssistantToolUseEntry([ + { id: 'tool-1', name: 'Read', input: { path: 'src/step.js' } }, + ], firstUserId), + uuid: firstAssistantId, + cwd: workDir, + sessionId, + }, + { + ...makeToolResultUserEntry('tool-1', 'first tool result', firstToolResultId, firstAssistantId, sessionId), + cwd: workDir, + }, + makeContentReplacementEntry(sessionId, [ + { kind: 'tool-result', toolUseId: 'tool-1', replacement: 'preview-1' }, + { kind: 'tool-result', toolUseId: 'tool-2', replacement: 'preview-2' }, + ]), + makeFileHistorySnapshotEntry(laterUserId, { + 'src/step.js': { + backupFileName: 'branch-source-step@v2', + version: 2, + backupTime: '2026-01-01T00:00:00.000Z', + }, + }), + { + ...makeUserEntry('later prompt', laterUserId), + parentUuid: firstToolResultId, + cwd: workDir, + sessionId, + }, + { + ...makeAssistantEntry('later reply', laterUserId), + uuid: laterAssistantId, + cwd: workDir, + sessionId, + }, + ]) + + const sourceBefore = await fs.readFile(sourcePath, 'utf-8') + + const branch = await createSessionBranch({ + sourceSessionId: sessionId, + sourceTranscriptPath: sourcePath, + targetMessageId: firstToolResultId, + title: 'Desktop branch', + sourceWorkDir: workDir, + sourceRepository: repository, + sourceWorktreeSession: { + originalCwd: workDir, + worktreePath, + worktreeName: 'desktop-main-12345678', + worktreeBranch: 'worktree-desktop-main-12345678', + originalBranch: 'main', + sessionId, + }, + }) + + const branchMessages = await service.getSessionMessages(branch.sessionId) + expect(branchMessages.map((message) => message.id)).toEqual([ + firstUserId, + firstAssistantId, + firstToolResultId, + ]) + expect(branch.title).toBe('Desktop branch (Branch)') + + const launchInfo = await service.getSessionLaunchInfo(branch.sessionId) + expect(launchInfo).toMatchObject({ + workDir, + repository, + worktreeSession: { + originalCwd: workDir, + worktreePath, + }, + }) + + const branchEntries = parseJSONL>(await fs.readFile(branch.forkPath)) + expect(branchEntries.some((entry) => ( + entry.type === 'content-replacement' && + entry.sessionId === branch.sessionId && + Array.isArray(entry.replacements) && + entry.replacements.length === 1 && + (entry.replacements[0] as { toolUseId?: string }).toolUseId === 'tool-1' + ))).toBe(true) + expect(branchEntries.some((entry) => ( + entry.type === 'file-history-snapshot' && + typeof (entry.snapshot as { messageId?: string } | undefined)?.messageId === 'string' && + (entry.snapshot as { messageId?: string }).messageId === firstUserId + ))).toBe(true) + expect(branchEntries.some((entry) => ( + entry.type === 'file-history-snapshot' && + typeof (entry.snapshot as { messageId?: string } | undefined)?.messageId === 'string' && + (entry.snapshot as { messageId?: string }).messageId === laterUserId + ))).toBe(false) + expect(branchEntries.some((entry) => ( + entry.type === 'custom-title' && + entry.customTitle === 'Desktop branch (Branch)' + ))).toBe(true) + expect(branchEntries.filter((entry) => ( + entry.type === 'user' || + entry.type === 'assistant' + )).every((entry) => ( + entry.sessionId === branch.sessionId && + typeof (entry.forkedFrom as { sessionId?: string } | undefined)?.sessionId === 'string' + ))).toBe(true) + + const sourceAfter = await fs.readFile(sourcePath, 'utf-8') + expect(sourceAfter).toBe(sourceBefore) + }) }) // ============================================================================ @@ -2731,6 +2919,147 @@ describe('Sessions API', () => { }) }) + it('POST /api/sessions/:id/branch should create a branched session up to the target message', async () => { + const sessionId = '11111111-1111-4111-8111-111111111111' + const workDir = path.join(tmpDir, 'branch-api-workdir') + const firstUserId = crypto.randomUUID() + const firstAssistantId = crypto.randomUUID() + const secondUserId = crypto.randomUUID() + const secondAssistantId = crypto.randomUUID() + + await writeSessionFile(sanitizePath(workDir), sessionId, [ + { + type: 'session-meta', + isMeta: true, + workDir, + timestamp: '2026-01-01T00:00:00.000Z', + }, + { + ...makeUserEntry('first prompt', firstUserId), + cwd: workDir, + sessionId, + }, + { + ...makeAssistantEntry('first reply', firstUserId), + uuid: firstAssistantId, + cwd: workDir, + sessionId, + }, + { + ...makeUserEntry('second prompt', secondUserId), + parentUuid: firstAssistantId, + cwd: workDir, + sessionId, + }, + { + ...makeAssistantEntry('second reply', secondUserId), + uuid: secondAssistantId, + cwd: workDir, + sessionId, + }, + ]) + + const res = await fetch(`${baseUrl}/api/sessions/${sessionId}/branch`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + targetMessageId: firstAssistantId, + title: 'API branch', + }), + }) + expect(res.status).toBe(201) + + const body = await res.json() as { + sessionId: string + title: string + workDir: string + sourceSessionId: string + targetMessageId: string + } + expect(body).toMatchObject({ + title: 'API branch (Branch)', + workDir, + sourceSessionId: sessionId, + targetMessageId: firstAssistantId, + }) + + const branchMessages = await service.getSessionMessages(body.sessionId) + expect(branchMessages.map((message) => message.id)).toEqual([ + firstUserId, + firstAssistantId, + ]) + }) + + it('POST /api/sessions/:id/branch should reject sidechain targets', async () => { + const sessionId = '22222222-2222-4222-8222-222222222222' + const rootUserId = crypto.randomUUID() + const rootAssistantId = crypto.randomUUID() + const sidechainId = crypto.randomUUID() + + await writeSessionFile('-tmp-api-branch-sidechain', sessionId, [ + makeSnapshotEntry(), + { + ...makeUserEntry('root prompt', rootUserId), + sessionId, + }, + { + ...makeAssistantEntry('root reply', rootUserId), + uuid: rootAssistantId, + sessionId, + }, + { + ...makeUserEntry('side question', sidechainId), + parentUuid: rootAssistantId, + isSidechain: true, + sessionId, + }, + ]) + + const res = await fetch(`${baseUrl}/api/sessions/${sessionId}/branch`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ targetMessageId: sidechainId }), + }) + + expect(res.status).toBe(400) + expect(await res.json()).toMatchObject({ + error: 'BAD_REQUEST', + }) + }) + + it('POST /api/sessions/:id/branch should validate request bodies and missing sessions', async () => { + const methodNotAllowedRes = await fetch(`${baseUrl}/api/sessions/33333333-3333-4333-8333-333333333333/branch`) + expect(methodNotAllowedRes.status).toBe(405) + + const missingTargetRes = await fetch(`${baseUrl}/api/sessions/branch-missing-target/branch`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({}), + }) + expect(missingTargetRes.status).toBe(400) + + const invalidJsonRes = await fetch(`${baseUrl}/api/sessions/branch-invalid-json/branch`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: '{', + }) + expect(invalidJsonRes.status).toBe(400) + + const invalidTitleRes = await fetch(`${baseUrl}/api/sessions/44444444-4444-4444-8444-444444444444/branch`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ targetMessageId: 'message-1', title: 123 }), + }) + expect(invalidTitleRes.status).toBe(400) + + const missingSessionRes = await fetch(`${baseUrl}/api/sessions/00000000-0000-0000-0000-000000000000/branch`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ targetMessageId: 'missing-target' }), + }) + expect(missingSessionRes.status).toBe(404) + }) + it('POST /api/sessions/:id/rewind should preview and trim the active conversation chain', async () => { const sessionId = 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee' const firstUserId = crypto.randomUUID() diff --git a/src/server/api/sessions.ts b/src/server/api/sessions.ts index d431d1c0..51d9e375 100644 --- a/src/server/api/sessions.ts +++ b/src/server/api/sessions.ts @@ -33,6 +33,10 @@ import { type RewindTargetSelector, } from '../services/sessionRewindService.js' import { SessionStore } from '../../../adapters/common/session-store.js' +import { + createSessionBranch, + SessionBranchingError, +} from '../../utils/sessionBranching.js' const workspaceService = new WorkspaceService( async (sessionId) => ( @@ -124,6 +128,16 @@ export async function handleSessionsApi( return await rewindSession(req, sessionId) } + if (subResource === 'branch') { + if (req.method !== 'POST') { + return Response.json( + { error: 'METHOD_NOT_ALLOWED', message: `Method ${req.method} not allowed` }, + { status: 405 } + ) + } + return await branchSession(req, sessionId) + } + if (subResource === 'turn-checkpoints') { if (req.method !== 'GET') { return Response.json( @@ -710,6 +724,56 @@ async function rewindSession(req: Request, sessionId: string): Promise return Response.json(result) } +async function branchSession(req: Request, sessionId: string): Promise { + let body: { targetMessageId?: unknown; title?: unknown } + try { + body = (await req.json()) as { targetMessageId?: unknown; title?: unknown } + } catch { + throw ApiError.badRequest('Invalid JSON body') + } + + if (typeof body.targetMessageId !== 'string' || body.targetMessageId.trim().length === 0) { + throw ApiError.badRequest('targetMessageId (string) is required in request body') + } + + if (body.title !== undefined && typeof body.title !== 'string') { + throw ApiError.badRequest('title must be a string') + } + + const launchInfo = await sessionService.getSessionLaunchInfo(sessionId) + if (!launchInfo) { + throw ApiError.notFound(`Session not found: ${sessionId}`) + } + + try { + const result = await createSessionBranch({ + sourceSessionId: sessionId, + sourceTranscriptPath: launchInfo.filePath, + targetMessageId: body.targetMessageId.trim(), + title: body.title?.trim() || undefined, + sourceWorkDir: launchInfo.workDir, + sourceRepository: launchInfo.repository, + sourceWorktreeSession: launchInfo.worktreeSession, + }) + + return Response.json({ + sessionId: result.sessionId, + title: result.title, + workDir: result.workDir ?? launchInfo.workDir, + sourceSessionId: sessionId, + targetMessageId: body.targetMessageId.trim(), + }, { status: 201 }) + } catch (error) { + if (error instanceof SessionBranchingError) { + if (error.code === 'SOURCE_NOT_FOUND') { + throw ApiError.notFound(error.message) + } + throw ApiError.badRequest(error.message) + } + throw error + } +} + async function getTurnCheckpoints(sessionId: string): Promise { const checkpoints = await listSessionTurnCheckpoints(sessionId) return Response.json({ checkpoints }) diff --git a/src/utils/__tests__/sessionBranching.test.ts b/src/utils/__tests__/sessionBranching.test.ts new file mode 100644 index 00000000..7c926e13 --- /dev/null +++ b/src/utils/__tests__/sessionBranching.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from 'bun:test' +import type { SerializedMessage } from '../../types/logs.js' +import { deriveFirstPrompt, SessionBranchingError } from '../sessionBranching.js' + +describe('sessionBranching', () => { + it('derives a compact branch title from multiline user text', () => { + const message = { + type: 'user', + message: { + role: 'user', + content: ' first line\n\nsecond line ', + }, + } as SerializedMessage + + expect(deriveFirstPrompt(message)).toBe('first line second line') + }) + + it('keeps branching errors identifiable by code', () => { + const error = new SessionBranchingError( + 'INVALID_TARGET', + 'targetMessageId must reference a main conversation message', + ) + + expect(error).toBeInstanceOf(Error) + expect(error.name).toBe('SessionBranchingError') + expect(error.code).toBe('INVALID_TARGET') + }) +}) diff --git a/src/utils/sessionBranching.ts b/src/utils/sessionBranching.ts new file mode 100644 index 00000000..8fbc3c51 --- /dev/null +++ b/src/utils/sessionBranching.ts @@ -0,0 +1,554 @@ +import { randomUUID, type UUID } from 'crypto' +import { mkdir, readFile, readdir, writeFile } from 'fs/promises' +import * as path from 'node:path' +import type { + AttributionSnapshotMessage, + PersistedWorktreeSession, + SerializedMessage, + TranscriptMessage, +} from '../types/logs.js' +import type { ContentReplacementRecord } from './toolResultStorage.js' +import { parseJSONL } from './json.js' +import { buildConversationChain, loadTranscriptFile } from './sessionStorage.js' +import { jsonStringify } from './slowOperations.js' +import { escapeRegExp } from './stringUtils.js' + +type SessionMetaEntry = { + type: 'session-meta' + isMeta?: boolean + workDir?: string + repository?: unknown + timestamp?: string + [key: string]: unknown +} + +type WorktreeStateEntry = { + type: 'worktree-state' + sessionId: string + worktreeSession: PersistedWorktreeSession | null + [key: string]: unknown +} + +type ModeEntry = { + type: 'mode' + sessionId: string + mode: string + [key: string]: unknown +} + +type PrLinkEntry = { + type: 'pr-link' + sessionId: string + prNumber: number + prUrl: string + prRepository: string + timestamp: string + [key: string]: unknown +} + +type FileHistorySnapshotEntry = { + type: 'file-history-snapshot' + messageId: string + snapshot?: { + messageId?: string + trackedFileBackups?: Record + timestamp?: string + } + [key: string]: unknown +} + +type ContentReplacementEntry = { + type: 'content-replacement' + sessionId: string + replacements: ContentReplacementRecord[] + [key: string]: unknown +} + +type TranscriptEntry = TranscriptMessage & { + forkedFrom?: { + sessionId: string + messageUuid: UUID + } +} + +type RawEntry = + | SessionMetaEntry + | WorktreeStateEntry + | ModeEntry + | PrLinkEntry + | FileHistorySnapshotEntry + | AttributionSnapshotMessage + | ContentReplacementEntry + | Record + +export type SessionBranchResult = { + sessionId: UUID + title: string + forkPath: string + workDir: string | null + serializedMessages: SerializedMessage[] + contentReplacementRecords: ContentReplacementRecord[] +} + +export type CreateSessionBranchOptions = { + sourceSessionId: string + sourceTranscriptPath: string + title?: string + targetMessageId?: string + sourceWorkDir?: string | null + sourceRepository?: unknown + sourceWorktreeSession?: PersistedWorktreeSession | null +} + +export class SessionBranchingError extends Error { + constructor( + readonly code: + | 'SOURCE_NOT_FOUND' + | 'INVALID_TARGET' + | 'NO_BRANCHABLE_MESSAGES', + message: string, + ) { + super(message) + this.name = 'SessionBranchingError' + } +} + +function isTranscriptEntry(entry: RawEntry): entry is TranscriptMessage { + return ( + typeof entry === 'object' && + entry !== null && + (entry.type === 'user' || + entry.type === 'assistant' || + entry.type === 'attachment' || + entry.type === 'system') && + typeof (entry as TranscriptMessage).uuid === 'string' + ) +} + +function isSessionMetaEntry(entry: RawEntry): entry is SessionMetaEntry { + return entry.type === 'session-meta' +} + +function isWorktreeStateEntry(entry: RawEntry): entry is WorktreeStateEntry { + return entry.type === 'worktree-state' +} + +function isModeEntry(entry: RawEntry): entry is ModeEntry { + return entry.type === 'mode' +} + +function isPrLinkEntry(entry: RawEntry): entry is PrLinkEntry { + return entry.type === 'pr-link' +} + +function isFileHistorySnapshotEntry(entry: RawEntry): entry is FileHistorySnapshotEntry { + return entry.type === 'file-history-snapshot' +} + +function isAttributionSnapshotEntry(entry: RawEntry): entry is AttributionSnapshotMessage { + return entry.type === 'attribution-snapshot' +} + +function getSnapshotMessageId(entry: FileHistorySnapshotEntry): string | null { + if (typeof entry.snapshot?.messageId === 'string') return entry.snapshot.messageId + if (typeof entry.messageId === 'string') return entry.messageId + return null +} + +function getAttributionMessageId(entry: AttributionSnapshotMessage): string | null { + return typeof entry.messageId === 'string' ? entry.messageId : null +} + +function findLatestActiveLeaf( + messages: Map, + leafUuids: Set, +): TranscriptMessage | undefined { + let latest: TranscriptMessage | undefined + let latestTimestamp = -Infinity + + const candidates = leafUuids.size > 0 + ? [...leafUuids] + .map((uuid) => messages.get(uuid)) + .filter((message): message is TranscriptMessage => !!message) + : [...messages.values()] + + for (const message of candidates) { + if (message.isSidechain) continue + const timestamp = Date.parse(message.timestamp) + if (timestamp >= latestTimestamp) { + latest = message + latestTimestamp = timestamp + } + } + + return latest +} + +function extractToolResultIds(message: TranscriptMessage): string[] { + if (message.type !== 'user' || !Array.isArray(message.message.content)) { + return [] + } + + return message.message.content.flatMap((block) => ( + block.type === 'tool_result' && typeof block.tool_use_id === 'string' + ? [block.tool_use_id] + : [] + )) +} + +async function readRawEntries(filePath: string): Promise { + const content = await readFile(filePath) + return parseJSONL(content) +} + +async function listExistingTitles(projectDirPath: string): Promise { + let dirents: import('fs').Dirent[] + try { + dirents = await readdir(projectDirPath, { withFileTypes: true }) + } catch { + return [] + } + + const titles: string[] = [] + for (const dirent of dirents) { + if (!dirent.isFile() || !dirent.name.endsWith('.jsonl')) continue + try { + const entries = await readRawEntries(path.join(projectDirPath, dirent.name)) + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i] + if ( + entry?.type === 'custom-title' && + typeof (entry as Record).customTitle === 'string' + ) { + titles.push((entry as Record).customTitle as string) + break + } + } + } catch { + continue + } + } + + return titles +} + +export function deriveFirstPrompt( + firstUserMessage: Extract | undefined, +): string { + const content = firstUserMessage?.message?.content + if (!content) return 'Branched conversation' + const raw = + typeof content === 'string' + ? content + : content.find( + (block): block is { type: 'text'; text: string } => + block.type === 'text', + )?.text + if (!raw) return 'Branched conversation' + return ( + raw.replace(/\s+/g, ' ').trim().slice(0, 100) || 'Branched conversation' + ) +} + +export async function getUniqueForkName( + baseName: string, + projectDirPath: string, +): Promise { + const existingTitles = new Set( + (await listExistingTitles(projectDirPath)) + .map((title) => title.trim()) + .filter(Boolean), + ) + + const candidateName = `${baseName} (Branch)` + if (!existingTitles.has(candidateName)) { + return candidateName + } + + const usedNumbers = new Set([1]) + const forkNumberPattern = new RegExp( + `^${escapeRegExp(baseName)} \\(Branch(?: (\\d+))?\\)$`, + ) + + for (const title of existingTitles) { + const match = title.match(forkNumberPattern) + if (!match) continue + usedNumbers.add(match[1] ? parseInt(match[1], 10) : 1) + } + + let nextNumber = 2 + while (usedNumbers.has(nextNumber)) { + nextNumber++ + } + return `${baseName} (Branch ${nextNumber})` +} + +function buildPreservedMetadataEntries( + sourceEntries: RawEntry[], + copiedMessageIds: Set, + forkSessionId: UUID, +): RawEntry[] { + return sourceEntries.flatMap((entry) => { + if (isSessionMetaEntry(entry)) { + return [entry] + } + + if (isWorktreeStateEntry(entry) && entry.worktreeSession) { + return [{ + ...entry, + sessionId: forkSessionId, + worktreeSession: { + ...entry.worktreeSession, + sessionId: forkSessionId, + }, + }] + } + + if (isModeEntry(entry)) { + return [{ ...entry, sessionId: forkSessionId }] + } + + if (isPrLinkEntry(entry)) { + return [{ ...entry, sessionId: forkSessionId }] + } + + if (isFileHistorySnapshotEntry(entry)) { + const messageId = getSnapshotMessageId(entry) + return messageId && copiedMessageIds.has(messageId) ? [entry] : [] + } + + if (isAttributionSnapshotEntry(entry)) { + const messageId = getAttributionMessageId(entry) + return messageId && copiedMessageIds.has(messageId) ? [entry] : [] + } + + return [] + }) +} + +function ensureSyntheticSessionMeta( + preservedMetadataEntries: RawEntry[], + sourceWorkDir: string | null | undefined, + sourceRepository: unknown, +): RawEntry[] { + const hasSessionMeta = preservedMetadataEntries.some(isSessionMetaEntry) + if (hasSessionMeta || (!sourceWorkDir && sourceRepository === undefined)) { + return preservedMetadataEntries + } + + return [{ + type: 'session-meta', + isMeta: true, + ...(sourceWorkDir ? { workDir: sourceWorkDir } : {}), + ...(sourceRepository !== undefined ? { repository: sourceRepository } : {}), + timestamp: new Date().toISOString(), + }, ...preservedMetadataEntries] +} + +function ensureSyntheticWorktreeState( + preservedMetadataEntries: RawEntry[], + forkSessionId: UUID, + sourceWorktreeSession: PersistedWorktreeSession | null | undefined, +): RawEntry[] { + const hasWorktreeState = preservedMetadataEntries.some(isWorktreeStateEntry) + if (hasWorktreeState || !sourceWorktreeSession) { + return preservedMetadataEntries + } + + return [{ + type: 'worktree-state', + sessionId: forkSessionId, + worktreeSession: { + ...sourceWorktreeSession, + sessionId: forkSessionId, + }, + }, ...preservedMetadataEntries] +} + +export async function createSessionBranch( + options: CreateSessionBranchOptions, +): Promise { + const { + sourceSessionId, + sourceTranscriptPath, + title, + targetMessageId, + sourceWorkDir, + sourceRepository, + sourceWorktreeSession, + } = options + + const projectDirPath = path.dirname(sourceTranscriptPath) + const sourceEntries = await readRawEntries(sourceTranscriptPath).catch((error) => { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') { + throw new SessionBranchingError( + 'SOURCE_NOT_FOUND', + `Session not found: ${sourceSessionId}`, + ) + } + throw error + }) + + const transcript = await loadTranscriptFile(sourceTranscriptPath, { + keepAllLeaves: true, + }) + const activeLeaf = findLatestActiveLeaf(transcript.messages, transcript.leafUuids) + if (!activeLeaf) { + throw new SessionBranchingError( + 'NO_BRANCHABLE_MESSAGES', + 'No conversation to branch', + ) + } + + const activeChain = buildConversationChain(transcript.messages, activeLeaf) + .filter((message) => !message.isSidechain) + + if (activeChain.length === 0) { + throw new SessionBranchingError( + 'NO_BRANCHABLE_MESSAGES', + 'No conversation to branch', + ) + } + + const targetIndex = targetMessageId + ? activeChain.findIndex((message) => message.uuid === targetMessageId) + : activeChain.length - 1 + + if (targetIndex < 0) { + throw new SessionBranchingError( + 'INVALID_TARGET', + 'targetMessageId must reference a main conversation message in the active chain', + ) + } + + const copiedMessages = activeChain.slice(0, targetIndex + 1) + if (copiedMessages.length === 0) { + throw new SessionBranchingError( + 'NO_BRANCHABLE_MESSAGES', + 'No messages to branch', + ) + } + + const copiedMessageIds = new Set(copiedMessages.map((message) => message.uuid)) + const sourceMessageEntriesById = new Map( + sourceEntries + .filter( + (entry): entry is TranscriptMessage => + isTranscriptEntry(entry) && + entry.isSidechain !== true && + copiedMessageIds.has(entry.uuid), + ) + .map((entry) => [entry.uuid, entry]), + ) + const branchMessageEntries = copiedMessages.flatMap((message) => { + const entry = sourceMessageEntriesById.get(message.uuid) + return entry ? [entry] : [] + }) + + if (branchMessageEntries.length === 0) { + throw new SessionBranchingError( + 'NO_BRANCHABLE_MESSAGES', + 'No messages to branch', + ) + } + + const copiedToolResultIds = new Set( + copiedMessages.flatMap((message) => extractToolResultIds(message)), + ) + const contentReplacementRecords = ( + transcript.contentReplacements.get(sourceSessionId as UUID) ?? [] + ).filter((record) => copiedToolResultIds.has(record.toolUseId)) + + const forkSessionId = randomUUID() as UUID + const forkPath = path.join(projectDirPath, `${forkSessionId}.jsonl`) + await mkdir(projectDirPath, { recursive: true, mode: 0o700 }) + + const serializedMessages: SerializedMessage[] = [] + const messageLines: string[] = [] + let parentUuid: UUID | null = null + + for (const entry of branchMessageEntries) { + const forkedEntry: TranscriptEntry = { + ...entry, + sessionId: forkSessionId, + parentUuid, + isSidechain: false, + forkedFrom: { + sessionId: sourceSessionId, + messageUuid: entry.uuid, + }, + } + + serializedMessages.push({ + ...entry, + sessionId: forkSessionId, + }) + messageLines.push(jsonStringify(forkedEntry)) + parentUuid = entry.uuid + } + + const firstPrompt = deriveFirstPrompt( + serializedMessages.find( + (message): message is Extract => + message.type === 'user', + ), + ) + const effectiveTitle = await getUniqueForkName(title ?? firstPrompt, projectDirPath) + + let metadataEntries = buildPreservedMetadataEntries( + sourceEntries, + copiedMessageIds, + forkSessionId, + ) + metadataEntries = ensureSyntheticSessionMeta( + metadataEntries, + sourceWorkDir ?? null, + sourceRepository, + ) + metadataEntries = ensureSyntheticWorktreeState( + metadataEntries, + forkSessionId, + sourceWorktreeSession, + ) + + const lines = [ + ...metadataEntries.map((entry) => jsonStringify(entry)), + ...messageLines, + ] + + if (contentReplacementRecords.length > 0) { + lines.push(jsonStringify({ + type: 'content-replacement', + sessionId: forkSessionId, + replacements: contentReplacementRecords, + } satisfies ContentReplacementEntry)) + } + + lines.push(jsonStringify({ + type: 'custom-title', + sessionId: forkSessionId, + customTitle: effectiveTitle, + })) + + await writeFile(forkPath, `${lines.join('\n')}\n`, { + encoding: 'utf8', + mode: 0o600, + }) + + const workDirFromMetadata = + ( + metadataEntries.find(isSessionMetaEntry)?.workDir as string | undefined + ) ?? + sourceWorkDir ?? + branchMessageEntries.findLast((entry) => typeof entry.cwd === 'string')?.cwd ?? + null + + return { + sessionId: forkSessionId, + title: effectiveTitle, + forkPath, + workDir: workDirFromMetadata, + serializedMessages, + contentReplacementRecords, + } +}