mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-17 13:13:35 +08:00
Merge desktop branch-from-here support into main
Bring the verified worktree implementation into local main while preserving the existing local main commits and unrelated dirty working tree changes. Constraint: Local main already had unrelated uncommitted provider/runtime edits, so the merge only touches the desktop branch feature paths. Confidence: high Scope-risk: moderate Tested: bun run verify before merge; merge completed without conflicts Not-tested: Post-merge full quality gate on dirty local main
This commit is contained in:
程序员阿江(Relakkes)
commit
fed9b6c9f6
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