mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-16 13:03:31 +08:00
Enable desktop branching from transcript messages
Desktop users need the same branch-from-here workflow that the CLI already exposed, so the branch creation logic now lives in a shared transcript utility and the desktop app routes completed message actions through the server API. The UI hydrates transcript ids after live completions so newly generated turns can be branched immediately without a refresh. Constraint: Source sessions must remain unmodified while branch sessions inherit the active transcript chain and persistence metadata. Rejected: Keep a desktop-only branch implementation | it would drift from CLI /branch semantics and duplicate transcript filtering rules Confidence: high Scope-risk: moderate Directive: Do not remove the post-completion transcript hydration without a real-model desktop E2E for just-finished messages Tested: bun run verify; Chrome Web UI E2E with real gpt-5.5 provider on ports 45678/45679 Not-tested: Provider-specific behavior beyond the configured Sub2API-ChatGPT route
This commit is contained in:
程序员阿江(Relakkes)
parent
e1c2f54f14
commit
73e915ac6b
42
desktop/src/api/sessions.test.ts
Normal file
42
desktop/src/api/sessions.test.ts
Normal file
@ -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',
|
||||
}),
|
||||
})
|
||||
})
|
||||
})
|
||||
@ -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<CreateSessionResponse>('/api/sessions', body)
|
||||
},
|
||||
|
||||
branch(sessionId: string, body: BranchSessionRequest) {
|
||||
return api.post<BranchSessionResponse>(`/api/sessions/${sessionId}/branch`, body)
|
||||
},
|
||||
|
||||
delete(sessionId: string) {
|
||||
return api.delete<{ ok: true }>(`/api/sessions/${sessionId}`)
|
||||
},
|
||||
|
||||
@ -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(<MessageList />)
|
||||
|
||||
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(<MessageList />)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@ -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<UIMessage, { type: 'user_text' }>['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 (
|
||||
<div
|
||||
className={`pointer-events-none relative z-10 -mt-3 h-0 w-full opacity-0 transition-opacity duration-200 group-hover:opacity-100 group-focus-within:opacity-100 ${
|
||||
align === 'user' ? 'flex justify-end pr-1' : 'flex justify-start pl-1'
|
||||
}`}
|
||||
>
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClick}
|
||||
disabled={loading}
|
||||
aria-label={label}
|
||||
title={label}
|
||||
className="pointer-events-auto inline-flex h-7 w-7 items-center justify-center rounded-full border border-[var(--color-border)]/70 bg-[var(--color-surface-container-low)] text-[var(--color-text-tertiary)] transition-colors hover:border-[var(--color-brand)]/35 hover:text-[var(--color-text-primary)] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-[var(--color-brand)]/35 disabled:cursor-wait disabled:opacity-60"
|
||||
>
|
||||
<GitBranch size={13} strokeWidth={2.2} aria-hidden="true" />
|
||||
</button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
function appendChildToolCall(
|
||||
childToolCallsByParent: Map<string, ToolCall[]>,
|
||||
parentToolUseId: string,
|
||||
@ -450,6 +487,48 @@ function isTurnResponseMessage(message: UIMessage) {
|
||||
)
|
||||
}
|
||||
|
||||
function getBranchableMessageTargets(messages: UIMessage[]): Map<string, BranchableMessageTarget> {
|
||||
const branchableTargets = new Map<string, BranchableMessageTarget>()
|
||||
let currentTurnCandidates: Array<Extract<UIMessage, { type: 'user_text' | 'assistant_text' }>> = []
|
||||
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<string | null>(null)
|
||||
const [turnActionErrors, setTurnActionErrors] = useState<Record<string, string>>({})
|
||||
const [isLoadingTurnChangeCards, setIsLoadingTurnChangeCards] = useState(false)
|
||||
const [branchingMessageId, setBranchingMessageId] = useState<string | null>(null)
|
||||
const [rewindingTurnId, setRewindingTurnId] = useState<string | null>(null)
|
||||
const [turnUndoConfirmTargetId, setTurnUndoConfirmTargetId] = useState<string | null>(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<string, BranchableMessageTarget>() : 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 (
|
||||
<div className="relative min-h-0 flex-1">
|
||||
<div
|
||||
@ -1024,6 +1140,17 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = {
|
||||
})()
|
||||
: null
|
||||
}
|
||||
branchAction={
|
||||
(() => {
|
||||
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<string, AgentTaskNotification>
|
||||
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}
|
||||
>
|
||||
<UserMessage
|
||||
content={message.content}
|
||||
attachments={message.attachments}
|
||||
/>
|
||||
<div className="group">
|
||||
<UserMessage
|
||||
content={message.content}
|
||||
attachments={message.attachments}
|
||||
/>
|
||||
{branchAction ? (
|
||||
<MessageBranchAction
|
||||
align="user"
|
||||
label={branchAction.label}
|
||||
loading={branchAction.loading}
|
||||
onClick={branchAction.onBranch}
|
||||
/>
|
||||
) : null}
|
||||
</div>
|
||||
</SelectableChatMessage>
|
||||
)
|
||||
case 'assistant_text':
|
||||
@ -1144,7 +1287,17 @@ export const MessageBlock = memo(function MessageBlock({
|
||||
role="assistant"
|
||||
content={message.content}
|
||||
>
|
||||
<AssistantMessage content={message.content} />
|
||||
<div className="group">
|
||||
<AssistantMessage content={message.content} />
|
||||
{branchAction ? (
|
||||
<MessageBranchAction
|
||||
align="assistant"
|
||||
label={branchAction.label}
|
||||
loading={branchAction.loading}
|
||||
onClick={branchAction.onBranch}
|
||||
/>
|
||||
) : null}
|
||||
</div>
|
||||
</SelectableChatMessage>
|
||||
)
|
||||
case 'thinking':
|
||||
|
||||
@ -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',
|
||||
|
||||
@ -953,6 +953,9 @@ export const zh: Record<TranslationKey, string> = {
|
||||
'chat.workspaceReferencesOnly': '已添加 {count} 个工作区引用',
|
||||
'chat.contextReferencesOnly': '已添加 {count} 个引用',
|
||||
'chat.addSelectionToChat': '添加到对话',
|
||||
'chat.branchFromHere': '从这里派生',
|
||||
'chat.branchSuccess': '已创建派生会话“{title}”。',
|
||||
'chat.branchError': '从该消息派生失败。详情：{detail}',
|
||||
'chat.userMessageReference': '用户消息',
|
||||
'chat.assistantMessageReference': 'AI 回复',
|
||||
'chat.slashCommands': '斜杠命令',
|
||||
|
||||
@ -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[] = [
|
||||
{
|
||||
|
||||
@ -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<UIMessage, { type: 'user_text' | 'assistant_text' }> =>
|
||||
(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<ChatStore>((set, get) => ({
|
||||
s.backgroundAgentTasks ?? {},
|
||||
restoredBackgroundTasks,
|
||||
),
|
||||
messages: mergeRestoredTerminalGoalEvents(
|
||||
messages: mergeRestoredHistoryIntoLiveMessages(
|
||||
mergeBackgroundTaskMessages(s.messages, restoredBackgroundTasks),
|
||||
uiMessages,
|
||||
),
|
||||
@ -1072,6 +1169,7 @@ export const useChatStore = create<ChatStore>((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,
|
||||
|
||||
@ -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<T>() {
|
||||
|
||||
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',
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@ -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<BranchSessionResponse, 'sessionId' | 'title' | 'workDir'>
|
||||
|
||||
type SessionStore = {
|
||||
sessions: SessionListItem[]
|
||||
activeSessionId: string | null
|
||||
@ -19,6 +26,11 @@ type SessionStore = {
|
||||
|
||||
fetchSessions: (project?: string) => Promise<void>
|
||||
createSession: (workDir?: string, options?: CreateSessionOptions) => Promise<string>
|
||||
branchSession: (
|
||||
sourceSessionId: string,
|
||||
targetMessageId: string,
|
||||
options?: { title?: string },
|
||||
) => Promise<BranchSessionResult>
|
||||
deleteSession: (id: string) => Promise<void>
|
||||
deleteSessions: (ids: string[]) => Promise<BatchDeleteSessionsResponse>
|
||||
enterBatchMode: () => void
|
||||
@ -96,6 +108,43 @@ export const useSessionStore = create<SessionStore>((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)
|
||||
|
||||
@ -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 }
|
||||
|
||||
@ -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<SerializedMessage, { type: 'user' }> | 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<Entry>(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<string> {
|
||||
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<number>([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<React.ReactNode> {
|
||||
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<typeof serializedMessages[number], { type: 'user' }> =>
|
||||
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}`,
|
||||
)
|
||||
|
||||
@ -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<string, unknown> {
|
||||
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<string, unknown> {
|
||||
return {
|
||||
parentUuid: null,
|
||||
@ -306,6 +335,17 @@ function makeWorktreeStateEntry(
|
||||
}
|
||||
}
|
||||
|
||||
function makeContentReplacementEntry(
|
||||
sessionId: string,
|
||||
replacements: Array<{ kind: 'tool-result'; toolUseId: string; replacement: string }>,
|
||||
): Record<string, unknown> {
|
||||
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<Record<string, unknown>>(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()
|
||||
|
||||
@ -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<Response>
|
||||
return Response.json(result)
|
||||
}
|
||||
|
||||
async function branchSession(req: Request, sessionId: string): Promise<Response> {
|
||||
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<Response> {
|
||||
const checkpoints = await listSessionTurnCheckpoints(sessionId)
|
||||
return Response.json({ checkpoints })
|
||||
|
||||
28
src/utils/__tests__/sessionBranching.test.ts
Normal file
28
src/utils/__tests__/sessionBranching.test.ts
Normal file
@ -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')
|
||||
})
|
||||
})
|
||||
554
src/utils/sessionBranching.ts
Normal file
554
src/utils/sessionBranching.ts
Normal file
@ -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<string, unknown>
|
||||
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<string, unknown>
|
||||
|
||||
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<UUID, TranscriptMessage>,
|
||||
leafUuids: Set<UUID>,
|
||||
): 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<RawEntry[]> {
|
||||
const content = await readFile(filePath)
|
||||
return parseJSONL<RawEntry>(content)
|
||||
}
|
||||
|
||||
async function listExistingTitles(projectDirPath: string): Promise<string[]> {
|
||||
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<string, unknown>).customTitle === 'string'
|
||||
) {
|
||||
titles.push((entry as Record<string, unknown>).customTitle as string)
|
||||
break
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
return titles
|
||||
}
|
||||
|
||||
export function deriveFirstPrompt(
|
||||
firstUserMessage: Extract<SerializedMessage, { type: 'user' }> | 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<string> {
|
||||
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<number>([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<string>,
|
||||
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<SessionBranchResult> {
|
||||
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<SerializedMessage, { type: 'user' }> =>
|
||||
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,
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user